fix/issue 7956 assistant cancellation hang

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py

@@ -1197,22 +1197,26 @@ async def _execute_tool_calls(
                 function_calls: List[FunctionCall],
                 stream_queue: asyncio.Queue[BaseAgentEvent | BaseChatMessage | None],
             ) -> List[Tuple[FunctionCall, FunctionExecutionResult]]:
-                results = await asyncio.gather(
-                    *[
-                        cls._execute_tool_call(
-                            tool_call=call,
-                            workbench=workbench,
-                            handoff_tools=handoff_tools,
-                            agent_name=agent_name,
-                            cancellation_token=cancellation_token,
-                            stream=stream_queue,
-                        )
-                        for call in function_calls
-                    ]
-                )
-                # Signal the end of streaming by putting None in the queue.
-                stream_queue.put_nowait(None)
-                return results
+                try:
+                    results = await asyncio.gather(
+                        *[
+                            cls._execute_tool_call(
+                                tool_call=call,
+                                workbench=workbench,
+                                handoff_tools=handoff_tools,
+                                agent_name=agent_name,
+                                cancellation_token=cancellation_token,
+                                stream=stream_queue,
+                            )
+                            for call in function_calls
+                        ]
+                    )
+                    return results
+                finally:
+                    # Signal the end of streaming by putting None in the queue.
+                    # Using finally ensures the queue always terminates, even when
+                    # the task is cancelled, preventing the consumer from hanging.
+                    stream_queue.put_nowait(None)
 
             task = asyncio.create_task(_execute_tool_calls(current_model_result.content, stream))
 

python/packages/autogen-agentchat/tests/test_assistant_agent.py

@@ -2816,6 +2816,71 @@ async def test_reset_with_cancellation_token(self) -> None:
         # Context clear should be called
         mock_context.clear.assert_called_once()
 
+    @pytest.mark.asyncio
+    async def test_tool_cancellation_terminates_stream(self) -> None:
+        """Regression test for #7956: cancelling a tool call terminates the stream.
+
+        Prior to the fix, _execute_tool_calls placed the stream terminator
+        (None) after the await, so when asyncio.gather raised CancelledError
+        the terminator was never enqueued and on_messages_stream hung forever.
+        The fix wraps gather in try/finally, ensuring the queue always
+        terminates.
+        """
+        # A tool that blocks indefinitely — cancellation is the only way out.
+        async def hanging_tool(param: str) -> str:
+            await asyncio.Event().wait()  # Never finishes on its own
+            return f"Result: {param}"
+
+        model_client = ReplayChatCompletionClient(
+            [
+                CreateResult(
+                    finish_reason="function_calls",
+                    content=[FunctionCall(id="1", arguments=json.dumps({"param": "test"}), name="hanging_tool")],
+                    usage=RequestUsage(prompt_tokens=10, completion_tokens=5),
+                    cached=False,
+                ),
+            ],
+            model_info={
+                "function_calling": True,
+                "vision": False,
+                "json_output": False,
+                "family": ModelFamily.GPT_4O,
+                "structured_output": False,
+            },
+        )
+
+        agent = AssistantAgent(
+            name="test_agent",
+            model_client=model_client,
+            tools=[hanging_tool],
+        )
+
+        cancellation_token = CancellationToken()
+        stream = agent.on_messages_stream(
+            [TextMessage(content="Test", source="user")],
+            cancellation_token,
+        )
+
+        async def _consume() -> None:
+            async for _ in stream:
+                pass
+
+        consume_task = asyncio.create_task(_consume())
+
+        # Let tool execution start (model returns function_calls → workbench starts tool).
+        await asyncio.sleep(0.2)
+
+        # Cancel mid-tool-execution.
+        cancellation_token.cancel()
+
+        # The stream MUST terminate within the timeout.
+        # Prior to the fix this would hang forever; after the fix the
+        # CancelledError propagates out (and is expected here).
+        try:
+            await asyncio.wait_for(consume_task, timeout=5.0)
+        except asyncio.CancelledError:
+            pass  # expected: cancellation propagates out of the consumer
+
 
 class TestAssistantAgentStreamingEdgeCases:
     """Test suite for streaming edge cases and error scenarios."""