From 468b61e513f012c952a35d72ade0aa69f710c786 Mon Sep 17 00:00:00 2001 From: nesquena-hermes Date: Tue, 30 Jun 2026 04:39:00 +0000 Subject: [PATCH] fix(auth): release STREAM_SESSION_OWNERS on worker early-return + stale-run prune MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gate hardening for #5198: the new stream-owner map leaked entries on three paths that bypass the teardown finally — worker early-return when the stream was cancelled before startup (streaming.py, gateway_chat.py) and the direct ACTIVE_RUNS stale-zombie prune (routes.py). Each now calls unregister_stream_owner. Adds a regression test proving the early-return path no longer leaks. Co-authored-by: starship-s --- api/gateway_chat.py | 4 +++ api/routes.py | 5 +++ api/streaming.py | 5 +++ ...st_session_active_profile_authorization.py | 31 +++++++++++++++++++ 4 files changed, 45 insertions(+) diff --git a/api/gateway_chat.py b/api/gateway_chat.py index 2c67e340..0a93eb4e 100644 --- a/api/gateway_chat.py +++ b/api/gateway_chat.py @@ -24,6 +24,7 @@ from api.config import ( gateway_supports_approval, register_active_run, unregister_active_run, + unregister_stream_owner, update_active_run, ) from api.helpers import _redact_text, redact_session_data @@ -513,6 +514,9 @@ def _run_gateway_chat_streaming( """ q = STREAMS.get(stream_id) if q is None: + # Cancelled before the worker started; release the owner entry the route + # layer registered so STREAM_SESSION_OWNERS does not leak (no teardown finally runs). + unregister_stream_owner(stream_id) return register_active_run( stream_id, diff --git a/api/routes.py b/api/routes.py index fd6491f6..01825de1 100644 --- a/api/routes.py +++ b/api/routes.py @@ -2351,6 +2351,7 @@ from api.config import ( ACTIVE_RUNS_LOCK, register_stream_owner, stream_owner_session_id, + unregister_stream_owner, CHAT_LOCK, _get_session_agent_lock, SESSION_AGENT_LOCKS, @@ -17479,6 +17480,10 @@ def _active_run_stream_for_session(session_id: str | None) -> str | None: return stream_id for stale_stream_id in stale_stream_ids: (_live_config.ACTIVE_RUNS or {}).pop(stale_stream_id, None) + # The zombie run is pruned directly here (not via the normal teardown + # finally / unregister_active_run), so release its stream-owner entry too + # or STREAM_SESSION_OWNERS leaks for every reconciled zombie. (#5198 gate) + unregister_stream_owner(stale_stream_id) except Exception: return None return None diff --git a/api/streaming.py b/api/streaming.py index e18c6aec..04211a24 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -32,6 +32,7 @@ from api.config import ( LOCK, SESSIONS, SESSIONS_MAX, SESSION_DIR, _get_session_agent_lock, _set_thread_env, _clear_thread_env, register_active_run, update_active_run, unregister_active_run, + unregister_stream_owner, SESSION_AGENT_LOCKS, SESSION_AGENT_LOCKS_LOCK, resolve_model_provider, resolve_custom_provider_connection, @@ -6056,6 +6057,10 @@ def _run_agent_streaming( """ q = STREAMS.get(stream_id) if q is None: + # The stream was cancelled before the worker started; the route layer + # already registered the stream owner, so release it here to avoid + # leaking a STREAM_SESSION_OWNERS entry that the teardown finally never sees. + unregister_stream_owner(stream_id) return register_active_run( stream_id, diff --git a/tests/test_session_active_profile_authorization.py b/tests/test_session_active_profile_authorization.py index 1451cbe8..aec6b3d8 100644 --- a/tests/test_session_active_profile_authorization.py +++ b/tests/test_session_active_profile_authorization.py @@ -511,3 +511,34 @@ def test_session_new_keeps_prev_session_commit_for_same_profile(monkeypatch): assert calls["commit"] == 1 assert calls["new"] == 1 assert cap["ok"]["session"]["session_id"] == "new_session" + + +def test_stream_owner_unregistered_on_worker_early_return(monkeypatch): + """A stream cancelled before the worker starts (q is None) must not leak its owner entry. + + The route layer registers the stream owner synchronously before launching the + worker. If the stream is cancelled before the worker reaches `q = STREAMS.get(...)`, + the worker early-returns before register_active_run / the teardown finally — so the + owner entry must be released on that early-return path or STREAM_SESSION_OWNERS leaks. + """ + from api import config + import api.streaming as streaming + + with config.STREAM_SESSION_OWNERS_LOCK: + previous = dict(config.STREAM_SESSION_OWNERS) + config.STREAM_SESSION_OWNERS.clear() + # Owner registered by the route layer, but the stream was already cancelled + # (never placed in STREAMS), so the worker will hit `q is None`. + config.register_stream_owner("leak-stream", "some_session") + try: + with config.STREAMS_LOCK: + config.STREAMS.pop("leak-stream", None) + streaming._run_agent_streaming( + "some_session", "hi", "m", "/tmp", "leak-stream", + ) + with config.STREAM_SESSION_OWNERS_LOCK: + assert "leak-stream" not in config.STREAM_SESSION_OWNERS + finally: + with config.STREAM_SESSION_OWNERS_LOCK: + config.STREAM_SESSION_OWNERS.clear() + config.STREAM_SESSION_OWNERS.update(previous)