refactor(a2a): resolve the relay's Entra hop bearer inside the a2a provider helper

litellm/llms/a2a/common_utils.py

@@ -2,7 +2,7 @@
 Common utilities for A2A (Agent-to-Agent) Protocol
 """
 
-from collections.abc import Mapping
+from collections.abc import Awaitable, Callable, Mapping
 from typing import Any, Final
 
 from pydantic import BaseModel
@@ -10,6 +10,7 @@ from pydantic import BaseModel
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     convert_content_list_to_str,
 )
+from litellm.llms.azure_ai.common_utils import has_azure_entra_params, resolve_azure_ai_agent_auth_header
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.types.llms.openai import AllMessageValues
 
@@ -142,3 +143,17 @@ def extract_text_from_a2a_response(response_dict: Mapping[str, object], max_dept
         return extract_text_from_a2a_message(first_artifact, depth=0, max_depth=max_depth)
 
     return ""
+
+
+AgentAuthHeaderResolver = Callable[[Mapping[str, object]], Awaitable[Mapping[str, str]]]
+
+
+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):
+        return None
+    return await resolve_entra_header(litellm_params)

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -24,7 +24,7 @@ from pydantic import ValidationError
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
-from litellm.llms.azure_ai.common_utils import has_azure_entra_params, resolve_azure_ai_agent_auth_header
+from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.a2a.version_convert import (
     A2AVersion,
@@ -162,12 +162,9 @@ async def _resolve_backend_auth_header(
     litellm_params: dict[str, object],
     custom_llm_provider: object,
 ) -> Mapping[str, str] | None:
-    """Entra credentials only authenticate the A2A hop; completion-bridge agents pass them to the model provider instead."""
     if litellm_params.get(DATABRICKS_OAUTH_PARAM):
         return await resolve_databricks_app_auth_header(litellm_params)
-    if not custom_llm_provider and has_azure_entra_params(litellm_params):
-        return await resolve_azure_ai_agent_auth_header(litellm_params)
-    return None
+    return await resolve_a2a_hop_auth_header(litellm_params, custom_llm_provider)
 
 
 def _forwarding_headers(

tests/test_litellm/llms/a2a/test_common_utils.py

@@ -0,0 +1,52 @@
+"""Tests for litellm/llms/a2a/common_utils.py."""
+
+from collections.abc import Mapping
+from types import MappingProxyType
+
+import pytest
+
+from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
+
+
+class _RecordingEntraResolver:
+    def __init__(self) -> None:
+        self.calls: list[Mapping[str, object]] = []
+
+    async def __call__(self, litellm_params: Mapping[str, object]) -> Mapping[str, str]:
+        self.calls.append(litellm_params)
+        return MappingProxyType({"Authorization": "Bearer minted-entra-token"})
+
+
+_SERVICE_PRINCIPAL = MappingProxyType({"tenant_id": "tenant", "client_id": "client", "client_secret": "sp-secret"})
+
+
+@pytest.mark.asyncio
+async def test_entra_agent_gets_a_minted_bearer_for_the_a2a_hop():
+    resolver = _RecordingEntraResolver()
+
+    header = await resolve_a2a_hop_auth_header(_SERVICE_PRINCIPAL, None, resolver)
+
+    assert header == {"Authorization": "Bearer minted-entra-token"}
+    assert resolver.calls == [_SERVICE_PRINCIPAL]
+
+
+@pytest.mark.asyncio
+async def test_completion_bridge_agent_keeps_its_entra_credentials_for_the_model_provider():
+    """A bridged agent's tenant_id/client_id/client_secret authenticate the model it bridges to, so the A2A hop
+    must not spend them on a bearer of its own."""
+    resolver = _RecordingEntraResolver()
+
+    header = await resolve_a2a_hop_auth_header(_SERVICE_PRINCIPAL, "azure_ai", resolver)
+
+    assert header is None
+    assert resolver.calls == []
+
+
+@pytest.mark.asyncio
+async def test_agent_without_entra_credentials_gets_no_bearer():
+    resolver = _RecordingEntraResolver()
+
+    header = await resolve_a2a_hop_auth_header({"api_base": "https://agent.example.com"}, None, resolver)
+
+    assert header is None
+    assert resolver.calls == []