Merge 66d7fb6570 into aeede53112

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -11,7 +11,8 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM
 """
 
 import json
-from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping
+from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Iterator, Mapping
+from itertools import chain
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
 from urllib.parse import urlparse
@@ -273,6 +274,32 @@ async def _log_agent_rejection(
         verbose_proxy_logger.exception("Failed to log rejected A2A request")
 
 
+_CONVENTION_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "proxy-authorization", "cookie", "x-api-key"})
+
+
+def _convention_headers(
+    headers: Mapping[str, str],
+    aliases: Iterable[str],
+    *,
+    block_credentials: bool,
+) -> Mapping[str, str]:
+    """Collect x-a2a-{agent_id_or_name}-{header_name} headers, case-insensitive.
+
+    Agents that set block_convention_credential_headers never receive caller
+    credentials this way, only their own static headers.
+    """
+    stripped: Final = chain.from_iterable(_strip_prefix(headers, f"x-a2a-{alias.lower()}-") for alias in aliases)
+    return {
+        name: val
+        for name, val in stripped
+        if name and not (block_credentials and name in _CONVENTION_CREDENTIAL_HEADERS)
+    }
+
+
+def _strip_prefix(headers: Mapping[str, str], prefix: str) -> Iterator[tuple[str, str]]:
+    return ((key.lower()[len(prefix) :], val) for key, val in headers.items() if key.lower().startswith(prefix))
+
+
 def _jsonrpc_error(
     request_id: object,
     code: int,
@@ -978,13 +1005,13 @@ async def invoke_agent_a2a(
 
         # 2. Convention-based forwarding: x-a2a-{agent_id_or_name}-{header_name}
         #    Matches both agent_id (UUID) and agent_name (alias), case-insensitive.
-        for alias in (agent.agent_id.lower(), agent.agent_name.lower()):
-            prefix = f"x-a2a-{alias}-"
-            for key, val in normalized.items():
-                if key.startswith(prefix):
-                    header_name = key[len(prefix) :]
-                    if header_name:
-                        dynamic_headers[header_name] = val
+        dynamic_headers.update(
+            _convention_headers(
+                normalized,
+                (agent.agent_id, agent.agent_name),
+                block_credentials=litellm_params.get("block_convention_credential_headers") is True,
+            )
+        )
 
         merged_agent_headers: Final = merge_agent_headers(
             dynamic_headers=dynamic_headers or None,

tests/unit/proxy/agent_endpoints/test_agent_headers.py

@@ -313,6 +313,75 @@ async def test_convention_unrelated_prefix_not_forwarded():
     assert headers == {"X-LiteLLM-User-Id": "u1"}
 
 
+@pytest.mark.parametrize(
+    "block_credentials, expected",
+    [
+        (
+            False,
+            {
+                "authorization": "Bearer caller",
+                "cookie": "session=caller",
+                "proxy-authorization": "Basic caller",
+                "x-api-key": "caller-key",
+                "x-request-id": "req-1",
+            },
+        ),
+        (True, {"x-request-id": "req-1"}),
+    ],
+)
+def test_convention_headers_drop_credentials_only_when_the_agent_opts_in(block_credentials, expected):
+    from litellm.proxy.agent_endpoints.a2a_endpoints import _convention_headers
+
+    headers: Final = {
+        "x-a2a-my-agent-authorization": "Bearer caller",
+        "x-a2a-my-agent-cookie": "session=caller",
+        "x-a2a-my-agent-proxy-authorization": "Basic caller",
+        "x-a2a-agent-123-x-api-key": "caller-key",
+        "x-a2a-my-agent-x-request-id": "req-1",
+        "x-a2a-other-agent-x-request-id": "not-for-this-agent",
+        "authorization": "Bearer sk-litellm-key",
+    }
+
+    forwarded: Final = _convention_headers(headers, ("agent-123", "my-agent"), block_credentials=block_credentials)
+    assert dict(forwarded) == expected
+
+
+@pytest.mark.asyncio
+async def test_agent_opted_in_never_receives_caller_credentials_via_convention():
+    mock_agent = _make_mock_agent()
+    mock_agent.agent_name = "my-agent"
+    mock_agent.litellm_params = {"block_convention_credential_headers": True}
+    mock_request = _make_mock_request(
+        extra_headers={
+            "x-a2a-my-agent-authorization": "Bearer forged-by-caller",
+            "x-a2a-my-agent-cookie": "session=forged-by-caller",
+            "x-a2a-my-agent-x-request-id": "req-abc",
+        }
+    )
+
+    mock_asend = await _invoke(mock_agent, mock_request, None)
+
+    assert mock_asend.call_args.kwargs.get("agent_extra_headers") == {
+        "x-request-id": "req-abc",
+        "X-LiteLLM-User-Id": "u1",
+    }
+
+
+@pytest.mark.asyncio
+async def test_agent_opted_in_still_sends_its_static_credentials():
+    mock_agent = _make_mock_agent(static_headers={"Authorization": "Bearer backend-static-secret"})
+    mock_agent.agent_name = "my-agent"
+    mock_agent.litellm_params = {"block_convention_credential_headers": True}
+    mock_request = _make_mock_request(extra_headers={"x-a2a-my-agent-authorization": "Bearer forged-by-caller"})
+
+    mock_asend = await _invoke(mock_agent, mock_request, None)
+
+    assert mock_asend.call_args.kwargs.get("agent_extra_headers") == {
+        "Authorization": "Bearer backend-static-secret",
+        "X-LiteLLM-User-Id": "u1",
+    }
+
+
 # ---------------------------------------------------------------------------
 # Databricks App OAuth M2M injection
 # ---------------------------------------------------------------------------