From 367be6b0addfda6a06a78ce0a64fdd479de7e093 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 12:30:11 +0000 Subject: [PATCH] fix(a2a): honor agent protocolVersion in the chat completions bridge --- litellm/llms/a2a/chat/transformation.py | 83 ++++++++++++++-- litellm/llms/a2a/common_utils.py | 3 +- .../a2a/chat/test_a2a_chat_transformation.py | 97 ++++++++++++++++++- 3 files changed, 173 insertions(+), 10 deletions(-) diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 113c000f352..6cc185c4fdf 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -19,6 +19,8 @@ from ..common_utils import ( ) from .streaming_iterator import A2AModelResponseIterator +A2A_PROTOCOL_VERSION_PARAM = "a2a_protocol_version" + class A2AConfig(BaseConfig): """ @@ -66,9 +68,10 @@ class A2AConfig(BaseConfig): agent = global_agent_registry.get_agent_by_name(agent_name) if agent: - # Get api_base from agent card URL - if api_base is None and agent.agent_card_params: - api_base = agent.agent_card_params.get("url") + # Get api_base from the agent card + if agent.agent_card_params: + if api_base is None: + api_base = agent.agent_card_params.get("url") # Get api_key, headers, and other params from litellm_params if agent.litellm_params: @@ -225,21 +228,85 @@ class A2AConfig(BaseConfig): "messageId": str(uuid.uuid4()), } - # Build JSON-RPC 2.0 request - # For A2A protocol, the method is "message/send" for non-streaming - # and "message/stream" for streaming stream = optional_params.get("stream", False) - method = "message/stream" if stream else "message/send" + protocol_version = ( + optional_params.get(A2A_PROTOCOL_VERSION_PARAM) + or litellm_params.get(A2A_PROTOCOL_VERSION_PARAM) + or self._pinned_protocol_version_from_registry(model) + ) + + if str(protocol_version) == "1.0": + method = "SendStreamingMessage" if stream else "SendMessage" + params = self._build_v1_send_params(a2a_message) + else: + method = "message/stream" if stream else "message/send" + params = {"message": a2a_message} request_data = { "jsonrpc": "2.0", "id": request_id, "method": method, - "params": {"message": a2a_message}, + "params": params, } return request_data + @staticmethod + def _pinned_protocol_version_from_registry(model: str) -> Optional[str]: + """Look up the agent's pinned ``protocolVersion`` from the proxy registry. + + The completion bridge injects the agent's URL as ``api_base`` and passes the + bare agent name as ``model`` (the ``a2a/`` prefix is stripped upstream), so + the pinned version is not otherwise threaded through to request building. + The registry is the single source of truth; it is only importable in a proxy + context, so this returns ``None`` when running as a plain SDK. + """ + agent_name = model[len("a2a/") :] if model.startswith("a2a/") else model + try: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + except ImportError: + return None + + agent = global_agent_registry.get_agent_by_name(agent_name) + if agent is None or not agent.agent_card_params: + return None + version = agent.agent_card_params.get("protocolVersion") + return str(version) if version is not None else None + + @staticmethod + def _build_v1_send_params(a2a_message: Dict[str, Any]) -> Dict[str, Any]: + """Build A2A 1.0 JSON-RPC ``params`` (protobuf-JSON) from a 0.3 message dict. + + Mirrors how the a2a-sdk 1.x JSON-RPC transport serializes a + ``SendMessageRequest``, so 1.0-only agents receive the envelope they expect. + """ + try: + from a2a.compat.v0_3.conversions import ( + MessageToDict, + to_core_send_message_request, + ) + from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest + except ImportError as e: + raise A2AError( + status_code=500, + message=( + "The 'a2a-sdk' package is required to call an A2A agent pinned to " + "protocolVersion '1.0'. Install it with: pip install a2a-sdk" + ), + ) from e + + request = SendMessageRequest( + id=str(uuid.uuid4()), + params=MessageSendParams.model_validate({"message": a2a_message}), + ) + params: Dict[str, Any] = MessageToDict( + to_core_send_message_request(request), + preserving_proto_field_name=False, + ) + return params + def transform_response( self, model: str, diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 4fc0ff2623e..94e1f7e2dbc 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -80,7 +80,8 @@ def extract_text_from_a2a_message(message: Dict[str, Any], depth: int = 0, max_d text_parts: List[str] = [] for part in parts: - if part.get("kind") == "text": + kind = part.get("kind") + if kind == "text" or (kind is None and "text" in part): text_parts.append(part.get("text", "")) # Handle nested parts if they exist elif "parts" in part: diff --git a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py b/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py index 2e11c68244c..b53188f31fb 100644 --- a/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py +++ b/tests/test_litellm/llms/a2a/chat/test_a2a_chat_transformation.py @@ -2,9 +2,11 @@ from unittest.mock import MagicMock -from litellm.llms.a2a.chat.transformation import A2AConfig +from litellm.llms.a2a.chat.transformation import A2A_PROTOCOL_VERSION_PARAM, A2AConfig from litellm.types.utils import ModelResponse +_MESSAGES = [{"role": "user", "content": "hi there agent"}] + def _raw_response(text: str) -> MagicMock: raw = MagicMock() @@ -21,6 +23,99 @@ def _raw_response(text: str) -> MagicMock: return raw +def _v1_raw_response(text: str) -> MagicMock: + """A2A 1.0 (a2a-sdk 1.x) protobuf-JSON send response: no ``kind`` fields.""" + raw = MagicMock() + raw.status_code = 200 + raw.headers = {} + raw.json.return_value = { + "jsonrpc": "2.0", + "id": "resp-1", + "result": {"message": {"messageId": "m1", "role": "ROLE_AGENT", "parts": [{"text": text}]}}, + } + return raw + + +def test_transform_request_defaults_to_0_3(): + """Without a pinned protocol version the bridge emits the legacy 0.3 method.""" + request = A2AConfig().transform_request("a2a/agent", _MESSAGES, {}, {}, {}) + assert request["method"] == "message/send" + assert request["params"]["message"]["parts"][0]["kind"] == "text" + + +def test_transform_request_v1_uses_send_message(): + """Regression for #32609: a 1.0-pinned agent must get the 1.0 method and the + protobuf-JSON envelope (uppercase role enum, parts without ``kind``).""" + request = A2AConfig().transform_request( + "a2a/agent", _MESSAGES, {A2A_PROTOCOL_VERSION_PARAM: "1.0"}, {}, {} + ) + assert request["method"] == "SendMessage" + message = request["params"]["message"] + assert message["role"] == "ROLE_USER" + assert message["parts"] == [{"text": "user: hi there agent"}] + assert "kind" not in message["parts"][0] + + +def test_transform_request_v1_streaming_uses_send_streaming_message(): + request = A2AConfig().transform_request( + "a2a/agent", _MESSAGES, {A2A_PROTOCOL_VERSION_PARAM: "1.0", "stream": True}, {}, {} + ) + assert request["method"] == "SendStreamingMessage" + + +def test_transform_request_v1_from_litellm_params(): + """The pinned version may arrive via litellm_params (direct SDK usage).""" + request = A2AConfig().transform_request( + "a2a/agent", _MESSAGES, {}, {A2A_PROTOCOL_VERSION_PARAM: "1.0"}, {} + ) + assert request["method"] == "SendMessage" + + +def test_transform_request_reads_pinned_version_from_registry(): + """Regression for #32609: the completion bridge strips the ``a2a/`` prefix and + injects only api_base, so transform_request must resolve the pinned version from + the registry (keyed by the bare agent name) rather than defaulting to 0.3.""" + try: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.types.agents import AgentResponse + except ImportError: + import pytest + + pytest.skip("Registry not available (not in proxy context)") + + agent = AgentResponse( + agent_id="v1-id", + agent_name="v1-agent", + agent_card_params={"url": "http://agent.example.com:9999", "protocolVersion": "1.0"}, + litellm_params={}, + ) + original_agents = global_agent_registry.agent_list.copy() + global_agent_registry.register_agent(agent) + try: + request = A2AConfig().transform_request("v1-agent", _MESSAGES, {}, {}, {}) + assert request["method"] == "SendMessage" + assert "kind" not in request["params"]["message"]["parts"][0] + finally: + global_agent_registry.agent_list = original_agents + + +def test_transform_response_extracts_text_from_v1_message(): + """Regression for #32609: text must be extracted from a 1.0 protobuf-JSON + message result whose parts carry no ``kind`` field.""" + result = A2AConfig().transform_response( + model="a2a/test-agent", + raw_response=_v1_raw_response("hello from a 1.0 agent"), + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=_MESSAGES, + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert result.choices[0].message.content == "hello from a 1.0 agent" + + def test_transform_response_sets_usage(): """Regression: A2AConfig.transform_response must populate usage so per-token pricing computes real cost and callers don't get usage 0/0/0."""