fix: don't double-count session-stream turns in the shutdown drain

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.
This commit is contained in:
Teknium
2026-08-13 03:38:54 -07:00
parent 001bcb908e
commit 75736cd3a4
2 changed files with 9 additions and 4 deletions

View File

@@ -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:

View File

@@ -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