diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 22cacc83c04..642ccbff904 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -43,7 +43,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.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, @@ -397,19 +400,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, ) @@ -634,7 +644,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( @@ -982,7 +992,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": diff --git a/tests/integration/compatibility/test_a2a_wire_versions.py b/tests/integration/compatibility/test_a2a_wire_versions.py index 981700080e4..934bf89fd7a 100644 --- a/tests/integration/compatibility/test_a2a_wire_versions.py +++ b/tests/integration/compatibility/test_a2a_wire_versions.py @@ -965,9 +965,18 @@ def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated {"taskId": "t1", "url": callback}, ): rejected = call("tasks/pushNotificationConfig/set", params) - assert (rejected.status_code, rejected.json()) == (400, {"detail": detail}), ( - f"{params}: {rejected.text}" - ) + assert (rejected.status_code, rejected.json()) == ( + 400, + { + "jsonrpc": "2.0", + "id": marker + "tasks/pushNotificationConfig/set", + "error": { + "code": -32000, + "message": detail, + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, + ), f"{params}: {rejected.text}" assert tuple(item for item in wire.drain() if item.method == "POST") == () public: Final = {"taskId": "t1", "pushNotificationConfig": {"url": "https://1.1.1.1/callback"}} @@ -1060,7 +1069,15 @@ def test_a2a_v1_task_methods_and_extended_card_lower_to_0_3_upstream_and_answer_ rejected = call("CreateTaskPushNotificationConfig", params) assert (rejected.status_code, rejected.json()) == ( 400, - {"detail": "Push notification URL must use HTTPS"}, + { + "jsonrpc": "2.0", + "id": marker + "CreateTaskPushNotificationConfig", + "error": { + "code": -32000, + "message": "Push notification URL must use HTTPS", + "data": {"type": "invalid_request_error", "param": None, "code": "400"}, + }, + }, ), f"{params}: {rejected.text}" assert forwarded() == () diff --git a/tests/integration/spend/test_agent_spend.py b/tests/integration/spend/test_agent_spend.py index cd92d0089ff..eeb6a3f8d61 100644 --- a/tests/integration/spend/test_agent_spend.py +++ b/tests/integration/spend/test_agent_spend.py @@ -366,6 +366,9 @@ def _assert_reused_jsonrpc_id_has_distinct_spend_identity( 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