diff --git a/hermes_cli/pty_session.py b/hermes_cli/pty_session.py index 59cede640a0f..0bc6f1104a83 100644 --- a/hermes_cli/pty_session.py +++ b/hermes_cli/pty_session.py @@ -8,6 +8,13 @@ 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.""" @@ -30,3 +37,75 @@ class RingBuffer: @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 diff --git a/tests/test_pty_session.py b/tests/test_pty_session.py index 1ab5bacdbaa5..aa477bedb70d 100644 --- a/tests/test_pty_session.py +++ b/tests/test_pty_session.py @@ -1,3 +1,8 @@ +import asyncio +import time + +import pytest + from hermes_cli.pty_session import RingBuffer @@ -22,3 +27,89 @@ def test_ringbuffer_truncation_across_appends(): 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()