mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
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>
This commit is contained in:
moe-berri
GitHub
parent
7143aeeb86
commit
a5f980600c
4 changed files with 214 additions and 15 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue