mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
fix(mcp): reject missing upstream authentication credentials
This commit is contained in:
parent
0b3e56448f
commit
13553473aa
12 changed files with 396 additions and 86 deletions
|
|
@ -346,13 +346,17 @@ class MCPClient:
|
|||
self.update_auth_value(auth_value)
|
||||
|
||||
async def discovery_auth_fingerprint(self) -> str:
|
||||
return self._hash_discovery_auth(await self.prepare_request_auth())
|
||||
|
||||
async def prepare_request_auth(self) -> httpx.Request:
|
||||
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
|
||||
request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
|
||||
if self._resolved_auth is None:
|
||||
return self._hash_discovery_auth(request)
|
||||
return request
|
||||
flow: Final = self._resolved_auth.async_auth_flow(request)
|
||||
try:
|
||||
authenticated: Final = await flow.__anext__()
|
||||
return self._hash_discovery_auth(authenticated)
|
||||
return authenticated
|
||||
finally:
|
||||
await flow.aclose()
|
||||
|
||||
|
|
|
|||
|
|
@ -132,6 +132,7 @@ 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,
|
||||
|
|
@ -4229,16 +4230,19 @@ class MCPServerManager:
|
|||
)
|
||||
|
||||
record_auth_resolution(server.server_id, AuthResolution.not_applicable)
|
||||
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,
|
||||
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,
|
||||
),
|
||||
)
|
||||
else:
|
||||
# For HTTP/SSE transports
|
||||
|
|
@ -4259,15 +4263,20 @@ class MCPServerManager:
|
|||
user_api_key_auth=user_api_key_auth,
|
||||
extra_headers=extra_headers,
|
||||
)
|
||||
return MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
|
||||
extra_headers=extra_headers,
|
||||
resolved_auth=resolved_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
return await prepare_mcp_client(
|
||||
resolved_server,
|
||||
MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
timeout=(
|
||||
resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT
|
||||
),
|
||||
extra_headers=extra_headers,
|
||||
resolved_auth=resolved_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
),
|
||||
)
|
||||
|
||||
# Create SigV4 auth if configured
|
||||
|
|
@ -4297,17 +4306,20 @@ class MCPServerManager:
|
|||
else AuthResolution.no_auth
|
||||
)
|
||||
record_auth_resolution(server.server_id, legacy_source)
|
||||
return MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
auth_value=auth_value,
|
||||
auth_header_name=auth_header_name,
|
||||
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
|
||||
extra_headers=extra_headers,
|
||||
aws_auth=aws_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
return await prepare_mcp_client(
|
||||
resolved_server,
|
||||
MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
auth_value=auth_value,
|
||||
auth_header_name=auth_header_name,
|
||||
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
|
||||
extra_headers=extra_headers,
|
||||
aws_auth=aws_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
),
|
||||
)
|
||||
|
||||
async def _get_tools_from_server(
|
||||
|
|
@ -6188,6 +6200,7 @@ 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.
|
||||
|
||||
|
|
@ -6211,9 +6224,12 @@ class MCPServerManager:
|
|||
"""
|
||||
spec: Final = to_server_spec(mcp_server)
|
||||
if spec is None:
|
||||
if oauth2_headers:
|
||||
return None, forwarded_headers
|
||||
stored_headers = await self._resolve_oauth2_headers_for_tool_call(mcp_server, None, user_api_key_auth)
|
||||
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)
|
||||
return stored_headers, forwarded_headers
|
||||
|
||||
subject_token: str | None = None
|
||||
|
|
@ -6232,7 +6248,9 @@ class MCPServerManager:
|
|||
user_api_key_auth=user_api_key_auth,
|
||||
extra_headers=forwarded_headers,
|
||||
)
|
||||
return await _materialize_auth_headers(resolved_auth), 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
|
||||
|
||||
async def _gather_openapi_tool_tasks(
|
||||
self,
|
||||
|
|
@ -6358,6 +6376,7 @@ 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():
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ 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
|
||||
|
|
@ -415,26 +416,9 @@ def _merge_openapi_tool_request_headers(
|
|||
Header names are compared case-insensitively so different casing cannot
|
||||
bypass the precedence rules.
|
||||
"""
|
||||
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
|
||||
return merge_openapi_headers(
|
||||
static_headers, _request_extra_headers.get(), _request_auth_header.get(), _request_resolved_auth_headers.get()
|
||||
)
|
||||
|
||||
|
||||
def _raise_for_upstream_failure(
|
||||
|
|
|
|||
|
|
@ -79,7 +79,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
|
|||
|
||||
BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just
|
||||
like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers
|
||||
to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later).
|
||||
to v1 for its static schemes. Declared OBO always stays with the exchange arm.
|
||||
|
||||
Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
|
||||
an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
|
||||
|
|
@ -90,8 +90,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
|
|||
modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough
|
||||
oauth2 and SigV4 return None and stay on v1.
|
||||
"""
|
||||
if server.is_byok:
|
||||
return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type)
|
||||
if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange:
|
||||
return None # per-user BYOK source not migrated yet -> defer to v1
|
||||
resource: Final = server.url or server.server_id
|
||||
auth_type: Final = server.auth_type
|
||||
match auth_type:
|
||||
|
|
@ -165,21 +165,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
|
|||
)
|
||||
|
||||
|
||||
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
|
||||
"""Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured.
|
||||
|
||||
An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the
|
||||
``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at
|
||||
the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the
|
||||
gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is
|
||||
nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect
|
||||
(``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value
|
||||
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
|
||||
forwarded only when the operator set it; a missing one is omitted, not derived.
|
||||
"""
|
||||
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec:
|
||||
"""Keep declared OBO owned by the resolver, including incomplete client configuration."""
|
||||
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
|
||||
if not server.client_id or not server.client_secret:
|
||||
return None
|
||||
profile: Final[Literal["rfc8693", "entra_obo"]] = (
|
||||
"entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693"
|
||||
)
|
||||
|
|
@ -193,7 +181,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
|
|||
token_exchange_endpoint=endpoint,
|
||||
audience=server.audience,
|
||||
client_id=server.client_id,
|
||||
client_secret=SecretStr(server.client_secret),
|
||||
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
|
||||
token_endpoint_auth_method=server.token_endpoint_auth_method,
|
||||
scopes=tuple(server.scopes or ()),
|
||||
),
|
||||
|
|
|
|||
|
|
@ -3141,6 +3141,7 @@ 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)
|
||||
|
|
|
|||
81
litellm/proxy/_experimental/mcp_server/upstream.py
Normal file
81
litellm/proxy/_experimental/mcp_server/upstream.py
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.experimental_mcp_client.client import MCPClient
|
||||
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 == MCPAuth.basic:
|
||||
parts: Final = value.split(None, 1)
|
||||
if len(parts) != 2 or parts[0].lower() != "basic":
|
||||
return False
|
||||
try:
|
||||
return bool(base64.b64decode(parts[1], validate=True).strip())
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def validate_static_credential(
|
||||
server: MCPServer, headers: Mapping[str, str], *, header_slot: str | None = None, openapi: bool = False
|
||||
) -> 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 (
|
||||
header_slot or server.upstream_token_header or default_slot,
|
||||
"Authorization" if openapi else default_slot,
|
||||
)
|
||||
)
|
||||
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
|
||||
if values and all(_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, openapi=True):
|
||||
case Error(error):
|
||||
raise_public(error)
|
||||
case Ok():
|
||||
return
|
||||
|
|
@ -756,6 +756,22 @@ 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,
|
||||
|
|
|
|||
|
|
@ -1934,3 +1934,18 @@ async def test_discovery_auth_fingerprint_tracks_effective_credentials(resolved:
|
|||
assert original != replaced
|
||||
assert len(original) == 64
|
||||
assert "private-original-credential" not in original
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_auth_preview_uses_the_same_effective_headers_as_egress() -> None:
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import StaticHeaderAuth
|
||||
|
||||
client: Final = MCPClient(
|
||||
server_url="https://upstream.example/mcp", auth_type=MCPAuth.bearer_token,
|
||||
resolved_auth=StaticHeaderAuth("Bearer resolved"), extra_headers={"X-Trace": "trace"},
|
||||
)
|
||||
request: Final = await client.prepare_request_auth()
|
||||
assert request.method == "POST"
|
||||
assert str(request.url) == "https://upstream.example/mcp"
|
||||
assert request.headers["Authorization"] == "Bearer resolved"
|
||||
assert request.headers["X-Trace"] == "trace"
|
||||
|
|
|
|||
|
|
@ -155,12 +155,6 @@ def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow):
|
|||
_server(auth_type=MCPAuth.api_key), # no token configured
|
||||
_server(auth_type=MCPAuth.bearer_token), # no token configured
|
||||
_server(auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True), # delegated upstream OAuth -> v1
|
||||
_server(auth_type=MCPAuth.oauth2_token_exchange), # no endpoint/client creds -> incomplete -> v1
|
||||
_server(
|
||||
auth_type=MCPAuth.oauth2_token_exchange,
|
||||
token_exchange_endpoint="https://idp/token",
|
||||
client_id="cid",
|
||||
), # missing client_secret -> incomplete -> v1
|
||||
_server(auth_type=MCPAuth.aws_sigv4),
|
||||
_server(auth_type=None, oauth_passthrough=True, extra_headers=["Authorization"]),
|
||||
],
|
||||
|
|
@ -802,3 +796,14 @@ def test_a_blank_header_name_means_unset_rather_than_an_error(blank):
|
|||
spec = to_server_spec(server)
|
||||
assert spec is not None
|
||||
assert spec.config.header_name == "Authorization"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client_secret", [None, ""])
|
||||
@pytest.mark.parametrize("is_byok", [False, True])
|
||||
def test_incomplete_obo_keeps_exchange_ownership(client_secret: str | None, is_byok: bool) -> None:
|
||||
spec = to_server_spec(_server(auth_type=MCPAuth.oauth2_token_exchange, client_id="client",
|
||||
client_secret=client_secret, is_byok=is_byok))
|
||||
assert spec is not None
|
||||
assert isinstance(spec.config, TokenExchangeConfig)
|
||||
assert spec.config.client_id == "client"
|
||||
assert spec.config.client_secret is None
|
||||
|
|
|
|||
|
|
@ -1391,6 +1391,7 @@ 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
|
||||
|
|
|
|||
|
|
@ -9401,12 +9401,13 @@ class TestCreateMcpClientV2Graft:
|
|||
assert "misconfigured" in str(exc_info.value.detail)
|
||||
assert "token_url" in str(exc_info.value.detail)
|
||||
|
||||
async def test_static_token_missing_defers_to_v1(self):
|
||||
client = await MCPServerManager()._create_mcp_client(
|
||||
self._http_server(auth_type=MCPAuth.api_key, authentication_token=None)
|
||||
)
|
||||
|
||||
assert client._resolved_auth is None
|
||||
async def test_static_token_missing_rejects_before_connecting(self):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MCPServerManager()._create_mcp_client(
|
||||
self._http_server(auth_type=MCPAuth.api_key, authentication_token=None)
|
||||
)
|
||||
assert exc.value.status_code == 500
|
||||
assert "credential" in str(exc.value.detail)
|
||||
|
||||
async def test_stdio_migrated_auth_type_still_defers_to_v1(self):
|
||||
client = await MCPServerManager()._create_mcp_client(
|
||||
|
|
@ -13467,3 +13468,180 @@ async def test_discovery_cache_returns_oversized_results_without_retaining_them(
|
|||
result: Final = await cache.get(("server", None), fetch)
|
||||
assert result[0].description == description
|
||||
assert fetch.await_count == 2
|
||||
|
||||
|
||||
class TestProtectedCredentialPreparation:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("transport", [MCPTransport.http, MCPTransport.sse])
|
||||
@pytest.mark.parametrize("client_secret", [None, ""])
|
||||
@pytest.mark.parametrize("subject", [None, "caller-subject"])
|
||||
async def test_incomplete_obo_rejects_caller_and_static_fallback(
|
||||
self, transport: MCPTransport, client_secret: str | None, subject: str | None
|
||||
) -> None:
|
||||
server = MCPServer(
|
||||
server_id="incomplete-obo", name="incomplete-obo", url="https://upstream.example/mcp",
|
||||
transport=transport, auth_type=MCPAuth.oauth2_token_exchange,
|
||||
client_id="gateway", client_secret=client_secret,
|
||||
token_exchange_endpoint="https://idp.example/token", authentication_token="static-fallback",
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MCPServerManager()._create_mcp_client(
|
||||
server, mcp_auth_header="Bearer override", subject_token=subject,
|
||||
)
|
||||
assert exc.value.status_code == (401 if subject is None else 500)
|
||||
assert "static-fallback" not in str(exc.value.detail)
|
||||
assert "override" not in str(exc.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("auth_type", [MCPAuth.api_key, MCPAuth.bearer_token])
|
||||
@pytest.mark.parametrize("credential", [None, "", " ", {"X-Trace": "trace"}])
|
||||
async def test_static_auth_without_usable_credential_rejects(
|
||||
self, auth_type: MCPAuthType, credential: str | dict[str, str] | None
|
||||
) -> None:
|
||||
server = MCPServer(
|
||||
server_id="empty-static", name="empty-static", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=auth_type,
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MCPServerManager()._create_mcp_client(server, mcp_auth_header=credential)
|
||||
assert exc.value.status_code == 500
|
||||
assert "credential" in str(exc.value.detail).lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("auth_type,headers", [
|
||||
(MCPAuth.api_key, {"X-API-Key": "key"}),
|
||||
(MCPAuth.bearer_token, {"Authorization": "Bearer token"}),
|
||||
])
|
||||
async def test_static_auth_accepts_actual_forwarded_credential(
|
||||
self, auth_type: MCPAuthType, headers: dict[str, str]
|
||||
) -> None:
|
||||
server = MCPServer(
|
||||
server_id="header-static", name="header-static", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=auth_type,
|
||||
)
|
||||
client = await MCPServerManager()._create_mcp_client(server, extra_headers=headers)
|
||||
assert client._get_auth_headers() == headers
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("auth_type", [MCPAuth.oauth2_token_exchange, MCPAuth.api_key, MCPAuth.bearer_token])
|
||||
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",
|
||||
transport=MCPTransport.http, auth_type=auth_type,
|
||||
token_exchange_endpoint="https://idp.example/token",
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
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=None,
|
||||
)
|
||||
assert exc.value.status_code in (401, 500)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("auth_type,slot", [(MCPAuth.api_key, "X-API-Key"), (MCPAuth.authorization, "Authorization")])
|
||||
async def test_raw_static_value_named_token_is_a_usable_credential(self, auth_type: MCPAuthType, slot: str) -> None:
|
||||
server = MCPServer(server_id="raw-key", name="raw-key", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=auth_type, authentication_token="token")
|
||||
client = await MCPServerManager()._create_mcp_client(server)
|
||||
assert client._resolved_auth is not None
|
||||
request = httpx.Request("GET", server.url)
|
||||
flow = client._resolved_auth.auth_flow(request)
|
||||
try:
|
||||
assert next(flow).headers[slot] == "token"
|
||||
finally:
|
||||
flow.close()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_byok_flag_cannot_bypass_incomplete_obo(self) -> None:
|
||||
server = MCPServer(server_id="obo-byok", name="obo-byok", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=MCPAuth.oauth2_token_exchange, is_byok=True,
|
||||
token_exchange_endpoint="https://idp.example/token")
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MCPServerManager()._create_mcp_client(server, mcp_auth_header="Bearer override")
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("configured,override", [(None, "Bearer usable"), ("shared", "Bearer usable")])
|
||||
async def test_bearer_override_remains_usable(self, configured: str | None, override: str) -> None:
|
||||
server = MCPServer(server_id="override", name="override", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=configured)
|
||||
client = await MCPServerManager()._create_mcp_client(server, mcp_auth_header=override)
|
||||
assert client._get_auth_headers()["Authorization"] == override
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("token", [None, "shared"])
|
||||
async def test_empty_injected_header_cannot_satisfy_protected_auth(self, token: str | None) -> None:
|
||||
server = MCPServer(server_id="empty-header", name="empty-header", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=MCPAuth.bearer_token, authentication_token=token)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MCPServerManager()._create_mcp_client(server, extra_headers={"authorization": " "})
|
||||
assert exc.value.status_code == 500
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_slot_uses_its_actual_credential(self) -> None:
|
||||
server = MCPServer(server_id="custom", name="custom", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=MCPAuth.api_key,
|
||||
upstream_token_header="X-Custom", authentication_token="key")
|
||||
client = await MCPServerManager()._create_mcp_client(server, extra_headers={"X-Trace": "trace"})
|
||||
assert client._credential_slot == "X-Custom"
|
||||
assert await client.discovery_auth_fingerprint()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("static,forwarded,caller", [
|
||||
({"X-API-Key": "static"}, {"x-api-key": "forwarded"}, None),
|
||||
({}, {"X-API-Key": "forwarded"}, None),
|
||||
({}, None, "ApiKey caller"),
|
||||
])
|
||||
async def test_openapi_static_credentials_remain_supported(
|
||||
self, 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,
|
||||
)
|
||||
assert resolved is None
|
||||
assert retained == forwarded
|
||||
|
||||
@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
|
||||
|
||||
class CancelledAuth(httpx.Auth):
|
||||
closed = False
|
||||
|
||||
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
|
||||
try:
|
||||
raise asyncio.CancelledError()
|
||||
yield request
|
||||
finally:
|
||||
self.closed = True
|
||||
|
||||
auth = CancelledAuth()
|
||||
server = MCPServer(server_id="cancel", name="cancel", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=MCPAuth.api_key)
|
||||
client = MCPClient(server_url=server.url, auth_type=MCPAuth.api_key, resolved_auth=auth)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await prepare_mcp_client(server, client)
|
||||
assert auth.closed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("auth_type", [MCPAuth.basic, MCPAuth.token, MCPAuth.authorization])
|
||||
async def test_other_static_schemes_reject_whitespace_credentials(self, auth_type: MCPAuthType) -> None:
|
||||
server = MCPServer(server_id="blank-static", name="blank-static", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=auth_type, authentication_token=" ")
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MCPServerManager()._create_mcp_client(server)
|
||||
assert exc.value.status_code == 500
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("header", ["Basic", "Basic @@@", "Other abc"])
|
||||
async def test_basic_headers_without_usable_credentials_reject(self, header: str) -> None:
|
||||
server = MCPServer(server_id="bad-basic", name="bad-basic", url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.http, auth_type=MCPAuth.basic)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await MCPServerManager()._create_mcp_client(server, extra_headers={"Authorization": header})
|
||||
assert exc.value.status_code == 500
|
||||
|
|
|
|||
|
|
@ -1458,3 +1458,21 @@ class TestBoundedOpenAPISpecLoading:
|
|||
else:
|
||||
assert await load_openapi_spec_async("https://93.184.216.34/spec.json", max_bytes=100) == {"paths": {}}
|
||||
assert destination.call_count == 1
|
||||
|
||||
|
||||
def test_openapi_generator_import_does_not_require_mcp_sdk() -> None:
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
script = """
|
||||
import builtins
|
||||
original_import = builtins.__import__
|
||||
def without_mcp(name, *args, **kwargs):
|
||||
if name == 'mcp' or name.startswith('mcp.'):
|
||||
raise ModuleNotFoundError('MCP SDK unavailable')
|
||||
return original_import(name, *args, **kwargs)
|
||||
builtins.__import__ = without_mcp
|
||||
import litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator
|
||||
"""
|
||||
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue