fix(a2a): stop writing per-caller state onto the shared cached httpx client (#35978)

create_a2a_client took the raw client off a process-wide cached handler and
called headers.update() on it, then leaned on folding the header set into the
cache key (through the unrelated disable_aiohttp_transport field) to keep one
caller's credentials away from the next.

Per-caller headers now ride with each request through the a2a SDK's call
context, and the agent card fetch gets them through resolver_http_kwargs, so
the shared client is never written to and its cache key no longer varies by
header set. Since the proxy puts a fresh trace id in every request's headers,
that key previously changed on every call, giving each request its own httpx
client and flushing the 200-entry client cache that every other provider
shares. All A2A callers on one timeout now reuse a single pooled client.

Sharing that client also means sharing its httpx cookie jar, which httpx fills
from every Set-Cookie and replays on any later request to a matching domain, so
one agent's session cookie would arrive at another agent on the same host. The
pooled client now carries a cookie policy that stores and sends nothing, which
neither litellm nor the a2a SDK relies on: the SDK's auth interceptor skips
cookie-borne API keys outright.

litellm/a2a_protocol/exception_mapping_utils.py

@@ -202,9 +202,9 @@ async def handle_a2a_localhost_retry(
     # Fix the agent card URL
     set_agent_card_url(agent_card, error.base_url)
 
-    # Reuse the httpx client LiteLLM attached at creation. It carries this agent's
-    # trace-id and auth headers, so a fresh client would drop them. Only clients built
-    # by ``create_a2a_client`` have it; an externally-supplied client cannot be retried.
+    # Reuse the httpx client and call context LiteLLM attached at creation, since the
+    # context carries this agent's trace-id/auth headers. Only clients built by
+    # ``create_a2a_client`` have them; an externally-supplied client cannot be retried.
     httpx_client: Final = getattr(a2a_client, "_litellm_httpx_client", None)
     if httpx_client is None:
         raise RuntimeError(
@@ -220,5 +220,8 @@ async def handle_a2a_localhost_retry(
         ),
     )
     new_client._litellm_httpx_client = httpx_client
+    new_client._litellm_call_context = getattr(  # pyright: ignore[reportAttributeAccessIssue]  # LiteLLM-owned stash
+        a2a_client, "_litellm_call_context", None
+    )
     new_client._litellm_agent_card = agent_card
     return new_client

litellm/a2a_protocol/main.py

@@ -13,6 +13,7 @@ import asyncio
 import datetime
 import uuid
 from collections.abc import AsyncIterator, Coroutine
+from http.cookiejar import DefaultCookiePolicy
 from typing import TYPE_CHECKING, Any, Final, Optional, cast
 
 import litellm
@@ -30,6 +31,7 @@ from litellm.utils import client
 
 if TYPE_CHECKING:
     from a2a.client import Client as A2AClientType
+    from a2a.client import ClientCallContext as A2ACallContextType
     from a2a.compat.v0_3.types import (
         AgentCard,
         Message,
@@ -45,7 +47,7 @@ A2A_SDK_AVAILABLE = False
 _a2a_conversions: Any = None
 
 try:
-    from a2a.client import Client, ClientConfig, create_client
+    from a2a.client import Client, ClientCallContext, ClientConfig, create_client
     from a2a.compat.v0_3 import conversions as _a2a_conversions
     from a2a.compat.v0_3.types import (
         Message,
@@ -60,6 +62,7 @@ try:
     A2A_SDK_AVAILABLE = True
 except ImportError:
     Client = None
+    ClientCallContext = None
     ClientConfig = None
     create_client = None
 
@@ -77,6 +80,8 @@ from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
 # Use our custom resolver instead of the default A2A SDK resolver
 A2ACardResolver: Final = LiteLLMA2ACardResolver
 
+_BLOCK_ALL_COOKIES: Final = DefaultCookiePolicy(allowed_domains=())
+
 
 def _set_usage_on_logging_obj(
     kwargs: dict[str, Any],
@@ -218,6 +223,10 @@ async def _send_message_via_completion_bridge(
     return LiteLLMSendMessageResponse.from_dict(response_dict, request_id=str(request.id))
 
 
+def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallContextType"]:
+    return getattr(a2a_client, "_litellm_call_context", None)
+
+
 async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse":
     """Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response."""
     if _a2a_conversions is None:
@@ -227,7 +236,7 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques
 
     pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
     last_event = None
-    async for event in a2a_client.send_message(pb_request):
+    async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
         last_event = event
     if last_event is None:
         raise RuntimeError("A2A send_message failed: no response received from agent.")
@@ -301,7 +310,7 @@ async def _stream_messages(
         )
 
     pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
-    async for event in a2a_client.send_message(pb_request):
+    async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
         compat_chunk = _a2a_conversions.to_compat_stream_response(
             event,
             request_id=request.id,
@@ -756,26 +765,13 @@ async def create_a2a_client(
 
     verbose_logger.info("Creating A2A client for %s", base_url)
 
-    # Use get_async_httpx_client with per-agent params so that different agents
-    # (with different extra_headers) get separate cached clients.  The params
-    # dict is hashed into the cache key, keeping agent auth isolated while
-    # still reusing connections within the same agent.
-    #
-    # Only pass params that AsyncHTTPHandler.__init__ accepts (e.g. timeout).
-    # Use "disable_aiohttp_transport" key for cache-key-only data (it's
-    # filtered out before reaching the constructor).
-    _client_params: Final[dict] = {"timeout": timeout}
-    if extra_headers:
-        # Encode headers into a cache-key-only param so each unique header
-        # set produces a distinct cache key.
-        _client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items()))
     _async_handler: Final = get_async_httpx_client(
         llm_provider=httpxSpecialProvider.A2AProvider,
-        params=_client_params,
+        params={"timeout": timeout},
     )
     httpx_client: Final = _async_handler.client
+    httpx_client.cookies.jar.set_policy(_BLOCK_ALL_COOKIES)
     if extra_headers:
-        httpx_client.headers.update(extra_headers)
         verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
 
     a2a_client: Final = await create_client(  # pyright: ignore[reportOptionalCall]
@@ -784,11 +780,17 @@ async def create_a2a_client(
             httpx_client=httpx_client,
             streaming=streaming,
         ),
+        resolver_http_kwargs={"headers": extra_headers} if extra_headers else None,
     )
     # Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse
-    # the configured httpx client (with this agent's trace-id/auth headers) without
-    # excavating a2a-sdk private internals.
+    # the configured httpx client and this agent's headers without excavating
+    # a2a-sdk private internals.
     a2a_client._litellm_httpx_client = httpx_client
+    a2a_client._litellm_call_context = (  # pyright: ignore[reportAttributeAccessIssue]  # LiteLLM-owned stash
+        ClientCallContext(service_parameters=extra_headers)  # pyright: ignore[reportOptionalCall]  # SDK checked above
+        if extra_headers
+        else None
+    )
     agent_card: Final = getattr(a2a_client, "_card", None)
     if agent_card is not None:
         a2a_client._litellm_agent_card = agent_card

tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py

@@ -58,6 +58,32 @@ async def test_localhost_retry_reuses_stashed_httpx_client():
     assert mock_create.await_count == 1
 
 
+@pytest.mark.asyncio
+async def test_localhost_retry_carries_the_agents_call_context_onto_the_new_client():
+    """Per-caller headers ride on the call context now, not on the shared httpx client,
+    so a retry that drops the context would replay the request unauthenticated."""
+    stashed_context = object()
+    a2a_client = MagicMock()
+    a2a_client._litellm_httpx_client = object()
+    a2a_client._litellm_call_context = stashed_context
+    new_client = MagicMock()
+
+    with (
+        patch.object(emu, "A2A_SDK_AVAILABLE", True),
+        patch.object(emu, "set_agent_card_url"),
+        patch.object(emu, "ClientConfig", side_effect=lambda **_: MagicMock()),
+        patch.object(emu, "create_client", new=AsyncMock(return_value=new_client)),
+    ):
+        result = await emu.handle_a2a_localhost_retry(
+            error=_localhost_error(),
+            agent_card=MagicMock(),
+            a2a_client=a2a_client,
+            is_streaming=False,
+        )
+
+    assert result._litellm_call_context is stashed_context
+
+
 @pytest.mark.asyncio
 async def test_localhost_retry_raises_when_no_stashed_client():
     """An externally-supplied client has no LiteLLM httpx handle; the retry must fail

tests/test_litellm/a2a_protocol/test_main.py

@@ -1,13 +1,26 @@
 """Tests for litellm/a2a_protocol/main.py non-streaming send behavior."""
 
+import httpx
 import pytest
 
 pytest.importorskip("a2a.compat.v0_3.conversions")
 
 from a2a.compat.v0_3 import conversions as _conv
-from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest
+from a2a.compat.v0_3.types import (
+    MessageSendParams,
+    SendMessageRequest,
+    SendStreamingMessageRequest,
+)
 
-from litellm.a2a_protocol.main import _send_message
+import litellm
+from litellm.a2a_protocol.main import _send_message, _stream_messages, create_a2a_client
+from litellm.caching.llm_caching_handler import LLMClientCache
+from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
+from litellm.llms.custom_httpx.http_handler import (
+    AsyncHTTPHandler,
+    get_async_httpx_client,
+    httpxSpecialProvider,
+)
 
 
 def _request() -> SendMessageRequest:
@@ -40,7 +53,8 @@ class _FakeClient:
     def __init__(self, *events):
         self._events = events
 
-    async def send_message(self, _pb_request):
+    async def send_message(self, _pb_request, context=None):
+        self.context = context
         for event in self._events:
             yield event
 
@@ -133,3 +147,203 @@ def test_streaming_logging_obj_carries_call_type_into_model_call_details():
     )
 
     assert logging_obj.model_call_details["call_type"] == "asend_message_streaming"
+
+
+_AGENT_CARD = {
+    "protocolVersion": "0.3.0",
+    "name": "recording-agent",
+    "url": "http://127.0.0.1:9/",
+    "preferredTransport": "JSONRPC",
+    "version": "1.0.0",
+    "capabilities": {"streaming": True},
+    "defaultInputModes": ["text/plain"],
+    "defaultOutputModes": ["text/plain"],
+    "skills": [],
+}
+
+_RPC_REPLY = {
+    "jsonrpc": "2.0",
+    "id": "reply",
+    "result": {
+        "messageId": "reply-1",
+        "role": "agent",
+        "parts": [{"kind": "text", "text": "pong"}],
+        "kind": "message",
+    },
+}
+
+_AGENT_A_HEADERS = {"x-agent-token": "token-for-a", "x-tenant": "tenant-a"}
+_AGENT_B_HEADERS = {"x-agent-token": "token-for-b", "x-tenant": "tenant-b"}
+_UPSTREAM_SESSION_COOKIE = "a2a_session=only-agent-a-may-hold-this; Path=/"
+
+
+class _RequestRecorder:
+    """Records the headers httpx put on the wire, per outbound request.
+
+    ``cookie_from_tenant`` makes that tenant's agent answer with a Set-Cookie, standing in
+    for an upstream that issues a session cookie.
+    """
+
+    def __init__(self, cookie_from_tenant: str | None = None):
+        self.card_requests = []
+        self.rpc_requests = []
+        self.client = None
+        self.cookie_from_tenant = cookie_from_tenant
+
+    def __call__(self, request: httpx.Request) -> httpx.Response:
+        headers = {k.lower(): v for k, v in request.headers.items()}
+        if request.method == "GET":
+            self.card_requests.append(headers)
+            return httpx.Response(200, json=_AGENT_CARD)
+        self.rpc_requests.append(headers)
+        if self.cookie_from_tenant is not None and headers.get("x-tenant") == self.cookie_from_tenant:
+            return httpx.Response(200, json=_RPC_REPLY, headers={"set-cookie": _UPSTREAM_SESSION_COOKIE})
+        return httpx.Response(200, json=_RPC_REPLY)
+
+
+def _a2a_client_cache_key(timeout: float) -> str:
+    return "async_httpx_client" + f"timeout_{timeout}" + httpxSpecialProvider.A2AProvider
+
+
+async def _seed_shared_a2a_client(cookie_from_tenant: str | None = None) -> _RequestRecorder:
+    """Put the one A2A client the cache will hand out behind a mock transport.
+
+    Seeding has to happen on the test's own event loop, because the client cache keys on
+    it. The injected client is a real httpx.AsyncClient, so the merge of per-request
+    headers over client defaults, and httpx's own cookie handling, which is what these
+    tests are about, stay real.
+    """
+    recorder = _RequestRecorder(cookie_from_tenant=cookie_from_tenant)
+    handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT)
+    owned_client = handler.client
+    handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder))
+    await owned_client.aclose()
+
+    litellm.in_memory_llm_clients_cache.set_cache(key=_a2a_client_cache_key(DEFAULT_A2A_AGENT_TIMEOUT), value=handler)
+    seeded = get_async_httpx_client(
+        llm_provider=httpxSpecialProvider.A2AProvider,
+        params={"timeout": DEFAULT_A2A_AGENT_TIMEOUT},
+    )
+    assert seeded is handler, "cache key drifted from get_async_httpx_client; these tests would test nothing"
+
+    recorder.client = handler.client
+    return recorder
+
+
+@pytest.fixture
+def isolated_client_cache():
+    previous = getattr(litellm, "in_memory_llm_clients_cache", None)
+    litellm.in_memory_llm_clients_cache = LLMClientCache()
+    yield litellm.in_memory_llm_clients_cache
+    litellm.in_memory_llm_clients_cache = previous
+
+
+def _send_request(request_id):
+    return SendMessageRequest(
+        id=request_id,
+        params=MessageSendParams(
+            message={"messageId": request_id, "role": "user", "parts": [{"kind": "text", "text": "hi"}]}
+        ),
+    )
+
+
+@pytest.mark.asyncio
+async def test_extra_headers_never_land_on_the_shared_cached_client(isolated_client_cache):
+    """get_async_httpx_client hands back a process-wide shared client, so a caller's
+    headers written onto it would outlive the request that supplied them."""
+    recorder = await _seed_shared_a2a_client()
+
+    await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS)
+
+    assert "x-agent-token" not in recorder.client.headers
+    assert "x-tenant" not in recorder.client.headers
+
+
+@pytest.mark.asyncio
+async def test_callers_with_different_headers_reuse_one_pooled_client(isolated_client_cache):
+    """Headers must not segregate the connection pool. Every A2A caller on one timeout
+    shares one cached client, so header sets cannot multiply cached clients."""
+    recorder = await _seed_shared_a2a_client()
+
+    client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS)
+    client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS)
+    client_none = await create_a2a_client(base_url="http://127.0.0.1:9")
+
+    assert client_a._litellm_httpx_client is recorder.client
+    assert client_b._litellm_httpx_client is recorder.client
+    assert client_none._litellm_httpx_client is recorder.client
+
+    cached = [key for key in isolated_client_cache.cache_dict if "a2a_provider" in key]
+    assert len(cached) == 1, f"expected one pooled A2A client, cached: {cached}"
+
+
+@pytest.mark.parametrize("order", [("a", "b", "none"), ("b", "none", "a"), ("none", "a", "b")])
+@pytest.mark.asyncio
+async def test_each_caller_sends_only_its_own_headers(order, isolated_client_cache):
+    """Whatever order callers arrive in, each request carries that caller's headers and
+    no other caller's, and a caller with no extra_headers sends none."""
+    recorder = await _seed_shared_a2a_client()
+    headers_by_caller = {"a": _AGENT_A_HEADERS, "b": _AGENT_B_HEADERS, "none": None}
+
+    for caller in order:
+        a2a_client = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=headers_by_caller[caller])
+        await _send_message(a2a_client, _send_request(caller))
+
+    received = dict(zip(order, recorder.rpc_requests, strict=True))
+    assert received["a"]["x-agent-token"] == "token-for-a"
+    assert received["a"]["x-tenant"] == "tenant-a"
+    assert received["b"]["x-agent-token"] == "token-for-b"
+    assert received["b"]["x-tenant"] == "tenant-b"
+    assert "x-agent-token" not in received["none"]
+    assert "x-tenant" not in received["none"]
+
+
+@pytest.mark.asyncio
+async def test_streaming_send_carries_only_its_own_caller_headers(isolated_client_cache):
+    """The streaming path shares the same pooled client, so it needs the same guard."""
+    recorder = await _seed_shared_a2a_client()
+
+    client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS, streaming=True)
+    await _send_message(client_a, _send_request("a"))
+
+    client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS, streaming=True)
+    streaming_request = SendStreamingMessageRequest(
+        id="b",
+        params=MessageSendParams(message={"messageId": "b", "role": "user", "parts": [{"kind": "text", "text": "hi"}]}),
+    )
+    async for _ in _stream_messages(client_b, streaming_request):
+        pass
+
+    received = dict(zip(("a", "b"), recorder.rpc_requests, strict=True))
+    assert received["a"]["x-agent-token"] == "token-for-a"
+    assert received["b"]["x-agent-token"] == "token-for-b"
+    assert received["b"]["x-tenant"] == "tenant-b"
+
+
+@pytest.mark.asyncio
+async def test_agent_card_fetch_carries_the_callers_headers(isolated_client_cache):
+    """Agent cards can sit behind the same auth as the agent, so the card fetch must stay
+    authenticated once the headers stop living on the client."""
+    recorder = await _seed_shared_a2a_client()
+
+    await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS)
+
+    assert recorder.card_requests, "no agent card request was made"
+    assert recorder.card_requests[-1]["x-agent-token"] == "token-for-a"
+
+
+@pytest.mark.asyncio
+async def test_one_agents_session_cookie_never_reaches_another_agent(isolated_client_cache):
+    """One pooled client is also one httpx cookie jar. httpx stores every Set-Cookie on the
+    client and replays it on any later request to a matching domain, so an agent's session
+    cookie would ride along on a different agent's call to the same host."""
+    recorder = await _seed_shared_a2a_client(cookie_from_tenant="tenant-a")
+
+    client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS)
+    await _send_message(client_a, _send_request("a"))
+    client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS)
+    await _send_message(client_b, _send_request("b"))
+
+    assert dict(recorder.client.cookies) == {}, "the shared client kept an agent's session cookie"
+    assert "cookie" not in recorder.card_requests[-1]
+    assert "cookie" not in recorder.rpc_requests[-1]

tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py

@@ -227,7 +227,7 @@ async def test_each_agent_gets_only_its_own_static_headers():
 
 
 # ---------------------------------------------------------------------------
-# Unit tests: create_a2a_client (httpx client per call + timeout defaults)
+# Unit tests: create_a2a_client (shared client left untouched + timeout defaults)
 # ---------------------------------------------------------------------------
 
 
@@ -253,17 +253,13 @@ async def _fake_create_client(base_url, client_config=None, **kwargs):
 
 
 @pytest.mark.asyncio
-async def test_create_a2a_client_uses_fresh_httpx_client():
+async def test_create_a2a_client_leaves_the_shared_client_untouched():
     """
-    Two calls to create_a2a_client with different extra_headers must produce
-    distinct underlying httpx clients — preventing header bleed between agents.
-
-    The test checks:
-    1. get_async_httpx_client was called twice (once per create_a2a_client call).
-    2. The two returned A2A clients carry distinct httpx client objects (direct
-       proof of header isolation, not just cache-key difference).
-    3. The cache-key param differs between calls (so the real LRU cache cannot
-       return the same httpx client even under load).
+    The client from get_async_httpx_client is shared across every A2A caller, so
+    create_a2a_client must neither write an agent's headers onto it nor derive its
+    cache key from them. Both would tie one agent's credentials to a cached object
+    that outlives the request. Header isolation on the wire is covered end to end in
+    tests/test_litellm/a2a_protocol/test_main.py.
     """
     pytest.importorskip("a2a.client")
     from litellm.a2a_protocol.main import create_a2a_client
