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

fix(a2a): restore ID passthrough and header precedence

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 20:01:12 +00:00
parent 7e53cd3c84
commit 34bb0a4744
3 changed files with 6 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/a2a_protocol/main.py
View file

@ -781,9 +781,9 @@ async def asend_message_streaming(
if api_base is None:
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
extra_headers: Final = {
"X-LiteLLM-Trace-Id": trace_id,
**({"X-LiteLLM-Agent-Id": agent_id} if agent_id else {}),
**(agent_extra_headers or {}),
"X-LiteLLM-Trace-Id": trace_id,
}
a2a_client = await create_a2a_client(
base_url=api_base,

9
litellm/litellm_core_utils/litellm_logging.py
View file

@ -364,7 +364,6 @@ def _get_provider_request_id(original_exception: Exception) -> str | None:
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
_METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object])
_RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object])
_STANDARD_LOGGING_ID_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
_UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset(
("user_api_key_auth", "user_api_key_budget_reservation")
@ -7744,11 +7743,11 @@ def get_standard_logging_object_payload(
else raw_usage_dict
)
litellm_call_id: Final[str | None] = _STANDARD_LOGGING_ID_ADAPTER.validate_python(kwargs.get("litellm_call_id"))
request_id: Final[str | None] = (
litellm_call_id
request_id: Final[str | None] = cast( # cast-ok: preserve dynamic IDs without runtime validation
str | None,
kwargs.get("litellm_call_id") # pyright: ignore[reportUnknownMemberType] # kwargs is an untyped legacy mapping
if isinstance(call_type, str) and call_type in A2A_CALL_TYPES
else _STANDARD_LOGGING_ID_ADAPTER.validate_python(response_obj.get("id")) or litellm_call_id
else response_obj.get("id") or kwargs.get("litellm_call_id"),
)
_model_id: Final = metadata.get("model_info", {}).get("id", "")

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

@ -185,7 +185,7 @@ def _assert_a2a_calls_are_billed(
lambda: read_rows(
"SELECT request_id, agent_id, spend, api_key, call_type, "
"""to_char("startTime", 'YYYY-MM-DD') AS day FROM "LiteLLM_SpendLogs" """
"WHERE api_key=%s ORDER BY request_id",
"WHERE api_key=%s ORDER BY call_type",
(digest,),
),
lambda rows: len(rows) == len(methods),
@ -371,21 +371,6 @@ def _assert_reused_jsonrpc_id_has_distinct_spend_identity(
assert len(trace_ids) == 2, trace_ids
assert len(set(trace_ids)) == 2 and all(trace_id != "1" for trace_id in trace_ids), trace_ids
if "stream" in method.lower():
stream_call_ids: Final = tuple(
(digest, response.headers.get("x-litellm-call-id"))
for digest, response in zip(digests, responses)
)
matching_rows: Final = tuple(
(call_id, tuple(row for row in rows if row["api_key"] == digest))
for digest, call_id in stream_call_ids
if call_id is not None
)
assert all(
len(matches) == 1 and str(matches[0]["request_id"]) == call_id
for call_id, matches in matching_rows
), matching_rows
@pytest.mark.parametrize(
("version", "method"),