scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

test(a2a): cover call identity for reused JSON-RPC ids

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Moe Khalil 2026-10-31 18:36:44 +00:00
parent aea0e4f2d4
commit 90381f7ced
7 changed files with 213 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
tests/integration/compatibility/_a2a_peers.py
View file

@ -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,),
)

2
tests/integration/compatibility/test_a2a_protocol_version_bridge.py
View file

@ -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

31
tests/integration/compatibility/test_a2a_protocol_version_override.py
View file

@ -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")

133
tests/integration/spend/test_agent_spend.py
View file

@ -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)

6
tests/unit/a2a_protocol/test_main.py
View file

@ -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",

42
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -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

14
tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -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