scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

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

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

The OAuth token endpoint stored a user's per-server token under the identity
returned by _extract_user_id_from_request, which read only the Authorization
header and did getattr(cached, "user_id") on a raw user_api_key_cache lookup
with no model_type rehydration and no DB fallback. That silently returned None
in two common cases on a multi-replica gateway: the LiteLLM key arrives on
x-litellm-api-key (what MCP clients such as Claude Desktop and Claude Code
send) rather than Authorization, and a cross-replica cache hit deserializes to
a plain dict rather than a UserAPIKeyAuth, so getattr finds no attribute. When
it returned None the token was not persisted.

This was survivable until the authorization_code v2 migration began stripping
the caller's Authorization for migrated per-user OAuth servers and routing the
preemptive 401 existence check through the stored token, so a persist miss now
hard-fails: the egress challenges with 401 on every reconnect (the client sees
"rejected them on reconnect" or a successful connect with zero tools).

Resolve identity through get_key_object, the canonical resolver that reads the
cache with model_type and falls back to the DB, and accept the key from
x-litellm-api-key as well as Authorization. The silent persist skip is now a
warning. The caller-Authorization stripping stays as is, since reinstating it
would reopen the cross-user credential override it was added to prevent.

* fix(mcp): reject blocked or expired keys when resolving the token-endpoint identity

The OAuth token endpoint is unauthenticated, and get_key_object resolves a key row without the
blocked/expiry checks the main user_api_key_auth pipeline runs (that pipeline is bypassed here). So
a holder of a revoked or expired LiteLLM key could POST a valid upstream authorization code with
that key in x-litellm-api-key/Authorization and write or overwrite the stored per-user OAuth token
for that key's user. The cache-only resolver this replaced incidentally dropped blocked keys
(blocking purges the cache entry), so moving to the authoritative cache-then-DB resolution removed
that accidental shield.

Validate the resolved key before trusting its identity: return None when blocked or expired, so the
upsert is skipped. Deleted keys are already rejected, since get_key_object raises on a missing row.
Regression tests cover the blocked and expired cases and fail without the guard.
This commit is contained in:
tin-berri 2026-06-29 18:09:19 -07:00 • committed by GitHub
parent e195532c14
commit ec808edece
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 238 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

106
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -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 <token>``.
``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,
)

153
tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
View file

@ -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