diff --git a/README.md b/README.md index 3d0f7282d7c..90d3e944fcc 100644 --- a/README.md +++ b/README.md @@ -156,35 +156,41 @@ response = await client.send_message(request) ### AI Gateway (Proxy Server) -**Step 1.** [Add your Agent to the AI Gateway](https://docs.litellm.ai/docs/a2a#adding-your-agent) +**Step 1.** [Add your Agent to the AI Gateway](https://docs.litellm.ai/docs/a2a#adding-your-agent) — set `protocolVersion` to `1.0` or `0.3` per agent -**Step 2.** Call Agent via A2A SDK +**Step 2.** Call Agent via A2A SDK (requires `a2a-sdk>=1.1.0`) ```python -from a2a.client import A2ACardResolver, A2AClient -from a2a.types import MessageSendParams, SendMessageRequest -from uuid import uuid4 import httpx +from a2a.client import A2ACardResolver, ClientConfig, ClientFactory +from a2a.types import Message, Part, Role, SendMessageRequest +from a2a.utils.constants import TransportProtocol +from uuid import uuid4 base_url = "http://localhost:4000/a2a/my-agent" # LiteLLM proxy + agent name headers = {"Authorization": "Bearer sk-1234"} # LiteLLM Virtual Key -async with httpx.AsyncClient(headers=headers) as httpx_client: - resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url) +async with httpx.AsyncClient(headers=headers, timeout=60.0) as http_client: + resolver = A2ACardResolver(httpx_client=http_client, base_url=base_url) agent_card = await resolver.get_agent_card() - client = A2AClient(httpx_client=httpx_client, agent_card=agent_card) + config = ClientConfig( + httpx_client=http_client, + streaming=False, + supported_protocol_bindings=[TransportProtocol.JSONRPC, TransportProtocol.HTTP_JSON], + ) + client = ClientFactory(config).create(agent_card) request = SendMessageRequest( - id=str(uuid4()), - params=MessageSendParams( - message={ - "role": "user", - "parts": [{"kind": "text", "text": "Hello!"}], - "messageId": uuid4().hex, - } + message=Message( + message_id=uuid4().hex, + role=Role.ROLE_USER, + parts=[Part(text="Hello!")], ) ) - response = await client.send_message(request) + async for event in client.send_message(request): + populated = event.ListFields() + if populated and populated[0][0].name in ("message", "msg"): + print("".join(getattr(p, "text", "") or "" for p in populated[0][1].parts)) ``` [**Docs: A2A Agent Gateway**](https://docs.litellm.ai/docs/a2a) diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 1955b5268e1..412c7a0897d 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -4,7 +4,7 @@ Custom A2A Card Resolver for LiteLLM. Extends the A2A SDK's card resolver to support multiple well-known paths. """ -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict from litellm._logging import verbose_logger from litellm.constants import LOCALHOST_URL_PATTERNS @@ -27,7 +27,7 @@ except ImportError: pass -def is_localhost_or_internal_url(url: Optional[str]) -> bool: +def is_localhost_or_internal_url(url: str | None) -> bool: """ Check if a URL is a localhost or internal URL. @@ -48,6 +48,29 @@ def is_localhost_or_internal_url(url: Optional[str]) -> bool: return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS) +def get_agent_card_url(agent_card: "AgentCard") -> str | None: + """Return the agent endpoint URL from the resolved SDK card.""" + url = getattr(agent_card, "url", None) + if url: + return url + + interfaces = getattr(agent_card, "supported_interfaces", None) + if interfaces: + return getattr(interfaces[0], "url", None) + return None + + +def set_agent_card_url(agent_card: "AgentCard", url: str) -> None: + """Set the agent endpoint URL on the resolved SDK card.""" + normalized = url.rstrip("/") + "/" + if hasattr(agent_card, "url"): + agent_card.url = normalized + + interfaces = getattr(agent_card, "supported_interfaces", None) + if interfaces: + interfaces[0].url = normalized + + def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard": """ Fix the agent card URL if it contains a localhost/internal address. @@ -70,6 +93,12 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard": fixed_url = base_url.rstrip("/") + "/" agent_card.url = fixed_url + interfaces = getattr(agent_card, "supported_interfaces", None) + if interfaces: + interface_url = getattr(interfaces[0], "url", None) + if interface_url and is_localhost_or_internal_url(interface_url): + interfaces[0].url = base_url.rstrip("/") + "/" + return agent_card @@ -84,8 +113,8 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] async def get_agent_card( self, - relative_card_path: Optional[str] = None, - http_kwargs: Optional[Dict[str, Any]] = None, + relative_card_path: str | None = None, + http_kwargs: Dict[str, Any] | None = None, ) -> "AgentCard": """ Fetch the agent card, trying multiple well-known paths. diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py index 99706e15cee..89b831351ab 100644 --- a/litellm/a2a_protocol/exception_mapping_utils.py +++ b/litellm/a2a_protocol/exception_mapping_utils.py @@ -8,8 +8,8 @@ from typing import TYPE_CHECKING, Any, Optional from litellm._logging import verbose_logger from litellm.a2a_protocol.card_resolver import ( - fix_agent_card_url, is_localhost_or_internal_url, + set_agent_card_url, ) from litellm.a2a_protocol.exceptions import ( A2AAgentCardError, @@ -20,17 +20,18 @@ from litellm.a2a_protocol.exceptions import ( from litellm.constants import CONNECTION_ERROR_PATTERNS if TYPE_CHECKING: - from a2a.client import A2AClient as A2AClientType + from a2a.client import Client as A2AClientType -# Runtime import -A2A_SDK_AVAILABLE = False try: - from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef] + from a2a.client import Client, ClientConfig, create_client A2A_SDK_AVAILABLE = True except ImportError: - _A2AClient = None # type: ignore[assignment, misc] + A2A_SDK_AVAILABLE = False + Client = None # type: ignore[misc, assignment] + ClientConfig = None # type: ignore[misc, assignment] + create_client = None # type: ignore[misc, assignment] class A2AExceptionCheckers: @@ -156,7 +157,7 @@ def map_a2a_exception( ) -def handle_a2a_localhost_retry( +async def handle_a2a_localhost_retry( error: A2ALocalhostURLError, agent_card: Any, a2a_client: "A2AClientType", @@ -180,8 +181,14 @@ def handle_a2a_localhost_retry( Raises: ImportError: If the A2A SDK is not installed """ - if not A2A_SDK_AVAILABLE or _A2AClient is None: - raise ImportError("A2A SDK is required for localhost retry handling. Install it with: pip install a2a") + if not A2A_SDK_AVAILABLE: + raise ImportError("A2A SDK is required for localhost retry handling. Install it with: pip install a2a-sdk") + + if agent_card is None: + raise RuntimeError( + "Cannot retry A2A localhost URL fix: no agent card is available to " + "rewrite, so the upstream URL cannot be corrected." + ) request_type = "streaming " if is_streaming else "" verbose_logger.warning( @@ -191,10 +198,25 @@ def handle_a2a_localhost_retry( ) # Fix the agent card URL - fix_agent_card_url(agent_card, error.base_url) + set_agent_card_url(agent_card, error.base_url) - # Create a new client with the fixed agent card (transport caches URL) - return _A2AClient( - httpx_client=a2a_client._transport.httpx_client, # type: ignore[union-attr] - agent_card=agent_card, + # Reuse the httpx client LiteLLM attached at creation. It carries this agent's + # trace-id and auth headers, so a fresh client would drop them. Only clients built + # by ``create_a2a_client`` have it; an externally-supplied client cannot be retried. + httpx_client = getattr(a2a_client, "_litellm_httpx_client", None) + if httpx_client is None: + raise RuntimeError( + "Cannot retry A2A localhost URL fix: the client was not created by " + "create_a2a_client, so no LiteLLM httpx client is attached." + ) + + new_client = await create_client( # pyright: ignore[reportOptionalCall] + agent_card, + client_config=ClientConfig( # pyright: ignore[reportOptionalCall] + httpx_client=httpx_client, + streaming=is_streaming, + ), ) + new_client._litellm_httpx_client = httpx_client # type: ignore[attr-defined] + new_client._litellm_agent_card = agent_card # type: ignore[attr-defined] + return new_client diff --git a/litellm/a2a_protocol/litellm_completion_bridge/README.md b/litellm/a2a_protocol/litellm_completion_bridge/README.md index a809e9bf55e..3359e75f6df 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/README.md +++ b/litellm/a2a_protocol/litellm_completion_bridge/README.md @@ -67,6 +67,8 @@ When an A2A request hits `/a2a/{agent_id}/message/send`, the bridge: 3. Calls `litellm.acompletion(model="langgraph/agent", api_base="http://localhost:2024")` 4. Transforms response → A2A format +The proxy then normalizes the client-facing response to the agent's pinned `protocolVersion` (`0.3` or `1.0`). No extra provider config is required for completion-bridge agents — pin `protocolVersion` only if your client expects a specific wire format. + ## Classes - `A2ACompletionBridgeTransformation` - Static methods for message format conversion diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 6694c5c4af3..37bf7c34f02 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -1,3 +1,8 @@ +# pyright: reportUnknownArgumentType=false +# a2a-sdk (and its protobuf-generated compat conversions) ships no usable types for +# the call surface used here, so SDK calls take Unknown-typed arguments. This module +# is dedicated to the A2A SDK boundary; the rule is off file-wide instead of +# scattering per-line ignores across every SDK call. """ LiteLLM A2A SDK functions. @@ -7,7 +12,16 @@ Provides standalone functions with @client decorator for LiteLLM logging integra import asyncio import datetime import uuid -from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union +from typing import ( + TYPE_CHECKING, + Any, + AsyncIterator, + Coroutine, + Dict, + Optional, + Union, + cast, +) import litellm from litellm._logging import verbose_logger, verbose_proxy_logger @@ -23,23 +37,45 @@ from litellm.types.agents import LiteLLMSendMessageResponse from litellm.utils import client if TYPE_CHECKING: - from a2a.client import A2AClient as A2AClientType - from a2a.types import AgentCard, SendMessageRequest, SendStreamingMessageRequest + from a2a.client import Client as A2AClientType + from a2a.compat.v0_3.types import ( + AgentCard, + Message, + SendMessageRequest, + SendMessageResponse, + SendStreamingMessageRequest, + SendStreamingMessageResponse, + Task, + ) -# Runtime imports with availability check +# Runtime imports — requires a2a-sdk>=1.1.0 A2A_SDK_AVAILABLE = False -A2ACardResolver: Any = None -_A2AClient: Any = None +_a2a_conversions: Any = None try: - from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef] + from a2a.client import Client, ClientConfig, create_client + from a2a.compat.v0_3 import conversions as _a2a_conversions + from a2a.compat.v0_3.types import ( + Message, + SendMessageRequest, + SendMessageResponse, + SendMessageSuccessResponse, + SendStreamingMessageRequest, + SendStreamingMessageResponse, + Task, + ) A2A_SDK_AVAILABLE = True except ImportError: - pass + Client = None # type: ignore[misc, assignment] + ClientConfig = None # type: ignore[misc, assignment] + create_client = None # type: ignore[misc, assignment] # Import our custom card resolver that supports multiple well-known paths -from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver +from litellm.a2a_protocol.card_resolver import ( + LiteLLMA2ACardResolver, + get_agent_card_url, +) from litellm.a2a_protocol.exception_mapping_utils import ( handle_a2a_localhost_retry, map_a2a_exception, @@ -75,7 +111,7 @@ def _set_usage_on_logging_obj( def _set_agent_id_on_logging_obj( kwargs: Dict[str, Any], - agent_id: Optional[str], + agent_id: str | None, ) -> None: """ Set agent_id on litellm_logging_obj for SpendLogs tracking. @@ -102,10 +138,7 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: """ agent_name = "unknown" - # Try to get agent card from our stored attribute first, then fallback to SDK attribute - agent_card = getattr(a2a_client, "_litellm_agent_card", None) - if agent_card is None: - agent_card = getattr(a2a_client, "agent_card", None) + agent_card = _get_a2a_client_agent_card(a2a_client) if agent_card is not None: agent_name = getattr(agent_card, "name", "unknown") or "unknown" @@ -125,12 +158,22 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: return agent_name +def _get_a2a_client_agent_card(a2a_client: Any) -> Optional["AgentCard"]: + agent_card = cast(Optional["AgentCard"], getattr(a2a_client, "_litellm_agent_card", None)) + if agent_card is not None: + return agent_card + agent_card = cast(Optional["AgentCard"], getattr(a2a_client, "agent_card", None)) + if agent_card is not None: + return agent_card + return cast(Optional["AgentCard"], getattr(a2a_client, "_card", None)) + + async def _send_message_via_completion_bridge( request: "SendMessageRequest", custom_llm_provider: str, - api_base: Optional[str], + api_base: str | None, litellm_params: Dict[str, Any], - agent_extra_headers: Optional[Dict[str, str]] = None, + agent_extra_headers: Dict[str, str] | None = None, ) -> LiteLLMSendMessageResponse: """ Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore). @@ -156,39 +199,71 @@ async def _send_message_via_completion_bridge( return LiteLLMSendMessageResponse.from_dict(response_dict, request_id=str(request.id)) +async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse": + """Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response.""" + if _a2a_conversions is None: + raise ImportError( + "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" + ) + + pb_request = _a2a_conversions.to_core_send_message_request(request) + last_event = None + async for event in a2a_client.send_message(pb_request): + last_event = event + if last_event is None: + raise RuntimeError("A2A send_message failed: no response received from agent.") + + stream_compat = _a2a_conversions.to_compat_stream_response( + last_event, + request_id=request.id, + ) + result = stream_compat.result + if not isinstance(result, (Message, Task)): + raise RuntimeError( + "A2A send_message failed: non-streaming message/send expects the " + "agent's final event to be a Message or Task result." + ) + return SendMessageResponse( + root=SendMessageSuccessResponse( + id=request.id, + result=result, + ) + ) + + async def _execute_a2a_send_with_retry( - a2a_client: Any, - request: Any, - agent_card: Any, - card_url: Optional[str], - api_base: Optional[str], - agent_name: Optional[str], -) -> Any: + a2a_client: "A2AClientType", + request: "SendMessageRequest", + agent_card: Optional["AgentCard"], + card_url: str | None, + api_base: str | None, + agent_name: str | None, +) -> "SendMessageResponse": """Send an A2A message with retry logic for localhost URL errors.""" a2a_response = None for _ in range(2): # max 2 attempts: original + 1 retry try: - a2a_response = await a2a_client.send_message(request) + a2a_response = await _send_message(a2a_client, request) break # success, exit retry loop except A2ALocalhostURLError as e: - a2a_client = handle_a2a_localhost_retry( + a2a_client = await handle_a2a_localhost_retry( error=e, agent_card=agent_card, a2a_client=a2a_client, is_streaming=False, ) - card_url = agent_card.url if agent_card else None + card_url = get_agent_card_url(agent_card) if agent_card else None except Exception as e: try: map_a2a_exception(e, card_url, api_base, model=agent_name) except A2ALocalhostURLError as localhost_err: - a2a_client = handle_a2a_localhost_retry( + a2a_client = await handle_a2a_localhost_retry( error=localhost_err, agent_card=agent_card, a2a_client=a2a_client, is_streaming=False, ) - card_url = agent_card.url if agent_card else None + card_url = get_agent_card_url(agent_card) if agent_card else None continue except Exception: raise @@ -197,14 +272,80 @@ async def _execute_a2a_send_with_retry( return a2a_response +async def _stream_messages( + a2a_client: "A2AClientType", request: "SendStreamingMessageRequest" +) -> AsyncIterator["SendStreamingMessageResponse"]: + """Stream message events via a2a-sdk 1.x and yield JSON-RPC chunks.""" + if _a2a_conversions is None: + raise ImportError( + "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" + ) + + pb_request = _a2a_conversions.to_core_send_message_request(request) + async for event in a2a_client.send_message(pb_request): + compat_chunk = _a2a_conversions.to_compat_stream_response( + event, + request_id=request.id, + ) + yield SendStreamingMessageResponse(root=compat_chunk) + + +async def _execute_a2a_stream_with_retry( + a2a_client: "A2AClientType", + request: "SendStreamingMessageRequest", + agent_card: Optional["AgentCard"], + card_url: str | None, + api_base: str | None, + agent_name: str | None, +) -> AsyncIterator["SendStreamingMessageResponse"]: + """Stream an A2A message with retry logic for localhost URL errors.""" + response_started = False + stream_succeeded = False + for _ in range(2): # max 2 attempts: original + 1 retry + try: + async for chunk in _stream_messages(a2a_client, request): + response_started = True + yield chunk + stream_succeeded = True + return + except A2ALocalhostURLError as e: + if response_started: + raise + a2a_client = await handle_a2a_localhost_retry( + error=e, + agent_card=agent_card, + a2a_client=a2a_client, + is_streaming=True, + ) + card_url = get_agent_card_url(agent_card) if agent_card else None + continue + except Exception as e: + if response_started: + raise + try: + map_a2a_exception(e, card_url, api_base, model=agent_name) + except A2ALocalhostURLError as localhost_err: + a2a_client = await handle_a2a_localhost_retry( + error=localhost_err, + agent_card=agent_card, + a2a_client=a2a_client, + is_streaming=True, + ) + card_url = get_agent_card_url(agent_card) if agent_card else None + continue + raise + if not stream_succeeded: + raise RuntimeError("A2A send_message_streaming failed: no response received after retry attempts.") + + @client async def asend_message( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendMessageRequest"] = None, - api_base: Optional[str] = None, - litellm_params: Optional[Dict[str, Any]] = None, - agent_id: Optional[str] = None, - agent_extra_headers: Optional[Dict[str, str]] = None, + api_base: str | None = None, + litellm_params: Dict[str, Any] | None = None, + agent_id: str | None = None, + agent_extra_headers: Dict[str, str] | None = None, **kwargs: Any, ) -> LiteLLMSendMessageResponse: """ @@ -301,8 +442,8 @@ async def asend_message( verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}") # Get agent card URL for localhost retry logic - agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(a2a_client, "agent_card", None) - card_url = getattr(agent_card, "url", None) if agent_card else None + agent_card = _get_a2a_client_agent_card(a2a_client) + card_url = get_agent_card_url(agent_card) if agent_card else None a2a_response = await _execute_a2a_send_with_retry( a2a_client=a2a_client, @@ -375,10 +516,10 @@ def send_message( def _build_streaming_logging_obj( request: "SendStreamingMessageRequest", agent_name: str, - agent_id: Optional[str], - litellm_params: Optional[Dict[str, Any]], - metadata: Optional[Dict[str, Any]], - proxy_server_request: Optional[Dict[str, Any]], + agent_id: str | None, + litellm_params: Dict[str, Any] | None, + metadata: Dict[str, Any] | None, + proxy_server_request: Dict[str, Any] | None, ) -> Logging: """Build logging object for streaming A2A requests.""" start_time = datetime.datetime.now() @@ -417,12 +558,13 @@ def _build_streaming_logging_obj( async def asend_message_streaming( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendStreamingMessageRequest"] = None, - api_base: Optional[str] = None, - litellm_params: Optional[Dict[str, Any]] = None, - agent_id: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - proxy_server_request: Optional[Dict[str, Any]] = None, - agent_extra_headers: Optional[Dict[str, str]] = None, + api_base: str | None = None, + litellm_params: Dict[str, Any] | None = None, + agent_id: str | None = None, + metadata: Dict[str, Any] | None = None, + proxy_server_request: Dict[str, Any] | None = None, + agent_extra_headers: Dict[str, str] | None = None, + **kwargs: object, ) -> AsyncIterator[Any]: """ Async: Send a streaming message to an A2A agent. @@ -491,99 +633,72 @@ async def asend_message_streaming( yield chunk return - # Standard A2A client flow if request is None: raise ValueError("request is required") - # Create A2A client if not provided but api_base is available + _raw_logging_obj = kwargs.get("litellm_logging_obj") + logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None + if a2a_client is None: if api_base is None: raise ValueError("Either a2a_client or api_base is required for standard A2A flow") - # Mirror the non-streaming path: always include trace and agent-id headers - streaming_extra_headers: Dict[str, str] = { - "X-LiteLLM-Trace-Id": str(request.id), - } + logging_trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None + trace_id = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4())) + extra_headers: dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: - streaming_extra_headers["X-LiteLLM-Agent-Id"] = agent_id + extra_headers["X-LiteLLM-Agent-Id"] = agent_id if agent_extra_headers: - streaming_extra_headers.update(agent_extra_headers) - a2a_client = await create_a2a_client(base_url=api_base, extra_headers=streaming_extra_headers) - - # Type assertion: a2a_client is guaranteed to be non-None here - assert a2a_client is not None - - verbose_logger.info(f"A2A send_message_streaming request_id={request.id}") - - # Build logging object for streaming completion callbacks - agent_card = getattr(a2a_client, "_litellm_agent_card", None) or getattr(a2a_client, "agent_card", None) - card_url = getattr(agent_card, "url", None) if agent_card else None - agent_name = getattr(agent_card, "name", "unknown") if agent_card else "unknown" - - logging_obj = _build_streaming_logging_obj( - request=request, - agent_name=agent_name, - agent_id=agent_id, - litellm_params=litellm_params, - metadata=metadata, - proxy_server_request=proxy_server_request, - ) - - # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL - # Connection errors in streaming typically occur on first chunk iteration - first_chunk = True - for attempt in range(2): # max 2 attempts: original + 1 retry - stream = a2a_client.send_message_streaming(request) - iterator = A2AStreamingIterator( - stream=stream, - request=request, - logging_obj=logging_obj, - agent_name=agent_name, + extra_headers.update(agent_extra_headers) + a2a_client = await create_a2a_client( + base_url=api_base, + extra_headers=extra_headers, + streaming=True, ) - try: - first_chunk = True - async for chunk in iterator: - if first_chunk: - first_chunk = False # connection succeeded - yield chunk - return # stream completed successfully - except A2ALocalhostURLError as e: - # Only retry on first chunk, not mid-stream - if first_chunk and attempt == 0: - a2a_client = handle_a2a_localhost_retry( - error=e, - agent_card=agent_card, - a2a_client=a2a_client, - is_streaming=True, - ) - card_url = agent_card.url if agent_card else None - else: - raise - except Exception as e: - # Only map exception on first chunk - if first_chunk and attempt == 0: - try: - map_a2a_exception(e, card_url, api_base, model=agent_name) - except A2ALocalhostURLError as localhost_err: - # Localhost URL error - fix and retry - a2a_client = handle_a2a_localhost_retry( - error=localhost_err, - agent_card=agent_card, - a2a_client=a2a_client, - is_streaming=True, - ) - card_url = agent_card.url if agent_card else None - continue - except Exception: - # Re-raise the mapped exception - raise - raise + assert a2a_client is not None + + agent_name = _get_a2a_model_info(a2a_client, kwargs) + + if logging_obj is None: + logging_obj = _build_streaming_logging_obj( + request=request, + agent_name=agent_name, + agent_id=agent_id, + litellm_params=litellm_params, + metadata=metadata, + proxy_server_request=proxy_server_request, + ) + + verbose_logger.info(f"A2A send_message_streaming request_id={request.id}, agent={agent_name}") + + agent_card = _get_a2a_client_agent_card(a2a_client) + card_url = get_agent_card_url(agent_card) if agent_card else None + + stream = _execute_a2a_stream_with_retry( + a2a_client=a2a_client, + request=request, + agent_card=agent_card, + card_url=card_url, + api_base=api_base, + agent_name=agent_name, + ) + + _set_agent_id_on_logging_obj(kwargs=kwargs, agent_id=agent_id) + + async for chunk in A2AStreamingIterator( + stream=stream, + request=request, + logging_obj=logging_obj, + agent_name=agent_name, + ): + yield chunk async def create_a2a_client( base_url: str, timeout: float = DEFAULT_A2A_AGENT_TIMEOUT, - extra_headers: Optional[Dict[str, str]] = None, + extra_headers: Dict[str, str] | None = None, + streaming: bool = False, ) -> "A2AClientType": """ Create an A2A client for the given agent URL. @@ -640,23 +755,20 @@ async def create_a2a_client( httpx_client.headers.update(extra_headers) verbose_proxy_logger.debug(f"A2A client created with extra_headers={list(extra_headers.keys())}") - # Resolve agent card - resolver = A2ACardResolver( - httpx_client=httpx_client, - base_url=base_url, + a2a_client = await create_client( # pyright: ignore[reportOptionalCall] + base_url, + client_config=ClientConfig( # pyright: ignore[reportOptionalCall] + httpx_client=httpx_client, + streaming=streaming, + ), ) - agent_card = await resolver.get_agent_card() - - verbose_logger.debug(f"Resolved agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}") - - # Create A2A client - a2a_client = _A2AClient( - httpx_client=httpx_client, - agent_card=agent_card, - ) - - # Store agent_card on client for later retrieval (SDK doesn't expose it) - a2a_client._litellm_agent_card = agent_card # type: ignore[attr-defined] + # Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse + # the configured httpx client (with this agent's trace-id/auth headers) without + # excavating a2a-sdk private internals. + a2a_client._litellm_httpx_client = httpx_client # type: ignore[attr-defined] + agent_card = getattr(a2a_client, "_card", None) + if agent_card is not None: + a2a_client._litellm_agent_card = agent_card # type: ignore[attr-defined] verbose_logger.info(f"A2A client created for {base_url}") @@ -666,7 +778,7 @@ async def create_a2a_client( async def aget_agent_card( base_url: str, timeout: float = DEFAULT_A2A_AGENT_TIMEOUT, - extra_headers: Optional[Dict[str, str]] = None, + extra_headers: Dict[str, str] | None = None, ) -> "AgentCard": """ Fetch the agent card from an A2A agent. diff --git a/litellm/proxy/a2a/agent_card.py b/litellm/proxy/a2a/agent_card.py index 79638341ac1..e97ab4a01ae 100644 --- a/litellm/proxy/a2a/agent_card.py +++ b/litellm/proxy/a2a/agent_card.py @@ -8,11 +8,24 @@ and uses LiteLLM auth. """ from copy import deepcopy -from typing import Any, Dict, List, Mapping, Optional +from typing import Any, Dict, List, Mapping -# Protocol version LiteLLM speaks. Bump when the proxy's A2A surface changes. +# Protocol versions LiteLLM can serve to A2A clients. The admin pins one per agent; +# responses are normalized to it regardless of the upstream agent's own version. +SUPPORTED_A2A_PROTOCOL_VERSIONS = ("0.3", "1.0") + +# Default served version when the agent card does not pin one. LITELLM_A2A_PROTOCOL_VERSION = "1.0" + +def resolve_served_protocol_version(card: Mapping[str, Any] | None) -> str: + """Return the validated protocol version an agent card pins, else the default.""" + version = card.get("protocolVersion") if card else None + if version in SUPPORTED_A2A_PROTOCOL_VERSIONS: + return version + return LITELLM_A2A_PROTOCOL_VERSION + + # Security scheme exposed by the LiteLLM-fronted agent card. Always replaces # whatever upstream advertised — the client must authenticate to the proxy, # not the upstream agent. @@ -106,12 +119,12 @@ def _default_litellm_provider(proxy_base_url: str) -> Dict[str, str]: def merge_agent_card( - upstream_card: Optional[Mapping[str, Any]], + upstream_card: Mapping[str, Any] | None, *, proxy_url: str, proxy_base_url: str, - name: Optional[str] = None, - description: Optional[str] = None, + name: str | None = None, + description: str | None = None, ) -> Dict[str, Any]: """ Build the LiteLLM-fronted agent card. @@ -139,7 +152,8 @@ def merge_agent_card( # proxy requests. The public well-known endpoint rewrites this field # to the proxy URL before exposing the card to clients. - base["protocolVersion"] = LITELLM_A2A_PROTOCOL_VERSION + served_version = resolve_served_protocol_version(upstream_card) + base["protocolVersion"] = served_version if name: base["name"] = name @@ -165,7 +179,7 @@ def merge_agent_card( { "url": proxy_url, "protocolBinding": "JSONRPC", - "protocolVersion": LITELLM_A2A_PROTOCOL_VERSION, + "protocolVersion": served_version, } ] diff --git a/litellm/proxy/a2a/version_convert.py b/litellm/proxy/a2a/version_convert.py new file mode 100644 index 00000000000..e8f49e6f6a9 --- /dev/null +++ b/litellm/proxy/a2a/version_convert.py @@ -0,0 +1,449 @@ +# pyright: reportUnknownArgumentType=false +# a2a-sdk's compat conversions (pb2_v10, ParseDict, MessageToDict, to_compat_*) +# are protobuf-generated/untyped, so every conversion call here takes Unknown-typed +# arguments. This module is the A2A 0.3<->1.0 boundary; the rule is off file-wide +# rather than scattering per-line ignores across every SDK call. +""" +Normalize A2A JSON-RPC payloads to the protocol version LiteLLM serves for an agent. + +LiteLLM fronts upstream agents and lets an admin pin the protocol version it speaks +to clients (``0.3`` or ``1.0``) per agent. Upstream responses may arrive in either +wire shape, so every response, stream event, forwarded request and extended card is +converted to the served version here. Conversion is shape-detecting (we infer the +payload's current version rather than trusting a stored one) and best-effort: any +failure falls back to returning the input unchanged so a conversion bug can never +break an otherwise-valid response. + +The two wire shapes: + +- ``0.3``: JSON dump of the compat pydantic types, discriminated by a ``kind`` field + (``message`` / ``task`` / ``status-update`` / ``artifact-update``). A send result is + the bare object. +- ``1.0``: protobuf JSON (``MessageToDict``), a oneof envelope keyed by + ``message`` / ``task`` / ``statusUpdate`` / ``artifactUpdate`` with no ``kind``. A + ``Task`` result is a bare object without ``kind``. +""" + +from types import ModuleType +from typing import Callable, Literal, Union + +from pydantic import BaseModel + +from litellm._logging import verbose_proxy_logger + +A2AVersion = Literal["0.3", "1.0"] +RequestId = Union[str, int, None] +JsonDict = dict[str, object] + +_V1_SEND_ENVELOPE_KEYS = frozenset({"message", "task"}) +_V1_STREAM_ENVELOPE_KEYS = frozenset({"message", "task", "statusUpdate", "artifactUpdate"}) + + +def _dump_03(model: BaseModel) -> JsonDict: + """Dump a compat (0.3) pydantic model to its camelCase wire dict.""" + return model.model_dump(by_alias=True, exclude_none=True, mode="json") + + +def _best_effort(convert: Callable[[], JsonDict], fallback: JsonDict, *, label: str) -> JsonDict: + """Run a conversion, returning ``fallback`` unchanged if it raises.""" + try: + return convert() + except Exception as e: # noqa: BLE001 - best-effort passthrough + verbose_proxy_logger.debug("A2A %s conversion failed: %s", label, e) + return fallback + + +def normalize_jsonrpc_response(content: JsonDict, target: A2AVersion, *, method: str) -> JsonDict: + """Convert a JSON-RPC response's ``result`` to ``target``. + + Errors and non-dict results pass through untouched. + """ + if content.get("error") is not None: + return content + result = content.get("result") + if not isinstance(result, dict): + return content + + converted = _convert_result(result, target, method=method, request_id=_as_request_id(content.get("id"))) + if converted is result: + return content + return {**content, "result": converted} + + +def normalize_stream_event(event: JsonDict, target: A2AVersion, *, request_id: RequestId) -> JsonDict: + """Convert a single streamed JSON-RPC event's ``result`` to ``target``.""" + if event.get("error") is not None: + return event + result = event.get("result") + if not isinstance(result, dict): + return event + + converted = _convert_stream_result(result, target, request_id=request_id) + if converted is result: + return event + return {**event, "result": converted} + + +def normalize_request_params(params: JsonDict, served: A2AVersion, *, method: str) -> JsonDict: + """Down-convert forwarded request ``params`` from the served version to 0.3. + + Upstream agents in this proxy pivot on 0.3 wire format, so when LiteLLM serves + 1.0 the inbound params must be lowered before forwarding. A no-op when the served + version is already 0.3. + """ + if served == "0.3": + return params + return _best_effort( + lambda: _lower_request_params(params, method=method), + params, + label=f"request params ({method})", + ) + + +def _detect_card_version(card: JsonDict) -> A2AVersion: + """Infer the wire version of an agent card dict. + + ``protocolVersion`` is the authoritative indicator; fall back to presence of + ``supportedInterfaces`` (a 1.0-only field) only when the explicit field is absent. + Cards that set ``protocolVersion: "0.3"`` or carry neither signal are treated as 0.3. + """ + pv = card.get("protocolVersion") + if pv == "1.0": + return "1.0" + if pv == "0.3": + return "0.3" + # No protocolVersion field: use structural heuristic. + return "1.0" if "supportedInterfaces" in card else "0.3" + + +def normalize_agent_card(card: JsonDict, target: A2AVersion) -> JsonDict: + """Convert an extended agent card to ``target``. + + When lowering to 0.3, ``additionalInterfaces`` is stripped so the conversion never + re-exposes upstream backend URLs that the LiteLLM-fronting merge deliberately drops. + """ + if not isinstance(card, dict): + return card + + current = _detect_card_version(card) + if current == target and not (target == "0.3" and "supportedInterfaces" in card): + return card + return _best_effort(lambda: _convert_agent_card(card, target), card, label="agent card") + + +def _as_request_id(value: object) -> RequestId: + return value if isinstance(value, (str, int)) else None + + +def _convert_result( + result: JsonDict, + target: A2AVersion, + *, + method: str, + request_id: RequestId, +) -> JsonDict: + if method == "message/send": + return _convert_send_result(result, target, request_id=request_id) + if method in ("tasks/get", "tasks/cancel"): + return _convert_task(result, target) + if method == "tasks/list": + return _convert_list_tasks_result(result, target) + return result + + +def _detect_send_version(result: JsonDict) -> A2AVersion | None: + if "kind" in result: + return "0.3" + if result.keys() & _V1_SEND_ENVELOPE_KEYS: + return "1.0" + return None + + +def _convert_send_result(result: JsonDict, target: A2AVersion, *, request_id: RequestId) -> JsonDict: + current = _detect_send_version(result) + if current is None or current == target: + return result + return _best_effort( + lambda: _send_result_to(result, target, request_id), + result, + label="send result", + ) + + +def _send_result_to(result: JsonDict, target: A2AVersion, request_id: RequestId) -> JsonDict: + from a2a.compat.v0_3.conversions import ( + MessageToDict, + ParseDict, + pb2_v10, + to_compat_send_message_response, + to_core_send_message_response, + types_v03, + ) + + if target == "1.0": + compat_result = _validate_message_or_task(result, types_v03) + response = types_v03.SendMessageResponse( + root=types_v03.SendMessageSuccessResponse( + id=str(request_id) if request_id is not None else "", + result=compat_result, # pyright: ignore[reportArgumentType] + ) + ) + return MessageToDict( + to_core_send_message_response(response), + preserving_proto_field_name=False, + ) + + pb = pb2_v10.SendMessageResponse() + ParseDict(result, pb, ignore_unknown_fields=True) + return _dump_03(to_compat_send_message_response(pb, request_id).root.result) + + +def _convert_task(result: JsonDict, target: A2AVersion) -> JsonDict: + current: A2AVersion = "0.3" if "kind" in result else "1.0" + if current == target: + return result + return _best_effort(lambda: _task_to(result, target), result, label="task") + + +def _detect_list_tasks_version(result: JsonDict) -> A2AVersion | None: + tasks = result.get("tasks") + if not isinstance(tasks, list) or not tasks: + return None + first = tasks[0] + if not isinstance(first, dict): + return None + return "0.3" if "kind" in first else "1.0" + + +def _convert_list_tasks_result(result: JsonDict, target: A2AVersion) -> JsonDict: + current = _detect_list_tasks_version(result) + if current is None or current == target: + return result + return _best_effort( + lambda: _list_tasks_result_to(result, target), + result, + label="list tasks result", + ) + + +def _list_tasks_result_to(result: JsonDict, target: A2AVersion) -> JsonDict: + tasks = result.get("tasks") + if not isinstance(tasks, list): + return result + return { + **result, + "tasks": [_task_to(item, target) if isinstance(item, dict) else item for item in tasks], + } + + +def _task_to(result: JsonDict, target: A2AVersion) -> JsonDict: + from a2a.compat.v0_3.conversions import ( + MessageToDict, + ParseDict, + pb2_v10, + to_compat_task, + to_core_task, + types_v03, + ) + + if target == "1.0": + core = to_core_task(types_v03.Task.model_validate(result)) + return MessageToDict(core, preserving_proto_field_name=False) + + pb = pb2_v10.Task() + ParseDict(result, pb, ignore_unknown_fields=True) + return _dump_03(to_compat_task(pb)) + + +def _detect_stream_version(result: JsonDict) -> A2AVersion | None: + if "kind" in result: + return "0.3" + if result.keys() & _V1_STREAM_ENVELOPE_KEYS: + return "1.0" + return None + + +def _convert_stream_result(result: JsonDict, target: A2AVersion, *, request_id: RequestId) -> JsonDict: + current = _detect_stream_version(result) + if current is None or current == target: + return result + return _best_effort( + lambda: _stream_result_to(result, target, request_id), + result, + label="stream event", + ) + + +def _stream_result_to(result: JsonDict, target: A2AVersion, request_id: RequestId) -> JsonDict: + from a2a.compat.v0_3.conversions import ( + MessageToDict, + ParseDict, + pb2_v10, + to_compat_stream_response, + to_core_stream_response, + types_v03, + ) + + if target == "1.0": + event = _validate_stream_event(result, types_v03) + wrapper = types_v03.SendStreamingMessageSuccessResponse( + id=str(request_id) if request_id is not None else "", + result=event, # pyright: ignore[reportArgumentType] + ) + return MessageToDict(to_core_stream_response(wrapper), preserving_proto_field_name=False) + + pb = pb2_v10.StreamResponse() + ParseDict(result, pb, ignore_unknown_fields=True) + return _dump_03(to_compat_stream_response(pb, request_id).result) + + +def _convert_agent_card(card: JsonDict, target: A2AVersion) -> JsonDict: + from a2a.compat.v0_3.conversions import ( + MessageToDict, + ParseDict, + pb2_v10, + to_compat_agent_card, + to_core_agent_card, + types_v03, + ) + + if target == "0.3": + pb = pb2_v10.AgentCard() + ParseDict(card, pb, ignore_unknown_fields=True) + lowered = _dump_03(to_compat_agent_card(pb)) + lowered.pop("additionalInterfaces", None) + return lowered + + core = to_core_agent_card(types_v03.AgentCard.model_validate(card)) + return MessageToDict(core, preserving_proto_field_name=False) + + +def _validate_message_or_task(result: JsonDict, types_v03: ModuleType) -> BaseModel: + if result.get("kind") == "task": + return types_v03.Task.model_validate(result) + return types_v03.Message.model_validate(result) + + +def _validate_stream_event(result: JsonDict, types_v03: ModuleType) -> BaseModel: + kind = result.get("kind") + if kind == "task": + return types_v03.Task.model_validate(result) + if kind == "status-update": + return types_v03.TaskStatusUpdateEvent.model_validate(result) + if kind == "artifact-update": + return types_v03.TaskArtifactUpdateEvent.model_validate(result) + return types_v03.Message.model_validate(result) + + +def _lower_request_params(params: JsonDict, *, method: str) -> JsonDict: + if method == "tasks/list": + return _lower_list_tasks_params(params) + + from a2a.compat.v0_3.conversions import ( + ParseDict, + pb2_v10, + to_compat_cancel_task_request, + to_compat_create_task_push_notification_config_request, + to_compat_delete_task_push_notification_config_request, + to_compat_get_task_push_notification_config_request, + to_compat_get_task_request, + to_compat_list_task_push_notification_config_request, + to_compat_subscribe_to_task_request, + ) + + lowerings: dict[str, Callable[[JsonDict], BaseModel]] = { + "tasks/get": lambda p: to_compat_get_task_request(_parse(ParseDict, p, pb2_v10.GetTaskRequest()), "").params, + "tasks/cancel": lambda p: ( + to_compat_cancel_task_request(_parse(ParseDict, p, pb2_v10.CancelTaskRequest()), "").params + ), + "tasks/resubscribe": lambda p: ( + to_compat_subscribe_to_task_request(_parse(ParseDict, p, pb2_v10.SubscribeToTaskRequest()), "").params + ), + "tasks/pushNotificationConfig/set": lambda p: ( + to_compat_create_task_push_notification_config_request( + _parse( + ParseDict, + _flatten_create_push_notification_params(p), + pb2_v10.TaskPushNotificationConfig(), + ), + "", + ).params + ), + "tasks/pushNotificationConfig/get": lambda p: ( + to_compat_get_task_push_notification_config_request( + _parse(ParseDict, p, pb2_v10.GetTaskPushNotificationConfigRequest()), "" + ).params + ), + "tasks/pushNotificationConfig/list": lambda p: ( + to_compat_list_task_push_notification_config_request( + _parse(ParseDict, p, pb2_v10.ListTaskPushNotificationConfigsRequest()), "" + ).params + ), + "tasks/pushNotificationConfig/delete": lambda p: ( + to_compat_delete_task_push_notification_config_request( + _parse(ParseDict, p, pb2_v10.DeleteTaskPushNotificationConfigRequest()), "" + ).params + ), + } + lower = lowerings.get(method) + if lower is None: + return params + return _dump_03(lower(params)) + + +def _lower_list_tasks_params(params: JsonDict) -> JsonDict: + from a2a.compat.v0_3.conversions import ( + MessageToDict, + ParseDict, + pb2_v10, + types_v03, + ) + + proto = pb2_v10.ListTasksRequest() + _parse(ParseDict, params, proto) + lowered = MessageToDict(proto, preserving_proto_field_name=False) + status_name = str(pb2_v10.TaskState.Name(proto.status)) + valid_0_3_values = frozenset(str(member.value) for member in types_v03.TaskState) + compat_status = _proto_task_state_name_to_0_3(status_name, valid_0_3_values) + if compat_status is None: + lowered.pop("status", None) + else: + lowered["status"] = compat_status + return lowered + + +def _proto_task_state_name_to_0_3(name: str, valid_0_3_values: frozenset[str]) -> str | None: + """Map a 1.0 protobuf ``TaskState`` enum name to its 0.3 wire string. + + The ``TASK_STATE_`` enum names line up with the 0.3 wire values once the + prefix is dropped and underscores become dashes, so no private SDK mapping is + needed. The result is validated against the 0.3 enum's own values; an unspecified + or unrecognized state yields ``None`` so the status filter is dropped. + """ + base = name.removeprefix("TASK_STATE_") + if base == "UNSPECIFIED": + return None + candidate = base.lower().replace("_", "-") + return candidate if candidate in valid_0_3_values else None + + +def _flatten_create_push_notification_params(params: JsonDict) -> JsonDict: + """Merge 1.x create envelope fields (parent/configId/config) into flat pb fields.""" + flat = dict(params) + config = flat.pop("config", None) + push_config = flat.pop("pushNotificationConfig", None) + nested = config if config is not None else push_config + if not isinstance(nested, dict): + return params + parent = flat.pop("parent", None) + if isinstance(parent, str) and parent.startswith("tasks/") and "taskId" not in flat: + flat["taskId"] = parent.removeprefix("tasks/").split("/")[0] + if (config_id := flat.pop("configId", None)) and "id" not in nested: + nested["id"] = config_id + flat.update(nested) + return flat + + +def _parse(parse_dict: Callable[..., object], data: JsonDict, message: object) -> object: + parse_dict(data, message, ignore_unknown_fields=True) + return message diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 78be6231ae0..76423fbebde 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -1,3 +1,8 @@ +# pyright: reportUnknownArgumentType=false +# This module forwards JSON-RPC payloads through the untyped a2a-sdk compat +# conversions (pb2_v10/ParseDict/MessageToDict/to_compat_*), so SDK and decoded-JSON +# values flow in as Unknown. The rule is off file-wide rather than scattering per-line +# ignores across every SDK and JSON-RPC call. """ A2A Protocol endpoints for LiteLLM Proxy. @@ -6,26 +11,43 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json -from typing import Any, AsyncGenerator, Dict, List, Optional +from copy import deepcopy +from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, List from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, StreamingResponse +from pydantic import ValidationError from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.a2a.version_convert import ( + A2AVersion, + normalize_agent_card, + normalize_jsonrpc_response, + normalize_request_params, + normalize_stream_event, +) from litellm.proxy.agent_endpoints.databricks_oauth import ( DATABRICKS_OAUTH_PARAM, resolve_databricks_app_auth_header, ) from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.utils import get_custom_url from litellm.types.utils import all_litellm_params +if TYPE_CHECKING: + from a2a.compat.v0_3.types import MessageSendParams + + from litellm.types.agents import AgentResponse + router = APIRouter() _PASCAL_TO_WIRE: Dict[str, str] = { + "SendMessage": "message/send", + "SendStreamingMessage": "message/stream", "GetTask": "tasks/get", "ListTasks": "tasks/list", "CancelTask": "tasks/cancel", @@ -38,6 +60,39 @@ _PASCAL_TO_WIRE: Dict[str, str] = { } +def _build_message_send_params(params: dict[str, Any]) -> "MessageSendParams": + """Build MessageSendParams from wire (0.3) or A2A 1.0 JSON-RPC params.""" + from a2a.compat.v0_3.types import MessageSendParams + + try: + return MessageSendParams(**params) + except ValidationError: + from a2a.compat.v0_3.conversions import pb2_v10, to_compat_send_message_request + from google.protobuf.json_format import ParseDict, ParseError + + pb = pb2_v10.SendMessageRequest() + try: + ParseDict(params, pb, ignore_unknown_fields=True) + except ParseError as e: + raise ValueError(f"Invalid message/send params: {e}") from e + return to_compat_send_message_request(pb, "").params + + +def _served_version(agent: "AgentResponse", request: Request, original_method: str | None = None) -> A2AVersion: + """Protocol version LiteLLM serves for this agent. + + The agent's configured version governs. For agents that pin no version, fall back + to the client's signal: PascalCase JSON-RPC methods and an ``a2a-version: 1.x`` + header both mark a 1.0 caller; otherwise default to 0.3. + """ + configured = (agent.agent_card_params or {}).get("protocolVersion") + if configured in ("0.3", "1.0"): + return configured + if original_method in _PASCAL_TO_WIRE: + return "1.0" + return "1.0" if request.headers.get("a2a-version", "").startswith("1.") else "0.3" + + def _validate_push_notification_url(url: str) -> None: parsed = urlparse(url) if parsed.scheme != "https": @@ -62,9 +117,9 @@ def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Dict[str, str def _forwarding_headers( user_api_key_dict: UserAPIKeyAuth, - request_data: dict, - agent_extra_headers: Optional[Dict[str, str]], -) -> Optional[Dict[str, str]]: + request_data: dict[str, Any], + agent_extra_headers: Dict[str, str] | None, +) -> Dict[str, str] | None: sanitized = ( {k: v for k, v in agent_extra_headers.items() if not k.lower().startswith("x-litellm-")} if agent_extra_headers @@ -80,7 +135,7 @@ def _forwarding_headers( def _jsonrpc_error( - request_id: Optional[Any], + request_id: Any | None, code: int, message: str, status_code: int = 400, @@ -125,9 +180,9 @@ def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: async def _forward_jsonrpc( agent_url: str, - body: dict, - extra_headers: Optional[Dict[str, str]] = None, -) -> dict: + body: dict[str, Any], + extra_headers: Dict[str, str] | None = None, +) -> dict[str, Any]: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -149,9 +204,10 @@ async def _forward_jsonrpc( async def _a2a_sse_event_source( agent_url: str, - body: dict, - request_id: Optional[Any] = None, - extra_headers: Optional[Dict[str, str]] = None, + body: dict[str, Any], + request_id: Any | None = None, + extra_headers: Dict[str, str] | None = None, + served_version: A2AVersion = "0.3", ) -> AsyncGenerator[dict, None]: """Stream an upstream A2A SSE response as parsed JSON-RPC event dicts. @@ -177,7 +233,7 @@ async def _a2a_sse_event_source( try: if not resp.is_success: error_body = await resp.aread() - error_event: Optional[dict] = None + error_event: dict[str, Any] | None = None try: parsed = json.loads(error_body) if isinstance(parsed, dict) and "error" in parsed: @@ -198,23 +254,33 @@ async def _a2a_sse_event_source( if not payload: continue try: - yield json.loads(payload) + event = json.loads(payload) except Exception: continue + if isinstance(event, dict): + event = normalize_stream_event(event, served_version, request_id=request_id) + yield event finally: await resp.aclose() async def _forward_jsonrpc_sse( agent_url: str, - body: dict, - request_id: Optional[Any] = None, - extra_headers: Optional[Dict[str, str]] = None, - proxy_logging_obj: Optional[Any] = None, - user_api_key_dict: Optional[Any] = None, - request_data: Optional[dict] = None, + body: dict[str, Any], + request_id: Any | None = None, + extra_headers: Dict[str, str] | None = None, + proxy_logging_obj: Any | None = None, + user_api_key_dict: Any | None = None, + request_data: dict[str, Any] | None = None, + served_version: A2AVersion = "0.3", ) -> StreamingResponse: - event_source = _a2a_sse_event_source(agent_url, body, request_id=request_id, extra_headers=extra_headers) + event_source = _a2a_sse_event_source( + agent_url, + body, + request_id=request_id, + extra_headers=extra_headers, + served_version=served_version, + ) def _serialize_chunk(chunk: Any) -> str: return f"data: {json.dumps(chunk)}\n\n" @@ -263,18 +329,19 @@ async def _forward_jsonrpc_sse( async def _handle_stream_message( - api_base: Optional[str], + api_base: str | None, request_id: Any, - params: dict, - litellm_params: Optional[dict] = None, - agent_id: Optional[str] = None, - metadata: Optional[dict] = None, - proxy_server_request: Optional[dict] = None, + params: dict[str, Any], + litellm_params: dict[str, Any] | None = None, + agent_id: str | None = None, + metadata: dict[str, Any] | None = None, + proxy_server_request: dict[str, Any] | None = None, *, - agent_extra_headers: Optional[Dict[str, str]] = None, - user_api_key_dict: Optional[UserAPIKeyAuth] = None, - request_data: Optional[dict] = None, - proxy_logging_obj: Optional[Any] = None, + agent_extra_headers: Dict[str, str] | None = None, + user_api_key_dict: UserAPIKeyAuth | None = None, + request_data: dict[str, Any] | None = None, + proxy_logging_obj: Any | None = None, + served_version: A2AVersion = "0.3", ) -> StreamingResponse: """Handle message/stream method via SDK functions. @@ -304,15 +371,34 @@ async def _handle_stream_message( return StreamingResponse(_error_stream(), media_type="application/x-ndjson") - from a2a.types import MessageSendParams, SendStreamingMessageRequest + from a2a.compat.v0_3.types import SendStreamingMessageRequest use_proxy_hooks = user_api_key_dict is not None and request_data is not None and proxy_logging_obj is not None + try: + message_send_params = _build_message_send_params(params) + except (ValidationError, ValueError) as e: + invalid_params_message = f"Invalid params: {e}" + + async def _invalid_params_stream(): + yield ( + json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32602, "message": invalid_params_message}, + } + ) + + "\n" + ) + + return StreamingResponse(_invalid_params_stream(), media_type="application/x-ndjson") + async def stream_response(): try: a2a_request = SendStreamingMessageRequest( id=request_id, - params=MessageSendParams(**params), + params=message_send_params, ) a2a_stream = asend_message_streaming( request=a2a_request, @@ -339,6 +425,8 @@ async def _handle_stream_message( obj = chunk.model_dump(mode="json", exclude_none=True) else: obj = chunk + if isinstance(obj, dict): + obj = normalize_stream_event(obj, served_version, request_id=request_id) return json.dumps(obj) + "\n" def _ndjson_error(proxy_exc: Any) -> str: @@ -372,9 +460,12 @@ async def _handle_stream_message( else: async for chunk in a2a_stream: if hasattr(chunk, "model_dump"): - yield (json.dumps(chunk.model_dump(mode="json", exclude_none=True)) + "\n") + obj = chunk.model_dump(mode="json", exclude_none=True) else: - yield json.dumps(chunk) + "\n" + obj = chunk + if isinstance(obj, dict): + obj = normalize_stream_event(obj, served_version, request_id=request_id) + yield json.dumps(obj) + "\n" except Exception as e: verbose_proxy_logger.exception(f"Error streaming A2A response: {e}") if ( @@ -460,13 +551,16 @@ async def get_agent_card( detail=f"Agent '{agent_id}' has no agent card configured", ) - # Copy and rewrite URL to point to LiteLLM proxy - agent_card = { - **agent.agent_card_params, - "url": f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}", - } + proxy_url = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}") + agent_card = deepcopy(agent.agent_card_params) + agent_card["url"] = proxy_url + interfaces = agent_card.get("supportedInterfaces") + if isinstance(interfaces, list) and interfaces: + interfaces[0]["url"] = proxy_url + served_version = _served_version(agent, request) + agent_card = normalize_agent_card(agent_card, served_version) - verbose_proxy_logger.debug(f"Returning agent card for '{agent_id}' with proxy URL: {agent_card['url']}") + verbose_proxy_logger.debug(f"Returning agent card for '{agent_id}' with proxy URL: {proxy_url}") return JSONResponse(content=agent_card) except HTTPException: @@ -526,8 +620,9 @@ async def invoke_agent_a2a( if body.get("jsonrpc") != "2.0": return _jsonrpc_error(body.get("id"), -32600, "Invalid Request: jsonrpc must be '2.0'") - request_id: Optional[Any] = body.get("id") - method: Optional[str] = body.get("method") + request_id: Any | None = body.get("id") + original_method: str | None = body.get("method") + method: str | None = original_method params = body.get("params", {}) if method: @@ -553,6 +648,8 @@ async def invoke_agent_a2a( if agent is None: return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404) + served_version = _served_version(agent, request, original_method) + is_allowed = await AgentRequestHandler.is_agent_allowed( agent_id=agent.agent_id, user_api_key_auth=user_api_key_dict, @@ -691,11 +788,16 @@ async def invoke_agent_a2a( "Server error: 'a2a' package not installed. Please install 'a2a-sdk'.", 500, ) - from a2a.types import MessageSendParams, SendMessageRequest + from a2a.compat.v0_3.types import SendMessageRequest + + try: + message_send_params = _build_message_send_params(params) + except (ValidationError, ValueError) as e: + return _jsonrpc_error(request_id, -32602, f"Invalid params: {e}") a2a_request = SendMessageRequest( id=request_id if request_id is not None else "", - params=MessageSendParams(**params), + params=message_send_params, ) # Defer spend-log until after post_call_success_hook so guardrail # results written by the unified_guardrail hook are captured. @@ -723,11 +825,18 @@ async def invoke_agent_a2a( logging_obj._enqueue_deferred_logging = None # type: ignore[union-attr] _enqueue_fn() + response_dict: Dict[str, Any] = ( + response.model_dump(mode="json", exclude_none=True) # type: ignore + if hasattr(response, "model_dump") + else response + if isinstance(response, dict) + else {} + ) return JSONResponse( - content=( - response.model_dump(mode="json", exclude_none=True) # type: ignore - if hasattr(response, "model_dump") - else response + content=normalize_jsonrpc_response( + response_dict, + served_version, + method="message/send", ) ) @@ -744,6 +853,7 @@ async def invoke_agent_a2a( user_api_key_dict=user_api_key_dict, request_data=data, proxy_logging_obj=proxy_logging_obj, + served_version=served_version, ) elif method in { "tasks/get", @@ -757,6 +867,8 @@ async def invoke_agent_a2a( }: if not agent_url: return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500) + if isinstance(params, dict): + params = normalize_request_params(params, served_version, method=method) if method == "tasks/pushNotificationConfig/set": if not isinstance(params, dict): raise HTTPException( @@ -791,8 +903,20 @@ async def invoke_agent_a2a( ) result = await _forward_jsonrpc(agent_url, forward_body, extra_headers=caller_headers) if method == "agent/getAuthenticatedExtendedCard": - if isinstance(result.get("result"), dict) and "url" in result["result"]: - result["result"]["url"] = f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}" + if isinstance(result.get("result"), dict): + card = result["result"] + proxy_url = get_custom_url(str(request.base_url), route=f"a2a/{agent_id}") + # Rewrite the upstream agent URL in both 0.3 (top-level `url`) + # and 1.0 (`supportedInterfaces[0].url`) wire formats so that + # downstream clients never see the upstream internal address. + if "url" in card: + card["url"] = proxy_url + interfaces = card.get("supportedInterfaces") + if isinstance(interfaces, list) and interfaces: + interfaces[0]["url"] = proxy_url + result["result"] = normalize_agent_card(card, served_version) + else: + result = normalize_jsonrpc_response(result, served_version, method=method) from litellm.types.agents import LiteLLMSendMessageResponse response = LiteLLMSendMessageResponse.from_dict(result, request_id=request_id) @@ -810,6 +934,8 @@ async def invoke_agent_a2a( elif method == "tasks/resubscribe": if not agent_url: return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500) + if isinstance(params, dict): + params = normalize_request_params(params, served_version, method=method) forward_body = { "jsonrpc": "2.0", "id": request_id, @@ -829,6 +955,7 @@ async def invoke_agent_a2a( proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, request_data=data, + served_version=served_version, ) else: diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 450c6414270..51e451efbb9 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -11,7 +11,7 @@ Follows the A2A Spec. import asyncio import os import uuid -from typing import Any, Dict, List, Mapping, Optional +from typing import Any, Dict, List, Mapping from fastapi import APIRouter, Depends, HTTPException, Query, Request @@ -20,10 +20,14 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.a2a.agent_card import merge_agent_card +from litellm.proxy.a2a.agent_card import ( + SUPPORTED_A2A_PROTOCOL_VERSIONS, + merge_agent_card, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.utils import get_custom_url from litellm.types.agents import ( AgentConfig, AgentKeySummary, @@ -40,19 +44,33 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( def _proxy_base_url(http_request: Request) -> str: - """Return the proxy's base URL as seen by the caller, without trailing slash.""" - return str(http_request.base_url).rstrip("/") + """Return the proxy's public base URL, preferring PROXY_BASE_URL when set.""" + return get_custom_url(str(http_request.base_url), route=None) + + +def _validate_protocol_version(upstream_card: Mapping[str, Any] | None) -> None: + """Reject an agent card pinning an unsupported A2A protocol version.""" + version = upstream_card.get("protocolVersion") if upstream_card else None + if version is not None and version not in SUPPORTED_A2A_PROTOCOL_VERSIONS: + raise HTTPException( + status_code=400, + detail=( + f"Unsupported protocolVersion '{version}'. " + f"Supported versions: {', '.join(SUPPORTED_A2A_PROTOCOL_VERSIONS)}." + ), + ) def _build_merged_agent_card( - upstream_card: Optional[Mapping[str, Any]], + upstream_card: Mapping[str, Any] | None, *, agent_id: str, http_request: Request, - agent_name: Optional[str] = None, + agent_name: str | None = None, ) -> Dict[str, Any]: """Apply the LiteLLM-fronting merge to ``upstream_card`` for ``agent_id``.""" proxy_base = _proxy_base_url(http_request) + _validate_protocol_version(upstream_card) # Prefer a card-supplied ``name`` (the discovery UI exposes an editable # "Name (shown to API clients)" field that flows into # ``agent_card_params.name``) over the internal ``agent_name`` identifier. @@ -382,7 +400,7 @@ async def create_agent( # schemes, default skills) the agent doesn't actually expose. upstream_card = request.get("agent_card_params") agent_to_create: AgentConfig = request - new_agent_id: Optional[str] = None + new_agent_id: str | None = None if upstream_card is not None: # Pre-generate the agent_id so the merged card can reference it # in ``supportedInterfaces`` before the DB row exists. @@ -988,14 +1006,14 @@ async def make_agents_public( response_model=SpendAnalyticsPaginatedResponse, ) async def get_agent_daily_activity( - agent_ids: Optional[str] = None, - start_date: Optional[str] = None, - end_date: Optional[str] = None, - model: Optional[str] = None, - api_key: Optional[str] = None, + agent_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + model: str | None = None, + api_key: str | None = None, page: int = 1, page_size: int = 10, - exclude_agent_ids: Optional[str] = None, + exclude_agent_ids: str | None = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1012,7 +1030,7 @@ async def get_agent_daily_activity( ) agent_ids_list = agent_ids.split(",") if agent_ids else None - exclude_agent_ids_list: Optional[List[str]] = None + exclude_agent_ids_list: List[str] | None = None if exclude_agent_ids: exclude_agent_ids_list = exclude_agent_ids.split(",") if exclude_agent_ids else None diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 91332480d75..76d211d2c8f 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.get_blog_posts import ( from litellm.proxy._types import ( CommonProxyErrors, ) +from litellm.proxy.utils import get_custom_url from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.agents import AgentCard from litellm.types.mcp import MCPPublicServer @@ -213,11 +214,10 @@ async def get_agents(request: Request): if litellm.public_agent_groups is None: return [] - proxy_base = str(request.base_url).rstrip("/") return [ { **(agent.agent_card_params or {}), - "url": f"{proxy_base}/a2a/{agent.agent_id}", + "url": get_custom_url(str(request.base_url), route=f"a2a/{agent.agent_id}"), } for agent in agents if agent.agent_id in litellm.public_agent_groups diff --git a/pyproject.toml b/pyproject.toml index 23f02cb5762..2ad96c4936b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -90,7 +90,7 @@ extra_proxy = [ # Not in PyPI proxy extra. "resend>=2.23.0,<3.0", "redisvl>=0.4.1,<1.0; python_version < '3.14'", - "a2a-sdk>=0.3.24,<1.0", + "a2a-sdk>=1.1.0,<2.0", ] utils = [ # Not in Docker or PyPI proxy extra. @@ -193,7 +193,7 @@ proxy-dev = [ "opentelemetry-exporter-otlp==1.28.0", "opentelemetry-instrumentation-fastapi==0.49b0", "azure-identity==1.25.2", - "a2a-sdk==0.3.24", + "a2a-sdk==1.1.0", ] ci = [ # These are lazily imported at call sites; keep them out of core deps to @@ -241,6 +241,11 @@ build-backend = "uv_build" constraint-dependencies = [ "tornado>=6.5.6", "aiohttp>=3.14.1,<4.0", + "packaging>=24.0", +] +override-dependencies = [ + # a2a-sdk 1.x requires packaging>=24.0; lunary 1.4.x still caps at <24.0. + "packaging>=24.0", ] default-groups = ["dev"] required-version = ">=0.10.9" diff --git a/tests/agent_tests/test_a2a_agent.py b/tests/agent_tests/test_a2a_agent.py index ed7a5ab9823..f5ad9601369 100644 --- a/tests/agent_tests/test_a2a_agent.py +++ b/tests/agent_tests/test_a2a_agent.py @@ -41,21 +41,24 @@ class MockA2AClient: ) async def send_message(self, request): - return MockA2AResponse(text="hello") + from a2a.compat.v0_3.conversions import pb2_v10 - def send_message_streaming(self, request): - async def _stream(): - yield MockA2AStreamingChunk(text="hel", state="in_progress") - yield MockA2AStreamingChunk(text="hello", state="completed") - - return _stream() + for text in ("hel", "hello"): + event = pb2_v10.StreamResponse() + message = event.message + message.message_id = uuid4().hex + message.role = pb2_v10.ROLE_AGENT + message.parts.add().text = text + yield event @pytest.fixture def mock_a2a_client(monkeypatch): import litellm.a2a_protocol.main as a2a_main - async def _fake_create_a2a_client(base_url, timeout=60.0, extra_headers=None): + async def _fake_create_a2a_client( + base_url, timeout=60.0, extra_headers=None, streaming=False + ): return MockA2AClient() monkeypatch.setattr(a2a_main, "create_a2a_client", _fake_create_a2a_client) @@ -64,7 +67,7 @@ def mock_a2a_client(monkeypatch): @pytest.mark.asyncio async def test_a2a_non_streaming(mock_a2a_client): """Test non-streaming A2A request.""" - from a2a.types import MessageSendParams, SendMessageRequest + from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest from litellm.a2a_protocol import asend_message request = SendMessageRequest( @@ -90,7 +93,7 @@ async def test_a2a_non_streaming(mock_a2a_client): @pytest.mark.asyncio async def test_a2a_streaming(mock_a2a_client): """Test streaming A2A request.""" - from a2a.types import MessageSendParams, SendStreamingMessageRequest + from a2a.compat.v0_3.types import MessageSendParams, SendStreamingMessageRequest from litellm.a2a_protocol import asend_message_streaming request = SendStreamingMessageRequest( diff --git a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py new file mode 100644 index 00000000000..44d2803a260 --- /dev/null +++ b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py @@ -0,0 +1,153 @@ +"""Tests for litellm/a2a_protocol/exception_mapping_utils.py.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.a2a_protocol import exception_mapping_utils as emu +from litellm.a2a_protocol.exceptions import A2ALocalhostURLError + + +def _localhost_error() -> A2ALocalhostURLError: + return A2ALocalhostURLError( + localhost_url="http://localhost:10001/", + base_url="https://agent.example", + original_error=ConnectionError("boom"), + ) + + +@pytest.mark.asyncio +async def test_localhost_retry_reuses_stashed_httpx_client(): + """The retry must reuse the httpx client LiteLLM attached at creation (it carries + the agent's trace-id/auth headers), passing it straight into the new ClientConfig. + """ + stashed_httpx_client = object() + a2a_client = MagicMock() + a2a_client._litellm_httpx_client = stashed_httpx_client + new_client = MagicMock() + + captured = {} + + def fake_client_config(*, httpx_client, streaming): + captured["httpx_client"] = httpx_client + captured["streaming"] = streaming + return MagicMock() + + with ( + patch.object(emu, "A2A_SDK_AVAILABLE", True), + patch.object(emu, "set_agent_card_url") as mock_set_url, + patch.object(emu, "ClientConfig", side_effect=fake_client_config), + patch.object( + emu, "create_client", new=AsyncMock(return_value=new_client) + ) as mock_create, + ): + result = await emu.handle_a2a_localhost_retry( + error=_localhost_error(), + agent_card=MagicMock(), + a2a_client=a2a_client, + is_streaming=True, + ) + + assert result is new_client + mock_set_url.assert_called_once() + # The exact stashed client is threaded through, not a freshly built one. + assert captured["httpx_client"] is stashed_httpx_client + assert captured["streaming"] is True + assert new_client._litellm_httpx_client is stashed_httpx_client + assert mock_create.await_count == 1 + + +@pytest.mark.asyncio +async def test_localhost_retry_raises_when_no_stashed_client(): + """An externally-supplied client has no LiteLLM httpx handle; the retry must fail + with a clear error instead of excavating a2a-sdk internals.""" + a2a_client = MagicMock(spec=[]) # no _litellm_httpx_client attribute + + with ( + patch.object(emu, "A2A_SDK_AVAILABLE", True), + patch.object(emu, "set_agent_card_url"), + patch.object(emu, "create_client", new=AsyncMock()) as mock_create, + ): + with pytest.raises(RuntimeError, match="not created by create_a2a_client"): + await emu.handle_a2a_localhost_retry( + error=_localhost_error(), + agent_card=MagicMock(), + a2a_client=a2a_client, + is_streaming=False, + ) + + mock_create.assert_not_called() + + +@pytest.mark.asyncio +async def test_localhost_retry_raises_when_agent_card_is_none(): + """With no agent card to rewrite, the retry must fail with a clear error instead + of calling create_client(None, ...) and surfacing an opaque SDK TypeError.""" + a2a_client = MagicMock() + a2a_client._litellm_httpx_client = MagicMock() + + with ( + patch.object(emu, "A2A_SDK_AVAILABLE", True), + patch.object(emu, "set_agent_card_url") as mock_set_url, + patch.object(emu, "create_client", new=AsyncMock()) as mock_create, + ): + with pytest.raises(RuntimeError, match="no agent card is available"): + await emu.handle_a2a_localhost_retry( + error=_localhost_error(), + agent_card=None, + a2a_client=a2a_client, + is_streaming=False, + ) + + mock_set_url.assert_not_called() + mock_create.assert_not_called() + + +def test_get_a2a_client_agent_card_reads_sdk_private_card(): + from litellm.a2a_protocol.main import _get_a2a_client_agent_card + + sdk_card = SimpleNamespace(name="Test Agent", url="http://localhost:10001/") + a2a_client = SimpleNamespace(_card=sdk_card) + + assert _get_a2a_client_agent_card(a2a_client) is sdk_card + + +@pytest.mark.asyncio +async def test_stream_with_retry_raises_after_localhost_retries_exhausted(): + """Exhausted localhost retries must not return a silent empty stream.""" + from litellm.a2a_protocol.main import _execute_a2a_stream_with_retry + + localhost_err = _localhost_error() + mock_request = MagicMock() + mock_request.id = "req-1" + mock_a2a_client = MagicMock() + + async def _always_fail_stream(a2a_client, request): + raise localhost_err + yield # pragma: no cover - makes this an async generator + + with ( + patch( + "litellm.a2a_protocol.main._stream_messages", + new=_always_fail_stream, + ), + patch( + "litellm.a2a_protocol.main.handle_a2a_localhost_retry", + new=AsyncMock(return_value=mock_a2a_client), + ), + ): + stream = _execute_a2a_stream_with_retry( + a2a_client=mock_a2a_client, + request=mock_request, + agent_card=MagicMock(), + card_url="http://localhost:10001/", + api_base="https://agent.example", + agent_name="test-agent", + ) + with pytest.raises( + RuntimeError, + match="no response received after retry attempts", + ): + async for _chunk in stream: + pytest.fail("expected retry exhaustion to raise before yielding") diff --git a/tests/test_litellm/a2a_protocol/test_card_resolver.py b/tests/test_litellm/a2a_protocol/test_card_resolver.py index 1bdab50860c..053f28c940f 100644 --- a/tests/test_litellm/a2a_protocol/test_card_resolver.py +++ b/tests/test_litellm/a2a_protocol/test_card_resolver.py @@ -4,7 +4,8 @@ Mock tests for LiteLLMA2ACardResolver. Tests that the card resolver tries both old and new well-known paths. """ -from unittest.mock import AsyncMock, MagicMock, patch +from types import SimpleNamespace +from unittest.mock import MagicMock, patch import pytest @@ -12,6 +13,7 @@ from litellm.a2a_protocol.card_resolver import ( LiteLLMA2ACardResolver, fix_agent_card_url, is_localhost_or_internal_url, + set_agent_card_url, ) @@ -88,3 +90,27 @@ def test_fix_agent_card_url_replaces_localhost(): # Verify localhost URL was replaced with base_url assert result.url == "https://my-public-agent.example.com/" + + +def test_set_agent_card_url_updates_top_level_and_supported_interface(): + card = SimpleNamespace( + url="http://localhost:10001/", + supported_interfaces=[SimpleNamespace(url="http://0.0.0.0:10001/")], + ) + + set_agent_card_url(card, "https://my-public-agent.example.com") + + assert card.url == "https://my-public-agent.example.com/" + assert card.supported_interfaces[0].url == "https://my-public-agent.example.com/" + + +def test_fix_agent_card_url_updates_interface_when_top_level_is_localhost(): + card = SimpleNamespace( + url="http://localhost:10001/", + supported_interfaces=[SimpleNamespace(url="http://0.0.0.0:10001/")], + ) + + result = fix_agent_card_url(card, "https://my-public-agent.example.com") + + assert result.url == "https://my-public-agent.example.com/" + assert result.supported_interfaces[0].url == "https://my-public-agent.example.com/" diff --git a/tests/test_litellm/a2a_protocol/test_cost_calculator.py b/tests/test_litellm/a2a_protocol/test_cost_calculator.py index d7bacaf39eb..bf03562f8ae 100644 --- a/tests/test_litellm/a2a_protocol/test_cost_calculator.py +++ b/tests/test_litellm/a2a_protocol/test_cost_calculator.py @@ -3,8 +3,8 @@ Test A2A cost calculator with cost_per_query parameter. """ import asyncio -from typing import Optional -from unittest.mock import AsyncMock, MagicMock +from typing import Any, AsyncIterator, Optional +from unittest.mock import MagicMock, patch import pytest @@ -12,6 +12,102 @@ import litellm from litellm.integrations.custom_logger import CustomLogger +def _make_send_message_request(request_id: str, user_text: str = "Hello"): + from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest + + return SendMessageRequest( + id=request_id, + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": user_text}], + "messageId": "msg-1", + } + ), + ) + + +async def _mock_execute_a2a_send( + a2a_client: Any, + request: Any, + **kwargs: Any, +) -> Any: + mock_response = MagicMock() + mock_response.model_dump = MagicMock( + return_value={ + "id": request.id, + "jsonrpc": "2.0", + "result": {"status": "completed"}, + } + ) + return mock_response + + +async def _mock_execute_a2a_send_with_assistant_reply( + a2a_client: Any, + request: Any, + **kwargs: Any, +) -> Any: + mock_response = MagicMock() + mock_response.model_dump = MagicMock( + return_value={ + "id": request.id, + "jsonrpc": "2.0", + "result": { + "status": {"state": "completed"}, + "message": { + "role": "assistant", + "parts": [ + { + "kind": "text", + "text": "Hello! I am your assistant. How can I help you today?", + } + ], + "messageId": "msg-456", + }, + }, + } + ) + return mock_response + + +def _make_streaming_request(request_id: str): + from a2a.compat.v0_3.types import MessageSendParams, SendStreamingMessageRequest + + return SendStreamingMessageRequest( + id=request_id, + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-1", + } + ), + ) + + +async def _mock_stream_messages(a2a_client: Any, request: Any) -> AsyncIterator[Any]: + from a2a.compat.v0_3.types import ( + Message, + Part, + Role, + SendStreamingMessageResponse, + SendStreamingMessageSuccessResponse, + TextPart, + ) + + msg = Message( + message_id="msg-agent", + role=Role.agent, + parts=[Part(root=TextPart(kind="text", text="hello"))], + kind="message", + ) + for _ in range(2): + yield SendStreamingMessageResponse( + root=SendStreamingMessageSuccessResponse(id=request.id, result=msg) + ) + + class CostLogger(CustomLogger): """Custom logger to capture response_cost.""" @@ -46,27 +142,18 @@ async def test_asend_message_uses_cost_per_query(): mock_client._litellm_agent_card = MagicMock() mock_client._litellm_agent_card.name = "test-agent" - # Mock response with required fields - mock_response = MagicMock() - mock_response.model_dump = MagicMock( - return_value={ - "id": "test-123", - "jsonrpc": "2.0", - "result": {"status": "completed"}, - } - ) - mock_client.send_message = AsyncMock(return_value=mock_response) - - # Mock request - mock_request = MagicMock() - mock_request.id = "test-123" + mock_request = _make_send_message_request("test-123") # Call asend_message with cost_per_query - await asend_message( - a2a_client=mock_client, - request=mock_request, - cost_per_query=0.05, - ) + with patch( + "litellm.a2a_protocol.main._execute_a2a_send_with_retry", + new=_mock_execute_a2a_send, + ): + await asend_message( + a2a_client=mock_client, + request=mock_request, + cost_per_query=0.05, + ) await asyncio.sleep(0.1) @@ -120,49 +207,24 @@ async def test_asend_message_uses_input_output_cost_per_token(): mock_client._litellm_agent_card = MagicMock() mock_client._litellm_agent_card.name = "test-agent" - # Realistic A2A response with message parts - mock_response = MagicMock() - mock_response.model_dump = MagicMock( - return_value={ - "id": "test-123", - "jsonrpc": "2.0", - "result": { - "status": {"state": "completed"}, - "message": { - "role": "assistant", - "parts": [ - { - "kind": "text", - "text": "Hello! I am your assistant. How can I help you today?", - } - ], - "messageId": "msg-456", - }, - }, - } + mock_request = _make_send_message_request( + "test-123", user_text="Hello, what can you do?" ) - mock_client.send_message = AsyncMock(return_value=mock_response) - - # Mock request with message parts - mock_request = MagicMock() - mock_request.id = "test-123" - mock_request.params = MagicMock() - mock_request.params.message = { - "role": "user", - "parts": [{"kind": "text", "text": "Hello, what can you do?"}], - "messageId": "msg-123", - } # Define specific cost per token values input_cost_per_token = 0.00001 # $0.01 per 1000 tokens output_cost_per_token = 0.00002 # $0.02 per 1000 tokens - await asend_message( - a2a_client=mock_client, - request=mock_request, - input_cost_per_token=input_cost_per_token, - output_cost_per_token=output_cost_per_token, - ) + with patch( + "litellm.a2a_protocol.main._execute_a2a_send_with_retry", + new=_mock_execute_a2a_send_with_assistant_reply, + ): + await asend_message( + a2a_client=mock_client, + request=mock_request, + input_cost_per_token=input_cost_per_token, + output_cost_per_token=output_cost_per_token, + ) await asyncio.sleep(0.1) @@ -225,29 +287,20 @@ async def test_asend_message_passes_agent_id_to_callback(): mock_client._litellm_agent_card = MagicMock() mock_client._litellm_agent_card.name = "test-agent" - # Mock response - mock_response = MagicMock() - mock_response.model_dump = MagicMock( - return_value={ - "id": "test-123", - "jsonrpc": "2.0", - "result": {"status": "completed"}, - } - ) - mock_client.send_message = AsyncMock(return_value=mock_response) - - # Mock request - mock_request = MagicMock() - mock_request.id = "test-123" + mock_request = _make_send_message_request("test-123") test_agent_id = "agent-uuid-12345" # Call asend_message with agent_id - await asend_message( - a2a_client=mock_client, - request=mock_request, - agent_id=test_agent_id, - ) + with patch( + "litellm.a2a_protocol.main._execute_a2a_send_with_retry", + new=_mock_execute_a2a_send, + ): + await asend_message( + a2a_client=mock_client, + request=mock_request, + agent_id=test_agent_id, + ) await asyncio.sleep(0.1) @@ -294,21 +347,7 @@ async def test_asend_message_streaming_propagates_metadata(): mock_client._litellm_agent_card = MagicMock() mock_client._litellm_agent_card.name = "test-agent" - # Mock streaming response - async def mock_stream(): - yield MagicMock(model_dump=lambda mode, exclude_none: {"chunk": 1}) - yield MagicMock(model_dump=lambda mode, exclude_none: {"chunk": 2}) - - mock_client.send_message_streaming = MagicMock(return_value=mock_stream()) - - # Mock request - mock_request = MagicMock() - mock_request.id = "test-stream-metadata" - mock_request.params = MagicMock() - mock_request.params.message = { - "role": "user", - "parts": [{"kind": "text", "text": "Hello"}], - } + mock_request = _make_streaming_request("test-stream-metadata") # Metadata from proxy (contains user_api_key, user_id, team_id for SpendLogs) test_metadata = { @@ -319,12 +358,16 @@ async def test_asend_message_streaming_propagates_metadata(): # Consume streaming response with metadata chunks = [] - async for chunk in asend_message_streaming( - a2a_client=mock_client, - request=mock_request, - metadata=test_metadata, + with patch( + "litellm.a2a_protocol.main._stream_messages", + new=_mock_stream_messages, ): - chunks.append(chunk) + async for chunk in asend_message_streaming( + a2a_client=mock_client, + request=mock_request, + metadata=test_metadata, + ): + chunks.append(chunk) await asyncio.sleep(0.2) @@ -352,32 +395,22 @@ async def test_asend_message_streaming_triggers_callbacks(): mock_client._litellm_agent_card = MagicMock() mock_client._litellm_agent_card.name = "test-agent" - # Mock streaming response - async def mock_stream(): - yield MagicMock(model_dump=lambda mode, exclude_none: {"chunk": 1}) - yield MagicMock(model_dump=lambda mode, exclude_none: {"chunk": 2}) - - mock_client.send_message_streaming = MagicMock(return_value=mock_stream()) - - # Mock request - mock_request = MagicMock() - mock_request.id = "test-stream-123" - mock_request.params = MagicMock() - mock_request.params.message = { - "role": "user", - "parts": [{"kind": "text", "text": "Hello"}], - } + mock_request = _make_streaming_request("test-stream-123") test_agent_id = "test-agent-id-streaming" # Consume streaming response chunks = [] - async for chunk in asend_message_streaming( - a2a_client=mock_client, - request=mock_request, - agent_id=test_agent_id, + with patch( + "litellm.a2a_protocol.main._stream_messages", + new=_mock_stream_messages, ): - chunks.append(chunk) + async for chunk in asend_message_streaming( + a2a_client=mock_client, + request=mock_request, + agent_id=test_agent_id, + ): + chunks.append(chunk) await asyncio.sleep(0.2) diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py new file mode 100644 index 00000000000..2a675616245 --- /dev/null +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -0,0 +1,106 @@ +"""Tests for litellm/a2a_protocol/main.py non-streaming send behavior.""" + +import pytest + +pytest.importorskip("a2a.compat.v0_3.conversions") + +from a2a.compat.v0_3 import conversions as _conv +from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest + +from litellm.a2a_protocol.main import _send_message + + +def _request() -> SendMessageRequest: + params = MessageSendParams( + message={ + "messageId": "m1", + "role": "user", + "parts": [{"kind": "text", "text": "hi"}], + } + ) + return SendMessageRequest(id="r1", params=params) + + +def _message_stream_response(): + sr = _conv.pb2_v10.StreamResponse() + sr.message.message_id = "reply-1" + sr.message.role = _conv.pb2_v10.Role.ROLE_AGENT + sr.message.parts.add().text = "hello back" + return sr + + +def _status_update_stream_response(): + sr = _conv.pb2_v10.StreamResponse() + sr.status_update.task_id = "t1" + sr.status_update.context_id = "c1" + return sr + + +class _FakeClient: + def __init__(self, *events): + self._events = events + + async def send_message(self, _pb_request): + for event in self._events: + yield event + + +@pytest.mark.asyncio +async def test_send_message_returns_message_result(): + response = await _send_message(_FakeClient(_message_stream_response()), _request()) + result = response.root.result + assert type(result).__name__ == "Message" + assert response.root.id == "r1" + + +@pytest.mark.asyncio +async def test_send_message_rejects_update_event_final_with_runtime_error(): + with pytest.raises(RuntimeError, match="Message or Task"): + await _send_message(_FakeClient(_status_update_stream_response()), _request()) + + +@pytest.mark.asyncio +async def test_streaming_trace_id_prefers_logging_trace_id(): + """The streaming X-LiteLLM-Trace-Id must use the logging object's trace id (same + as the non-streaming path), not the JSON-RPC request id, so traces correlate.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from a2a.compat.v0_3.types import ( + MessageSendParams, + SendStreamingMessageRequest, + ) + + from litellm.a2a_protocol import main as a2a_main + from litellm.litellm_core_utils.litellm_logging import Logging + + request = SendStreamingMessageRequest( + id="rpc-1", + params=MessageSendParams( + message={ + "messageId": "m1", + "role": "user", + "parts": [{"kind": "text", "text": "hi"}], + } + ), + ) + logging_obj = MagicMock(spec=Logging) + logging_obj.litellm_trace_id = "trace-from-logging" + + captured: dict = {} + + async def _capture(*, base_url, extra_headers=None, streaming=False, **_): + captured["extra_headers"] = extra_headers + raise RuntimeError("stop") + + with patch.object( + a2a_main, "create_a2a_client", new=AsyncMock(side_effect=_capture) + ): + with pytest.raises(RuntimeError, match="stop"): + async for _ in a2a_main.asend_message_streaming( + request=request, + api_base="http://upstream.local", + litellm_logging_obj=logging_obj, + ): + pass + + assert captured["extra_headers"]["X-LiteLLM-Trace-Id"] == "trace-from-logging" diff --git a/tests/test_litellm/proxy/a2a/test_agent_card.py b/tests/test_litellm/proxy/a2a/test_agent_card.py index 0022053d8d1..d302bde7895 100644 --- a/tests/test_litellm/proxy/a2a/test_agent_card.py +++ b/tests/test_litellm/proxy/a2a/test_agent_card.py @@ -49,13 +49,31 @@ def test_preserves_top_level_url_for_runtime_invocation(): assert merged["url"] == "http://internal:9999/" -def test_overrides_protocol_version(): +def test_unsupported_protocol_version_defaults_to_1_0(): + # The fixture card pins "0.9", which LiteLLM does not serve; it falls back to + # the default rather than advertising a version the proxy can't honor. merged = merge_agent_card( _full_upstream_card(), proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE ) assert merged["protocolVersion"] == LITELLM_A2A_PROTOCOL_VERSION +def test_serves_pinned_protocol_version(): + for version in ("0.3", "1.0"): + card = _full_upstream_card() + card["protocolVersion"] = version + merged = merge_agent_card(card, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE) + assert merged["protocolVersion"] == version + assert merged["supportedInterfaces"][0]["protocolVersion"] == version + + +def test_absent_protocol_version_defaults_to_1_0(): + card = _full_upstream_card() + card.pop("protocolVersion", None) + merged = merge_agent_card(card, proxy_url=PROXY_URL, proxy_base_url=PROXY_BASE) + assert merged["protocolVersion"] == "1.0" + + def test_overrides_name_and_description_when_provided(): merged = merge_agent_card( _full_upstream_card(), diff --git a/tests/test_litellm/proxy/a2a/test_version_convert.py b/tests/test_litellm/proxy/a2a/test_version_convert.py new file mode 100644 index 00000000000..f3c51ca6b72 --- /dev/null +++ b/tests/test_litellm/proxy/a2a/test_version_convert.py @@ -0,0 +1,315 @@ +"""Unit tests for A2A protocol version normalization in +litellm/proxy/a2a/version_convert.py. + +These assert the conversion actually changes wire shape in the right direction and +preserves core fields on a round trip, so a mutation that no-ops or flips the direction +fails the suite. +""" + +import pytest + +from litellm.proxy.a2a.version_convert import ( + normalize_agent_card, + normalize_jsonrpc_response, + normalize_request_params, + normalize_stream_event, +) + +a2a = pytest.importorskip("a2a.compat.v0_3.conversions") + + +def _rpc(result: dict, request_id: str = "1") -> dict: + return {"jsonrpc": "2.0", "id": request_id, "result": result} + + +V03_MESSAGE = { + "kind": "message", + "messageId": "m1", + "role": "agent", + "parts": [{"kind": "text", "text": "hi"}], +} + +V03_TASK = { + "kind": "task", + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, +} + +V03_STATUS_UPDATE = { + "kind": "status-update", + "taskId": "t1", + "contextId": "c1", + "status": {"state": "working"}, + "final": False, +} + +V03_ARTIFACT_UPDATE = { + "kind": "artifact-update", + "taskId": "t1", + "contextId": "c1", + "artifact": {"artifactId": "a1", "parts": [{"kind": "text", "text": "out"}]}, +} + + +def test_send_result_0_3_to_1_0_wraps_in_envelope(): + out = normalize_jsonrpc_response(_rpc(V03_MESSAGE), "1.0", method="message/send") + assert "message" in out["result"] + assert "kind" not in out["result"] + assert out["result"]["message"]["messageId"] == "m1" + + +def test_send_result_1_0_to_0_3_unwraps_to_bare_kind(): + v1 = normalize_jsonrpc_response(_rpc(V03_MESSAGE), "1.0", method="message/send") + out = normalize_jsonrpc_response(v1, "0.3", method="message/send") + assert out["result"]["kind"] == "message" + assert out["result"]["messageId"] == "m1" + assert out["result"]["parts"][0]["text"] == "hi" + + +def test_send_result_same_version_is_identity_passthrough(): + rpc = _rpc(V03_MESSAGE) + out = normalize_jsonrpc_response(rpc, "0.3", method="message/send") + assert out is rpc + + +def test_task_result_round_trip_preserves_ids(): + v1 = normalize_jsonrpc_response(_rpc(V03_TASK), "1.0", method="tasks/get") + assert "kind" not in v1["result"] + assert v1["result"]["id"] == "t1" + back = normalize_jsonrpc_response(v1, "0.3", method="tasks/get") + assert back["result"]["kind"] == "task" + assert back["result"]["id"] == "t1" + assert back["result"]["contextId"] == "c1" + + +def test_error_response_passes_through_untouched(): + err = {"jsonrpc": "2.0", "id": "1", "error": {"code": -32600, "message": "bad"}} + assert normalize_jsonrpc_response(err, "1.0", method="message/send") is err + + +def test_malformed_result_falls_back_to_passthrough(): + # A 0.3 message missing required fields can't validate; conversion must not raise. + rpc = _rpc({"kind": "message"}) + out = normalize_jsonrpc_response(rpc, "1.0", method="message/send") + assert out["result"] == {"kind": "message"} + + +def test_unknown_shape_passes_through(): + rpc = _rpc({"unexpected": "shape"}) + out = normalize_jsonrpc_response(rpc, "1.0", method="message/send") + assert out is rpc + + +@pytest.mark.parametrize("event", [V03_STATUS_UPDATE, V03_ARTIFACT_UPDATE]) +def test_stream_event_round_trip_preserves_kind(event): + v1 = normalize_stream_event(_rpc(event), "1.0", request_id="1") + assert "kind" not in v1["result"] + back = normalize_stream_event(v1, "0.3", request_id="1") + assert back["result"]["kind"] == event["kind"] + assert back["result"]["taskId"] == "t1" + + +def test_stream_event_envelope_key_for_status_update(): + v1 = normalize_stream_event(_rpc(V03_STATUS_UPDATE), "1.0", request_id="1") + assert "statusUpdate" in v1["result"] + + +def test_request_params_lowering_is_noop_for_0_3(): + params = {"id": "t1", "historyLength": 5} + assert normalize_request_params(params, "0.3", method="tasks/get") is params + + +def test_request_params_lowering_get_task_to_0_3(): + out = normalize_request_params( + {"id": "t1", "historyLength": 5}, "1.0", method="tasks/get" + ) + assert out["id"] == "t1" + assert out["historyLength"] == 5 + + +def test_request_params_lowering_create_push_notification_config_preserves_task_id(): + out = normalize_request_params( + { + "parent": "tasks/task-1", + "configId": "cfg-1", + "config": {"url": "https://webhook.example.com"}, + }, + "1.0", + method="tasks/pushNotificationConfig/set", + ) + assert out["taskId"] == "task-1" + assert out["pushNotificationConfig"]["url"] == "https://webhook.example.com" + assert out["pushNotificationConfig"]["id"] == "cfg-1" + + +def test_flatten_create_push_notification_drops_redundant_envelope_key(): + from litellm.proxy.a2a.version_convert import ( + _flatten_create_push_notification_params, + ) + + flat = _flatten_create_push_notification_params( + { + "parent": "tasks/task-1", + "config": {"url": "https://chosen.example.com"}, + "pushNotificationConfig": {"url": "https://ignored.example.com"}, + } + ) + assert flat["url"] == "https://chosen.example.com" + assert "pushNotificationConfig" not in flat + assert "config" not in flat + + +def test_request_params_lowering_list_tasks_to_0_3(): + out = normalize_request_params( + { + "contextId": "ctx-1", + "pageSize": 10, + "status": "TASK_STATE_COMPLETED", + }, + "1.0", + method="tasks/list", + ) + assert out["contextId"] == "ctx-1" + assert out["pageSize"] == 10 + assert out["status"] == "completed" + + +@pytest.mark.parametrize( + "proto_status, expected", + [ + ("TASK_STATE_COMPLETED", "completed"), + ("TASK_STATE_INPUT_REQUIRED", "input-required"), + ("TASK_STATE_AUTH_REQUIRED", "auth-required"), + ("TASK_STATE_CANCELED", "canceled"), + ], +) +def test_list_tasks_status_filter_lowers_to_0_3_wire_value(proto_status, expected): + out = normalize_request_params( + {"status": proto_status}, + "1.0", + method="tasks/list", + ) + assert out["status"] == expected + + +def test_list_tasks_unspecified_status_is_dropped(): + out = normalize_request_params( + {"contextId": "ctx-1", "status": "TASK_STATE_UNSPECIFIED"}, + "1.0", + method="tasks/list", + ) + assert "status" not in out + assert out["contextId"] == "ctx-1" + + +@pytest.mark.parametrize( + "method, result", + [ + ( + "message/send", + { + "task": { + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, + }, + "vendorExtraField": "x", + }, + ), + ( + "tasks/get", + { + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, + "vendorExtraField": "x", + }, + ), + ], +) +def test_lowering_1_0_to_0_3_tolerates_unknown_upstream_fields(method, result): + out = normalize_jsonrpc_response(_rpc(result), "0.3", method=method) + lowered = out["result"] + assert lowered["kind"] == "task" + assert lowered["id"] == "t1" + assert "vendorExtraField" not in lowered + + +def test_stream_event_lowering_1_0_to_0_3_tolerates_unknown_fields(): + event = { + "task": {"id": "t1", "contextId": "c1", "status": {"state": "completed"}}, + "vendorExtraField": "x", + } + out = normalize_stream_event(_rpc(event), "0.3", request_id="1") + lowered = out["result"] + assert lowered["kind"] == "task" + assert lowered["id"] == "t1" + + +def test_list_tasks_result_round_trip_preserves_task_ids(): + rpc = _rpc( + { + "tasks": [ + { + "kind": "task", + "id": "t1", + "contextId": "c1", + "status": {"state": "completed"}, + } + ], + "nextPageToken": "tok", + } + ) + v1 = normalize_jsonrpc_response(rpc, "1.0", method="tasks/list") + assert "kind" not in v1["result"]["tasks"][0] + assert v1["result"]["tasks"][0]["id"] == "t1" + back = normalize_jsonrpc_response(v1, "0.3", method="tasks/list") + assert back["result"]["tasks"][0]["kind"] == "task" + assert back["result"]["tasks"][0]["id"] == "t1" + + +def _extended_card_1_0() -> dict: + return { + "name": "Card", + "description": "d", + "version": "1.0.0", + "supportedInterfaces": [ + { + "url": "https://upstream.example", + "protocolBinding": "JSONRPC", + "protocolVersion": "0.3", + }, + { + "url": "http://internal:9999", + "protocolBinding": "JSONRPC", + "protocolVersion": "0.3", + }, + ], + } + + +def test_agent_card_lowered_to_0_3_drops_additional_interfaces(): + # A 1.0 card with multiple interfaces would lower into a 0.3 card carrying the + # secondary backend URLs in ``additionalInterfaces``; those must be stripped so + # the conversion never re-exposes an upstream backend to A2A clients. + out = normalize_agent_card(_extended_card_1_0(), "0.3") + assert out["url"] == "https://upstream.example" + assert "additionalInterfaces" not in out + assert "supportedInterfaces" not in out + assert "http://internal:9999" not in str(out) + + +def test_agent_card_with_0_3_pin_and_supported_interfaces_is_lowered(): + card = _extended_card_1_0() + card["protocolVersion"] = "0.3" + + out = normalize_agent_card(card, "0.3") + + assert out["protocolVersion"] == "0.3" + assert "supportedInterfaces" not in out + + +def test_agent_card_same_version_passthrough(): + card = _extended_card_1_0() + assert normalize_agent_card(card, "1.0") is card diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 07e878401e0..09a73e076bb 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -967,6 +967,193 @@ async def test_get_extended_agent_card_rewrites_url(): assert body["result"]["name"] == "Test Agent" +@pytest.mark.asyncio +async def test_get_agent_card_uses_proxy_base_url_when_set(monkeypatch): + """Regression: discovery must expose the public proxy URL, not the internal one.""" + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + agent = _make_agent_mock() + agent.agent_card_params["protocolVersion"] = "1.0" + agent.agent_card_params["supportedInterfaces"] = [ + { + "url": "http://old-proxy.example.com/a2a/test-agent", + "protocolBinding": "JSONRPC", + "protocolVersion": "1.0", + } + ] + mock_request = MagicMock() + mock_request.base_url = "http://litellm-internal:4000/" + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + + from litellm.proxy.agent_endpoints.a2a_endpoints import get_agent_card + + response = await get_agent_card( + agent_id="test-agent", + request=mock_request, + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert body["url"] == "https://litellm.example.com/a2a/test-agent" + assert ( + body["supportedInterfaces"][0]["url"] + == "https://litellm.example.com/a2a/test-agent" + ) + + +@pytest.mark.asyncio +async def test_get_agent_card_normalizes_0_3_discovery_card(): + from litellm.proxy._types import UserAPIKeyAuth + + agent = _make_agent_mock() + agent.agent_card_params["protocolVersion"] = "0.3" + agent.agent_card_params["supportedInterfaces"] = [ + { + "url": "http://localhost:4000/a2a/test-agent", + "protocolBinding": "JSONRPC", + "protocolVersion": "0.3", + } + ] + mock_request = MagicMock() + mock_request.base_url = "http://localhost:4000/" + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + + from litellm.proxy.agent_endpoints.a2a_endpoints import get_agent_card + + response = await get_agent_card( + agent_id="test-agent", + request=mock_request, + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert body["protocolVersion"] == "0.3" + assert body["url"] == "http://localhost:4000/a2a/test-agent" + assert "supportedInterfaces" not in body + + +@pytest.mark.asyncio +async def test_get_agent_card_0_3_card_with_a2a_version_1_0_header(): + """Regression: 0.3 card normalized to 1.0 must not KeyError on debug log.""" + from litellm.proxy._types import UserAPIKeyAuth + + agent = _make_agent_mock() + agent.agent_card_params = { + "name": "Test Agent", + "description": "A test agent", + "url": "http://backend-agent:10001", + "version": "1.0.0", + "capabilities": {"streaming": True}, + "skills": [ + {"id": "s1", "name": "skill one", "description": "d", "tags": ["t"]} + ], + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + } + mock_request = MagicMock() + mock_request.base_url = "http://localhost:4000/" + mock_request.headers = {"a2a-version": "1.0"} + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + + from litellm.proxy.agent_endpoints.a2a_endpoints import get_agent_card + + response = await get_agent_card( + agent_id="test-agent", + request=mock_request, + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert "url" not in body + assert body["supportedInterfaces"][0]["url"] == ( + "http://localhost:4000/a2a/test-agent" + ) + + +@pytest.mark.asyncio +async def test_get_extended_agent_card_uses_proxy_base_url_when_set(monkeypatch): + """Regression: proxied extended cards must rewrite url to the public proxy base.""" + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + agent = _make_agent_mock() + mock_request = _make_request_mock("GetExtendedAgentCard", {}) + mock_request.base_url = "http://litellm-internal:4000/" + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + + upstream_card = { + "name": "Test Agent", + "url": "http://backend-agent:10001", + "description": "A test agent", + } + upstream_response = {"jsonrpc": "2.0", "id": "req-1", "result": upstream_card} + + mock_http_response = MagicMock() + mock_http_response.json.return_value = upstream_response + mock_http_response.is_success = True + mock_http_response.raise_for_status = MagicMock() + + mock_handler = MagicMock() + mock_handler.post = AsyncMock(return_value=mock_http_response) + mock_handler.client = MagicMock() + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + stack.enter_context( + patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=mock_handler, + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert body["result"]["url"] == "https://litellm.example.com/a2a/test-agent" + + +def test_build_merged_agent_card_uses_proxy_base_url_for_supported_interfaces( + monkeypatch, +): + """Regression: agent create/update must front supportedInterfaces with the public base.""" + from litellm.proxy.agent_endpoints.endpoints import _build_merged_agent_card + + monkeypatch.setenv("PROXY_BASE_URL", "https://litellm.example.com") + mock_request = MagicMock() + mock_request.base_url = "http://litellm-internal:4000/" + + merged = _build_merged_agent_card( + {"name": "My Agent", "url": "http://upstream:8080"}, + agent_id="jenkins_agent", + http_request=mock_request, + ) + + assert merged["supportedInterfaces"][0]["url"] == ( + "https://litellm.example.com/a2a/jenkins_agent" + ) + + @pytest.mark.asyncio async def test_unknown_method_returns_jsonrpc_error(): from litellm.proxy._types import UserAPIKeyAuth @@ -1076,6 +1263,173 @@ async def test_pascal_method_names_normalize_to_wire_format( ) +@pytest.mark.parametrize( + "params", + [ + { + "message": { + "messageId": "msg-1", + "role": "ROLE_USER", + "parts": [{"text": "hello"}], + }, + "configuration": {}, + }, + { + "message": { + "messageId": "msg-2", + "role": "user", + "parts": [{"kind": "text", "text": "hello"}], + }, + }, + ], +) +def test_build_message_send_params_accepts_wire_and_a2a_10(params): + from litellm.proxy.agent_endpoints.a2a_endpoints import _build_message_send_params + + result = _build_message_send_params(params) + assert result.message.role.value == "user" + assert result.message.parts[0].root.text == "hello" + + +def test_build_message_send_params_proto_fallback_ignores_unknown_fields(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _build_message_send_params + + result = _build_message_send_params( + { + "message": { + "messageId": "msg-1", + "role": "ROLE_USER", + "parts": [{"text": "hello"}], + }, + "configuration": {}, + "futureField": "ignored", + } + ) + assert result.message.role.value == "user" + + +@pytest.mark.asyncio +async def test_handle_stream_message_rejects_invalid_params_with_32602(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message + + response = await _handle_stream_message( + api_base="http://upstream.local", + request_id="req-1", + params={"message": 12345}, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join( + chunk.decode() if isinstance(chunk, bytes) else chunk for chunk in chunks + ) + payload = json.loads(body.strip()) + assert payload["error"]["code"] == -32602 + assert payload["id"] == "req-1" + + +@pytest.mark.asyncio +async def test_send_message_pascal_case_routes_to_asend_message(): + from litellm.proxy._types import UserAPIKeyAuth + + agent = _make_agent_mock() + params = { + "message": { + "messageId": "msg-123", + "role": "ROLE_USER", + "parts": [{"text": "Hello"}], + }, + "configuration": {}, + } + mock_request = _make_request_mock("SendMessage", params) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") + captured = {} + + async def capture_asend_message(request, **kwargs): + captured["method"] = request.method + captured["role"] = request.params.message.role.value + response = MagicMock() + response.model_dump.return_value = { + "jsonrpc": "2.0", + "id": request.id, + "result": { + "contextId": "ctx-1", + "kind": "message", + "messageId": "msg-123", + "parts": [{"kind": "text", "text": "Hello"}], + "role": "agent", + }, + } + return response + + with ExitStack() as stack: + for p in _base_patches(agent): + stack.enter_context(p) + stack.enter_context(patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True)) + stack.enter_context( + patch( + "litellm.a2a_protocol.asend_message", + new=AsyncMock(side_effect=capture_asend_message), + ) + ) + + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=MagicMock(), + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert "error" not in body, f"Got error: {body}" + assert captured["method"] == "message/send" + assert captured["role"] == "user" + assert "message" in body["result"] + assert body["result"]["message"]["role"] == "ROLE_AGENT" + + +def test_normalize_response_wraps_flat_message_result_for_1_0(): + from litellm.proxy.a2a.version_convert import normalize_jsonrpc_response + + wire_response = { + "jsonrpc": "2.0", + "id": "req-1", + "result": { + "contextId": "ctx-1", + "kind": "message", + "messageId": "msg-1", + "parts": [{"kind": "text", "text": "hello"}], + "role": "agent", + "taskId": "task-1", + }, + } + formatted = normalize_jsonrpc_response(wire_response, "1.0", method="message/send") + assert "message" in formatted["result"] + assert formatted["result"]["message"]["role"] == "ROLE_AGENT" + assert formatted["result"]["message"]["parts"] == [{"text": "hello"}] + assert "contextId" not in formatted["result"] + + +def test_normalize_response_keeps_wire_format_for_0_3(): + from litellm.proxy.a2a.version_convert import normalize_jsonrpc_response + + wire_response = { + "jsonrpc": "2.0", + "id": "req-1", + "result": { + "contextId": "ctx-1", + "kind": "message", + "messageId": "msg-1", + "parts": [{"kind": "text", "text": "hello"}], + "role": "agent", + }, + } + assert ( + normalize_jsonrpc_response(wire_response, "0.3", method="message/send") + is wire_response + ) + + @pytest.mark.asyncio async def test_task_method_upstream_jsonrpc_error_on_http_4xx_is_relayed(): """When upstream returns HTTP 4xx with a JSON-RPC error body, the error body @@ -1560,3 +1914,36 @@ async def test_caller_identity_headers_cannot_be_spoofed_via_forwarded_headers() assert ( posted_headers.get("X-LiteLLM-Team-Id") == "real-team" ), "authenticated team id must not be overridden by forwarded client headers" + + +def _agent(protocol_version): + agent = MagicMock() + agent.agent_card_params = ( + {"protocolVersion": protocol_version} if protocol_version is not None else {} + ) + return agent + + +def _request_with_a2a_header(value): + request = MagicMock() + request.headers = {"a2a-version": value} if value is not None else {} + return request + + +def test_served_version_config_governs_over_header(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _served_version + + # A 0.3-configured agent serves 0.3 even when the client asks for 1.0. + agent = _agent("0.3") + request = _request_with_a2a_header("1.0") + assert _served_version(agent, request) == "0.3" + + # A 1.0-configured agent serves 1.0 even when the client asks for 0.3. + assert _served_version(_agent("1.0"), _request_with_a2a_header("0.3")) == "1.0" + + +def test_served_version_falls_back_to_header_when_unconfigured(): + from litellm.proxy.agent_endpoints.a2a_endpoints import _served_version + + assert _served_version(_agent(None), _request_with_a2a_header("1.0")) == "1.0" + assert _served_version(_agent(None), _request_with_a2a_header(None)) == "0.3" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_version_e2e.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_version_e2e.py new file mode 100644 index 00000000000..069c72af53a --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_version_e2e.py @@ -0,0 +1,326 @@ +""" +Near-E2E tests for A2A 0.3/1.0 version routing through the proxy. + +Runs invoke_agent_a2a -> asend_message -> a2a-sdk 1.x -> ASGI mock upstream. +Only proxy auth/registry/pre-call plumbing is patched; version normalization +runs on the real response path. +""" + +from __future__ import annotations + +import json +from contextlib import ExitStack +from typing import Any, AsyncIterator, Dict, List, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from httpx import ASGITransport +from starlette.applications import Starlette +from starlette.responses import JSONResponse, StreamingResponse +from starlette.routing import Route + +pytest.importorskip("a2a.compat.v0_3.types") + +from litellm.proxy._types import UserAPIKeyAuth + +UPSTREAM_BASE = "http://testserver" + +_UPSTREAM_CALLS: List[Dict[str, Any]] = [] + + +def _upstream_card_payload() -> Dict[str, Any]: + return { + "protocolVersion": "0.3", + "name": "mock-agent", + "url": f"{UPSTREAM_BASE}/", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + + +def _message_result(request_id: Any) -> Dict[str, Any]: + return { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "kind": "message", + "role": "agent", + "messageId": "m-out", + "parts": [{"kind": "text", "text": "pong"}], + }, + } + + +def _sse_stream(request_id: Any) -> AsyncIterator[bytes]: + events = [ + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "kind": "task", + "id": "t1", + "contextId": "c1", + "status": {"state": "submitted"}, + }, + }, + _message_result(request_id), + ] + + async def _gen() -> AsyncIterator[bytes]: + for event in events: + yield f"data: {json.dumps(event)}\n\n".encode() + + return _gen() + + +async def _serve_upstream_agent_card(request: Any) -> JSONResponse: + return JSONResponse(_upstream_card_payload()) + + +async def _upstream_jsonrpc(request: Any) -> JSONResponse | StreamingResponse: + body = await request.json() + _UPSTREAM_CALLS.append(body) + request_id = body.get("id", "req-1") + method = body.get("method") + + if method == "message/stream": + return StreamingResponse( + _sse_stream(request_id), + media_type="text/event-stream", + ) + + return JSONResponse(_message_result(request_id)) + + +def _build_upstream_app() -> Starlette: + return Starlette( + routes=[ + Route( + "/.well-known/agent-card.json", + _serve_upstream_agent_card, + methods=["GET"], + ), + Route("/.well-known/agent.json", _serve_upstream_agent_card, methods=["GET"]), + Route("/", _upstream_jsonrpc, methods=["POST"]), + ] + ) + + +def _fake_get_async_httpx_client( + llm_provider: Any = None, params: Optional[Dict[str, Any]] = None +) -> MagicMock: + handler = MagicMock() + handler.client = httpx.AsyncClient( + transport=ASGITransport(_build_upstream_app()), + base_url=UPSTREAM_BASE, + ) + return handler + + +def _make_agent(*, protocol_version: str) -> MagicMock: + agent = MagicMock() + agent.agent_id = "test-agent" + agent.agent_name = "test-agent" + agent.agent_card_params = { + "url": f"{UPSTREAM_BASE}/", + "name": "Test Agent", + "protocolVersion": protocol_version, + } + agent.litellm_params = {} + agent.static_headers = None + agent.extra_headers = None + return agent + + +def _make_request( + method: str, + params: Dict[str, Any], + *, + headers: Optional[Dict[str, str]] = None, + request_id: str = "req-1", +) -> MagicMock: + request = MagicMock() + request.headers = headers or {} + request.json = AsyncMock( + return_value={ + "jsonrpc": "2.0", + "id": request_id, + "method": method, + "params": params, + } + ) + return request + + +async def _add_proxy_data(data: Dict[str, Any], **_: Any) -> Dict[str, Any]: + data["proxy_server_request"] = { + "url": "http://localhost:4000/a2a/test-agent", + "method": "POST", + "headers": {}, + "body": {}, + } + data.setdefault("metadata", {}) + return data + + +def _proxy_patches(agent: MagicMock) -> List[Any]: + from litellm.proxy.agent_endpoints import a2a_endpoints as a2a_endpoints_mod + + return [ + patch.object(a2a_endpoints_mod, "_get_agent", return_value=agent), + patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler" + ".AgentRequestHandler.is_agent_allowed", + new=AsyncMock(return_value=True), + ), + patch( + "litellm.proxy.common_request_processing.add_litellm_data_to_request", + new=AsyncMock(side_effect=_add_proxy_data), + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "1.0.0"), + patch( + "litellm.a2a_protocol.main.get_async_httpx_client", + side_effect=_fake_get_async_httpx_client, + ), + patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True), + ] + + +def _wire_send_params() -> Dict[str, Any]: + return { + "message": { + "role": "user", + "messageId": "m-in", + "parts": [{"kind": "text", "text": "ping"}], + } + } + + +def _a2a10_send_params() -> Dict[str, Any]: + return { + "message": { + "role": "ROLE_USER", + "messageId": "m-in", + "parts": [{"text": "ping"}], + }, + "configuration": {}, + } + + +@pytest.fixture(autouse=True) +def _clear_upstream_calls() -> None: + _UPSTREAM_CALLS.clear() + + +@pytest.mark.asyncio +async def test_proxy_serves_1_0_when_agent_pinned_and_upstream_speaks_03(): + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + agent = _make_agent(protocol_version="1.0") + request = _make_request( + "SendMessage", + _a2a10_send_params(), + headers={"a2a-version": "1.0"}, + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", user_id="test-user", team_id="test-team" + ) + + with ExitStack() as stack: + for item in _proxy_patches(agent): + stack.enter_context(item) + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert "error" not in body, body + assert "message" in body["result"] + assert "kind" not in body["result"] + assert body["result"]["message"]["parts"][0]["text"] == "pong" + assert _UPSTREAM_CALLS, "expected upstream to receive a JSON-RPC call" + assert _UPSTREAM_CALLS[0]["method"] == "message/send" + + +@pytest.mark.asyncio +async def test_proxy_serves_0_3_when_agent_pinned_passthrough(): + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + agent = _make_agent(protocol_version="0.3") + request = _make_request("message/send", _wire_send_params()) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", user_id="test-user", team_id="test-team" + ) + + with ExitStack() as stack: + for item in _proxy_patches(agent): + stack.enter_context(item) + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=user_api_key_dict, + ) + + body = json.loads(response.body.decode()) + assert "error" not in body, body + assert body["result"]["kind"] == "message" + assert body["result"]["parts"][0]["text"] == "pong" + assert "message" not in body["result"] + + +@pytest.mark.asyncio +async def test_proxy_streaming_serves_1_0_envelopes(): + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + agent = _make_agent(protocol_version="1.0") + request = _make_request( + "SendStreamingMessage", + _a2a10_send_params(), + headers={"a2a-version": "1.0"}, + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test", user_id="test-user", team_id="test-team" + ) + + with ExitStack() as stack: + for item in _proxy_patches(agent): + stack.enter_context(item) + + response = await invoke_agent_a2a( + agent_id="test-agent", + request=request, + fastapi_response=MagicMock(), + user_api_key_dict=user_api_key_dict, + ) + + lines: List[Dict[str, Any]] = [] + async for raw_line in response.body_iterator: + line = ( + raw_line.decode().strip() + if isinstance(raw_line, (bytes, bytearray)) + else str(raw_line).strip() + ) + if line: + lines.append(json.loads(line)) + + assert lines, "expected at least one streamed JSON-RPC event" + message_events = [ + line + for line in lines + if isinstance(line.get("result"), dict) and "message" in line["result"] + ] + assert message_events, f"expected a 1.0 message envelope, got: {lines}" + assert message_events[-1]["result"]["message"]["parts"][0]["text"] == "pong" + assert _UPSTREAM_CALLS, "expected upstream streaming call" + assert _UPSTREAM_CALLS[0]["method"] == "message/stream" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py index 554f98d7209..a51d6f6abcc 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py @@ -10,7 +10,7 @@ per call; default timeout uses DEFAULT_A2A_AGENT_TIMEOUT). """ import sys -from unittest.mock import AsyncMock, MagicMock, call, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -231,57 +231,95 @@ async def test_each_agent_gets_only_its_own_static_headers(): # --------------------------------------------------------------------------- +def _fake_get_async_httpx_client_factory(captured_calls: list): + """Return a side_effect that records every (params, client) pair.""" + + def _fake_get_async_httpx_client(llm_provider, params, **kwargs): + client = MagicMock() + client.headers = MagicMock() + handler = MagicMock() + handler.client = client + captured_calls.append({"params": params.copy(), "client": client}) + return handler + + return _fake_get_async_httpx_client + + +async def _fake_create_client(base_url, client_config=None, **kwargs): + client = MagicMock() + if client_config is not None: + client._litellm_httpx_client = client_config.httpx_client + return client + + @pytest.mark.asyncio async def test_create_a2a_client_uses_fresh_httpx_client(): """ - Two calls to create_a2a_client with different extra_headers must NOT - share the same underlying httpx.AsyncClient instance. - """ - import httpx + Two calls to create_a2a_client with different extra_headers must produce + distinct underlying httpx clients — preventing header bleed between agents. + The test checks: + 1. get_async_httpx_client was called twice (once per create_a2a_client call). + 2. The two returned A2A clients carry distinct httpx client objects (direct + proof of header isolation, not just cache-key difference). + 3. The cache-key param differs between calls (so the real LRU cache cannot + return the same httpx client even under load). + """ + pytest.importorskip("a2a.client") from litellm.a2a_protocol.main import create_a2a_client - created_clients = [] - - fake_agent_card = MagicMock() - fake_agent_card.name = "test-agent" - - class FakeResolver: - def __init__(self, **kw): - created_clients.append(kw.get("httpx_client")) - - async def get_agent_card(self): - return fake_agent_card - - class FakeA2AClient: - def __init__(self, httpx_client, agent_card): - self._client = httpx_client - self._litellm_agent_card = agent_card + captured_calls: list = [] with ( patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True), - patch("litellm.a2a_protocol.main.A2ACardResolver", FakeResolver), - patch("litellm.a2a_protocol.main._A2AClient", FakeA2AClient), + patch( + "litellm.a2a_protocol.main.get_async_httpx_client", + side_effect=_fake_get_async_httpx_client_factory(captured_calls), + ), + patch( + "litellm.a2a_protocol.main.create_client", + new=AsyncMock(side_effect=_fake_create_client), + ), ): - await create_a2a_client( + a2a_client_a = await create_a2a_client( base_url="http://agent-a:9999", extra_headers={"Authorization": "Bearer a"}, ) - await create_a2a_client( + a2a_client_b = await create_a2a_client( base_url="http://agent-b:9999", extra_headers={"Authorization": "Bearer b"}, ) - assert len(created_clients) == 2 - # Must be distinct objects assert ( - created_clients[0] is not created_clients[1] - ), "create_a2a_client reused a cached httpx client — headers will bleed between agents" + len(captured_calls) == 2 + ), "create_a2a_client should call get_async_httpx_client once per invocation" + + # Direct proof: the two A2A clients must carry distinct httpx client objects. + # If they share one, mutating agent-B's Authorization header would bleed into A. + httpx_a = getattr(a2a_client_a, "_litellm_httpx_client", None) + httpx_b = getattr(a2a_client_b, "_litellm_httpx_client", None) + assert httpx_a is not None, "a2a_client_a missing _litellm_httpx_client" + assert httpx_b is not None, "a2a_client_b missing _litellm_httpx_client" + assert httpx_a is not httpx_b, ( + "create_a2a_client returned the same httpx client for two agents with " + "different headers — Authorization header will bleed between agents" + ) + + # Also verify the cache-key param differs so the LRU cache never conflates them. + key_a = captured_calls[0]["params"].get("disable_aiohttp_transport") + key_b = captured_calls[1]["params"].get("disable_aiohttp_transport") + assert key_a is not None, "cache-key param 'disable_aiohttp_transport' missing" + assert key_b is not None, "cache-key param 'disable_aiohttp_transport' missing" + assert key_a != key_b, ( + f"create_a2a_client used the same cache key for two agents with different " + f"headers — headers will bleed: key_a={key_a!r}, key_b={key_b!r}" + ) @pytest.mark.asyncio async def test_create_a2a_client_default_timeout_matches_constant(): """When timeout is omitted, httpx client params must use DEFAULT_A2A_AGENT_TIMEOUT.""" + pytest.importorskip("a2a.client") from litellm.a2a_protocol.main import create_a2a_client captured: dict = {} @@ -293,28 +331,16 @@ async def test_create_a2a_client_default_timeout_matches_constant(): handler.client.headers = MagicMock() return handler - fake_agent_card = MagicMock() - fake_agent_card.name = "test-agent" - - class _FakeResolver: - def __init__(self, **kw): - pass - - async def get_agent_card(self): - return fake_agent_card - - class _FakeA2AClient: - def __init__(self, httpx_client, agent_card): - pass - with ( patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True), patch( "litellm.a2a_protocol.main.get_async_httpx_client", side_effect=_capture_get_async_httpx_client, ), - patch("litellm.a2a_protocol.main.A2ACardResolver", _FakeResolver), - patch("litellm.a2a_protocol.main._A2AClient", _FakeA2AClient), + patch( + "litellm.a2a_protocol.main.create_client", + new=AsyncMock(side_effect=_fake_create_client), + ), ): await create_a2a_client(base_url="http://127.0.0.1:9") @@ -324,6 +350,7 @@ async def test_create_a2a_client_default_timeout_matches_constant(): @pytest.mark.asyncio async def test_create_a2a_client_explicit_timeout_overrides_default(): """Explicit timeout= must be passed through to the httpx client params.""" + pytest.importorskip("a2a.client") from litellm.a2a_protocol.main import create_a2a_client captured: dict = {} @@ -335,28 +362,16 @@ async def test_create_a2a_client_explicit_timeout_overrides_default(): handler.client.headers = MagicMock() return handler - fake_agent_card = MagicMock() - fake_agent_card.name = "test-agent" - - class _FakeResolver: - def __init__(self, **kw): - pass - - async def get_agent_card(self): - return fake_agent_card - - class _FakeA2AClient: - def __init__(self, httpx_client, agent_card): - pass - with ( patch("litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", True), patch( "litellm.a2a_protocol.main.get_async_httpx_client", side_effect=_capture_get_async_httpx_client, ), - patch("litellm.a2a_protocol.main.A2ACardResolver", _FakeResolver), - patch("litellm.a2a_protocol.main._A2AClient", _FakeA2AClient), + patch( + "litellm.a2a_protocol.main.create_client", + new=AsyncMock(side_effect=_fake_create_client), + ), ): await create_a2a_client(base_url="http://127.0.0.1:9", timeout=42.5) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 0cb03023e0f..3740c01b7fc 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -786,3 +786,32 @@ class TestCheckAgentUrlHealth: ) result = await _check_agent_url_health(agent) assert result["healthy"] is True + + +@pytest.mark.parametrize( + "base_url", + ["http://0.0.0.0:4000/", "http://localhost:4000/", "https://api.example.com/"], +) +def test_merged_agent_card_url_has_no_double_slash_without_proxy_base_url( + monkeypatch, base_url +): + """Without PROXY_BASE_URL, request.base_url carries a trailing slash; the merged + card's supportedInterfaces URL must still join cleanly (no `//a2a`).""" + from litellm.proxy.agent_endpoints.endpoints import _build_merged_agent_card + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + + http_request = MagicMock() + http_request.base_url = base_url + + merged = _build_merged_agent_card( + _sample_agent_card_params(), + agent_id="agent-xyz", + http_request=http_request, + agent_name="Test Agent", + ) + + interface_url = merged["supportedInterfaces"][0]["url"] + assert interface_url == f"{base_url.rstrip('/')}/a2a/agent-xyz" + assert "//a2a" not in interface_url diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index 14b6729bf93..442dcd48f66 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -6,13 +6,15 @@ export interface FieldConfig { name: string; label: string; - type: "text" | "textarea" | "url" | "switch" | "list"; + type: "text" | "textarea" | "url" | "switch" | "list" | "select"; required?: boolean; tooltip?: string; placeholder?: string; defaultValue?: any; rows?: number; validation?: any[]; + options?: string[]; + helpText?: string; } export interface SectionConfig { @@ -69,9 +71,13 @@ export const AGENT_FORM_CONFIG: { { name: "protocolVersion", label: "Protocol Version", - type: "text", - placeholder: "1.0", + type: "select", + options: ["1.0", "0.3"], defaultValue: "1.0", + tooltip: + "The A2A protocol version LiteLLM serves to clients for this agent. LiteLLM converts the upstream agent's responses to this version, so clients always see the version you pick here regardless of the original agent's version.", + helpText: + "LiteLLM serves this version to clients and converts the upstream agent's responses to match it, regardless of the original agent's version.", }, ], }, diff --git a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx index 7965af6322d..ff103218346 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx @@ -47,9 +47,18 @@ const AgentFormFields: React.FC = ({ showAgentName = true, : undefined } tooltip={field.tooltip} + extra={field.helpText} > {field.type === "textarea" ? ( + ) : field.type === "select" ? ( + ) : ( )} diff --git a/uv.lock b/uv.lock index da44ad25715..f76be69505f 100644 --- a/uv.lock +++ b/uv.lock @@ -20,23 +20,29 @@ members = [ ] constraints = [ { name = "aiohttp", specifier = ">=3.14.1,<4.0" }, + { name = "packaging", specifier = ">=24.0" }, { name = "tornado", specifier = ">=6.5.6" }, ] +overrides = [{ name = "packaging", specifier = ">=24.0" }] [[package]] name = "a2a-sdk" -version = "0.3.24" +version = "1.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "culsans", marker = "python_full_version < '3.13'" }, { name = "google-api-core" }, + { name = "googleapis-common-protos" }, { name = "httpx" }, { name = "httpx-sse" }, + { name = "json-rpc" }, + { name = "packaging" }, { name = "protobuf" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ad/76/cefa956fb2d3911cb91552a1da8ce2dbb339f1759cb475e2982f0ae2332b/a2a_sdk-0.3.24.tar.gz", hash = "sha256:3581e6e8a854cd725808f5732f90b7978e661b6d4e227a4755a8f063a3c1599d", size = 255550, upload-time = "2026-02-20T10:05:43.423Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/7e/8ac10bbf8b15b16574355f39b17dbdf617a282c27b41c7ff2116e30336df/a2a_sdk-1.1.0.tar.gz", hash = "sha256:e8102dad1b36709dbdc3d19319e38e6dfa3b3a79c30416030eb2d482576be204", size = 375726, upload-time = "2026-05-29T09:34:43.015Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/10/6e/cae5f0caea527b39c0abd7204d9416768764573c76649ca03cc345a372be/a2a_sdk-0.3.24-py3-none-any.whl", hash = "sha256:7b248767096bb55311f57deebf6b767349388d94c1b376c60cb8f6b715e053f6", size = 145752, upload-time = "2026-02-20T10:05:41.729Z" }, + { url = "https://files.pythonhosted.org/packages/d4/ea/3a5b160cfd51c67759b08748051094d9365ceff18127633d0021950c9860/a2a_sdk-1.1.0-py3-none-any.whl", hash = "sha256:d7f5846caf18033d8bf3108b11ec827dd8dd32f867c98848ede0e39474be93be", size = 241886, upload-time = "2026-05-29T09:34:41.484Z" }, ] [[package]] @@ -165,6 +171,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/22/0a/62e7232dc9484fbec112ceb32efb6a624cc7994ec6e2b019286f17c4e8f2/aiohttp-3.14.1-cp313-cp313-win_arm64.whl", hash = "sha256:250d14af67f6b6a1a4a811049b1afa69d61d617fca6bf33149b3ab1a6dbcf7b8", size = 447723, upload-time = "2026-06-07T21:08:00.154Z" }, ] +[[package]] +name = "aiologic" +version = "0.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "sniffio", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "wrapt", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/53/a7/809482759f40079f4c4328c7318bf569ae25d457f5017aad30a1b9aafedc/aiologic-0.17.0.tar.gz", hash = "sha256:65aa058e858c94cd208badb188e7f00b54dcabb3ba85b34f794db98074d108b9", size = 251625, upload-time = "2026-06-14T12:24:35.367Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/6b/5f75d6194b597ac32bbdbb7b524a28fb1fa98bd0ddcefce94b313a818cc0/aiologic-0.17.0-py3-none-any.whl", hash = "sha256:1bf4d3e4314df2bcb06a9e696417204e206ab50e10ec98d28d157e2e57634f74", size = 161084, upload-time = "2026-06-14T12:24:34.146Z" }, +] + [[package]] name = "aiosignal" version = "1.4.0" @@ -1191,6 +1211,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5d/8c/ce3823c06c2804f194f9e64f0d67fa3f4094a39f2bb1a990cd03603af8fc/cryptography-48.0.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:6184ca7b174f28d7c703f1290d4b297217c45355f77a98f67e9b7f14549ac54a", size = 3742204, upload-time = "2026-06-09T22:31:34.773Z" }, ] +[[package]] +name = "culsans" +version = "0.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiologic", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d9/e3/49afa1bc180e0d28008ec6bcdf82a4072d1c7a41032b5b759b60814ca4b0/culsans-0.11.0.tar.gz", hash = "sha256:0b43d0d05dce6106293d114c86e3fb4bfc63088cfe8ff08ed3fe36891447fe33", size = 107546, upload-time = "2025-12-31T23:15:38.196Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/5d/9fb19fb38f6d6120422064279ea5532e22b84aa2be8831d49607194feda3/culsans-0.11.0-py3-none-any.whl", hash = "sha256:278d118f63fc75b9db11b664b436a1b83cc30d9577127848ba41420e66eb5a47", size = 21811, upload-time = "2025-12-31T23:15:37.189Z" }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -2746,6 +2779,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl", hash = "sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713", size = 309071, upload-time = "2025-12-15T08:41:44.973Z" }, ] +[[package]] +name = "json-rpc" +version = "1.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6d/9e/59f4a5b7855ced7346ebf40a2e9a8942863f644378d956f68bcef2c88b90/json-rpc-1.15.0.tar.gz", hash = "sha256:e6441d56c1dcd54241c937d0a2dcd193bdf0bdc539b5316524713f554b7f85b9", size = 28854, upload-time = "2023-06-11T09:45:49.078Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/9e/820c4b086ad01ba7d77369fb8b11470a01fac9b4977f02e18659cf378b6b/json_rpc-1.15.0-py2.py3-none-any.whl", hash = "sha256:4a4668bbbe7116feb4abbd0f54e64a4adcf4b8f648f19ffa0848ad0f6606a9bf", size = 39450, upload-time = "2023-06-11T09:45:47.136Z" }, +] + [[package]] name = "jsonlines" version = "4.0.0" @@ -3428,7 +3470,7 @@ proxy-dev = [ [package.metadata] requires-dist = [ - { name = "a2a-sdk", marker = "extra == 'extra-proxy'", specifier = ">=0.3.24,<1.0" }, + { name = "a2a-sdk", marker = "extra == 'extra-proxy'", specifier = ">=1.1.0,<2.0" }, { name = "aiohttp", specifier = ">=3.10,<4.0" }, { name = "anthropic", extras = ["vertex"], marker = "extra == 'proxy-runtime'", specifier = ">=0.84.0,<1.0" }, { name = "apscheduler", marker = "extra == 'proxy'", specifier = ">=3.11.2,<4.0" }, @@ -3584,7 +3626,7 @@ healthcheck = [ { name = "pyyaml", specifier = "==6.0.3" }, ] proxy-dev = [ - { name = "a2a-sdk", specifier = "==0.3.24" }, + { name = "a2a-sdk", specifier = "==1.1.0" }, { name = "azure-identity", specifier = "==1.25.2" }, { name = "hypercorn", specifier = "==0.17.3" }, { name = "opentelemetry-api", specifier = "==1.28.0" }, @@ -5144,11 +5186,11 @@ wheels = [ [[package]] name = "packaging" -version = "23.2" +version = "26.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/fb/2b/9b9c33ffed44ee921d0967086d653047286054117d584f1b1a7c22ceaf7b/packaging-23.2.tar.gz", hash = "sha256:048fb0e9405036518eaaf48a55953c750c11e1a1b68e0dd1a9d62ed0c092cfc5", size = 146714, upload-time = "2023-10-01T13:50:05.279Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ec/1a/610693ac4ee14fcdf2d9bf3c493370e4f2ef7ae2e19217d7a237ff42367d/packaging-23.2-py3-none-any.whl", hash = "sha256:8c491190033a9af7e1d931d0b5dacc2ef47509b34dd0de67ed209b5203fc88c7", size = 53011, upload-time = "2023-10-01T13:50:03.745Z" }, + { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, ] [[package]]