fix(tui-gateway): close WS disconnect/reconnect session race

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,246 @@ 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_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
+    this function's per-session claim, the old transport's teardown must not
+    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 = []
+    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()
+    server._sessions.clear()
+    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 == []  # 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()
+
+
 def test_session_create_records_close_on_disconnect_flag(monkeypatch):
     monkeypatch.setattr(server, "_start_agent_build", lambda sid, session: None)
     server._sessions.clear()

tui_gateway/server.py

@@ -1167,13 +1167,18 @@ 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]:
     """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,25 +1187,76 @@ 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."""
+    # Defense-in-depth (start of flow): snapshot owned sessions and inspect initial state
     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)
+
+    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.
+                    return None, None
+                if session.get("close_on_disconnect"):
+                    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_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:
-            # 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
+        elif to_detach is not None:
             detached += 1
             try:
-                _schedule_ws_orphan_reap(sid)
+                _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