fix:task execution cancelled by client disconnect

Author: ovidiutaralesca

src/a2a/server/request_handlers/default_request_handler.py

@@ -385,7 +385,7 @@ async def on_message_send_stream(
                 )
                 yield event
         finally:
-            await self._cleanup_producer(producer_task, task_id)
+            asyncio.create_task(self._cleanup_producer(producer_task, task_id))
 
     async def _register_producer(
         self, task_id: str, producer_task: asyncio.Task

tests/server/request_handlers/test_default_request_handler.py

@@ -904,6 +904,14 @@ async def test_on_message_send_stream_with_push_notification():
         configuration=message_config,
     )
 
+    # Latch to ensure background execute is scheduled before asserting
+    execute_called = asyncio.Event()
+
+    async def exec_side_effect(*args, **kwargs):
+        execute_called.set()
+
+    mock_agent_executor.execute.side_effect = exec_side_effect
+
     # Mock ResultAggregator and its consume_and_emit
     mock_result_aggregator_instance = MagicMock(
         spec=ResultAggregator
@@ -1117,6 +1125,8 @@ def sync_get_event_stream_gen_for_prop_test(*args, **kwargs):
         ):
             pass
 
+    await asyncio.wait_for(execute_called.wait(), timeout=0.1)
+
     # Assertions
     # 1. set_info called once at the beginning if task exists (or after task is created from message)
     mock_push_config_store.set_info.assert_any_call(task_id, push_config)
@@ -1129,6 +1139,202 @@ def sync_get_event_stream_gen_for_prop_test(*args, **kwargs):
     mock_agent_executor.execute.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_stream_disconnect_then_resubscribe_receives_future_events():
+    """Start streaming, disconnect, then resubscribe and ensure subsequent events are streamed."""
+    # Arrange
+    mock_task_store = AsyncMock(spec=TaskStore)
+    mock_agent_executor = AsyncMock(spec=AgentExecutor)
+
+    # Use a real queue manager so taps receive future events
+    queue_manager = InMemoryQueueManager()
+
+    task_id = 'reconn_task_1'
+    context_id = 'reconn_ctx_1'
+
+    # Task exists and is non-final
+    task_for_resub = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.working
+    )
+    mock_task_store.get.return_value = task_for_resub
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=mock_agent_executor,
+        task_store=mock_task_store,
+        queue_manager=queue_manager,
+    )
+
+    params = MessageSendParams(
+        message=Message(
+            role=Role.user,
+            message_id='msg_reconn',
+            parts=[],
+            task_id=task_id,
+            context_id=context_id,
+        )
+    )
+
+    # Producer behavior: emit one event, then later emit second event
+    exec_started = asyncio.Event()
+    allow_second_event = asyncio.Event()
+    allow_finish = asyncio.Event()
+
+    first_event = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.working
+    )
+    second_event = create_sample_task(
+        task_id=task_id, context_id=context_id, status_state=TaskState.completed
+    )
+
+    async def exec_side_effect(_request, queue: EventQueue):
+        exec_started.set()
+        await queue.enqueue_event(first_event)
+        await allow_second_event.wait()
+        await queue.enqueue_event(second_event)
+        await allow_finish.wait()
+
+    mock_agent_executor.execute.side_effect = exec_side_effect
+
+    # Start streaming and consume first event
+    agen = request_handler.on_message_send_stream(
+        params, create_server_call_context()
+    )
+    first = await agen.__anext__()
+    assert first == first_event
+
+    # Simulate client disconnect
+    await asyncio.wait_for(agen.aclose(), timeout=0.1)
+
+    # Resubscribe and start consuming future events
+    resub_gen = request_handler.on_resubscribe_to_task(
+        TaskIdParams(id=task_id), create_server_call_context()
+    )
+
+    # Allow producer to emit the next event
+    allow_second_event.set()
+
+    received = await resub_gen.__anext__()
+    assert received == second_event
+
+    # Finish producer to allow cleanup paths to complete
+    allow_finish.set()
+
+
+@pytest.mark.asyncio
+async def test_on_message_send_stream_client_disconnect_triggers_background_cleanup_and_producer_continues():
+    """Simulate client disconnect: stream stops early, cleanup is scheduled in background,
+    producer keeps running, and cleanup completes after producer finishes."""
+    # Arrange
+    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 = 'disc_task_1'
+    context_id = 'disc_ctx_1'
+
+    # RequestContext with IDs
+    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
+
+    # Queue used by _run_event_stream; must support close()
+    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='mid',
+            parts=[],
+            task_id=task_id,
+            context_id=context_id,
+        )
+    )
+
+    # Agent executor runs in background until we allow it to finish
+    execute_started = asyncio.Event()
+    execute_finish = asyncio.Event()
+
+    async def exec_side_effect(*_args, **_kwargs):
+        execute_started.set()
+        await execute_finish.wait()
+
+    mock_agent_executor.execute.side_effect = exec_side_effect
+
+    # ResultAggregator emits one Task event (so the stream yields once)
+    first_event = create_sample_task(task_id=task_id, context_id=context_id)
+
+    async def single_event_stream():
+        yield first_event
+        # will never yield again; client will disconnect
+
+    mock_result_aggregator_instance = MagicMock(spec=ResultAggregator)
+    mock_result_aggregator_instance.consume_and_emit.return_value = (
+        single_event_stream()
+    )
+
+    produced_task: asyncio.Task | None = None
+    cleanup_task: asyncio.Task | None = None
+
+    orig_create_task = asyncio.create_task
+
+    def create_task_spy(coro):
+        nonlocal produced_task, cleanup_task
+        task = orig_create_task(coro)
+        if produced_task is None:
+            produced_task = task
+        else:
+            cleanup_task = task
+        return task
+
+    with (
+        patch(
+            'a2a.server.request_handlers.default_request_handler.ResultAggregator',
+            return_value=mock_result_aggregator_instance,
+        ),
+        patch('asyncio.create_task', side_effect=create_task_spy),
+    ):
+        # Act: start stream and consume only the first event, then disconnect
+        agen = request_handler.on_message_send_stream(
+            params, create_server_call_context()
+        )
+        first = await agen.__anext__()
+        assert first == first_event
+        # Simulate client disconnect
+        await asyncio.wait_for(agen.aclose(), timeout=0.1)
+
+    # Assert cleanup was scheduled and producer was started
+    assert produced_task is not None
+    assert cleanup_task is not None
+
+    # execute should have started
+    await asyncio.wait_for(execute_started.wait(), timeout=0.1)
+
+    # Producer should still be running (not finished immediately on disconnect)
+    assert not produced_task.done()
+
+    # Allow executor to finish, which should complete producer and then cleanup
+    execute_finish.set()
+    await asyncio.wait_for(produced_task, timeout=0.2)
+    await asyncio.wait_for(cleanup_task, timeout=0.2)
+
+    # Queue close awaited by _run_event_stream
+    mock_queue.close.assert_awaited_once()
+    # QueueManager close called by _cleanup_producer
+    mock_queue_manager.close.assert_awaited_once_with(task_id)
+    # Running agents is cleared
+    assert task_id not in request_handler._running_agents
+
+
 @pytest.mark.asyncio
 async def test_on_message_send_stream_task_id_mismatch():
     """Test on_message_send_stream raises error if yielded task ID mismatches."""

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -1,3 +1,4 @@
+import asyncio
 import unittest
 import unittest.async_case
 
