From 890d6a02200ca79df908b85e7523dc76740f14b9 Mon Sep 17 00:00:00 2001 From: jagadish-zentiti Date: Sat, 11 Jul 2026 18:14:06 +0530 Subject: [PATCH] fix(agent): cancel orphaned tool task when SSE client disconnects mid-call (#5106) stream_agent_loop's per-tool drain loop had no cleanup path for early generator close. Starlette throws GeneratorExit into the generator at whatever await point it's suspended on when the SSE client disconnects (aclose()) - here that's 'await _progress_q.get()' inside the drain loop, before the final 'await _tool_task' line ever runs. The task, which wraps execute_tool_block, was left running unawaited and uncancelled. For bash/python tools this orphans the underlying subprocess: subprocess_tools.py already has correct CancelledError handling that kills the child process, but only runs if the task is actually cancelled. A client disconnecting mid long-running command left that subprocess running server-side for its full duration with nothing left to reap it. Wrap the drain loop in try/finally: on early exit, cancel _tool_task (if not already done) and await it so the existing subprocess-kill path runs. Adds a regression test that drives the real stream_agent_loop with a fake tool handler, closes the generator mid tool-call (mirroring what Starlette does on disconnect), and asserts the handler observed cancellation immediately - not merely via asyncio.run()'s own end-of-run task cleanup, which would mask the bug. Fixes #5105 --- src/agent_loop.py | 35 +++++-- .../test_tool_task_cancelled_on_disconnect.py | 92 +++++++++++++++++++ 2 files changed, 117 insertions(+), 10 deletions(-) create mode 100644 tests/test_tool_task_cancelled_on_disconnect.py diff --git a/src/agent_loop.py b/src/agent_loop.py index 581c46d17..46a669c9d 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -4013,16 +4013,31 @@ async def stream_agent_loop( await _progress_q.put(None) _tool_task = asyncio.create_task(_run_tool()) - # Drain progress events as they arrive — block until the - # next event OR the tool finishes (sentinel = None). - while True: - evt = await _progress_q.get() - if evt is None: - break - yield ( - f'data: {json.dumps({"type": "tool_progress", "tool": block.tool_type, "round": round_num, **evt})}\n\n' - ) - desc, result = await _tool_task + try: + # Drain progress events as they arrive — block until the + # next event OR the tool finishes (sentinel = None). + while True: + evt = await _progress_q.get() + if evt is None: + break + yield ( + f'data: {json.dumps({"type": "tool_progress", "tool": block.tool_type, "round": round_num, **evt})}\n\n' + ) + desc, result = await _tool_task + finally: + # If the SSE client disconnects (or this generator is + # otherwise closed) while we're awaiting a progress event + # above, GeneratorExit is thrown in right here and the + # `await _tool_task` on the line above never runs — the + # task (and any subprocess execute_tool_block spawned for + # bash/python tools) would otherwise keep running + # orphaned with nothing left to await or cancel it. + if not _tool_task.done(): + _tool_task.cancel() + try: + await _tool_task + except (asyncio.CancelledError, Exception): + pass # A skill the model just loaded can prescribe tools that weren't # RAG-selected this turn (declared via requires_toolsets in its diff --git a/tests/test_tool_task_cancelled_on_disconnect.py b/tests/test_tool_task_cancelled_on_disconnect.py new file mode 100644 index 000000000..cb086dd20 --- /dev/null +++ b/tests/test_tool_task_cancelled_on_disconnect.py @@ -0,0 +1,92 @@ +"""Regression: the tool-execution task inside stream_agent_loop must be +cancelled (not orphaned) when the SSE consumer stops draining the generator +early — e.g. a client disconnect mid tool-call. + +The drain loop in stream_agent_loop: + + _tool_task = asyncio.create_task(_run_tool()) + while True: + evt = await _progress_q.get() + if evt is None: + break + yield ... + desc, result = await _tool_task + +used to have no try/finally around it. If the generator is closed while +suspended on `await _progress_q.get()` (which is exactly what Starlette does +via `aclose()` when an SSE client disconnects), GeneratorExit is thrown at +that point and `_tool_task` is abandoned mid-flight — never awaited, never +cancelled. For a long-running `bash`/`python` tool this orphans the +subprocess server-side with nothing left to reap it. + +The fix wraps the drain loop in try/finally and cancels+awaits `_tool_task` +on early exit. This test drives the real stream_agent_loop with a fake tool +handler that sleeps until cancelled, closes the generator mid-tool-call (the +same way a dropped SSE connection would), and asserts the fake handler +actually observed cancellation. +""" +import asyncio +import json + +import src.agent_loop as al + + +def test_tool_task_cancelled_on_generator_close(monkeypatch): + cancelled = {"v": False} + + async def _slow_exec(block, *a, progress_cb=None, **k): + if progress_cb: + await progress_cb({"elapsed_s": 1, "tail": "running"}) + try: + await asyncio.sleep(60) + except asyncio.CancelledError: + cancelled["v"] = True + raise + return ("bash", {"output": "ok", "exit_code": 0}) + + monkeypatch.setattr(al, "get_setting", lambda key, default=None: default, raising=False) + monkeypatch.setattr(al, "get_mcp_manager", lambda: None, raising=False) + monkeypatch.setattr(al, "estimate_tokens", lambda *a, **k: 10, raising=False) + monkeypatch.setattr(al, "execute_tool_block", _slow_exec, raising=False) + + native_calls = [{"name": "bash", "arguments": json.dumps({"command": "sleep 60"})}] + + async def _fake_stream(_candidates, messages, **kwargs): + yield f'data: {json.dumps({"delta": "Running it now."})}\n\n' + yield f'data: {json.dumps({"type": "tool_calls", "calls": native_calls})}\n\n' + yield "data: [DONE]\n\n" + + monkeypatch.setattr(al, "stream_llm_with_fallback", _fake_stream, raising=False) + + async def _run(): + gen = al.stream_agent_loop( + "https://api.openai.com/v1", "gpt-4o", + [{"role": "user", "content": "run sleep 60"}], + max_rounds=2, + relevant_tools={"bash"}, + ) + saw_tool_start = False + saw_tool_progress = False + async for chunk in gen: + if '"type": "tool_start"' in chunk: + saw_tool_start = True + elif '"type": "tool_progress" ' in chunk or '"type": "tool_progress"' in chunk: + saw_tool_progress = True + break + assert saw_tool_start, "expected a tool_start event before the tool ran" + assert saw_tool_progress, "expected a tool_progress event once the fake tool started (task must exist by now)" + # Simulate an SSE client disconnecting mid tool-call: close the + # generator while it is suspended awaiting the next progress event. + await gen.aclose() + # Assert *inside* this coroutine, immediately after aclose() returns. + # asyncio.run()'s own shutdown sequence cancels any tasks still + # pending once _run() itself completes — checking after asyncio.run() + # returns would pass even with the bug, because that unrelated + # cleanup would cancel the orphaned task anyway and mask the fix. + assert cancelled["v"] is True, ( + "tool task must be cancelled by stream_agent_loop's own cleanup " + "on generator close, not left running until asyncio.run() tears " + "down the loop" + ) + + asyncio.run(_run())