scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

refactor(a2a): share the A2A method map with the applied-guardrails stream check

Browse source 
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
This commit is contained in:
yucheng 2026-10-27 16:35:41 -07:00
parent 0676e6a460
commit 80b95eb8c0
3 changed files with 27 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/a2a/version_convert.py
View file

@ -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."""

20
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -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'

9
litellm/proxy/utils.py
View file

@ -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)