Merge 551d7a8d4b into b1b236826a

litellm/litellm_core_utils/url_utils.py

@@ -467,6 +467,7 @@ def safe_get(client: _UrlFetcher, url: str, **kwargs: Unpack[_FetchOptions]) ->
 
     When ``litellm.user_url_validation`` is False, validation is bypassed
     and this function delegates to ``client.get(url, follow_redirects=True)``.
+    ``follow_redirects=False`` returns the first response without following it.
 
     Args:
         client: An httpx.Client (sync).
@@ -494,7 +495,7 @@ def safe_get(client: _UrlFetcher, url: str, **kwargs: Unpack[_FetchOptions]) ->
             "follow_redirects": False,
         }
         response = fetcher.get(validated_url, **hop_options)
-        if not response.is_redirect:
+        if not response.is_redirect or not kwargs.get("follow_redirects", True):
             return response
         # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so
         # relative Location headers keep the original hostname.
@@ -522,7 +523,7 @@ async def async_safe_get(client: _AsyncUrlFetcher, url: str, **kwargs: Unpack[_A
             "follow_redirects": False,
         }
         response = await fetcher.get(validated_url, **hop_options)
-        if not response.is_redirect:
+        if not response.is_redirect or not kwargs.get("follow_redirects", True):
             return response
         # Resolve the next hop against the ORIGINAL (pre-rewrite) URL so
         # relative Location headers keep the original hostname.

litellm/proxy/_lazy_openapi_snapshot.json

@@ -572,6 +572,18 @@
       "schemas": {
         "DiscoverAgentRequest": {
           "properties": {
+            "api_key": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "The agent's ``litellm_params.api_key``, sent as ``Authorization: Bearer <api_key>``.",
+              "title": "Api Key"
+            },
             "discovery_mode": {
               "$ref": "#/components/schemas/DiscoveryMode",
               "default": "well_known_fallback",
@@ -590,6 +602,21 @@
               "description": "Mode-specific parameters. ``langgraph_platform`` requires ``{'assistant_id': <id>}``. ``well_known_fallback`` ignores this.",
               "title": "Params"
             },
+            "static_headers": {
+              "anyOf": [
+                {
+                  "additionalProperties": {
+                    "type": "string"
+                  },
+                  "type": "object"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "The agent's ``static_headers``, combined with ``api_key`` the way the agent's runtime does.",
+              "title": "Static Headers"
+            },
             "url": {
               "description": "Base URL of the upstream agent. Behavior depends on ``discovery_mode``: ``well_known_fallback`` (default) tries /.well-known/agent-card.json, /.well-known/agent.json, /agent.json under this URL in order; ``langgraph_platform`` hits ``/.well-known/agent-card.json?assistant_id=<id>`` instead.",
               "title": "Url",

litellm/proxy/a2a/discovery.py

@@ -16,15 +16,17 @@ fetcher dispatches by ``discovery_mode``:
 
 from collections.abc import Mapping
 from enum import Enum
+from types import MappingProxyType
 from typing import Final
 from urllib.parse import urlencode
 
+import httpx
 from pydantic import TypeAdapter
 
 from litellm._logging import verbose_proxy_logger
 from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get
-from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
 from litellm.types.llms.custom_http import httpxSpecialProvider
 
 
@@ -56,6 +58,47 @@ class AgentCardDiscoveryError(Exception):
     """Raised when none of the well-known paths returned a usable agent card."""
 
 
+def card_request_headers(
+    api_key: str | None, static_headers: Mapping[str, str], discovery_mode: DiscoveryMode
+) -> Mapping[str, str]:
+    """Authenticate like the agent's runtime: a LangGraph key overrides a static Authorization header,
+    while an A2A agent's static Authorization header overrides the key."""
+    bearer: Final = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+    if bearer and discovery_mode is DiscoveryMode.LANGGRAPH_PLATFORM:
+        unauthorized: Final = {name: value for name, value in static_headers.items() if name.lower() != "authorization"}
+        return MappingProxyType({**unauthorized, **bearer})
+    sets_authorization: Final = any(name.lower() == "authorization" for name in static_headers)
+    return MappingProxyType({**({} if sets_authorization else bearer), **static_headers})
+
+
+def _is_card(response: httpx.Response) -> bool:
+    """Whether discovery accepts the response as an agent card."""
+    if response.status_code >= 400:
+        return False
+    try:
+        return isinstance(as_str_mapping(_DECODED_JSON.validate_python(response.json())), dict)
+    except (ValueError, RecursionError):
+        return False
+
+
+async def _get_card(
+    client: AsyncHTTPHandler, url: str, headers: Mapping[str, str], credentials: Mapping[str, str]
+) -> httpx.Response:
+    """Make discovery's usual request; when it yields no card, retry once with the agent's credentials added,
+    never following a redirect with them, and keep the first answer unless the retry yields a card."""
+    response: Final = await async_safe_get(client, url, headers=dict(headers))
+    if not credentials or _is_card(response):
+        return response
+    try:
+        credentialed: Final = await async_safe_get(
+            client, url, headers={**headers, **credentials}, follow_redirects=False
+        )
+    except (httpx.HTTPError, ValueError) as exc:  # ValueError: the aiohttp transport rejecting a header value
+        verbose_proxy_logger.debug("A2A discovery with credentials raised %s", type(exc).__name__)
+        return response
+    return credentialed if _is_card(credentialed) and not credentialed.is_redirect else response
+
+
 def _normalize_base_url(base_url: str) -> str:
     return base_url.rstrip("/")
 
@@ -92,11 +135,13 @@ async def fetch_well_known_card(
     params: Mapping[str, object] | None = None,
     timeout: float = DEFAULT_DISCOVERY_TIMEOUT_SECONDS,
     headers: dict[str, str] | None = None,
+    credentials: Mapping[str, str] | None = None,
 ) -> dict[str, object]:
     """
     Fetch an agent card from ``base_url`` using the strategy chosen by
     ``discovery_mode``. Returns the parsed JSON from the first path that
-    responds with a JSON body.
+    responds with a JSON body. ``headers`` go with every request;
+    ``credentials`` are added only to retry a path that yielded no card.
 
     Raises:
         AgentCardDiscoveryError: if every path fails (network error, non-2xx,
@@ -126,7 +171,7 @@ async def fetch_well_known_card(
             # returns ``None`` when the key is present-but-None, then crashes
             # on ``{**None, "Host": ...}``. Default the kwarg to an empty
             # dict so production (``user_url_validation=True``) doesn't 500.
-            response = await async_safe_get(client, url, headers=headers or {})
+            response = await _get_card(client, url, headers or {}, credentials or {})
         except SSRFError as exc:
             last_error = f"{url}: {exc}"
             verbose_proxy_logger.debug("A2A discovery blocked by SSRF guard for %s: %s", url, exc)

litellm/proxy/a2a/endpoints.py

@@ -9,6 +9,7 @@ admin pick which ones to expose through the proxy. The actual merge into a
 LiteLLM-fronted card happens when the agent is saved via ``POST /v1/agents``.
 """
 
+from collections.abc import Mapping
 from typing import Final
 
 from fastapi import APIRouter, Depends, HTTPException
@@ -21,6 +22,7 @@ from litellm.proxy.a2a.discovery import (
     AGENT_CARD_WELL_KNOWN_PATHS,
     AgentCardDiscoveryError,
     DiscoveryMode,
+    card_request_headers,
     fetch_well_known_card,
 )
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@@ -57,6 +59,14 @@ class DiscoverAgentRequest(LiteLLMBaseModel):
             "``{'assistant_id': <id>}``. ``well_known_fallback`` ignores this."
         ),
     )
+    api_key: str | None = Field(
+        default=None,
+        description="The agent's ``litellm_params.api_key``, sent as ``Authorization: Bearer <api_key>``.",
+    )
+    static_headers: Mapping[str, str] | None = Field(
+        default=None,
+        description="The agent's ``static_headers``, combined with ``api_key`` the way the agent's runtime does.",
+    )
 
 
 class DiscoverAgentResponse(LiteLLMBaseModel):
@@ -100,6 +110,7 @@ async def discover_agent_card(
             request.url,
             discovery_mode=request.discovery_mode,
             params=request.params,
+            credentials=card_request_headers(request.api_key, request.static_headers or {}, request.discovery_mode),
         )
     except AgentCardDiscoveryError as exc:
         raise HTTPException(status_code=400, detail=str(exc))

tests/unit/litellm_core_utils/test_url_utils.py

@@ -1,7 +1,8 @@
 import asyncio
 import socket
 import threading
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
+from typing import Final
 from urllib.parse import urlsplit
 
 import httpx
@@ -320,6 +321,48 @@ class TestRedirectHostnamePreservation:
         assert hops[1]["url"].endswith("/redirected")
 
 
+def _redirect(url: str) -> httpx.Response:
+    return httpx.Response(
+        302, headers={"Location": "https://elsewhere.example/next"}, request=httpx.Request("GET", url)
+    )
+
+
+class _RedirectingClient:
+    def get(self, url: str, headers: Mapping[str, str] | None = None, follow_redirects: bool = True) -> httpx.Response:
+        return _redirect(url)
+
+
+class _AsyncRedirectingClient:
+    async def get(
+        self, url: str, headers: Mapping[str, str] | None = None, follow_redirects: bool = True
+    ) -> httpx.Response:
+        return _redirect(url)
+
+
+def _public_dns(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        url_utils.socket,
+        "getaddrinfo",
+        lambda host, port, *a, **kw: [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port))],
+    )
+
+
+def test_safe_get_returns_the_redirect_when_told_not_to_follow(monkeypatch: pytest.MonkeyPatch) -> None:
+    _public_dns(monkeypatch)
+    response: Final = url_utils.safe_get(_RedirectingClient(), "https://agent.example/card", follow_redirects=False)
+
+    assert response.status_code == 302
+
+
+def test_async_safe_get_returns_the_redirect_when_told_not_to_follow(monkeypatch: pytest.MonkeyPatch) -> None:
+    _public_dns(monkeypatch)
+    response: Final = asyncio.run(
+        url_utils.async_safe_get(_AsyncRedirectingClient(), "https://agent.example/card", follow_redirects=False)
+    )
+
+    assert response.status_code == 302
+
+
 class TestValidationMasterSwitch:
     def test_disabled_bypasses_fetch_in_safe_get(self, monkeypatch):
         """When user_url_validation is False, safe_get delegates to client.get without validation."""

tests/unit/proxy/a2a/test_discovery.py

@@ -1,8 +1,12 @@
 """Tests for the well-known card fetcher and the discovery endpoint."""
 
+from collections.abc import Mapping
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
+import respx
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
@@ -12,6 +16,7 @@ from litellm.proxy.a2a.discovery import (
     AGENT_CARD_WELL_KNOWN_PATHS,
     AgentCardDiscoveryError,
     DiscoveryMode,
+    card_request_headers,
     fetch_well_known_card,
 )
 from litellm.proxy.a2a.endpoints import router as a2a_router
@@ -281,3 +286,213 @@ def test_discover_rejects_unknown_mode():
         json={"url": "http://localhost:2024", "discovery_mode": "bogus"},
     )
     assert resp.status_code == 422
+
+
+# ---------------------------------------------------------------------------
+# Discovery authenticates like the saved agent
+# ---------------------------------------------------------------------------
+
+_UPSTREAM: Final = "https://upstream.example"
+_CARD: Final = {"name": "upstream-agent"}
+_BASIC: Final = {"authorization": "Basic dXNlcjpwdw=="}
+
+
+@pytest.mark.parametrize(
+    ("api_key", "static_headers", "mode", "expected"),
+    [
+        (
+            "sk-typed",
+            {"x-tenant": "t1"},
+            DiscoveryMode.WELL_KNOWN_FALLBACK,
+            {"Authorization": "Bearer sk-typed", "x-tenant": "t1"},
+        ),
+        ("sk-typed", _BASIC, DiscoveryMode.WELL_KNOWN_FALLBACK, _BASIC),
+        (
+            "sk-typed",
+            {**_BASIC, "x-tenant": "t1"},
+            DiscoveryMode.LANGGRAPH_PLATFORM,
+            {"x-tenant": "t1", "Authorization": "Bearer sk-typed"},
+        ),
+        (None, _BASIC, DiscoveryMode.LANGGRAPH_PLATFORM, _BASIC),
+        ("", {}, DiscoveryMode.WELL_KNOWN_FALLBACK, {}),
+    ],
+)
+def test_card_request_headers_authenticate_like_the_agents_runtime(
+    api_key: str | None, static_headers: Mapping[str, str], mode: DiscoveryMode, expected: Mapping[str, str]
+) -> None:
+    assert dict(card_request_headers(api_key, static_headers, mode)) == expected
+
+
+def _card_behind_auth(request: httpx.Request) -> httpx.Response:
+    if "authorization" in request.headers:
+        return httpx.Response(200, json=_CARD)
+    return httpx.Response(401, json={"detail": "Unauthorized"})
+
+
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+def test_discover_sends_the_key_and_static_headers_upstream() -> None:
+    route: Final = respx.get(f"{_UPSTREAM}/.well-known/agent-card.json").mock(side_effect=_card_behind_auth)
+
+    response: Final = _client_for_role(LitellmUserRoles.PROXY_ADMIN).post(
+        "/v1/a2a/discover",
+        json={"url": _UPSTREAM, "api_key": "sk-typed", "static_headers": {"x-tenant": "t1"}},
+    )
+
+    assert (response.status_code, response.json()["agent_card"]) == (200, _CARD)
+    sent: Final = [(call.request.headers.get("authorization"), call.request.headers.get("x-tenant")) for call in route.calls]
+    assert sent == [(None, None), ("Bearer sk-typed", "t1")]
+
+
+def _public_card_unless_authorized(request: httpx.Request) -> httpx.Response:
+    return httpx.Response(200, json=_CARD if "authorization" in request.headers else {"name": "public-agent"})
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+async def test_fetch_sends_its_headers_with_every_request_as_before() -> None:
+    route: Final = respx.get(f"{_UPSTREAM}/.well-known/agent-card.json").mock(side_effect=_public_card_unless_authorized)
+
+    card: Final = await fetch_well_known_card(_UPSTREAM, headers={"Authorization": "Bearer sk-1"})
+
+    assert card == _CARD
+    assert [call.request.headers.get("authorization") for call in route.calls] == ["Bearer sk-1"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+async def test_fetch_keeps_its_headers_on_the_credentialed_retry() -> None:
+    route: Final = respx.get(f"{_UPSTREAM}/.well-known/agent-card.json").mock(side_effect=_card_behind_auth)
+
+    card: Final = await fetch_well_known_card(
+        _UPSTREAM, headers={"x-tenant": "t1"}, credentials={"Authorization": "Bearer sk-1"}
+    )
+
+    assert card == _CARD
+    sent: Final = [(call.request.headers.get("authorization"), call.request.headers.get("x-tenant")) for call in route.calls]
+    assert sent == [(None, "t1"), ("Bearer sk-1", "t1")]
+
+
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+def test_discover_returns_the_card_the_credentials_unlock_before_a_later_public_one() -> None:
+    """The saved agent's own client also asks this path first, with its credentials."""
+    respx.get(f"{_UPSTREAM}/.well-known/agent-card.json").mock(side_effect=_card_behind_auth)
+    respx.get(f"{_UPSTREAM}/.well-known/agent.json").respond(200, json={"name": "public-agent"})
+
+    response: Final = _client_for_role(LitellmUserRoles.PROXY_ADMIN).post(
+        "/v1/a2a/discover", json={"url": _UPSTREAM, "api_key": "sk-typed"}
+    )
+
+    assert (response.status_code, response.json()["agent_card"]) == (200, _CARD)
+
+
+_LOCAL_AGENT: Final = "http://127.0.0.1:8123"
+_OTHER_ORIGIN: Final = "http://127.0.0.1:8124"
+
+
+@pytest.mark.parametrize("validate_urls", [True, False])
+@pytest.mark.usefixtures("httpx_transport")
+def test_discover_never_sends_credentials_to_a_redirect_target(
+    monkeypatch: pytest.MonkeyPatch, validate_urls: bool
+) -> None:
+    monkeypatch.setattr(litellm, "user_url_validation", validate_urls)
+    monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["127.0.0.1:8123", "127.0.0.1:8124"])
+
+    def agent(request: httpx.Request) -> httpx.Response:
+        if "x-api-key" in request.headers:
+            return httpx.Response(302, headers={"Location": f"{_OTHER_ORIGIN}/card"})
+        return httpx.Response(401)
+
+    with respx.mock(assert_all_called=False) as router:
+        router.get(url__startswith=_LOCAL_AGENT).mock(side_effect=agent)
+        elsewhere: Final = router.get(url__startswith=_OTHER_ORIGIN).respond(200, json=_CARD)
+        response: Final = _client_for_role(LitellmUserRoles.PROXY_ADMIN).post(
+            "/v1/a2a/discover", json={"url": _LOCAL_AGENT, "static_headers": {"x-api-key": "secret"}}
+        )
+
+    assert response.status_code == 400
+    assert elsewhere.call_count == 0
+
+
+@pytest.mark.parametrize("status", [200, 203])
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+def test_discover_loads_a_public_card_exactly_as_before_without_sending_credentials(status: int) -> None:
+    route: Final = respx.get(f"{_UPSTREAM}/.well-known/agent-card.json").respond(status, json=_CARD)
+
+    response: Final = _client_for_role(LitellmUserRoles.PROXY_ADMIN).post(
+        "/v1/a2a/discover", json={"url": _UPSTREAM, "api_key": "sk-1", "static_headers": {"x-tenant": "t1"}}
+    )
+
+    assert (response.status_code, response.json()["agent_card"]) == (200, _CARD)
+    assert [dict(call.request.headers).keys() & {"authorization", "x-tenant"} for call in route.calls] == [set()]
+
+
+@pytest.mark.parametrize("anonymous_body", ["<html>sign in</html>", "[" * 100_000 + "]" * 100_000])
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+def test_discover_sends_credentials_when_the_anonymous_answer_is_not_a_card(anonymous_body: str) -> None:
+    respx.get(f"{_UPSTREAM}/.well-known/agent-card.json").mock(
+        side_effect=lambda request: (
+            httpx.Response(200, json=_CARD)
+            if "authorization" in request.headers
+            else httpx.Response(200, text=anonymous_body)
+        )
+    )
+
+    response: Final = _client_for_role(LitellmUserRoles.PROXY_ADMIN).post(
+        "/v1/a2a/discover", json={"url": _UPSTREAM, "api_key": "sk-1"}
+    )
+
+    assert (response.status_code, response.json()["agent_card"]) == (200, _CARD)
+
+
+@pytest.mark.parametrize(
+    "rejection",
+    [
+        httpx.Response(401),
+        httpx.Response(200, text="<html>sign in</html>"),
+        httpx.Response(200, json=["not", "a card"]),
+        httpx.Response(302, headers={"Location": f"{_UPSTREAM}/elsewhere"}, json={"moved": "elsewhere"}),
+        httpx.LocalProtocolError("Illegal header value"),
+        ValueError("Newline in header value"),
+        httpx.ReadTimeout("timed out"),
+    ],
+)
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+def test_discover_reports_the_anonymous_answer_when_the_credentials_fail(rejection: httpx.Response | Exception) -> None:
+    def card(request: httpx.Request) -> httpx.Response:
+        if "authorization" not in request.headers:
+            return httpx.Response(403)
+        if isinstance(rejection, Exception):
+            raise rejection
+        return rejection
+
+    respx.route().mock(side_effect=card)
+
+    response: Final = _client_for_role(LitellmUserRoles.PROXY_ADMIN).post(
+        "/v1/a2a/discover", json={"url": _UPSTREAM, "api_key": "sk-1"}
+    )
+
+    assert (response.status_code, response.json()["detail"]) == (
+        400,
+        "Could not fetch agent card from https://upstream.example (mode=well_known_fallback). "
+        "Last error: https://upstream.example/agent.json: HTTP 403",
+    )
+
+
+@pytest.mark.usefixtures("httpx_transport")
+@respx.mock
+def test_discover_does_not_wait_twice_for_an_unresponsive_agent() -> None:
+    timing_out: Final = respx.route().mock(side_effect=httpx.ReadTimeout("timed out"))
+
+    response: Final = _client_for_role(LitellmUserRoles.PROXY_ADMIN).post(
+        "/v1/a2a/discover", json={"url": _UPSTREAM, "api_key": "sk-1"}
+    )
+
+    assert response.status_code == 400
+    assert timing_out.call_count == len(AGENT_CARD_WELL_KNOWN_PATHS)

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -32932,6 +32932,11 @@ export interface components {
         };
         /** DiscoverAgentRequest */
         DiscoverAgentRequest: {
+            /**
+             * Api Key
+             * @description The agent's ``litellm_params.api_key``, sent as ``Authorization: Bearer <api_key>``.
+             */
+            api_key?: string | null;
             /**
              * @description How to locate the upstream card. ``well_known_fallback`` for pure A2A agents (try standard paths); ``langgraph_platform`` for LangGraph Platform deployments where the card is shared across assistants and disambiguated by a query parameter.
              * @default well_known_fallback
@@ -32944,6 +32949,13 @@ export interface components {
             params?: {
                 [key: string]: unknown;
             } | null;
+            /**
+             * Static Headers
+             * @description The agent's ``static_headers``, combined with ``api_key`` the way the agent's runtime does.
+             */
+            static_headers?: {
+                [key: string]: string;
+            } | null;
             /**
              * Url
              * @description Base URL of the upstream agent. Behavior depends on ``discovery_mode``: ``well_known_fallback`` (default) tries /.well-known/agent-card.json, /.well-known/agent.json, /agent.json under this URL in order; ``langgraph_platform`` hits ``/.well-known/agent-card.json?assistant_id=<id>`` instead.