From 75736cd3a4dcf50a370252eee36e26ac1936d3b2 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Thu, 13 Aug 2026 03:38:54 -0700
Subject: [PATCH] fix: don't double-count session-stream turns in the shutdown
 drain
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The session chat stream registered its wrapper task in _active_run_tasks,
but that turn is already counted by active_agent_work_count() via
_inflight_agent_runs (_run_agent) — the drain saw 2 for one turn
(test_session_chat_sse_turn_is_interrupted). Keep only the agent-ref
registration; run-scoped steer control doesn't need the task entry.
---
 gateway/platforms/api_server.py   | 7 +++++--
 tests/gateway/test_session_api.py | 6 ++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py
index 415a7dd8e5..53fb4a34f1 100644
--- a/gateway/platforms/api_server.py
+++ b/gateway/platforms/api_server.py
@@ -3982,12 +3982,15 @@ class APIServerAdapter(BasePlatformAdapter):
                 await queue.put(_event_payload("error", {"message": _redact_api_error_text(exc)}))
             finally:
                 self._active_run_agents.pop(run_id, None)
-                self._active_run_tasks.pop(run_id, None)
                 await queue.put(_event_payload("done", {}))
                 await queue.put(None)
 
+        # NOTE: deliberately NOT registered in _active_run_tasks — this turn
+        # is already counted by active_agent_work_count() via
+        # _inflight_agent_runs (_run_agent), and a second task-based entry
+        # would double-count it in the shutdown drain. Run-scoped control
+        # needs only the agent ref, registered by _run_agent(active_run_id).
         task = asyncio.create_task(_run_and_signal())
-        self._active_run_tasks[run_id] = task
         try:
             self._background_tasks.add(task)
         except TypeError:
diff --git a/tests/gateway/test_session_api.py b/tests/gateway/test_session_api.py
index 189010e193..5e93d74371 100644
--- a/tests/gateway/test_session_api.py
+++ b/tests/gateway/test_session_api.py
@@ -286,14 +286,16 @@ async def test_session_chat_stream_disconnect_keeps_control_refs_until_executor_
 
         assert interrupt_called.is_set()
         assert run_id in adapter._active_run_agents
-        assert run_id in adapter._active_run_tasks
+        # Not in _active_run_tasks: session-stream turns are counted via
+        # _inflight_agent_runs; a task entry would double-count them in the
+        # shutdown drain (active_agent_work_count).
+        assert run_id not in adapter._active_run_tasks
         assert not handler_task.done()
 
         allow_finish.set()
         await handler_task
 
     assert run_id not in adapter._active_run_agents
-    assert run_id not in adapter._active_run_tasks
 
 
 @pytest.mark.asyncio