diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 7a8df83f9f9..655b49e90a3 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2,7 +2,8 @@ import asyncio import html as _html import json import time -from typing import Any, Dict, Optional, Tuple +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx @@ -29,6 +30,9 @@ from litellm.proxy.utils import get_server_root_path from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + # TTL cache for upstream OAuth metadata fetched from pass-through MCP servers. # Keeps us from hammering the upstream IdP on each discovery request. # Keyed by (server_id, resource_url) → (expires_at_epoch, payload). @@ -228,28 +232,87 @@ def _validate_token_response( ) -async def _extract_user_id_from_request(request: Request) -> Optional[str]: - """Best-effort extraction of LiteLLM user_id from the request's Authorization header. +def _litellm_key_from_request(request: Request) -> Optional[str]: + """Return the LiteLLM API key presented on the request, or ``None``. - Called at the OAuth token endpoint so that per-user tokens can be stored - server-side. Uses a read-only cache lookup to avoid re-running the full - auth pipeline (which has side effects such as rate-limit increments and - spend logging). Returns ``None`` if no cached credential is found. + Accepts the key from ``x-litellm-api-key`` (what MCP clients such as Claude Desktop/Code + send) as well as ``Authorization``; either may carry a bare token or ``Bearer ``. + ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry + an OAuth/upstream bearer. """ - auth_header = request.headers.get("Authorization") or request.headers.get("authorization") - if not auth_header: + for header_value in ( + request.headers.get("x-litellm-api-key"), + request.headers.get("Authorization") or request.headers.get("authorization"), + ): + if not header_value: + continue + value = header_value.strip() + if value.lower().startswith("bearer "): + value = value[7:].strip() + if value: + return value + return None + + +def _active_key_user_id(key_obj: "UserAPIKeyAuth") -> Optional[str]: + """The key's ``user_id``, or ``None`` if the key is blocked or expired. + + The OAuth token endpoint is unauthenticated, so the presented key is validated here before its + identity is trusted to key a stored credential; a revoked or expired key must not be able to + write or overwrite the per-user OAuth token. ``get_key_object`` resolves a row without these + checks (the main ``user_api_key_auth`` pipeline enforces them downstream, which this endpoint + bypasses), so they are applied here. Deleted keys are already rejected upstream, where + ``get_key_object`` raises on a row that no longer exists. + """ + if key_obj.blocked is True: return None - lower = auth_header.lower() - if not lower.startswith("bearer "): + expires = key_obj.expires + if expires is not None: + expiry = expires if isinstance(expires, datetime) else datetime.fromisoformat(expires) + if expiry.tzinfo is None or expiry.tzinfo.utcoffset(expiry) is None: + expiry = expiry.replace(tzinfo=timezone.utc) + if expiry < datetime.now(timezone.utc): + return None + return key_obj.user_id + + +async def _extract_user_id_from_request(request: Request) -> Optional[str]: + """Resolve the LiteLLM ``user_id`` at the OAuth token endpoint so a per-user token is stored + under the same identity the egress later reads it by (``user_api_key_auth.user_id``). + + Resolves authoritatively via ``get_key_object`` (cache first, then DB) instead of a raw cache + peek. On a multi-replica gateway the token-exchange request can land on a worker whose in-memory + cache never saw the key, and a cross-replica Redis hit deserializes to a plain ``dict`` rather + than a ``UserAPIKeyAuth``; the previous code read only ``Authorization`` and did + ``getattr(cached, "user_id")`` with no ``model_type`` rehydration and no DB fallback, so it + silently returned ``None`` and the token was never persisted, which makes the egress 401 on every + reconnect. The resolved key is validated (``_active_key_user_id``) before its identity is trusted, + so a blocked or expired key cannot write. Returns ``None`` when no key is present, the key cannot + be resolved, or it is blocked/expired. + """ + token = _litellm_key_from_request(request) + if not token: return None - token = auth_header[7:].strip() try: from litellm.proxy._types import hash_token # noqa: PLC0415 - from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + from litellm.proxy.auth.auth_checks import get_key_object # noqa: PLC0415 + from litellm.proxy.proxy_server import ( # noqa: PLC0415 + prisma_client, + user_api_key_cache, + ) - cached = await user_api_key_cache.async_get_cache(hash_token(token)) - return getattr(cached, "user_id", None) - except Exception: + key_obj = await get_key_object( + hashed_token=hash_token(token), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return _active_key_user_id(key_obj) + except Exception as exc: + verbose_logger.debug( + "_extract_user_id_from_request: could not resolve a LiteLLM user_id for the presented " + "key (%s); per-user token will not be stored server-side.", + type(exc).__name__, + ) return None @@ -477,11 +540,12 @@ async def exchange_token_with_server( exc, ) else: - verbose_logger.debug( - "exchange_token_with_server: no LiteLLM user_id found in request; " - "per-user token for server=%s will not be stored server-side. " - "The client should call POST /mcp/server/{id}/oauth-user-credential " - "to store it manually.", + verbose_logger.warning( + "exchange_token_with_server: could not resolve a LiteLLM user_id for the request, " + "so the per-user token for server=%s was NOT stored. The authorization_code egress " + "requires the stored token, so the client will be challenged with 401 on reconnect. " + "Ensure the request carries a valid LiteLLM key (x-litellm-api-key or Authorization), " + "or store it via POST /mcp/server/{id}/oauth-user-credential.", mcp_server.server_id, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 6fd935e3364..69158b00a1a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -2733,3 +2733,156 @@ async def test_token_exchange_passes_through_upstream_expires_in(): {"access_token": "tok", "token_type": "Bearer", "expires_in": 43200} ) assert body["expires_in"] == 43200 + + +def _token_request(headers): + """A real Starlette request with case-insensitive headers (matches production).""" + from starlette.requests import Request + + raw = [(k.lower().encode(), v.encode()) for k, v in headers.items()] + return Request({"type": "http", "method": "POST", "path": "/token", "headers": raw, "query_string": b""}) + + +@pytest.fixture +def proxy_globals(): + """Inject the cache/prisma the OAuth token endpoint resolves identity through, and restore + them afterward. These module globals are the proxy's real wiring points, so setting them is + dependency injection rather than monkeypatching a class.""" + import litellm.proxy.proxy_server as ps + + saved = (ps.user_api_key_cache, ps.prisma_client) + try: + yield ps + finally: + ps.user_api_key_cache, ps.prisma_client = saved + + +@pytest.mark.asyncio +async def test_extract_user_id_reads_x_litellm_api_key_header(proxy_globals): + """The LiteLLM key arrives on x-litellm-api-key (what Claude Desktop/Code send), not + Authorization. Reading only Authorization dropped the identity, so the per-user token was + never stored and the egress 401'd forever. Resolution must honor x-litellm-api-key.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _extract_user_id_from_request, + ) + from litellm.proxy._types import UserAPIKeyAuth, hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-alice-key" + cache = UserApiKeyCache() + await cache.async_set_cache( + hash_token(key), + UserAPIKeyAuth(token=hash_token(key), user_id="alice"), + model_type=UserAPIKeyAuth, + ) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"x-litellm-api-key": f"Bearer {key}"}) + assert await _extract_user_id_from_request(request) == "alice" + + +@pytest.mark.asyncio +async def test_extract_user_id_rehydrates_cross_replica_dict_cache(proxy_globals): + """Cross-replica, async_get_cache hands back a serialized dict, not a UserAPIKeyAuth. + Resolution must rehydrate it; the old getattr(cached, "user_id") returned None on a dict, + which is exactly why a multi-replica gateway never found the stored token.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _extract_user_id_from_request, + ) + from litellm.proxy._types import hash_token + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-alice-key" + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache(hash_token(key), {"token": hash_token(key), "user_id": "alice"}) + proxy_globals.user_api_key_cache = cache + proxy_globals.prisma_client = object() + + request = _token_request({"Authorization": f"Bearer {key}"}) + assert await _extract_user_id_from_request(request) == "alice" + + +@pytest.mark.asyncio +async def test_extract_user_id_falls_back_to_db_on_cache_miss(proxy_globals): + """A cache miss must read the key from the DB rather than returning None; the old code did a + cache-only peek and skipped the DB, so any replica that hadn't just authenticated the key + failed to store the token.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _extract_user_id_from_request, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + key = "sk-bob-key" + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="db-bob") + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": key}) + assert await _extract_user_id_from_request(request) == "db-bob" + + +@pytest.mark.asyncio +async def test_extract_user_id_none_without_litellm_key(proxy_globals): + """No LiteLLM key on the request resolves to None without consulting the resolver.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _extract_user_id_from_request, + ) + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = object() + + request = _token_request({"content-type": "application/json"}) + assert await _extract_user_id_from_request(request) is None + + +@pytest.mark.asyncio +async def test_extract_user_id_rejects_blocked_key(proxy_globals): + """A blocked LiteLLM key must not resolve an identity. get_key_object returns the DB row without + checking blocked/expiry (the main auth pipeline does, and the public token endpoint bypasses it), + so a revoked key could otherwise overwrite the stored per-user OAuth token for its user.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _extract_user_id_from_request, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="blocked-user", blocked=True) + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-blocked-key"}) + assert await _extract_user_id_from_request(request) is None + + +@pytest.mark.asyncio +async def test_extract_user_id_rejects_expired_key(proxy_globals): + """An expired LiteLLM key must not resolve an identity, for the same reason as a blocked key.""" + from datetime import datetime, timedelta, timezone + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _extract_user_id_from_request, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + expired = (datetime.now(timezone.utc) - timedelta(hours=1)).isoformat() + + class _FakePrisma: + async def get_data(self, token, table_name, parent_otel_span=None, proxy_logging_obj=None): + return UserAPIKeyAuth(token=token, user_id="expired-user", expires=expired) + + proxy_globals.user_api_key_cache = UserApiKeyCache() + proxy_globals.prisma_client = _FakePrisma() + + request = _token_request({"x-litellm-api-key": "sk-expired-key"}) + assert await _extract_user_id_from_request(request) is None