From b3b1e895f077b49d9fd9681b8f0ab3e9de65dbd6 Mon Sep 17 00:00:00 2001 From: Mykyta Netipa Date: Tue, 3 Nov 2026 12:30:06 +0000 Subject: [PATCH] fix(server): forward CancelTaskRequest.metadata to AgentExecutor.cancel Parse the optional JSON request body in RestDispatcher.on_cancel_task while preserving the URL path task ID and allowing empty bodies, accept CancelTaskRequest in RequestContext so context.metadata exposes cancel metadata without affecting message-only fields, thread CancelTaskRequest through DefaultRequestHandlerV2/ActiveTask.cancel and LegacyRequestHandler, and register the optional CancelTaskRequest body in FastAPI OpenAPI routes. Fixes #1415 --- src/a2a/server/agent_execution/active_task.py | 8 ++- src/a2a/server/agent_execution/context.py | 16 +++-- .../default_request_handler.py | 2 +- .../default_request_handler_v2.py | 8 +-- src/a2a/server/routes/_proto_schema.py | 7 +- src/a2a/server/routes/fastapi_routes.py | 10 ++- src/a2a/server/routes/rest_dispatcher.py | 10 +-- tests/integration/test_end_to_end.py | 66 ++++++++++++++++++- .../agent_execution/test_active_task.py | 58 ++++++++++++++++ tests/server/agent_execution/test_context.py | 47 +++++++++++++ .../test_default_request_handler.py | 7 +- .../test_default_request_handler_v2.py | 35 ++++++++++ tests/server/routes/test_fastapi_routes.py | 9 ++- tests/server/routes/test_proto_schema.py | 3 +- tests/server/routes/test_rest_dispatcher.py | 56 ++++++++++++++++ 15 files changed, 318 insertions(+), 24 deletions(-) diff --git a/src/a2a/server/agent_execution/active_task.py b/src/a2a/server/agent_execution/active_task.py index 121900fa2..71929fe1d 100644 --- a/src/a2a/server/agent_execution/active_task.py +++ b/src/a2a/server/agent_execution/active_task.py @@ -71,6 +71,7 @@ ) from a2a.server.tasks import PushNotificationEvent from a2a.types.a2a_pb2 import ( + CancelTaskRequest, Message, Task, TaskArtifactUpdateEvent, @@ -877,7 +878,11 @@ async def subscribe( # noqa: PLR0915 await self._maybe_cleanup() self.schedule_idle_release() - async def cancel(self, call_context: ServerCallContext) -> Task: + async def cancel( + self, + call_context: ServerCallContext, + request: CancelTaskRequest | None = None, + ) -> Task: """Cancels the running active task. The returned task carries a terminal state, which is written to the @@ -899,6 +904,7 @@ async def cancel(self, call_context: ServerCallContext) -> Task: task = await self._task_manager.get_task() request_context = RequestContext( call_context=call_context, + request=request, task_id=self._task_id, context_id=task.context_id if task else None, task=task, diff --git a/src/a2a/server/agent_execution/context.py b/src/a2a/server/agent_execution/context.py index 7c47bc7d8..5d1ca5f16 100644 --- a/src/a2a/server/agent_execution/context.py +++ b/src/a2a/server/agent_execution/context.py @@ -10,6 +10,7 @@ UUIDGenerator, ) from a2a.types.a2a_pb2 import ( + CancelTaskRequest, Message, SendMessageConfiguration, SendMessageRequest, @@ -29,19 +30,19 @@ class RequestContext: def __init__( # noqa: PLR0913 self, call_context: ServerCallContext, - request: SendMessageRequest | None = None, + request: SendMessageRequest | CancelTaskRequest | None = None, task_id: str | None = None, context_id: str | None = None, task: Task | None = None, related_tasks: list[Task] | None = None, task_id_generator: IDGenerator | None = None, context_id_generator: IDGenerator | None = None, - ): + ) -> None: """Initializes the RequestContext. Args: call_context: The server call context associated with this request. - request: The incoming `SendMessageRequest` request payload. + request: The incoming `SendMessageRequest` or `CancelTaskRequest` payload. task_id: The ID of the task explicitly provided in the request or path. context_id: The ID of the context explicitly provided in the request or path. task: The existing `Task` object retrieved from the store, if any. @@ -52,7 +53,10 @@ def __init__( # noqa: PLR0913 if related_tasks is None: related_tasks = [] self._call_context = call_context - self._params = request + self._request: SendMessageRequest | CancelTaskRequest | None = request + self._params: SendMessageRequest | None = ( + None if isinstance(request, CancelTaskRequest) else request + ) self._task_id = task_id self._context_id = context_id self._current_task = task @@ -149,10 +153,10 @@ def call_context(self) -> ServerCallContext: @property def metadata(self) -> dict[str, Any]: """Metadata associated with the request, if available.""" - if self._params: + if self._request: # MessageToDict recurses into nested Struct/ListValue fields; # dict would leak raw protobuf objects to callers. - return json_format.MessageToDict(self._params.metadata) + return json_format.MessageToDict(self._request.metadata) return {} @property diff --git a/src/a2a/server/request_handlers/default_request_handler.py b/src/a2a/server/request_handlers/default_request_handler.py index 9bbb5310d..709d738f4 100644 --- a/src/a2a/server/request_handlers/default_request_handler.py +++ b/src/a2a/server/request_handlers/default_request_handler.py @@ -244,7 +244,7 @@ async def on_cancel_task( await self.agent_executor.cancel( RequestContext( call_context=context, - request=None, + request=params, task_id=task.id, context_id=task.context_id, task=task, diff --git a/src/a2a/server/request_handlers/default_request_handler_v2.py b/src/a2a/server/request_handlers/default_request_handler_v2.py index 5655fcfdc..4a42a2781 100644 --- a/src/a2a/server/request_handlers/default_request_handler_v2.py +++ b/src/a2a/server/request_handlers/default_request_handler_v2.py @@ -295,21 +295,21 @@ async def on_cancel_task( # noqa: D102 # This replica holds the task, running or waiting for input. if await self._active_task_registry.get(task_id) is not None: - return await self._cancel_local(task_id, context) + return await self._cancel_local(params, context) # Running on another replica (or nowhere): record CANCELED in shared # state; the owner sees the version bump on its next save and aborts. return await self._cancel_remote(task_id, task, version, context) async def _cancel_local( - self, task_id: str, context: ServerCallContext + self, params: CancelTaskRequest, context: ServerCallContext ) -> Task: try: active_task = await self._active_task_registry.get_or_create( - task_id, call_context=context, create_task_if_missing=False + params.id, call_context=context, create_task_if_missing=False ) try: - result = await active_task.cancel(context) + result = await active_task.cancel(context, request=params) finally: # get_or_create stopped the timer; a failed cancel leaves the task parked. active_task.schedule_idle_release() diff --git a/src/a2a/server/routes/_proto_schema.py b/src/a2a/server/routes/_proto_schema.py index bed284aa6..1f066a151 100644 --- a/src/a2a/server/routes/_proto_schema.py +++ b/src/a2a/server/routes/_proto_schema.py @@ -6,13 +6,18 @@ from google.protobuf.descriptor import Descriptor, FieldDescriptor from google.protobuf.message import Message -from a2a.types.a2a_pb2 import SendMessageRequest, TaskPushNotificationConfig +from a2a.types.a2a_pb2 import ( + CancelTaskRequest, + SendMessageRequest, + TaskPushNotificationConfig, +) from a2a.utils.proto_utils import _field_is_repeated REST_BODY_TYPES: dict[tuple[str, str], type[Message]] = { ('/message:send', 'POST'): SendMessageRequest, ('/message:stream', 'POST'): SendMessageRequest, + ('/tasks/{id}:cancel', 'POST'): CancelTaskRequest, ('/tasks/{id}/pushNotificationConfigs', 'POST'): TaskPushNotificationConfig, } diff --git a/src/a2a/server/routes/fastapi_routes.py b/src/a2a/server/routes/fastapi_routes.py index 7efa90a49..77c0c0711 100644 --- a/src/a2a/server/routes/fastapi_routes.py +++ b/src/a2a/server/routes/fastapi_routes.py @@ -54,13 +54,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: } +_OPTIONAL_REST_BODIES: frozenset[tuple[str, str]] = frozenset( + {('/tasks/{id}:cancel', 'POST')} +) + + def _request_body_extra( - ref: dict[str, Any], description: str + ref: dict[str, Any], description: str, *, required: bool = True ) -> dict[str, Any]: return { 'requestBody': { 'description': description, - 'required': True, + 'required': required, 'content': {'application/json': {'schema': ref}}, }, } @@ -151,6 +156,7 @@ def add_a2a_routes_to_fastapi( key: _request_body_extra( message_schema(cls.DESCRIPTOR, components), f'A2A {cls.__name__}', + required=key not in _OPTIONAL_REST_BODIES, ) for key, cls in REST_BODY_TYPES.items() } diff --git a/src/a2a/server/routes/rest_dispatcher.py b/src/a2a/server/routes/rest_dispatcher.py index df1ab991e..80964e5bf 100644 --- a/src/a2a/server/routes/rest_dispatcher.py +++ b/src/a2a/server/routes/rest_dispatcher.py @@ -218,10 +218,12 @@ async def on_cancel_task(self, request: Request) -> Response: @validate_version(constants.PROTOCOL_VERSION_1_0) async def _handler(context: ServerCallContext) -> a2a_pb2.Task: - task_id = request.path_params['id'] - task = await self.request_handler.on_cancel_task( - CancelTaskRequest(id=task_id), context - ) + params = CancelTaskRequest() + body = await request.body() + if body.strip(): + proto_utils.parse_lenient(body, params) + params.id = request.path_params['id'] + task = await self.request_handler.on_cancel_task(params, context) if task: return task raise TaskNotFoundError diff --git a/tests/integration/test_end_to_end.py b/tests/integration/test_end_to_end.py index eeaf46416..d4276183d 100644 --- a/tests/integration/test_end_to_end.py +++ b/tests/integration/test_end_to_end.py @@ -1,3 +1,5 @@ +import json + from collections.abc import AsyncGenerator from typing import NamedTuple @@ -156,8 +158,19 @@ async def execute(self, context: RequestContext, event_queue: EventQueue): message=task_updater.new_agent_message([Part(text='done')]), ) - async def cancel(self, context: RequestContext, event_queue: EventQueue): - raise NotImplementedError('Cancellation is not supported') + async def cancel( + self, context: RequestContext, event_queue: EventQueue + ) -> None: + assert context.task_id is not None + assert context.context_id is not None + task_updater = TaskUpdater( + event_queue, context.task_id, context.context_id + ) + await task_updater.cancel( + message=task_updater.new_agent_message( + [Part(text=json.dumps(context.metadata, sort_keys=True))] + ) + ) @pytest.fixture @@ -831,3 +844,52 @@ async def test_end_to_end_extensions_propagation(transport_setups, streaming): response.message, Role.ROLE_AGENT, 'extensions echoed' ) assert set(response.message.extensions) == set(SUPPORTED_EXTENSION_URIS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + 'cancel_metadata', + [ + pytest.param( + {'reason': 'user_aborted', 'nested': {'retryable': False}}, + id='with_metadata', + ), + pytest.param({}, id='without_metadata'), + ], +) +async def test_end_to_end_cancel_task_metadata_propagation( + transport_setups: ClientSetup, + cancel_metadata: dict[str, object], +) -> None: + """Test CancelTaskRequest.metadata reaches AgentExecutor.cancel on all bindings.""" + client = transport_setups.client + client._config.streaming = False + + events = [ + event + async for event in client.send_message( + request=SendMessageRequest( + message=Message( + role=Role.ROLE_USER, + message_id='msg-cancel-metadata', + parts=[Part(text='Need input')], + ) + ) + ) + ] + assert len(events) == 1 + parked_task = events[0].task + assert parked_task.status.state == TaskState.TASK_STATE_INPUT_REQUIRED + + cancel_req = ( + CancelTaskRequest(id=parked_task.id, metadata=cancel_metadata) + if cancel_metadata + else CancelTaskRequest(id=parked_task.id) + ) + cancelled_task = await client.cancel_task(request=cancel_req) + + assert cancelled_task.id == parked_task.id + assert cancelled_task.status.state == TaskState.TASK_STATE_CANCELED + assert json.loads(cancelled_task.status.message.parts[0].text) == ( + cancel_metadata + ) diff --git a/tests/server/agent_execution/test_active_task.py b/tests/server/agent_execution/test_active_task.py index 9f82d1685..20af08749 100644 --- a/tests/server/agent_execution/test_active_task.py +++ b/tests/server/agent_execution/test_active_task.py @@ -15,6 +15,7 @@ from a2a.server.tasks.push_notification_sender import PushNotificationSender from a2a.server.tasks.task_manager import TaskManager from a2a.types.a2a_pb2 import ( + CancelTaskRequest, Message, Task, TaskState, @@ -131,6 +132,63 @@ async def execute_mock(req, q): agent_executor.cancel.assert_called_once() stop_event.set() + @pytest.mark.asyncio + async def test_active_task_cancel_forwards_request_metadata( + self, + active_task: ActiveTask, + agent_executor: Mock, + request_context: Mock, + task_manager: Mock, + ) -> None: + """Test ActiveTask.cancel passes CancelTaskRequest metadata to AgentExecutor.cancel.""" + started = asyncio.Event() + stop_event = asyncio.Event() + + async def execute_mock( + _ctx: RequestContext, _queue: EventQueue + ) -> None: + started.set() + await stop_event.wait() + + agent_executor.execute = AsyncMock(side_effect=execute_mock) + agent_executor.cancel = AsyncMock() + task_manager.get_task.side_effect = [ + Task( + id='test-task-id', + context_id='ctx-1', + status=TaskStatus(state=TaskState.TASK_STATE_WORKING), + ) + ] * 5 + [ + Task( + id='test-task-id', + context_id='ctx-1', + status=TaskStatus(state=TaskState.TASK_STATE_CANCELED), + ) + ] * 5 + + await active_task.enqueue_request(request_context) + await active_task.start( + call_context=ServerCallContext(), create_task_if_missing=True + ) + await asyncio.wait_for(started.wait(), timeout=5) + + cancel_req = CancelTaskRequest( + id='test-task-id', + metadata={'reason': 'user_aborted', 'actor': 'alice'}, + ) + await active_task.cancel(ServerCallContext(), request=cancel_req) + + agent_executor.cancel.assert_awaited_once() + cancel_ctx: RequestContext = agent_executor.cancel.call_args.args[0] + assert cancel_ctx.metadata == { + 'reason': 'user_aborted', + 'actor': 'alice', + } + assert cancel_ctx.message is None + assert cancel_ctx.task_id == 'test-task-id' + assert cancel_ctx.context_id == 'ctx-1' + stop_event.set() + @pytest.mark.asyncio async def test_active_task_cancel_producer_cancelled_on_cancellederror( self, diff --git a/tests/server/agent_execution/test_context.py b/tests/server/agent_execution/test_context.py index 42f4178a1..77c92a202 100644 --- a/tests/server/agent_execution/test_context.py +++ b/tests/server/agent_execution/test_context.py @@ -9,6 +9,7 @@ from a2a.server.context import ServerCallContext from a2a.server.id_generator import IDGenerator from a2a.types.a2a_pb2 import ( + CancelTaskRequest, Message, SendMessageRequest, Task, @@ -346,3 +347,49 @@ def test_extension_handling(self) -> None: context = RequestContext(call_context=call_context) assert context.requested_extensions == {'foo', 'bar'} + + def test_init_with_cancel_task_request_metadata( + self, mock_task: Mock + ) -> None: + """Test RequestContext exposes metadata from CancelTaskRequest safely.""" + cancel_request = CancelTaskRequest( + id=mock_task.id, + metadata={ + 'reason': 'user_aborted', + 'nested': {'attempt': 2, 'flags': ['urgent', 'manual']}, + }, + ) + context = RequestContext( + call_context=ServerCallContext(), + request=cancel_request, + task_id=mock_task.id, + context_id=mock_task.context_id, + task=mock_task, + ) + + assert context.metadata == { + 'reason': 'user_aborted', + 'nested': {'attempt': 2.0, 'flags': ['urgent', 'manual']}, + } + assert context.message is None + assert context.configuration is None + assert context.get_user_input() == '' + assert context.task_id == mock_task.id + assert context.context_id == mock_task.context_id + assert context.current_task == mock_task + + def test_init_with_cancel_task_request_without_metadata( + self, mock_task: Mock + ) -> None: + """Test RequestContext returns empty metadata when CancelTaskRequest has none.""" + cancel_request = CancelTaskRequest(id=mock_task.id) + context = RequestContext( + call_context=ServerCallContext(), + request=cancel_request, + task_id=mock_task.id, + context_id=mock_task.context_id, + task=mock_task, + ) + + assert context.metadata == {} + assert context.message is None diff --git a/tests/server/request_handlers/test_default_request_handler.py b/tests/server/request_handlers/test_default_request_handler.py index 27bd4f16d..7c3986cd3 100644 --- a/tests/server/request_handlers/test_default_request_handler.py +++ b/tests/server/request_handlers/test_default_request_handler.py @@ -423,11 +423,16 @@ async def test_on_cancel_task_cancels_running_agent(agent_card): 'a2a.server.request_handlers.default_request_handler.ResultAggregator', return_value=mock_result_aggregator_instance, ): - params = CancelTaskRequest(id=f'{task_id}') + params = CancelTaskRequest( + id=f'{task_id}', + metadata={'reason': 'user_aborted'}, + ) await request_handler.on_cancel_task(params, context) mock_producer_task.cancel.assert_called_once() mock_agent_executor.cancel.assert_awaited_once() + cancel_ctx: RequestContext = mock_agent_executor.cancel.call_args.args[0] + assert cancel_ctx.metadata == {'reason': 'user_aborted'} @pytest.mark.asyncio 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 9043df77a..3eb90353b 100644 --- a/tests/server/request_handlers/test_default_request_handler_v2.py +++ b/tests/server/request_handlers/test_default_request_handler_v2.py @@ -2506,6 +2506,7 @@ def __init__(self) -> None: self.working = asyncio.Event() self.release = asyncio.Event() self.cancel_called = asyncio.Event() + self.cancel_metadata: dict[str, Any] | None = None async def execute( self, context: RequestContext, event_queue: EventQueue @@ -2520,6 +2521,7 @@ async def execute( async def cancel( self, context: RequestContext, event_queue: EventQueue ) -> None: + self.cancel_metadata = context.metadata self.cancel_called.set() updater = TaskUpdater(event_queue, context.task_id, context.context_id) await updater.cancel() @@ -2598,6 +2600,39 @@ async def test_on_cancel_task_is_owner_scoped_for_live_task(): await handler.aclose() +@pytest.mark.timeout(10) +@pytest.mark.asyncio +async def test_on_cancel_task_forwards_metadata_to_executor() -> None: + """Issue #1415: CancelTaskRequest.metadata reaches AgentExecutor.cancel.""" + agent = _HangingAgent() + handler = DefaultRequestHandlerV2( + agent_executor=agent, + task_store=InMemoryTaskStore(), + agent_card=create_default_agent_card(), + ) + alice = _ctx('alice') + + task_id = await _start_live_task(handler, alice, 'work') + await asyncio.wait_for(agent.working.wait(), timeout=5) + + result = await handler.on_cancel_task( + CancelTaskRequest( + id=task_id, + metadata={'reason': 'user_aborted', 'actor': 'alice'}, + ), + alice, + ) + assert result is not None + assert result.status.state == TaskState.TASK_STATE_CANCELED + assert agent.cancel_metadata == { + 'reason': 'user_aborted', + 'actor': 'alice', + } + + agent.release.set() + await handler.aclose() + + @pytest.mark.timeout(10) @pytest.mark.asyncio async def test_on_subscribe_to_task_is_owner_scoped_for_live_task(): diff --git a/tests/server/routes/test_fastapi_routes.py b/tests/server/routes/test_fastapi_routes.py index 293d4b94a..817862039 100644 --- a/tests/server/routes/test_fastapi_routes.py +++ b/tests/server/routes/test_fastapi_routes.py @@ -130,10 +130,17 @@ def test_request_body_schemas_are_attached( } send_body = schema['paths']['/message:send']['post']['requestBody'] + assert send_body['required'] is True assert send_body['content']['application/json']['schema'] == { '$ref': '#/components/schemas/SendMessageRequest' } + cancel_body = schema['paths']['/tasks/{id}:cancel']['post']['requestBody'] + assert cancel_body['required'] is False + assert cancel_body['content']['application/json']['schema'] == { + '$ref': '#/components/schemas/CancelTaskRequest' + } + assert 'Message' in components assert 'Part' in components assert components['Message']['properties']['role']['enum'] == [ @@ -152,7 +159,7 @@ def test_routes_without_body_have_no_request_body( assert 'requestBody' not in paths[AGENT_CARD_WELL_KNOWN_PATH]['get'] assert 'requestBody' not in paths['/tasks']['get'] - assert 'requestBody' not in paths['/tasks/{id}:cancel']['post'] + assert 'requestBody' not in paths['/tasks/{id}:subscribe']['post'] assert ( 'requestBody' not in paths['/tasks/{id}/pushNotificationConfigs/{push_id}']['delete'] diff --git a/tests/server/routes/test_proto_schema.py b/tests/server/routes/test_proto_schema.py index 8191c37ae..e11ea9c20 100644 --- a/tests/server/routes/test_proto_schema.py +++ b/tests/server/routes/test_proto_schema.py @@ -170,9 +170,10 @@ def test_field_schema_map_entry(): assert schema == {'type': 'object'} -def test_rest_body_types_coverage(): +def test_rest_body_types_coverage() -> None: assert ('/message:send', 'POST') in REST_BODY_TYPES assert ('/message:stream', 'POST') in REST_BODY_TYPES + assert ('/tasks/{id}:cancel', 'POST') in REST_BODY_TYPES assert ('/tasks/{id}/pushNotificationConfigs', 'POST') in REST_BODY_TYPES diff --git a/tests/server/routes/test_rest_dispatcher.py b/tests/server/routes/test_rest_dispatcher.py index 4865b0cb5..3d7bb3986 100644 --- a/tests/server/routes/test_rest_dispatcher.py +++ b/tests/server/routes/test_rest_dispatcher.py @@ -17,6 +17,7 @@ AgentCapabilities, AgentCard, Artifact, + CancelTaskRequest, ListTaskPushNotificationConfigsResponse, ListTasksResponse, Message, @@ -194,6 +195,61 @@ async def test_on_cancel_task_not_found( response = await rest_dispatcher_instance.on_cancel_task(req) assert response.status_code == 404 # TaskNotFoundError maps to 404 + async def test_on_cancel_task_forwards_metadata_and_keeps_path_id( + self, + rest_dispatcher_instance: RestDispatcher, + mock_handler: AsyncMock, + ) -> None: + body = json.dumps( + { + 'id': 'body_id_ignored', + 'metadata': {'reason': 'user_aborted', 'nested': {'step': 3}}, + } + ).encode('utf-8') + req = make_mock_request( + method='POST', path_params={'id': 'test_task'}, body=body + ) + + response = await rest_dispatcher_instance.on_cancel_task(req) + assert response.status_code == 200 + mock_handler.on_cancel_task.assert_awaited_once() + cancel_params = mock_handler.on_cancel_task.call_args.args[0] + assert cancel_params == CancelTaskRequest( + id='test_task', + metadata={'reason': 'user_aborted', 'nested': {'step': 3}}, + ) + + @pytest.mark.parametrize('body', [b'', b' \n\t ']) + async def test_on_cancel_task_accepts_empty_body( + self, + rest_dispatcher_instance: RestDispatcher, + mock_handler: AsyncMock, + body: bytes, + ) -> None: + req = make_mock_request( + method='POST', path_params={'id': 'test_task'}, body=body + ) + + response = await rest_dispatcher_instance.on_cancel_task(req) + assert response.status_code == 200 + mock_handler.on_cancel_task.assert_awaited_once() + cancel_params = mock_handler.on_cancel_task.call_args.args[0] + assert cancel_params == CancelTaskRequest(id='test_task') + assert not cancel_params.HasField('metadata') + + async def test_on_cancel_task_invalid_json_body_returns_400( + self, + rest_dispatcher_instance: RestDispatcher, + mock_handler: AsyncMock, + ) -> None: + req = make_mock_request( + method='POST', path_params={'id': 'test_task'}, body=b'{not-json' + ) + + response = await rest_dispatcher_instance.on_cancel_task(req) + assert response.status_code == 400 + mock_handler.on_cancel_task.assert_not_awaited() + async def test_on_get_task_success( self, rest_dispatcher_instance, mock_handler ):