fix: resolve rebase conflict in _teardown_session worker cleanup

Main folded slash_worker.close() into _finalize_session (the single
_finalized-guarded chokepoint) while #42143 was open. The rebase
conflicted with the PR's worker-close in _teardown_session. Keep both —
they target the same #38095 leak and _SlashWorker.close() is
idempotent (_closed/poll()-guarded) — so callers reaching
_teardown_session without the real _finalize_session (and the PR's own
tests, which monkeypatch _finalize_session out) still reap the worker.
Same for _shutdown_sessions, now routed through the unified
_close_session_by_id funnel.
This commit is contained in:
teknium1
2026-06-08 10:02:05 -07:00
parent ae94ed1728
commit 365813a72b
3 changed files with 37 additions and 12 deletions
+15 -6
View File
@@ -6200,10 +6200,16 @@ def test_close_session_by_id_is_idempotent_and_full(monkeypatch):
def close(self): def close(self):
calls["agent"] += 1 calls["agent"] += 1
monkeypatch.setattr( def _fake_finalize(s, end_reason="tui_close"):
server, "_finalize_session", # Real _finalize_session is the single chokepoint that closes the
lambda s, end_reason="tui_close": calls.__setitem__("finalize", calls["finalize"] + 1), # slash-worker; mirror that here so the test exercises the actual
) # teardown contract (worker close lives in finalize, not the caller).
calls["finalize"] += 1
w = s.get("slash_worker")
if w:
w.close()
monkeypatch.setattr(server, "_finalize_session", _fake_finalize)
monkeypatch.setattr( monkeypatch.setattr(
"tools.approval.unregister_gateway_notify", "tools.approval.unregister_gateway_notify",
lambda key: calls.__setitem__("unreg", calls["unreg"] + 1), raising=False, lambda key: calls.__setitem__("unreg", calls["unreg"] + 1), raising=False,
@@ -6307,6 +6313,9 @@ def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch):
server, "_close_session_by_id", server, "_close_session_by_id",
lambda sid, *, end_reason: bool(seen.append((sid, end_reason))) or True, lambda sid, *, end_reason: bool(seen.append((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.
monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)
transport = object() # the disconnecting transport transport = object() # the disconnecting transport
server._sessions.clear() server._sessions.clear()
server._sessions["a"] = {"transport": transport, "close_on_disconnect": True} server._sessions["a"] = {"transport": transport, "close_on_disconnect": True}
@@ -6314,7 +6323,7 @@ def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch):
try: try:
server._close_sessions_for_transport(transport, end_reason="ws_disconnect") server._close_sessions_for_transport(transport, end_reason="ws_disconnect")
assert seen == [("a", "ws_disconnect")] # only the flagged one closed assert seen == [("a", "ws_disconnect")] # only the flagged one closed
assert server._sessions["b"]["transport"] is server._stdio_transport # re-pointed assert server._sessions["b"]["transport"] is server._detached_ws_transport # re-pointed
finally: finally:
server._sessions.clear() server._sessions.clear()
@@ -6360,7 +6369,7 @@ def _idle_evictable_session(now):
return { return {
"running": False, "running": False,
"agent_ready": ready, "agent_ready": ready,
"transport": server._stdio_transport, # dead/detached "transport": server._detached_ws_transport, # dead/detached
"last_active": old, "last_active": old,
"created_at": old, "created_at": old,
} }
+18 -2
View File
@@ -7,7 +7,23 @@ from tui_gateway import ws as ws_mod
def _run_disconnect(monkeypatch, seed): def _run_disconnect(monkeypatch, seed):
"""Drive handle_ws to its disconnect `finally`, seeding sessions against the """Drive handle_ws to its disconnect `finally`, seeding sessions against the
live WSTransport the moment it exists. Returns nothing; inspect _sessions.""" live WSTransport the moment it exists. Returns nothing; inspect _sessions."""
monkeypatch.setattr(server, "_finalize_session", lambda s, end_reason="tui_close": None) # Disable the grace-reap Timer: detached sessions normally schedule a
# threading.Timer via _schedule_ws_orphan_reap, which would outlive the test
# and fire _reap during interpreter teardown — touching _sessions/DB and
# producing spurious post-run errors under the per-file CI runner. Grace=0
# short-circuits the Timer (see _schedule_ws_orphan_reap) so the test leaves
# no lingering thread.
monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0)
# Mirror the real _finalize_session chokepoint: it is the single place that
# closes the slash-worker (#38095). Stub it but keep that behavior so the
# disconnect-reap path still exercises worker teardown.
def _fake_finalize(s, end_reason="tui_close"):
w = s.get("slash_worker")
if w:
w.close()
monkeypatch.setattr(server, "_finalize_session", _fake_finalize)
created = [] created = []
real_transport = ws_mod.WSTransport real_transport = ws_mod.WSTransport
@@ -68,6 +84,6 @@ def test_ws_disconnect_preserves_and_repoints_reconnectable_session(monkeypatch)
plain={"transport": t, "close_on_disconnect": False, "session_key": "k"} plain={"transport": t, "close_on_disconnect": False, "session_key": "k"}
), ),
) )
assert server._sessions["plain"]["transport"] is server._stdio_transport assert server._sessions["plain"]["transport"] is server._detached_ws_transport
finally: finally:
server._sessions.clear() server._sessions.clear()
+4 -4
View File
@@ -425,10 +425,10 @@ def _teardown_session(session: dict | None, *, end_reason: str = "tui_close") ->
agent.close() agent.close()
except Exception: except Exception:
pass pass
# NOTE: the slash-worker subprocess is already closed inside # NOTE: the slash-worker is closed inside _finalize_session (the single
# _finalize_session (the single _finalized-guarded chokepoint on main). # _finalized-guarded chokepoint that main folded it into), exactly once.
# No second worker.close() here — it would be redundant (poll()-guarded, # We deliberately do NOT re-close it here — _teardown_session's job beyond
# harmless, but dead). # finalize is unregistering the notifier and closing the in-process agent.
def _attach_worker(sid: str, session: dict, worker) -> None: def _attach_worker(sid: str, session: dict, worker) -> None: