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:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user