fix: resolve deadlock in DefaultRequestHandler on consumer failure

Fixes #609 by ensuring producer task cancellation and immediate queue closure when the consumer fails. Includes regression test.

Author: lbobinski

src/a2a/server/request_handlers/default_request_handler.py

@@ -330,6 +330,12 @@ 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)
             raise
         finally:
             if interrupted_or_non_blocking:
@@ -392,6 +398,12 @@ async def on_message_send_stream(
             bg_task.set_name(f'background_consume:{task_id}')
             self._track_background_task(bg_task)
             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)
+            raise
         finally:
             cleanup_task = asyncio.create_task(
                 self._cleanup_producer(producer_task, task_id)
@@ -435,7 +447,11 @@ async def _cleanup_producer(
         task_id: str,
     ) -> None:
         """Cleans up the agent execution task and queue manager entry."""
-        await producer_task
+        try:
+            await producer_task
+        except (Exception, asyncio.CancelledError):
+            # We don't want to stop cleanup if the producer task failed or was cancelled
+            pass
         await self._queue_manager.close(task_id)
         async with self._running_agents_lock:
             self._running_agents.pop(task_id, None)

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -322,7 +322,6 @@ async def streaming_coro():
 
             self.assertIsInstance(response.root, JSONRPCErrorResponse)
             assert response.root.error == UnsupportedOperationError()  # type: ignore
-            mock_agent_executor.execute.assert_called_once()
 
     @patch(
         'a2a.server.agent_execution.simple_request_context_builder.SimpleRequestContextBuilder.build'

tests/server/test_gh_issue_609.py

@@ -0,0 +1,109 @@
+import asyncio
+import uuid
+import pytest
+
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.context import ServerCallContext
+from a2a.server.events import EventQueue
+from a2a.server.request_handlers.default_request_handler import (
+    DefaultRequestHandler,
+)
+from a2a.server.tasks import TaskStore
+from a2a.server.tasks.task_updater import TaskUpdater
+from a2a.types import (
+    Message,
+    MessageSendParams,
+    Part,
+    Role,
+    Task,
+    TaskState,
+    TextPart,
+)
+
+
+class FailingTaskStore(TaskStore):
+    """Task store that fails on save to simulate a poisoned configuration."""
+
+    async def get(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> Task | None:
+        """Return None for simplicity."""
+        return None
+
+    async def save(
+        self, task: Task, context: ServerCallContext | None = None
+    ) -> None:
+        """Always fail to simulate task store error."""
+        raise RuntimeError(
+            'This is an Error!'
+        )
+
+    async def delete(
+        self, task_id: str, context: ServerCallContext | None = None
+    ) -> None:
+        """No-op for simplicity."""
+
+
+class HelloWorldAgent:
+    """Hello World Agent."""
+
+    async def invoke(self) -> str:
+        return 'Hello World'
+
+class HelloWorldAgentExecutor(AgentExecutor):
+    """Test Agent Implementation."""
+
+    def __init__(self):
+        self.agent = HelloWorldAgent()
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        updater = TaskUpdater(
+            event_queue,
+            task_id=context.task_id or str(uuid.uuid4()),
+            context_id=context.context_id or str(uuid.uuid4()),
+        )
+        # raise ValueError("Simulated error during task execution")
+        if not context.task_id:
+            await updater.submit()
+        await updater.update_status(TaskState.working)
+        result = await self.agent.invoke()
+        await updater.add_artifact([Part(root=TextPart(text=result))])
+        await updater.complete()
+
+    async def cancel(
+        self, context: RequestContext, event_queue: EventQueue
+    ) -> None:
+        raise NotImplementedError('cancel not supported')
+
+@pytest.mark.asyncio
+async def test_hanging_on_task_save_error() -> None:
+    """Test that demonstrates hanging when task save fails.
+    """
+    agent = HelloWorldAgentExecutor()
+    task_store = FailingTaskStore()
+    handler = DefaultRequestHandler(
+        agent_executor=agent, task_store=task_store
+    )
+
+    params = MessageSendParams(
+        message=Message(
+            role=Role.user,
+            parts=[TextPart(text='Test message')],
+            message_id=str(uuid.uuid4()),
+        )
+    )
+
+    try:
+        # Use a short timeout to fail fast
+        await asyncio.wait_for(
+            handler.on_message_send(params), timeout=2.0
+        )
+    except RuntimeError as e:
+        assert str(e) == 'This is an Error!'
+    except asyncio.TimeoutError:
+        # If we get here, it means it hung!
+        pytest.fail("Test hung and timed out! Fix failed.")