fix(server,client): return Task/Message directly from message/send

src/a2a/client/transports/jsonrpc.py

@@ -9,25 +9,26 @@
 from google.protobuf import json_format
 from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response
 
 from a2a.client.client import ClientCallContext
 from a2a.client.errors import A2AClientError
 from a2a.client.transports.base import ClientTransport
 from a2a.client.transports.http_helpers import (
     get_http_args,
     send_http_request,
     send_http_stream_request,
 )
 from a2a.types.a2a_pb2 import (
     AgentCard,
     CancelTaskRequest,
     DeleteTaskPushNotificationConfigRequest,
     GetExtendedAgentCardRequest,
     GetTaskPushNotificationConfigRequest,
     GetTaskRequest,
     ListTaskPushNotificationConfigsRequest,
     ListTaskPushNotificationConfigsResponse,
     ListTasksRequest,
     ListTasksResponse,
+    Message,
     SendMessageRequest,
     SendMessageResponse,
     StreamResponse,
@@ -47,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."""
@@ -70,20 +88,54 @@
     ) -> SendMessageResponse:
         """Sends a non-streaming message request to the agent."""
         rpc_request = JSONRPC20Request(
             method='SendMessage',
             params=json_format.MessageToDict(request),
             _id=str(uuid4()),
         )
         response_data = await self._send_request(
             dict(rpc_request.data), context
         )
         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
+        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": {...}}.
+        if isinstance(result, dict) and 'task' in result:
+            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(
+                _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 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 result:
+                kind = 'task'
+        payload = _strip_kind(result)
+        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(payload, Task())
+        except json_format.ParseError:
+            message = json_format.ParseDict(payload, Message())
+            return SendMessageResponse(message=message)
+        return SendMessageResponse(task=task)
 
     async def send_message_streaming(
         self,

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

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)
 

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()
@@ -157,6 +155,167 @@ 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.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.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()
     )
@@ -527,11 +686,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()

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(

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',
             },

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()