From 551d7a8d4b32358e5a70cd0b50c212f8edaf03de Mon Sep 17 00:00:00 2001 From: shrey kharbanda Date: Tue, 3 Nov 2026 10:35:09 -0800 Subject: [PATCH] feat(a2a): let card discovery send the agent's credentials POST /v1/a2a/discover now takes the agent's api_key and static_headers and combines them the way the agent's runtime does. Each card path still gets the anonymous request discovery always made. Only when that is answered without the card does discovery retry once with the credentials, never following a redirect with them, so every card that loaded before loads the same way and its errors read the same. --- litellm/litellm_core_utils/url_utils.py | 5 +- litellm/proxy/_lazy_openapi_snapshot.json | 27 +++ litellm/proxy/a2a/discovery.py | 51 ++++- litellm/proxy/a2a/endpoints.py | 11 + .../unit/litellm_core_utils/test_url_utils.py | 45 +++- tests/unit/proxy/a2a/test_discovery.py | 215 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 + 7 files changed, 360 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0db6a41b503..70ad98bb1db 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -459,6 +459,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). @@ -486,7 +487,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. @@ -514,7 +515,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. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5cc1c53d396..85978195184 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/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 ``.", + "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': }``. ``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=`` instead.", "title": "Url", diff --git a/litellm/proxy/a2a/discovery.py b/litellm/proxy/a2a/discovery.py index 03d1be7d185..fe95fe7db40 100644 --- a/litellm/proxy/a2a/discovery.py +++ b/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) diff --git a/litellm/proxy/a2a/endpoints.py b/litellm/proxy/a2a/endpoints.py index 730d40b46d0..fd86ce77305 100644 --- a/litellm/proxy/a2a/endpoints.py +++ b/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': }``. ``well_known_fallback`` ignores this." ), ) + api_key: str | None = Field( + default=None, + description="The agent's ``litellm_params.api_key``, sent as ``Authorization: Bearer ``.", + ) + 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)) diff --git a/tests/unit/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py index 9b29eac208e..98f069f2ec0 100644 --- a/tests/unit/litellm_core_utils/test_url_utils.py +++ b/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 @@ -306,6 +307,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.""" diff --git a/tests/unit/proxy/a2a/test_discovery.py b/tests/unit/proxy/a2a/test_discovery.py index ac1e7dfbb56..325e4737c43 100644 --- a/tests/unit/proxy/a2a/test_discovery.py +++ b/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", ["sign in", "[" * 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="sign in"), + 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) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 66fe9b17a82..265a49a37cb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -32338,6 +32338,11 @@ export interface components { }; /** DiscoverAgentRequest */ DiscoverAgentRequest: { + /** + * Api Key + * @description The agent's ``litellm_params.api_key``, sent as ``Authorization: Bearer ``. + */ + 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 @@ -32350,6 +32355,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=`` instead.