diff --git a/hermes_cli/pty_session.py b/hermes_cli/pty_session.py new file mode 100644 index 000000000000..9b02bd8ef68f --- /dev/null +++ b/hermes_cli/pty_session.py @@ -0,0 +1,185 @@ +"""Keep-alive PTY sessions for dashboard terminals. + +A PTY process outlives the WebSocket that created it: a single drain task +always reads the PTY into a bounded RingBuffer and forwards to the attached +socket when present. Reconnecting with the same opaque token replays the +buffer and resumes live. See +docs/superpowers/specs/2026-06-20-pty-keepalive-reattach-design.md. +""" +from __future__ import annotations + +import asyncio +import time +from typing import Optional + +WS_CLOSE_PROCESS_EXITED = 4410 +WS_CLOSE_SUPERSEDED = 4409 + + +class RingBuffer: + """Keeps only the most recent ``capacity`` bytes appended to it.""" + + def __init__(self, capacity: int) -> None: + self._cap = capacity + self._buf = bytearray() + self._truncated = False + + def append(self, data: bytes) -> None: + self._buf.extend(data) + overflow = len(self._buf) - self._cap + if overflow > 0: + del self._buf[:overflow] + self._truncated = True + + def snapshot(self) -> bytes: + return bytes(self._buf) + + @property + def truncated(self) -> bool: + return self._truncated + + +class PtySession: + def __init__(self, key: str, bridge, *, buffer_cap: int, read_timeout: float) -> None: + self.key = key + self.bridge = bridge + self.buffer = RingBuffer(buffer_cap) + self.alive = True + self.attached = False + self.last_detached_at: Optional[float] = None + self._read_timeout = read_timeout + self._ws = None + self._drain_task: Optional[asyncio.Task] = None + + async def start(self) -> None: + self._drain_task = asyncio.create_task(self._drain()) + + async def _drain(self) -> None: + loop = asyncio.get_running_loop() + while True: + chunk = await loop.run_in_executor(None, self.bridge.read, self._read_timeout) + if chunk is None: # EOF — the agent process exited + self.alive = False + ws = self._ws + if ws is not None: + try: + await ws.close(code=WS_CLOSE_PROCESS_EXITED) + except Exception: + pass + return + if not chunk: # idle tick + await asyncio.sleep(0) + continue + self.buffer.append(chunk) + ws = self._ws + if ws is not None: + try: + await ws.send_bytes(chunk) + except Exception: + pass # detached mid-send; keep buffering + + async def attach(self, ws) -> None: + old = self._ws + if old is not None and old is not ws: + try: + await old.close(code=WS_CLOSE_SUPERSEDED) + except Exception: + pass + self._ws = ws + self.attached = True + self.last_detached_at = None + snap = self.buffer.snapshot() + if snap: + await ws.send_bytes(snap) + + def detach(self, ws) -> None: + if self._ws is ws: + self._ws = None + self.attached = False + self.last_detached_at = time.monotonic() + + async def close(self) -> None: + if self._drain_task is not None: + self._drain_task.cancel() + try: + await self._drain_task + except (asyncio.CancelledError, Exception): + pass + try: + self.bridge.close() + except Exception: + pass + + +from typing import Callable, Dict, Tuple + + +class RegistryFull(Exception): + pass + + +async def run_reaper(registry: "PtySessionRegistry", *, interval: float = 60.0) -> None: + """Periodically reap idle/dead keep-alive sessions. Cancelled on shutdown.""" + while True: + await asyncio.sleep(interval) + try: + await registry.reap_idle() + except Exception: + pass + + +class PtySessionRegistry: + def __init__(self, *, ttl: float, max_sessions: int, + buffer_cap: int, read_timeout: float) -> None: + self._ttl = ttl + self._max = max_sessions + self._buffer_cap = buffer_cap + self._read_timeout = read_timeout + self._sessions: Dict[str, PtySession] = {} + + async def attach_or_spawn(self, key: str, *, spawn: Callable[[], object] + ) -> Tuple[PtySession, bool]: + await self.reap_idle() + existing = self._sessions.get(key) + if existing is not None and existing.alive: + return existing, False + if existing is not None: # dead remnant + await existing.close() + self._sessions.pop(key, None) + if len(self._sessions) >= self._max: + self._reap_one_idle_or_raise() + bridge = spawn() + session = PtySession(key, bridge, buffer_cap=self._buffer_cap, + read_timeout=self._read_timeout) + await session.start() + self._sessions[key] = session + return session, True + + def detach(self, key: str, ws) -> None: + s = self._sessions.get(key) + if s is not None: + s.detach(ws) + + async def reap_idle(self, now: Optional[float] = None) -> None: + now = time.monotonic() if now is None else now + doomed = [ + key for key, s in self._sessions.items() + if (not s.alive) + or (not s.attached and s.last_detached_at is not None + and (now - s.last_detached_at) > self._ttl) + ] + for key in doomed: + await self._sessions.pop(key).close() + + def _reap_one_idle_or_raise(self) -> None: + idle = [s for s in self._sessions.values() + if not s.attached and s.last_detached_at is not None] + if not idle: + raise RegistryFull() + oldest = min(idle, key=lambda s: s.last_detached_at or 0.0) + self._sessions.pop(oldest.key, None) + asyncio.create_task(oldest.close()) + + async def close_all(self) -> None: + for key in list(self._sessions): + await self._sessions.pop(key).close() diff --git a/hermes_cli/web_server.py b/hermes_cli/web_server.py index 398e61772f08..ac2b7c54c76d 100644 --- a/hermes_cli/web_server.py +++ b/hermes_cli/web_server.py @@ -166,9 +166,14 @@ async def _lifespan(app: "FastAPI"): ) cron_thread.start() + # Reap idle/dead keep-alive PTY sessions in the background (30-min TTL). + pty_reaper_task = asyncio.create_task(run_reaper(PTY_REGISTRY)) + try: yield finally: + pty_reaper_task.cancel() + await PTY_REGISTRY.close_all() if cron_stop is not None: cron_stop.set() @@ -10925,6 +10930,65 @@ class PtyUnavailableError(RuntimeError): # type: ignore[no-redef] _RESIZE_RE = re.compile(rb"\x1b\[RESIZE:(\d+);(\d+)\]") _PTY_READ_CHUNK_TIMEOUT = 0.2 + +# Keep-alive PTY sessions: a terminal connecting with ``?attach=`` is +# bound to a process that survives disconnect/refresh and is reattachable. +from hermes_cli.pty_session import PtySessionRegistry, RegistryFull, run_reaper # noqa: E402 + +PTY_REGISTRY = PtySessionRegistry( + ttl=30 * 60, + max_sessions=16, + buffer_cap=1 * 1024 * 1024, + read_timeout=_PTY_READ_CHUNK_TIMEOUT, +) + + +async def _legacy_pump(ws: "WebSocket", bridge) -> None: + """Original 1:1 socket<->PTY pump: stream until disconnect, then close the + bridge. Used when no ``?attach=`` token is supplied (keep-alive opt-in).""" + loop = asyncio.get_running_loop() + + async def pump_pty_to_ws() -> None: + while True: + chunk = await loop.run_in_executor(None, bridge.read, _PTY_READ_CHUNK_TIMEOUT) + if chunk is None: + return + if not chunk: + await asyncio.sleep(0) + continue + try: + await ws.send_bytes(chunk) + except Exception: + return + + reader_task = asyncio.create_task(pump_pty_to_ws()) + try: + while True: + msg = await ws.receive() + if msg.get("type") == "websocket.disconnect": + break + raw = msg.get("bytes") + if raw is None: + text = msg.get("text") + raw = text.encode("utf-8") if isinstance(text, str) else b"" + if not raw: + continue + match = _RESIZE_RE.match(raw) + if match and match.end() == len(raw): + bridge.resize(cols=int(match.group(1)), rows=int(match.group(2))) + continue + bridge.write(raw) + except WebSocketDisconnect: + pass + finally: + reader_task.cancel() + try: + await reader_task + except (asyncio.CancelledError, Exception): + pass + bridge.close() + + _VALID_CHANNEL_RE = re.compile(r"^[A-Za-z0-9._-]{1,128}$") # Starlette's TestClient reports the peer as "testclient"; treat it as # loopback so tests don't need to rewrite request scope. @@ -11439,43 +11503,43 @@ async def pty_ws(ws: WebSocket) -> None: return + attach_token = ws.query_params.get("attach") or None + + def _spawn(): + return PtyBridge.spawn(argv, cwd=cwd, env=env) + + if attach_token is None: + # Legacy path: 1:1 socket<->PTY, killed on disconnect (unchanged). + try: + bridge = _spawn() + except PtyUnavailableError as exc: + await ws.send_text(f"\r\n\x1b[31mChat unavailable: {exc}\x1b[0m\r\n") + await ws.close(code=1011) + return + except (FileNotFoundError, OSError) as exc: + await ws.send_text(f"\r\n\x1b[31mChat failed to start: {exc}\x1b[0m\r\n") + await ws.close(code=1011) + return + await _legacy_pump(ws, bridge) + return + + # Keep-alive path: the PTY outlives this socket; reattach by token. try: - bridge = PtyBridge.spawn(argv, cwd=cwd, env=env) + session, _created = await PTY_REGISTRY.attach_or_spawn(attach_token, spawn=_spawn) except PtyUnavailableError as exc: await ws.send_text(f"\r\n\x1b[31mChat unavailable: {exc}\x1b[0m\r\n") await ws.close(code=1011) return - except (FileNotFoundError, OSError) as exc: - await ws.send_text(f"\r\n\x1b[31mChat failed to start: {exc}\x1b[0m\r\n") + except (FileNotFoundError, OSError, RegistryFull) as exc: + await ws.send_text(f"\r\n\x1b[31mChat unavailable: {exc}\x1b[0m\r\n") await ws.close(code=1011) return - loop = asyncio.get_running_loop() - - # --- reader task: PTY master → WebSocket ---------------------------- - async def pump_pty_to_ws() -> None: - while True: - chunk = await loop.run_in_executor( - None, bridge.read, _PTY_READ_CHUNK_TIMEOUT - ) - if chunk is None: # EOF - return - if not chunk: # no data this tick; yield control and retry - await asyncio.sleep(0) - continue - try: - await ws.send_bytes(chunk) - except Exception: - return - - reader_task = asyncio.create_task(pump_pty_to_ws()) - - # --- writer loop: WebSocket → PTY master ---------------------------- + await session.attach(ws) try: while True: msg = await ws.receive() - msg_type = msg.get("type") - if msg_type == "websocket.disconnect": + if msg.get("type") == "websocket.disconnect": break raw = msg.get("bytes") if raw is None: @@ -11487,21 +11551,14 @@ async def pump_pty_to_ws() -> None: # Resize escape is consumed locally, never written to the PTY. match = _RESIZE_RE.match(raw) if match and match.end() == len(raw): - cols = int(match.group(1)) - rows = int(match.group(2)) - bridge.resize(cols=cols, rows=rows) + session.bridge.resize(cols=int(match.group(1)), rows=int(match.group(2))) continue - bridge.write(raw) + session.bridge.write(raw) except WebSocketDisconnect: pass finally: - reader_task.cancel() - try: - await reader_task - except (asyncio.CancelledError, Exception): - pass - bridge.close() + PTY_REGISTRY.detach(attach_token, ws) # --------------------------------------------------------------------------- diff --git a/tests/test_pty_keepalive_ws.py b/tests/test_pty_keepalive_ws.py new file mode 100644 index 000000000000..782967ef2070 --- /dev/null +++ b/tests/test_pty_keepalive_ws.py @@ -0,0 +1,53 @@ +import pytest + +from hermes_cli import web_server + + +@pytest.mark.asyncio +async def test_attach_token_reuses_same_session(monkeypatch): + """Two connects with the same ?attach= token hit one spawned bridge.""" + spawned = [] + + class FakeBridge: + def __init__(self): + self.alive = True + + def read(self, timeout): + return b"" # idle forever + + def write(self, data): + pass + + def resize(self, cols, rows): + pass + + def close(self): + self.alive = False + + def fake_spawn(argv, cwd=None, env=None): + b = FakeBridge() + spawned.append(b) + return b + + monkeypatch.setattr(web_server.PtyBridge, "spawn", staticmethod(fake_spawn)) + # bypass auth + argv resolution for the test + monkeypatch.setattr(web_server, "_ws_auth_reason", lambda ws: (None, "test")) + monkeypatch.setattr(web_server, "_ws_host_origin_reason", lambda ws: None) + monkeypatch.setattr(web_server, "_ws_client_reason", lambda ws: None) + + async def fake_argv(**kw): + return (["x"], "/tmp", {}) + + monkeypatch.setattr(web_server, "_resolve_chat_argv_async", fake_argv) + + from starlette.testclient import TestClient + + try: + client = TestClient(web_server.app) + with client.websocket_connect("/api/pty?attach=TOK1") as ws1: + ws1.send_bytes(b"hi") + with client.websocket_connect("/api/pty?attach=TOK1") as ws2: + ws2.send_bytes(b"again") + assert len(spawned) == 1 # reattached, did not respawn + finally: + web_server.PTY_REGISTRY._sessions.clear() diff --git a/tests/test_pty_session.py b/tests/test_pty_session.py new file mode 100644 index 000000000000..4fcce10c1a81 --- /dev/null +++ b/tests/test_pty_session.py @@ -0,0 +1,182 @@ +import asyncio +import time + +import pytest + +from hermes_cli.pty_session import RingBuffer + + +def test_ringbuffer_keeps_everything_under_capacity(): + rb = RingBuffer(10) + rb.append(b"abc") + rb.append(b"def") + assert rb.snapshot() == b"abcdef" + assert rb.truncated is False + + +def test_ringbuffer_drops_oldest_over_capacity(): + rb = RingBuffer(4) + rb.append(b"abcdef") # 6 bytes into a 4-byte buffer + assert rb.snapshot() == b"cdef" + assert rb.truncated is True + + +def test_ringbuffer_truncation_across_appends(): + rb = RingBuffer(3) + rb.append(b"ab") + rb.append(b"cd") # now "abcd" -> keep "bcd" + assert rb.snapshot() == b"bcd" + assert rb.truncated is True + + +class FakeBridge: + """Implements the bridge contract PtySession depends on.""" + + def __init__(self, chunks): + self._chunks = list(chunks) # bytes; b"" = idle tick; None = EOF + self.written = bytearray() + self.closed = False + self.resized = None + + def read(self, timeout): + if not self._chunks: + return b"" # idle + return self._chunks.pop(0) + + def write(self, data): + self.written.extend(data) + + def resize(self, cols, rows): + self.resized = (cols, rows) + + def close(self): + self.closed = True + + +class FakeWS: + def __init__(self): + self.sent = [] # list of ("bytes"|"text", payload) + self.close_code = None + + async def send_bytes(self, data): + self.sent.append(("bytes", bytes(data))) + + async def send_text(self, text): + self.sent.append(("text", text)) + + async def close(self, code=1000, reason=""): + self.close_code = code + + +@pytest.mark.asyncio +async def test_attach_replays_buffer_then_streams_live(): + from hermes_cli.pty_session import PtySession + bridge = FakeBridge([b"hello ", b"world", None]) + s = PtySession("k", bridge, buffer_cap=1024, read_timeout=0.01) + await s.start() + await asyncio.sleep(0.05) # drain consumes "hello world" + ws = FakeWS() + await s.attach(ws) + replay = b"".join(p for kind, p in ws.sent if kind == "bytes") + assert replay == b"hello world" + await s.close() + + +@pytest.mark.asyncio +async def test_detach_keeps_draining_into_buffer(): + from hermes_cli.pty_session import PtySession + bridge = FakeBridge([b"one", b"", b"two"]) + s = PtySession("k", bridge, buffer_cap=1024, read_timeout=0.01) + await s.start() + ws = FakeWS() + await s.attach(ws) + s.detach(ws) + assert s.attached is False + assert s.last_detached_at is not None + await asyncio.sleep(0.05) # "two" drains while detached + ws2 = FakeWS() + await s.attach(ws2) + replay = b"".join(p for kind, p in ws2.sent if kind == "bytes") + assert replay == b"onetwo" + await s.close() + + +@pytest.mark.asyncio +async def test_eof_marks_dead_and_closes_socket_4410(): + from hermes_cli.pty_session import PtySession + bridge = FakeBridge([b"bye", None]) + s = PtySession("k", bridge, buffer_cap=1024, read_timeout=0.01) + await s.start() + ws = FakeWS() + await s.attach(ws) + await asyncio.sleep(0.05) # drain hits None (EOF) + assert s.alive is False + assert ws.close_code == 4410 + await s.close() + + +from hermes_cli.pty_session import PtySessionRegistry, RegistryFull + + +def make_registry(ttl=1800.0, max_sessions=16): + return PtySessionRegistry(ttl=ttl, max_sessions=max_sessions, + buffer_cap=1024, read_timeout=0.01) + + +@pytest.mark.asyncio +async def test_same_key_reattaches_same_session(): + reg = make_registry() + b1 = FakeBridge([b"", b"", b""]) + s1, created1 = await reg.attach_or_spawn("tok", spawn=lambda: b1) + s2, created2 = await reg.attach_or_spawn("tok", spawn=lambda: FakeBridge([])) + assert created1 is True and created2 is False + assert s1 is s2 + assert s2.bridge is b1 # second spawn callable was NOT used + await reg.close_all() + + +@pytest.mark.asyncio +async def test_reap_idle_closes_sessions_past_ttl(): + reg = make_registry(ttl=10.0) + b = FakeBridge([b"", b""]) + s, _ = await reg.attach_or_spawn("tok", spawn=lambda: b) + ws = FakeWS() + await s.attach(ws) + s.detach(ws) + s.last_detached_at = time.monotonic() - 11.0 # detached 11s ago, ttl 10s + await reg.reap_idle() + assert b.closed is True + s2, created = await reg.attach_or_spawn("tok", spawn=lambda: FakeBridge([])) + assert created is True + await reg.close_all() + + +@pytest.mark.asyncio +async def test_new_key_at_capacity_raises_when_none_reapable(): + reg = make_registry(max_sessions=1) + b = FakeBridge([b"", b""]) + s, _ = await reg.attach_or_spawn("a", spawn=lambda: b) + await s.attach(FakeWS()) # attached → not reapable + with pytest.raises(RegistryFull): + await reg.attach_or_spawn("b", spawn=lambda: FakeBridge([])) + await reg.close_all() + + +@pytest.mark.asyncio +async def test_reaper_loop_invokes_reap(monkeypatch): + from hermes_cli.pty_session import run_reaper + reg = make_registry() + calls = {"n": 0} + + async def fake_reap(now=None): + calls["n"] += 1 + + monkeypatch.setattr(reg, "reap_idle", fake_reap) + task = asyncio.create_task(run_reaper(reg, interval=0.01)) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + assert calls["n"] >= 2 diff --git a/web/src/pages/ChatPage.tsx b/web/src/pages/ChatPage.tsx index 2a135ed1a57d..001ea066f02d 100644 --- a/web/src/pages/ChatPage.tsx +++ b/web/src/pages/ChatPage.tsx @@ -40,6 +40,30 @@ import { PluginSlot } from "@/plugins"; import { useTheme } from "@/themes"; import { useProfileScope } from "@/contexts/useProfileScope"; +// Stable per-browser token identifying THIS chat tab's keep-alive PTY session. +// Sent as ?attach=; lets a refresh/disconnect reattach to the same live process +// instead of spawning a fresh one. Per-localStorage, so other devices can't grab it. +function ptyAttachToken(): string { + const KEY = "hermes.pty.token.chat"; + let t = ""; + try { + t = window.localStorage.getItem(KEY) ?? ""; + } catch { + /* private mode / storage blocked */ + } + if (!t) { + const a = new Uint8Array(16); + crypto.getRandomValues(a); + t = Array.from(a, (b) => b.toString(16).padStart(2, "0")).join(""); + try { + window.localStorage.setItem(KEY, t); + } catch { + /* ignore */ + } + } + return t; +} + function buildWsUrl( authParam: [string, string], resume: string | null, @@ -52,6 +76,8 @@ function buildWsUrl( // ``_ws_auth_ok`` picks whichever shape matches the current gate state. const qs = new URLSearchParams({ [authParam[0]]: authParam[1], channel }); if (resume) qs.set("resume", resume); + // Keep-alive identity: reattach to this tab's living PTY across refresh. + qs.set("attach", ptyAttachToken()); // Profile-scoped chat: the PTY child gets HERMES_HOME pointed at the // selected profile, so the conversation runs with that profile's model, // skills, memory, and sessions (see web_server._resolve_chat_argv). @@ -140,6 +166,8 @@ export default function ChatPage({ isActive = true }: { isActive?: boolean }) { ); const [copyState, setCopyState] = useState<"idle" | "copied">("idle"); const copyResetRef = useRef | null>(null); + // Pending auto-reconnect after a transient PTY socket drop (keep-alive). + const reconnectTimerRef = useRef | null>(null); // NS-504: when the agent process exits cleanly (the user typed `/exit`, or // started a new session that ended the current PTY child), the PTY socket // closes with a normal code. Before this fix the terminal just printed @@ -618,6 +646,11 @@ export default function ChatPage({ isActive = true }: { isActive?: boolean }) { ws.onopen = () => { setBanner(null); setSessionEnded(false); + // Connected — cancel any pending reconnect from a prior transient drop. + if (reconnectTimerRef.current) { + clearTimeout(reconnectTimerRef.current); + reconnectTimerRef.current = null; + } // Send the initial RESIZE immediately so Ink has *a* size to lay // out against on its first paint. The double-rAF block above will // follow up with the authoritative measurement — at worst Ink @@ -679,14 +712,25 @@ export default function ChatPage({ isActive = true }: { isActive?: boolean }) { // Server already wrote an ANSI error frame. return; } - // Normal/clean exit: the agent process ended (e.g. the user typed - // `/exit`, or started a new session). NS-504: surface an explicit - // restart affordance instead of leaving a dead terminal that only a - // full page refresh could recover. - term.write( - `\r\n\x1b[90m[session ended (code ${ev.code})]\x1b[0m\r\n`, - ); - setSessionEnded(true); + // Keep-alive close-code contract (web_server.pty_ws + pty_session): + // 4410 = the agent PROCESS exited (real end) → restart affordance. + // 4409 = superseded by a newer tab attaching the same token → stay quiet. + // anything else = transient transport drop (refresh, signal loss) → + // reattach to the still-living PTY with the same ?attach= token. + if (ev.code === 4410) { + term.write(`\r\n\x1b[90m[session ended]\x1b[0m\r\n`); + setSessionEnded(true); + return; + } + if (ev.code === 4409) { + return; + } + // Transient: reconnect by re-running the connect effect (reconnectNonce). + if (reconnectTimerRef.current) clearTimeout(reconnectTimerRef.current); + reconnectTimerRef.current = setTimeout(() => { + reconnectTimerRef.current = null; + setReconnectNonce((n) => n + 1); + }, 400); }; // Keystrokes → PTY. @@ -753,6 +797,10 @@ export default function ChatPage({ isActive = true }: { isActive?: boolean }) { clearTimeout(copyResetRef.current); copyResetRef.current = null; } + if (reconnectTimerRef.current) { + clearTimeout(reconnectTimerRef.current); + reconnectTimerRef.current = null; + } }; }, [channel, resumeParam, scopedProfile, reconnectNonce]);