fix: close pending tool calls on stream error#4963
fix: close pending tool calls on stream error#4963proever wants to merge 4 commits intopydantic:mainfrom
Conversation
When an exception (e.g. UnexpectedModelBehavior from retry exhaustion) kills an agent run mid-tool-call, the UI adapter emits an ErrorChunk but never closes the pending tool call with a ToolOutputError. This leaves the tool showing as "running" forever in the frontend. Track pending function tool calls in _pending_tool_calls on the base UIEventStream. On error, synthesize failed FunctionToolResultEvents for all pending calls before emitting the stream-level error. Also drain _final_result_event (pending output-tool calls) on the error path, matching the existing success-path cleanup. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
| _pending_tool_calls: dict[str, str] = field(default_factory=dict) | ||
| """Function tool calls dispatched but not yet completed. Maps tool_call_id to tool_name.""" |
There was a problem hiding this comment.
🚩 Dangling tool calls from PartStartEvent are not covered by this fix
The _pending_tool_calls tracking only covers tool calls that reach the FunctionToolCallEvent stage (i.e., the agent has dispatched the call). Tool calls that only reach the PartStartEvent/PartEndEvent stage (model response says 'I want to call tool X' but the agent throws before dispatching) remain dangling in the UI. This is visible in test_run_stream_response_error (both tests/test_vercel_ai.py:2258-2264 and tests/test_ui.py:531-533) where the second unknown_tool call shows tool-input-start + tool-input-available but no closing tool-output-error. This is a pre-existing limitation and not a regression from this PR, but it means the fix is partial — some UI frameworks may still show spinner states for tool calls that the agent never actually dispatched.
Was this helpful? React with 👍 or 👎 to provide feedback.
68d51cd to
c3868d1
Compare
Add test_run_stream_output_tool_error in both test_ui.py
and test_vercel_ai.py, covering the _final_result_event
drain and _turn_to('request') yield in error cleanup.
Remove pragmas — DummyUIAdapter yields from turn hooks,
so the yield body is exercised by the test_ui.py test.
| if ( | ||
| self._final_result_event | ||
| and (tool_call_id := self._final_result_event.tool_call_id) | ||
| and (tool_name := self._final_result_event.tool_name) | ||
| ): | ||
| self._final_result_event = None | ||
| self._pending_tool_calls[tool_call_id] = tool_name |
There was a problem hiding this comment.
🚩 _final_result_event drain may re-add already-resolved output tool in edge case
In the error handler at lines 219-225, when _final_result_event is set with a tool_call_id, that ID is unconditionally added to _pending_tool_calls. In theory, if an output tool had a retry cycle (FunctionToolCallEvent + FunctionToolResultEvent balanced pair) and THEN the run fails for an unrelated reason (e.g., a function tool fails in exhaustive end_strategy), the output tool's ID was already removed from _pending_tool_calls by the FunctionToolResultEvent pop at line 206, but the _final_result_event drain at line 225 would re-add it, causing a spurious error result to be emitted for an already-resolved tool call. This requires a very specific combination: exhaustive strategy + output tool with retries + subsequent function tool failure. I verified that common scenarios (retries=0 failure, first-try success + on_complete error) are all handled correctly.
Was this helpful? React with 👍 or 👎 to provide feedback.
| elif isinstance(event, FinalResultEvent): | ||
| self._final_result_event = event | ||
|
|
||
| if isinstance(event, FunctionToolResultEvent): |
There was a problem hiding this comment.
Let's use elif here and the one below
| result=ToolReturnPart( | ||
| tool_call_id=tool_call_id, | ||
| tool_name=tool_name, | ||
| content=str(exc), |
There was a problem hiding this comment.
I don't think we should send the exception content to the frontend by default, I'd rather have a generic message. If the user/model are meant to see the exception text, there's ModelRetry and ToolFailed (coming in #2586). There are example generic messages in _agent_graph.py
Use `elif` for mutually exclusive event type checks and replace raw `str(exc)` with a generic message in pending tool call closures so exception details aren't leaked to the frontend. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
Thanks for taking a look @DouweM, should be all fixed now. Let me know if there's anything else! |
Summary
_pending_tool_callson the baseUIEventStreamFunctionToolResultEvents for all pending calls before emitting the stream-level error_final_result_event(pending output-tool calls) on the error path, matching the existing success-path cleanupTest plan
Updated
test_run_stream_request_errorsnapshot intest_ui.py(base adapter now emits tool result before error)Updated
test_run_stream_request_errorsnapshot intest_vercel_ai.py(now expectstool-output-errorbeforeerror)Added
test_run_stream_tool_retry_exhaustionintest_vercel_ai.py— verifies every started tool call gets closed on retry exhaustionAll 127 tests pass (
test_ui.py+test_vercel_ai.py)Ruff lint and format clean
AG-UI adapter test coverage (follow-up — fix applies automatically via base class)
AI generated code