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.
This commit is contained in:
mateo-berri 2026-09-17 14:58:50 -07:00
parent c2f77fd358
commit 4669041ae8
2 changed files with 44 additions and 11 deletions

View file

@ -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,

View file

@ -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