fix(tui-gateway): WS disconnect/reconnect TOCTOU fix + real regression test (RAH-05 + RAH-06)

tests/test_tui_gateway_server.py

@@ -14204,8 +14204,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.
@@ -14217,11 +14217,79 @@ 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.
+
+    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",
+        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()
+    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 "x" in server._sessions  # not closed
+        assert server._sessions["x"]["transport"] is new_transport  # not stomped back
+    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

@@ -1066,9 +1066,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
@@ -1077,23 +1076,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