diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index dcbe7a6506e..fd7a34e9f95 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-Trace-Id": trace_id, + **({"X-LiteLLM-Agent-Id": agent_id} if agent_id else {}), + **(agent_extra_headers or {}), + } 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 b43f6d9f358..67bfe47a831 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -680,6 +680,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) +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 9b6ec8d28d6..659c41e7840 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -43,6 +43,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, @@ -7836,7 +7837,12 @@ def get_standard_logging_object_payload( else raw_usage_dict ) - id = response_obj.get("id") or kwargs.get("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 response_obj.get("id") or kwargs.get("litellm_call_id"), + ) _model_id: Final = metadata.get("model_info", {}).get("id", "") _model_group: Final = metadata.get("model_group", "") @@ -7874,7 +7880,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, ) @@ -7883,8 +7889,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 f0a4d47f1e4..a50c2a573b3 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -45,7 +45,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.custom_openapi_spec import inline_request_body -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, @@ -464,19 +467,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, ) @@ -559,6 +569,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. @@ -629,6 +641,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 ( @@ -697,7 +711,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( @@ -1074,7 +1088,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": @@ -1092,6 +1107,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 e4492ee8992..ffb7e114142 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/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..eeb6a3f8d61 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 @@ -184,15 +185,17 @@ 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), 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,117 @@ 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 + ) + 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") + 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) + 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 + assert len(set(trace_ids)) == 2 and all(trace_id != "1" for trace_id in trace_ids), trace_ids + + +@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 e8ea0c04f6f..2adefc997bc 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14233,6 +14233,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 b9411eacc8f..2e6ec5b32ba 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5379,6 +5379,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