diff --git a/src/a2a/client/transports/jsonrpc.py b/src/a2a/client/transports/jsonrpc.py index b6a687998..e21a4cb62 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, @@ -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.""" @@ -80,10 +98,44 @@ 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 + 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, 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..48dd0bcb4 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() @@ -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() 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()