diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index eae8ac04e8e..6079dcc3d32 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/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. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 5cb63c98d56..66024c758b7 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 958960af27a..6ddb7df0a0d 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 @@ -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.""" 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 2272b4bea43..2966b582e31 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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?: 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=`` instead.