From 4669041ae83cb8e327c4136b399c8a0766f16bd7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 17 Sep 2026 14:58:50 -0700 Subject: [PATCH] fix(a2a): copy registered agent headers so one caller's bearer never reaches the next The chat route handed the registry's stored headers dict straight to validate_environment, which wrote the caller's bearer into it, so the next caller of the same agent with no key of their own sent the previous caller's token. The registry lookup now copies the stored headers and validate_environment returns a new dict instead of mutating its input. A regression test drives two completions through one registered agent and asserts the second carries no Authorization and the stored agent is unchanged. --- litellm/llms/a2a/chat/transformation.py | 18 ++++----- .../test_litellm/test_a2a_registry_lookup.py | 37 +++++++++++++++++++ 2 files changed, 44 insertions(+), 11 deletions(-) diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index b185db1b69f..7b91cb780d9 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -103,7 +103,7 @@ class A2AConfig(BaseConfig): if not headers: agent_headers: Final = agent.litellm_params.get("headers") if agent_headers: - headers = agent_headers + headers = dict(agent_headers) # Merge other litellm_params (timeout, max_retries, etc.) registry_params: Final = tuple( @@ -174,17 +174,13 @@ class A2AConfig(BaseConfig): api_base: API base URL Returns: - Updated headers dict + A new headers dict; the caller's dict is left untouched """ - # Ensure Content-Type is set to application/json for JSON-RPC 2.0 - if "content-type" not in headers and "Content-Type" not in headers: - headers["Content-Type"] = "application/json" - - # Add Authorization header if API key is provided - if api_key is not None: - headers["Authorization"] = f"Bearer {api_key}" - - return headers + content_type_default: Final = ( + () if "content-type" in headers or "Content-Type" in headers else (("Content-Type", "application/json"),) + ) + bearer: Final = () if api_key is None else (("Authorization", f"Bearer {api_key}"),) + return dict((*headers.items(), *content_type_default, *bearer)) def get_complete_url( self, diff --git a/tests/test_litellm/test_a2a_registry_lookup.py b/tests/test_litellm/test_a2a_registry_lookup.py index 68cdd3f4995..5f371d69059 100644 --- a/tests/test_litellm/test_a2a_registry_lookup.py +++ b/tests/test_litellm/test_a2a_registry_lookup.py @@ -75,6 +75,43 @@ def test_a2a_registry_integration(): assert post.call_args.kwargs["headers"]["X-Agent"] == "static" +def test_one_callers_bearer_never_reaches_another_caller_of_the_same_registered_agent(): + """The registered headers dict is shared by every request to the agent, so the bearer one caller + supplies must be written to that request alone and never persisted onto the agent for the next + caller, who has no key of their own.""" + from litellm.llms.custom_httpx.http_handler import HTTPHandler + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.types.agents import AgentResponse + + shared_agent = AgentResponse( + agent_id="shared-id", + agent_name="shared-agent", + agent_card_params={"url": "http://registry-url.example.com:9999"}, + litellm_params={"headers": {"X-Agent": "static"}}, + ) + client = HTTPHandler() + agent_reply = httpx.Response( + 200, + json={"jsonrpc": "2.0", "id": "1", "result": {"kind": "message", "parts": [{"kind": "text", "text": "ok"}]}}, + ) + messages = [{"role": "user", "content": "hi"}] + original_agents = global_agent_registry.agent_list.copy() + global_agent_registry.register_agent(shared_agent) + + try: + with patch.object(client, "post", return_value=agent_reply) as post: # test-quality-ok: injected client + litellm.completion(model="a2a/shared-agent", messages=messages, api_key="caller-one-key", client=client) + litellm.completion(model="a2a/shared-agent", messages=messages, client=client) + finally: + global_agent_registry.agent_list = original_agents + + first_call_headers, second_call_headers = (call.kwargs["headers"] for call in post.call_args_list) + assert first_call_headers["Authorization"] == "Bearer caller-one-key" + assert "Authorization" not in second_call_headers + assert second_call_headers["X-Agent"] == "static" + assert shared_agent.litellm_params == {"headers": {"X-Agent": "static"}} + + def _foundry_card_stored_through_the_agents_api() -> dict: from litellm.proxy.a2a.agent_card import merge_agent_card