From 80b95eb8c0910fe834c59d1fb765e15fcee17a6a Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 27 Oct 2026 16:35:41 -0700 Subject: [PATCH] refactor(a2a): share the A2A method map with the applied-guardrails stream check The applied-guardrails header check now reads the A2A streaming methods from the same map the A2A endpoint uses to translate method names, instead of its own copy --- litellm/proxy/a2a/version_convert.py | 23 ++++++++++++++++++- .../proxy/agent_endpoints/a2a_endpoints.py | 20 +++------------- litellm/proxy/utils.py | 9 ++------ 3 files changed, 27 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/a2a/version_convert.py b/litellm/proxy/a2a/version_convert.py index bde2ff45a88..eb55bff50bc 100644 --- a/litellm/proxy/a2a/version_convert.py +++ b/litellm/proxy/a2a/version_convert.py @@ -24,7 +24,8 @@ The two wire shapes: ``Task`` result is a bare object without ``kind``. """ -from collections.abc import Callable +from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import Final, Literal from pydantic import BaseModel @@ -39,6 +40,26 @@ JsonDict = dict[str, object] _V1_SEND_ENVELOPE_KEYS: Final = frozenset({"message", "task"}) _V1_STREAM_ENVELOPE_KEYS: Final = frozenset({"message", "task", "statusUpdate", "artifactUpdate"}) +A2A_PASCAL_TO_WIRE: Final[Mapping[str, str]] = MappingProxyType( + { + "SendMessage": "message/send", + "SendStreamingMessage": "message/stream", + "GetTask": "tasks/get", + "ListTasks": "tasks/list", + "CancelTask": "tasks/cancel", + "SubscribeToTask": "tasks/resubscribe", + "CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set", + "GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get", + "ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list", + "DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete", + "GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard", + } +) +_STREAMING_WIRE_METHODS: Final = frozenset({"message/stream", "tasks/resubscribe"}) +A2A_STREAMING_METHODS: Final = _STREAMING_WIRE_METHODS | frozenset( + pascal for pascal, wire in A2A_PASCAL_TO_WIRE.items() if wire in _STREAMING_WIRE_METHODS +) + def _dump_03(model: BaseModel) -> JsonDict: """Dump a compat (0.3) pydantic model to its camelCase wire dict.""" diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 4691c7b0bc5..85a5ca30816 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.a2a.version_convert import ( + A2A_PASCAL_TO_WIRE, A2AVersion, normalize_agent_card, normalize_jsonrpc_response, @@ -65,21 +66,6 @@ _SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType( } ) -_PASCAL_TO_WIRE: Final[Mapping[str, str]] = { - "SendMessage": "message/send", - "SendStreamingMessage": "message/stream", - "GetTask": "tasks/get", - "ListTasks": "tasks/list", - "CancelTask": "tasks/cancel", - "SubscribeToTask": "tasks/resubscribe", - "CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set", - "GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get", - "ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list", - "DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete", - "GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard", -} - - _DECODED_JSON: Final = TypeAdapter(object) @@ -132,7 +118,7 @@ def _served_version(agent: "AgentResponse", request: Request, original_method: s configured: Final = (agent.agent_card_params or {}).get("protocolVersion") if configured in ("0.3", "1.0"): return configured - if original_method in _PASCAL_TO_WIRE: + if original_method in A2A_PASCAL_TO_WIRE: return "1.0" return "1.0" if request.headers.get("a2a-version", "").startswith("1.") else "0.3" @@ -699,7 +685,7 @@ async def invoke_agent_a2a( params = body.get("params", {}) if method: - method = _PASCAL_TO_WIRE.get(method, method) + method = A2A_PASCAL_TO_WIRE.get(method, method) if isinstance(params, dict): # extract any litellm params from the params - eg. 'guardrails' diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9a960d30d3d..36e5c95c79c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -179,6 +179,7 @@ from litellm.proxy._types import ( SpecialModelNames, UserAPIKeyAuth, ) +from litellm.proxy.a2a.version_convert import A2A_STREAMING_METHODS from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header @@ -721,12 +722,6 @@ def _ran_or_runs_outside_pipelines( ) -# The A2A JSON-RPC methods answered with a stream, as both protocol versions spell them -_A2A_STREAMING_METHODS: Final = frozenset( - ("message/stream", "tasks/resubscribe", "SendStreamingMessage", "SubscribeToTask") -) - - def _post_call_steps_named_before_headers( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> frozenset[str]: @@ -735,7 +730,7 @@ def _post_call_steps_named_before_headers( headers go out: every step of a whole response, and of a streamed one the steps of the pipelines that hold the stream until they finish. """ - if data.get("stream") is not True and data.get("method") not in _A2A_STREAMING_METHODS: + if data.get("stream") is not True and data.get("method") not in A2A_STREAMING_METHODS: return pipeline_managed_guardrail_names(data, "post_call") gated: Final = _stream_gated_pipelines(data, user_api_key_dict) return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated)