"""Interactive shells, one per agent chat, held open behind the panel. The other half of `ssh.py`. There a connection lives for one command, because a runner is a request-and-answer and holding state would be the wrong shape. Here the connection *is* the state: a PTY with a shell on the far side, its scrollback, and whoever is currently watching it. Shaped after `services/generation.py` -- a registry, a background task that owns the work, and a socket that merely follows it -- and it differs in three ways worth knowing: * **Keyed on the chat, not on a session of its own.** A reload is indistinguishable from a second tab, so anything finer needs an id in the browser's storage, and then an abandoned tab leaks a shell nothing in the UI can find. One chat, one shell. Two tabs share it, like `tmux attach` twice, which is the only reading under which "it is still there when you come back" means anything. They also share a size, and the smaller one wins. * **Nothing here ends by itself.** A generation finishes, so `generation.ensure` can prune inside itself. A shell sits at a prompt forever and nothing calls in again, so there is a reaper task instead. Copying the generation shape here would mean nothing was ever swept. * **A slow reader is dropped, not buffered.** Every viewer has a bounded queue; one that fills is disconnected and reattaches with the scrollback. Blocking the pump instead would stall every other viewer and buffer without bound inside the server -- and `yes` is one word to type. What a person types here is deliberately not run past `agent/policy.py`. The modes and the two lists govern a *model*, which reads untrusted pages and files and can be talked into things. Somebody at a keyboard holds the credential already and could open the same shell with an ssh client; asking them to approve their own keystrokes would be theatre. """ from __future__ import annotations import asyncio import contextlib import logging import time import uuid from collections import deque from dataclasses import dataclass, field from typing import Any from lembas.services.agent import capture, shell_marks from lembas.services.agent.base import ExecError from lembas.services.agent.ssh import available, connect_kwargs log = logging.getLogger(__name__) # What one shell keeps to hand a returning viewer. Bytes rather than lines: a # line budget is dishonest about a program that writes one very long line. SCROLLBACK_BYTES = 256 * 1024 # How much output one viewer may fall behind by before it is dropped. Frames # are whatever the far side wrote, so this is generous in wall-clock terms and # only reached by a browser that has genuinely stopped reading. VIEWER_QUEUE = 512 # Read size. Large enough that `cat` of a big file is not a million wakeups, # small enough that a prompt appears the instant it is written. READ_BYTES = 64 * 1024 # A terminal nobody has ever heard of gets no colours; this one every shell # knows and it is what an ordinary ssh client announces. TERM_TYPE = "xterm-256color" # A size is a number the browser sends. `change_terminal_size(100000, 100000)` # is a way to ask the far side to allocate. MAX_COLS = 500 MAX_ROWS = 300 MIN_COLS = 20 MIN_ROWS = 5 # How long a closed session stays in the registry. A tab attaching a second # after the shell exited should be told what happened rather than silently # handed a fresh one. KEEP_CLOSED = 60.0 # How often the reaper looks. Nothing here is urgent: the idle timeout is # measured in minutes. REAP_INTERVAL = 30.0 # Why a session ended. The browser is told, and the wording differs enough to be # worth the constants. CLOSED_EXITED = "exited" CLOSED_IDLE = "idle" CLOSED_SHUTDOWN = "shutdown" CLOSED_REVOKED = "revoked" CLOSED_ERROR = "error" # Whether this shell tells us where its commands begin and end. # live -- it does # loading -- the hooks went in; the first prompt has not arrived yet # none -- it never will: an unknown shell, or a dotfile that replaced it INTEGRATION_LIVE = "live" INTEGRATION_LOADING = "loading" INTEGRATION_NONE = "none" # How long a shell may produce output without ever marking a prompt before we # conclude it is not going to. This is what catches a `.bashrc` ending in `exec # tmux`: the hooks were installed and then the shell replaced itself. Without # it the buttons stay greyed out forever with no explanation. INTEGRATION_GRACE = 10.0 @dataclass class Viewer: """One browser watching one shell.""" cols: int = 80 rows: int = 24 id: str = field(default_factory=lambda: uuid.uuid4().hex) queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=VIEWER_QUEUE)) # Everything the shell has said so far, handed over in the same synchronous # call that subscribes. Reading the buffer and subscribing as two awaits # loses whatever arrives between them. snapshot: bytes = b"" # Set when this viewer fell behind. It is woken with the sentinel below and # told to reconnect, which costs it nothing: the scrollback is the state. dropped: bool = False class Session: """A shell on the far side of one chat's connection.""" def __init__( self, chat_id: str, *, owner_id: str, profile_id: str, label: str, project_dir: str = "", idle_timeout: float = 1800.0, cols: int = 80, rows: int = 24, integrate: bool = True, ) -> None: self.chat_id = chat_id self.owner_id = owner_id self.profile_id = profile_id self.label = label self.project_dir = project_dir self.idle_timeout = idle_timeout self.viewers: dict[str, Viewer] = {} self._scrollback: deque[bytes] = deque() self._scrollback_bytes = 0 # --- Command boundaries --------------------------------------------- # Three states, and the middle one matters: INTEGRATION_LOADING means # the hooks were installed and no marker has arrived yet, which is a # different thing to tell somebody than "this shell will never mark". self.integrate = integrate self.integration = INTEGRATION_LOADING if integrate else INTEGRATION_NONE self.shell = "" self._marks = shell_marks.Marks(self._on_mark, self._on_text) # At most two, ever. The one being written and the last finished one -- # a history would be a second scrollback with none of the bounding. self.current: capture.Capture | None = None self.last: capture.Capture | None = None self._captures = 0 # Where the shell says it is, which the panel header shows live and a # capture records. Seeded from the chat so it says something sensible # before the first prompt. self.cwd = project_dir # Set by the socket layer, which knows how to shape a frame. Called # when a command finishes so a panel can enable its buttons without # polling for something that happens a few times a minute. self.on_command: Any = None self._conn: Any = None self._process: Any = None self._pump: asyncio.Task | None = None self.closed = False self.closed_reason = "" self.closed_at = 0.0 self.started_at = time.monotonic() # Bumped by a keystroke and by a viewer coming or going. Idle is this # going quiet *with nobody attached*: a build running behind a closed # panel is the case this whole lifetime exists for. self.last_active = time.monotonic() # The size the PTY is *created* with, which matters: a shell prints its # prompt before anything could resize it, and a prompt drawn at 80 # columns inside a 140-column window stays wrong until the next one. self.cols = _clamp(cols, MIN_COLS, MAX_COLS) self.rows = _clamp(rows, MIN_ROWS, MAX_ROWS) # --- Opening ------------------------------------------------------------- async def start(self, spec: dict[str, Any]) -> None: """Connect, ask for a PTY, and start pumping what it says. The credential is used here and not kept. `spec` is a decrypted snapshot of a profile and the connection outlives the request that made it, so holding a private key in memory for the hour a shell sits at a prompt buys nothing. """ if problem := available(): raise ExecError(problem) import asyncssh try: self._conn = await asyncssh.connect( spec["host"], **connect_kwargs(spec), # asyncssh sends no keepalives by default. `SshExecutor` never # needed them because its connections live for one command; a # shell held open behind NAT otherwise gets dropped with no FIN # and no exception, and the pump simply never returns -- the # panel looks alive and answers nothing. keepalive_interval=30, keepalive_count_max=3, ) self._process = await self._conn.create_process( self._command(), term_type=TERM_TYPE, term_size=(self.cols, self.rows), # Bytes in both directions. A read lands mid-character often # enough to matter, and the browser's decoder is stateful across # writes while a per-frame decode here is not: it would corrupt # every boundary. Nothing decodes, so nothing can split. encoding=None, stderr=asyncssh.STDOUT, ) except asyncssh.HostKeyNotVerifiable as exc: await self._teardown() raise ExecError( f"{self.label} presented a different host key than the one that " "was confirmed. Nothing was sent." ) from exc except asyncssh.PermissionDenied as exc: await self._teardown() raise ExecError(f"{self.label} refused the credential.") from exc except (OSError, asyncssh.Error) as exc: await self._teardown() raise ExecError(f"Could not reach {self.label}: {exc}") from exc self._pump = asyncio.create_task(self._read_forever()) def _command(self) -> str | None: """What the PTY runs, or None for the account's plain login shell. A shell has no notion of "start here" that SSH can carry, so the chat's project directory has to be a `cd` -- run before the shell rather than typed into it, so the scrollback opens on a prompt instead of on a command nobody entered. It is single-quoted, and a failure is ignored: a directory that has been deleted should leave somebody at a shell to find out why, not with a connection that closes as it opens. With integration on, the same string also writes the shell-integration files and execs through them; see `shell_marks` for why it is done in the command rather than over SFTP or through the environment. """ return shell_marks.command_for(self.project_dir, integrate=self.integrate) # --- Where one command ends and the next begins -------------------------- def _on_mark(self, kind: str, value: str) -> None: """One marker, from the scanner in the pump. Advisory, never trusted: a program can print these itself and move a boundary. It is not a way in -- the text is sanitised and fenced either way, and a program could already print anything on screen -- but that is why nothing here validates them, and why none of it decides anything a person could not already do at the keyboard. """ if kind == shell_marks.MARK_READY: self.integration = INTEGRATION_LIVE self.shell = value.split(";")[0][:32] return if self.integration != INTEGRATION_LIVE and kind in ( shell_marks.MARK_PROMPT, shell_marks.MARK_OUTPUT, ): self.integration = INTEGRATION_LIVE if kind == shell_marks.MARK_CWD: path = shell_marks.unescape(value.partition("=")[2])[:1000] self.cwd = path return if kind == shell_marks.MARK_COMMAND: self._captures += 1 self.current = capture.Capture( seq=self._captures, command=capture.trim_command(shell_marks.unescape(value)), cwd=self.cwd, ) return if kind == shell_marks.MARK_DONE and self.current is not None: try: self.current.exit_status = int(value.strip() or 0) except ValueError: self.current.exit_status = 0 self.current.ended = time.monotonic() self.last = self.current self.current = None if self.on_command is not None: self.on_command(self.last) def _on_text(self, data: bytes) -> None: """Everything that was not a marker, while a command is running. Interleaved with `_on_mark` rather than applied to the whole chunk afterwards: a shell often writes the command marker, the output and the finished marker in one read, and absorbing after the scan would find the capture already closed and keep nothing at all. """ if self.current is not None: self.current.absorb(data) def latest(self) -> capture.Capture | None: """The command to act on: the one still running, else the last one. In-flight counts. "Copy the last command and its output" while `make` is still going should give what has been printed so far, marked as still running -- not "nothing yet". """ return self.current or self.last # --- Following ----------------------------------------------------------- def attach(self, cols: int = 80, rows: int = 24) -> Viewer: """Subscribe, and take the scrollback, in one synchronous step. One step on purpose: reading the buffer and subscribing as two awaits loses whatever the shell says between them, which is exactly the moment somebody reattaches to a build that is still writing. """ viewer = Viewer( cols=_clamp(cols, MIN_COLS, MAX_COLS), rows=_clamp(rows, MIN_ROWS, MAX_ROWS), ) viewer.snapshot = b"".join(self._scrollback) self.viewers[viewer.id] = viewer self.last_active = time.monotonic() self.apply_size() return viewer def detach(self, viewer: Viewer) -> None: self.viewers.pop(viewer.id, None) # Counts as activity: the timeout is "nobody has been here and nothing # has happened for a while", so it starts when the last viewer leaves. self.last_active = time.monotonic() self.apply_size() async def send(self, data: bytes) -> None: """Type into the shell.""" if self.closed or self._process is None: return self.last_active = time.monotonic() try: self._process.stdin.write(data) except (BrokenPipeError, ConnectionResetError, OSError): await self._finish(CLOSED_EXITED) def resize(self, viewer: Viewer, cols: int, rows: int) -> None: """Record this viewer's size and give the far side the smallest. Two tabs on one PTY cannot each have their own geometry. The smaller wins in both directions, so nothing is drawn off the edge of the smaller window -- the larger one gets an unused margin, which is the harmless half of the trade. """ viewer.cols = _clamp(cols, MIN_COLS, MAX_COLS) viewer.rows = _clamp(rows, MIN_ROWS, MAX_ROWS) self.apply_size() def apply_size(self) -> None: """Synchronous: `change_terminal_size` only queues a window-change message, so there is nothing to await and no reason to make every caller a coroutine.""" if self.closed or self._process is None or not self.viewers: return cols = min(v.cols for v in self.viewers.values()) rows = min(v.rows for v in self.viewers.values()) if (cols, rows) == (self.cols, self.rows): return self.cols, self.rows = cols, rows with contextlib.suppress(Exception): self._process.change_terminal_size(cols, rows) # --- The pump ------------------------------------------------------------ async def _read_forever(self) -> None: assert self._process is not None try: while True: data = await self._process.stdout.read(READ_BYTES) if not data: break self._remember(data) # Before the fan-out, so a "this command finished" frame can # never reach a browser after the output it describes. self._observe(data) self._fan_out(data) except asyncio.CancelledError: raise except Exception as exc: # noqa: BLE001 - one shell dying is not a crash log.info("terminal %s ended: %s", self.chat_id, exc) await self._finish(CLOSED_ERROR) return await self._finish(CLOSED_EXITED) def _observe(self, data: bytes) -> None: """Watch the stream for markers, and feed the command being captured. Server-side rather than in the browser, for five reasons. The `behind` path calls `term.reset()` and replays a *truncated* scrollback, so a client parser routinely sees a "finished" with no matching "started". Two tabs share one shell and two parsers can disagree about what "the last command" is. The server sees the stream once however many are watching. And what comes out of this ends up inside a prompt -- deriving it here means there is nothing to disbelieve later. The bytes are still fanned out unchanged, markers and all: xterm consumes an OSC it has no handler for and never draws it, and rewriting frames on the hot path would break the "nothing decodes, so nothing can split" property the pump depends on. """ self._marks.feed(data) if self.current is None and ( self.integration == INTEGRATION_LOADING and time.monotonic() - self.started_at > INTEGRATION_GRACE ): # Output arrived, the grace period passed, and no marker ever came. # Output arrived, the grace period passed, and no marker ever came. # Something replaced the shell -- a dotfile ending in `exec tmux` is # the usual one. Say so rather than leaving the buttons greyed. self.integration = INTEGRATION_NONE def _remember(self, data: bytes) -> None: self._scrollback.append(data) self._scrollback_bytes += len(data) while self._scrollback_bytes > SCROLLBACK_BYTES and len(self._scrollback) > 1: self._scrollback_bytes -= len(self._scrollback.popleft()) def announce(self, text: str) -> None: """Put one text frame in front of every viewer. Through the same queues as the output so ordering is preserved: a "finished" that overtook the last of the output it describes would have a panel offering a capture the screen has not caught up with. Dropped rather than blocking on a full queue -- that viewer is already being disconnected and will be told again on reattach. """ for viewer in list(self.viewers.values()): with contextlib.suppress(asyncio.QueueFull): viewer.queue.put_nowait(text) def _fan_out(self, data: bytes) -> None: for viewer in list(self.viewers.values()): try: viewer.queue.put_nowait(data) except asyncio.QueueFull: # Emptied first so the sentinel fits and so the socket does not # spend its last moments writing frames nobody will see. _drain(viewer.queue) viewer.dropped = True with contextlib.suppress(asyncio.QueueFull): viewer.queue.put_nowait(None) self.viewers.pop(viewer.id, None) # --- Closing ------------------------------------------------------------- async def close(self, reason: str = CLOSED_SHUTDOWN) -> None: pump, self._pump = self._pump, None if pump is not None: pump.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await pump await self._finish(reason) async def _finish(self, reason: str) -> None: """Mark this session over and wake everybody watching. Called from the pump when the shell exits, and from `close` after the pump has been cancelled -- which is why it does not cancel the pump itself. The entry stays in the registry for KEEP_CLOSED so a late attachment gets an explanation. """ if self.closed: return self.closed = True self.closed_reason = reason self.closed_at = time.monotonic() for viewer in list(self.viewers.values()): with contextlib.suppress(asyncio.QueueFull): viewer.queue.put_nowait(None) await self._teardown() log.info( "terminal closed chat=%s owner=%s profile=%s reason=%s after=%.0fs", self.chat_id, self.owner_id, self.profile_id, reason, time.monotonic() - self.started_at, ) async def _teardown(self) -> None: process, self._process = self._process, None conn, self._conn = self._conn, None if process is not None: with contextlib.suppress(Exception): process.terminate() if conn is not None: with contextlib.suppress(Exception): conn.close() with contextlib.suppress(Exception): await conn.wait_closed() @property def idle_for(self) -> float: if self.viewers: return 0.0 return time.monotonic() - self.last_active # --- The registry ------------------------------------------------------------ _SESSIONS: dict[str, Session] = {} _REAPER: asyncio.Task | None = None def get(chat_id: str) -> Session | None: """The live session for a chat, if there is one. Closed ones do not count.""" session = _SESSIONS.get(chat_id) if session is None or session.closed: return None return session def peek(chat_id: str) -> Session | None: """As `get`, but a recently closed session too -- it carries the reason.""" return _SESSIONS.get(chat_id) def count() -> int: return sum(1 for s in _SESSIONS.values() if not s.closed) def count_for(owner_id: str) -> int: return sum(1 for s in _SESSIONS.values() if not s.closed and s.owner_id == owner_id) async def open_session( chat_id: str, *, owner_id: str, profile_id: str, label: str, spec: dict[str, Any], project_dir: str = "", idle_timeout: float = 1800.0, max_sessions: int = 20, max_per_user: int = 3, cols: int = 80, rows: int = 24, integrate: bool = True, ) -> Session: """The shell for this chat, opening one if it is not already there. Idempotent for the same reason `generation.ensure` is: a second tab, or the same tab after a reload, must attach to what is running rather than start a second shell on the same machine. """ existing = get(chat_id) if existing is not None: return existing _reap() if count() >= max_sessions: raise ExecError( "This instance already has as many terminals open as it allows. " "Close one, or ask an administrator to raise the limit." ) if count_for(owner_id) >= max_per_user: raise ExecError( f"You already have {max_per_user} terminal" f"{'' if max_per_user == 1 else 's'} open. Close one first." ) session = Session( chat_id, owner_id=owner_id, profile_id=profile_id, label=label, integrate=integrate, project_dir=project_dir, idle_timeout=idle_timeout, cols=cols, rows=rows, ) await session.start(spec) _SESSIONS[chat_id] = session _ensure_reaper() log.info( "terminal opened chat=%s owner=%s profile=%s host=%s dir=%s", chat_id, owner_id, profile_id, label, project_dir or "~", ) return session async def close_chat(chat_id: str, reason: str = CLOSED_REVOKED) -> bool: session = _SESSIONS.pop(chat_id, None) if session is None: return False await session.close(reason) return True async def close_for_profile(profile_id: str) -> int: """End every shell opened on one connection. `session.profile_for` re-checks the profile on every reply, so deleting or disabling one stops the model at once. A terminal resolves the profile when it opens and then holds the connection, so without this "I disabled that connection" would simply not be true of the shell already on screen. """ doomed = [s for s in _SESSIONS.values() if s.profile_id == profile_id and not s.closed] for session in doomed: _SESSIONS.pop(session.chat_id, None) await session.close(CLOSED_REVOKED) return len(doomed) async def close_for_owner(owner_id: str) -> int: doomed = [s for s in _SESSIONS.values() if s.owner_id == owner_id and not s.closed] for session in doomed: _SESSIONS.pop(session.chat_id, None) await session.close(CLOSED_REVOKED) return len(doomed) async def shutdown() -> None: """End every shell. Called from the lifespan, beside stop_generations.""" global _REAPER reaper, _REAPER = _REAPER, None if reaper is not None: reaper.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await reaper for session in list(_SESSIONS.values()): await session.close(CLOSED_SHUTDOWN) _SESSIONS.clear() def _reap() -> None: """Drop sessions that have been closed long enough to stop explaining.""" now = time.monotonic() for chat_id, session in list(_SESSIONS.items()): if session.closed and now - session.closed_at > KEEP_CLOSED: _SESSIONS.pop(chat_id, None) def _ensure_reaper() -> None: global _REAPER if _REAPER is None or _REAPER.done(): _REAPER = asyncio.create_task(_reaper_loop()) async def _reaper_loop() -> None: """Close idle shells, then forget closed ones. A task rather than a sweep inside `open_session`, which is the shape `generation` uses. That works there because a generation ends on its own and something calls in again; a shell at a prompt does neither, so a lazy sweep would run only when somebody opened the *next* terminal. """ while True: try: await asyncio.sleep(REAP_INTERVAL) for session in list(_SESSIONS.values()): if not session.closed and session.idle_for > session.idle_timeout: _SESSIONS.pop(session.chat_id, None) await session.close(CLOSED_IDLE) _reap() except asyncio.CancelledError: raise except Exception: # noqa: BLE001 - the reaper must outlive one bad sweep log.exception("the terminal reaper raised") def _drain(queue: asyncio.Queue) -> None: while True: try: queue.get_nowait() except asyncio.QueueEmpty: return def _clamp(value: int, low: int, high: int) -> int: try: number = int(value) except (TypeError, ValueError): return low return min(max(number, low), high) __all__ = [ "CLOSED_EXITED", "CLOSED_IDLE", "CLOSED_REVOKED", "CLOSED_SHUTDOWN", "Session", "Viewer", "close_chat", "close_for_owner", "close_for_profile", "count", "count_for", "get", "open_session", "peek", "shutdown", ]