fix(client): raise A2A errors for malformed JSON-RPC responses

src/a2a/client/transports/jsonrpc.py

@@ -1,3 +1,4 @@
+import json
 import logging
 
 from collections.abc import AsyncGenerator
@@ -7,35 +8,35 @@
 import httpx
 
 from google.protobuf import json_format
-from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response
+from jsonrpc.jsonrpc2 import JSONRPC20Request
 
 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,
     SendMessageRequest,
     SendMessageResponse,
     StreamResponse,
     SubscribeToTaskRequest,
     Task,
     TaskPushNotificationConfig,
 )
 from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP
 from a2a.utils.proto_utils import parse_lenient
 from a2a.utils.telemetry import SpanKind, trace_class
 
@@ -80,11 +81,8 @@
         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 = parse_lenient(
-            json_rpc_response.result, SendMessageResponse()
+            response_data['result'], SendMessageResponse()
         )
         return response
 
@@ -114,20 +112,17 @@
     ) -> Task:
         """Retrieves the current state and history of a specific task."""
         rpc_request = JSONRPC20Request(
             method='GetTask',
             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: Task = parse_lenient(json_rpc_response.result, Task())
+        response: Task = parse_lenient(response_data['result'], Task())
         return response
 
     async def list_tasks(
         self,
         request: ListTasksRequest,
         *,
@@ -142,11 +137,8 @@
         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: ListTasksResponse = parse_lenient(
-            json_rpc_response.result, ListTasksResponse()
+            response_data['result'], ListTasksResponse()
         )
         return response
 
@@ -158,20 +150,17 @@
     ) -> Task:
         """Requests the agent to cancel a specific task."""
         rpc_request = JSONRPC20Request(
             method='CancelTask',
             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: Task = parse_lenient(json_rpc_response.result, Task())
+        response: Task = parse_lenient(response_data['result'], Task())
         return response
 
     async def create_task_push_notification_config(
         self,
         request: TaskPushNotificationConfig,
         *,
@@ -179,22 +168,19 @@
     ) -> TaskPushNotificationConfig:
         """Sets or updates the push notification configuration for a specific task."""
         rpc_request = JSONRPC20Request(
             method='CreateTaskPushNotificationConfig',
             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: TaskPushNotificationConfig = parse_lenient(
-            json_rpc_response.result, TaskPushNotificationConfig()
+            response_data['result'], TaskPushNotificationConfig()
         )
         return response
 
     async def get_task_push_notification_config(
         self,
         request: GetTaskPushNotificationConfigRequest,
         *,
@@ -202,22 +188,19 @@
     ) -> TaskPushNotificationConfig:
         """Retrieves the push notification configuration for a specific task."""
         rpc_request = JSONRPC20Request(
             method='GetTaskPushNotificationConfig',
             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: TaskPushNotificationConfig = parse_lenient(
-            json_rpc_response.result, TaskPushNotificationConfig()
+            response_data['result'], TaskPushNotificationConfig()
         )
         return response
 
     async def list_task_push_notification_configs(
         self,
         request: ListTaskPushNotificationConfigsRequest,
         *,
@@ -232,11 +215,8 @@
         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: ListTaskPushNotificationConfigsResponse = parse_lenient(
-            json_rpc_response.result, ListTaskPushNotificationConfigsResponse()
+            response_data['result'], ListTaskPushNotificationConfigsResponse()
         )
         return response
 
@@ -252,12 +232,7 @@
             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)
+        await self._send_request(dict(rpc_request.data), context)
 
     async def subscribe(
         self,
@@ -297,36 +272,30 @@
             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)
+        result = response_data['result']
         # Validate type of the response
-        if not isinstance(json_rpc_response.result, dict):
-            raise A2AClientError(
-                f'Invalid response type: {type(json_rpc_response.result)}'
-            )
-        response: AgentCard = parse_lenient(
-            json_rpc_response.result, AgentCard()
-        )
+        if not isinstance(result, dict):
+            raise A2AClientError(f'Invalid response type: {type(result)}')
+        response: AgentCard = parse_lenient(result, AgentCard())
 
         return response
 
     async def close(self) -> None:
         """Closes the httpx client unless it is shared."""
         if self.is_shared_httpx:
             logger.debug(
                 'Cannot close a shared httpx client; only the provider '
                 'that supplied it via ClientConfig.httpx_client can close it.'
             )
             return
         if self.httpx_client.is_closed:
             raise A2AClientError('HTTP client is already closed')
         try:
             await self.httpx_client.aclose()
         except Exception as e:
             raise A2AClientError(f'Failed to close HTTP client: {e}') from e
 
     def _create_jsonrpc_error(self, error_dict: dict[str, Any]) -> Exception:
         """Creates the appropriate A2AError from a JSON-RPC error dictionary."""
         code = error_dict.get('code')
         message = error_dict.get('message', str(error_dict))
@@ -357,7 +326,37 @@
         request = self.httpx_client.build_request(
             'POST', self.url, json=payload, **(http_kwargs or {})
         )
-        return await send_http_request(self.httpx_client, request)
+        response_data = await send_http_request(self.httpx_client, request)
+        return self._check_jsonrpc_response(response_data)
+
+    def _check_jsonrpc_response(self, response_data: Any) -> dict[str, Any]:
+        """Returns a JSON-RPC success response, raising for anything else.
+
+        An error response raises the A2A error mapped from its code, even when
+        the error object lacks a ``message`` or carries extra members. A body
+        that is not a JSON-RPC 2.0 response object (a list, a string, an object
+        with neither ``result`` nor ``error``, or an ``error`` that is not an
+        object) raises ``A2AClientError`` instead of leaking a ``TypeError`` or
+        ``ValueError`` from the ``json-rpc`` package.
+        """
+        if not isinstance(response_data, dict):
+            raise A2AClientError(
+                'Invalid JSON-RPC response: expected an object, got '
+                f'{type(response_data).__name__}'
+            )
+        error = response_data.get('error')
+        if error:
+            if not isinstance(error, dict):
+                raise A2AClientError(
+                    f'Invalid JSON-RPC response: error is not an object: '
+                    f'{error!r}'
+                )
+            raise self._create_jsonrpc_error(error)
+        if 'result' not in response_data:
+            raise A2AClientError(
+                "Invalid JSON-RPC response: it has neither 'result' nor 'error'"
+            )
+        return response_data
 
     async def _send_stream_request(
         self,
@@ -375,17 +374,26 @@
             json=rpc_request_payload,
             **http_kwargs,
         ):
-            json_rpc_response = JSONRPC20Response.from_json(sse_data)
-            if json_rpc_response.error:
-                raise self._create_jsonrpc_error(json_rpc_response.error)
+            try:
+                response_data = json.loads(sse_data)
+            except json.JSONDecodeError as e:
+                raise A2AClientError(
+                    f'Invalid JSON-RPC response in SSE stream: {e}'
+                ) from e
+            response_data = self._check_jsonrpc_response(response_data)
             response: StreamResponse = parse_lenient(
-                json_rpc_response.result, StreamResponse()
+                response_data['result'], StreamResponse()
             )
             yield response
 
     def _handle_sse_error(self, sse_data: str) -> NoReturn:
         """Handles SSE error events by parsing JSON-RPC error payload and raising the appropriate domain error."""
-        json_rpc_response = JSONRPC20Response.from_json(sse_data)
-        if json_rpc_response.error:
-            raise self._create_jsonrpc_error(json_rpc_response.error)
+        try:
+            response_data = json.loads(sse_data)
+        except json.JSONDecodeError:
+            response_data = None
+        if isinstance(response_data, dict) and isinstance(
+            response_data.get('error'), dict
+        ):
+            raise self._create_jsonrpc_error(response_data['error'])
         raise A2AClientError(f'SSE stream error: {sse_data}')

tests/client/transports/test_jsonrpc_client.py

@@ -30,7 +30,7 @@
     TaskPushNotificationConfig,
     TaskState,
 )
-from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP
+from a2a.utils.errors import JSON_RPC_ERROR_CODE_MAP, TaskNotFoundError
 from google.protobuf import json_format
 
 
@@ -828,3 +828,134 @@ async def test_send_message_streaming_ignores_unrecognized_fields(
         assert len(events) == 1
         assert events[0].message.message_id == 'r1'
         assert events[0].message.parts[0].text == 'hello'
+
+
+# A JSON-RPC error object with a member beyond code, message and data.
+_ERROR_WITH_EXTRA_MEMBER = {
+    'jsonrpc': '2.0',
+    'id': '1',
+    'error': {'code': -32001, 'message': 'Task not found', 'reason': 'x'},
+}
+
+
+class TestMalformedJsonRpcResponses:
+    """Responses that are not valid JSON-RPC 2.0 raise A2A errors.
+
+    The transport used to build ``JSONRPC20Response`` objects from the raw
+    body, so these responses leaked ``TypeError``/``ValueError`` from the
+    ``json-rpc`` package instead of an ``A2AClientError``.
+    """
+
+    @pytest.mark.parametrize(
+        'body',
+        [
+            pytest.param(
+                [{'jsonrpc': '2.0', 'id': '1', 'result': {}}], id='list'
+            ),
+            pytest.param('Bad Gateway', id='string'),
+            pytest.param(
+                {'jsonrpc': '2.0', 'id': '1'}, id='no-result-or-error'
+            ),
+            pytest.param(
+                {'jsonrpc': '2.0', 'id': '1', 'error': 'boom'},
+                id='error-not-an-object',
+            ),
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_invalid_envelope_raises_client_error(
+        self, transport, mock_httpx_client, body
+    ):
+        mock_response = MagicMock()
+        mock_response.json.return_value = body
+        mock_response.raise_for_status = MagicMock()
+        mock_httpx_client.send.return_value = mock_response
+
+        with pytest.raises(A2AClientError, match='Invalid JSON-RPC response'):
+            await transport.get_task(GetTaskRequest(id='task-1'))
+
+    @pytest.mark.parametrize(
+        'error',
+        [
+            pytest.param(_ERROR_WITH_EXTRA_MEMBER['error'], id='extra-member'),
+            pytest.param({'code': -32001}, id='no-message'),
+        ],
+    )
+    @pytest.mark.asyncio
+    async def test_nonstandard_error_object_maps_to_a2a_error(
+        self, transport, mock_httpx_client, error
+    ):
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            'jsonrpc': '2.0',
+            'id': '1',
+            'error': error,
+        }
+        mock_response.raise_for_status = MagicMock()
+        mock_httpx_client.send.return_value = mock_response
+
+        with pytest.raises(TaskNotFoundError):
+            await transport.get_task(GetTaskRequest(id='task-1'))
+
+    @pytest.mark.asyncio
+    async def test_delete_push_config_error_still_raises(
+        self, transport, mock_httpx_client
+    ):
+        """The delete call ignores the result but must still raise errors."""
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            'jsonrpc': '2.0',
+            'id': '1',
+            'error': {'code': -32001, 'message': 'Task not found'},
+        }
+        mock_response.raise_for_status = MagicMock()
+        mock_httpx_client.send.return_value = mock_response
+
+        with pytest.raises(TaskNotFoundError):
+            await transport.delete_task_push_notification_config(
+                DeleteTaskPushNotificationConfigRequest(
+                    task_id='task-1', id='config-1'
+                )
+            )
+
+    @pytest.mark.parametrize(
+        ('lines', 'expected_error'),
+        [
+            pytest.param(['data: not json', ''], A2AClientError, id='not-json'),
+            pytest.param(
+                [f'data: {json.dumps({"jsonrpc": "2.0", "id": "1"})}', ''],
+                A2AClientError,
+                id='no-result-or-error',
+            ),
+            pytest.param(
+                [f'data: {json.dumps(_ERROR_WITH_EXTRA_MEMBER)}', ''],
+                TaskNotFoundError,
+                id='error-extra-member',
+            ),
+            pytest.param(
+                [
+                    'event: error',
+                    f'data: {json.dumps(_ERROR_WITH_EXTRA_MEMBER)}',
+                    '',
+                ],
+                TaskNotFoundError,
+                id='error-event-extra-member',
+            ),
+        ],
+    )
+    @pytest.mark.asyncio
+    @patch('a2a.client.transports.http_helpers._SSEEventSource')
+    async def test_streaming_malformed_event_raises_a2a_error(
+        self, mock_sse_source: AsyncMock, transport, lines, expected_error
+    ):
+        mock_response = AsyncMock(spec=httpx.Response)
+        mock_response.raise_for_status = MagicMock()
+        mock_response.headers = {'content-type': 'text/event-stream'}
+        mock_response.aiter_lines.return_value = async_iterable_from_list(lines)
+        mock_sse_source.return_value.__aenter__.return_value = mock_response
+
+        with pytest.raises(expected_error):
+            async for _ in transport.send_message_streaming(
+                create_send_message_request()
+            ):
+                pass