scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge 03285fc6f6 into aeede53112

Browse source
This commit is contained in:
moe-berri 2026-11-06 15:51:07 -05:00 • committed by GitHub
commit edcca7440b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 268 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
litellm/a2a_protocol/main.py
View file

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

1
litellm/constants.py
View file

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

14
litellm/litellm_core_utils/litellm_logging.py
View file

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

29
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

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

3
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

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

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

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

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

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

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

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