diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 7b91cb780d9..77f26b65de0 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -8,11 +8,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx -from litellm.llms.azure_ai.common_utils import ( - AZURE_ENTRA_LITELLM_PARAM_KEYS, - get_azure_ai_agent_entra_token, - has_azure_entra_params, -) +from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_LITELLM_PARAM_KEYS, get_azure_ai_agent_entra_token from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues @@ -20,6 +16,7 @@ from litellm.types.utils import Choices, Message, ModelResponse, Usage from ..common_utils import ( A2AError, + a2a_hop_uses_entra, convert_messages_to_prompt, extract_text_from_a2a_response, ) @@ -41,13 +38,27 @@ def _card_declares_no_streaming(agent_card_params: Mapping[str, object]) -> bool return isinstance(capabilities, Mapping) and not capabilities.get("streaming") -def _registry_api_key(agent_litellm_params: dict[str, object]) -> str | None: - configured_api_key: Final = agent_litellm_params.get("api_key") - if isinstance(configured_api_key, str): - return configured_api_key - if has_azure_entra_params(agent_litellm_params): +def _agent_authenticates_with_entra(agent_litellm_params: Mapping[str, object]) -> bool: + return a2a_hop_uses_entra(agent_litellm_params, agent_litellm_params.get("custom_llm_provider")) + + +def _registry_api_key(agent_litellm_params: Mapping[str, object]) -> str | None: + if _agent_authenticates_with_entra(agent_litellm_params): return get_azure_ai_agent_entra_token(agent_litellm_params) - return None + configured_api_key: Final = agent_litellm_params.get("api_key") + return configured_api_key if isinstance(configured_api_key, str) else None + + +def _registry_headers(agent_litellm_params: Mapping[str, object]) -> dict[str, Any] | None: + stored_headers: Final = agent_litellm_params.get("headers") + if not isinstance(stored_headers, Mapping): + return None + entra_owns_authorization: Final = _agent_authenticates_with_entra(agent_litellm_params) + return { # mutable-ok: completion() and httpx take the request headers as a dict + name: value + for name, value in stored_headers.items() + if not (entra_owns_authorization and str(name).lower() == "authorization") + } class A2AConfig(BaseConfig): @@ -101,9 +112,7 @@ class A2AConfig(BaseConfig): api_key = _registry_api_key(agent.litellm_params) if not headers: - agent_headers: Final = agent.litellm_params.get("headers") - if agent_headers: - headers = dict(agent_headers) + headers = _registry_headers(agent.litellm_params) or headers # Merge other litellm_params (timeout, max_retries, etc.) registry_params: Final = tuple( diff --git a/litellm/llms/a2a/common_utils.py b/litellm/llms/a2a/common_utils.py index 0cbc137c998..030c5bc222e 100644 --- a/litellm/llms/a2a/common_utils.py +++ b/litellm/llms/a2a/common_utils.py @@ -148,12 +148,16 @@ def extract_text_from_a2a_response(response_dict: Mapping[str, object], max_dept AgentAuthHeaderResolver = Callable[[Mapping[str, object]], Awaitable[Mapping[str, str]]] +def a2a_hop_uses_entra(litellm_params: Mapping[str, object], custom_llm_provider: object) -> bool: + return not custom_llm_provider and has_azure_entra_params(litellm_params) + + async def resolve_a2a_hop_auth_header( litellm_params: Mapping[str, object], custom_llm_provider: object, resolve_entra_header: AgentAuthHeaderResolver = resolve_azure_ai_agent_auth_header, ) -> Mapping[str, str] | None: """Entra credentials authenticate the A2A hop only; a completion-bridge agent hands them to the model provider it bridges to.""" - if custom_llm_provider or not has_azure_entra_params(litellm_params): + if not a2a_hop_uses_entra(litellm_params, custom_llm_provider): return None return await resolve_entra_header(litellm_params) diff --git a/tests/test_litellm/test_a2a_registry_lookup.py b/tests/test_litellm/test_a2a_registry_lookup.py index 5f371d69059..5ba0b84cdbf 100644 --- a/tests/test_litellm/test_a2a_registry_lookup.py +++ b/tests/test_litellm/test_a2a_registry_lookup.py @@ -247,6 +247,73 @@ def test_registry_entra_agent_authenticates_with_the_entra_token_and_keeps_its_s assert optional_params == {"timeout": 30} +_STORED_STATIC_CREDENTIALS: dict = { + "api_key": "stored-key", + "headers": {"authorization": "Bearer stored-header", "X-Agent": "static"}, +} + + +@pytest.mark.parametrize( + ("litellm_params", "expected_authorization_lines"), + [ + ( + {**_STORED_STATIC_CREDENTIALS, "azure_ad_token": "entra-token"}, + {"Authorization": "Bearer entra-token"}, + ), + ( + _STORED_STATIC_CREDENTIALS, + {"authorization": "Bearer stored-header", "Authorization": "Bearer stored-key"}, + ), + ( + {**_STORED_STATIC_CREDENTIALS, "azure_ad_token": "model-provider-token", "custom_llm_provider": "azure_ai"}, + {"authorization": "Bearer stored-header", "Authorization": "Bearer stored-key"}, + ), + ], + ids=[ + "entra agent: the minted bearer is the only authorization line", + "agent without entra credentials: static credentials sent as before", + "bridge agent: its entra credentials belong to the model provider, never to the a2a hop", + ], +) +def test_entra_credentials_beat_the_static_credentials_stored_next_to_them_on_the_chat_route( + litellm_params: dict, expected_authorization_lines: dict +): + """The relay sends the minted Entra bearer over any static Authorization stored on the agent; the chat + route must agree, or an api_key or authorization header left next to the Entra fields makes the same + agent answer on /a2a and fail with the backend's 401 on /v1/chat/completions.""" + 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 + + agent = AgentResponse( + agent_id="mixed-credentials-id", + agent_name="mixed-credentials-agent", + agent_card_params={"url": "https://foundry.example.com/a2a"}, + litellm_params=litellm_params, + ) + client = HTTPHandler() + agent_reply = httpx.Response( + 200, + json={"jsonrpc": "2.0", "id": "1", "result": {"kind": "message", "parts": [{"kind": "text", "text": "ok"}]}}, + ) + original_agents = global_agent_registry.agent_list.copy() + global_agent_registry.register_agent(agent) + + try: + with patch.object(client, "post", return_value=agent_reply) as post: # test-quality-ok: injected client + litellm.completion( + model="a2a/mixed-credentials-agent", messages=[{"role": "user", "content": "hi"}], client=client + ) + finally: + global_agent_registry.agent_list = original_agents + + sent_headers = post.call_args.kwargs["headers"] + assert { + name: value for name, value in sent_headers.items() if name.lower() == "authorization" + } == expected_authorization_lines + assert sent_headers["X-Agent"] == "static" + + def test_registry_entra_agent_with_an_unresolvable_credential_fails_the_chat_call(monkeypatch): """The chat route mints the Foundry bearer from the registered credentials; when they resolve to nothing the caller must get the credential error instead of an unauthenticated backend call."""