From 9d8f36b232a6dbbf8d201732503158b149c3c04f Mon Sep 17 00:00:00 2001 From: Richardson Lima Date: Fri, 21 Aug 2026 11:08:38 -0300 Subject: [PATCH 1/4] fix(server,client): return Task/Message directly from message/send instead of wrapping in SendMessageResponse oneof Also updates the JSON-RPC client transport to parse the unwrapped result, since the server no longer nests it under task/message. --- src/a2a/client/transports/jsonrpc.py | 14 ++++++++++---- src/a2a/server/routes/jsonrpc_dispatcher.py | 6 +----- tests/client/test_auth_interceptor.py | 4 +--- tests/client/transports/test_jsonrpc_client.py | 16 ++++++---------- .../test_client_server_integration.py | 4 ++-- tests/integration/test_tenant.py | 2 +- tests/server/test_integration.py | 8 +++----- 7 files changed, 24 insertions(+), 30 deletions(-) diff --git a/src/a2a/client/transports/jsonrpc.py b/src/a2a/client/transports/jsonrpc.py index b6a687998..ea96a6b6f 100644 --- a/src/a2a/client/transports/jsonrpc.py +++ b/src/a2a/client/transports/jsonrpc.py @@ -28,6 +28,7 @@ ListTaskPushNotificationConfigsResponse, ListTasksRequest, ListTasksResponse, + Message, SendMessageRequest, SendMessageResponse, StreamResponse, @@ -80,10 +81,15 @@ async def send_message( json_rpc_response = JSONRPC20Response(**response_data) if json_rpc_response.error: raise self._create_jsonrpc_error(json_rpc_response.error) - response: SendMessageResponse = json_format.ParseDict( - json_rpc_response.result, SendMessageResponse() - ) - return response + # The server returns the Task or Message directly (not wrapped in + # the streaming SendMessageResponse oneof), so try each in turn and + # re-wrap for callers that expect a SendMessageResponse. + try: + task = json_format.ParseDict(json_rpc_response.result, Task()) + except json_format.ParseError: + message = json_format.ParseDict(json_rpc_response.result, Message()) + return SendMessageResponse(message=message) + return SendMessageResponse(task=task) async def send_message_streaming( self, diff --git a/src/a2a/server/routes/jsonrpc_dispatcher.py b/src/a2a/server/routes/jsonrpc_dispatcher.py index b59ed7551..6515202ef 100644 --- a/src/a2a/server/routes/jsonrpc_dispatcher.py +++ b/src/a2a/server/routes/jsonrpc_dispatcher.py @@ -36,9 +36,7 @@ ListTaskPushNotificationConfigsRequest, ListTasksRequest, SendMessageRequest, - SendMessageResponse, SubscribeToTaskRequest, - Task, TaskPushNotificationConfig, ) from a2a.utils import constants, json_utils, proto_utils @@ -405,9 +403,7 @@ async def _handle_send_message( task_or_message = await self.request_handler.on_message_send( request_obj, context ) - if isinstance(task_or_message, Task): - return MessageToDict(SendMessageResponse(task=task_or_message)) - return MessageToDict(SendMessageResponse(message=task_or_message)) + return MessageToDict(task_or_message, preserving_proto_field_name=False) async def _handle_cancel_task( self, request_obj: CancelTaskRequest, context: ServerCallContext diff --git a/tests/client/test_auth_interceptor.py b/tests/client/test_auth_interceptor.py index f7979bff7..6ca8d15b5 100644 --- a/tests/client/test_auth_interceptor.py +++ b/tests/client/test_auth_interceptor.py @@ -33,7 +33,6 @@ SecurityRequirement, SecurityScheme, SendMessageRequest, - SendMessageResponse, StringList, ) from a2a.utils.constants import TransportProtocol @@ -49,11 +48,10 @@ def build_success_response(request: httpx.Request) -> httpx.Response: role=Role.ROLE_AGENT, parts=[], ) - response = SendMessageResponse(message=message) response_payload = { 'id': request_payload['id'], 'jsonrpc': '2.0', - 'result': json_format.MessageToDict(response), + 'result': json_format.MessageToDict(message), } return httpx.Response(200, json=response_payload) diff --git a/tests/client/transports/test_jsonrpc_client.py b/tests/client/transports/test_jsonrpc_client.py index 7e6680e5f..34da4219e 100644 --- a/tests/client/transports/test_jsonrpc_client.py +++ b/tests/client/transports/test_jsonrpc_client.py @@ -137,11 +137,9 @@ async def test_send_message_success(self, transport, mock_httpx_client): 'jsonrpc': '2.0', 'id': '1', 'result': { - 'task': { - 'id': task_id, - 'contextId': 'ctx-123', - 'status': {'state': 'TASK_STATE_COMPLETED'}, - } + 'id': task_id, + 'contextId': 'ctx-123', + 'status': {'state': 'TASK_STATE_COMPLETED'}, }, } mock_response.raise_for_status = MagicMock() @@ -527,11 +525,9 @@ async def test_extensions_added_to_request( 'jsonrpc': '2.0', 'id': '1', 'result': { - 'task': { - 'id': 'task-123', - 'contextId': 'ctx-123', - 'status': {'state': 'TASK_STATE_COMPLETED'}, - } + 'id': 'task-123', + 'contextId': 'ctx-123', + 'status': {'state': 'TASK_STATE_COMPLETED'}, }, } mock_response.raise_for_status = MagicMock() diff --git a/tests/integration/test_client_server_integration.py b/tests/integration/test_client_server_integration.py index b1962b54b..126a63c34 100644 --- a/tests/integration/test_client_server_integration.py +++ b/tests/integration/test_client_server_integration.py @@ -650,11 +650,11 @@ async def test_json_transport_base_client_send_message_with_extensions( with patch.object( transport, '_send_request', new_callable=AsyncMock ) as mock_send_request: - # Mock returns a JSON-RPC response with SendMessageResponse structure + # Mock returns a JSON-RPC response with the Task returned directly. mock_send_request.return_value = { 'id': '123', 'jsonrpc': '2.0', - 'result': {'task': MessageToDict(TASK_FROM_BLOCKING)}, + 'result': MessageToDict(TASK_FROM_BLOCKING), } service_params = ServiceParametersFactory.create( diff --git a/tests/integration/test_tenant.py b/tests/integration/test_tenant.py index a617d26c4..f8be375ab 100644 --- a/tests/integration/test_tenant.py +++ b/tests/integration/test_tenant.py @@ -90,7 +90,7 @@ async def test_tenant_decorator_jsonrpc(self, agent_card): mock_httpx.send.return_value = MagicMock( status_code=200, json=lambda: { - 'result': {'message': {}}, + 'result': {}, 'id': '1', 'jsonrpc': '2.0', }, diff --git a/tests/server/test_integration.py b/tests/server/test_integration.py index cc0678c22..0f2f4043a 100644 --- a/tests/server/test_integration.py +++ b/tests/server/test_integration.py @@ -350,9 +350,8 @@ def test_send_message(client: TestClient, handler: mock.AsyncMock): assert response.status_code == 200 data = response.json() assert 'result' in data - # Result is wrapped in SendMessageResponse with task field - assert data['result']['task']['id'] == 'task1' - assert data['result']['task']['status']['state'] == 'TASK_STATE_SUBMITTED' + assert data['result']['id'] == 'task1' + assert data['result']['status']['state'] == 'TASK_STATE_SUBMITTED' # Verify handler was called handler.on_message_send.assert_awaited_once() @@ -536,8 +535,7 @@ async def authenticate( assert response.status_code == 200 data = response.json() assert 'result' in data - # Result is wrapped in SendMessageResponse with message field - assert data['result']['message']['parts'][0]['text'] == 'test_user' + assert data['result']['parts'][0]['text'] == 'test_user' # Verify handler was called handler.on_message_send.assert_awaited_once() From 3e40d4988d001996f7e1f8cfafcd677e69fe815b Mon Sep 17 00:00:00 2001 From: Richardson Lima Date: Fri, 21 Aug 2026 11:35:58 -0300 Subject: [PATCH 2/4] fix(client): keep backward compat with peers still wrapping message/send in SendMessageResponse The ITK cross-SDK test caught a real gap: when this SDK's client calls a peer that hasn't adopted the unwrapped response shape yet (older SDKs, other language implementations), it now unwraps the legacy {task: {...}} / {message: {...}} shape explicitly before falling back to parsing the direct Task/Message payload this SDK's own server returns. --- src/a2a/client/transports/jsonrpc.py | 19 +++++-- .../client/transports/test_jsonrpc_client.py | 57 +++++++++++++++++++ 2 files changed, 71 insertions(+), 5 deletions(-) diff --git a/src/a2a/client/transports/jsonrpc.py b/src/a2a/client/transports/jsonrpc.py index ea96a6b6f..2414a8eba 100644 --- a/src/a2a/client/transports/jsonrpc.py +++ b/src/a2a/client/transports/jsonrpc.py @@ -81,13 +81,22 @@ async def send_message( json_rpc_response = JSONRPC20Response(**response_data) if json_rpc_response.error: raise self._create_jsonrpc_error(json_rpc_response.error) - # The server returns the Task or Message directly (not wrapped in - # the streaming SendMessageResponse oneof), so try each in turn and - # re-wrap for callers that expect a SendMessageResponse. + result = json_rpc_response.result + # Servers that still nest the payload under the streaming + # SendMessageResponse oneof (older SDKs, other language + # implementations) send {"task": {...}} or {"message": {...}}. + # Unwrap explicitly for those before falling back to the direct + # Task/Message shape this SDK's own server now returns. + if isinstance(result, dict) and 'task' in result: + task = json_format.ParseDict(result['task'], Task()) + return SendMessageResponse(task=task) + if isinstance(result, dict) and 'message' in result: + message = json_format.ParseDict(result['message'], Message()) + return SendMessageResponse(message=message) try: - task = json_format.ParseDict(json_rpc_response.result, Task()) + task = json_format.ParseDict(result, Task()) except json_format.ParseError: - message = json_format.ParseDict(json_rpc_response.result, Message()) + message = json_format.ParseDict(result, Message()) return SendMessageResponse(message=message) return SendMessageResponse(task=task) diff --git a/tests/client/transports/test_jsonrpc_client.py b/tests/client/transports/test_jsonrpc_client.py index 34da4219e..19e82c975 100644 --- a/tests/client/transports/test_jsonrpc_client.py +++ b/tests/client/transports/test_jsonrpc_client.py @@ -155,6 +155,63 @@ async def test_send_message_success(self, transport, mock_httpx_client): payload = call_args[1]['json'] assert payload['method'] == 'SendMessage' + @pytest.mark.asyncio + async def test_send_message_legacy_wrapped_task( + self, transport, mock_httpx_client + ): + """A peer that still nests the payload under the streaming + SendMessageResponse oneof (older SDKs, other language + implementations) should still be parsed correctly. + """ + task_id = str(uuid4()) + mock_response = MagicMock() + mock_response.json.return_value = { + 'jsonrpc': '2.0', + 'id': '1', + 'result': { + 'task': { + 'id': task_id, + 'contextId': 'ctx-123', + 'status': {'state': 'TASK_STATE_COMPLETED'}, + } + }, + } + mock_response.raise_for_status = MagicMock() + mock_httpx_client.send.return_value = mock_response + + request = create_send_message_request() + response = await transport.send_message(request) + + assert response.HasField('task') + assert response.task.id == task_id + assert response.task.status.state == TaskState.TASK_STATE_COMPLETED + + @pytest.mark.asyncio + async def test_send_message_legacy_wrapped_message( + self, transport, mock_httpx_client + ): + """Same as above, but for a peer returning a wrapped Message.""" + mock_response = MagicMock() + mock_response.json.return_value = { + 'jsonrpc': '2.0', + 'id': '1', + 'result': { + 'message': { + 'messageId': 'msg-1', + 'role': 'ROLE_AGENT', + 'parts': [{'text': 'hi'}], + } + }, + } + mock_response.raise_for_status = MagicMock() + mock_httpx_client.send.return_value = mock_response + + request = create_send_message_request() + response = await transport.send_message(request) + + assert response.HasField('message') + assert response.message.message_id == 'msg-1' + @pytest.mark.parametrize( 'error_cls, error_code', JSON_RPC_ERROR_CODE_MAP.items() ) From 9eb85b557fc29ada9e2d5d4a35ebaf9bfcc5641d Mon Sep 17 00:00:00 2001 From: Richardson Lima Date: Fri, 21 Aug 2026 11:59:07 -0300 Subject: [PATCH 3/4] fix(client): handle the spec kind discriminator on unwrapped message/send responses Peers that already return the direct Task/Message shape (per the A2A spec) include a top-level "kind" field to disambiguate the two, since that's not encoded via a protobuf oneof on the wire. Our generated Task/Message types don't declare that field, so ParseDict rejected it on both attempts. Strip and use it (falling back to the same messageId/id field-presence heuristic the v0.3 compat transport already uses) before parsing. --- src/a2a/client/transports/jsonrpc.py | 24 ++++++-- .../client/transports/test_jsonrpc_client.py | 55 +++++++++++++++++++ 2 files changed, 75 insertions(+), 4 deletions(-) diff --git a/src/a2a/client/transports/jsonrpc.py b/src/a2a/client/transports/jsonrpc.py index 2414a8eba..bf5f294a9 100644 --- a/src/a2a/client/transports/jsonrpc.py +++ b/src/a2a/client/transports/jsonrpc.py @@ -85,18 +85,34 @@ async def send_message( # Servers that still nest the payload under the streaming # SendMessageResponse oneof (older SDKs, other language # implementations) send {"task": {...}} or {"message": {...}}. - # Unwrap explicitly for those before falling back to the direct - # Task/Message shape this SDK's own server now returns. if isinstance(result, dict) and 'task' in result: task = json_format.ParseDict(result['task'], Task()) return SendMessageResponse(task=task) if isinstance(result, dict) and 'message' in result: message = json_format.ParseDict(result['message'], Message()) return SendMessageResponse(message=message) + # Otherwise the payload is the Task/Message itself, per spec + # possibly carrying a "kind" discriminator field that our + # protobuf-generated Task/Message types don't declare. Strip it + # and use it (falling back to field-presence, same heuristic the + # v0.3 compat transport already uses) to pick which type to parse. + payload = dict(result) if isinstance(result, dict) else result + kind = payload.pop('kind', None) if isinstance(payload, dict) else None + if not kind and isinstance(payload, dict): + if 'messageId' in payload: + kind = 'message' + elif 'id' in payload: + kind = 'task' + if kind == 'message': + message = json_format.ParseDict(payload, Message()) + return SendMessageResponse(message=message) + if kind == 'task': + task = json_format.ParseDict(payload, Task()) + return SendMessageResponse(task=task) try: - task = json_format.ParseDict(result, Task()) + task = json_format.ParseDict(payload, Task()) except json_format.ParseError: - message = json_format.ParseDict(result, Message()) + message = json_format.ParseDict(payload, Message()) return SendMessageResponse(message=message) return SendMessageResponse(task=task) diff --git a/tests/client/transports/test_jsonrpc_client.py b/tests/client/transports/test_jsonrpc_client.py index 19e82c975..e62f28b84 100644 --- a/tests/client/transports/test_jsonrpc_client.py +++ b/tests/client/transports/test_jsonrpc_client.py @@ -212,6 +212,61 @@ async def test_send_message_legacy_wrapped_message( assert response.HasField('message') assert response.message.message_id == 'msg-1' + @pytest.mark.asyncio + async def test_send_message_unwrapped_with_kind_task( + self, transport, mock_httpx_client + ): + """A spec-compliant peer (e.g. another language SDK) sends the + Task unwrapped but with a "kind" discriminator field, which our + protobuf-generated Task type doesn't declare. It must be + stripped rather than break parsing. + """ + task_id = str(uuid4()) + mock_response = MagicMock() + mock_response.json.return_value = { + 'jsonrpc': '2.0', + 'id': '1', + 'result': { + 'id': task_id, + 'kind': 'task', + 'contextId': 'ctx-123', + 'status': {'state': 'TASK_STATE_COMPLETED'}, + }, + } + mock_response.raise_for_status = MagicMock() + mock_httpx_client.send.return_value = mock_response + + request = create_send_message_request() + response = await transport.send_message(request) + + assert response.HasField('task') + assert response.task.id == task_id + + @pytest.mark.asyncio + async def test_send_message_unwrapped_with_kind_message( + self, transport, mock_httpx_client + ): + """Same as above, but for a peer's unwrapped Message with kind.""" + mock_response = MagicMock() + mock_response.json.return_value = { + 'jsonrpc': '2.0', + 'id': '1', + 'result': { + 'messageId': 'msg-1', + 'kind': 'message', + 'role': 'ROLE_AGENT', + 'parts': [{'text': 'hi'}], + }, + } + mock_response.raise_for_status = MagicMock() + mock_httpx_client.send.return_value = mock_response + + request = create_send_message_request() + response = await transport.send_message(request) + + assert response.HasField('message') + assert response.message.message_id == 'msg-1' + @pytest.mark.parametrize( 'error_cls, error_code', JSON_RPC_ERROR_CODE_MAP.items() ) From e597084b5b0a6c746db5be63099f16e750dbf0c1 Mon Sep 17 00:00:00 2001 From: Richardson Lima Date: Fri, 21 Aug 2026 13:20:56 -0300 Subject: [PATCH 4/4] fix(client): strip the kind discriminator recursively, not just at the top level The previous fix only popped "kind" off the root of the result dict. Spec-compliant peers stamp it on every nested Task/Message/Part too (TaskStatus.message, Task.history[], Message.parts[]), so ParseDict still rejected the payload as soon as it recursed into any of those. Added _strip_kind() to clean the whole tree before parsing. --- src/a2a/client/transports/jsonrpc.py | 43 +++++++++++----- .../client/transports/test_jsonrpc_client.py | 49 +++++++++++++++++++ 2 files changed, 81 insertions(+), 11 deletions(-) diff --git a/src/a2a/client/transports/jsonrpc.py b/src/a2a/client/transports/jsonrpc.py index bf5f294a9..e21a4cb62 100644 --- a/src/a2a/client/transports/jsonrpc.py +++ b/src/a2a/client/transports/jsonrpc.py @@ -48,6 +48,23 @@ _ERROR_INFO_TYPE = 'type.googleapis.com/google.rpc.ErrorInfo' +def _strip_kind(value: Any) -> Any: + """Recursively drops "kind" discriminator keys from a decoded JSON value. + + Spec-compliant peers stamp "kind" on Task, Message, and Part + objects wherever they appear (top level, TaskStatus.message, + Task.history[], Message.parts[]), but this SDK's protobuf-generated + types don't declare that field, so json_format.ParseDict rejects it. + """ + if isinstance(value, dict): + return { + key: _strip_kind(val) for key, val in value.items() if key != 'kind' + } + if isinstance(value, list): + return [_strip_kind(item) for item in value] + return value + + @trace_class(kind=SpanKind.CLIENT) class JsonRpcTransport(ClientTransport): """A JSON-RPC transport for the A2A client.""" @@ -86,23 +103,27 @@ async def send_message( # SendMessageResponse oneof (older SDKs, other language # implementations) send {"task": {...}} or {"message": {...}}. if isinstance(result, dict) and 'task' in result: - task = json_format.ParseDict(result['task'], Task()) + task = json_format.ParseDict(_strip_kind(result['task']), Task()) return SendMessageResponse(task=task) if isinstance(result, dict) and 'message' in result: - message = json_format.ParseDict(result['message'], Message()) + message = json_format.ParseDict( + _strip_kind(result['message']), Message() + ) return SendMessageResponse(message=message) # Otherwise the payload is the Task/Message itself, per spec - # possibly carrying a "kind" discriminator field that our - # protobuf-generated Task/Message types don't declare. Strip it - # and use it (falling back to field-presence, same heuristic the - # v0.3 compat transport already uses) to pick which type to parse. - payload = dict(result) if isinstance(result, dict) else result - kind = payload.pop('kind', None) if isinstance(payload, dict) else None - if not kind and isinstance(payload, dict): - if 'messageId' in payload: + # possibly carrying a "kind" discriminator field on itself and on + # every nested Message/Part (status.message, history[], parts[]), + # none of which this SDK's protobuf-generated types declare. + # Read the top-level kind (or fall back to the same field-presence + # heuristic the v0.3 compat transport already uses) before + # stripping it throughout the tree. + kind = result.get('kind') if isinstance(result, dict) else None + if not kind and isinstance(result, dict): + if 'messageId' in result: kind = 'message' - elif 'id' in payload: + elif 'id' in result: kind = 'task' + payload = _strip_kind(result) if kind == 'message': message = json_format.ParseDict(payload, Message()) return SendMessageResponse(message=message) diff --git a/tests/client/transports/test_jsonrpc_client.py b/tests/client/transports/test_jsonrpc_client.py index e62f28b84..48dd0bcb4 100644 --- a/tests/client/transports/test_jsonrpc_client.py +++ b/tests/client/transports/test_jsonrpc_client.py @@ -267,6 +267,55 @@ async def test_send_message_unwrapped_with_kind_message( assert response.HasField('message') assert response.message.message_id == 'msg-1' + @pytest.mark.asyncio + async def test_send_message_unwrapped_with_nested_kind( + self, transport, mock_httpx_client + ): + """A peer stamps "kind" on every Task/Message/Part it emits, not + just the top-level object, e.g. TaskStatus.message, Task.history + entries, and Message.parts entries. All of them must be stripped, + not just the one at the root. + """ + task_id = str(uuid4()) + mock_response = MagicMock() + mock_response.json.return_value = { + 'jsonrpc': '2.0', + 'id': '1', + 'result': { + 'id': task_id, + 'kind': 'task', + 'contextId': 'ctx-123', + 'status': { + 'state': 'TASK_STATE_COMPLETED', + 'message': { + 'messageId': 'msg-1', + 'kind': 'message', + 'role': 'ROLE_AGENT', + 'parts': [{'kind': 'text', 'text': 'hi'}], + }, + }, + 'history': [ + { + 'messageId': 'msg-0', + 'kind': 'message', + 'role': 'ROLE_USER', + 'parts': [{'kind': 'text', 'text': 'hello'}], + } + ], + }, + } + mock_response.raise_for_status = MagicMock() + mock_httpx_client.send.return_value = mock_response + + request = create_send_message_request() + response = await transport.send_message(request) + + assert response.HasField('task') + assert response.task.id == task_id + assert response.task.status.message.message_id == 'msg-1' + assert response.task.status.message.parts[0].text == 'hi' + assert response.task.history[0].message_id == 'msg-0' + @pytest.mark.parametrize( 'error_cls, error_code', JSON_RPC_ERROR_CODE_MAP.items() )