From 2b38991df91054aec1293ed28a03f2d00290e7ad Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 6 Aug 2026 18:58:26 -0700 Subject: [PATCH] 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. --- .../a2a_protocol/exception_mapping_utils.py | 9 +- litellm/a2a_protocol/main.py | 42 ++-- .../test_a2a_exception_mapping_utils.py | 26 +++ tests/test_litellm/a2a_protocol/test_main.py | 220 +++++++++++++++++- .../test_agent_header_isolation.py | 52 ++--- 5 files changed, 287 insertions(+), 62 deletions(-) diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py index 16c295f469c..62bef6e02ae 100644 --- a/litellm/a2a_protocol/exception_mapping_utils.py +++ b/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 diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 4b931e84427..fe0e6837586 100644 --- a/litellm/a2a_protocol/main.py +++ b/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 diff --git a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py index 44d2803a260..06191d1a370 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py +++ b/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 diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index c5e48620fde..59cb8c2c438 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/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] diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py index a51d6f6abcc..5193ad989dc 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py +++ b/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