From e10e4bca825fff6b3c690857ad18c56aabaf1a5f Mon Sep 17 00:00:00 2001 From: Benn Denton <80194363+TinkerOfThings@users.noreply.github.com> Date: Sun, 21 Jun 2026 08:48:21 +0000 Subject: [PATCH] feat(chat): reattach /api/pty sessions via ?attach= token Keep-alive path when ?attach= is present: PTY outlives the socket via PTY_REGISTRY, reattaches on reconnect. No token = unchanged legacy pump (_legacy_pump). detach (not close) on disconnect. Co-Authored-By: Claude Opus 4.8 --- hermes_cli/pty_session.py | 18 ++- hermes_cli/web_server.py | 198 +++++++++++++++++++++++---------- tests/test_pty_keepalive_ws.py | 53 +++++++++ 3 files changed, 206 insertions(+), 63 deletions(-) create mode 100644 tests/test_pty_keepalive_ws.py diff --git a/hermes_cli/pty_session.py b/hermes_cli/pty_session.py index fd7b614d334..1a6e9b78bd2 100644 --- a/hermes_cli/pty_session.py +++ b/hermes_cli/pty_session.py @@ -93,8 +93,14 @@ class PtySession: await ws.send_bytes(snap) def detach(self, ws) -> None: - if self._ws is ws: - self._ws = None + # Only the currently-attached socket may mark the session detached. + # A superseded socket's handler also calls detach on its way out + # (its ``finally`` runs after the new tab attached); flipping + # ``attached`` then would make a session with a live viewer look + # idle and reapable. + if self._ws is not ws: + return + self._ws = None self.attached = False self.last_detached_at = time.monotonic() @@ -106,7 +112,9 @@ class PtySession: except (asyncio.CancelledError, Exception): pass try: - self.bridge.close() + # bridge.close() joins the child — blocking; keep it off the + # event loop (#53227). + await asyncio.to_thread(self.bridge.close) except Exception: pass @@ -138,7 +146,9 @@ class PtySessionRegistry: self._sessions.pop(key, None) if len(self._sessions) >= self._max: self._reap_one_idle_or_raise() - bridge = spawn() + # PTY spawn does blocking fork/exec work — keep it off the event + # loop (#53227). + bridge = await asyncio.to_thread(spawn) session = PtySession(key, bridge, buffer_cap=self._buffer_cap, read_timeout=self._read_timeout) await session.start() diff --git a/hermes_cli/web_server.py b/hermes_cli/web_server.py index d8b9d4ac7af..325d9353de1 100644 --- a/hermes_cli/web_server.py +++ b/hermes_cli/web_server.py @@ -12540,6 +12540,105 @@ else: _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 # 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). + + Behavior is identical to the pre-keep-alive ``pty_ws`` body, including the + #54028 half-open-socket protection (reader EOF → close the WS so the + writer's ``ws.receive()`` unparks) and the #53227 ``to_thread`` offloads + for the blocking ``bridge.close()``. + """ + loop = asyncio.get_running_loop() + + # --- reader task: PTY master → WebSocket ---------------------------- + async def pump_pty_to_ws() -> None: + try: + 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 + finally: + # The child has exited (EOF) or the send side broke. Close the + # WebSocket so the writer loop's ``ws.receive()`` returns instead + # of blocking forever — otherwise, when the browser's socket is + # half-open (no FIN delivered, common on macOS/launchd) the + # handler never reaches its ``finally`` and the PTY's fds leak. + # With dashboard auto-reconnect (#52962) every dropped socket then + # stacks a fresh PTY on top of the orphaned one, exhausting fds. + # + # Reap the bridge here too (close() is idempotent): on child EOF the + # writer loop's ``finally`` is the usual closer, but if the handler + # task is cancelled the instant we close the WS, that ``finally`` + # can be skipped, leaking the PTY. Closing from the EOF path makes + # the reap independent of that cancellation race (#54028). + try: + await asyncio.to_thread(bridge.close) + except Exception: + pass + try: + await ws.close() + except Exception: + pass + + reader_task = asyncio.create_task(pump_pty_to_ws()) + + # --- writer loop: WebSocket → PTY master ---------------------------- + try: + while True: + try: + msg = await ws.receive() + except RuntimeError: + # Raised when ws.receive() is called after the socket is + # already disconnected (e.g. closed by the reader task above). + break + 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 + # Resize escape is consumed locally, never written to the PTY. + 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 + await asyncio.to_thread(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. @@ -13745,71 +13844,57 @@ 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 = await asyncio.to_thread(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: - try: - 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 - finally: - # The child has exited (EOF) or the send side broke. Close the - # WebSocket so the writer loop's ``ws.receive()`` returns instead - # of blocking forever — otherwise, when the browser's socket is - # half-open (no FIN delivered, common on macOS/launchd) the - # handler never reaches its ``finally`` and the PTY's fds leak. - # With dashboard auto-reconnect (#52962) every dropped socket then - # stacks a fresh PTY on top of the orphaned one, exhausting fds. - # - # Reap the bridge here too (close() is idempotent): on child EOF the - # writer loop's ``finally`` is the usual closer, but if the handler - # task is cancelled the instant we close the WS, that ``finally`` - # can be skipped, leaking the PTY. Closing from the EOF path makes - # the reap independent of that cancellation race (#54028). - try: - await asyncio.to_thread(bridge.close) - except Exception: - pass - try: - await ws.close() - except Exception: - pass - - reader_task = asyncio.create_task(pump_pty_to_ws()) + await session.attach(ws) # --- writer loop: WebSocket → PTY master ---------------------------- + # No reader task here: the session's drain task (spawned once per PTY, + # inside the registry) forwards PTY output to whichever socket is + # attached and rings-buffers it while detached. On child EOF the drain + # closes the attached socket with 4410, which unparks ``ws.receive()`` + # below — same half-open-socket protection the legacy pump has (#54028). try: while True: try: msg = await ws.receive() except RuntimeError: - # Raised when ws.receive() is called after the socket is - # already disconnected (e.g. closed by the reader task above). + # ws.receive() after the socket is already disconnected + # (e.g. closed by the drain task on process exit). break - 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: @@ -13821,21 +13906,16 @@ async def pty_ws(ws: WebSocket) -> 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 - await asyncio.to_thread(bridge.close) + # Detach only — the PTY keeps running for a reattach; the registry + # reaper closes it after the TTL (or immediately on process exit). + 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 00000000000..782967ef207 --- /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()