From 90381f7ced355db4719a67233bd047da9d59a234 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 31 Oct 2026 18:36:44 +0000 Subject: [PATCH 1/5] test(a2a): cover call identity for reused JSON-RPC ids Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/compatibility/_a2a_peers.py | 5 +- .../test_a2a_protocol_version_bridge.py | 2 +- .../test_a2a_protocol_version_override.py | 31 ++-- tests/integration/spend/test_agent_spend.py | 133 +++++++++++++++++- tests/unit/a2a_protocol/test_main.py | 6 + .../test_litellm_logging.py | 42 ++++++ .../test_spend_tracking_utils.py | 14 ++ 7 files changed, 213 insertions(+), 20 deletions(-) diff --git a/tests/integration/compatibility/_a2a_peers.py b/tests/integration/compatibility/_a2a_peers.py index 6a34bd78f41..2abce82fe46 100644 --- a/tests/integration/compatibility/_a2a_peers.py +++ b/tests/integration/compatibility/_a2a_peers.py @@ -277,7 +277,8 @@ def sse_frames(text: str) -> tuple[dict[str, JsonValue], ...]: ) -def spend_rows(request_id: str) -> list[dict[str, JsonValue]]: +def spend_rows(agent_id: str) -> list[dict[str, JsonValue]]: return read_rows( - 'SELECT request_id, model, call_type, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,) + 'SELECT request_id, model, call_type, status FROM "LiteLLM_SpendLogs" WHERE agent_id=%s', + (agent_id,), ) diff --git a/tests/integration/compatibility/test_a2a_protocol_version_bridge.py b/tests/integration/compatibility/test_a2a_protocol_version_bridge.py index 6b65f1280c5..320c9b2ffea 100644 --- a/tests/integration/compatibility/test_a2a_protocol_version_bridge.py +++ b/tests/integration/compatibility/test_a2a_protocol_version_bridge.py @@ -53,7 +53,7 @@ def test_completion_bridge_drops_protocol_version_before_the_provider(gateway: G assert sent["model"] == "gpt-4o-mini" and sent["messages"] == [{"role": "user", "content": text}], sent assert sent.get("stream", False) is (method == "message/stream"), sent if method == "message/send": - rows: Final = eventually(lambda: spend_rows(marker), lambda values: len(values) == 1, seconds=70) + rows: Final = eventually(lambda: spend_rows(identity), lambda values: len(values) == 1, seconds=70) assert rows[0]["call_type"] == "asend_message" and rows[0]["status"] == "success", rows diff --git a/tests/integration/compatibility/test_a2a_protocol_version_override.py b/tests/integration/compatibility/test_a2a_protocol_version_override.py index cc36fd58e89..abad85396ef 100644 --- a/tests/integration/compatibility/test_a2a_protocol_version_override.py +++ b/tests/integration/compatibility/test_a2a_protocol_version_override.py @@ -65,8 +65,8 @@ def _assert_v1_wire(sent: tuple[dict[str, JsonValue], ...], marker: str, method: assert message["messageId"] == marker + "-in", sent -def _one_spend_row(request_id: str, call_type: str, status: str) -> None: - rows: Final = eventually(lambda: spend_rows(request_id), lambda values: len(values) == 1, seconds=70) +def _one_spend_row(agent_id: str, call_type: str, status: str) -> None: + rows: Final = eventually(lambda: spend_rows(agent_id), lambda values: len(values) == 1, seconds=70) assert rows[0]["call_type"] == call_type and rows[0]["status"] == status, rows @@ -93,7 +93,7 @@ def test_pinned_langgraph_agent_answers_message_send_in_the_legacy_dialect(gatew identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED) _send_ok(gateway, identity, agent.marker, route) _assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") def test_pinned_langgraph_agent_streams_message_stream_to_the_final_frame(gateway: Gateway) -> None: @@ -113,7 +113,7 @@ def test_pinned_langgraph_agent_streams_message_stream_to_the_final_frame(gatewa assert results[1]["artifactUpdate"]["artifact"]["parts"] == [{"text": "echo: synthetic ping"}] # type: ignore[index] assert results[2]["statusUpdate"]["status"]["state"] == "TASK_STATE_COMPLETED", response.text # type: ignore[index] _assert_legacy_wire(posts(agent.wire), agent.marker, "message/stream") # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message_streaming", "success") + _one_spend_row(identity, "asend_message_streaming", "success") def test_v1_caller_reaches_pinned_legacy_agent_in_its_own_envelope(gateway: Gateway) -> None: @@ -142,7 +142,7 @@ def test_v1_caller_reaches_pinned_legacy_agent_in_its_own_envelope(gateway: Gate assert "kind" not in task and task["status"]["state"] == "TASK_STATE_COMPLETED", response.text assert task["artifacts"][0]["parts"] == [{"text": "echo: synthetic ping"}], response.text _assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") @pytest.mark.parametrize("value", ("0.3.0", 0.3, " 0.3 "), ids=("dotted", "number", "spaced")) @@ -153,7 +153,7 @@ def test_protocol_version_aliases_pin_the_legacy_transport(gateway: Gateway, val ) _send_ok(gateway, identity, agent.marker) _assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") def test_genuine_v1_agent_pinned_one_point_zero_is_unchanged(gateway: Gateway) -> None: @@ -163,7 +163,7 @@ def test_genuine_v1_agent_pinned_one_point_zero_is_unchanged(gateway: Gateway) - ) _send_ok(gateway, identity, agent.marker) _assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") def test_genuine_v1_agent_unpinned_is_unchanged(gateway: Gateway) -> None: @@ -171,7 +171,7 @@ def test_genuine_v1_agent_unpinned_is_unchanged(gateway: Gateway) -> None: identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, None) _send_ok(gateway, identity, agent.marker) _assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") @pytest.mark.parametrize("pin", (None, "0.3"), ids=("unpinned", "pinned-legacy")) @@ -184,7 +184,7 @@ def test_lowercase_binding_still_downgrades_to_the_legacy_transport(gateway: Gat identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, params) _send_ok(gateway, identity, agent.marker) _assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") def test_explicit_one_point_zero_opts_out_of_the_lowercase_downgrade(gateway: Gateway) -> None: @@ -194,7 +194,7 @@ def test_explicit_one_point_zero_opts_out_of_the_lowercase_downgrade(gateway: Ga ) _send_ok(gateway, identity, agent.marker) _assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") def test_one_point_zero_pin_on_a_legacy_card_sends_v1(gateway: Gateway) -> None: @@ -205,7 +205,7 @@ def test_one_point_zero_pin_on_a_legacy_card_sends_v1(gateway: Gateway) -> None: ) _send_ok(gateway, identity, agent.marker) _assert_v1_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") def test_override_rewrites_only_canonical_interfaces(gateway: Gateway) -> None: @@ -219,7 +219,7 @@ def test_override_rewrites_only_canonical_interfaces(gateway: Gateway) -> None: _send_ok(gateway, identity, agent.marker) _assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] assert decoy.wire.drain() == (), "the HTTP+JSON and custom bindings must not be selected" # type: ignore[union-attr] - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") def test_override_added_and_removed_by_patch_takes_effect_on_both_workers(gateway: Gateway) -> None: @@ -267,8 +267,9 @@ def test_repeated_pinned_sends_log_one_row_each(gateway: Gateway) -> None: assert [item["params"]["message"]["messageId"] for item in sent] == [m + "-in" for m in markers], sent # type: ignore[index] assert len({item["id"] for item in sent}) == len(markers), sent assert {item["method"] for item in sent} == {"message/send"}, sent - for marker in markers: - _one_spend_row(marker, "asend_message", "success") + rows: Final = eventually(lambda: spend_rows(identity), lambda values: len(values) == len(markers), seconds=70) + assert all(row["call_type"] == "asend_message" and row["status"] == "success" for row in rows), rows + assert len({str(row["request_id"]) for row in rows}) == len(markers), rows def test_registered_agent_persists_protocol_version(gateway: Gateway) -> None: @@ -295,4 +296,4 @@ def test_pinned_agent_survives_localhost_retry(gateway: Gateway) -> None: traffic ) _assert_legacy_wire(sent, agent.marker) - _one_spend_row(agent.marker, "asend_message", "success") + _one_spend_row(identity, "asend_message", "success") diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index fca848fbb43..0a101fcb0c2 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -7,6 +7,7 @@ import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server +from integration.compatibility._a2a_peers import legacy_send, peer, register, sse_frames from pydantic import BaseModel, JsonValue @@ -190,9 +191,11 @@ def _assert_a2a_calls_are_billed( lambda rows: len(rows) == len(methods), seconds=30, ) - assert [{name: value for name, value in row.items() if name != "day"} for row in logs] == [ + assert [ + {name: value for name, value in row.items() if name not in {"day", "request_id"}} + for row in logs + ] == [ { - "request_id": f"{marker}-{method}", "agent_id": agent, "spend": cost, "api_key": digest, @@ -200,6 +203,9 @@ def _assert_a2a_calls_are_billed( } for method in sorted(methods) ], logs + request_ids: Final = tuple(str(row["request_id"]) for row in logs) + assert len(set(request_ids)) == len(methods), logs + assert all(str(uuid.UUID(request_id)) == request_id for request_id in request_ids), logs total: Final = cost * len(methods) assert eventually( lambda: read_rows('SELECT spend FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent,)), @@ -270,3 +276,126 @@ def test_a2a_send_bills_input_and_output_tokens_at_the_agent_per_token_prices(ga def test_a2a_stream_bills_input_and_output_tokens_at_the_agent_per_token_prices(gateway: Gateway) -> None: _assert_a2a_calls_are_billed(gateway, _PER_TOKEN, _PER_TOKEN_COST, (("message/stream", "/a2a/{agent}"),)) + + +def _assert_reused_jsonrpc_id_has_distinct_spend_identity( + gateway: Gateway, + *, + version: str, + method: str, +) -> None: + marker: Final = "a2aidentity" + uuid.uuid4().hex + dialect: Final = "v1" if version == "1.0" else "langgraph" + params: Final = None if version == "1.0" else {"a2a_protocol_version": "0.3"} + + with peer(marker, dialect) as agent, gateway.scenario() as scenario: + identity: Final = register(gateway, scenario.cleanups.callback, marker, agent.url, params) + user_ids: Final = tuple(scenario.user(user_id=f"{marker}-user-{index}") for index in range(2)) + alias_user_pairs: Final = tuple( + (f"{marker}-key-{index}", user_id) + for index, user_id in enumerate(user_ids) + ) + keys: Final = tuple( + scenario.key(user_id=user_id, key_alias=alias) + for alias, user_id in alias_user_pairs + ) + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + body: Final = ( + legacy_send("1", method=method) + if version == "0.3" + else { + "jsonrpc": "2.0", + "id": "1", + "method": method, + "params": { + "message": { + "messageId": "client-message", + "parts": [{"text": "reused id"}], + "role": "ROLE_USER", + } + }, + } + ) + + responses: Final = tuple( + gateway.request( + "POST", + f"/a2a/{identity}", + body, + key=key, + headers={"a2a-version": version}, + ) + for key in keys + ) + for response in responses: + assert response.status_code == 200, response.text + if "stream" in method.lower(): + frames: Final = sse_frames(response.text) + assert frames and all(frame["id"] == "1" for frame in frames), response.text + assert all("error" not in frame for frame in frames), response.text + else: + response_body: Final = response.json() + assert response_body["id"] == "1" and "error" not in response_body, response.text + + assert agent.wire is not None + upstream_requests: Final = tuple(item for item in agent.wire.drain() if item.method == "POST") + upstream: Final = tuple(json.loads(item.body) for item in upstream_requests) + assert len(upstream) == 2, upstream + trace_ids: Final = tuple( + str(item.headers["x-litellm-trace-id"]) + for item in upstream_requests + ) + + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend.request_id, spend.api_key, spend.call_type, token.key_alias, token.user_id ' + 'FROM "LiteLLM_SpendLogs" AS spend ' + 'JOIN "LiteLLM_VerificationToken" AS token ON spend.api_key=token.token ' + "WHERE spend.api_key=%s OR spend.api_key=%s ORDER BY token.key_alias", + (digests[0], digests[1]), + ), + lambda found: len(found) == 2, + seconds=30, + return_last_on_timeout=True, + ) + assert len(rows) == 2, rows + assert tuple((str(row["key_alias"]), str(row["user_id"])) for row in rows) == tuple( + sorted(alias_user_pairs) + ), rows + request_ids: Final = tuple(str(row["request_id"]) for row in rows) + assert len(set(request_ids)) == 2 and all(request_id != "1" for request_id in request_ids), rows + assert all(str(uuid.UUID(request_id)) == request_id for request_id in request_ids), rows + assert len(trace_ids) == 2, trace_ids + assert len(set(trace_ids)) == 2 and all(trace_id != "1" for trace_id in trace_ids), trace_ids + + if "stream" in method.lower(): + stream_call_ids: Final = tuple( + (digest, response.headers.get("x-litellm-call-id")) + for digest, response in zip(digests, responses) + ) + matching_rows: Final = tuple( + (call_id, tuple(row for row in rows if row["api_key"] == digest)) + for digest, call_id in stream_call_ids + if call_id is not None + ) + assert all( + len(matches) == 1 and str(matches[0]["request_id"]) == call_id + for call_id, matches in matching_rows + ), matching_rows + + +@pytest.mark.parametrize( + ("version", "method"), + ( + ("1.0", "SendMessage"), + ("1.0", "SendStreamingMessage"), + ("0.3", "message/send"), + ("0.3", "message/stream"), + ), +) +def test_a2a_reused_jsonrpc_ids_have_distinct_litellm_spend_ids( + gateway: Gateway, + version: str, + method: str, +) -> None: + _assert_reused_jsonrpc_id_has_distinct_spend_identity(gateway, version=version, method=method) diff --git a/tests/unit/a2a_protocol/test_main.py b/tests/unit/a2a_protocol/test_main.py index c0cabee1127..25cf86d8a7d 100644 --- a/tests/unit/a2a_protocol/test_main.py +++ b/tests/unit/a2a_protocol/test_main.py @@ -159,6 +159,8 @@ def test_streaming_logging_obj_carries_call_type_into_model_call_details(): request=request, agent_name="some-agent", agent_id=None, + litellm_call_id="litellm-call-id", + litellm_trace_id="litellm-trace-id", litellm_params=None, metadata=None, proxy_server_request=None, @@ -166,6 +168,8 @@ def test_streaming_logging_obj_carries_call_type_into_model_call_details(): ) assert logging_obj.model_call_details["call_type"] == "asend_message_streaming" + assert logging_obj.litellm_call_id == "litellm-call-id" + assert logging_obj.litellm_trace_id == "litellm-trace-id" _AGENT_CARD = { @@ -780,6 +784,8 @@ def test_streaming_logging_obj_keeps_agent_credentials_out_of_logging_params(): request=request, agent_name="foundry-agent", agent_id="agent-1", + litellm_call_id="litellm-call-id", + litellm_trace_id="litellm-trace-id", litellm_params={ "client_secret": "sp-secret", "azure_ad_token": "entra-token", diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 93a0b0735da..cd0d564b54a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14161,6 +14161,48 @@ def test_a_websocket_turn_row_is_keyed_by_the_response_id_or_else_its_call_id( assert payload["id"] == expected_row_id +@pytest.mark.parametrize("call_type", ["asend_message", "send_message", "asend_message_streaming"]) +def test_a2a_standard_logging_payload_uses_the_litellm_call_id( + logging_obj: LitellmLogging, + call_type: str, +) -> None: + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + payload: Final = get_standard_logging_object_payload( + kwargs={ + "call_type": call_type, + "litellm_call_id": "litellm-call-id", + "litellm_params": {"metadata": {}}, + }, + init_response_obj=ModelResponse(id="1", model="a2a_agent/test-agent"), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None and payload["id"] == "litellm-call-id" + + +def test_completion_standard_logging_payload_keeps_the_response_id( + logging_obj: LitellmLogging, +) -> None: + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + payload: Final = get_standard_logging_object_payload( + kwargs={ + "call_type": "acompletion", + "litellm_call_id": "litellm-call-id", + "litellm_params": {"metadata": {}}, + }, + init_response_obj=ModelResponse(id="chatcmpl-x", model="gpt-4o-mini"), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None and payload["id"] == "chatcmpl-x" + + @pytest.mark.parametrize("for_async_request, expected_s3_writes", [(True, 1), (False, 0)]) def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_async_side( logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch, for_async_request: bool, expected_s3_writes: int diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index f52f40ad55d..5be3c08e911 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5186,6 +5186,20 @@ def test_get_spend_logs_id_prefers_the_response_id_over_the_standard_logging_id( ) +def test_get_spend_logs_id_uses_litellm_call_id_for_a2a_responses(): + assert ( + get_spend_logs_id( + "asend_message", + {"id": "1"}, + { + "litellm_call_id": "litellm-call-id", + "standard_logging_object": {"id": "litellm-call-id"}, + }, + ) + == "litellm-call-id" + ) + + @pytest.mark.asyncio async def test_spend_log_request_id_is_the_message_id_a_bridged_streaming_caller_was_streamed(): """A streaming /v1/messages call against a non-Anthropic model is served a msg_ id the From 7e53cd3c848c2d145c2de6a9d2f224fb27d73f11 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 31 Oct 2026 19:35:09 +0000 Subject: [PATCH 2/5] fix(a2a): use LiteLLM-owned call IDs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/a2a_protocol/main.py | 38 +++++++++++++------ litellm/constants.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 15 ++++++-- .../proxy/agent_endpoints/a2a_endpoints.py | 6 +++ .../spend_tracking/spend_tracking_utils.py | 3 +- tests/integration/spend/test_agent_spend.py | 21 +++++----- 6 files changed, 59 insertions(+), 25 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index dcbe7a6506e..8e13cc1b8a6 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -608,6 +608,8 @@ def _build_streaming_logging_obj( request: "SendStreamingMessageRequest", agent_name: str, agent_id: str | None, + litellm_call_id: str, + litellm_trace_id: str, litellm_params: dict[str, object] | None, metadata: dict[str, object] | None, proxy_server_request: dict[str, object] | None, @@ -622,8 +624,9 @@ def _build_streaming_logging_obj( stream=False, call_type="asend_message_streaming", start_time=start_time, - litellm_call_id=str(request.id), - function_id=str(request.id), + litellm_call_id=litellm_call_id, + litellm_trace_id=litellm_trace_id, + function_id=litellm_call_id, ) logging_obj.model = model logging_obj.custom_llm_provider = "a2a_agent" @@ -656,6 +659,8 @@ async def asend_message_streaming( proxy_server_request: dict[str, object] | None = None, agent_extra_headers: dict[str, str] | None = None, agent_name: str | None = None, + litellm_call_id: str | None = None, + litellm_trace_id: str | None = None, **kwargs: object, ) -> AsyncIterator["SendStreamingMessageResponse | Mapping[str, object]"]: """ @@ -725,13 +730,22 @@ async def asend_message_streaming( agent_extra_headers=agent_extra_headers, ) caller_logging_obj: Final = kwargs.get("litellm_logging_obj") + bridge_logging_obj_from_caller: Final = caller_logging_obj if isinstance(caller_logging_obj, Logging) else None + bridge_call_id: Final = litellm_call_id or str(uuid.uuid4()) + bridge_trace_id: Final = ( + (bridge_logging_obj_from_caller.litellm_trace_id if bridge_logging_obj_from_caller else None) + or litellm_trace_id + or str(uuid.uuid4()) + ) bridge_logging_obj: Final = ( - caller_logging_obj - if isinstance(caller_logging_obj, Logging) + bridge_logging_obj_from_caller + if bridge_logging_obj_from_caller is not None else _build_streaming_logging_obj( request=request, agent_name=agent_name or "unknown", agent_id=agent_id, + litellm_call_id=bridge_call_id, + litellm_trace_id=bridge_trace_id, litellm_params=litellm_params, metadata=metadata, proxy_server_request=proxy_server_request, @@ -760,17 +774,17 @@ async def asend_message_streaming( _raw_logging_obj: Final = kwargs.get("litellm_logging_obj") logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None + call_id: Final = litellm_call_id or str(uuid.uuid4()) + trace_id: Final = (logging_obj.litellm_trace_id if logging_obj else None) or litellm_trace_id or str(uuid.uuid4()) if a2a_client is None: if api_base is None: raise ValueError("Either a2a_client or api_base is required for standard A2A flow") - logging_trace_id: Final = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None - trace_id: Final = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4())) - extra_headers: Final[dict[str, str]] = {"X-LiteLLM-Trace-Id": trace_id} - if agent_id: - extra_headers["X-LiteLLM-Agent-Id"] = agent_id - if agent_extra_headers: - extra_headers.update(agent_extra_headers) + extra_headers: Final = { + **({"X-LiteLLM-Agent-Id": agent_id} if agent_id else {}), + **(agent_extra_headers or {}), + "X-LiteLLM-Trace-Id": trace_id, + } a2a_client = await create_a2a_client( base_url=api_base, extra_headers=extra_headers, @@ -788,6 +802,8 @@ async def asend_message_streaming( request=request, agent_name=card_agent_name, agent_id=agent_id, + litellm_call_id=call_id, + litellm_trace_id=trace_id, litellm_params=litellm_params, metadata=metadata, proxy_server_request=proxy_server_request, diff --git a/litellm/constants.py b/litellm/constants.py index e21751c7fd1..e412284512c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -676,6 +676,7 @@ request_timeout: float = ( ) request_timeout_explicitly_set: bool = _request_timeout_from_env is not None DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = get_env_float("DEFAULT_A2A_AGENT_TIMEOUT", 6000.0) # 10 minutes +A2A_CALL_TYPES: Final = frozenset({"asend_message", "send_message", "asend_message_streaming"}) AGENT_KILL_SWITCH_TIMEOUT_SECONDS: Final = 10.0 AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS: Final = 2000 # Patterns that indicate a localhost/internal URL in A2A agent cards that should be diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b6574a7787d..d79d67df384 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -41,6 +41,7 @@ from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import ( + A2A_CALL_TYPES, DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, @@ -363,6 +364,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object]) _RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_STANDARD_LOGGING_ID_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") @@ -7742,7 +7744,12 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = response_obj.get("id") or kwargs.get("litellm_call_id") + litellm_call_id: Final[str | None] = _STANDARD_LOGGING_ID_ADAPTER.validate_python(kwargs.get("litellm_call_id")) + request_id: Final[str | None] = ( + litellm_call_id + if isinstance(call_type, str) and call_type in A2A_CALL_TYPES + else _STANDARD_LOGGING_ID_ADAPTER.validate_python(response_obj.get("id")) or litellm_call_id + ) _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") @@ -7780,7 +7787,7 @@ def get_standard_logging_object_payload( usage_object=usage_dict, proxy_server_request=proxy_server_request, start_time=start_time, - response_id=id, + response_id=request_id, custom_llm_provider=selected_provider if isinstance(selected_provider, str) else None, is_silent_experiment=logging_obj.is_silent_experiment is True, ) @@ -7789,8 +7796,8 @@ def get_standard_logging_object_payload( "user", None ) # maintain backwards compatibility with old request body check - if cache_hit is True or served_from_cache(init_response_obj): - id = f"{id}_cache_hit{time.time()}" # do not duplicate the request id + is_cached: Final = cache_hit is True or served_from_cache(init_response_obj) + id: Final[str | None] = f"{request_id}_cache_hit{time.time()}" if is_cached else request_id uncached_price: Final = logging_obj.uncached_price(result=init_response_obj) if cache_hit is True else None saved_cache_cost: Final = (uncached_price.cost if uncached_price is not None else None) or 0.0 diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 5e1d861e7e9..19b84b4b5be 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -436,6 +436,8 @@ async def _handle_stream_message( request_data: dict[str, object] | None = None, proxy_logging_obj: ProxyLogging | None = None, served_version: A2AVersion = "0.3", + litellm_call_id: str | None = None, + litellm_trace_id: str | None = None, ) -> StreamingResponse: """Handle message/stream method via SDK functions. @@ -506,6 +508,8 @@ async def _handle_stream_message( proxy_server_request=proxy_server_request, agent_extra_headers=agent_extra_headers, agent_name=agent_name, + litellm_call_id=litellm_call_id, + litellm_trace_id=litellm_trace_id, ) if ( @@ -939,6 +943,8 @@ async def invoke_agent_a2a( request_data=data, proxy_logging_obj=proxy_logging_obj, served_version=served_version, + litellm_call_id=logging_obj.litellm_call_id, + litellm_trace_id=logging_obj.litellm_trace_id, ) elif method in { "tasks/get", diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index ba5f873c4e9..781d9298b23 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -16,6 +16,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + A2A_CALL_TYPES, BILLED_RETRY_ATTEMPT_METADATA_KEY, CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, @@ -368,7 +369,7 @@ def get_spend_logs_id( ) -> str | None: standard_logging_payload = kwargs.get("standard_logging_object") candidate_ids: Final = ( - response_obj.get("id"), + None if call_type in A2A_CALL_TYPES else response_obj.get("id"), standard_logging_payload.get("id") if isinstance(standard_logging_payload, dict) else None, kwargs.get("litellm_call_id"), ) diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index 0a101fcb0c2..635e99531af 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -327,15 +327,18 @@ def _assert_reused_jsonrpc_id_has_distinct_spend_identity( ) for key in keys ) - for response in responses: - assert response.status_code == 200, response.text - if "stream" in method.lower(): - frames: Final = sse_frames(response.text) - assert frames and all(frame["id"] == "1" for frame in frames), response.text - assert all("error" not in frame for frame in frames), response.text - else: - response_body: Final = response.json() - assert response_body["id"] == "1" and "error" not in response_body, response.text + assert all(response.status_code == 200 for response in responses), responses + if "stream" in method.lower(): + frames_by_response: Final = tuple(sse_frames(response.text) for response in responses) + assert all( + frames and all(frame["id"] == "1" for frame in frames) and all("error" not in frame for frame in frames) + for frames in frames_by_response + ), responses + else: + response_bodies: Final = tuple(response.json() for response in responses) + assert all( + response_body["id"] == "1" and "error" not in response_body for response_body in response_bodies + ), responses assert agent.wire is not None upstream_requests: Final = tuple(item for item in agent.wire.drain() if item.method == "POST") From 34bb0a47443c4fd9d9c3d60ec6db553a044b4497 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 31 Oct 2026 20:01:12 +0000 Subject: [PATCH 3/5] fix(a2a): restore ID passthrough and header precedence Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/a2a_protocol/main.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 9 ++++----- tests/integration/spend/test_agent_spend.py | 17 +---------------- 3 files changed, 6 insertions(+), 22 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 8e13cc1b8a6..fd7a34e9f95 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -781,9 +781,9 @@ async def asend_message_streaming( if api_base is None: raise ValueError("Either a2a_client or api_base is required for standard A2A flow") extra_headers: Final = { + "X-LiteLLM-Trace-Id": trace_id, **({"X-LiteLLM-Agent-Id": agent_id} if agent_id else {}), **(agent_extra_headers or {}), - "X-LiteLLM-Trace-Id": trace_id, } a2a_client = await create_a2a_client( base_url=api_base, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d79d67df384..ac5f59d72b6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -364,7 +364,6 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object]) _RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object]) -_STANDARD_LOGGING_ID_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") @@ -7744,11 +7743,11 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - litellm_call_id: Final[str | None] = _STANDARD_LOGGING_ID_ADAPTER.validate_python(kwargs.get("litellm_call_id")) - request_id: Final[str | None] = ( - litellm_call_id + request_id: Final[str | None] = cast( # cast-ok: preserve dynamic IDs without runtime validation + str | None, + kwargs.get("litellm_call_id") # pyright: ignore[reportUnknownMemberType] # kwargs is an untyped legacy mapping if isinstance(call_type, str) and call_type in A2A_CALL_TYPES - else _STANDARD_LOGGING_ID_ADAPTER.validate_python(response_obj.get("id")) or litellm_call_id + else response_obj.get("id") or kwargs.get("litellm_call_id"), ) _model_id: Final = metadata.get("model_info", {}).get("id", "") diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index 635e99531af..cd92d0089ff 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -185,7 +185,7 @@ def _assert_a2a_calls_are_billed( lambda: read_rows( "SELECT request_id, agent_id, spend, api_key, call_type, " """to_char("startTime", 'YYYY-MM-DD') AS day FROM "LiteLLM_SpendLogs" """ - "WHERE api_key=%s ORDER BY request_id", + "WHERE api_key=%s ORDER BY call_type", (digest,), ), lambda rows: len(rows) == len(methods), @@ -371,21 +371,6 @@ def _assert_reused_jsonrpc_id_has_distinct_spend_identity( assert len(trace_ids) == 2, trace_ids assert len(set(trace_ids)) == 2 and all(trace_id != "1" for trace_id in trace_ids), trace_ids - if "stream" in method.lower(): - stream_call_ids: Final = tuple( - (digest, response.headers.get("x-litellm-call-id")) - for digest, response in zip(digests, responses) - ) - matching_rows: Final = tuple( - (call_id, tuple(row for row in rows if row["api_key"] == digest)) - for digest, call_id in stream_call_ids - if call_id is not None - ) - assert all( - len(matches) == 1 and str(matches[0]["request_id"]) == call_id - for call_id, matches in matching_rows - ), matching_rows - @pytest.mark.parametrize( ("version", "method"), From 54e193d1a1d8897436bd23e1fb0912a57ed269b5 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 31 Oct 2026 20:50:26 +0000 Subject: [PATCH 4/5] fix(a2a): return call IDs in successful response headers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/agent_endpoints/a2a_endpoints.py | 23 ++++++++++++----- .../compatibility/test_a2a_wire_versions.py | 25 ++++++++++++++++--- tests/integration/spend/test_agent_spend.py | 3 +++ 3 files changed, 41 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 22cacc83c04..642ccbff904 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -43,7 +43,10 @@ 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.openai_error_payload import ( + headers_with_litellm_call_id, + status_typed_proxy_exception, +) from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING, coerce_keepalive_interval, @@ -397,19 +400,26 @@ async def _a2a_sse_event_source( await resp.aclose() -def _sse_streaming_response(generator: AsyncGenerator[str, None]) -> StreamingResponse: +def _sse_streaming_response( + generator: AsyncGenerator[str, None], + litellm_call_id: str | None = None, +) -> StreamingResponse: # The upstream agent is only contacted once this generator is first pulled, so # a slow first event leaves the response body idle for its whole # time-to-first-token and an intermediary with an idle read timeout drops a # healthy connection. Off until an operator sets an interval, and the # buffering hint only goes out when there are keepalives to protect. keepalive_interval: Final = coerce_keepalive_interval(litellm.sse_keepalive_ping_interval_seconds) + base_headers: Final[Mapping[str, str] | None] = _SSE_KEEPALIVE_HEADERS if keepalive_interval is not None else None + response_headers: Final = ( + headers_with_litellm_call_id(base_headers, litellm_call_id) if litellm_call_id is not None else base_headers + ) if keepalive_interval is None: - return StreamingResponse(generator, media_type="text/event-stream") + return StreamingResponse(generator, media_type="text/event-stream", headers=response_headers) return StreamingResponse( wrap_sse_stream_with_keepalive_pings(generator, keepalive_interval, ping_chunk=SSE_COMMENT_PING), media_type="text/event-stream", - headers=_SSE_KEEPALIVE_HEADERS, + headers=response_headers, ) @@ -634,7 +644,7 @@ async def _handle_stream_message( } ) - return _sse_streaming_response(stream_response()) + return _sse_streaming_response(stream_response(), litellm_call_id=litellm_call_id) @router.get( @@ -982,7 +992,8 @@ async def invoke_agent_a2a( response_dict, served_version, method="message/send", - ) + ), + headers=headers_with_litellm_call_id(None, logging_obj.litellm_call_id), ) elif method == "message/stream": diff --git a/tests/integration/compatibility/test_a2a_wire_versions.py b/tests/integration/compatibility/test_a2a_wire_versions.py index 981700080e4..934bf89fd7a 100644 --- a/tests/integration/compatibility/test_a2a_wire_versions.py +++ b/tests/integration/compatibility/test_a2a_wire_versions.py @@ -965,9 +965,18 @@ def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated {"taskId": "t1", "url": callback}, ): rejected = call("tasks/pushNotificationConfig/set", params) - assert (rejected.status_code, rejected.json()) == (400, {"detail": detail}), ( - f"{params}: {rejected.text}" - ) + assert (rejected.status_code, rejected.json()) == ( + 400, + { + "jsonrpc": "2.0", + "id": marker + "tasks/pushNotificationConfig/set", + "error": { + "code": -32000, + "message": detail, + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, + ), f"{params}: {rejected.text}" assert tuple(item for item in wire.drain() if item.method == "POST") == () public: Final = {"taskId": "t1", "pushNotificationConfig": {"url": "https://1.1.1.1/callback"}} @@ -1060,7 +1069,15 @@ def test_a2a_v1_task_methods_and_extended_card_lower_to_0_3_upstream_and_answer_ rejected = call("CreateTaskPushNotificationConfig", params) assert (rejected.status_code, rejected.json()) == ( 400, - {"detail": "Push notification URL must use HTTPS"}, + { + "jsonrpc": "2.0", + "id": marker + "CreateTaskPushNotificationConfig", + "error": { + "code": -32000, + "message": "Push notification URL must use HTTPS", + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, ), f"{params}: {rejected.text}" assert forwarded() == () diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index cd92d0089ff..eeb6a3f8d61 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -366,6 +366,9 @@ def _assert_reused_jsonrpc_id_has_distinct_spend_identity( sorted(alias_user_pairs) ), rows request_ids: Final = tuple(str(row["request_id"]) for row in rows) + response_call_ids: Final = tuple(response.headers.get("x-litellm-call-id") for response in responses) + assert all(response_call_id is not None for response_call_id in response_call_ids), responses + assert response_call_ids == request_ids, (response_call_ids, rows) assert len(set(request_ids)) == 2 and all(request_id != "1" for request_id in request_ids), rows assert all(str(uuid.UUID(request_id)) == request_id for request_id in request_ids), rows assert len(trace_ids) == 2, trace_ids From 03285fc6f68bebd8f456576d9ed8e22cbc7ce73b Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Sat, 31 Oct 2026 23:16:34 +0000 Subject: [PATCH 5/5] style(a2a): sort a2a_endpoints imports after main merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/a2a_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 9e803937a0e..0614cb2d6e1 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -43,11 +43,11 @@ 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.custom_openapi_spec import inline_request_body from litellm.proxy.common_utils.openai_error_payload import ( headers_with_litellm_call_id, status_typed_proxy_exception, ) -from litellm.proxy.common_utils.custom_openapi_spec import inline_request_body from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING, coerce_keepalive_interval,