mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
refactor(mcp): share v1's OAuth egress core; make V1PerUserTokenStore refresh-capable
Extract v1's per-user OAuth egress (Redis cache, else DB read with the refresh_token grant, then re-cache) from _get_user_oauth_extra_headers_from_db into resolve_user_oauth_access_token in db.py; the v1 header builder is now a thin wrapper over it and its callers are unchanged. V1PerUserTokenStore (the v2 OAuthTokenStore adapter) resolves through that same core via an injected server lookup, so the authorization_code arm injects exactly the token v1 would, with the same silent refresh, rather than a Redis-only read that can never refresh. One resolution implementation, two thin adapters (header dict and OAuthToken). Behavior-preserving: the existing v1 egress tests pass unchanged, and the arm is not wired into the live path yet (that lands with to_server_spec + the manager).
This commit is contained in:
Tin Chi Lo
parent
f406b3a626
commit
13b1dc18fb
4 changed files with 156 additions and 140 deletions
|
|
@ -3,7 +3,7 @@ import binascii
|
|||
import hashlib
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Iterable, List, Optional, Set, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Union, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -39,6 +39,9 @@ from litellm.repositories.verification_token_repository import (
|
|||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.mcp import MCPCredentials
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
||||
def _is_global_env_var_scope(scope: Any) -> bool:
|
||||
"""``scope="user"`` entries are placeholders the user fills in; everything
|
||||
|
|
@ -1222,6 +1225,91 @@ async def resolve_valid_user_oauth_token(
|
|||
return refreshed
|
||||
|
||||
|
||||
async def resolve_user_oauth_access_token(
|
||||
user_id: str | None,
|
||||
server: "MCPServer",
|
||||
prefetched_creds: dict[str, dict[str, object]] | None = None,
|
||||
) -> str | None:
|
||||
"""Resolve a user's valid OAuth2 access token for a server: Redis cache, else DB + refresh.
|
||||
|
||||
The egress token-resolution core shared by v1's header builder and the v2 ``OAuthTokenStore``
|
||||
adapter. Redis fast-path (skipped when ``prefetched_creds`` is supplied), else a DB read through
|
||||
``resolve_valid_user_oauth_token`` (which refreshes an expired token when a ``refresh_token`` is
|
||||
stored), re-warming the Redis cache with the per-server TTL. Returns ``None`` when there is no
|
||||
usable token; any error is swallowed to ``None`` so a transient failure reads as "not
|
||||
authorized" rather than raising.
|
||||
"""
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if not user_id or not server_id:
|
||||
return None
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
|
||||
_compute_per_user_token_ttl,
|
||||
mcp_per_user_token_cache,
|
||||
)
|
||||
|
||||
if prefetched_creds is None:
|
||||
cached_token = await mcp_per_user_token_cache.get(user_id, server_id)
|
||||
if cached_token is not None:
|
||||
return cached_token
|
||||
|
||||
prisma_client = None
|
||||
if prefetched_creds is not None:
|
||||
cred = prefetched_creds.get(server_id)
|
||||
else:
|
||||
from litellm.proxy.utils import get_prisma_client_or_throw
|
||||
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Connect a database to use OAuth2 MCP tools."
|
||||
)
|
||||
cred = await get_user_oauth_credential(prisma_client, user_id, server_id)
|
||||
|
||||
if not cred or not cred.get("access_token"):
|
||||
return None
|
||||
|
||||
cred = await resolve_valid_user_oauth_token(
|
||||
user_id=user_id,
|
||||
server=server,
|
||||
cred=cred,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if cred is None:
|
||||
# Refresh failed or token expired with no usable refresh_token — clear the stale
|
||||
# Redis entry so the next request doesn't reuse it.
|
||||
await mcp_per_user_token_cache.delete(user_id, server_id)
|
||||
return None
|
||||
|
||||
access_token: str = cred["access_token"]
|
||||
if prefetched_creds is None:
|
||||
ttl = _compute_per_user_token_ttl(
|
||||
server, _remaining_token_seconds(cred.get("expires_at"))
|
||||
)
|
||||
await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl)
|
||||
return access_token
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"resolve_user_oauth_access_token: failed for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _remaining_token_seconds(expires_at: str | None) -> int | None:
|
||||
"""Seconds until ``expires_at`` (ISO 8601), or None when absent/past/unparseable."""
|
||||
if not expires_at:
|
||||
return None
|
||||
try:
|
||||
exp_dt = datetime.fromisoformat(expires_at)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if exp_dt.tzinfo is None:
|
||||
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
|
||||
remaining = int((exp_dt - datetime.now(timezone.utc)).total_seconds())
|
||||
return remaining if remaining > 0 else None
|
||||
|
||||
|
||||
async def approve_mcp_server(
|
||||
prisma_client: PrismaClient,
|
||||
server_id: str,
|
||||
|
|
|
|||
|
|
@ -1,37 +1,49 @@
|
|||
"""v1-backed ``OAuthTokenStore`` source for the ``authorization_code`` mode.
|
||||
|
||||
Reads the user's stored access token through v1's ``mcp_per_user_token_cache`` (a Redis-backed,
|
||||
encrypted-at-rest per-user cache). This is a temporary adapter: step 1b replaces it with a
|
||||
v2-native token store that also tracks expiry and refresh, behind the same ``OAuthTokenStore`` seam.
|
||||
It imports v1, so it is kept out of the package ``__init__`` like the rest of the adapter layer.
|
||||
Resolves the user's access token through v1's egress core (``resolve_user_oauth_access_token``:
|
||||
Redis cache, else DB read with refresh), so the v2 arm injects exactly the token v1 would, with the
|
||||
same silent refresh. This is the strangler adapter; step 1b swaps the core for a v2-native store
|
||||
behind the same ``OAuthTokenStore`` seam. It imports v1, so it is kept out of the package
|
||||
``__init__`` like the rest of the adapter layer.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
||||
class V1PerUserTokenStore:
|
||||
"""``OAuthTokenStore`` backed by v1's per-user token cache.
|
||||
"""``OAuthTokenStore`` backed by v1's per-user OAuth egress.
|
||||
|
||||
v1 stores only the access token (the cache TTL is its lifetime), so the ``OAuthToken`` carries
|
||||
no ``expires_at`` or ``refresh_token``: the v2 cache holds it for its default TTL, and the OAuth
|
||||
challenge drives re-auth once v1's cache drops it. v1's ``get`` swallows errors as a miss, so
|
||||
this never raises ``TokenStoreUnavailable``.
|
||||
``fetch`` looks the server up by id (injected ``server_lookup``) and resolves the token through
|
||||
v1's ``resolve_user_oauth_access_token``, which refreshes an expired token when a refresh_token
|
||||
is stored and re-warms the Redis cache. Returns ``None`` when the user has no usable token (the
|
||||
arm turns that into a challenge); v1's core swallows store errors as a miss, so this never
|
||||
raises ``TokenStoreUnavailable``. The returned ``OAuthToken`` carries only the access token —
|
||||
refresh happens inside the core, not via ``RefreshingTokenStore``.
|
||||
"""
|
||||
|
||||
def __init__(self, server_lookup: Callable[[str], MCPServer | None]) -> None:
|
||||
self._server_lookup = server_lookup
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
if not user_id:
|
||||
return None
|
||||
server = self._server_lookup(server_id)
|
||||
if server is None:
|
||||
return None
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
|
||||
mcp_per_user_token_cache,
|
||||
from litellm.proxy._experimental.mcp_server.db import (
|
||||
resolve_user_oauth_access_token,
|
||||
)
|
||||
|
||||
access_token = await mcp_per_user_token_cache.get(user_id, server_id)
|
||||
if not access_token:
|
||||
return None
|
||||
return OAuthToken(access_token=access_token)
|
||||
access_token = await resolve_user_oauth_access_token(user_id, server)
|
||||
return OAuthToken(access_token=access_token) if access_token else None
|
||||
|
|
|
|||
|
|
@ -1360,115 +1360,21 @@ if MCP_AVAILABLE:
|
|||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Look up stored OAuth2 token for (user, server) and return as extra_headers dict.
|
||||
"""Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None.
|
||||
|
||||
Lookup order:
|
||||
1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied
|
||||
2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query
|
||||
3. Auto-refresh when the stored token is expired and a refresh_token exists
|
||||
|
||||
Args:
|
||||
prefetched_creds: Optional dict keyed by server_id with credential payloads.
|
||||
When provided, the Redis and individual DB lookups are
|
||||
skipped in favour of the pre-fetched batch result.
|
||||
Thin wrapper over ``resolve_user_oauth_access_token`` (Redis cache, else DB + refresh);
|
||||
``prefetched_creds`` skips the per-server Redis/DB lookups for the batch path.
|
||||
"""
|
||||
if server.auth_type != MCPAuth.oauth2:
|
||||
if server.auth_type != MCPAuth.oauth2 or user_api_key_auth is None:
|
||||
return None
|
||||
if user_api_key_auth is None:
|
||||
return None
|
||||
user_id = getattr(user_api_key_auth, "user_id", None)
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if not user_id or not server_id:
|
||||
return None
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
get_user_oauth_credential,
|
||||
resolve_valid_user_oauth_token,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415
|
||||
_compute_per_user_token_ttl,
|
||||
mcp_per_user_token_cache,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
resolve_user_oauth_access_token,
|
||||
)
|
||||
|
||||
# ── Fast path: Redis cache ────────────────────────────────────────
|
||||
# Only used when prefetched_creds is not supplied (individual lookup).
|
||||
if prefetched_creds is None:
|
||||
cached_token = await mcp_per_user_token_cache.get(user_id, server_id)
|
||||
if cached_token is not None:
|
||||
verbose_logger.debug(
|
||||
"_get_user_oauth_extra_headers_from_db: Redis hit for user=%s server=%s",
|
||||
user_id,
|
||||
server_id,
|
||||
)
|
||||
return {"Authorization": f"Bearer {cached_token}"}
|
||||
|
||||
# ── Slow path: DB lookup ──────────────────────────────────────────
|
||||
prisma_client = None
|
||||
if prefetched_creds is not None:
|
||||
cred = prefetched_creds.get(server_id)
|
||||
else:
|
||||
from litellm.proxy.utils import ( # noqa: PLC0415
|
||||
get_prisma_client_or_throw,
|
||||
)
|
||||
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
"Database not connected. Connect a database to use OAuth2 MCP tools."
|
||||
)
|
||||
cred = await get_user_oauth_credential(
|
||||
prisma_client, user_id, server_id
|
||||
)
|
||||
|
||||
if not cred or not cred.get("access_token"):
|
||||
return None
|
||||
|
||||
cred = await resolve_valid_user_oauth_token(
|
||||
user_id=user_id,
|
||||
server=server,
|
||||
cred=cred,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if cred is None:
|
||||
# Refresh failed or token expired with no usable refresh_token —
|
||||
# clear the stale Redis entry so the next request doesn't reuse it.
|
||||
await mcp_per_user_token_cache.delete(user_id, server_id)
|
||||
return None
|
||||
|
||||
access_token: str = cred["access_token"]
|
||||
|
||||
# Warm (or re-warm) the Redis cache from the DB result.
|
||||
# Always write regardless of whether expires_at is present — tokens
|
||||
# without an expiry are still valid and should be cached using the
|
||||
# server/default TTL so subsequent requests are fast.
|
||||
if prefetched_creds is None:
|
||||
raw_expires = None
|
||||
expires_at = cred.get("expires_at")
|
||||
if expires_at:
|
||||
from datetime import datetime, timezone # noqa: PLC0415
|
||||
|
||||
try:
|
||||
exp_dt = datetime.fromisoformat(expires_at)
|
||||
if exp_dt.tzinfo is None:
|
||||
exp_dt = exp_dt.replace(tzinfo=timezone.utc)
|
||||
remaining = int(
|
||||
(exp_dt - datetime.now(timezone.utc)).total_seconds()
|
||||
)
|
||||
raw_expires = max(remaining, 0) if remaining > 0 else None
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
ttl = _compute_per_user_token_ttl(server, raw_expires)
|
||||
await mcp_per_user_token_cache.set(
|
||||
user_id, server_id, access_token, ttl
|
||||
)
|
||||
|
||||
return {"Authorization": f"Bearer {access_token}"}
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for user=%s server=%s: %s",
|
||||
user_id,
|
||||
server_id,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
token = await resolve_user_oauth_access_token(
|
||||
getattr(user_api_key_auth, "user_id", None), server, prefetched_creds
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"} if token else None
|
||||
|
||||
async def _prefetch_oauth_creds_for_user(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
|
|
|
|||
|
|
@ -6,32 +6,42 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.v1_token_store
|
|||
V1PerUserTokenStore,
|
||||
)
|
||||
|
||||
_GET = (
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache."
|
||||
"mcp_per_user_token_cache.get"
|
||||
)
|
||||
_RESOLVE = "litellm.proxy._experimental.mcp_server.db.resolve_user_oauth_access_token"
|
||||
|
||||
|
||||
async def test_wraps_the_v1_access_token():
|
||||
with patch(_GET, new=AsyncMock(return_value="at-123")):
|
||||
token = await V1PerUserTokenStore().fetch("alice", "s")
|
||||
def _store_for(server: object) -> V1PerUserTokenStore:
|
||||
return V1PerUserTokenStore(server_lookup=lambda _server_id: server)
|
||||
|
||||
|
||||
async def test_wraps_the_resolved_access_token():
|
||||
with patch(_RESOLVE, new=AsyncMock(return_value="at-123")):
|
||||
token = await _store_for(object()).fetch("alice", "s")
|
||||
assert token is not None and token.access_token == "at-123"
|
||||
|
||||
|
||||
async def test_missing_token_is_none():
|
||||
with patch(_GET, new=AsyncMock(return_value=None)):
|
||||
assert await V1PerUserTokenStore().fetch("alice", "s") is None
|
||||
with patch(_RESOLVE, new=AsyncMock(return_value=None)):
|
||||
assert await _store_for(object()).fetch("alice", "s") is None
|
||||
|
||||
|
||||
async def test_empty_user_short_circuits_without_hitting_v1():
|
||||
get = AsyncMock(return_value="at")
|
||||
with patch(_GET, new=get):
|
||||
assert await V1PerUserTokenStore().fetch("", "s") is None
|
||||
get.assert_not_called()
|
||||
async def test_empty_user_short_circuits_without_resolving():
|
||||
resolve = AsyncMock(return_value="at")
|
||||
with patch(_RESOLVE, new=resolve):
|
||||
assert await _store_for(object()).fetch("", "s") is None
|
||||
resolve.assert_not_called()
|
||||
|
||||
|
||||
async def test_passes_user_and_server_through_to_v1():
|
||||
get = AsyncMock(return_value="at")
|
||||
with patch(_GET, new=get):
|
||||
await V1PerUserTokenStore().fetch("alice", "srv-1")
|
||||
get.assert_awaited_once_with("alice", "srv-1")
|
||||
async def test_unknown_server_is_none_without_resolving():
|
||||
resolve = AsyncMock(return_value="at")
|
||||
store = V1PerUserTokenStore(server_lookup=lambda _server_id: None)
|
||||
with patch(_RESOLVE, new=resolve):
|
||||
assert await store.fetch("alice", "missing") is None
|
||||
resolve.assert_not_called()
|
||||
|
||||
|
||||
async def test_passes_user_and_resolved_server_to_the_core():
|
||||
server = object()
|
||||
resolve = AsyncMock(return_value="at")
|
||||
with patch(_RESOLVE, new=resolve):
|
||||
await _store_for(server).fetch("alice", "srv-1")
|
||||
resolve.assert_awaited_once_with("alice", server)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue