ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix : a2a bugs 030626 (#29566)

Browse Source 
* Fix error code and context id injection bug

* Add support for all A2A methods

* Add logging

* address greptile review: relay upstream JSON-RPC errors, move _PASCAL_TO_WIRE to module level, add error path tests

* fix(a2a): run pre_call_hook for tasks/resubscribe SSE path to enforce guardrails

tasks/resubscribe was returning the raw SSE stream without calling proxy_logging_obj.pre_call_hook, silently bypassing any guardrails configured on the agent. This patch calls pre_call_hook before streaming begins and wires post_call_failure_hook into the SSE generator so errors are logged. Adds a regression test verifying the hook is called.

* fix(a2a): use get_async_httpx_client instead of creating httpx clients per request

Creating httpx.AsyncClient instances per-request adds ~500ms latency. Switch _forward_jsonrpc and _forward_jsonrpc_sse to use the shared client from get_async_httpx_client(httpxSpecialProvider.A2A).

* fix(a2a): forward caller identity headers on task ops; validate push notification URL

Two security fixes for task management methods:

1. All task operations (tasks/get, tasks/list, tasks/cancel, tasks/resubscribe, push notification config methods) now forward X-LiteLLM-User-Id and X-LiteLLM-Team-Id headers to the upstream agent, so the agent can scope task access to the authenticated caller.

2. tasks/pushNotificationConfig/set validates the callback URL before forwarding: requires HTTPS and rejects private/loopback/reserved IP ranges and localhost hostnames to prevent SSRF.

* Fix A2A task hook and push URL handling

* fix(a2a): fix mypy type errors for request_id and header_name dict key types

* Fix A2A request id and params forwarding

* Forward trace IDs for A2A task calls

* fix(a2a): strip client-forwarded X-LiteLLM-* headers before applying authenticated identity

A client could send x-a2a-<agent>-x-litellm-user-id in their request and have it forwarded to the upstream agent as an authenticated identity header. Fix: sanitize any X-LiteLLM-* headers from agent_extra_headers before merging, then apply the authenticated identity headers last so they always override client-supplied values.

* Fix A2A SSE fallback JSON-RPC error code

* Fix A2A SSE error id backfill

* fix(a2a): validate both push notification url fields to close SSRF bypass

* fix(a2a): widen request_id annotation to match JSON-RPC id call sites

* fix(a2a): run post-call streaming hook for tasks/resubscribe so agent guardrails apply

tasks/resubscribe returned the raw upstream SSE stream without routing events
through the post-call streaming hook, so output guardrails configured on the
agent were silently skipped for streaming task subscriptions while every other
task method and message/stream applied them. Parse upstream JSON-RPC SSE events
and feed them through async_streaming_data_generator, matching message/stream,
so guardrails inspect the streamed task content. Adds a regression test that
fails when the streamed events bypass the guardrail hook.

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Sameer Kankute 2026-06-03 23:44:15 +05:30 committed by GitHub
parent c7ab9adde5
commit 48c9fabb26
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 1639 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
litellm/a2a_protocol/main.py
View File

@ -159,7 +159,9 @@ async def _send_message_via_completion_bridge(
api_base=api_base,
)
return LiteLLMSendMessageResponse.from_dict(response_dict)
return LiteLLMSendMessageResponse.from_dict(
response_dict, request_id=str(request.id)
)
async def _execute_a2a_send_with_retry(
@ -317,15 +319,6 @@ async def asend_message(
)
card_url = getattr(agent_card, "url", None) if agent_card else None
context_id = trace_id or str(uuid.uuid4())
message = request.params.message
if isinstance(message, dict):
if message.get("context_id") is None:
message["context_id"] = context_id
else:
if getattr(message, "context_id", None) is None:
message.context_id = context_id
a2a_response = await _execute_a2a_send_with_retry(
a2a_client=a2a_client,
request=request,
@ -338,7 +331,9 @@ async def asend_message(
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
# Wrap in LiteLLM response type for _hidden_params support
response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response)
response = LiteLLMSendMessageResponse.from_a2a_response(
a2a_response, request_id=str(request.id)
)
# Calculate token usage from request and response
response_dict = a2a_response.model_dump(mode="json", exclude_none=True)

367
litellm/proxy/agent_endpoints/a2a_endpoints.py
View File

@ -6,12 +6,14 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
"""
import json
from typing import Any, Dict, List, Optional
from typing import Any, AsyncGenerator, Dict, List, Optional
from urllib.parse import urlparse
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse, StreamingResponse
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.agent_endpoints.utils import merge_agent_headers
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -19,9 +21,66 @@ from litellm.types.utils import all_litellm_params
router = APIRouter()
_PASCAL_TO_WIRE: Dict[str, str] = {
"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",
}
def _validate_push_notification_url(url: str) -> None:
parsed = urlparse(url)
if parsed.scheme != "https":
raise HTTPException(
status_code=400,
detail="Push notification URL must use HTTPS",
)
try:
validate_url(url)
except (SSRFError, ValueError) as e:
raise HTTPException(status_code=400, detail=str(e)) from e
def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Dict[str, str]:
headers: Dict[str, str] = {}
if user_api_key_dict.user_id:
headers["X-LiteLLM-User-Id"] = user_api_key_dict.user_id
if user_api_key_dict.team_id:
headers["X-LiteLLM-Team-Id"] = user_api_key_dict.team_id
return headers
def _forwarding_headers(
user_api_key_dict: UserAPIKeyAuth,
request_data: dict,
agent_extra_headers: Optional[Dict[str, str]],
) -> Optional[Dict[str, str]]:
sanitized = (
{
k: v
for k, v in agent_extra_headers.items()
if not k.lower().startswith("x-litellm-")
}
if agent_extra_headers
else None
)
merged = merge_agent_headers(dynamic_headers=sanitized, static_headers=None) or {}
identity = _caller_identity_headers(user_api_key_dict)
trace_id = request_data.get("litellm_trace_id")
if trace_id:
identity["X-LiteLLM-Trace-Id"] = str(trace_id)
merged.update(identity)
return merged or None
def _jsonrpc_error(
request_id: Optional[str],
request_id: Optional[Any],
code: int,
message: str,
status_code: int = 400,
@ -67,9 +126,158 @@ 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:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
headers = {"Content-Type": "application/json", **(extra_headers or {})}
handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.A2A,
params={"timeout": 60.0},
)
resp = await handler.post(agent_url, json=body, headers=headers)
try:
result = resp.json()
except Exception:
resp.raise_for_status()
raise
if not resp.is_success and "error" not in result:
resp.raise_for_status()
return result
async def _a2a_sse_event_source(
agent_url: str,
body: dict,
request_id: Optional[Any] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncGenerator[dict, None]:
"""Stream an upstream A2A SSE response as parsed JSON-RPC event dicts.
Upstream HTTP/JSON-RPC errors are surfaced as a single JSON-RPC error event
so the caller can relay them instead of breaking the stream.
"""
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.agents import _normalize_a2a_jsonrpc_response
from litellm.types.llms.custom_http import httpxSpecialProvider
headers = {
"Content-Type": "application/json",
"Accept": "text/event-stream",
**(extra_headers or {}),
}
handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.A2A,
params={"timeout": None},
)
async_client = handler.client
req = async_client.build_request("POST", agent_url, json=body, headers=headers)
resp = await async_client.send(req, stream=True)
try:
if not resp.is_success:
error_body = await resp.aread()
error_event: Optional[dict] = None
try:
parsed = json.loads(error_body)
if isinstance(parsed, dict) and "error" in parsed:
error_event = _normalize_a2a_jsonrpc_response(
parsed, request_id=request_id
)
except Exception:
error_event = None
yield error_event or {
"jsonrpc": "2.0",
"id": request_id,
"error": {"code": -32603, "message": resp.reason_phrase},
}
return
async for line in resp.aiter_lines():
stripped = line.strip()
if not stripped.startswith("data:"):
continue
payload = stripped[len("data:") :].strip()
if not payload:
continue
try:
yield json.loads(payload)
except Exception:
continue
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,
) -> StreamingResponse:
event_source = _a2a_sse_event_source(
agent_url, body, request_id=request_id, extra_headers=extra_headers
)
def _serialize_chunk(chunk: Any) -> str:
return f"data: {json.dumps(chunk)}\n\n"
def _serialize_error(proxy_exc: Any) -> str:
return (
"data: "
+ json.dumps(
{
"jsonrpc": "2.0",
"id": request_id,
"error": {
"code": -32603,
"message": getattr(proxy_exc, "message", str(proxy_exc)),
},
}
)
+ "\n\n"
)
if (
proxy_logging_obj is not None
and user_api_key_dict is not None
and request_data is not None
):
# Route streamed events through the shared streaming generator so the
# post-call streaming hook (and therefore agent guardrails) inspects
# tasks/resubscribe output the same way message/stream does.
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
generator: AsyncGenerator[str, None] = (
ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
response=event_source,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
proxy_logging_obj=proxy_logging_obj,
serialize_chunk=_serialize_chunk,
serialize_error=_serialize_error,
)
)
else:
async def _passthrough() -> AsyncGenerator[str, None]:
async for chunk in event_source:
yield _serialize_chunk(chunk)
generator = _passthrough()
return StreamingResponse(generator, media_type="text/event-stream")
async def _handle_stream_message(
api_base: Optional[str],
request_id: str,
request_id: Any,
params: dict,
litellm_params: Optional[dict] = None,
agent_id: Optional[str] = None,
@ -310,8 +518,6 @@ async def invoke_agent_a2a( # noqa: PLR0915
- message/send: Send a message and get a response
- message/stream: Send a message and stream the response
"""
from litellm.a2a_protocol import asend_message
from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
)
@ -322,9 +528,11 @@ async def invoke_agent_a2a( # noqa: PLR0915
version,
)
body = {}
body: Dict[str, Any] = {}
request_data: Dict[str, Any] = body
try:
body = await request.json()
request_data = body
verbose_proxy_logger.debug(f"A2A request for agent '{agent_id}': {body}")
@ -334,11 +542,14 @@ async def invoke_agent_a2a( # noqa: PLR0915
body.get("id"), -32600, "Invalid Request: jsonrpc must be '2.0'"
)
request_id = body.get("id")
method = body.get("method")
request_id: Optional[Any] = body.get("id")
method: Optional[str] = body.get("method")
params = body.get("params", {})
if params:
if method:
method = _PASCAL_TO_WIRE.get(method, method)
if isinstance(params, dict):
# extract any litellm params from the params - eg. 'guardrails'
# ``metadata`` is intentionally excluded: it's a first-class A2A
# ``MessageSendParams`` field that the completion bridge forwards
@ -347,20 +558,12 @@ async def invoke_agent_a2a( # noqa: PLR0915
# silently drop the caller's A2A request-level metadata.
params_to_remove = []
for key, value in params.items():
if key in all_litellm_params and key != "metadata":
if key in all_litellm_params and key not in {"id", "metadata"}:
params_to_remove.append(key)
body[key] = value
for key in params_to_remove:
params.pop(key)
if not A2A_SDK_AVAILABLE:
return _jsonrpc_error(
request_id,
-32603,
"Server error: 'a2a' package not installed. Please install 'a2a-sdk'.",
500,
)
# Find the agent
agent = _get_agent(agent_id)
if agent is None:
@ -441,6 +644,7 @@ async def invoke_agent_a2a( # noqa: PLR0915
route_type="asend_message",
version=version,
)
request_data = data
# Build merged headers for the backend agent
static_headers: Dict[str, str] = dict(agent.static_headers or {})
@ -453,9 +657,10 @@ async def invoke_agent_a2a( # noqa: PLR0915
# 1. Admin-configured extra_headers: forward named headers from client request
if agent.extra_headers:
for header_name in agent.extra_headers:
val = normalized.get(header_name.lower())
header_name_str = str(header_name)
val = normalized.get(header_name_str.lower())
if val is not None:
dynamic_headers[header_name] = val
dynamic_headers[header_name_str] = val
# 2. Convention-based forwarding: x-a2a-{agent_id_or_name}-{header_name}
# Matches both agent_id (UUID) and agent_name (alias), case-insensitive.
@ -489,10 +694,20 @@ async def invoke_agent_a2a( # noqa: PLR0915
# Route through SDK functions
if method == "message/send":
from litellm.a2a_protocol import asend_message
from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE
if not A2A_SDK_AVAILABLE:
return _jsonrpc_error(
request_id,
-32603,
"Server error: 'a2a' package not installed. Please install 'a2a-sdk'.",
500,
)
from a2a.types import MessageSendParams, SendMessageRequest
a2a_request = SendMessageRequest(
id=request_id,
id=request_id if request_id is not None else "",
params=MessageSendParams(**params),
)
# Defer spend-log until after post_call_success_hook so guardrail
@ -532,7 +747,7 @@ async def invoke_agent_a2a( # noqa: PLR0915
elif method == "message/stream":
return await _handle_stream_message(
api_base=agent_url,
request_id=request_id,
request_id=request_id if request_id is not None else "",
params=params,
litellm_params=litellm_params,
agent_id=agent.agent_id,
@ -543,6 +758,106 @@ async def invoke_agent_a2a( # noqa: PLR0915
request_data=data,
proxy_logging_obj=proxy_logging_obj,
)
elif method in {
"tasks/get",
"tasks/list",
"tasks/cancel",
"tasks/pushNotificationConfig/set",
"tasks/pushNotificationConfig/get",
"tasks/pushNotificationConfig/list",
"tasks/pushNotificationConfig/delete",
"agent/getAuthenticatedExtendedCard",
}:
if not agent_url:
return _jsonrpc_error(
request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500
)
if method == "tasks/pushNotificationConfig/set":
if not isinstance(params, dict):
raise HTTPException(
status_code=400,
detail="params must be an object",
)
push_config = params.get("pushNotificationConfig", {})
if "pushNotificationConfig" in params and not isinstance(
push_config, dict
):
raise HTTPException(
status_code=400,
detail="pushNotificationConfig must be an object",
)
for callback_url in (params.get("url"), push_config.get("url")):
if not callback_url:
continue
if not isinstance(callback_url, str):
raise HTTPException(
status_code=400,
detail="Push notification URL must be a string",
)
_validate_push_notification_url(callback_url)
forward_body = {
"jsonrpc": "2.0",
"id": request_id,
"method": method,
"params": params,
}
caller_headers = _forwarding_headers(
user_api_key_dict=user_api_key_dict,
request_data=data,
agent_extra_headers=agent_extra_headers,
)
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}"
from litellm.types.agents import LiteLLMSendMessageResponse
response = LiteLLMSendMessageResponse.from_dict(
result, request_id=request_id
)
response = await proxy_logging_obj.post_call_success_hook(
user_api_key_dict=user_api_key_dict,
data=data,
response=response,
)
return JSONResponse(
content=(
response.model_dump(mode="json", exclude_none=True)
if hasattr(response, "model_dump")
else response
)
)
elif method == "tasks/resubscribe":
if not agent_url:
return _jsonrpc_error(
request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500
)
forward_body = {
"jsonrpc": "2.0",
"id": request_id,
"method": method,
"params": params,
}
sse_caller_headers = _forwarding_headers(
user_api_key_dict=user_api_key_dict,
request_data=data,
agent_extra_headers=agent_extra_headers,
)
return await _forward_jsonrpc_sse(
agent_url,
forward_body,
request_id=request_id,
extra_headers=sse_caller_headers,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=data,
)
else:
return _jsonrpc_error(request_id, -32601, f"Method '{method}' not found")
@ -550,4 +865,12 @@ async def invoke_agent_a2a( # noqa: PLR0915
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error invoking agent: {e}")
try:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data=request_data,
)
except Exception:
pass
return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {str(e)}", 500)

38
litellm/types/agents.py
View File

@ -298,6 +298,23 @@ class MakeAgentsPublicRequest(BaseModel):
agent_ids: List[str]
def _normalize_a2a_jsonrpc_response(
response_dict: Dict[str, Any],
request_id: Optional[Any] = None,
) -> Dict[str, Any]:
"""
Ensure JSON-RPC responses include ``id`` when the caller supplied one.
The a2a SDK may omit ``id`` on error payloads even when the upstream agent
returned it. Backfill from the outbound request id so LiteLLM can surface the
agent error instead of failing Pydantic validation.
"""
normalized = dict(response_dict)
if normalized.get("id") is None and request_id is not None:
normalized["id"] = str(request_id)
return normalized
class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase):
"""
LiteLLM wrapper for A2A SendMessageResponse.
@ -322,31 +339,42 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase):
@classmethod
def from_a2a_response(
cls, response: "SendMessageResponse"
cls,
response: "SendMessageResponse",
request_id: Optional[Any] = None,
) -> "LiteLLMSendMessageResponse":
"""
Create a LiteLLMSendMessageResponse from an a2a SDK SendMessageResponse.
Args:
response: The a2a SDK SendMessageResponse
request_id: JSON-RPC request id to backfill when the SDK omits it on errors
Returns:
LiteLLMSendMessageResponse with _hidden_params support
"""
# Convert the a2a response to a dict
response_dict = response.model_dump(mode="json", exclude_none=True)
response_dict = _normalize_a2a_jsonrpc_response(
response_dict, request_id=request_id
)
return cls(**response_dict)
@classmethod
def from_dict(cls, response_dict: Dict[str, Any]) -> "LiteLLMSendMessageResponse":
def from_dict(
cls,
response_dict: Dict[str, Any],
request_id: Optional[Any] = None,
) -> "LiteLLMSendMessageResponse":
"""
Create a LiteLLMSendMessageResponse from a dict.
Args:
response_dict: Dict with A2A response structure
request_id: JSON-RPC request id to backfill when missing on error payloads
Returns:
LiteLLMSendMessageResponse with _hidden_params support
"""
return cls(**response_dict)
return cls(
**_normalize_a2a_jsonrpc_response(response_dict, request_id=request_id)
)

43
tests/test_litellm/a2a_protocol/test_send_message_response.py Normal file
View File

@ -0,0 +1,43 @@
"""Tests for LiteLLMSendMessageResponse JSON-RPC normalization."""
from litellm.types.agents import LiteLLMSendMessageResponse
def test_from_dict_backfills_id_on_agent_error_response():
agent_error = {
"jsonrpc": "2.0",
"error": {"code": -32054, "message": "Session not found"},
}
response = LiteLLMSendMessageResponse.from_dict(
agent_error, request_id="r1"
)
assert response.id == "r1"
assert response.error == {"code": -32054, "message": "Session not found"}
assert response.result is None
def test_from_dict_preserves_existing_id():
payload = {
"id": "upstream-id",
"jsonrpc": "2.0",
"error": {"code": -32001, "message": "Task not found"},
}
response = LiteLLMSendMessageResponse.from_dict(
payload, request_id="r1"
)
assert response.id == "upstream-id"
def test_from_dict_without_request_id_still_requires_id():
try:
LiteLLMSendMessageResponse.from_dict(
{"jsonrpc": "2.0", "error": {"code": -32054, "message": "x"}}
)
except Exception as exc:
assert "id" in str(exc).lower()
else:
raise AssertionError("expected validation error when id and request_id missing")

1212
tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py
View File

File diff suppressed because it is too large Load Diff