@@ -281,39 +277,23 @@ async def test_create_a2a_client_uses_fresh_httpx_client():
             new=AsyncMock(side_effect=_fake_create_client),
         ),
     ):
-        a2a_client_a = await create_a2a_client(
+        await create_a2a_client(
             base_url="http://agent-a:9999",
             extra_headers={"Authorization": "Bearer a"},
         )
-        a2a_client_b = await create_a2a_client(
+        await create_a2a_client(
             base_url="http://agent-b:9999",
             extra_headers={"Authorization": "Bearer b"},
         )
 
-    assert (
-        len(captured_calls) == 2
-    ), "create_a2a_client should call get_async_httpx_client once per invocation"
+    assert len(captured_calls) == 2
 
-    # Direct proof: the two A2A clients must carry distinct httpx client objects.
-    # If they share one, mutating agent-B's Authorization header would bleed into A.
-    httpx_a = getattr(a2a_client_a, "_litellm_httpx_client", None)
-    httpx_b = getattr(a2a_client_b, "_litellm_httpx_client", None)
-    assert httpx_a is not None, "a2a_client_a missing _litellm_httpx_client"
-    assert httpx_b is not None, "a2a_client_b missing _litellm_httpx_client"
-    assert httpx_a is not httpx_b, (
-        "create_a2a_client returned the same httpx client for two agents with "
-        "different headers — Authorization header will bleed between agents"
-    )
-
-    # Also verify the cache-key param differs so the LRU cache never conflates them.
-    key_a = captured_calls[0]["params"].get("disable_aiohttp_transport")
-    key_b = captured_calls[1]["params"].get("disable_aiohttp_transport")
-    assert key_a is not None, "cache-key param 'disable_aiohttp_transport' missing"
-    assert key_b is not None, "cache-key param 'disable_aiohttp_transport' missing"
-    assert key_a != key_b, (
-        f"create_a2a_client used the same cache key for two agents with different "
-        f"headers — headers will bleed: key_a={key_a!r}, key_b={key_b!r}"
-    )
+    for call in captured_calls:
+        call["client"].headers.update.assert_not_called()
+        assert list(call["params"]) == ["timeout"], (
+            f"create_a2a_client passed extra client params {sorted(call['params'])}; "
+            "anything header-derived here gives every agent its own cached client"
+        )
 
 
 @pytest.mark.asyncio