scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge a0d6d74066 into 156bda3e5b

Browse source
This commit is contained in:
so1 2026-11-07 02:31:45 +03:00 • committed by GitHub
commit 84479a824f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 207 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -277,9 +277,14 @@ def _jsonrpc_error(
request_id: object,
code: int,
message: str,
status_code: int = 400,
status_code: int = 200,
) -> JSONResponse:
"""Create a JSON-RPC 2.0 error response."""
"""Create a JSON-RPC 2.0 error response.
JSON-RPC errors travel in the body with HTTP 200, like the A2A reference server
answers them, so a2a-sdk clients reach the typed error instead of failing on the
HTTP status. Server-side failures keep an explicit 5xx.
"""
return JSONResponse(
content={
"jsonrpc": "2.0",
@ -290,6 +295,11 @@ def _jsonrpc_error(
)
def _with_request_id(content: Mapping[str, object], request_id: object) -> Mapping[str, object]:
"""Echo the caller's id verbatim: ``exclude_none`` dumps drop a null id and the SDK request type has no null."""
return {**content, "id": request_id}
class _JsonRpcRequestId(BaseModel):
model_config = ConfigDict(frozen=True)
@ -864,7 +874,7 @@ async def invoke_agent_a2a(
target_agent_id=agent_id,
model=f"a2a_agent/{agent_id}",
)
return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404)
return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found")
served_version: Final = _served_version(agent, request, original_method)
@ -1070,10 +1080,9 @@ async def invoke_agent_a2a(
else {}
)
return JSONResponse(
content=normalize_jsonrpc_response(
response_dict,
served_version,
method="message/send",
content=_with_request_id(
normalize_jsonrpc_response(response_dict, served_version, method="message/send"),
request_id,
)
)
@ -1160,7 +1169,9 @@ async def invoke_agent_a2a(
)
return JSONResponse(
content=(
response.model_dump(mode="json", exclude_none=True) if hasattr(response, "model_dump") else response
_with_request_id(response.model_dump(mode="json", exclude_none=True), request_id)
if hasattr(response, "model_dump")
else response
)
)
@ -1202,7 +1213,7 @@ async def invoke_agent_a2a(
except Exception:
pass
if isinstance(e, litellm.BadRequestError):
return _jsonrpc_error(body.get("id"), -32602, e.message, 400)
return _jsonrpc_error(body.get("id"), -32602, e.message)
if isinstance(e, litellm.AuthenticationError) and caused_by_entra_rejection(e):
return _jsonrpc_error(body.get("id"), -32000, e.message, 401)
return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {e}", 500)

24
tests/integration/authorization/test_agent_denial_logging.py
View file

@ -482,7 +482,7 @@ def test_denied_completion_bridge_agent_logs_one_failure_and_never_reaches_the_b
pytest.param("/a2a/{}", "message/stream", None, id="regression-stream-0.3"),
],
)
def test_unknown_agent_keeps_its_jsonrpc_404_and_logs_one_zero_spend_failure_for_the_caller(
def test_unknown_agent_answers_its_jsonrpc_error_with_http_200_and_logs_one_zero_spend_404_failure_for_the_caller(
rig: _Rig, route: str, method: str, version: str | None
) -> None:
unknown: Final = "no-such-" + uuid.uuid4().hex
@ -499,7 +499,7 @@ def test_unknown_agent_keeps_its_jsonrpc_404_and_logs_one_zero_spend_failure_for
request_id=request_id,
message_id=message,
)
assert (response.status_code, response.content) == (404, _not_found(request_id, unknown))
assert (response.status_code, response.content) == (200, _not_found(request_id, unknown))
assert response.headers["content-type"] == "application/json"
_assert_never_dispatched(rig, scenario, message)
_assert_one_failure(rig, caller, status=404, agent_id=unknown, model=f"a2a_agent/{unknown}")
@ -513,7 +513,7 @@ def test_unknown_agent_echoes_the_jsonrpc_id_with_its_type_and_logs_one_failure(
with rig.proxy.scenario() as scenario:
caller: Final = _caller(rig, scenario)
response: Final = _send(rig, caller.key, f"/a2a/{unknown}", request_id=request_id)
assert (response.status_code, response.content) == (404, _not_found(request_id, unknown))
assert (response.status_code, response.content) == (200, _not_found(request_id, unknown))
assert type(JSON_OBJECT.validate_json(response.content)["id"]) is type(request_id)
_assert_one_failure(rig, caller, status=404, agent_id=unknown, model=f"a2a_agent/{unknown}")
@ -554,9 +554,10 @@ def test_client_supplied_attribution_cannot_rewrite_a_denial_failure_row(rig: _R
f"/a2a/{target}", headers={"Authorization": f"Bearer {caller.key}", "A2A-Version": "1.0"}, json=body
)
expected: Final = _refusal(target, response) if outcome == "denied" else _not_found(request_id, unknown)
assert (response.status_code, response.content) == (403 if outcome == "denied" else 404, expected)
assert (response.status_code, response.content) == (403 if outcome == "denied" else 200, expected)
model: Final = f"a2a_agent/{rig.denied.name}" if outcome == "denied" else f"a2a_agent/{unknown}"
row: Final = _assert_one_failure(rig, caller, status=response.status_code, agent_id=target, model=model)
logged_status: Final = 403 if outcome == "denied" else 404
row: Final = _assert_one_failure(rig, caller, status=logged_status, agent_id=target, model=model)
assert row["request_id"] not in {success_request, "spoof2", request_id}, row
metadata: Final = _metadata(row)
assert (metadata["user_api_key_alias"], metadata["user_api_key_team_id"]) == (caller.alias, caller.team_id)
@ -591,7 +592,7 @@ def test_failing_or_transforming_failure_callbacks_never_change_the_denial_respo
with hostile.proxy.scenario() as scenario:
caller: Final = _caller(hostile, scenario)
response: Final = _send(hostile, caller.key, f"/a2a/{target}", request_id=request_id)
expected: Final = (403, _refusal(target, response)) if outcome == "denied" else (404, _not_found(request_id, unknown))
expected: Final = (403, _refusal(target, response)) if outcome == "denied" else (200, _not_found(request_id, unknown))
assert (response.status_code, response.content) == expected
@ -604,9 +605,10 @@ def test_failing_or_transforming_failure_callbacks_do_not_stop_the_denial_from_b
with hostile.proxy.scenario() as scenario:
caller: Final = _caller(hostile, scenario)
response: Final = _send(hostile, caller.key, f"/a2a/{target}")
assert response.status_code == (403 if outcome == "denied" else 404), response.text
assert response.status_code == (403 if outcome == "denied" else 200), response.text
model: Final = f"a2a_agent/{hostile.denied.name}" if outcome == "denied" else f"a2a_agent/{unknown}"
_assert_one_failure(hostile, caller, status=response.status_code, agent_id=target, model=model)
logged_status: Final = 403 if outcome == "denied" else 404
_assert_one_failure(hostile, caller, status=logged_status, agent_id=target, model=model)
def test_denials_follow_the_model_denied_chat_logging_under_error_log_flags(flagged: _Rig) -> None:
@ -630,7 +632,7 @@ def test_denials_follow_the_model_denied_chat_logging_under_error_log_flags(flag
denied: Final = _send(flagged, denied_caller.key, f"/a2a/{flagged.denied.agent_id}")
assert (denied.status_code, denied.content) == (403, _refusal(flagged.denied.agent_id, denied))
missing: Final = _send(flagged, unknown_caller.key, f"/a2a/{unknown}", request_id=request_id)
assert (missing.status_code, missing.content) == (404, _not_found(request_id, unknown))
assert (missing.status_code, missing.content) == (200, _not_found(request_id, unknown))
_assert_never_dispatched(flagged, scenario, "never-sent")
chat_hooks: Final = held(
lambda: _hooks(flagged.events.of(chat_alias)),
@ -662,7 +664,7 @@ def test_client_error_denials_never_raise_an_llm_exceptions_alert(alerting: tupl
with rig.proxy.scenario() as scenario, wire_server(failing) as broken:
caller: Final = _caller(rig, scenario)
assert _send(rig, caller.key, f"/a2a/{rig.denied.agent_id}").status_code == 403
assert _send(rig, caller.key, f"/a2a/{unknown}").status_code == 404
assert JSON_OBJECT.validate_json(_send(rig, caller.key, f"/a2a/{unknown}").content)["error"]["code"] == -32000
model: Final = scenario.model(api_base=broken.url + "/v1", num_retries=0)
failed: Final = rig.proxy.request(
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "x"}]}
@ -798,7 +800,7 @@ def test_guard_jsonrpc_1_0_envelope_keeps_its_invalid_request_body(rig: _Rig) ->
extra={"jsonrpc": "1.0"},
)
assert (response.status_code, response.content) == (
400,
200,
_compact(
{
"jsonrpc": "2.0",

180
tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py
View file

@ -1416,7 +1416,7 @@ async def test_upstream_authentication_error_without_an_entra_rejection_stays_an
@pytest.mark.asyncio
async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400():
async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_with_http_200():
from litellm.proxy._types import UserAPIKeyAuth
agent = _make_agent_mock()
@ -1463,7 +1463,7 @@ async def test_agentcore_invalid_context_id_returns_jsonrpc_invalid_params_400()
)
body = json.loads(response.body.decode())
assert response.status_code == 400
assert response.status_code == 200
assert body["id"] == "req-1"
assert body["error"]["code"] == -32602
assert "Invalid AgentCore runtime session id" in body["error"]["message"]
@ -1720,10 +1720,169 @@ async def test_unknown_method_returns_jsonrpc_error():
)
body = json.loads(response.body.decode())
assert response.status_code == 200
assert body["error"]["code"] == -32601
assert "SomeUnknownMethod" in body["error"]["message"]
@pytest.mark.asyncio
async def test_wrong_jsonrpc_version_returns_jsonrpc_error_with_http_200():
"""a2a-sdk clients call raise_for_status() before reading the body, so a JSON-RPC
error behind a 4xx never reaches them as the typed error."""
mock_request = _make_request_mock("message/send", _HELLO_MESSAGE_PARAMS)
mock_request.json = AsyncMock(
return_value={"jsonrpc": "1.0", "id": "req-1", "method": "message/send", "params": _HELLO_MESSAGE_PARAMS}
)
with ExitStack() as stack:
for p in _base_patches(_make_agent_mock()):
stack.enter_context(p)
from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
response = await invoke_agent_a2a(
agent_id="test-agent",
request=mock_request,
fastapi_response=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"),
)
assert response.status_code == 200
assert json.loads(response.body.decode()) == {
"jsonrpc": "2.0",
"id": "req-1",
"error": {"code": -32600, "message": "Invalid Request: jsonrpc must be '2.0'"},
}
_SCRIPTED_AGENT_CARD: Final = {
"protocolVersion": "0.3.0",
"name": "Test Agent",
"url": "http://backend-agent:10001/",
"preferredTransport": "JSONRPC",
"version": "1.0.0",
"capabilities": {"streaming": True},
"defaultInputModes": ["text/plain"],
"defaultOutputModes": ["text/plain"],
"skills": [],
}
_SCRIPTED_AGENT_REPLY: Final = {
"jsonrpc": "2.0",
"id": "",
"result": {"messageId": "reply-1", "role": "agent", "parts": [{"kind": "text", "text": "pong"}], "kind": "message"},
}
class _ScriptedAgent:
"""Answers the card GET and the JSON-RPC POST of a real a2a-sdk dispatch, recording the RPC bodies."""
def __init__(self) -> None:
self.rpc_bodies = []
def __call__(self, request: httpx.Request) -> httpx.Response:
if request.method == "GET":
return httpx.Response(200, json=dict(_SCRIPTED_AGENT_CARD), request=request)
self.rpc_bodies.append(json.loads(request.content))
return httpx.Response(200, json=dict(_SCRIPTED_AGENT_REPLY), request=request)
async def _seed_scripted_agent(monkeypatch: pytest.MonkeyPatch) -> _ScriptedAgent:
"""Put the pooled A2A httpx client behind a scripted upstream, so invoke_agent_a2a runs the real
asend_message, card resolution and response wrapping with no LiteLLM function replaced."""
from litellm.caching.llm_caching_handler import LLMClientCache
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, httpxSpecialProvider
agent: Final = _ScriptedAgent()
handler: Final = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT)
pooled: Final = handler.client
handler.client = httpx.AsyncClient(transport=httpx.MockTransport(agent))
await pooled.aclose()
monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache())
litellm.in_memory_llm_clients_cache.set_cache(
key="async_httpx_client" + f"timeout_{DEFAULT_A2A_AGENT_TIMEOUT}" + httpxSpecialProvider.A2AProvider,
value=handler,
)
return agent
@pytest.mark.asyncio
async def test_message_send_echoes_a_null_request_id(monkeypatch: pytest.MonkeyPatch):
"""JSON-RPC 2.0 requires the response id to be null when the request's was. The SDK request type
cannot carry a null, so the real dispatch sends "" upstream and the proxy must restore the null."""
upstream = await _seed_scripted_agent(monkeypatch)
mock_request = _make_request_mock("message/send", _HELLO_MESSAGE_PARAMS, request_id=None)
mock_proxy_logging = MagicMock()
mock_proxy_logging.pre_call_hook = AsyncMock(
side_effect=lambda user_api_key_dict, data, call_type, skip_guardrails=False: data
)
mock_proxy_logging.post_call_success_hook = AsyncMock(
side_effect=lambda user_api_key_dict, data, response: response
)
with ExitStack() as stack:
for p in _base_patches(_make_agent_mock()):
stack.enter_context(p)
stack.enter_context(
patch( # test-quality-ok: same proxy_logging_obj injection the sibling send tests use; the dispatch itself is real
"litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging
)
)
from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
response = await invoke_agent_a2a(
agent_id="test-agent",
request=mock_request,
fastapi_response=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"),
)
body = json.loads(response.body.decode())
assert response.status_code == 200, body
assert "id" in body and body["id"] is None
assert body["result"]["parts"][0]["text"] == "pong"
assert upstream.rpc_bodies[-1]["method"] == "message/send"
@pytest.mark.asyncio
async def test_task_methods_echo_a_null_request_id():
mock_request = _make_request_mock("tasks/get", {"id": "nope"}, request_id=None)
mock_http_response = MagicMock()
mock_http_response.json.return_value = {
"jsonrpc": "2.0",
"id": None,
"error": {"code": -32001, "message": "Task not found"},
}
mock_http_response.is_success = True
mock_http_response.raise_for_status = MagicMock()
mock_handler = MagicMock()
mock_handler.post = AsyncMock(return_value=mock_http_response)
mock_handler.client = MagicMock()
with ExitStack() as stack:
for p in _base_patches(_make_agent_mock()):
stack.enter_context(p)
stack.enter_context(
patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=mock_handler)
)
from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a
response = await invoke_agent_a2a(
agent_id="test-agent",
request=mock_request,
fastapi_response=MagicMock(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1"),
)
body = json.loads(response.body.decode())
assert response.status_code == 200
assert "id" in body and body["id"] is None
assert body["error"]["code"] == -32001
@pytest.mark.asyncio
@pytest.mark.parametrize(
"pascal_method,expected_wire_method",
@ -3199,6 +3358,7 @@ _EXPECTED_REJECTIONS: Final = (
pytest.param(
True,
403,
403,
"Agent 'denied-agent' is not allowed for your key/team. Contact proxy admin for access.",
"registered-denied-id",
"a2a_agent/Test Agent",
@ -3206,6 +3366,7 @@ _EXPECTED_REJECTIONS: Final = (
),
pytest.param(
False,
200,
404,
{"jsonrpc": "2.0", "id": "rej-1", "error": {"code": -32000, "message": "Agent 'no-such-agent' not found"}},
"no-such-agent",
@ -3216,9 +3377,15 @@ _EXPECTED_REJECTIONS: Final = (
@pytest.mark.asyncio
@pytest.mark.parametrize(("registered", "status", "body", "target", "model"), _EXPECTED_REJECTIONS)
@pytest.mark.parametrize(("registered", "status", "logged_status", "body", "target", "model"), _EXPECTED_REJECTIONS)
async def test_rejected_agent_call_is_logged_once_as_the_caller_without_dispatch(
monkeypatch: pytest.MonkeyPatch, registered: bool, status: int, body: object, target: str, model: str
monkeypatch: pytest.MonkeyPatch,
registered: bool,
status: int,
logged_status: int,
body: object,
target: str,
model: str,
) -> None:
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
@ -3251,16 +3418,17 @@ async def test_rejected_agent_call_is_logged_once_as_the_caller_without_dispatch
assert len(recorder.hook_calls) == 1
hook_data, hook_exception = recorder.hook_calls[0]
assert (hook_data["agent_id"], hook_data["metadata"]["agent_id"]) == (target, target)
assert isinstance(hook_exception, HTTPException) and hook_exception.status_code == status
assert isinstance(hook_exception, HTTPException) and hook_exception.status_code == logged_status
@pytest.mark.asyncio
@pytest.mark.parametrize(("registered", "status", "body", "target", "model"), _EXPECTED_REJECTIONS)
@pytest.mark.parametrize(("registered", "status", "logged_status", "body", "target", "model"), _EXPECTED_REJECTIONS)
async def test_rejected_agent_call_response_survives_a_failing_failure_hook(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
registered: bool,
status: int,
logged_status: int,
body: object,
target: str,
model: str,