Simplify fix for consumer error handling

Author: Łukasz Bobiński

src/a2a/server/request_handlers/default_request_handler.py

@@ -329,13 +329,7 @@ async def push_notification_callback() -> None:
             )
 
         except Exception:
-            logger.exception('Agent execution failed')
-            # If the consumer fails, we must cancel the producer to prevent it from hanging
-            # on queue operations (e.g., waiting for the queue to drain).
-            producer_task.cancel()
-            # Force the queue to close immediately, discarding any pending events.
-            # This ensures that any producers waiting on the queue are unblocked.
-            await queue.close(immediate=True)
+            await self._handle_execution_failure(producer_task, queue)
             raise
         finally:
             if interrupted_or_non_blocking:
@@ -400,9 +394,7 @@ async def on_message_send_stream(
             raise
         except Exception:
             # If the consumer fails (e.g. database error), we must cleanup.
-            logger.exception('Agent execution failed during streaming')
-            producer_task.cancel()
-            await queue.close(immediate=True)
+            await self._handle_execution_failure(producer_task, queue)
             raise
         finally:
             cleanup_task = asyncio.create_task(
@@ -441,6 +433,18 @@ def _on_done(completed: asyncio.Task) -> None:
 
         task.add_done_callback(_on_done)
 
+    async def _handle_execution_failure(
+        self, producer_task: asyncio.Task, queue: EventQueue
+    ) -> None:
+        """Cancels the producer and closes the queue immediately on failure."""
+        logger.exception('Agent execution failed')
+        # If the consumer fails, we must cancel the producer to prevent it from hanging
+        # on queue operations (e.g., waiting for the queue to drain).
+        producer_task.cancel()
+        # Force the queue to close immediately, discarding any pending events.
+        # This ensures that any producers waiting on the queue are unblocked.
+        await queue.close(immediate=True)
+
     async def _cleanup_producer(
         self,
         producer_task: asyncio.Task,

tests/server/request_handlers/test_default_request_handler.py

@@ -2732,3 +2732,80 @@ async def spy_register_producer(tid, task):
 
     # Verify queue closed immediately
     mock_queue.close.assert_awaited_with(immediate=True)
+
+@pytest.mark.asyncio
+async def test_on_message_send_consumer_error_cancels_producer_and_closes_queue():
+    """Test that if the consumer raises an exception during blocking wait, the producer is cancelled."""
+    mock_task_store = AsyncMock(spec=TaskStore)
+    mock_queue_manager = AsyncMock(spec=QueueManager)
+    mock_agent_executor = AsyncMock(spec=AgentExecutor)
+    mock_request_context_builder = AsyncMock(spec=RequestContextBuilder)
+
+    task_id = 'error_cleanup_blocking_task'
+    context_id = 'error_cleanup_blocking_ctx'
+
+    mock_request_context = MagicMock(spec=RequestContext)
+    mock_request_context.task_id = task_id
+    mock_request_context.context_id = context_id
+    mock_request_context_builder.build.return_value = mock_request_context
+
+    mock_queue = AsyncMock(spec=EventQueue)
+    mock_queue_manager.create_or_tap.return_value = mock_queue
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=mock_agent_executor,
+        task_store=mock_task_store,
+        queue_manager=mock_queue_manager,
+        request_context_builder=mock_request_context_builder,
+    )
+
+    params = MessageSendParams(
+        message=Message(
+            role=Role.user,
+            message_id='msg_error_blocking',
+            parts=[],
+        )
+    )
+
+    # Mock ResultAggregator to raise exception
+    mock_result_aggregator_instance = MagicMock(spec=ResultAggregator)
+    mock_result_aggregator_instance.consume_and_break_on_interrupt.side_effect = ValueError('Consumer failed!')
+
+    # Capture the producer task to verify cancellation
+    captured_producer_task = None
+    original_register = request_handler._register_producer
+
+    async def spy_register_producer(tid, task):
+        nonlocal captured_producer_task
+        captured_producer_task = task
+        # Wrap the cancel method to spy on it
+        task.cancel = MagicMock(wraps=task.cancel)
+        await original_register(tid, task)
+
+    with (
+        patch(
+            'a2a.server.request_handlers.default_request_handler.ResultAggregator',
+            return_value=mock_result_aggregator_instance,
+        ),
+        patch(
+            'a2a.server.request_handlers.default_request_handler.TaskManager.get_task',
+            return_value=None,
+        ),
+        patch.object(
+            request_handler,
+            '_register_producer',
+            side_effect=spy_register_producer,
+        ),
+    ):
+        # Act
+        with pytest.raises(ValueError, match='Consumer failed!'):
+            await request_handler.on_message_send(
+                params, create_server_call_context()
+            )
+
+    assert captured_producer_task is not None
+    # Verify producer was cancelled
+    captured_producer_task.cancel.assert_called()
+
+    # Verify queue closed immediately
+    mock_queue.close.assert_awaited_with(immediate=True)