From d1b042c2bc1780bd6713e1e878a69ba80f5f10ca Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Sun, 2 Aug 2026 19:18:42 -0300 Subject: [PATCH 1/3] fix(tui-gateway): close TOCTOU race between WS disconnect teardown and session.resume reattach _close_sessions_for_transport() snapshotted sessions owned by the disconnecting transport under _sessions_lock, released the lock, then mutated each session (close or repoint to the detached sentinel) without re-checking ownership. session.resume's warm-reuse path (_reuse_live_payload -> _live_session_payload) rebinds session["transport"] under _session_resume_lock independently, so a reconnect landing in that window got silently undone: the old transport's teardown either force-closed the just-reattached session or stomped its new transport back to _detached_ws_transport, making a live reconnect look orphaned and eligible for grace-reap. Revalidate transport ownership under the same resume_lock -> sessions_lock ordering already used by the orphan-reap timer, immediately before claiming (close) or repointing (detach) each session. Slow teardown work still runs after releasing both locks. Fixes #HPA-01 --- tests/test_tui_gateway_server.py | 39 ++++++++++++++++++++++++-- tui_gateway/server.py | 48 ++++++++++++++++++++++++-------- 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 2f15a65015c3..05c1bc59cbe9 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -16723,8 +16723,8 @@ def test_session_close_rpc_claims_then_tears_down(monkeypatch): def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch): seen = [] monkeypatch.setattr( - server, "_close_session_by_id", - lambda sid, *, end_reason: bool(seen.append((sid, end_reason))) or True, + server, "_teardown_popped_session", + lambda session, *, end_reason: bool(seen.append((session["_sid"], end_reason))) or True, ) # Detached session "b" would schedule a real grace-reap threading.Timer that # outlives the test; grace=0 short-circuits it so no thread lingers. @@ -16736,11 +16736,46 @@ def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch): try: server._close_sessions_for_transport(transport, end_reason="ws_disconnect") assert seen == [("a", "ws_disconnect")] # only the flagged one closed + assert "a" not in server._sessions # claimed/popped assert server._sessions["b"]["transport"] is server._detached_ws_transport # re-pointed finally: server._sessions.clear() +def test_close_sessions_for_transport_skips_session_reattached_mid_teardown(monkeypatch): + """Regression for the disconnect/reconnect race: if session.resume rebinds + a session onto a new (live) transport between the ownership snapshot and + this function's per-session claim, the old transport's teardown must not + close it or stomp its transport back to the detached sentinel.""" + seen = [] + monkeypatch.setattr( + server, "_teardown_popped_session", + lambda session, *, end_reason: bool(seen.append((session["_sid"], end_reason))) or True, + ) + monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) + old_transport = object() + new_transport = object() + server._sessions.clear() + # "reattached" simulates session.resume's warm-reuse winning the race and + # rebinding the transport (under _session_resume_lock) before this + # function's snapshot is acted on. + server._sessions["reattached"] = {"transport": new_transport, "close_on_disconnect": True} + server._sessions["detached-reattached"] = { + "transport": new_transport, "close_on_disconnect": False, + } + try: + reaped, detached = server._close_sessions_for_transport( + old_transport, end_reason="ws_disconnect" + ) + assert reaped == 0 and detached == 0 + assert seen == [] + assert "reattached" in server._sessions # not closed + assert server._sessions["reattached"]["transport"] is new_transport + assert server._sessions["detached-reattached"]["transport"] is new_transport # untouched + finally: + server._sessions.clear() + + def test_session_create_records_close_on_disconnect_flag(monkeypatch): monkeypatch.setattr(server, "_start_agent_build", lambda sid, session: None) server._sessions.clear() diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 11d2cb17910b..83e10785318d 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1171,9 +1171,8 @@ def _close_sessions_for_transport( transport, *, end_reason: str = "ws_disconnect" ) -> tuple[int, int]: """On transport disconnect, reap the sessions that opted into - close_on_disconnect (sidecar/dashboard) immediately via the unified - ``_close_session_by_id`` path, and re-point the rest back to stdio so later - emits don't hit a dead socket. + close_on_disconnect (sidecar/dashboard) immediately, and re-point the rest + back to the drop sentinel so later emits don't hit a dead socket. Non-flagged detached sessions are handed to the grace-windowed WS-orphan reaper (``_schedule_ws_orphan_reap``): a quick reconnect / session.resume @@ -1182,23 +1181,48 @@ def _close_sessions_for_transport( the single WS-disconnect teardown entry point — there is no second independent reap loop in ``handle_ws``. + The initial snapshot below is taken outside any lock, so each session's + close/detach decision is re-validated (transport still matches) under + ``_session_resume_lock`` immediately before acting. session.resume's + warm-reuse rebind (``_reuse_live_payload`` / ``_live_session_payload``) + takes the same lock to repoint ``session["transport"]``, so a reconnect + that wins the race is never silently closed or stomped back to the + detached sentinel by this (now-stale) transport's teardown. + Returns ``(reaped, detached)`` counts for disconnect-path observability.""" with _sessions_lock: - owned = [(sid, s) for sid, s in _sessions.items() if s.get("transport") is transport] + owned_sids = [sid for sid, s in _sessions.items() if s.get("transport") is transport] reaped = 0 detached = 0 - for sid, session in owned: - if session.get("close_on_disconnect"): - _close_session_by_id(sid, end_reason=end_reason) + for sid in owned_sids: + with _session_resume_lock: + with _sessions_lock: + session = _sessions.get(sid) + if session is None or session.get("transport") is not transport: + # Already torn down, or reattached to a new transport + # since the snapshot above — leave it alone. + continue + if session.get("close_on_disconnect"): + del _sessions[sid] + session["_sid"] = sid + to_teardown, to_detach = session, None + else: + # Point detached sessions at the drop sentinel (NOT real + # stdio) so _ws_session_is_orphaned recognizes them and the + # grace-reap can actually fire; a standalone `hermes --tui` + # keeps real _stdio. + session["transport"] = _detached_ws_transport + to_teardown, to_detach = None, sid + # Slow teardown/timer scheduling happens after releasing both locks — + # see the module note above _pop_session_by_id about keeping that work + # off _session_resume_lock. + if to_teardown is not None: + _teardown_popped_session(to_teardown, end_reason=end_reason) reaped += 1 else: - # Point detached sessions at the drop sentinel (NOT real stdio) so - # _ws_session_is_orphaned recognizes them and the grace-reap can - # actually fire; a standalone `hermes --tui` keeps real _stdio. - session["transport"] = _detached_ws_transport detached += 1 try: - _schedule_ws_orphan_reap(sid) + _schedule_ws_orphan_reap(to_detach) except Exception: pass return reaped, detached From eb053911e02a1ddb7032607bc793197fb2dbd6d9 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Sun, 2 Aug 2026 22:33:03 -0300 Subject: [PATCH 2/3] test(tui-gateway): reproduce actual WS reattach-race interleaving MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existing regression test for this fix started both sessions already on new_transport, so they never entered owned_sids (filtered by old_transport) and the revalidation-under-lock logic this PR adds was never exercised. The test passed identically before and after the fix and proved nothing about it. Rewrite it to start the session on old_transport (so the snapshot captures it) and inject the reattach strictly between the snapshot and the per-sid claim under _session_resume_lock, via a thin wrapper around the real lock that performs the reattach on first acquire — modeling session.resume winning the lock race before teardown's revalidation runs. Confirmed this fails against the pre-fix implementation (reaped == 1, not 0) and passes against the fix. Co-Authored-By: Claude Sonnet 5 --- tests/test_tui_gateway_server.py | 57 +++++++++++++++++++++++++------- 1 file changed, 45 insertions(+), 12 deletions(-) diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 05c1bc59cbe9..ecef03305ea1 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -16746,7 +16746,21 @@ def test_close_sessions_for_transport_skips_session_reattached_mid_teardown(monk """Regression for the disconnect/reconnect race: if session.resume rebinds a session onto a new (live) transport between the ownership snapshot and this function's per-session claim, the old transport's teardown must not - close it or stomp its transport back to the detached sentinel.""" + close it or stomp its transport back to the detached sentinel. + + Unlike a naive version of this test that starts both sessions already on + ``new_transport`` (which never even enters ``owned_sids`` and exercises + nothing beyond the initial filter), this drives the actual interleaving + the fix revalidates against: the session starts on ``old_transport`` so + the snapshot captures it, and the reattach happens strictly between that + snapshot and this function's per-sid claim under ``_session_resume_lock`` + — the exact TOCTOU window closed by the WS disconnect/reconnect fix. A + ``_RaceLock`` stand-in for the module's real resume lock performs the + reattach the first time the loop acquires it, modeling session.resume + winning the lock race before teardown's revalidation runs. Against the + pre-fix implementation (no per-sid lock/revalidation at all) the injected + mutation never fires and the session is torn down/stomped regardless — + this test fails there and passes only once the race window is closed.""" seen = [] monkeypatch.setattr( server, "_teardown_popped_session", @@ -16756,22 +16770,41 @@ def test_close_sessions_for_transport_skips_session_reattached_mid_teardown(monk old_transport = object() new_transport = object() server._sessions.clear() - # "reattached" simulates session.resume's warm-reuse winning the race and - # rebinding the transport (under _session_resume_lock) before this - # function's snapshot is acted on. - server._sessions["reattached"] = {"transport": new_transport, "close_on_disconnect": True} - server._sessions["detached-reattached"] = { - "transport": new_transport, "close_on_disconnect": False, - } + server._sessions["x"] = {"transport": old_transport, "close_on_disconnect": True} + + real_resume_lock = server._session_resume_lock + + class _RaceLock: + """Wraps the real resume lock. The first acquire simulates + session.resume winning the race: it rebinds "x" onto new_transport + right after the snapshot above already captured it as owned by + old_transport, but before this function's own per-sid claim (which + also needs this lock) gets to revalidate.""" + + def __init__(self): + self._fired = False + + def __enter__(self): + real_resume_lock.acquire() + if not self._fired: + self._fired = True + with server._sessions_lock: + server._sessions["x"]["transport"] = new_transport + return self + + def __exit__(self, *exc_info): + real_resume_lock.release() + return False + + monkeypatch.setattr(server, "_session_resume_lock", _RaceLock()) try: reaped, detached = server._close_sessions_for_transport( old_transport, end_reason="ws_disconnect" ) assert reaped == 0 and detached == 0 - assert seen == [] - assert "reattached" in server._sessions # not closed - assert server._sessions["reattached"]["transport"] is new_transport - assert server._sessions["detached-reattached"]["transport"] is new_transport # untouched + assert seen == [] # teardown must not have claimed the reattached session + assert "x" in server._sessions # not closed + assert server._sessions["x"]["transport"] is new_transport # not stomped back finally: server._sessions.clear() From 25e05dc57556831ae0915347d49c83012bab8c3b Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Wed, 19 Aug 2026 14:42:31 -0300 Subject: [PATCH 3/3] fix(tui-gateway): preserve canonical teardown claim and close detach-race window (#77129) - Claim teardown via _pop_session_by_id(sid) under _session_resume_lock to preserve the session['_closing'] = True lifecycle invariant, preventing _drain_queued_prompt from dispatching work while teardown settles. - Add deterministic regression test covering the detach/orphan path (close_on_disconnect=False) under the snapshot -> reattach -> claim interleaving. - Add end-of-flow defense-in-depth safety sweep for straggler sessions and _count_orphaned_ws_sessions helper verification to prevent orphan leaks. --- tests/test_tui_gateway_server.py | 167 +++++++++++++++++++++++++++++++ tui_gateway/server.py | 46 +++++++-- 2 files changed, 206 insertions(+), 7 deletions(-) diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index ecef03305ea1..d23b0745c663 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -16742,6 +16742,49 @@ def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch): server._sessions.clear() +def test_close_sessions_for_transport_close_claims_via_pop_sets_closing_blocks_queued_prompt(monkeypatch): + """Regression: _close_sessions_for_transport must claim close_on_disconnect + sessions via _pop_session_by_id under lock, which marks session["_closing"] = True. + This ensures that while _teardown_popped_session is settling active threads, + a concurrent or finishing turn cannot dispatch queued prompts via _drain_queued_prompt.""" + seen = [] + transport = object() + session = { + "transport": transport, + "close_on_disconnect": True, + "history_lock": threading.Lock(), + "queued_prompt": {"text": "should not drain"}, + "running": False, + } + server._sessions.clear() + server._sessions["closing_test"] = session + + def fake_teardown(popped_session, *, end_reason): + seen.append((popped_session["_sid"], end_reason)) + assert popped_session.get("_closing") is True + # Verify that _drain_queued_prompt honors the _closing invariant and refuses to dispatch + dispatched = server._drain_queued_prompt("req-1", "closing_test", popped_session) + assert dispatched is False + assert popped_session.get("running") is False + assert popped_session.get("queued_prompt") == {"text": "should not drain"} + return True + + monkeypatch.setattr(server, "_teardown_popped_session", fake_teardown) + monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) + + try: + reaped, detached = server._close_sessions_for_transport( + transport, end_reason="ws_disconnect" + ) + assert reaped == 1 and detached == 0 + assert seen == [("closing_test", "ws_disconnect")] + assert "closing_test" not in server._sessions + assert session["_sid"] == "closing_test" + assert session.get("_closing") is True + finally: + server._sessions.clear() + + def test_close_sessions_for_transport_skips_session_reattached_mid_teardown(monkeypatch): """Regression for the disconnect/reconnect race: if session.resume rebinds a session onto a new (live) transport between the ownership snapshot and @@ -16762,10 +16805,15 @@ def test_close_sessions_for_transport_skips_session_reattached_mid_teardown(monk mutation never fires and the session is torn down/stomped regardless — this test fails there and passes only once the race window is closed.""" seen = [] + reap_scheduled = [] monkeypatch.setattr( server, "_teardown_popped_session", lambda session, *, end_reason: bool(seen.append((session["_sid"], end_reason))) or True, ) + monkeypatch.setattr( + server, "_schedule_ws_orphan_reap", + lambda sid: reap_scheduled.append(sid), + ) monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) old_transport = object() new_transport = object() @@ -16803,8 +16851,127 @@ def __exit__(self, *exc_info): ) assert reaped == 0 and detached == 0 assert seen == [] # teardown must not have claimed the reattached session + assert reap_scheduled == [] assert "x" in server._sessions # not closed assert server._sessions["x"]["transport"] is new_transport # not stomped back + assert server._sessions["x"].get("_closing") is not True + finally: + server._sessions.clear() + + +def test_close_sessions_for_transport_skips_reattached_session_on_detach_path(monkeypatch): + """Regression for the detach/orphan branch of the disconnect/reconnect race: + when close_on_disconnect is False, if session.resume rebinds the session to a + new transport between snapshot and locked claim, teardown must neither close it, + nor stomp its transport back to _detached_ws_transport, nor schedule orphan reap.""" + teardown_called = [] + reap_scheduled = [] + monkeypatch.setattr( + server, "_teardown_popped_session", + lambda session, *, end_reason: bool(teardown_called.append((session["_sid"], end_reason))) or True, + ) + monkeypatch.setattr( + server, "_schedule_ws_orphan_reap", + lambda sid: reap_scheduled.append(sid), + ) + monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) + old_transport = object() + new_transport = object() + server._sessions.clear() + server._sessions["detach_race"] = {"transport": old_transport, "close_on_disconnect": False} + + real_resume_lock = server._session_resume_lock + + class _RaceLock: + def __init__(self): + self._fired = False + + def __enter__(self): + real_resume_lock.acquire() + if not self._fired: + self._fired = True + with server._sessions_lock: + server._sessions["detach_race"]["transport"] = new_transport + return self + + def __exit__(self, *exc_info): + real_resume_lock.release() + return False + + monkeypatch.setattr(server, "_session_resume_lock", _RaceLock()) + try: + reaped, detached = server._close_sessions_for_transport( + old_transport, end_reason="ws_disconnect" + ) + assert reaped == 0 and detached == 0 + assert teardown_called == [] + assert reap_scheduled == [] + assert "detach_race" in server._sessions + assert server._sessions["detach_race"]["transport"] is new_transport + assert server._sessions["detach_race"].get("_closing") is not True + finally: + server._sessions.clear() + + +def test_close_sessions_for_transport_sweeps_stragglers_defense_in_depth(monkeypatch): + """Defense-in-depth: if a straggler session attaches to the disconnecting + transport mid-teardown, the final sweep at the end of the flow must catch + and process it so no session is left pointing to a dead transport.""" + reap_scheduled = [] + seen_teardown = [] + transport = object() + + def fake_teardown(session, *, end_reason): + seen_teardown.append((session["_sid"], end_reason)) + # Simulate a concurrent action injecting a straggler attached to the same transport + with server._sessions_lock: + server._sessions["straggler"] = { + "transport": transport, + "close_on_disconnect": False, + } + return True + + monkeypatch.setattr(server, "_teardown_popped_session", fake_teardown) + monkeypatch.setattr( + server, "_schedule_ws_orphan_reap", + lambda sid: reap_scheduled.append(sid), + ) + monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) + + server._sessions.clear() + server._sessions["s1"] = {"transport": transport, "close_on_disconnect": True} + + try: + reaped, detached = server._close_sessions_for_transport( + transport, end_reason="ws_disconnect" + ) + assert reaped == 1 + assert detached == 1 + assert seen_teardown == [("s1", "ws_disconnect")] + assert reap_scheduled == ["straggler"] + assert "s1" not in server._sessions + assert server._sessions["straggler"]["transport"] is server._detached_ws_transport + # Verify no sessions remain attached to the dead transport + assert not any(s.get("transport") is transport for s in server._sessions.values()) + finally: + server._sessions.clear() + + +def test_count_orphaned_ws_sessions(): + """Verify _count_orphaned_ws_sessions correctly counts idle detached WS sessions.""" + server._sessions.clear() + try: + assert server._count_orphaned_ws_sessions() == 0 + server._sessions["live_stdio"] = {"transport": object(), "running": False} + server._sessions["live_detached"] = {"transport": server._detached_ws_transport, "running": False} + server._sessions["running_detached"] = {"transport": server._detached_ws_transport, "running": True} + server._sessions["finalized_detached"] = { + "transport": server._detached_ws_transport, + "running": False, + "_finalized": True, + } + # Only live_detached meets all criteria for orphaned WS session + assert server._count_orphaned_ws_sessions() == 1 finally: server._sessions.clear() diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 83e10785318d..0fceb1560b3f 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1167,6 +1167,12 @@ def _reap() -> None: timer.start() +def _count_orphaned_ws_sessions() -> int: + """Count how many registered sessions are currently detached and orphaned.""" + with _sessions_lock: + return sum(1 for s in _sessions.values() if _ws_session_is_orphaned(s)) + + def _close_sessions_for_transport( transport, *, end_reason: str = "ws_disconnect" ) -> tuple[int, int]: @@ -1190,41 +1196,67 @@ def _close_sessions_for_transport( detached sentinel by this (now-stale) transport's teardown. Returns ``(reaped, detached)`` counts for disconnect-path observability.""" + # Defense-in-depth (start of flow): snapshot owned sessions and inspect initial state with _sessions_lock: owned_sids = [sid for sid, s in _sessions.items() if s.get("transport") is transport] reaped = 0 detached = 0 - for sid in owned_sids: + + def _process_sid(sid: str) -> tuple[dict | None, str | None]: + to_td = None + to_dt = None with _session_resume_lock: with _sessions_lock: session = _sessions.get(sid) if session is None or session.get("transport") is not transport: # Already torn down, or reattached to a new transport # since the snapshot above — leave it alone. - continue + return None, None if session.get("close_on_disconnect"): - del _sessions[sid] - session["_sid"] = sid - to_teardown, to_detach = session, None + to_td = _pop_session_by_id(sid) else: # Point detached sessions at the drop sentinel (NOT real # stdio) so _ws_session_is_orphaned recognizes them and the # grace-reap can actually fire; a standalone `hermes --tui` # keeps real _stdio. session["transport"] = _detached_ws_transport - to_teardown, to_detach = None, sid + to_dt = sid + return to_td, to_dt + + for sid in owned_sids: + to_teardown, to_detach = _process_sid(sid) # Slow teardown/timer scheduling happens after releasing both locks — # see the module note above _pop_session_by_id about keeping that work # off _session_resume_lock. if to_teardown is not None: _teardown_popped_session(to_teardown, end_reason=end_reason) reaped += 1 - else: + elif to_detach is not None: detached += 1 try: _schedule_ws_orphan_reap(to_detach) except Exception: pass + + # Defense-in-depth (end of flow): safety sweep for stragglers that may have + # attached to the disconnecting transport during the teardown window + with _sessions_lock: + stragglers = [ + sid for sid, s in _sessions.items() + if s.get("transport") is transport + ] + for sid in stragglers: + to_teardown, to_detach = _process_sid(sid) + if to_teardown is not None: + _teardown_popped_session(to_teardown, end_reason=end_reason) + reaped += 1 + elif to_detach is not None: + detached += 1 + try: + _schedule_ws_orphan_reap(to_detach) + except Exception: + pass + return reaped, detached