"""The socket behind the terminal panel. A WebSocket rather than SSE, because SSE is one-directional and a terminal is not: keystrokes have to go up, and an HTTP round trip per keypress is not a terminal. It is the only WebSocket in LLeMbas, and it is worth saying what that costs -- a cross-site page that could reach this endpoint would have a shell on somebody's machine, not merely a copy of their chat. So there are two locks on the door, and this module is mostly them. **Where a refusal happens is load-bearing.** A browser tells a page nothing about a handshake that *failed*: `new WebSocket()` fires `error` with no status and no reason. So the socket is accepted first and the reason sent as a frame for everything a person could act on -- no permission, the connection is disabled, its host key was never confirmed -- and refused before accepting only for the two cases where accepting is itself the risk. It holds no database session. A dependency would keep one open for the hour a shell sits at a prompt; `session_scope()` opens one for the authorisation and closes it, exactly as `generation._run` does. """ from __future__ import annotations import asyncio import contextlib import json import logging from urllib.parse import urlsplit from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect, status from lembas.api.deps import Db, RequiredUser from lembas.db.models import KIND_AGENT, Chat from lembas.db.session import session_scope from lembas.security import permissions from lembas.security.sessions import COOKIE_NAME, resolve_session from lembas.services import settings_store from lembas.services.agent import draft as draft_service from lembas.services.agent import session as agent_session from lembas.services.agent import terminal as terminal_service from lembas.services.agent.base import ExecError log = logging.getLogger(__name__) router = APIRouter(prefix="/api/chats", tags=["terminal"]) # Nothing a keyboard produces is anywhere near this. Paste is the only thing # that comes close, and a megabyte pasted into a shell is a mistake either way. MAX_INPUT_BYTES = 256 * 1024 # 1008 is "policy violation", the closest thing the protocol has to "no". CLOSE_POLICY = 1008 # What the far side is told when a shell ends, in words rather than a code. CLOSED_WORDS = { terminal_service.CLOSED_EXITED: "The shell exited.", terminal_service.CLOSED_IDLE: "This terminal was closed after sitting idle.", terminal_service.CLOSED_SHUTDOWN: "LLeMbas restarted, so this shell was closed.", terminal_service.CLOSED_REVOKED: "The connection behind this terminal was closed.", terminal_service.CLOSED_ERROR: "The connection to the machine was lost.", } def _same_origin(websocket: WebSocket) -> bool: """Whether this handshake came from a page served by this site. Required, not merely checked when present. The session cookie is SameSite Lax, which already withholds it from a handshake a foreign page starts, and this is the belt to that brace -- an absent Origin is not a browser, and a non-browser client has no business here. """ origin = websocket.headers.get("origin") host = websocket.headers.get("host") if not origin or not host: return False return urlsplit(origin).netloc.lower() == host.lower() def _chat_or_draft(db, user, chat_id: str): """The chat this panel belongs to, real or still being decided. A draft resolves to a transient `Chat` -- see services/agent/draft.py -- which is what lets the terminal open on the new-chat screen without `_prepare` or `agent_session.resolve` learning that drafts exist. """ if draft_service.is_draft(chat_id): draft = draft_service.get(chat_id, user.id) return draft_service.as_chat(draft) if draft is not None else None chat = db.get(Chat, chat_id) return chat if chat is not None and chat.user_id == user.id else None def _prepare(db, user, chat_id: str) -> tuple[str, dict]: """Everything that has to be true, and what opening needs. One or the other. Returns a message to show, or the arguments for `open_session`. The order is the order somebody would ask the questions in, and every "no" is a sentence rather than a silence. """ if not permissions.has(db, user, "agent.terminal"): return "You do not have permission to open a terminal.", {} chat = _chat_or_draft(db, user, chat_id) if chat is None: return "That chat no longer exists.", {} if chat.kind != KIND_AGENT: return "This is an ordinary chat, so it has no machine to open a shell on.", {} values = settings_store.agents(db) if not values.get("terminal_enabled", True): return "The terminal is switched off on this instance.", {} context = agent_session.resolve(db, chat, user) if context is None: return ( "This chat's connection is not usable: it may have been deleted, " "disabled, or agent chats may be switched off here.", {}, ) return "", { "owner_id": user.id, "profile_id": chat.ssh_profile_id or "", "label": context.label, "spec": context.spec, "project_dir": context.project_dir, "idle_timeout": float(values["terminal_idle_timeout"]), "max_sessions": int(values["terminal_max_sessions"]), "max_per_user": int(values["terminal_max_per_user"]), "integrate": bool(values.get("terminal_integration", True)), } @router.websocket("/{chat_id}/terminal/ws") async def terminal_socket( websocket: WebSocket, chat_id: str, cols: int = 80, rows: int = 24, ) -> None: if not _same_origin(websocket): await websocket.close(code=CLOSE_POLICY) return with session_scope() as db: user = resolve_session(db, websocket.cookies.get(COOKIE_NAME)) if user is None: await websocket.close(code=CLOSE_POLICY) return problem, opening = _prepare(db, user, chat_id) owner_email = user.email await websocket.accept() if problem: await _refuse(websocket, problem) return try: session = await terminal_service.open_session(chat_id, cols=cols, rows=rows, **opening) except ExecError as exc: await _refuse(websocket, str(exc)) return except Exception: # noqa: BLE001 - a failure here is one socket, not the app log.exception("could not open a terminal for %s", owner_email) await _refuse(websocket, "The shell could not be started.") return # Shaping a frame is this layer's job, not the session's; the session only # knows it finished something. Reassigned per socket and harmless: every # socket on this session would build the identical frame. session.on_command = lambda found: session.announce( json.dumps({"t": "command", "command": _command_frame(found)}) ) viewer = session.attach(cols, rows) await websocket.send_text( json.dumps( { "t": "ready", "label": session.label, "dir": session.project_dir, "cols": session.cols, "rows": session.rows, # Two tabs share one shell, and a size neither of them chose is # otherwise a mystery. "shared": len(session.viewers) > 1, # Whether this shell will tell us where commands begin and end, # which is what the Copy and Send buttons are made of. "integration": session.integration, "last": _command_frame(session.latest()), } ) ) if viewer.snapshot: await websocket.send_bytes(viewer.snapshot) downward = asyncio.create_task(_to_browser(websocket, session, viewer)) upward = asyncio.create_task(_from_browser(websocket, session, viewer)) try: await asyncio.wait({downward, upward}, return_when=asyncio.FIRST_COMPLETED) finally: for task in (downward, upward): task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await task # The session is deliberately left running. Closing the panel, or # navigating away, is not "I am finished with this machine" -- a build # carries on and the scrollback is still there on the way back. The # idle timeout is what eventually ends it. session.detach(viewer) @router.get("/{chat_id}/terminal/last") async def last_command(db: Db, user: RequiredUser, chat_id: str) -> dict: """The last command and its output, rendered ready to paste. The *server* renders the text, so Copy and Send are a fetch and a clipboard write with no formatting logic in the browser -- and the block a model eventually reads exists in exactly one place. The panel's own screen buffer could not produce it anyway: it holds what is on screen, hard-wrapped at the terminal's width, with no way to tell a wrap from a newline. """ if _chat_or_draft(db, user, chat_id) is None: raise HTTPException(status.HTTP_404_NOT_FOUND, "That chat no longer exists.") if not permissions.has(db, user, "agent.terminal"): raise HTTPException(status.HTTP_403_FORBIDDEN, "You cannot open a terminal.") session = terminal_service.get(chat_id) found = session.latest() if session is not None else None if session is None or found is None: return { "ok": False, "message": "Nothing has been run in this shell yet." if session is not None else "This terminal is not open.", } return { "ok": True, "command": found.command, "cwd": found.cwd, "exit": found.exit_status, "running": found.running, "summary": found.summary(), "text": found.as_text(label=session.label), } def _command_frame(found) -> dict | None: """A finished command, small enough to push at every viewer. Tens of bytes, and deliberately *not* the output: a 64KB text frame would compete with PTY bytes on the one path that has to stay responsive, and the two buttons are pressed by a person, where a request is the natural shape. """ if found is None: return None return { "seq": found.seq, "command": found.command, "cwd": found.cwd, "exit": found.exit_status, "running": found.running, "ms": found.duration_ms, "summary": found.summary(), } async def _to_browser(websocket: WebSocket, session, viewer) -> None: """Everything the shell says, plus the one frame that says it stopped.""" while True: chunk = await viewer.queue.get() if chunk is None: reason = terminal_service.CLOSED_EXITED if viewer.dropped else session.closed_reason payload = {"t": "closed", "reason": reason, "message": _words(reason)} if viewer.dropped: # Not the session's doing: this browser stopped reading and was # disconnected so the others kept up. Reconnecting costs it # nothing, because the scrollback is the state. payload = {"t": "behind", "message": "Reconnecting: output arrived faster than " "this window could draw it."} with contextlib.suppress(Exception): await websocket.send_text(json.dumps(payload)) return # A string in the queue is a control frame that had to keep its place # in the stream -- see `Session.announce`. if isinstance(chunk, str): await websocket.send_text(chunk) continue await websocket.send_bytes(chunk) async def _from_browser(websocket: WebSocket, session, viewer) -> None: """Keystrokes as binary, everything else as JSON. Binary for the hot path is what makes multi-byte characters safe: a read on the far side lands mid-sequence often enough to matter, and decoding each frame here would corrupt every boundary. Nothing decodes, so nothing splits. """ while True: try: message = await websocket.receive() except WebSocketDisconnect: return if message["type"] == "websocket.disconnect": return data = message.get("bytes") if data is not None: if len(data) > MAX_INPUT_BYTES: continue await session.send(data) continue text = message.get("text") if text: _control(session, viewer, text) def _control(session, viewer, text: str) -> None: try: payload = json.loads(text) except ValueError: return if not isinstance(payload, dict) or payload.get("t") != "resize": return session.resize(viewer, payload.get("cols", 80), payload.get("rows", 24)) def _words(reason: str) -> str: return CLOSED_WORDS.get(reason, "This terminal closed.") async def _refuse(websocket: WebSocket, message: str) -> None: """Say why, then close. Sent as a frame because a browser cannot read a rejected handshake -- the reason would be lost exactly when it is needed.""" with contextlib.suppress(Exception): await websocket.send_text(json.dumps({"t": "error", "message": message})) with contextlib.suppress(Exception): await websocket.close()