fix(mcp): resolve per-user OAuth identity authoritatively at the token endpoint

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -2,7 +2,8 @@
 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.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``.
+
+    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 <token>``.
+    ``x-litellm-api-key`` wins when both are present, since ``Authorization`` may instead carry
+    an OAuth/upstream bearer.
+    """
+    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.
 
-    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.
+    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.
     """
-    auth_header = request.headers.get("Authorization") or request.headers.get("authorization")
-    if not auth_header:
+    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,
             )
 

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