diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 95562fcae8c..7e499dde642 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,7 +1,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal -from pydantic import BaseModel, PrivateAttr +from pydantic import BaseModel, PrivateAttr, StrictInt from typing_extensions import Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -315,10 +315,23 @@ def _normalize_a2a_jsonrpc_response( 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. + + JSON-RPC 2.0 requires the response id to equal the request id, so a string or + integer request id is carried over as-is. Anything else is stringified, which + is the only representation the response model accepts. + + A caller that supplied no id leaves the response id null, which is what the + spec requires for an error that cannot be correlated to a request. ``bool`` counts + as "anything else" despite subclassing ``int``, so ``true`` is never relayed as + ``1``, where it would collide with a real integer id. """ normalized: Final = dict(response_dict) - if normalized.get("id") is None and request_id is not None: - normalized["id"] = str(request_id) + if isinstance(normalized.get("id"), bool): + normalized["id"] = str(normalized["id"]) + elif normalized.get("id") is None and request_id is not None: + normalized["id"] = ( + request_id if isinstance(request_id, (str, int)) and not isinstance(request_id, bool) else str(request_id) + ) return normalized @@ -331,7 +344,7 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): """ # A2A response fields - id: str + id: str | StrictInt | None = None jsonrpc: str = "2.0" result: dict[str, Any] | None = None error: dict[str, Any] | None = None @@ -360,8 +373,9 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): Returns: LiteLLMSendMessageResponse with _hidden_params support """ - response_dict = response.model_dump(mode="json", exclude_none=True) - response_dict = _normalize_a2a_jsonrpc_response(response_dict, request_id=request_id) + response_dict: Final = _normalize_a2a_jsonrpc_response( + response.model_dump(mode="json", exclude_none=True), request_id=request_id + ) return cls(**response_dict) @classmethod diff --git a/tests/test_litellm/a2a_protocol/test_send_message_response.py b/tests/test_litellm/a2a_protocol/test_send_message_response.py index 832aa288c7a..ade7c72fc2e 100644 --- a/tests/test_litellm/a2a_protocol/test_send_message_response.py +++ b/tests/test_litellm/a2a_protocol/test_send_message_response.py @@ -32,12 +32,102 @@ def test_from_dict_preserves_existing_id(): 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") +def test_from_dict_preserves_integer_id_echoed_by_upstream(): + """JSON-RPC 2.0 types ``id`` as string|integer|null, and pydantic v2 does not + coerce int to str, so a str-only annotation rejects an upstream agent that + echoes an integer id. The value AND the type must survive.""" + payload = { + "id": 42, + "jsonrpc": "2.0", + "result": {"kind": "task"}, + } + + response = LiteLLMSendMessageResponse.from_dict(payload, request_id="r1") + + assert response.id == 42 + assert isinstance(response.id, int) + + +def test_from_dict_preserves_falsy_integer_id(): + """``0`` is a legal JSON-RPC id and is falsy, so it must not be mistaken for an + absent id and backfilled from the request id.""" + payload = {"id": 0, "jsonrpc": "2.0", "result": {}} + + response = LiteLLMSendMessageResponse.from_dict(payload, request_id="r1") + + assert response.id == 0 + + +def test_backfilled_id_keeps_the_request_id_type(): + """The proxy's A2A endpoint reads the caller's ``id`` straight off the request + body, so it can be an integer. JSON-RPC requires the response id to equal the + request id, so backfilling an omitted id must not stringify it: a caller that + sent ``7`` cannot correlate a response carrying ``"7"``. One test, both + directions, so neither can regress unnoticed.""" + agent_error = { + "jsonrpc": "2.0", + "error": {"code": -32054, "message": "Session not found"}, + } + + from_int = LiteLLMSendMessageResponse.from_dict(agent_error, request_id=7) + from_str = LiteLLMSendMessageResponse.from_dict(agent_error, request_id="7") + + assert from_int.id == 7 + assert isinstance(from_int.id, int) + assert from_str.id == "7" + assert isinstance(from_str.id, str) + + +def test_from_dict_accepts_null_id_when_the_error_cannot_be_correlated(): + """JSON-RPC 2.0 section 5 requires ``id`` to be null on an error that cannot be + matched to a request, which is exactly the case where the caller supplied no id + for the backfill to use. Rejecting it turned an agent's error into a proxy 500.""" + response = LiteLLMSendMessageResponse.from_dict( + {"jsonrpc": "2.0", "error": {"code": -32054, "message": "x"}} + ) + + assert response.id is None + assert response.error == {"code": -32054, "message": "x"} + + +def test_from_dict_accepts_null_id_echoed_by_upstream(): + """An agent may answer an uncorrelatable request with an explicit ``"id": null``. + That is a well-formed response, not a validation failure.""" + response = LiteLLMSendMessageResponse.from_dict( + {"id": None, "jsonrpc": "2.0", "error": {"code": -32600, "message": "bad"}} + ) + + assert response.id is None + + +def test_id_accepts_every_member_of_the_json_rpc_union_and_nothing_else(): + """One test pinning the whole ``string | integer | null`` union the spec defines, + so widening the annotation cannot silently become "accept anything".""" + for accepted in ("s1", 42, 0, None): + assert LiteLLMSendMessageResponse(id=accepted).id == accepted + + # ``True``/``False`` are in here because bool subclasses int: a non-strict integer + # half would accept them and relay them as 1/0. Direct construction bypasses + # normalization, so the model has to hold this line on its own. + for rejected in (True, False, 1.5, ["a"], {"a": 1}): + try: + LiteLLMSendMessageResponse(id=rejected) + except Exception: + continue + raise AssertionError(f"id={rejected!r} is outside the JSON-RPC union and must be rejected") + + +def test_boolean_id_is_never_relayed_as_an_integer(): + """``bool`` subclasses ``int``, so widening the annotation to accept integers also + made pydantic coerce a boolean id to 1 or 0. That is worse than rejecting it: an id + of ``1`` collides with a real integer id another in-flight request may be using. + Both directions in one test, since either alone leaves the other free to regress.""" + agent_error = {"jsonrpc": "2.0", "error": {"code": -32054, "message": "x"}} + + echoed = LiteLLMSendMessageResponse.from_dict({"id": True, "jsonrpc": "2.0", "result": {}}) + backfilled = LiteLLMSendMessageResponse.from_dict(agent_error, request_id=True) + + assert echoed.id == "True" + assert backfilled.id == "True" + assert not isinstance(echoed.id, int) + assert not isinstance(backfilled.id, int)