From 5dcc6ae1c466b2cdf94465389602d2d1b34abaef Mon Sep 17 00:00:00 2001 From: Dmitry Voropaev Date: Thu, 29 Oct 2026 12:09:00 +0300 Subject: [PATCH] fix(server): propagate caller cancellation out of ActiveTask.subscribe ActiveTask.subscribe caught asyncio.CancelledError next to QueueShutDown and broke out of its loop. When the task iterating the subscription was cancelled, the generator ended normally instead of re-raising, so the cancellation was lost. For DefaultRequestHandlerV2.on_message_send this meant asyncio.wait_for never raised TimeoutError and Task.cancel() was ignored. The call returned the task in whatever non-terminal state it had at that moment as if it were the result. A cancelled consumer of on_message_send_stream or on_subscribe_to_task saw its stream end quietly instead of being cancelled. Only QueueShutDown is caught now. The producer and consumer only enqueue Exception instances and queue shutdown raises QueueShutDown, so a CancelledError at this point always comes from the caller. The finally block still detaches the tap and drops the reference count, so cleanup is unchanged. --- src/a2a/server/agent_execution/active_task.py | 2 +- .../agent_execution/test_active_task.py | 6 +- .../test_default_request_handler_v2.py | 91 +++++++++++++++++++ 3 files changed, 95 insertions(+), 4 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index c5dcbb9dd..0f07a2721 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -840,7 +840,7 @@ async def subscribe( # noqa: PLR0915 yield event finally: tapped_queue.task_done() - except (QueueShutDown, asyncio.CancelledError): + except QueueShutDown: break finally: logger.debug('Subscribe[%s]: Unsubscribing', self._task_id) diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 254423e18..d51240ea8 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -534,9 +534,9 @@ async def slow_execute(req, q): task.cancel() - # In python 3.10+ cancelling an async generator next() might raise StopAsyncIteration - # if the generator handles the cancellation by closing. - with pytest.raises((asyncio.CancelledError, StopAsyncIteration)): + # The cancellation belongs to the caller, so subscribe() must re-raise + # it rather than end the stream as if the task had finished. + with pytest.raises(asyncio.CancelledError): await task await it.aclose() diff --git a/tests/server/request_handlers/test_default_request_handler_v2.py b/tests/server/request_handlers/test_default_request_handler_v2.py index 66cf9f4ea..9396938a1 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -2614,6 +2614,97 @@ async def test_on_cancel_of_parked_task_is_owner_scoped(): await handler.aclose() +def _hanging_send_request(text: str) -> SendMessageRequest: + return SendMessageRequest( + message=Message( + message_id=f'msg-{text}', + role=Role.ROLE_USER, + parts=[Part(text=text)], + ) + ) + + +@pytest.mark.timeout(10) +@pytest.mark.asyncio +async def test_on_message_send_wait_for_times_out_while_task_works(): + """A caller timeout on a blocking send raises instead of returning the + unfinished task.""" + agent = _HangingAgent() + handler = DefaultRequestHandlerV2( + agent_executor=agent, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for( + handler.on_message_send( + _hanging_send_request('timeout'), create_server_call_context() + ), + timeout=0.5, + ) + + agent.release.set() + await handler.aclose() + + +@pytest.mark.timeout(10) +@pytest.mark.asyncio +async def test_on_message_send_propagates_caller_cancellation(): + """Cancelling the task that awaits a blocking send cancels it.""" + agent = _HangingAgent() + handler = DefaultRequestHandlerV2( + agent_executor=agent, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + + send = asyncio.create_task( + handler.on_message_send( + _hanging_send_request('cancel'), create_server_call_context() + ) + ) + await asyncio.wait_for(agent.working.wait(), timeout=5) + send.cancel() + + with pytest.raises(asyncio.CancelledError): + await send + assert send.cancelled() + + agent.release.set() + await handler.aclose() + + +@pytest.mark.timeout(10) +@pytest.mark.asyncio +async def test_on_message_send_stream_propagates_consumer_cancellation(): + """Cancelling a stream consumer cancels it instead of ending the stream + as if the task had finished.""" + agent = _HangingAgent() + handler = DefaultRequestHandlerV2( + agent_executor=agent, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + + async def consume() -> None: + async for _ in handler.on_message_send_stream( + _hanging_send_request('stream'), create_server_call_context() + ): + pass + + reader = asyncio.create_task(consume()) + await asyncio.wait_for(agent.working.wait(), timeout=5) + reader.cancel() + + with pytest.raises(asyncio.CancelledError): + await reader + assert reader.cancelled() + + agent.release.set() + await handler.aclose() + + class _AsksForInputAgent(AgentExecutor): """Always asks for more input; counts execute() and cancel() calls."""