diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0254f79cbcc..6881956595c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import ( UpstreamCredentialProvider, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + prepare_mcp_client, raise_public, raise_token_exchange_challenge, raise_user_oauth_challenge, @@ -132,7 +133,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( from litellm.proxy._experimental.mcp_server.sampling_handler import ( MCP_SAMPLING_AVAILABLE, ) -from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client, validate_openapi_credentials from litellm.proxy._experimental.mcp_server.utils import ( MCP_TOOL_PREFIX_SEPARATOR, MCPMissingUserEnvVarsError, @@ -2805,6 +2805,8 @@ class MCPServerManager: headers=headers, server_label=server.name or server.server_name or server.alias or server.server_id, relays_upstream_auth=server.is_client_forwarded_token, + auth_type=server.auth_type, + upstream_token_header=server.upstream_token_header, ) tool_func.__name__ = prefixed_tool_name tool_func.__doc__ = description @@ -4230,19 +4232,16 @@ class MCPServerManager: ) record_auth_resolution(server.server_id, AuthResolution.not_applicable) - return await prepare_mcp_client( - resolved_server, - MCPClient( - server_url="", # Not used for stdio - transport_type=transport, - auth_type=resolved_server.auth_type, - auth_value=auth_value, - timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), - stdio_config=stdio_config, - extra_headers=extra_headers, - sampling_callback=sampling_cb, - elicitation_callback=elicitation_cb, - ), + return MCPClient( + server_url="", # Not used for stdio + transport_type=transport, + auth_type=resolved_server.auth_type, + auth_value=auth_value, + timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT), + stdio_config=stdio_config, + extra_headers=extra_headers, + sampling_callback=sampling_cb, + elicitation_callback=elicitation_cb, ) else: # For HTTP/SSE transports @@ -6200,7 +6199,6 @@ class MCPServerManager: mcp_auth_header: str | dict[str, str] | None, user_api_key_auth: UserAPIKeyAuth | None, forwarded_headers: dict[str, str] | None, - caller_authorization: str | None = None, ) -> tuple[dict[str, str] | None, dict[str, str] | None]: """Resolve the gateway-owned upstream credential for a spec_path (OpenAPI) tool call. @@ -6224,12 +6222,9 @@ class MCPServerManager: """ spec: Final = to_server_spec(mcp_server) if spec is None: - stored_headers = ( - None - if oauth2_headers - else await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) - ) - validate_openapi_credentials(mcp_server, stored_headers, forwarded_headers, caller_authorization) + if oauth2_headers: + return None, forwarded_headers + stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth) return stored_headers, forwarded_headers subject_token: str | None = None @@ -6248,9 +6243,7 @@ class MCPServerManager: user_api_key_auth=user_api_key_auth, extra_headers=forwarded_headers, ) - resolved_headers: Final = await _materialize_auth_headers(resolved_auth) - validate_openapi_credentials(mcp_server, resolved_headers, forwarded_headers, caller_authorization) - return resolved_headers, forwarded_headers + return await _materialize_auth_headers(resolved_auth), forwarded_headers async def _gather_openapi_tool_tasks( self, @@ -6376,7 +6369,6 @@ class MCPServerManager: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, - caller_authorization=auth_header_value, ) async def _call_openapi_via_handler(): diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index 66712e97a34..477d86ab436 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -20,7 +20,6 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( MCPOpenApiUpstreamError, MCPUpstreamAuthError, ) -from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers # Tool names emitted from OpenAPI specs must work across all major LLM providers. # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to @@ -55,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) -from litellm.types.mcp import credential_redirect_hook, custom_credential_slot +from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot class _OpenAPIJSONSchema(TypedDict, total=False): @@ -416,9 +415,26 @@ def _merge_openapi_tool_request_headers( Header names are compared case-insensitively so different casing cannot bypass the precedence rules. """ - return merge_openapi_headers( - static_headers, _request_extra_headers.get(), _request_auth_header.get(), _request_resolved_auth_headers.get() - ) + request_extra: Final = _request_extra_headers.get() or {} + static: Final = static_headers or {} + + static_lower_names: Final = {k.lower() for k in static} + effective_headers: dict[str, str] = {k: v for k, v in request_extra.items() if k.lower() not in static_lower_names} + effective_headers.update(static) + + override_auth: Final = _request_auth_header.get() + if override_auth: + for existing in [k for k in effective_headers if k.lower() == "authorization"]: + del effective_headers[existing] + effective_headers["Authorization"] = override_auth + + resolved_auth_headers: Final = _request_resolved_auth_headers.get() or {} + for name, value in resolved_auth_headers.items(): + for existing in [k for k in effective_headers if k.lower() == name.lower()]: + del effective_headers[existing] + effective_headers[name] = value + + return effective_headers def _raise_for_upstream_failure( @@ -455,6 +471,8 @@ def create_tool_function( headers: dict[str, str] | None = None, server_label: str | None = None, relays_upstream_auth: bool = False, + auth_type: MCPAuthType = None, + upstream_token_header: str | None = None, ): """Create a tool function for an OpenAPI operation. @@ -487,6 +505,18 @@ def create_tool_function( by using **kwargs instead of named parameters. """ effective_headers: Final = _merge_openapi_tool_request_headers(headers) + if auth_type is not None: + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( + raise_public, + validate_static_credential, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok + + match validate_static_credential(auth_type, effective_headers, upstream_token_header): + case Error(error): + raise_public(error) + case Ok(): + pass # Build URL from base_url and path url = base_url + path diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py index d25946d81d0..5358878a248 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py @@ -13,15 +13,17 @@ from __future__ import annotations import base64 import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Literal, NoReturn from fastapi import HTTPException from pydantic import SecretStr from typing_extensions import assert_never -from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials +from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource +from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( DEFAULT_CREDENTIAL_HEADER, ApiKeyConfig, @@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import ( Subject, TokenExchangeConfig, ) -from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth +from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport if TYPE_CHECKING: from litellm.proxy._types import UserAPIKeyAuth @@ -385,3 +387,64 @@ def raise_token_exchange_challenge( detail="Unauthorized", headers={"WWW-Authenticate": www_authenticate}, ) + + +_STATIC_MODES: Final = frozenset( + (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) +) + + +def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: + if not value: + return False + if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): + return True + if value.lower() in ("bearer", "basic", "token", "apikey"): + return False + if auth_type in (MCPAuth.bearer_token, MCPAuth.token): + scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" + credential: Final = strip_auth_scheme(value, scheme).strip() + return bool(credential) and credential.lower() != scheme.lower() + if auth_type == MCPAuth.basic: + parts: Final = value.split(None, 1) + if len(parts) != 2 or parts[0].lower() != "basic": + return False + try: + decoded: Final = base64.b64decode(parts[1], validate=True).strip() + return b":" in decoded + except ValueError: + return False + return True + + +def validate_static_credential( + auth_type: MCPAuthType, + headers: Mapping[str, str], + upstream_token_header: str | None = None, +) -> Result[None, CredError]: + if auth_type not in _STATIC_MODES: + return Ok(None) + default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization" + slots: Final = frozenset( + name.lower() + for name in ( + upstream_token_header or default_slot, + default_slot, + "Authorization", + ) + ) + values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) + if any(_usable_credential_value(auth_type, name, value) for name, value in values): + return Ok(None) + return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential")) + + +async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: + if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: + return client + request: Final = await client.prepare_request_auth() + match validate_static_credential(server.auth_type, request.headers, server.upstream_token_header): + case Error(error): + raise_public(error) + case Ok(): + return client diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a3aaada41f7..7feb1fd468d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -3141,7 +3141,6 @@ if MCP_AVAILABLE: mcp_auth_header=upstream_credential, user_api_key_auth=user_api_key_auth, forwarded_headers=openapi_forwarded_headers, - caller_authorization=auth_header_value, ) _auth_token: Final = _request_auth_header.set(auth_header_value) diff --git a/litellm/proxy/_experimental/mcp_server/upstream.py b/litellm/proxy/_experimental/mcp_server/upstream.py deleted file mode 100644 index 66840db21ce..00000000000 --- a/litellm/proxy/_experimental/mcp_server/upstream.py +++ /dev/null @@ -1,85 +0,0 @@ -from __future__ import annotations - -import base64 -from collections.abc import Mapping -from typing import Final - -from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme -from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import raise_public -from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result -from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError -from litellm.proxy._experimental.mcp_server.utils import merge_openapi_headers -from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport -from litellm.types.mcp_server.mcp_server_manager import MCPServer - -_STATIC_MODES: Final = frozenset( - (MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization) -) - - -def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool: - if not value: - return False - if auth_type == MCPAuth.authorization or (auth_type == MCPAuth.api_key and name != "authorization"): - return True - if value.lower() in ("bearer", "basic", "token", "apikey"): - return False - if auth_type in (MCPAuth.bearer_token, MCPAuth.token): - scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token" - credential: Final = strip_auth_scheme(value, scheme).strip() - return bool(credential) and credential.lower() != scheme.lower() - if auth_type == MCPAuth.basic: - parts: Final = value.split(None, 1) - if len(parts) != 2 or parts[0].lower() != "basic": - return False - try: - decoded: Final = base64.b64decode(parts[1], validate=True).strip() - return b":" in decoded - except ValueError: - return False - return True - - -def validate_static_credential(server: MCPServer, headers: Mapping[str, str]) -> Result[None, CredError]: - if server.auth_type not in _STATIC_MODES or server.transport == MCPTransport.stdio: - return Ok(None) - default_slot: Final = "X-API-Key" if server.auth_type == MCPAuth.api_key else "Authorization" - slots: Final = frozenset( - name.lower() - for name in ( - server.upstream_token_header or default_slot, - default_slot, - "Authorization", - ) - ) - values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots) - if any(_usable_credential_value(server.auth_type, name, value) for name, value in values): - return Ok(None) - return Error(CredError.of_misconfigured(f"{server.auth_type} requires a usable upstream credential")) - - -async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient: - if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio: - return client - request: Final = await client.prepare_request_auth() - match validate_static_credential(server, request.headers): - case Error(error): - raise_public(error) - case Ok(): - return client - - -def validate_openapi_credentials( - server: MCPServer, - resolved_headers: Mapping[str, str] | None, - forwarded_headers: Mapping[str, str] | None, - caller_authorization: str | None, -) -> None: - headers: Final = merge_openapi_headers( - server.static_headers or {}, forwarded_headers, caller_authorization, resolved_headers - ) - match validate_static_credential(server, headers): - case Error(error): - raise_public(error) - case Ok(): - return diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index bea74d36b34..fb3eb06fd15 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -756,22 +756,6 @@ def build_env_var_setup_url(server_id: str) -> str: return f"{base}{path}" if base else path -def merge_openapi_headers( - static_headers: Mapping[str, str], - extra_headers: Mapping[str, str] | None, - caller_authorization: str | None, - resolved_headers: Mapping[str, str] | None, -) -> dict[str, str]: - sources: Final = ( - extra_headers or {}, - static_headers, - {"Authorization": caller_authorization} if caller_authorization else {}, - resolved_headers or {}, - ) - entries: Final = {name.lower(): (name, value) for source in sources for name, value in source.items()} - return dict(entries.values()) - - def merge_mcp_headers( *, extra_headers: Mapping[str, str] | None = None, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 5e3a26fb4ac..28faf375ab8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1391,7 +1391,6 @@ class TestOpenApiResolvedUpstreamAuth: mcp_auth_header="user-byok-key", user_api_key_auth=UserAPIKeyAuth(user_id="alice", api_key="sk-user"), forwarded_headers=None, - caller_authorization="ApiKey user-byok-key", ) assert resolved is None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 54add273c24..4c4c45162ca 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5,11 +5,13 @@ import logging import os import sys from datetime import datetime +from pathlib import Path from typing import Any, Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from respx import MockRouter from litellm.proxy._experimental.mcp_server.exceptions import ( MCPServerListError, @@ -5127,7 +5129,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers captured["server_label"] = server_label @@ -5212,7 +5215,8 @@ class TestMCPServerManager: captured: dict = {} def fake_create_tool_function( - path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False + path, method, operation, base_url, headers=None, server_label=None, relays_upstream_auth=False, + auth_type=None, upstream_token_header=None, ): captured["headers"] = headers @@ -13471,6 +13475,41 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them( class TestProtectedCredentialPreparation: + @pytest.mark.asyncio + @pytest.mark.parametrize("auth_type,credential", [ + (MCPAuth.bearer_token, None), + (MCPAuth.bearer_token, "Bearer"), + (MCPAuth.api_key, None), + (MCPAuth.basic, "Basic"), + ]) + @pytest.mark.parametrize("dispatch", ["managed", "local"]) + async def test_openapi_dispatch_rejects_unusable_effective_credentials( + self, tmp_path: Path, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, credential: str | None, dispatch: str, + ) -> None: + from litellm.proxy._experimental.mcp_server.server import _handle_local_mcp_tool + from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_name, get_server_prefix + + spec_path: Final = tmp_path / "openapi.json" + spec_path.write_text(json.dumps({"openapi": "3.0.0", "info": {"title": "Auth", "version": "1"}, + "paths": {"/echo": {"get": {"operationId": "echo"}}}})) + server: Final = MCPServer( + server_id="dispatch-auth", name="dispatch-auth", url="https://upstream.example", + transport=MCPTransport.http, auth_type=auth_type, authentication_token=credential, + ) + manager: Final = MCPServerManager() + await manager._register_openapi_tools(str(spec_path), server, server.url) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="unexpected success") + result: Final = ( + await manager._call_openapi_tool_handler(server, "echo", {}) + if dispatch == "managed" + else await _handle_local_mcp_tool(add_server_prefix_to_name("echo", get_server_prefix(server)), {}) + ) + assert result.isError is True + assert "requires a usable upstream credential" in result.content[0].text + assert destination.call_count == 0 + @pytest.mark.asyncio @pytest.mark.parametrize("transport", [MCPTransport.http, MCPTransport.sse]) @pytest.mark.parametrize("client_secret", [None, ""]) @@ -13523,7 +13562,7 @@ class TestProtectedCredentialPreparation: assert client._get_auth_headers() == headers @pytest.mark.asyncio - @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.api_key, MCPAuth.bearer_token]) + @pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange]) async def test_openapi_protected_auth_rejects_missing_credentials(self, auth_type: MCPAuthType) -> None: server = MCPServer( server_id="openapi-empty", name="openapi-empty", url="https://upstream.example/mcp", @@ -13594,22 +13633,35 @@ class TestProtectedCredentialPreparation: ({"X-API-Key": "static"}, {"Authorization": ""}, None), ]) async def test_openapi_static_credentials_remain_supported( - self, static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None + self, respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None ) -> None: - server = MCPServer(server_id="openapi-static", name="openapi-static", url="https://upstream.example", - transport=MCPTransport.http, auth_type=MCPAuth.api_key, static_headers=static) - resolved, retained = await MCPServerManager().resolve_openapi_upstream_auth( - mcp_server=server, oauth2_headers=None, raw_headers=None, mcp_auth_header=None, - user_api_key_auth=None, forwarded_headers=forwarded, caller_authorization=caller, + from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( + _request_auth_header, _request_extra_headers, create_tool_function, ) - assert resolved is None - assert retained == forwarded + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.api_key, + ) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(caller) + extra_token: Final = _request_extra_headers.set(forwarded) + try: + assert await tool() == "authenticated" + sent: Final = destination.calls.last.request.headers + assert sent.get("x-api-key") == static.get("X-API-Key", (forwarded or {}).get("X-API-Key")) + if caller: + assert sent["authorization"] == caller + assert destination.call_count == 1 + finally: + _request_auth_header.reset(caller_token) + _request_extra_headers.reset(extra_token) @pytest.mark.asyncio async def test_static_resolution_cancellation_closes_flow(self) -> None: from collections.abc import AsyncGenerator from litellm.experimental_mcp_client.client import MCPClient - from litellm.proxy._experimental.mcp_server.upstream import prepare_mcp_client + from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import prepare_mcp_client class CancelledAuth(httpx.Auth): closed = False diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index 66c5627bc94..979199d0dc9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -10,9 +10,14 @@ This test suite ensures that: """ from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, patch import pytest +from fastapi import HTTPException +from respx import MockRouter + +from litellm.types.mcp import MCPAuth, MCPAuthType from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import ( _request_auth_header, @@ -35,6 +40,85 @@ from litellm.proxy._experimental.mcp_server.exceptions import ( GET_ASYNC_CLIENT_TARGET = "litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator.get_async_httpx_client" +@pytest.mark.asyncio +@pytest.mark.parametrize("static,forwarded,caller,resolved,expected", [ + ({"Authorization": "Bearer configured"}, {"authorization": "Bearer forwarded"}, None, None, "Bearer configured"), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", None, "Bearer caller"), + ({"Authorization": "Bearer configured"}, None, "Bearer", None, None), + ({"Authorization": "Bearer configured"}, None, "Bearer caller", {"authorization": " "}, None), + ({"Authorization": "Bearer configured"}, None, "Bearer", {"authorization": "Bearer resolved"}, "Bearer resolved"), +]) +async def test_static_auth_validates_headers_after_existing_precedence( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + static: dict[str, str], forwarded: dict[str, str] | None, caller: str | None, + resolved: dict[str, str] | None, expected: str | None, +) -> None: + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers=static, auth_type=MCPAuth.bearer_token, + ) + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + caller_token: Final = _request_auth_header.set(caller) + extra_token: Final = _request_extra_headers.set(forwarded) + resolved_token: Final = _request_resolved_auth_headers.set(resolved) + try: + if expected is None: + with pytest.raises(HTTPException, match="requires a usable upstream credential") as exc: + await tool() + assert exc.value.status_code == 500 + assert destination.call_count == 0 + else: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["authorization"] == expected + finally: + _request_auth_header.reset(caller_token) + _request_extra_headers.reset(extra_token) + _request_resolved_auth_headers.reset(resolved_token) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("credential", ["custom-key", ""]) +async def test_static_auth_uses_configured_custom_header( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, credential: str, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function( + "/echo", "get", {}, "https://upstream.example", headers={"x-custom": credential}, + auth_type=MCPAuth.api_key, upstream_token_header="X-Custom", + ) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="authenticated") + if credential: + assert await tool() == "authenticated" + assert destination.call_count == 1 + assert destination.calls.last.request.headers["x-custom"] == credential + else: + with pytest.raises(HTTPException, match="requires a usable upstream credential"): + await tool() + assert destination.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type,resolved", [ + (MCPAuth.none, None), + (MCPAuth.oauth2, {"Authorization": "Bearer user-oauth"}), +]) +async def test_static_validation_preserves_no_auth_and_resolved_oauth( + respx_mock: MockRouter, monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthType, resolved: dict[str, str] | None, +) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + tool: Final = create_tool_function("/echo", "get", {}, "https://upstream.example", auth_type=auth_type) + destination: Final = respx_mock.get("https://upstream.example/echo").respond(200, text="echo") + token: Final = _request_resolved_auth_headers.set(resolved) + try: + assert await tool() == "echo" + assert destination.call_count == 1 + assert destination.calls.last.request.headers.get("authorization") == (resolved or {}).get("Authorization") + finally: + _request_resolved_auth_headers.reset(token) + + def _create_mock_client(method: str, response_text: str, status_code: int = 200) -> AsyncMock: """Utility to create a mocked async httpx client for the given method.