@@ -361,6 +362,14 @@ async def streaming_coro():
             for event in events:
                 yield event
 
+        # Latch to ensure background execute is scheduled before asserting
+        execute_called = asyncio.Event()
+
+        async def exec_side_effect(*args, **kwargs):
+            execute_called.set()
+
+        mock_agent_executor.execute.side_effect = exec_side_effect
+
         with patch(
             'a2a.server.request_handlers.default_request_handler.EventConsumer.consume_all',
             return_value=streaming_coro(),
@@ -382,6 +391,7 @@ async def streaming_coro():
                     event.root, SendStreamingMessageSuccessResponse
                 )
                 assert event.root.result == events[i]
+            await asyncio.wait_for(execute_called.wait(), timeout=0.1)
             mock_agent_executor.execute.assert_called_once()
 
     async def test_on_message_stream_new_message_existing_task_success(
@@ -418,6 +428,14 @@ async def streaming_coro():
             for event in events:
                 yield event
 
+        # Latch to ensure background execute is scheduled before asserting
+        execute_called = asyncio.Event()
+
+        async def exec_side_effect(*args, **kwargs):
+            execute_called.set()
+
+        mock_agent_executor.execute.side_effect = exec_side_effect
+
         with patch(
             'a2a.server.request_handlers.default_request_handler.EventConsumer.consume_all',
             return_value=streaming_coro(),
@@ -438,6 +456,7 @@ async def streaming_coro():
             assert isinstance(response, AsyncGenerator)
             collected_events = [item async for item in response]
             assert len(collected_events) == len(events)
+            await asyncio.wait_for(execute_called.wait(), timeout=0.1)
             mock_agent_executor.execute.assert_called_once()
             assert mock_task.history is not None and len(mock_task.history) == 1