diff --git a/hermes_cli/pty_session.py b/hermes_cli/pty_session.py index 0bc6f1104a8..fd7b614d334 100644 --- a/hermes_cli/pty_session.py +++ b/hermes_cli/pty_session.py @@ -109,3 +109,67 @@ class PtySession: self.bridge.close() except Exception: pass + + +from typing import Callable, Dict, Tuple + + +class RegistryFull(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/tests/test_pty_session.py b/tests/test_pty_session.py index aa477bedb70..f291bef134c 100644 --- a/tests/test_pty_session.py +++ b/tests/test_pty_session.py @@ -113,3 +113,50 @@ async def test_eof_marks_dead_and_closes_socket_4410(): 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()