feat: implement session management and terminal handling for VM sessions

nest/server.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Any
 
@@ -21,12 +22,29 @@ def create_app(application: Application | None = None, *, host: str | None = Non
     from fastapi.staticfiles import StaticFiles
 
     from .dashboard.router import router as dashboard_router
+    from .session_routes import router as session_router
+    from .sessions import SessionManager
+
+    sessions = SessionManager(api.state.uvm_application)
+    api.state.nest_sessions = sessions
+    previous_lifespan = api.router.lifespan_context
+
+    @asynccontextmanager
+    async def lifespan(app):
+        async with previous_lifespan(app):
+            try:
+                yield
+            finally:
+                await sessions.shutdown()
+
+    api.router.lifespan_context = lifespan
 
     api.title = "Nest"
     api.description = "Nest workspace API, with local Firecracker machine management powered by UVM."
     static_dir = Path(__file__).parent / "dashboard" / "static"
     api.mount("/dashboard/static", StaticFiles(directory=static_dir), name="dashboard_static")
     api.include_router(dashboard_router, prefix="/dashboard")
+    api.include_router(session_router)
 
     @api.get("/")
     async def root() -> Any:

nest/session_routes.py

@@ -0,0 +1,159 @@
+"""VM-scoped session history, creation, and authenticated browser terminal streams."""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import hmac
+import json
+from contextlib import suppress
+from typing import Literal
+from urllib.parse import urlsplit
+
+from fastapi import APIRouter, Depends, Request, WebSocket, WebSocketDisconnect
+from pydantic import BaseModel, ConfigDict, Field, StrictInt
+
+from uvm.app import Application
+from uvm.errors import UvmError
+from uvm.guest_agent import MAX_TERMINAL_FRAME, ProtocolError, terminal_data
+from uvm.routers.dependencies import get_authorized_application
+
+from .sessions import SessionManager
+
+
+router = APIRouter(prefix="/vms/{vm_id}/sessions", tags=["sessions"])
+
+
+class SessionCreate(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    rows: StrictInt = Field(default=24, ge=2, le=200)
+    cols: StrictInt = Field(default=80, ge=2, le=400)
+
+
+class SessionResponse(BaseModel):
+    id: str
+    vm_id: str
+    name: str
+    username: str
+    status: Literal["starting", "active", "closed", "interrupted", "failed"]
+    created_at: float
+    updated_at: float
+    ended_at: float | None
+    rows: int
+    cols: int
+    exit_code: int | None
+    error: str | None
+    output_start: int
+    output_end: int
+    truncated: bool
+
+
+def manager(request: Request, _application: Application = Depends(get_authorized_application)) -> SessionManager:
+    return request.app.state.nest_sessions
+
+
+@router.get("", response_model=list[SessionResponse])
+async def list_sessions(vm_id: str, sessions: SessionManager = Depends(manager)):
+    return await sessions.list(vm_id)
+
+
+@router.post("", response_model=SessionResponse, status_code=201)
+async def create_session(vm_id: str, body: SessionCreate, sessions: SessionManager = Depends(manager)):
+    return await sessions.create(vm_id, body.rows, body.cols)
+
+
+@router.get("/{session_id}", response_model=SessionResponse)
+async def get_session(vm_id: str, session_id: str, sessions: SessionManager = Depends(manager)):
+    return await sessions.get(vm_id, session_id)
+
+
+@router.post("/{session_id}/close", response_model=SessionResponse)
+async def close_session(vm_id: str, session_id: str, sessions: SessionManager = Depends(manager)):
+    return await sessions.end(vm_id, session_id)
+
+
+def _same_origin(websocket: WebSocket) -> bool:
+    origin = websocket.headers.get("origin")
+    if origin is None:
+        return True  # Non-browser API clients must still authenticate.
+    parsed = urlsplit(origin)
+    scheme = "https" if websocket.url.scheme == "wss" else "http"
+    return parsed.scheme == scheme and parsed.netloc.lower() == websocket.headers.get("host", "").lower() and not parsed.path
+
+
+async def _message(websocket: WebSocket, *, limit: int = MAX_TERMINAL_FRAME) -> dict:
+    text = await websocket.receive_text()
+    if len(text) > limit:
+        raise ProtocolError("terminal message exceeds its size limit")
+    message = json.loads(text)
+    if not isinstance(message, dict):
+        raise ProtocolError("terminal message must be an object")
+    return message
+
+
+@router.websocket("/{session_id}/terminal")
+async def terminal(websocket: WebSocket, vm_id: str, session_id: str):
+    if not _same_origin(websocket):
+        await websocket.close(code=1008)
+        return
+    await websocket.accept()
+    tasks = []
+    try:
+        # Browser WebSockets cannot set X-UVM-Token. Authenticate in the first
+        # frame so tokens never appear in URLs, access logs, or session history.
+        auth = await asyncio.wait_for(_message(websocket, limit=4096), timeout=10)
+        token = auth.get("token")
+        expected = websocket.app.state.uvm_application.settings.api_token
+        if (auth.get("type") != "auth" or not isinstance(token, str) or not token.isascii()
+                or not expected or not hmac.compare_digest(token, expected)):
+            await websocket.close(code=4401)
+            return
+        sessions = websocket.app.state.nest_sessions
+        await sessions.get(vm_id, session_id)
+        offset = auth.get("offset", 0)
+
+        async def output():
+            cursor = offset
+            previous = None
+            while True:
+                record, start, data = await sessions.read(vm_id, session_id, cursor)
+                if previous != record or start != cursor:
+                    await websocket.send_json({"type": "session", "session": record, "offset": start})
+                    previous = record
+                cursor = start
+                if data:
+                    await websocket.send_json({"type": "output", "offset": cursor, "data": base64.b64encode(data).decode("ascii")})
+                    cursor += len(data)
+                    continue
+                if record["status"] not in {"active", "starting"}:
+                    return
+                await sessions.wait(session_id)
+
+        async def controls():
+            while True:
+                message = await _message(websocket)
+                if set(message) == {"type", "data"} and message["type"] == "input":
+                    await sessions.write(vm_id, session_id, terminal_data(message["data"]))
+                elif set(message) == {"type", "rows", "cols"} and message["type"] == "resize":
+                    await sessions.resize(vm_id, session_id, message["rows"], message["cols"])
+                else:
+                    raise ProtocolError("invalid terminal control message")
+
+        tasks = [asyncio.create_task(output()), asyncio.create_task(controls())]
+        done, _pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+        for task in done:
+            task.result()
+    except WebSocketDisconnect:
+        pass
+    except (ValueError, KeyError, TypeError, UnicodeError, RecursionError, TimeoutError):
+        with suppress(WebSocketDisconnect, RuntimeError):
+            await websocket.send_json({"type": "error", "message": "Invalid terminal message or authentication timeout."})
+    except UvmError as error:
+        with suppress(WebSocketDisconnect, RuntimeError):
+            await websocket.send_json({"type": "error", "message": str(error)})
+    finally:
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        with suppress(WebSocketDisconnect, RuntimeError):
+            await websocket.close()

nest/session_store.py

@@ -0,0 +1,169 @@
+"""Private, durable session metadata and bounded PTY transcripts, separate from VM state."""
+
+from __future__ import annotations
+
+import fcntl
+import os
+import sqlite3
+import stat
+import threading
+import time
+import uuid
+from contextlib import contextmanager
+from pathlib import Path
+
+from uvm.errors import ConflictError, NotFoundError, StateError, ValidationError
+
+
+MAX_TRANSCRIPT_BYTES = 8 * 1024 * 1024
+
+
+class SessionStore:
+    def __init__(self, path: Path, *, max_output: int = MAX_TRANSCRIPT_BYTES) -> None:
+        self.path = path
+        self.max_output = max_output
+        self._mutex = threading.RLock()
+        self._owner = None
+        try:
+            path.parent.mkdir(parents=True, mode=0o700, exist_ok=True)
+            if path.parent.is_symlink():
+                raise StateError("session history directory must not be a symlink")
+            path.parent.chmod(0o700)
+            self._owner = self._private_file(path.with_suffix(".lock"))
+            try:
+                fcntl.flock(self._owner, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            except BlockingIOError as error:
+                raise ConflictError("Another Nest server is already managing sessions for this VM registry.") from error
+            descriptor = self._private_file(path)
+            os.close(descriptor)
+            with self._db() as db:
+                db.executescript("""
+                    CREATE TABLE IF NOT EXISTS sessions (
+                        id TEXT PRIMARY KEY, vm_id TEXT NOT NULL, name TEXT NOT NULL,
+                        username TEXT NOT NULL, status TEXT NOT NULL,
+                        created_at REAL NOT NULL, updated_at REAL NOT NULL, ended_at REAL,
+                        rows INTEGER NOT NULL, cols INTEGER NOT NULL, exit_code INTEGER, error TEXT,
+                        output_start INTEGER NOT NULL DEFAULT 0, output_end INTEGER NOT NULL DEFAULT 0,
+                        truncated INTEGER NOT NULL DEFAULT 0
+                    );
+                    CREATE INDEX IF NOT EXISTS sessions_vm ON sessions(vm_id, created_at);
+                    CREATE TABLE IF NOT EXISTS output (
+                        session_id TEXT NOT NULL REFERENCES sessions(id), offset INTEGER NOT NULL,
+                        data BLOB NOT NULL, PRIMARY KEY (session_id, offset)
+                    );
+                """)
+                now = time.time()
+                db.execute("""UPDATE sessions SET status='interrupted', ended_at=?, updated_at=?,
+                    error='Nest restarted; this shell is no longer connected.'
+                    WHERE status IN ('starting', 'active')""", (now, now))
+        except BaseException:
+            self.close()
+            raise
+
+    @staticmethod
+    def _private_file(path: Path) -> int:
+        descriptor = os.open(path, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
+        metadata = os.fstat(descriptor)
+        if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1:
+            os.close(descriptor)
+            raise StateError("session history must use private regular files")
+        os.fchmod(descriptor, 0o600)
+        return descriptor
+
+    @contextmanager
+    def _db(self):
+        with self._mutex:
+            db = None
+            try:
+                db = sqlite3.connect(self.path, timeout=5)
+                db.row_factory = sqlite3.Row
+                with db:
+                    yield db
+            except sqlite3.Error as error:
+                raise StateError("Could not read or save terminal session history.") from error
+            finally:
+                if db is not None:
+                    db.close()
+
+    @staticmethod
+    def _record(row) -> dict:
+        if row is None:
+            raise NotFoundError("Session not found for this VM.")
+        result = dict(row)
+        result["truncated"] = bool(result["truncated"])
+        return result
+
+    def create(self, vm_id: str, rows: int, cols: int) -> dict:
+        session_id = "session-" + uuid.uuid4().hex
+        now = time.time()
+        with self._db() as db:
+            number = db.execute("SELECT COUNT(*) FROM sessions WHERE vm_id=?", (vm_id,)).fetchone()[0] + 1
+            db.execute("""INSERT INTO sessions
+                (id, vm_id, name, username, status, created_at, updated_at, rows, cols)
+                VALUES (?, ?, ?, '', 'starting', ?, ?, ?, ?)""",
+                       (session_id, vm_id, f"Session {number}", now, now, rows, cols))
+        return self.get(vm_id, session_id)
+
+    def list(self, vm_id: str) -> list[dict]:
+        with self._db() as db:
+            return [self._record(row) for row in db.execute(
+                "SELECT * FROM sessions WHERE vm_id=? ORDER BY created_at DESC, id DESC", (vm_id,),
+            )]
+
+    def get(self, vm_id: str, session_id: str) -> dict:
+        with self._db() as db:
+            return self._record(db.execute("SELECT * FROM sessions WHERE vm_id=? AND id=?", (vm_id, session_id)).fetchone())
+
+    def activate(self, session_id: str, username: str) -> None:
+        with self._db() as db:
+            db.execute("UPDATE sessions SET status='active', username=?, updated_at=? WHERE id=?",
+                       (username, time.time(), session_id))
+
+    def resize(self, session_id: str, rows: int, cols: int) -> None:
+        with self._db() as db:
+            db.execute("UPDATE sessions SET rows=?, cols=? WHERE id=?", (rows, cols, session_id))
+
+    def finish(self, session_id: str, status: str, exit_code: int | None = None, error: str | None = None) -> None:
+        now = time.time()
+        with self._db() as db:
+            db.execute("""UPDATE sessions SET status=?, ended_at=?, updated_at=?, exit_code=?, error=?
+                WHERE id=? AND status IN ('starting', 'active')""", (status, now, now, exit_code, error, session_id))
+
+    def append(self, session_id: str, data: bytes) -> None:
+        if not data:
+            return
+        with self._db() as db:
+            row = db.execute("SELECT output_end FROM sessions WHERE id=?", (session_id,)).fetchone()
+            if row is None:
+                raise NotFoundError("Session not found.")
+            offset = row[0]
+            end = offset + len(data)
+            start = max(0, end - self.max_output)
+            db.execute("INSERT INTO output VALUES (?, ?, ?)", (session_id, offset, data))
+            db.execute("DELETE FROM output WHERE session_id=? AND offset + length(data)<=?", (session_id, start))
+            first = db.execute("SELECT offset, data FROM output WHERE session_id=? AND offset<?", (session_id, start)).fetchone()
+            if first is not None:
+                db.execute("UPDATE output SET offset=?, data=? WHERE session_id=? AND offset=?",
+                           (start, first["data"][start - first["offset"]:], session_id, first["offset"]))
+            db.execute("UPDATE sessions SET output_start=?, output_end=?, truncated=?, updated_at=? WHERE id=?",
+                       (start, end, int(start > 0), time.time(), session_id))
+
+    def read(self, vm_id: str, session_id: str, offset: int, limit: int = 65536) -> tuple[dict, int, bytes]:
+        with self._db() as db:
+            session = self._record(db.execute("SELECT * FROM sessions WHERE vm_id=? AND id=?", (vm_id, session_id)).fetchone())
+            if type(offset) is not int or not 0 <= offset <= session["output_end"]:
+                raise ValidationError("invalid terminal output offset")
+            start = max(offset, session["output_start"])
+            data = bytearray()
+            for row in db.execute("""SELECT offset, data FROM output
+                WHERE session_id=? AND offset + length(data)>? ORDER BY offset""", (session_id, start)):
+                part = row["data"][max(0, start - row["offset"]):]
+                data.extend(part[:limit - len(data)])
+                if len(data) >= limit:
+                    break
+            return session, start, bytes(data)
+
+    def close(self) -> None:
+        if self._owner is not None:
+            os.close(self._owner)
+            self._owner = None

nest/sessions.py

@@ -0,0 +1,174 @@
+"""Nest-owned live terminal sessions; browser attachments do not own guest shells."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from uvm.app import Application
+from uvm.errors import ConflictError, GuestExecutionError, StateError, UvmError
+from uvm.guest_terminal import GuestTerminal, validate_size
+
+from .session_store import SessionStore
+
+
+@dataclass
+class _LiveSession:
+    terminal: GuestTerminal
+    changed: asyncio.Event = field(default_factory=asyncio.Event)
+    task: asyncio.Task | None = None
+
+
+class SessionManager:
+    def __init__(self, application: Application) -> None:
+        self.application = application
+        self._store: SessionStore | None = None
+        self._live: dict[str, _LiveSession] = {}
+        self._pending: set[asyncio.Task] = set()
+        self._create_lock = asyncio.Lock()
+        self._closing = False
+
+    def _history(self) -> SessionStore:
+        if self._closing:
+            raise ConflictError("Nest is shutting down.")
+        if self._store is None:
+            try:
+                self._store = SessionStore(Path(self.application.settings.base) / "nest" / "sessions.sqlite3")
+            except OSError as error:
+                raise StateError("Could not open private terminal history; run Nest with access to the VM registry.") from error
+        return self._store
+
+    async def _vm(self, vm_id: str) -> None:
+        # Exact IDs prevent a session from following a reused VM name or IP.
+        await asyncio.to_thread(self.application.state_store.get, vm_id)
+
+    async def list(self, vm_id: str) -> list[dict]:
+        await self._vm(vm_id)
+        return await asyncio.to_thread(self._history().list, vm_id)
+
+    async def get(self, vm_id: str, session_id: str) -> dict:
+        await self._vm(vm_id)
+        return await asyncio.to_thread(self._history().get, vm_id, session_id)
+
+    async def create(self, vm_id: str, rows: int, cols: int) -> dict:
+        validate_size(rows, cols)
+        await self._vm(vm_id)
+        history = self._history()
+        # A dropped HTTP request must not abandon a just-created guest connection.
+        task = asyncio.create_task(self._create(history, vm_id, rows, cols))
+        self._pending.add(task)
+
+        def completed(future):
+            self._pending.discard(future)
+            if not future.cancelled():
+                future.exception()
+
+        task.add_done_callback(completed)
+        return await asyncio.shield(task)
+
+    async def _create(self, history: SessionStore, vm_id: str, rows: int, cols: int) -> dict:
+        async with self._create_lock:
+            if self._closing:
+                raise ConflictError("Nest is shutting down.")
+            if len(self._live) >= 32:
+                raise ConflictError("There are already 32 active terminal sessions. End a session before launching another.")
+            record = await asyncio.to_thread(history.create, vm_id, rows, cols)
+            terminal = None
+            try:
+                terminal, username = await asyncio.to_thread(self.application.lifecycle.open_terminal, vm_id, rows=rows, cols=cols)
+                await asyncio.to_thread(history.activate, record["id"], username)
+                live = _LiveSession(terminal)
+                self._live[record["id"]] = live
+                live.task = asyncio.create_task(self._read(history, record["id"], live))
+            except BaseException as error:
+                if terminal is not None:
+                    terminal.close()
+                message = str(error) if isinstance(error, UvmError) else "Could not launch the guest terminal."
+                await asyncio.to_thread(history.finish, record["id"], "failed", None, message)
+                raise
+            return await asyncio.to_thread(history.get, vm_id, record["id"])
+
+    async def _read(self, history: SessionStore, session_id: str, live: _LiveSession) -> None:
+        status, exit_code, error = "interrupted", None, None
+        try:
+            while True:
+                message = await live.terminal.read()
+                if message["type"] == "exit":
+                    status, exit_code = "closed", message["exit_code"]
+                    break
+                await asyncio.to_thread(history.append, session_id, message["data"])
+                live.changed.set()
+        except asyncio.CancelledError:
+            error = "Nest stopped; this shell is no longer connected." if self._closing else "The guest terminal connection was interrupted."
+        except (UvmError, OSError) as exception:
+            error = str(exception) if isinstance(exception, UvmError) else "The guest terminal connection was interrupted."
+        finally:
+            live.terminal.close()
+            try:
+                await asyncio.to_thread(history.finish, session_id, status, exit_code, error)
+            finally:
+                live.changed.set()
+                self._live.pop(session_id, None)
+
+    async def read(self, vm_id: str, session_id: str, offset: int) -> tuple[dict, int, bytes]:
+        return await asyncio.to_thread(self._history().read, vm_id, session_id, offset)
+
+    async def wait(self, session_id: str) -> None:
+        live = self._live.get(session_id)
+        if live is None:
+            await asyncio.sleep(0.1)
+            return
+        try:
+            await asyncio.wait_for(live.changed.wait(), timeout=1)
+        except TimeoutError:
+            pass
+        live.changed.clear()
+
+    async def _active(self, vm_id: str, session_id: str) -> _LiveSession:
+        record = await self.get(vm_id, session_id)
+        live = self._live.get(session_id)
+        if record["status"] != "active" or live is None:
+            raise ConflictError("This session has ended. Launch a new session to run commands.")
+        return live
+
+    async def write(self, vm_id: str, session_id: str, data: bytes) -> None:
+        live = await self._active(vm_id, session_id)
+        try:
+            await live.terminal.write(data)
+        except OSError as error:
+            live.task.cancel()
+            raise GuestExecutionError("The terminal connection was interrupted.", outcome_unknown=True) from error
+
+    async def resize(self, vm_id: str, session_id: str, rows: int, cols: int) -> None:
+        validate_size(rows, cols)
+        live = await self._active(vm_id, session_id)
+        try:
+            await live.terminal.resize(rows, cols)
+            await asyncio.to_thread(self._history().resize, session_id, rows, cols)
+        except OSError as error:
+            live.task.cancel()
+            raise GuestExecutionError("The terminal connection was interrupted.", outcome_unknown=True) from error
+
+    async def end(self, vm_id: str, session_id: str) -> dict:
+        record = await self.get(vm_id, session_id)
+        live = self._live.get(session_id)
+        if live is not None:
+            try:
+                await live.terminal.end()
+                await asyncio.wait_for(asyncio.shield(live.task), timeout=8)
+            except (OSError, UvmError):
+                live.task.cancel()
+                await asyncio.gather(live.task, return_exceptions=True)
+            record = await self.get(vm_id, session_id)
+        return record
+
+    async def shutdown(self) -> None:
+        self._closing = True
+        await asyncio.gather(*self._pending, return_exceptions=True)
+        tasks = [live.task for live in self._live.values()]
+        for task in tasks:
+            task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+        if self._store is not None:
+            self._store.close()