From 3680da238c4d9cd68123452d2bc3f8f524046e9d Mon Sep 17 00:00:00 2001 From: Max Freedom Pollard <272618364+MaxFreedomPollard@users.noreply.github.com> Date: Sat, 25 Jul 2026 00:13:03 -0400 Subject: [PATCH] fix(tui): close agents orphaned during session build Ported to the refactored build path: the attach is now _attach_built_agent, which takes the liveness check under _sessions_lock and returns whether it attached; _build closes the orphan and returns before any worker/notify/ callback is wired for the dead session. test_session_db_ownership_teardown's reaped-midbuild case asserted the discarded agent through session["agent"]; that record no longer holds it, so it reads ownership off the agent the build returned instead. --- tests/test_tui_gateway_server.py | 44 +++++++++++++------ .../test_session_db_ownership_teardown.py | 10 ++++- tui_gateway/server.py | 27 +++++++++--- 3 files changed, 59 insertions(+), 22 deletions(-) diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 0441de28e1a3..6e669d935f56 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -14531,24 +14531,29 @@ def test_mirror_slash_compress_honors_here_argument(monkeypatch): # --------------------------------------------------------------------------- # session.create / session.close race: fast /new churn must not orphan the -# global approval-notify registration. (Slash workers are no longer pre-warmed -# by the build thread — slash.exec spawns them on demand — so the build thread -# must ALSO never construct one here.) +# agent or the global approval-notify registration. (Slash workers are no +# longer pre-warmed by the build thread - slash.exec spawns them on demand - +# so the build thread must ALSO never construct one here.) # --------------------------------------------------------------------------- @pytest.mark.real_agent_prewarm -def test_session_create_close_race_does_not_orphan_worker(monkeypatch): +def test_session_create_close_race_does_not_orphan_resources(monkeypatch): """Regression guard: if session.close runs while session.create's _build thread is still constructing the agent, the build thread - must detect the orphan and unregister the notify registration it's - about to install. It must also never pre-warm a slash worker (each - worker forks the full stdio MCP fleet; spawn is on-demand in - slash.exec) — a worker constructed here would be a regression.""" + must detect the orphan, close the just-built agent, and unregister + the notify registration it's about to install - avoiding installing + a slash_worker or notify callback for the dead session. It must also + never pre-warm a slash worker (each worker forks the full stdio MCP + fleet; spawn is on-demand in slash.exec) - a worker constructed here + would be a regression. Without the early abort the agent outlives + the session until gateway shutdown.""" import threading created_workers: list[str] = [] closed_workers: list[str] = [] + closed_agents: list[str] = [] + registered_keys: list[str] = [] unregistered_keys: list[str] = [] class _FakeWorker: @@ -14568,6 +14573,9 @@ def __init__(self): self.base_url = "" self.api_key = "" + def close(self): + closed_agents.append("closed") + # Make _build block until we release it — simulates slow agent init. # Also signal when _build actually reaches _make_agent so the test # can close the session at the right moment: session.create now @@ -14601,7 +14609,11 @@ def _slow_make_agent(sid, key, session_id=None, session_db=None, **_kwargs): # Shim register/unregister to observe leaks import tools.approval as _approval - monkeypatch.setattr(_approval, "register_gateway_notify", lambda key, cb: None) + monkeypatch.setattr( + _approval, + "register_gateway_notify", + lambda key, cb: registered_keys.append(key), + ) monkeypatch.setattr( _approval, "unregister_gateway_notify", @@ -14640,20 +14652,24 @@ def _slow_make_agent(sid, key, session_id=None, session_db=None, **_kwargs): ) assert close_resp.get("result", {}).get("closed") is True - # At this point session.close saw slash_worker=None (never eagerly - # installed) so it had nothing to close. Release the build thread - # and let it finish — it should detect the orphan and unregister - # the notify, without ever having constructed a worker. + # At this point session.close saw agent=None and slash_worker=None (never + # eagerly installed) so it had nothing to close. Release the build thread: + # it must close the agent that finishes late and abort before registering + # any more session resources - without ever having constructed a worker. release_build.set() # Give the build thread a moment to run through its finally. for _ in range(100): - if own_key in unregistered_keys: + if closed_agents: break import time time.sleep(0.02) + assert closed_agents == ["closed"], ( + "agent built after session.close was never closed - " + f"closed_agents={closed_agents}" + ) assert created_workers == [], ( f"build thread pre-warmed a slash worker (spawn must stay on-demand " f"in slash.exec) — created_workers={created_workers}" diff --git a/tests/tui_gateway/test_session_db_ownership_teardown.py b/tests/tui_gateway/test_session_db_ownership_teardown.py index 63aa67cbfd4e..e8e646959873 100644 --- a/tests/tui_gateway/test_session_db_ownership_teardown.py +++ b/tests/tui_gateway/test_session_db_ownership_teardown.py @@ -407,11 +407,14 @@ def test_deferred_build_closes_the_handle_when_the_session_is_reaped_midbuild( handle has to be closed right here instead of handed over. """ + built = [] + def _fake_make_agent(sid, key, session_db=None, **_kwargs): # Simulate a concurrent reap landing while the agent was being built. with server._sessions_lock: server._sessions[sid] = {"session_key": "someone-else"} - return types.SimpleNamespace(_session_db=session_db, _owns_session_db=False) + built.append(types.SimpleNamespace(_session_db=session_db, _owns_session_db=False)) + return built[-1] monkeypatch.setattr(server, "_make_agent", _fake_make_agent) sid, session = "sid-reaped", _session(build_env.profile_home) @@ -421,7 +424,10 @@ def _fake_make_agent(sid, key, session_db=None, **_kwargs): db = build_env.opened[0] assert db.closed == 1 - assert session["agent"]._owns_session_db is False + # The orphaned agent is closed and dropped rather than attached to the reaped record + # (#49852), so ownership is read off the agent itself. + assert built[0]._owns_session_db is False + assert "agent" not in session def test_deferred_build_never_opens_or_closes_for_the_launch_profile( diff --git a/tui_gateway/server.py b/tui_gateway/server.py index cbfe41572dbc..691a62723f86 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -959,15 +959,22 @@ def _await_resume_history(sid: str, current: dict) -> bool: return _sessions.get(sid) is current -def _attach_built_agent(current: dict, agent) -> None: - """Attach a freshly built agent to its live record (session DB row deferred to first run_conversation()).""" +def _attach_built_agent(sid: str, current: dict, agent) -> bool: + """Attach a freshly built agent to its live record (session DB row deferred to first run_conversation()). + False when ``session.close`` popped this record mid-build: teardown saw ``agent=None`` and closed + nothing, so the caller owns closing the orphan (#49852).""" # Bot Mode gate hint: the DB title lands post-first-turn but the system prompt builds at turn START. if _title_hint := str(current.get("pending_title") or "").strip(): agent._session_title_hint = _title_hint - current["agent"] = agent + # Under the same lock session.close takes to pop the record: no window between "still live" and "attached". + with _sessions_lock: + if _sessions.get(sid) is not current: + return False + current["agent"] = agent _session_todo_state(current) # Baseline for the per-turn config sync (profile home override still active). current["config_model_seen"] = _config_model_target() + return True def _announce_built_agent(sid: str, key: str, current: dict, agent) -> None: @@ -989,8 +996,8 @@ def _finish_agent_build(sid: str, key: str, current: dict, *, notify_registered: """Release build scopes and settle ownership of the late notify registration + dedicated db handle.""" if scopes is not None: _release_build_profile_scopes(scopes) - # Reaped mid-build: _attach_worker closed the worker; only a late notify registration can still - # leak (session.close unregistered before _build registered). + # Reaped after the agent was attached: _attach_worker closed the worker; only a late notify + # registration can still leak (session.close unregistered before _build registered). with _sessions_lock: replaced = _sessions.get(sid) is not current if replaced and notify_registered: @@ -1049,7 +1056,15 @@ def _build() -> None: agent = _make_agent(sid, key, **_deferred_build_agent_kwargs(current, session_db)) finally: _clear_session_context(tokens) - _attach_built_agent(current, agent) + # Attach atomically against session teardown: ``session.close`` may have popped this + # session while the expensive build was in flight, in which case teardown could not close + # an agent that did not exist yet. Release the orphan immediately and do not keep wiring + # workers/callbacks for a dead session (#49852). + if not _attach_built_agent(sid, current, agent): + with contextlib.suppress(Exception): + if hasattr(agent, "close"): + agent.close() + return # No eager slash-worker pre-warm (slash.exec spawns on demand): each worker forks the full stdio # MCP fleet, and live-transport sessions are never reaped, so fleets would accumulate. notify_registered = _wire_session_agent(sid, key, agent)