From a5f980600c6d66bdb645707f16bddd6dc690c5ac Mon Sep 17 00:00:00 2001 From: moe-berri Date: Sat, 31 Oct 2026 12:54:10 -0700 Subject: [PATCH] fix(a2a): answer proxy refusals on /a2a POST routes as JSON-RPC errors (#49476) * test(a2a): expect JSON-RPC errors for proxy refusals on /a2a POST routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(a2a): answer proxy refusals on /a2a POST routes as JSON-RPC errors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(a2a): build JSON-RPC refusal data from typed fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(a2a): stamp the OTel server span on JSON-RPC refusals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(a2a): cover JSON-RPC refusal route in unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/agent_endpoints/a2a_endpoints.py | 68 +++++++++++-- .../authorization/test_agent_access.py | 98 +++++++++++++++++-- .../test_agent_least_privilege.py | 9 +- .../agent_endpoints/test_a2a_endpoints.py | 54 ++++++++++ 4 files changed, 214 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 11a63a6a536..227a1b2bdf0 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -11,7 +11,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping from copy import deepcopy from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol @@ -19,7 +19,8 @@ from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, StreamingResponse -from pydantic import TypeAdapter, ValidationError +from fastapi.routing import APIRoute +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -27,7 +28,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping 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.llms.custom_httpx.sse_lines import aiter_sse_lines -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.a2a.version_convert import ( A2AVersion, normalize_agent_card, @@ -42,6 +43,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import ( ) 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.common_utils.openai_error_payload import status_typed_proxy_exception from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING, coerce_keepalive_interval, @@ -221,6 +223,57 @@ def _jsonrpc_error( ) +class _JsonRpcRequestId(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str | int | None = None + + +async def _jsonrpc_request_id(request: Request) -> str | int | None: + try: + return _JsonRpcRequestId.model_validate_json(await request.body()).id + except ValidationError: + return None + + +async def _jsonrpc_refusal(request: Request, refusal: ProxyException | HTTPException) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; a route that answers an error itself must stamp the OTel server span like the global ProxyException handler does + ) + + error: Final = refusal if isinstance(refusal, ProxyException) else status_typed_proxy_exception(refusal) + _close_dangling_otel_server_span(request, int(error.code), exc=refusal) + return JSONResponse( + content={ + "jsonrpc": "2.0", + "id": await _jsonrpc_request_id(request), + "error": { + "code": -32000, + "message": error.message, + "data": {"type": error.type, "param": error.param, "code": error.code}, + }, + }, + status_code=int(error.code), + headers=error.headers if isinstance(refusal, ProxyException) else refusal.headers, + ) + + +class _JsonRpcRefusalRoute(APIRoute): + def get_route_handler(self) -> Callable[[Request], Coroutine[object, object, Response]]: + handler: Final = super().get_route_handler() + + async def answer_refusals_as_jsonrpc(request: Request) -> Response: + try: + return await handler(request) + except (ProxyException, HTTPException) as refusal: + return await _jsonrpc_refusal(request, refusal) + + return answer_refusals_as_jsonrpc + + +jsonrpc_router: Final = APIRouter(route_class=_JsonRpcRefusalRoute) + + async def _get_agent(agent_id: str) -> "AgentResponse | None": """Look up an agent by ID or name. Returns None if not found.""" from litellm.proxy.common_utils.registry_read_through import ( @@ -649,17 +702,17 @@ async def get_agent_card( raise HTTPException(status_code=500, detail=str(e)) -@router.post( +@jsonrpc_router.post( "/a2a/{agent_id}", tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], ) -@router.post( +@jsonrpc_router.post( "/a2a/{agent_id}/message/send", tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], ) -@router.post( +@jsonrpc_router.post( "/v1/a2a/{agent_id}/message/send", tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], @@ -1055,3 +1108,6 @@ async def invoke_agent_a2a( if isinstance(e, litellm.BadRequestError): return _jsonrpc_error(body.get("id"), -32602, e.message, 400) return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500) + + +router.include_router(jsonrpc_router) diff --git a/tests/integration/authorization/test_agent_access.py b/tests/integration/authorization/test_agent_access.py index 09e2afd289d..8c4d5714c50 100644 --- a/tests/integration/authorization/test_agent_access.py +++ b/tests/integration/authorization/test_agent_access.py @@ -1,7 +1,8 @@ +import itertools import json import uuid from collections.abc import Callable -from typing import Final +from typing import Final, Literal import pytest from integration._support.client import Gateway, Scenario @@ -18,6 +19,22 @@ class _JsonRpcResult(BaseModel): result: dict[str, JsonValue] +class _JsonRpcErrorBody(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + code: int + message: str + data: dict[str, JsonValue] + + +class _JsonRpcError(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + jsonrpc: Literal["2.0"] + id: str | int | None + error: _JsonRpcErrorBody + + _STREAM_EVENTS: Final = ( {"kind": "status-update", "taskId": "t1", "contextId": "c1", "status": {"state": "working"}, "final": False}, { @@ -163,6 +180,14 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car separators=(",", ":"), ) + def jsonrpc_refusal(text: str, identity: str) -> bool: + error: Final = _JsonRpcError.model_validate_json(text) + return (error.id, error.error.code, error.error.message) == ( + marker, + -32000, + f"Agent '{identity}' is not allowed for your key/team. Contact proxy admin for access.", + ) + spellings: Final = ( ("send by id", lambda agent, name: f"/a2a/{agent}", send_payload, "message/send"), ("send by name", lambda agent, name: f"/a2a/{name}", send_payload, "message/send"), @@ -224,13 +249,13 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car for label, key in callers.items(): for spelling, status, text in attempt(key, denied, denied_name): - refused: Final = refusal(denied_name if "by name" in spelling else denied) - assert (status, text) == (403, refused), f"{label} {spelling}: {status} {text}" + identity: Final = denied_name if "by name" in spelling else denied + assert status == 403 and jsonrpc_refusal(text, identity), f"{label} {spelling}: {status} {text}" for spelling, payload in denied_tasks: response: Final = gateway.client.post( f"/a2a/{denied}", headers={"Authorization": f"Bearer {key}"}, json=payload ) - assert (response.status_code, response.text) == (403, refusal(denied)), ( + assert response.status_code == 403 and jsonrpc_refusal(response.text, denied), ( f"{label} {spelling}: {response.status_code} {response.text}" ) assert card(key, denied) == (403, refusal(denied)), label @@ -278,13 +303,72 @@ def test_agent_permissions_on_keys_teams_and_access_groups_gate_a2a_send_and_car emptied: Final = gateway.request("PUT", f"/v1/access_group/{group_id}", {"access_agent_ids": []}) assert emptied.status_code == 200, emptied.text emptied_outcomes: Final = attempt(callers["access group"], allowed, allowed_name) - assert emptied_outcomes == tuple( - (spelling, 403, refusal(allowed_name if "by name" in spelling else allowed)) - for spelling, _, _ in emptied_outcomes + assert all( + status == 403 and jsonrpc_refusal(text, allowed_name if "by name" in spelling else allowed) + for spelling, status, text in emptied_outcomes ), emptied_outcomes assert wire.drain() == (), emptied_outcomes +def test_a2a_refusals_answer_as_jsonrpc_errors_that_echo_the_callers_request_id(gateway: Gateway) -> None: + marker: Final = "a2arefuse" + uuid.uuid4().hex[:12] + + def upstream(request: Request) -> Reply: + return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode()) + + with wire_server(upstream) as wire, gateway.scenario() as scenario: + agent: Final = _register_agent(gateway, scenario, marker, wire.url) + refused_callers: Final = ( + ( + "over budget", + scenario.key(object_permission={"agents": [agent]}, max_budget=0.0001, spend=1.0), + 422, + "Budget has been exceeded!", + "budget_exceeded", + ), + ( + "agent not granted", + scenario.key(object_permission={"agents": [marker + "-other"]}), + 403, + f"Agent '{agent}' is not allowed for your key/team. Contact proxy admin for access.", + None, + ), + ("unknown key", "sk-" + uuid.uuid4().hex, 401, "Authentication Error", None), + ) + message: Final = { + "kind": "message", + "role": "user", + "messageId": marker + "-in", + "parts": [{"kind": "text", "text": "ping"}], + } + requests: Final = tuple( + (path, method, request_id) + for path, (method, request_id) in itertools.product( + (f"/a2a/{agent}", f"/a2a/{agent}/message/send", f"/v1/a2a/{agent}/message/send"), + (("message/send", marker), ("message/stream", 7), ("SendMessage", "l10")), + ) + ) + for label, key, status, message_start, error_type in refused_callers: + for path, method, request_id in requests: + response = gateway.client.post( + path, + headers={"Authorization": f"Bearer {key}", "Accept": "application/json, text/event-stream"}, + json={"jsonrpc": "2.0", "id": request_id, "method": method, "params": {"message": message}}, + ) + context = f"{label} {path} {method}: {response.status_code} {response.text}" + assert response.status_code == status, context + error = _JsonRpcError.model_validate_json(response.text) + assert (error.id, error.error.code) == (request_id, -32000), context + assert error.error.message.startswith(message_start), context + assert error_type is None or error.error.data["type"] == error_type, context + unparseable: Final = gateway.client.post( + f"/a2a/{agent}", headers={"Authorization": "Bearer sk-" + uuid.uuid4().hex}, content=b"not json" + ) + assert unparseable.status_code == 400, unparseable.text + assert _JsonRpcError.model_validate_json(unparseable.text).id is None, unparseable.text + assert wire.drain() == () + + def test_legacy_agent_json_card_path_applies_agent_permissions_to_non_admin_keys(gateway: Gateway) -> None: pytest.skip( "BUG: GET /a2a/{agent_id}/.well-known/agent.json returns 401 'Only proxy admin can be used' for non-admin keys because agent_inference_routes lists only agent-card.json" diff --git a/tests/integration/authorization/test_agent_least_privilege.py b/tests/integration/authorization/test_agent_least_privilege.py index d993920a32e..a08c4047b9e 100644 --- a/tests/integration/authorization/test_agent_least_privilege.py +++ b/tests/integration/authorization/test_agent_least_privilege.py @@ -221,8 +221,13 @@ def test_agent_access_default_deny_denies_ungranted_key_but_keeps_explicit_grant denied_sends: Final = tuple( _send(candidate, agent_id, key_a, f"denied-{index}-{uuid.uuid4().hex}") for index in range(8) ) - _assert_denial(denied_sends[0], agent_id) - assert all(response.json() == denied_sends[0].json() for response in denied_sends), [ + assert denied_sends[0].status_code == 403, denied_sends[0].text + assert denied_sends[0].json()["error"] == { + "code": -32000, + "message": f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.", + "data": denied_sends[0].json()["error"]["data"], + }, denied_sends[0].text + assert all(response.json()["error"] == denied_sends[0].json()["error"] for response in denied_sends), [ response.text for response in denied_sends ] denied_completion: Final = _chat_completion(candidate, agent_name, key_a) diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index d6d6774d414..016a5151e16 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from fastapi.testclient import TestClient from pydantic import BaseModel, ConfigDict import litellm @@ -2941,3 +2942,56 @@ async def test_failed_upstream_stream_relays_its_body_only_when_that_is_a_json_o "id": "req-1", "method": "tasks/resubscribe", } + + +def _a2a_app_refusing_with(refusal: Exception) -> TestClient: + from fastapi import FastAPI + + from litellm.proxy.agent_endpoints.a2a_endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + def refuse() -> UserAPIKeyAuth: + raise refusal + + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = refuse + return TestClient(app) + + +@pytest.mark.parametrize("path", ["/a2a/agent-1", "/a2a/agent-1/message/send", "/v1/a2a/agent-1/message/send"]) +def test_a2a_auth_refusals_answer_as_jsonrpc_errors_with_the_callers_id(path: str): + from litellm.proxy._types import ProxyException + + client: Final = _a2a_app_refusing_with( + ProxyException(message="Budget has been exceeded!", type="budget_exceeded", param=None, code=422) + ) + + response: Final = client.post(path, json={"jsonrpc": "2.0", "id": "req-7", "method": "message/send"}) + + assert response.status_code == 422 + assert response.json() == { + "jsonrpc": "2.0", + "id": "req-7", + "error": { + "code": -32000, + "message": "Budget has been exceeded!", + "data": {"type": "budget_exceeded", "param": None, "code": "422"}, + }, + } + + +def test_a2a_http_refusals_answer_as_jsonrpc_errors_with_a_null_id_for_unparseable_bodies(): + from fastapi import HTTPException + + client: Final = _a2a_app_refusing_with(HTTPException(status_code=403, detail="Agent 'agent-1' is not allowed")) + + response: Final = client.post("/a2a/agent-1", content=b"not json") + + assert response.status_code == 403 + body: Final = response.json() + assert body["jsonrpc"] == "2.0" + assert body["id"] is None + assert body["error"]["code"] == -32000 + assert body["error"]["message"] == "Agent 'agent-1' is not allowed" + assert body["error"]["data"]["code"] == "403"