fix(a2a): stop writing per-caller headers onto the shared cached httpx client - #35978
Conversation
|
|
Greptile SummaryThe PR moves caller-specific A2A headers from cached HTTP clients into per-call SDK context while retaining a shared connection pool
Confidence Score: 5/5The PR appears safe to merge The previously reported socket-based test dependency is removed, and no blocking failure remains
|
| Filename | Overview |
|---|---|
| litellm/a2a_protocol/main.py | Moves A2A headers into resolver and call-context arguments while sharing the cached HTTP client and disabling cookie persistence |
| litellm/a2a_protocol/exception_mapping_utils.py | Carries the stored per-call context onto the client created for localhost retry |
| tests/test_litellm/a2a_protocol/test_main.py | Uses a cached client backed by MockTransport to verify header isolation, pooling, streaming, card discovery, and cookie handling |
| tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py | Verifies that localhost retry preserves the original call context |
| tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py | Updates the regression contract to ensure shared cached clients are not mutated or keyed by caller headers |
Reviews (3): Last reviewed commit: "fix(a2a): stop writing per-caller state ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallContextType"]: | ||
| """Return the per-caller header context LiteLLM stashed at client creation, if any.""" |
There was a problem hiding this comment.
This helper docstring and the additional explanatory comments in the A2A creation and retry paths violate the repository policy against newly written comments unless explicitly requested. Remove the new prose and keep the behavior self-explanatory through code structure and naming.
Context Used: CLAUDE.md (source)
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
97fccf2 to
767319d
Compare
|
@greptileai please re-review 767319d Both findings are addressed in that commit. The localhost server in Every comment and docstring added to the production files is removed. The only comments remaining there are the two |
…client 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.
767319d to
fcedef4
Compare
|
@greptileai please review the current head fcedef4. The commit was squashed and now also blocks cookie persistence on the pooled client |
|
@tin-berri your approval was against 767319d. Since then this also blocks cookie persistence on the pooled client, so worth a re-look |
|
osv-scan is red on litellm_internal_staging itself and every open branch: h2 and js-yaml in lockfiles this PR does not touch |
|
Correction: my cookie proof used 127.0.0.1, and aiohttp refuses cookies for IP hosts, so it hid a second jar. Do not merge, fixing |
2b38991
into
litellm_internal_staging
|
Follow-up in #36149: this only blocked the httpx jar. aiohttp is the default transport and its session keeps a second jar, so the leak is still live |
#35978 stopped the pooled A2A client replaying one upstream's Set-Cookie to another by installing a blocking policy on that client's httpx cookie jar. That covers only one of the two jars on the request path. AiohttpTransport is the default transport unless it is explicitly disabled, and the aiohttp ClientSession behind it keeps its own cookie jar which no httpx-level assertion can observe, so the leak is still live on the default path: a live proxy on that commit still delivers agent-alpha's session cookie to agent-beta's card fetch and JSON-RPC call. The reason it looked fixed is that aiohttp's default CookieJar is built with unsafe=False and refuses to store cookies for IP hosts, so a proof addressed to 127.0.0.1 comes back clean whether or not that jar is blocked. Cookie persistence is now blocked where the clients are built rather than at one call site: blocked_cookie_jar() gives every httpx client, async and sync, a jar whose DefaultCookiePolicy(allowed_domains=()) rejects every domain in both directions, and both ClientSession constructions litellm owns, the transport's session factory and the proxy's shared startup session, get a DummyCookieJar. LiteLLM reads a response cookie nowhere, and an explicitly supplied Cookie header still goes out, so passthrough forwarding and an agent's extra_headers are unaffected. The A2A-scoped policy #35978 added is removed, since it is now dead. The two suites that drive the aiohttp session factory synchronously mock ClientSession because a real one needs a running event loop; DummyCookieJar has the same requirement, so they mock it for the same reason.
…too (#36149) #35978 stopped the pooled A2A client replaying one upstream's Set-Cookie to another by installing a blocking policy on that client's httpx cookie jar. That covers only one of the two jars on the request path. AiohttpTransport is the default transport unless it is explicitly disabled, and the aiohttp ClientSession behind it keeps its own cookie jar which no httpx-level assertion can observe, so the leak is still live on the default path: a live proxy on that commit still delivers agent-alpha's session cookie to agent-beta's card fetch and JSON-RPC call. The reason it looked fixed is that aiohttp's default CookieJar is built with unsafe=False and refuses to store cookies for IP hosts, so a proof addressed to 127.0.0.1 comes back clean whether or not that jar is blocked. Cookie persistence is now blocked where the clients are built rather than at one call site: blocked_cookie_jar() gives every httpx client, async and sync, a jar whose DefaultCookiePolicy(allowed_domains=()) rejects every domain in both directions, and both ClientSession constructions litellm owns, the transport's session factory and the proxy's shared startup session, get a DummyCookieJar. LiteLLM reads a response cookie nowhere, and an explicitly supplied Cookie header still goes out, so passthrough forwarding and an agent's extra_headers are unaffected. The A2A-scoped policy #35978 added is removed, since it is now dead. The two suites that drive the aiohttp session factory synchronously mock ClientSession because a real one needs a running event loop; DummyCookieJar has the same requirement, so they mock it for the same reason.
3 participants
TLDR
Problem this solves:
How it solves it:
Relevant issues
Linear ticket
Resolves LIT-5229
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Before: the two changed source files restored to their
2792887e47content, reproducible withgit checkout 2792887e47 -- litellm/a2a_protocol/. After: commitfcedef4627. Section 3's before leg is narrower, since it isolates the one behaviour that section 1 and 2 introduce: it restores onlylitellm/a2a_protocol/main.pyto767319d1b7, the head this PR carried before the cookie changeThe upstream here is a local recording agent rather than a hosted one, because the measurement is which connection, which headers and which cookies reach the far side, and only an endpoint that reports its own inbound requests can show that. This is the live proof only; the test suite itself opens no sockets
1. The shared client cache, driving the real
create_a2a_client250 calls, each with a distinct
X-LiteLLM-Trace-Idexactly as the proxy generates them, with one non-A2A client seeded into the cache beforehand. The harness asserts it really did generate 250 distinct header sets, otherwise the number below would be measuring one cache keyBefore:
After:
2. Connection reuse through a live proxy
Two agents against one upstream, 20 alternating A2A calls through a proxy on port 15229, upstream on 25229, with this config:
Before:
After:
All 20 calls succeed on both runs and each agent's own header still arrives on its own requests, so the difference is purely that 20 requests stop costing 20 connections and 40 held sockets
Those two numbers were captured at
767319d1b7, before the cookie change. Section 3's burst re-measures the connection count on the current head and gets the same 23. Cookie isolation between agents, through a live proxy
Same shape, on ports 15230 and 25230. The upstream answers
agent-alpha's JSON-RPC call withSet-Cookie: a2a_session=only-alpha-may-hold-thisand records theCookieheader on every request it receives, agent card fetches included. Two agents, one host, so both land on the same cookie domainBefore:
After:
Pooling is untouched by that. 20 further alternating calls on the same proxy, after the two above:
Type
🐛 Bug Fix
Changes
create_a2a_clientfolded the caller's header set into the httpx client cache key, by way ofstr(sorted(extra_headers.items()))stuffed into the unrelateddisable_aiohttp_transportparam, so that it could then callheaders.update(extra_headers)on the returned client without one caller's headers reaching the next. That works, but it pays for the isolation by giving every distinct header set its own client, and the proxy puts a freshX-LiteLLM-Trace-Idintoextra_headerson every request. So the key changed on every call. Each A2A request built a newhttpx.AsyncClient, cached it under a key nothing would ever look up again, and pushed an entry into the 200-entryin_memory_llm_clients_cachethat every provider sharesMeasured above: 250 A2A requests take that cache from 5 entries to its cap and evict a previously cached non-A2A client. A2A traffic therefore destroys connection reuse for the whole proxy and gets none itself, holding two sockets per request open instead of pooling. Writing to a shared cached object was the other half of the same design, and it is only harmless today because of the cache key trick, so a future change that drops that one line would silently start bleeding one caller's headers into the next
The headers now travel with the request instead of living on the client.
create_a2a_clientasks the cache for a client keyed on timeout alone, hands the caller's headers to the agent card fetch throughcreate_client(resolver_http_kwargs=...), and stashes aClientCallContext(service_parameters=extra_headers)next to the existing_litellm_httpx_clienthandle._send_messageand_stream_messagespass that context down, and the SDK turns it into per-requestheaders=, which httpx merges over the client defaults.handle_a2a_localhost_retrycarries the context onto the retry client the same way it already carried the httpx client, so a retried request keeps its headersHeaders are not the only per-caller state an
httpx.AsyncClientcarries.AsyncClient._send_single_requestcallsself.cookies.extract_cookies(response)on every response, and_merge_cookiesreplays the jar onto every outgoing request, so the moment all A2A callers share one client they share one cookie jar. Section 3 shows the consequence on a live proxy: an agent's session cookie arrives at a different agent on the same host, on that agent's card fetch and on its JSON-RPC call. The pooled client now carriesDefaultCookiePolicy(allowed_domains=()), which rejects every domain in both directions, so nothing is stored and nothing is replayed. Nothing in litellm or in the a2a SDK reads cookies for A2A: the SDK's own auth interceptor skips API keys declaredin: cookieand says so, andgrep -ri cookie litellm/a2a_protocol/matches only the three lines this PR adds. Blocking is also what the pre-PR proxy behaviour amounted to, since a per-request client jar never outlived its requestThis is adjacent to LIT-4883, which covers memory growth in that same client cache. A workload minting one client per request is the profile that makes cache churn expensive, and this PR removes the churn at its source. It does not change any caching or eviction behaviour, so the two are independent
On tests, the claim that the SDK's call context becomes per-request headers is a claim about what reaches the transport, so
tests/test_litellm/a2a_protocol/test_main.pyseeds the shared client cache with a realhttpx.AsyncClientbacked byhttpx.MockTransportand asserts on thehttpx.Requestobjects that arrive there. No socket is opened, and the merge of per-request headers over client defaults, and httpx's own cookie handling, still happen for real inside httpx. Seeding has to run on the test's own event loop, since the client cache keys on it, and the fixture asserts the seeded handler is the oneget_async_httpx_clienthands back, so a drift in the cache key formula fails loudly instead of silently testing nothingBeing precise about which tests are regression evidence and which are guards, since they fail on the parent commit for different reasons:
test_create_a2a_client_leaves_the_shared_client_untouchedis the regression test. On2792887e47it fails on its own assertion,Expected 'update' to not have been called. Called 1 times., with no network and nothing environmental. It pins both halves of the invariant: the shared client is never written to, and no header-derived param reaches the cache keytest_one_agents_session_cookie_never_reaches_another_agentis the regression test for the cookie half. Its mock upstream answers one tenant with aSet-Cookie, and it asserts the shared client's jar stays empty and that the next agent's card fetch and RPC carry noCookieheader. Two mutants, dropping the policy install and widening it to a plainDefaultCookiePolicy(), each fail it on their own runThe
test_main.pygroup also fails on2792887e47, but on connection refused rather than on an assertion, because the old cache key means the code never reaches the seeded client at all. That failure is a fair statement of the defect, and it is deterministic and offline since the URL is the discard port, but it is not an assertion-level catch and is not presented as one. Their real job is as guards: once every caller genuinely shares one client,test_each_caller_sends_only_its_own_headersacross all three arrival orderings, plus its streaming and agent-card counterparts, fail loudly on any reintroduction of a write to the shared client's headers. Before this PR they could not do that, because each header set had its own clienttest_create_a2a_client_uses_fresh_httpx_clientin the proxy suite asserted the old contract, and was vacuous besides: its mock handed back a new client object per call whatever the code did, so it would have passed against code that ignoredextra_headersentirely. It is the test replaced by the regression test aboveMutation check over eight mutants, one per changed call site, mutated and run separately so no mutant can hide behind another. Each edit asserts it changed the file and each run asserts pytest reported a test count, so a silently skipped mutation cannot read as a pass. Baseline 257 passed, restored 257 passed, 8 killed, 0 survived
QA runbook
Final Attestation