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