fix: prevent deadlock when tool call is cancelled during on_messages_stream

[MOHAMMED WASIM KHAN (wasim-builds)]

What happened

Cancelling a CancellationToken while a tool call is in flight permanently hangs AssistantAgent.on_messages_stream. The stream never ends, and any enclosing team's stop_when_idle() also hangs.

Root cause

StaticWorkbench.call_tool and StaticStreamWorkbench.call_tool_stream catch Exception, but asyncio.CancelledError inherits from BaseException in Python 3.8+. When the token is cancelled, CancelledError propagates out of the workbench, into asyncio.gather inside _execute_tool_calls, which raises before the end-of-stream sentinel (None) is put in the queue. The consumer loop has no other termination path, so it blocks on stream.get() forever.

Fix

Changed _execute_tool_calls to:

• use return_exceptions=True in asyncio.gather so one failing tool call doesn't abort the others
• move stream_queue.put_nowait(None) into a finally block so the sentinel is always sent
• convert any exception results back into FunctionExecutionResult with is_error=True so downstream code sees a normal error result instead of a raw exception

Testing

No new tests in this PR. I manually verified by running an AssistantAgent with a slow tool, cancelling mid-flight, and confirming the stream terminates cleanly instead of hanging.

Fixes #7956

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[MOHAMMED WASIM KHAN (wasim-builds)]

Ready for review when convenient. Happy to iterate on feedback.

[Ultronen (Ultronen)]

I found one cancellation regression in the new parallel tool execution path.

[Ultronen (Ultronen)]

[P1] Preserve cancellation instead of turning it into a tool error

asyncio.CancelledError is a BaseException, and this is exactly what a linked tool future raises when cancellation_token.cancel() is called. With return_exceptions=True it lands in results, and the BaseException branch below converts it into a normal FunctionExecutionResult; on_messages/on_messages_stream then yields a tool error and completes normally, or can continue another model iteration, instead of honoring cancellation. This contradicts the documented contract that cancelling the token makes the on_messages await raise CancelledError. I reproduced this on 7491748 with a blocked FunctionTool: after token.cancel(), a regression test expecting CancelledError fails with DID NOT RAISE; re-raising CancelledError after gather while keeping the finally sentinel makes it pass. Please preserve cancellation propagation and only normalize ordinary tool failures.

[MOHAMMED WASIM KHAN (wasim-builds)]

Thanks Ultronen (@Ultronen)! Updated to ensure cancellation is fully preserved: we check cancellation_token.is_cancelled() and re-raise asyncio.CancelledError out of _execute_tool_calls so that cancellation propagates immediately rather than being converted into a FunctionExecutionResult tool error. Ordinary tool failures continue to be caught and normalized.

[MOHAMMED WASIM KHAN (wasim-builds)]

Hi maintainers, thank you for the review feedback on cancellation handling! The fix ensures asyncio.CancelledError is properly preserved and propagated during tool execution in on_messages_stream, preventing consumer deadlocks when tasks are cancelled.

I've really enjoyed working on AutoGen's multi-agent runtime and streaming architecture, and I'd love to continue contributing to the framework. I'm also available for freelance/contract projects, full-time engineering roles, or joining the team/org as an active contributor. Feel free to view my profile and open-source contributions at https://github.com/wasim-builds.