diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 6ac88d3a430..7f1e68fff0b 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -9,6 +9,7 @@ import datetime import uuid from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union +import httpx import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator @@ -512,9 +513,14 @@ async def asend_message_streaming( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - streaming_extra_headers: Optional[Dict[str, str]] = None + # Mirror the non-streaming path: always include trace and agent-id headers + streaming_extra_headers: Dict[str, str] = { + "X-LiteLLM-Trace-Id": str(request.id), + } + if agent_id: + streaming_extra_headers["X-LiteLLM-Agent-Id"] = agent_id if agent_extra_headers: - streaming_extra_headers = dict(agent_extra_headers) + streaming_extra_headers.update(agent_extra_headers) a2a_client = await create_a2a_client( base_url=api_base, extra_headers=streaming_extra_headers ) @@ -631,17 +637,17 @@ async def create_a2a_client( verbose_logger.info(f"Creating A2A client for {base_url}") - # Use LiteLLM's cached httpx client - http_handler = get_async_httpx_client( - llm_provider=httpxSpecialProvider.A2A, - params={"timeout": timeout}, + # Always create a fresh httpx client per A2A call so that per-agent auth + # headers (extra_headers) are never shared across agents or requests. + # Mutating a cached shared client would cause headers from one agent to + # bleed into requests made to a different agent. + httpx_client = httpx.AsyncClient( + timeout=httpx.Timeout(timeout), + headers=extra_headers or {}, ) - httpx_client = http_handler.client - if extra_headers: - httpx_client.headers.update(extra_headers) verbose_proxy_logger.debug( - f"A2A client created with extra_headers={extra_headers}" + f"A2A client created with extra_headers={list(extra_headers.keys())}" ) # Resolve agent card diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 550182f966f..555e9ad0a6d 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -226,12 +226,16 @@ class AgentRegistry: update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) - if agent.get("static_headers") is not None: + if "static_headers" in agent: + headers_value = agent.get("static_headers") update_data["static_headers"] = safe_dumps( - dict(agent.get("static_headers")) # type: ignore + dict(headers_value) if headers_value is not None else {} + ) + if "extra_headers" in agent: + extra_headers_value = agent.get("extra_headers") + update_data["extra_headers"] = ( + extra_headers_value if extra_headers_value is not None else [] ) - if agent.get("extra_headers") is not None: - update_data["extra_headers"] = agent.get("extra_headers") if agent.get("object_permission") is not None: agent_copy = dict(augment_agent) existing_object_permission_id = existing_agent.get(