feat: enhance terminal session handling with output acknowledgment and error handling

nest/session_routes.py

@@ -76,7 +76,10 @@ def _same_origin(websocket: WebSocket) -> bool:
     origin = websocket.headers.get("origin")
     if origin is None:
         return True  # Non-browser API clients must still authenticate.
+    try:
         parsed = urlsplit(origin)
+    except ValueError:
+        return False
     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
 
@@ -111,6 +114,7 @@ async def terminal(websocket: WebSocket, vm_id: str, session_id: str):
         sessions = websocket.app.state.nest_sessions
         await sessions.get(vm_id, session_id)
         offset = auth.get("offset", 0)
+        acknowledgements = asyncio.Queue(maxsize=1)
 
         async def output():
             cursor = offset
@@ -124,6 +128,10 @@ async def terminal(websocket: WebSocket, vm_id: str, session_id: str):
                 if data:
                     await websocket.send_json({"type": "output", "offset": cursor, "data": base64.b64encode(data).decode("ascii")})
                     cursor += len(data)
+                    # Only one output chunk may be waiting for the terminal renderer.
+                    acknowledged = await asyncio.wait_for(acknowledgements.get(), timeout=30)
+                    if type(acknowledged) is not int or acknowledged != cursor:
+                        raise ProtocolError("invalid terminal output acknowledgement")
                     continue
                 if record["status"] not in {"active", "starting"}:
                     return
@@ -136,6 +144,10 @@ async def terminal(websocket: WebSocket, vm_id: str, session_id: str):
                     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"])
+                elif set(message) == {"type", "offset"} and message["type"] == "ack":
+                    if acknowledgements.full():
+                        raise ProtocolError("unexpected terminal output acknowledgement")
+                    acknowledgements.put_nowait(message["offset"])
                 else:
                     raise ProtocolError("invalid terminal control message")
 

nest/sessions.py

@@ -7,7 +7,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 
 from uvm.app import Application
-from uvm.errors import ConflictError, GuestExecutionError, StateError, UvmError
+from uvm.errors import ConflictError, GuestExecutionError, NotFoundError, StateError, UvmError
 from uvm.guest_terminal import GuestTerminal, validate_size
 
 from .session_store import SessionStore
@@ -15,6 +15,7 @@ from .session_store import SessionStore
 
 @dataclass
 class _LiveSession:
+    vm_id: str
     terminal: GuestTerminal
     changed: asyncio.Event = field(default_factory=asyncio.Event)
     task: asyncio.Task | None = None
@@ -78,7 +79,7 @@ class SessionManager:
             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)
+                live = _LiveSession(vm_id, terminal)
                 self._live[record["id"]] = live
                 live.task = asyncio.create_task(self._read(history, record["id"], live))
             except BaseException as error:
@@ -126,10 +127,12 @@ class SessionManager:
         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:
+        if live is None:
+            await self.get(vm_id, session_id)
             raise ConflictError("This session has ended. Launch a new session to run commands.")
+        if live.vm_id != vm_id:
+            raise NotFoundError("Session not found for this VM.")
         return live
 
     async def write(self, vm_id: str, session_id: str, data: bytes) -> None:

uvm/guest_agent.py

@@ -353,7 +353,7 @@ def run_terminal(connection: socket.socket, username: str, rows: int, cols: int,
                             selector.modify(master, selectors.EVENT_READ | (selectors.EVENT_WRITE if pending else 0), "pty")
                 if _exited(process):
                     # Drain the leader's final output before terminating any background jobs.
-                    while True:
+                    for _ in range(16):
                         try:
                             data = os.read(master, MAX_TERMINAL_BYTES)
                         except (BlockingIOError, OSError):