mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
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:
parent
c2f77fd358
commit
4669041ae8
2 changed files with 44 additions and 11 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue