mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
feat(mcp): migrate authorization_code MCP to the v2 resolver (single-replica) [1/2] (#31473)
* feat(mcp): implement the authorization_code resolver arm
Resolve a user's authorization_code token through the injected OAuthTokenStore: present ->
Authorization: Bearer <access_token>; absent -> the RFC 9728 WWW-Authenticate OAuth challenge;
store unavailable -> the same challenge (not a 500), since a transient outage is not a definite
absence. UpstreamCredentialProvider gains the oauth_token_store collaborator (fail-closed null
default); per-subject isolation comes from keying the fetch on subject_id. Not live until
to_server_spec maps authorization_code and a v1-backed token source is wired (next steps).
* feat(mcp): v1-backed OAuth token source for authorization_code
V1PerUserTokenStore reads the user's stored access token through v1's mcp_per_user_token_cache
(Redis-backed, encrypted) and wraps it in an OAuthToken. v1 holds only the access token (its
cache TTL is the lifetime), so no expires_at/refresh_token yet; the v2 cache holds it for its
default TTL and the OAuth challenge drives re-auth once v1's cache drops it. Additive: nothing
wires it yet, so no behavior change. Step 1b swaps it for a v2-native token store behind the
OAuthTokenStore seam.
* style(mcp): modern type annotations in the authorization_code arm and source
* 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).
* feat(mcp): route oauth2 per-user (authorization_code) servers through the v2 resolver
to_server_spec maps an oauth2 server to AuthorizationCodeConfig when it relies on per-user tokens
(needs_user_oauth_token and not delegate_auth_to_upstream); client_credentials (M2M), delegated
upstream OAuth, token exchange, and SigV4 still defer to v1. The manager injects V1PerUserTokenStore
(resolving through v1's shared egress core) into the credential provider. The v2 path is live but
still defers to a token v1 places in extra_headers; the cutover that makes v1 step aside lands next,
alongside the unified challenge.
* feat(mcp): per-server fail-closed OAuth challenge at the v2 egress
When an authorization_code server has no usable per-user token, the arm returns a semantic
unauthorized and the graft builds the 401 where the full MCPServer is in hand: a relative,
per-server RFC 9728 resource_metadata pointer (/.well-known/oauth-protected-resource/mcp/{name})
that names the server's own authorization server, instead of the resolver's earlier root pointer
which resolved to the gateway's generic PRM. Relative, so it is correct behind a reverse proxy
without request context. The listing-phase 401 still emits the RFC 8414 authorization_uri form;
both now target the same server, so the remaining difference is cosmetic and unifies in a later PR.
* feat(mcp): cut the call_tool egress over to v2 for authorization_code servers
_resolve_oauth2_headers_for_tool_call steps aside (builds no header) when to_server_spec maps the
server, so the v2 resolver drives the token-present case instead of being shadowed by a token v1
places in extra_headers. Non-migrated oauth2 (delegate, client_credentials) and BYOK still build
their header on v1. With this, v2 owns the authorization_code egress end to end: inject the
refreshed per-user token when present, raise the per-server fail-closed 401 when absent.
* feat(mcp): cut the tools/list connection over to v2 for authorization_code servers
The listing connection's per-user OAuth header is no longer built by v1 for migrated servers; the
v2 resolver drives it at connect time, ending the double-resolution where v1 built the token into
extra_headers and the v2 graft then deferred to it. Safe because the preemptive 401 (in the
streamable-http and SSE handlers) already challenges a missing token before the listing connection
runs, so the connection is only reached with a token present. Non-migrated oauth2 (delegate) and
the rest still build their header on v1. With this, resolve_credentials' result is honored on every
authorization_code upstream path: tool calls and listing.
* feat(mcp): route the preemptive 401 existence check through the v2 resolver
The discovery-phase 401 no longer calls v1's _get_user_oauth_extra_headers_from_db to decide
whether a migrated server has a token; it asks the v2 resolver via a new has_user_oauth_token
manager method (to_server_spec + to_subject + resolve_credentials, Ok means a token exists). With
this, every authorization_code resolution runs through the v2 resolver: the call_tool egress, the
listing connection, and the discovery challenge. Delegate servers short-circuit before the check
(the client completes PKCE with the upstream). The challenge itself still emits the RFC 8414
authorization_uri form; the format unification stays a follow-up.
* refactor(mcp): extract the authorization_code arm into a helper
Mirror the api_key arm's structure: the inline AuthorizationCodeConfig body moves into
_authorization_code(subject, server), keeping resolve_credentials a flat one-line-per-arm dispatch.
The helper is annotated with the concrete StaticHeaderAuth it returns rather than the abstract
httpx.Auth (which api_key uses) because a new method carrying the unresolved httpx.Auth return
would add reportUnknownMemberType; the concrete type is both precise and budget-neutral.
* fix(mcp): emit the canonical WWW-Authenticate header name in the OAuth challenge
raise_user_oauth_challenge emitted the header lowercase while the sibling raise_public and every
resource_metadata (RFC 9728) emitter use the canonical WWW-Authenticate; align it. HTTP header names
are case-insensitive on the wire so this is cosmetic for compliant clients, but it keeps the challenge
builders consistent and matches RFC 6750.
* feat(mcp): v2-native per-user token read store (step 1b inner store)
Reads the user's persisted authorization_code credential and returns a typed OAuthToken (access
token, epoch expiry, refresh token), validating the decoded blob at this boundary so no Any leaks
past it. The raw inner store that RefreshingTokenStore/CachedOAuthTokenStore wrap; the DB read +
decode collaborator is injected so it stays testable. Not yet wired - V1PerUserTokenStore is still
the composition-root store until the refresher and cross-worker cache land.
* feat(mcp): v2-native authorization_code token refresher (step 1b)
The refresh_token grant for the authorization_code mode: POSTs the RFC 6749 refresh_token grant to
the server's token endpoint, persists the rotated triple, and returns the new typed OAuthToken for
RefreshingTokenStore to cache. HTTP post and persist are injected so the grant + response parsing
are testable without a live IdP/DB. Also extends the TokenRefresher seam with (user_id, server_id),
which the foundation's refresh(token) lacked but the grant (server config) and persist (key) need.
* feat(mcp): wire the v2-native per-user OAuth store into the resolver (step 1b piece 4)
Assemble Cached(Refreshing(V2PerUserTokenStore)) at the composition root and replace
V1PerUserTokenStore in mcp_server_manager. The chain is built lazily on first fetch (its cache/DB/
Redis collaborators are LiteLLM globals not ready at import); when Redis is wired it uses the
cross-replica path (DualCache cache + SET NX PX coordinator), else the in-process defaults. The DB
read, refresh-grant POST, and persist acquire their globals per call like v1. authorization_code
resolution now reads/refreshes through the v2-native lifecycle, not v1's core.
* refactor(mcp): delete the unwired V1PerUserTokenStore adapter (step 1b piece 5)
Piece 4 replaced V1PerUserTokenStore with the v2-native chain at the composition root, leaving the
adapter with no callers, so remove it and its test. The shared v1 read/refresh core
(resolve_user_oauth_access_token and friends) stays - delegate's egress in server.py still uses it -
and comes out with the delegate migration.
* fix(mcp): green CI for authz_code dispatch (format + UTC expiry + v2-seam tests)
- ruff format per_user_oauth_store.py (clears the lint check)
- v2_token_store._iso_to_epoch: anchor a tz-naive expiry to UTC before
.timestamp(), matching v1's db.py _remaining_token_seconds (Greptile P1) so a
non-UTC host doesn't read the expiry as local time and skew refresh timing
- test_mcp_stale_session: repoint the 3 discovery tests off the removed v1
_get_user_oauth_extra_headers_from_db onto the v2 has_user_oauth_token seam;
the delegate test now asserts the existence check is never consulted (delegate
short-circuits to the resource_metadata 401 before any token lookup)
- test_mcp_server_manager: repoint test_deferred_mode_uses_v1_auth_value at M2M
(oauth2 client_credentials), which is still a deferred mode, since per-user
oauth2 (authorization_code) now routes to the v2 resolver
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(mcp): caller Authorization must not override the stored per-user OAuth token
A caller with a valid x-litellm-api-key could include their own
"Authorization: Bearer <chosen>" header and have the proxy execute tools against
that bearer instead of the user's stored OAuth credential. For a v2-migrated
authorization_code server the caller's Authorization was seeded into
extra_headers, and the graft's apply-if-absent then dropped the resolved
per-user token in its favor. v1 prevented this by overwriting a stale client
Authorization with the stored token; this restores that precedence on both
egress paths (connect + call_tool).
- _should_strip_caller_authorization: also strip for migrated per-user OAuth
(authorization_code) servers - the v2 resolver injects the stored token, so a
caller-forwarded Authorization must not be forwarded upstream. Delegate /
pass-through (to_server_spec is None) keep forwarding the caller's bearer.
- both seed sites (_prepare_mcp_server_headers, _call_regular_mcp_tool) drop only
the Authorization from the caller's oauth2_headers (via _without_authorization),
keeping any other forwarded header and any hook/static Authorization (which
still wins, as in v1).
- regression test for the call_tool path; updated the two tests that asserted the
old (vulnerable) forwarding to assert the secure behavior.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(mcp): preserve recorded OAuth scopes across authorization_code refresh
When a refresh response omits `scope` (RFC 6749 §5.1, where omission means unchanged), the v2 refresher persisted scopes=None and overwrote the user's recorded grant. v1 carried the prior scopes forward via `or cred.get("scopes")`; the v2 path lost that because OAuthToken did not model scopes
OAuthToken now carries scopes, V2PerUserTokenStore populates them on read, and AuthorizationCodeRefresher carries them forward for both the persisted write and the returned/cached token, so repeated refreshes do not erode them. A present `scope` in the response still replaces the prior grant
Adds regression tests: a refresh omitting `scope` preserves the prior scopes, and a present `scope` overrides them
* fix(mcp): keep user token in authorization_code tools preview
After to_server_spec maps oauth2 onto the v2 resolver, the interactive tools preview for an unsaved authorization_code server read the per-user token store, found nothing, and fail-closed with a 401, so the create/test tab could no longer list tools
The preview now routes the just-authorized token (forwarded in oauth2_headers) through mcp_auth_header, so _create_mcp_client takes the per-request-override v1 path and uses it directly, matching v1's preview. Gated to the v2-mapped oauth2 case; M2M, delegate/passthrough, and token-exchange keep their existing preview path
Adds tests: interactive oauth routes the forwarded token to mcp_auth_header, M2M and token-exchange do not
* fix(mcp): stop caller-supplied auth from overriding stored authorization_code tokens
A caller-supplied per-request override (mcp_auth_header / x-mcp-auth / x-mcp-<alias>-authorization) disabled the v2 resolver in _create_mcp_client for any spec, so an authenticated user with a stored authorization_code token could force an arbitrary upstream bearer and bypass the stored credential and its save-time validation. _create_mcp_client now keeps the v2 spec for authorization_code and ignores the override; other modes keep the client-side-credentials override
The create/test tools preview no longer relies on that override path. It resolves the just-authorized, not-yet-persisted token through the v2 resolver via a one-shot PresentedOAuthTokenStore passed as cred_provider - the same path runtime uses for the stored token - so preview and runtime resolve identically. This replaces the mcp_auth_header routing added earlier
Adds tests: a caller override cannot bypass the v2 resolver for authorization_code; the interactive preview resolves via the presented store rather than a caller header; M2M and token-exchange build no presented provider
* feat(mcp): cross-replica single-flight refresh for the v2 per-user OAuth store [2/2] (#31474)
* feat(mcp): encrypt+serialize codec for caching OAuth tokens in Redis (step 1b §1.5)
The serialize+encrypt boundary a cross-replica cache needs: a plaintext bearer in Redis is a leak, so
encode() encrypts (NaCl in prod via the injected encrypt, identity in tests). Caches only access_token
and expires_at, never the refresh_token - the hot path needs just the bearer, and the long-lived
refresh_token stays in the DB (the refresh path is always a cache miss), matching v1. A decoded token
always has refresh_token=None. Undecryptable (key rotation) or corrupt entries read as a miss.
* feat(mcp): DualCache-backed token cache backend (step 1b §1.5)
The cross-replica TokenCacheBackend implementation that plugs into the foundation's
CachedOAuthTokenStore seam: encrypts+serializes the token via the codec and stores it in LiteLLM's
shared DualCache under the same per-(user,server) key v1 used, so workers share one refresh and a
token cached by v1 or v2 is readable by the other across the cutover. Cache and codec are injected;
a non-positive TTL (already-expired token) is not cached, and a missing/corrupt entry reads as a miss.
* feat(mcp): Redis SET NX PX refresh coordinator (step 1b §1.5)
The cross-replica RefreshCoordinator that plugs into the foundation's RefreshingTokenStore seam: a SET
NX PX lock elects one worker to refresh per (user, server) while the rest wait for it and re-read the
token it persisted, so a rotating refresh_token is used once across the fleet, not once per worker. The
lock self-expires (PX) so a crashed holder can't wedge refresh; a loser falls back to a bounded re-read
and the surrounding store re-checks expiry next fetch, so a crash self-heals. The lock (a thin Redis
SET NX/DEL/EXISTS wrapper in prod) is injected, so the single-flight logic is testable without Redis.
* feat(mcp): Redis SET NX PX distributed lock (step 1b §1.5)
The concrete DistributedLock the RedisRefreshCoordinator elects refreshers with: acquire is an atomic
SET key NX PX ttl (first caller wins, entry self-expires so a crashed holder can't wedge refresh),
release is DEL, is_held is EXISTS. The async Redis client is injected (the client from LiteLLM's
RedisCache in prod), so it is unit-testable with a fake. Any Redis error degrades to not-acquired /
not-held so a cache blip causes an extra refresh, never a crash on the resolve path.
* feat(mcp): wire the cross-replica cache + coordinator into the per-user store (step 1b §1.5)
Upgrade the composition root to use the DualCache-backed cache and SET NX PX refresh
coordinator when Redis is wired, falling back to the foundation's in-process defaults on a
single replica. Layers the cross-replica path on top of the single-replica dispatch store.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* fix(mcp): refresh on lock-backend error instead of serving a stale token
The cross-replica refresh coordinator elected refreshers with a boolean acquire:
a Redis transport error was caught and returned as False, which is
indistinguishable from "another worker holds the lock". On a total Redis
outage every worker therefore took the wait-then-reread branch and served the
still-expired token upstream (the upstream then 401s), even though the lock and
coordinator docstrings claimed a Redis blip "degrades to an extra refresh".
Make acquire tristate (LockAcquisition: ACQUIRED / HELD / ERROR) so the
coordinator can tell a busy holder from a dead backend, and refresh anyway on
ERROR. This single-flight lock is a load optimization, not a correctness mutex,
so failing open is correct: it degrades a lock-backend outage to the
no-coordinator behavior (an extra refresh), never a stale bearer.
Add a regression test asserting an acquire error refreshes rather than
re-reading the expired token, and update the docstrings to match.
* style(mcp): wrap redis lock signatures at line-length 88 for CI ruff format
* fix(mcp): a refresh loser surfaces None, not a stale token, when the winner failed
The cross-replica coordinator's losers re-read the token the winner persisted.
If the winner's refresh failed, the store still holds the expired token, so the
loser re-read it and RefreshingTokenStore handed that expired bearer to the
caller (the upstream then 401s) instead of the re-auth challenge the winner
returned via None.
Make the loser's re-read expiry-aware, mirroring refresh_latest_token: a
re-read that is still expired surfaces None so the arm challenges. This only
affects the loser path; the winner's freshly refreshed token is returned
directly by the coordinator and is unaffected.
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Revert "feat(mcp): cross-replica single-flight refresh for the v2 per-user OA…" (#31492)
This reverts commit cd2fb6b0f2.
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
tin-berri
GitHub
parent
0216c969b8
commit
4b398ef6d4
20 changed files with 1573 additions and 195 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,
|
||||
|
|
|
|||
|
|
@ -63,9 +63,16 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
|
|||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
|
||||
raise_public,
|
||||
raise_user_oauth_challenge,
|
||||
to_server_spec,
|
||||
to_subject,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import (
|
||||
LazyPerUserOAuthTokenStore,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
||||
AuthorizationCodeConfig,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
MCP_TOOL_PREFIX_SEPARATOR,
|
||||
MCPMissingUserEnvVarsError,
|
||||
|
|
@ -186,6 +193,10 @@ def _should_strip_caller_authorization(
|
|||
Strip rules:
|
||||
- **M2M (client_credentials) servers**: never forward the caller's
|
||||
``Authorization`` — the proxy fetches its own upstream token.
|
||||
- **Migrated per-user OAuth (authorization_code) servers**: never forward
|
||||
the caller's ``Authorization`` — the v2 resolver injects the stored
|
||||
per-user token, so a caller-supplied bearer cannot override another
|
||||
user's stored credential. Delegate / pass-through keep forwarding it.
|
||||
- **OAuth pass-through servers**: strip when the ``Authorization``
|
||||
header is actually the LiteLLM API key — either because admission
|
||||
validated it (``user_api_key_auth.api_key`` is set) and the caller
|
||||
|
|
@ -198,6 +209,15 @@ def _should_strip_caller_authorization(
|
|||
"""
|
||||
if mcp_server.has_client_credentials:
|
||||
return True
|
||||
if (
|
||||
mcp_server.auth_type == MCPAuth.oauth2
|
||||
and to_server_spec(mcp_server) is not None
|
||||
):
|
||||
# Migrated per-user OAuth (authorization_code): the v2 resolver injects the
|
||||
# stored token, so a caller-forwarded Authorization must not be forwarded
|
||||
# upstream — it would override another user's stored credential. Delegate and
|
||||
# pass-through return None from to_server_spec and keep forwarding the bearer.
|
||||
return True
|
||||
if not mcp_server.is_oauth_passthrough:
|
||||
return False
|
||||
|
||||
|
|
@ -217,6 +237,18 @@ def _should_strip_caller_authorization(
|
|||
)
|
||||
|
||||
|
||||
def _without_authorization(
|
||||
headers: Optional[dict[str, str]],
|
||||
) -> Optional[dict[str, str]]:
|
||||
"""A copy of ``headers`` with any ``Authorization`` key removed (case-insensitive), or
|
||||
None if nothing remains. Drops only the credential, keeping other forwarded headers.
|
||||
"""
|
||||
if not headers:
|
||||
return None
|
||||
filtered = {k: v for k, v in headers.items() if k.lower() != "authorization"}
|
||||
return filtered or None
|
||||
|
||||
|
||||
def _extract_upstream_auth_failure(
|
||||
exc: BaseException,
|
||||
) -> Optional[Tuple[int, Optional[str]]]:
|
||||
|
|
@ -523,7 +555,9 @@ class MCPServerManager:
|
|||
return None
|
||||
|
||||
def __init__(self, cred_provider: Optional[UpstreamCredentialProvider] = None):
|
||||
self._cred_provider = cred_provider or UpstreamCredentialProvider()
|
||||
self._cred_provider = cred_provider or UpstreamCredentialProvider(
|
||||
oauth_token_store=LazyPerUserOAuthTokenStore(self.get_mcp_server_by_id)
|
||||
)
|
||||
self.registry: Dict[str, MCPServer] = {}
|
||||
self.config_mcp_servers: Dict[str, MCPServer] = {}
|
||||
"""
|
||||
|
|
@ -1934,6 +1968,7 @@ class MCPServerManager:
|
|||
stdio_env: Optional[Dict[str, str]] = None,
|
||||
subject_token: Optional[str] = None,
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
cred_provider: Optional[UpstreamCredentialProvider] = None,
|
||||
) -> MCPClient:
|
||||
"""
|
||||
Create an MCPClient instance for the given server.
|
||||
|
|
@ -1957,11 +1992,17 @@ class MCPServerManager:
|
|||
"""
|
||||
transport = server.transport or MCPTransport.sse
|
||||
spec = None if transport == MCPTransport.stdio else to_server_spec(server)
|
||||
# A per-request override is the caller-supplied credential v1 turns into the upstream
|
||||
# auth, so it must win; defer those to v1 (this defer falls away once the per-user modes
|
||||
# stop writing mcp_auth_header). An inbound header already in extra_headers is handled on
|
||||
# the v2 path below, not here.
|
||||
if spec is not None and mcp_auth_header:
|
||||
provider = cred_provider or self._cred_provider
|
||||
# A caller-supplied per-request override (mcp_auth_header / x-mcp-*) defers to the v1 path
|
||||
# so it wins - except for authorization_code, whose per-user token the v2 resolver owns. A
|
||||
# caller must not be able to substitute another user's stored credential, so we keep the v2
|
||||
# spec and ignore the override there; the REST tools preview supplies its not-yet-persisted
|
||||
# token through the resolver (cred_provider), never this path.
|
||||
if (
|
||||
spec is not None
|
||||
and mcp_auth_header
|
||||
and not isinstance(spec.config, AuthorizationCodeConfig)
|
||||
):
|
||||
spec = None
|
||||
auth_value = (
|
||||
await resolve_mcp_auth(server, mcp_auth_header, subject_token=subject_token)
|
||||
|
|
@ -2039,7 +2080,7 @@ class MCPServerManager:
|
|||
server_url = server.url or ""
|
||||
|
||||
if spec is not None:
|
||||
match await self._cred_provider.resolve_credentials(
|
||||
match await provider.resolve_credentials(
|
||||
to_subject(user_api_key_auth, subject_token), spec
|
||||
):
|
||||
case Ok(auth):
|
||||
|
|
@ -2059,6 +2100,10 @@ class MCPServerManager:
|
|||
):
|
||||
resolved_auth = None
|
||||
case Error(err):
|
||||
if err.tag == "unauthorized":
|
||||
# The arm signals a missing per-user token semantically; raise the
|
||||
# per-server OAuth challenge here, where the full MCPServer is in hand.
|
||||
raise_user_oauth_challenge(server)
|
||||
raise_public(err)
|
||||
return MCPClient(
|
||||
server_url=server_url,
|
||||
|
|
@ -3499,6 +3544,16 @@ class MCPServerManager:
|
|||
extra_headers = None
|
||||
else:
|
||||
extra_headers = oauth2_headers
|
||||
# Migrated authorization_code: the v2 resolver injects the stored per-user
|
||||
# token, so drop the caller-forwarded Authorization (apply-if-absent would
|
||||
# otherwise let it shadow the resolved token). Delegate keeps it. Centralized
|
||||
# via _should_strip_caller_authorization to match _prepare_mcp_server_headers.
|
||||
if extra_headers and _should_strip_caller_authorization(
|
||||
mcp_server=mcp_server,
|
||||
raw_headers=raw_headers,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
):
|
||||
extra_headers = _without_authorization(extra_headers)
|
||||
|
||||
if mcp_server.extra_headers and raw_headers:
|
||||
if extra_headers is None:
|
||||
|
|
@ -3670,6 +3725,22 @@ class MCPServerManager:
|
|||
|
||||
return mcp_server
|
||||
|
||||
async def has_user_oauth_token(
|
||||
self, server: MCPServer, user_api_key_auth: Optional[UserAPIKeyAuth]
|
||||
) -> bool:
|
||||
"""Whether the v2 resolver can produce a per-user token for this server right now.
|
||||
|
||||
This is the preemptive 401's existence check, routed through the same resolver that drives
|
||||
the egress so every authorization_code resolution (egress and the discovery challenge) runs
|
||||
through v2. Returns False for a server the resolver does not own (a None spec).
|
||||
"""
|
||||
spec = to_server_spec(server)
|
||||
if spec is None:
|
||||
return False
|
||||
return await self._cred_provider.has_user_token(
|
||||
to_subject(user_api_key_auth, None), spec
|
||||
)
|
||||
|
||||
async def _resolve_oauth2_headers_for_tool_call(
|
||||
self,
|
||||
mcp_server: MCPServer,
|
||||
|
|
@ -3684,6 +3755,12 @@ class MCPServerManager:
|
|||
):
|
||||
return oauth2_headers
|
||||
|
||||
if to_server_spec(mcp_server) is not None:
|
||||
# Migrated to v2: the resolver owns this server's per-user token (inject or fail-closed
|
||||
# 401). Building it into extra_headers here would let the v2 graft defer to it and
|
||||
# shadow the resolver, double-resolving and hiding the per-server challenge.
|
||||
return oauth2_headers
|
||||
|
||||
user_id = getattr(user_api_key_auth, "user_id", None)
|
||||
if not user_id:
|
||||
return oauth2_headers
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from typing_extensions import assert_never
|
|||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
||||
ApiKeyConfig,
|
||||
AuthorizationCodeConfig,
|
||||
CredError,
|
||||
NoneConfig,
|
||||
ServerSpec,
|
||||
|
|
@ -61,8 +62,9 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]:
|
|||
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
|
||||
explicitly mapped or explicitly deferred, rather than silently falling through to v1. Live
|
||||
modes: ``none`` and the static-header family (``api_key`` plus the Authorization schemes),
|
||||
all shared-key; every other mode returns None and stays on v1.
|
||||
modes: ``none``, the static-header family (``api_key`` plus the Authorization schemes,
|
||||
all shared-key), and ``oauth2`` per-user tokens (``authorization_code``); client_credentials
|
||||
(M2M), delegated/passthrough oauth2, token exchange, and SigV4 return None and stay on v1.
|
||||
"""
|
||||
if server.is_byok:
|
||||
return (
|
||||
|
|
@ -89,8 +91,17 @@ def to_server_spec(server: MCPServer) -> Optional[ServerSpec]:
|
|||
return _shared_key_spec(
|
||||
server, resource, "Authorization", "Basic", encode=True
|
||||
)
|
||||
case MCPAuth.oauth2 | MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4:
|
||||
return None # OAuth grants and SigV4 are not migrated yet -> defer to v1
|
||||
case MCPAuth.oauth2:
|
||||
if server.needs_user_oauth_token and not server.delegate_auth_to_upstream:
|
||||
return ServerSpec(
|
||||
server_id=server.server_id,
|
||||
resource=resource,
|
||||
config=AuthorizationCodeConfig(),
|
||||
)
|
||||
# client_credentials (M2M) and delegate/passthrough oauth2 stay on v1
|
||||
return None
|
||||
case MCPAuth.oauth2_token_exchange | MCPAuth.aws_sigv4:
|
||||
return None # token exchange and SigV4 are not migrated yet -> defer to v1
|
||||
assert_never(auth_type)
|
||||
|
||||
|
||||
|
|
@ -147,3 +158,25 @@ def raise_public(error: CredError) -> NoReturn:
|
|||
case "not_implemented":
|
||||
raise HTTPException(status_code=501, detail=error.summary)
|
||||
assert_never(error.tag)
|
||||
|
||||
|
||||
def raise_user_oauth_challenge(server: MCPServer) -> NoReturn:
|
||||
"""Raise the 401 an ``authorization_code`` server returns at egress when the user has no token.
|
||||
|
||||
Points at the server's RFC 9728 Protected Resource Metadata (``resource_metadata``), which names
|
||||
the upstream authorization server the client must complete OAuth with. The URL is per-server and
|
||||
relative, so it resolves against the caller's own host (correct even behind a reverse proxy)
|
||||
without needing request context. The listing-phase 401 still emits the RFC 8414 ``authorization_uri``
|
||||
form pending the format unification; both target the same server, so the difference is cosmetic.
|
||||
"""
|
||||
from litellm.proxy.utils import get_server_root_path # noqa: PLC0415
|
||||
|
||||
root = get_server_root_path()
|
||||
prefix = "" if root == "/" else root
|
||||
name = server.alias or server.server_name or server.name or server.server_id
|
||||
resource_metadata = f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}"
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Unauthorized",
|
||||
headers={"WWW-Authenticate": f'Bearer resource_metadata="{resource_metadata}"'},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,121 @@
|
|||
"""v2-native refresher for the ``authorization_code`` mode: the refresh_token grant, then persist.
|
||||
|
||||
Mints a fresh access token from a stored refresh_token by POSTing the RFC 6749 refresh_token grant to
|
||||
the server's token endpoint, persists the rotated triple, and returns the new typed ``OAuthToken`` for
|
||||
``RefreshingTokenStore`` to cache. The HTTP post and the persist are injected, so the orchestration
|
||||
and the (untyped) response parsing stay testable without a live IdP or DB. Replaces v1's
|
||||
``refresh_user_oauth_token`` as part of step 1b; rotation safety - one refresh per (user, server)
|
||||
across replicas - is the wrapping store's distributed single-flight, not this refresher's concern.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Protocol
|
||||
|
||||
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
|
||||
|
||||
ServerLookup = Callable[[str], "MCPServer | None"]
|
||||
TokenEndpointPost = Callable[
|
||||
[str, dict[str, str]], Awaitable["dict[str, object] | None"]
|
||||
]
|
||||
|
||||
|
||||
class CredentialPersist(Protocol):
|
||||
async def __call__(
|
||||
self,
|
||||
user_id: str,
|
||||
server_id: str,
|
||||
access_token: str,
|
||||
refresh_token: str | None,
|
||||
expires_in: int | None,
|
||||
scopes: tuple[str, ...] | None,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
def _parse_expires_in(raw: object) -> int | None:
|
||||
if isinstance(raw, bool):
|
||||
return None
|
||||
if isinstance(raw, int):
|
||||
return raw
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
return int(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _parse_scopes(raw: object) -> tuple[str, ...] | None:
|
||||
return tuple(raw.split()) if isinstance(raw, str) and raw else None
|
||||
|
||||
|
||||
class AuthorizationCodeRefresher:
|
||||
"""``TokenRefresher`` for authorization_code: refresh_token grant against the server, then persist.
|
||||
|
||||
``token_endpoint`` POSTs the OAuth form and returns the parsed JSON body (``None`` on any
|
||||
transport/HTTP failure, mirroring v1: a failed refresh is a miss, not a 500). ``persist`` writes
|
||||
the rotated triple for ``(user, server)`` - the v1 ``store_user_oauth_credential`` write, which
|
||||
stays. Returns ``None`` (the arm challenges) when there is no refresh_token, the server lacks a
|
||||
token endpoint, or the grant fails; never a stale or partial token. A rotated refresh_token from
|
||||
the response replaces the old one; an omitted one is carried forward, as are the recorded scopes
|
||||
when the response omits ``scope``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_lookup: ServerLookup,
|
||||
token_endpoint: TokenEndpointPost,
|
||||
persist: CredentialPersist,
|
||||
*,
|
||||
clock: Callable[[], float] = time.time,
|
||||
) -> None:
|
||||
self._server_lookup = server_lookup
|
||||
self._token_endpoint = token_endpoint
|
||||
self._persist = persist
|
||||
self._clock = clock
|
||||
|
||||
async def refresh(
|
||||
self, user_id: str, server_id: str, token: OAuthToken
|
||||
) -> OAuthToken | None:
|
||||
if token.refresh_token is None:
|
||||
return None
|
||||
server = self._server_lookup(server_id)
|
||||
if server is None or not server.token_url:
|
||||
return None
|
||||
|
||||
form = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": token.refresh_token,
|
||||
**({"client_id": server.client_id} if server.client_id else {}),
|
||||
**({"client_secret": server.client_secret} if server.client_secret else {}),
|
||||
}
|
||||
body = await self._token_endpoint(server.token_url, form)
|
||||
if body is None:
|
||||
return None
|
||||
access_token = body.get("access_token")
|
||||
if not isinstance(access_token, str) or not access_token:
|
||||
return None
|
||||
|
||||
rotated = body.get("refresh_token")
|
||||
new_refresh = (
|
||||
rotated if isinstance(rotated, str) and rotated else token.refresh_token
|
||||
)
|
||||
expires_in = _parse_expires_in(body.get("expires_in"))
|
||||
scopes = _parse_scopes(body.get("scope")) or token.scopes
|
||||
|
||||
await self._persist(
|
||||
user_id, server_id, access_token, new_refresh, expires_in, scopes or None
|
||||
)
|
||||
return OAuthToken(
|
||||
access_token=access_token,
|
||||
expires_at=self._clock() + expires_in if expires_in is not None else None,
|
||||
refresh_token=new_refresh,
|
||||
scopes=scopes,
|
||||
)
|
||||
|
|
@ -31,15 +31,20 @@ class OAuthToken:
|
|||
with each mode); it is never minted into a header directly. ``repr`` masks both secrets so a
|
||||
stray log line cannot leak them (the values are still plain ``str`` for the header path, since
|
||||
``SecretStr`` resolves as unknown under this repo's basedpyright).
|
||||
|
||||
``scopes`` is the recorded grant. A refresh response that omits ``scope`` (RFC 6749 §5.1: an
|
||||
omitted ``scope`` means unchanged) carries the prior value forward, so a refresh never silently
|
||||
drops it; the resolver itself does not read it.
|
||||
"""
|
||||
|
||||
access_token: str
|
||||
expires_at: float | None = None
|
||||
refresh_token: str | None = None
|
||||
scopes: tuple[str, ...] = ()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
has_refresh = self.refresh_token is not None
|
||||
return f"OAuthToken(access_token=***, expires_at={self.expires_at!r}, has_refresh_token={has_refresh})"
|
||||
return f"OAuthToken(access_token=***, expires_at={self.expires_at!r}, has_refresh_token={has_refresh}, scopes={self.scopes!r})"
|
||||
|
||||
|
||||
class TokenStoreUnavailable(Exception):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,129 @@
|
|||
"""Composition root for the v2-native authorization_code per-user OAuth token store (step 1b).
|
||||
|
||||
Assembles ``Cached(Refreshing(V2PerUserTokenStore))`` and replaces ``V1PerUserTokenStore`` in the
|
||||
resolver. The runtime collaborators (DB, HTTP) are LiteLLM globals not ready at import time, so the
|
||||
chain is built lazily on first use. The cache and refresh coordinator use the foundation's in-process
|
||||
defaults (correct for a single replica); the cross-replica path is layered on separately. The DB
|
||||
read/refresh-grant/persist collaborators acquire their globals per call, mirroring v1's lazy-import
|
||||
pattern.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import (
|
||||
AuthorizationCodeRefresher,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
CachedOAuthTokenStore,
|
||||
OAuthToken,
|
||||
RefreshingTokenStore,
|
||||
TokenStoreUnavailable,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import (
|
||||
V2PerUserTokenStore,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
# A token with no declared expiry is cached for this long; one with an expiry is cached until then.
|
||||
_DEFAULT_TTL_SECONDS = 300.0
|
||||
|
||||
ServerLookup = Callable[[str], "MCPServer | None"]
|
||||
|
||||
|
||||
async def _read_credential(user_id: str, server_id: str) -> dict[str, object] | None:
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
get_user_oauth_credential,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415
|
||||
|
||||
if prisma_client is None:
|
||||
raise TokenStoreUnavailable("Database not connected")
|
||||
return await get_user_oauth_credential(prisma_client, user_id, server_id)
|
||||
|
||||
|
||||
async def _persist_credential(
|
||||
user_id: str,
|
||||
server_id: str,
|
||||
access_token: str,
|
||||
refresh_token: str | None,
|
||||
expires_in: int | None,
|
||||
scopes: tuple[str, ...] | None,
|
||||
) -> None:
|
||||
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415
|
||||
store_user_oauth_credential,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415
|
||||
|
||||
if prisma_client is None:
|
||||
return
|
||||
await store_user_oauth_credential(
|
||||
prisma_client=prisma_client,
|
||||
user_id=user_id,
|
||||
server_id=server_id,
|
||||
access_token=access_token,
|
||||
refresh_token=refresh_token,
|
||||
expires_in=expires_in,
|
||||
scopes=list(scopes) if scopes else None,
|
||||
skip_byok_guard=True,
|
||||
)
|
||||
|
||||
|
||||
async def _post_token_endpoint(
|
||||
url: str, form: dict[str, str]
|
||||
) -> dict[str, object] | None:
|
||||
from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415
|
||||
get_async_httpx_client, # pyright: ignore
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415
|
||||
|
||||
# litellm's httpx handler and httpx.Response are only partially typed; the IdP returns a JSON
|
||||
# object and the refresher validates each field, so the untyped boundary is contained here.
|
||||
provider = httpxSpecialProvider.Oauth2Check
|
||||
headers = {"Accept": "application/json"}
|
||||
# A failed refresh is a miss, not a 500 (matches v1), so any error becomes None.
|
||||
try:
|
||||
client = get_async_httpx_client(llm_provider=provider) # pyright: ignore
|
||||
response = await client.post(url, headers=headers, data=form) # pyright: ignore
|
||||
response.raise_for_status() # pyright: ignore
|
||||
body: dict[str, object] = response.json() # pyright: ignore
|
||||
except Exception as exc: # noqa: BLE001
|
||||
verbose_logger.warning("MCP OAuth refresh request failed: %s", exc)
|
||||
return None
|
||||
else:
|
||||
return body # pyright: ignore
|
||||
|
||||
|
||||
def build_per_user_oauth_token_store(
|
||||
server_lookup: ServerLookup,
|
||||
) -> CachedOAuthTokenStore:
|
||||
refresher = AuthorizationCodeRefresher(
|
||||
server_lookup, _post_token_endpoint, _persist_credential
|
||||
)
|
||||
# Cache and refresh coordinator use the foundation's in-process defaults (a single replica needs
|
||||
# no shared cache or lock); the cross-replica path is layered on separately.
|
||||
refreshing = RefreshingTokenStore(V2PerUserTokenStore(_read_credential), refresher)
|
||||
return CachedOAuthTokenStore(refreshing, default_ttl_seconds=_DEFAULT_TTL_SECONDS)
|
||||
|
||||
|
||||
class LazyPerUserOAuthTokenStore:
|
||||
"""``OAuthTokenStore`` that builds the v2-native chain on first ``fetch``.
|
||||
|
||||
The chain's cache/lock collaborators are LiteLLM runtime globals not available when the resolver
|
||||
is constructed at import time, so construction is deferred to the first request (by when they are
|
||||
wired). Built once, then reused.
|
||||
"""
|
||||
|
||||
def __init__(self, server_lookup: ServerLookup) -> None:
|
||||
self._server_lookup = server_lookup
|
||||
self._store: CachedOAuthTokenStore | None = None
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
if self._store is None:
|
||||
self._store = build_per_user_oauth_token_store(self._server_lookup)
|
||||
return await self._store.fetch(user_id, server_id)
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
"""One-shot ``OAuthTokenStore`` for the create/test tools preview.
|
||||
|
||||
The preview tests an unsaved server, so no per-user credential is persisted yet. The operator holds
|
||||
the just-authorized token; this serves it through the same v2 resolver path runtime uses for the
|
||||
stored token, so the preview never relies on the caller-credential-override path that
|
||||
``_create_mcp_client`` refuses for ``authorization_code``. It backs a single preview call, so it
|
||||
returns its one token regardless of the lookup key.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PresentedOAuthTokenStore:
|
||||
"""Serves one in-hand token for the single preview call it backs (no DB, no cache)."""
|
||||
|
||||
token: OAuthToken
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
return self.token
|
||||
|
|
@ -7,9 +7,9 @@ no precedence cascade. It is wildcard-free with an `assert_never` tail, so addin
|
|||
an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly
|
||||
at runtime instead of returning `None`.
|
||||
|
||||
`none` and `api_key` (shared-key source) are live; the remaining arms are `not_implemented`
|
||||
stubs that each land in a follow-up PR with their injected seam. The self-contained arms read
|
||||
straight from the config and need no collaborator. Pure v2: no imports from v1.
|
||||
`none` and `api_key` (shared-key source) are live, as is `authorization_code`, which reads the
|
||||
user's token from the injected `OAuthTokenStore`. The remaining arms are `not_implemented` stubs
|
||||
that each land in a follow-up PR with their seam. Pure v2: no imports from v1.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -21,6 +21,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth impo
|
|||
NoOpAuth,
|
||||
StaticHeaderAuth,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
OAuthTokenStore,
|
||||
TokenStoreUnavailable,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
|
||||
Error,
|
||||
Ok,
|
||||
|
|
@ -43,13 +48,26 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
|||
)
|
||||
|
||||
|
||||
class _NullOAuthTokenStore:
|
||||
"""Fail-closed default: with no token store wired, every user reads as not authorized."""
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
return None
|
||||
|
||||
|
||||
class UpstreamCredentialProvider:
|
||||
"""Produces the one `httpx.Auth` for a `(subject, upstream)` pair, per declared mode.
|
||||
|
||||
Collaborators (the per-mode credential stores and token fetchers) are injected as each arm
|
||||
is built; the live `none` and `api_key`-shared arms read from the config and need none.
|
||||
Collaborators (the per-mode credential stores and token fetchers) are injected as each arm is
|
||||
built; the live `none` and `api_key`-shared arms read from the config and need none, while
|
||||
`authorization_code` reads the user's token from the injected `OAuthTokenStore`.
|
||||
"""
|
||||
|
||||
def __init__(self, oauth_token_store: OAuthTokenStore | None = None) -> None:
|
||||
self._oauth_token_store: OAuthTokenStore = (
|
||||
oauth_token_store or _NullOAuthTokenStore()
|
||||
)
|
||||
|
||||
async def resolve_credentials(
|
||||
self, subject: Subject, server: ServerSpec
|
||||
) -> Result[httpx.Auth, CredError]:
|
||||
|
|
@ -65,11 +83,21 @@ class UpstreamCredentialProvider:
|
|||
case TokenExchangeConfig():
|
||||
return _not_implemented(AuthSpecKind.token_exchange)
|
||||
case AuthorizationCodeConfig():
|
||||
return _not_implemented(AuthSpecKind.authorization_code)
|
||||
return await self._authorization_code(subject, server)
|
||||
case AwsSigV4Config():
|
||||
return _not_implemented(AuthSpecKind.aws_sigv4)
|
||||
assert_never(server.config)
|
||||
|
||||
async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool:
|
||||
"""Whether a usable per-user token exists for this server (the preemptive 401's check).
|
||||
|
||||
Reads from the same per-user store as the ``authorization_code`` arm, so the discovery
|
||||
challenge and the egress agree on whether the user is authorized. Returns a typed ``bool``
|
||||
(no ``httpx.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the
|
||||
store, so it reads as False without a per-mode branch here.
|
||||
"""
|
||||
return await self._authz_token(subject, server) is not None
|
||||
|
||||
def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]:
|
||||
match config.key_source:
|
||||
case SharedKey() as source:
|
||||
|
|
@ -86,6 +114,37 @@ class UpstreamCredentialProvider:
|
|||
)
|
||||
assert_never(config.key_source)
|
||||
|
||||
async def _authorization_code(
|
||||
self, subject: Subject, server: ServerSpec
|
||||
) -> Result[StaticHeaderAuth, CredError]:
|
||||
token = await self._authz_token(subject, server)
|
||||
if token is None:
|
||||
return Error(
|
||||
CredError.of_unauthorized(
|
||||
"Authorization required: complete the OAuth flow for this server."
|
||||
)
|
||||
)
|
||||
return Ok(
|
||||
StaticHeaderAuth(
|
||||
f"Bearer {token.access_token}", header_name="Authorization"
|
||||
)
|
||||
)
|
||||
|
||||
async def _authz_token(
|
||||
self, subject: Subject, server: ServerSpec
|
||||
) -> OAuthToken | None:
|
||||
"""The user's authorization_code token, or None when absent or the store is unreachable.
|
||||
|
||||
A store outage is mapped to None (the OAuth challenge), not raised, so a transient outage
|
||||
does not 500; it is the store, not this resolver, that declines to cache the failure.
|
||||
"""
|
||||
try:
|
||||
return await self._oauth_token_store.fetch(
|
||||
subject.subject_id, server.server_id
|
||||
)
|
||||
except TokenStoreUnavailable:
|
||||
return None
|
||||
|
||||
|
||||
def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]:
|
||||
return Error(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,75 @@
|
|||
"""v2-native per-user OAuth token read store for the ``authorization_code`` mode.
|
||||
|
||||
The raw "inner" store that ``RefreshingTokenStore`` and ``CachedOAuthTokenStore`` wrap: it reads the
|
||||
user's persisted credential and returns a typed ``OAuthToken`` (access token, epoch expiry, refresh
|
||||
token), validating the decoded credential blob at this boundary so no ``Any`` leaks past it. It does
|
||||
not cache or refresh - those are the decorators. This replaces ``V1PerUserTokenStore`` (which handed
|
||||
the whole read + cache + refresh to v1's core) as step 1b: the ``read_credential`` collaborator is
|
||||
injected, so the DB/decoding plumbing stays testable and out of this seam.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
|
||||
CredentialReader = Callable[[str, str], Awaitable["dict[str, object] | None"]]
|
||||
|
||||
|
||||
def _iso_to_epoch(expires_at: str) -> float | None:
|
||||
try:
|
||||
dt = datetime.fromisoformat(expires_at)
|
||||
except ValueError:
|
||||
return None
|
||||
# A timezone-naive expiry is stored as UTC (db.py writes ``datetime.now(timezone.utc)``),
|
||||
# so anchor it to UTC before ``.timestamp()`` - otherwise a non-UTC host would read it as
|
||||
# local time and skew the expiry, diverging from v1's ``_remaining_token_seconds``.
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return dt.timestamp()
|
||||
|
||||
|
||||
def _to_scopes(raw: object) -> tuple[str, ...]:
|
||||
if isinstance(raw, (list, tuple)):
|
||||
return tuple(s for s in raw if isinstance(s, str))
|
||||
return ()
|
||||
|
||||
|
||||
def _to_oauth_token(payload: dict[str, object]) -> OAuthToken | None:
|
||||
access_token = payload.get("access_token")
|
||||
if not isinstance(access_token, str):
|
||||
return None
|
||||
refresh_token = payload.get("refresh_token")
|
||||
expires_at = payload.get("expires_at")
|
||||
return OAuthToken(
|
||||
access_token=access_token,
|
||||
expires_at=_iso_to_epoch(expires_at) if isinstance(expires_at, str) else None,
|
||||
refresh_token=refresh_token if isinstance(refresh_token, str) else None,
|
||||
scopes=_to_scopes(payload.get("scopes")),
|
||||
)
|
||||
|
||||
|
||||
class V2PerUserTokenStore:
|
||||
"""``OAuthTokenStore`` that reads the user's persisted authorization_code credential, typed.
|
||||
|
||||
The injected ``read_credential`` returns the decoded credential payload for a ``(user, server)``
|
||||
pair, or ``None`` when the user has not completed OAuth. A backing-store outage surfaces as
|
||||
``TokenStoreUnavailable`` from the reader, which the arm turns into a challenge rather than a
|
||||
500, so ``fetch`` lets it propagate. Refresh is the wrapping ``RefreshingTokenStore``'s job, so
|
||||
the returned token carries ``expires_at`` and ``refresh_token`` for it to act on.
|
||||
"""
|
||||
|
||||
def __init__(self, read_credential: CredentialReader) -> None:
|
||||
self._read_credential = read_credential
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None:
|
||||
if not user_id:
|
||||
return None
|
||||
payload = await self._read_credential(user_id, server_id)
|
||||
if payload is None:
|
||||
return None
|
||||
return _to_oauth_token(payload)
|
||||
|
|
@ -1034,8 +1034,54 @@ if MCP_AVAILABLE:
|
|||
None if server_model.has_client_credentials else oauth2_headers
|
||||
)
|
||||
|
||||
# Interactive authorization_code tools preview: the operator holds a just-authorized
|
||||
# token but it is not persisted yet. Resolve it through the v2 resolver via a one-shot
|
||||
# presented store - the same path runtime uses for the stored token - rather than the
|
||||
# caller-override path _create_mcp_client refuses for authorization_code. The bare token
|
||||
# becomes the upstream credential, so it is not also forwarded as a caller header. Gated
|
||||
# to the v2-mapped oauth2 case (to_server_spec non-None); M2M (client_credentials),
|
||||
# delegate/passthrough, and token-exchange are unaffected.
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials import ( # noqa: PLC0415
|
||||
UpstreamCredentialProvider,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415
|
||||
to_server_spec,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( # noqa: PLC0415
|
||||
OAuthToken,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.presented_token_store import ( # noqa: PLC0415
|
||||
PresentedOAuthTokenStore,
|
||||
)
|
||||
|
||||
forwarded_authorization = (
|
||||
effective_oauth2_headers.get("Authorization")
|
||||
if effective_oauth2_headers
|
||||
else None
|
||||
)
|
||||
is_interactive_authz_code = (
|
||||
server_model.auth_type == MCPAuth.oauth2
|
||||
and forwarded_authorization is not None
|
||||
and to_server_spec(server_model) is not None
|
||||
)
|
||||
preview_cred_provider = (
|
||||
UpstreamCredentialProvider(
|
||||
oauth_token_store=PresentedOAuthTokenStore(
|
||||
OAuthToken(
|
||||
access_token=forwarded_authorization[7:]
|
||||
if forwarded_authorization[:7].lower() == "bearer "
|
||||
else forwarded_authorization
|
||||
)
|
||||
)
|
||||
)
|
||||
if is_interactive_authz_code
|
||||
else None
|
||||
)
|
||||
|
||||
merged_headers = merge_mcp_headers(
|
||||
extra_headers=effective_oauth2_headers,
|
||||
extra_headers=(
|
||||
None if preview_cred_provider else effective_oauth2_headers
|
||||
),
|
||||
static_headers=request.static_headers,
|
||||
)
|
||||
|
||||
|
|
@ -1044,6 +1090,7 @@ if MCP_AVAILABLE:
|
|||
mcp_auth_header=mcp_auth_header,
|
||||
extra_headers=merged_headers,
|
||||
stdio_env=stdio_env,
|
||||
cred_provider=preview_cred_provider,
|
||||
)
|
||||
|
||||
return await operation(client)
|
||||
|
|
|
|||
|
|
@ -283,6 +283,7 @@ if MCP_AVAILABLE:
|
|||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
MCPServerManager,
|
||||
_should_strip_caller_authorization,
|
||||
_without_authorization,
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
|
||||
|
|
@ -1360,115 +1361,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],
|
||||
|
|
@ -1529,6 +1436,16 @@ if MCP_AVAILABLE:
|
|||
else:
|
||||
# Copy to avoid mutating the original dict (important for parallel fetching)
|
||||
extra_headers = oauth2_headers.copy() if oauth2_headers else None
|
||||
# Migrated authorization_code: the v2 resolver injects the stored per-user
|
||||
# token, so drop the caller-forwarded Authorization (apply-if-absent would
|
||||
# otherwise let it shadow the resolved token). Delegate keeps it. Centralized
|
||||
# via _should_strip_caller_authorization to match _call_regular_mcp_tool.
|
||||
if extra_headers and _should_strip_caller_authorization(
|
||||
mcp_server=server,
|
||||
raw_headers=raw_headers,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
):
|
||||
extra_headers = _without_authorization(extra_headers)
|
||||
|
||||
if server.extra_headers and raw_headers:
|
||||
if extra_headers is None:
|
||||
|
|
@ -1778,8 +1695,17 @@ if MCP_AVAILABLE:
|
|||
# Prefer server-stored per-user OAuth when configured, so a stale
|
||||
# Authorization header from the MCP client cannot override Redis/DB
|
||||
# (same issue as call_tool in mcp_server_manager: VS Code caches tokens).
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415
|
||||
to_server_spec,
|
||||
)
|
||||
|
||||
# A server migrated to the v2 resolver gets its token from the resolver at connect
|
||||
# time; building it here would double-resolve and be shadowed by the v2 graft. The
|
||||
# preemptive 401 already challenged a missing token, so one exists for the connect.
|
||||
migrated_to_v2 = to_server_spec(server) is not None
|
||||
if (
|
||||
server.auth_type == MCPAuth.oauth2
|
||||
not migrated_to_v2
|
||||
and server.auth_type == MCPAuth.oauth2
|
||||
and getattr(server, "needs_user_oauth_token", False)
|
||||
and user_api_key_auth is not None
|
||||
):
|
||||
|
|
@ -1792,7 +1718,11 @@ if MCP_AVAILABLE:
|
|||
extra_headers = db_headers
|
||||
|
||||
# If still no OAuth2 token, fall back to pre-fetched creds (non-stale-client path)
|
||||
elif extra_headers is None and server.auth_type == MCPAuth.oauth2:
|
||||
elif (
|
||||
not migrated_to_v2
|
||||
and extra_headers is None
|
||||
and server.auth_type == MCPAuth.oauth2
|
||||
):
|
||||
extra_headers = await _get_user_oauth_extra_headers_from_db(
|
||||
server,
|
||||
user_api_key_auth,
|
||||
|
|
@ -3501,12 +3431,6 @@ if MCP_AVAILABLE:
|
|||
# If no stored token exists, fail fast with 401 so clients can
|
||||
# kick off PKCE/interactive OAuth flow immediately.
|
||||
if server.needs_user_oauth_token:
|
||||
stored_oauth_headers = await _get_user_oauth_extra_headers_from_db(
|
||||
server=server,
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
)
|
||||
if stored_oauth_headers:
|
||||
continue
|
||||
if getattr(server, "delegate_auth_to_upstream", False) is True:
|
||||
# Delegate-auth servers run upstream PKCE: challenge with
|
||||
# the proxied resource_metadata (RFC 9728), not the
|
||||
|
|
@ -3521,6 +3445,12 @@ if MCP_AVAILABLE:
|
|||
detail="Unauthorized",
|
||||
headers={"www-authenticate": www_authenticate},
|
||||
)
|
||||
# The v2 resolver owns the existence check, so every authorization_code
|
||||
# resolution (egress and this discovery challenge) runs through it.
|
||||
if await global_mcp_server_manager.has_user_oauth_token(
|
||||
server, user_api_key_auth
|
||||
):
|
||||
continue
|
||||
|
||||
request = StarletteRequest(scope)
|
||||
base_url = get_request_base_url(request)
|
||||
|
|
|
|||
|
|
@ -7,17 +7,20 @@ maps each CredError onto its HTTP status. These pin the parity-critical mapping
|
|||
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
|
||||
raise_public,
|
||||
raise_user_oauth_challenge,
|
||||
to_server_spec,
|
||||
to_subject,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
||||
ApiKeyConfig,
|
||||
AuthorizationCodeConfig,
|
||||
CredError,
|
||||
NoneConfig,
|
||||
SharedKey,
|
||||
|
|
@ -71,12 +74,27 @@ def test_basic_scheme_base64_encodes_the_token():
|
|||
assert spec.config.key_source.value.get_secret_value() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"oauth2_flow",
|
||||
[None, "authorization_code"],
|
||||
)
|
||||
def test_oauth2_user_token_maps_to_authorization_code(oauth2_flow):
|
||||
# oauth2 without client_credentials is the per-user authorization_code mode.
|
||||
spec = to_server_spec(_server(auth_type=MCPAuth.oauth2, oauth2_flow=oauth2_flow))
|
||||
assert spec is not None and isinstance(spec.config, AuthorizationCodeConfig)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"server",
|
||||
[
|
||||
_server(auth_type=MCPAuth.api_key), # no token configured
|
||||
_server(auth_type=MCPAuth.bearer_token), # no token configured
|
||||
_server(auth_type=MCPAuth.oauth2),
|
||||
_server(
|
||||
auth_type=MCPAuth.oauth2, oauth2_flow="client_credentials"
|
||||
), # M2M -> v1
|
||||
_server(
|
||||
auth_type=MCPAuth.oauth2, delegate_auth_to_upstream=True
|
||||
), # delegated upstream OAuth -> v1
|
||||
_server(auth_type=MCPAuth.oauth2_token_exchange),
|
||||
_server(auth_type=MCPAuth.aws_sigv4),
|
||||
_server(
|
||||
|
|
@ -162,3 +180,43 @@ def test_raise_public_plain_unauthorized_has_no_challenge():
|
|||
assert exc.status_code == 401
|
||||
assert exc.detail == "unauthorized: nope"
|
||||
assert exc.headers is None
|
||||
|
||||
|
||||
_ROOT_PATH = "litellm.proxy.utils.get_server_root_path"
|
||||
|
||||
|
||||
def test_raise_user_oauth_challenge_points_at_per_server_prm():
|
||||
with patch(_ROOT_PATH, return_value="/"), pytest.raises(HTTPException) as exc_info:
|
||||
raise_user_oauth_challenge(_server(alias="my-srv"))
|
||||
exc = exc_info.value
|
||||
assert exc.status_code == 401
|
||||
assert (
|
||||
exc.headers["WWW-Authenticate"]
|
||||
== 'Bearer resource_metadata="/.well-known/oauth-protected-resource/mcp/my-srv"'
|
||||
)
|
||||
|
||||
|
||||
def test_raise_user_oauth_challenge_includes_server_root_path():
|
||||
with (
|
||||
patch(_ROOT_PATH, return_value="/api/v1"),
|
||||
pytest.raises(HTTPException) as exc_info,
|
||||
):
|
||||
raise_user_oauth_challenge(_server(alias="my-srv"))
|
||||
assert (
|
||||
exc_info.value.headers["WWW-Authenticate"]
|
||||
== 'Bearer resource_metadata="/.well-known/oauth-protected-resource/api/v1/mcp/my-srv"'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs, expected_name",
|
||||
[
|
||||
({"alias": "a", "server_name": "sn"}, "a"), # alias wins
|
||||
({"server_name": "sn"}, "sn"), # then server_name
|
||||
({}, "n"), # then the name field (server_id is the last fallback)
|
||||
],
|
||||
)
|
||||
def test_raise_user_oauth_challenge_name_fallback(kwargs, expected_name):
|
||||
with patch(_ROOT_PATH, return_value="/"), pytest.raises(HTTPException) as exc_info:
|
||||
raise_user_oauth_challenge(_server(**kwargs))
|
||||
assert f'/mcp/{expected_name}"' in exc_info.value.headers["WWW-Authenticate"]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,184 @@
|
|||
"""Tests for the authorization_code refresher: the refresh_token grant, parsing, and persist."""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.authz_code_refresher import (
|
||||
AuthorizationCodeRefresher,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
|
||||
|
||||
class _Server:
|
||||
def __init__(
|
||||
self,
|
||||
token_url="https://idp.example.com/token",
|
||||
client_id="cid",
|
||||
client_secret="sec",
|
||||
):
|
||||
self.token_url = token_url
|
||||
self.client_id = client_id
|
||||
self.client_secret = client_secret
|
||||
|
||||
|
||||
def _lookup(server):
|
||||
return lambda server_id: server
|
||||
|
||||
|
||||
def _endpoint(body, sink=None):
|
||||
async def post(url, form):
|
||||
if sink is not None:
|
||||
sink.append((url, form))
|
||||
return body
|
||||
|
||||
return post
|
||||
|
||||
|
||||
def _recording_persist(sink):
|
||||
async def persist(
|
||||
user_id, server_id, access_token, refresh_token, expires_in, scopes
|
||||
):
|
||||
sink.append(
|
||||
(user_id, server_id, access_token, refresh_token, expires_in, scopes)
|
||||
)
|
||||
|
||||
return persist
|
||||
|
||||
|
||||
def _refresher(
|
||||
server=None, body=None, *, post_sink=None, persist_sink=None, clock=lambda: 1000.0
|
||||
):
|
||||
return AuthorizationCodeRefresher(
|
||||
_lookup(server if server is not None else _Server()),
|
||||
_endpoint(body, post_sink),
|
||||
_recording_persist(persist_sink if persist_sink is not None else []),
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refreshes_persists_and_returns_typed_token():
|
||||
persisted = []
|
||||
posted = []
|
||||
refresher = _refresher(
|
||||
body={
|
||||
"access_token": "new-at",
|
||||
"expires_in": 3600,
|
||||
"refresh_token": "new-rt",
|
||||
"scope": "a b",
|
||||
},
|
||||
post_sink=posted,
|
||||
persist_sink=persisted,
|
||||
)
|
||||
token = await refresher.refresh(
|
||||
"alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")
|
||||
)
|
||||
|
||||
assert token is not None
|
||||
assert token.access_token == "new-at"
|
||||
assert token.refresh_token == "new-rt"
|
||||
assert token.expires_at == 1000.0 + 3600 # clock + expires_in -> epoch
|
||||
# the rotated triple is persisted for (user, server) with parsed scopes
|
||||
assert persisted == [("alice", "srv", "new-at", "new-rt", 3600, ("a", "b"))]
|
||||
# the grant carried the refresh_token + client credentials
|
||||
url, form = posted[0]
|
||||
assert url == "https://idp.example.com/token"
|
||||
assert form == {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": "old-rt",
|
||||
"client_id": "cid",
|
||||
"client_secret": "sec",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_refresh_token_is_not_refreshable():
|
||||
posted = []
|
||||
refresher = _refresher(body={"access_token": "x"}, post_sink=posted)
|
||||
assert (
|
||||
await refresher.refresh("alice", "srv", OAuthToken(access_token="old")) is None
|
||||
)
|
||||
assert posted == [] # never hit the IdP
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_server_or_no_token_url_yields_none():
|
||||
assert (
|
||||
await _refresher(server=None).refresh(
|
||||
"a", "s", OAuthToken("old", refresh_token="rt")
|
||||
)
|
||||
is None
|
||||
)
|
||||
no_url = _Server(token_url=None)
|
||||
assert (
|
||||
await _refresher(server=no_url).refresh(
|
||||
"a", "s", OAuthToken("old", refresh_token="rt")
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_grant_failure_does_not_persist():
|
||||
persisted = []
|
||||
refresher = _refresher(
|
||||
body=None, persist_sink=persisted
|
||||
) # token_endpoint signals failure
|
||||
assert (
|
||||
await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
|
||||
)
|
||||
assert persisted == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_without_access_token_does_not_persist():
|
||||
persisted = []
|
||||
refresher = _refresher(body={"expires_in": 60}, persist_sink=persisted)
|
||||
assert (
|
||||
await refresher.refresh("a", "s", OAuthToken("old", refresh_token="rt")) is None
|
||||
)
|
||||
assert persisted == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unrotated_refresh_token_is_carried_forward():
|
||||
persisted = []
|
||||
refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted)
|
||||
token = await refresher.refresh(
|
||||
"a", "s", OAuthToken("old", refresh_token="keep-rt")
|
||||
)
|
||||
assert token is not None
|
||||
assert (
|
||||
token.refresh_token == "keep-rt"
|
||||
) # response omitted refresh_token -> reuse the old one
|
||||
assert token.expires_at is None # no expires_in -> no known expiry
|
||||
assert persisted[0][3] == "keep-rt"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unrecorded_scope_is_carried_forward():
|
||||
persisted = []
|
||||
refresher = _refresher(body={"access_token": "new-at"}, persist_sink=persisted)
|
||||
token = await refresher.refresh(
|
||||
"a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write"))
|
||||
)
|
||||
assert token is not None
|
||||
# response omitted "scope" -> the user's recorded grant is preserved, not dropped
|
||||
assert token.scopes == ("read", "write")
|
||||
assert persisted[0][5] == ("read", "write")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returned_scope_overrides_prior_when_present():
|
||||
persisted = []
|
||||
refresher = _refresher(
|
||||
body={"access_token": "new-at", "scope": "read"},
|
||||
persist_sink=persisted,
|
||||
)
|
||||
token = await refresher.refresh(
|
||||
"a", "s", OAuthToken("old", refresh_token="rt", scopes=("read", "write"))
|
||||
)
|
||||
assert token is not None
|
||||
assert token.scopes == ("read",) # a present scope replaces the prior grant
|
||||
assert persisted[0][5] == ("read",)
|
||||
|
|
@ -0,0 +1,19 @@
|
|||
"""Tests for the create/test-preview presented OAuth token store."""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.presented_token_store import (
|
||||
PresentedOAuthTokenStore,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_serves_the_presented_token_regardless_of_key():
|
||||
token = OAuthToken(access_token="at", scopes=("read",))
|
||||
store = PresentedOAuthTokenStore(token)
|
||||
# one-shot store backs a single preview call, so the lookup key is irrelevant
|
||||
assert await store.fetch("alice", "srv-1") is token
|
||||
assert await store.fetch("", "") is token
|
||||
|
|
@ -1,9 +1,9 @@
|
|||
"""Tests for the resolver dispatch: live arms produce auth, stubbed arms fail closed.
|
||||
|
||||
`none` and `api_key` (shared-key source) are implemented; every other arm, plus the `api_key`
|
||||
BYOK source, returns a typed `not_implemented` error until its mode lands. Parametrizing the
|
||||
stubs over one config each also guards reachability: a dropped `case` would hit `assert_never`
|
||||
and raise instead of returning the stub.
|
||||
`none`, `api_key` (shared-key source), and `authorization_code` are implemented; every other arm,
|
||||
plus the `api_key` BYOK source, returns a typed `not_implemented` error until its mode lands.
|
||||
Parametrizing the stubs over one config each also guards reachability: a dropped `case` would hit
|
||||
`assert_never` and raise instead of returning the stub.
|
||||
"""
|
||||
|
||||
import httpx
|
||||
|
|
@ -28,6 +28,10 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
|
|||
TokenExchangeConfig,
|
||||
UpstreamCredentialProvider,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
TokenStoreUnavailable,
|
||||
)
|
||||
|
||||
_SUBJECT = Subject(tenant_id="", subject_id="")
|
||||
|
||||
|
|
@ -84,12 +88,114 @@ async def test_api_key_shared_honors_authorization_scheme():
|
|||
assert _emitted(result.ok)["Authorization"] == "Bearer tok"
|
||||
|
||||
|
||||
class _FakeTokenStore:
|
||||
"""An OAuthTokenStore returning a canned per-user token (None == not authorized)."""
|
||||
|
||||
def __init__(self, by_user: dict) -> None:
|
||||
self._by_user = by_user
|
||||
|
||||
async def fetch(self, user_id: str, server_id: str):
|
||||
return self._by_user.get((user_id, server_id))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorization_code_emits_bearer_for_a_stored_token():
|
||||
store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="at-alice")})
|
||||
result = await UpstreamCredentialProvider(
|
||||
oauth_token_store=store
|
||||
).resolve_credentials(
|
||||
Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig())
|
||||
)
|
||||
assert isinstance(result, Ok)
|
||||
assert _emitted(result.ok)["Authorization"] == "Bearer at-alice"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorization_code_without_token_is_semantically_unauthorized():
|
||||
result = await UpstreamCredentialProvider(
|
||||
oauth_token_store=_FakeTokenStore({})
|
||||
).resolve_credentials(
|
||||
Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig())
|
||||
)
|
||||
assert isinstance(result, Error)
|
||||
assert result.error.tag == "unauthorized"
|
||||
# Semantic only: the per-server challenge is built at the edge, not in the request-free arm.
|
||||
assert "Authorization required" in result.error.unauthorized.detail
|
||||
assert result.error.unauthorized.www_authenticate is None
|
||||
assert result.error.unauthorized.body is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorization_code_store_unavailable_is_unauthorized():
|
||||
class _Unavailable:
|
||||
async def fetch(self, user_id: str, server_id: str):
|
||||
raise TokenStoreUnavailable("down")
|
||||
|
||||
result = await UpstreamCredentialProvider(
|
||||
oauth_token_store=_Unavailable()
|
||||
).resolve_credentials(
|
||||
Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig())
|
||||
)
|
||||
assert isinstance(result, Error)
|
||||
assert result.error.tag == "unauthorized"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorization_code_with_no_store_wired_is_unauthorized():
|
||||
result = await UpstreamCredentialProvider().resolve_credentials(
|
||||
Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig())
|
||||
)
|
||||
assert isinstance(result, Error)
|
||||
assert result.error.tag == "unauthorized"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorization_code_isolates_by_subject():
|
||||
store = _FakeTokenStore({("alice", "s"): OAuthToken(access_token="at-alice")})
|
||||
provider = UpstreamCredentialProvider(oauth_token_store=store)
|
||||
alice = await provider.resolve_credentials(
|
||||
Subject(tenant_id="", subject_id="alice"), _spec(AuthorizationCodeConfig())
|
||||
)
|
||||
bob = await provider.resolve_credentials(
|
||||
Subject(tenant_id="", subject_id="bob"), _spec(AuthorizationCodeConfig())
|
||||
)
|
||||
assert (
|
||||
isinstance(alice, Ok)
|
||||
and _emitted(alice.ok)["Authorization"] == "Bearer at-alice"
|
||||
)
|
||||
assert isinstance(bob, Error) and bob.error.tag == "unauthorized"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_has_user_token_reflects_the_stored_token():
|
||||
present = UpstreamCredentialProvider(
|
||||
oauth_token_store=_FakeTokenStore(
|
||||
{("alice", "s"): OAuthToken(access_token="at")}
|
||||
)
|
||||
)
|
||||
absent = UpstreamCredentialProvider(oauth_token_store=_FakeTokenStore({}))
|
||||
spec = _spec(AuthorizationCodeConfig())
|
||||
subject = Subject(tenant_id="", subject_id="alice")
|
||||
assert await present.has_user_token(subject, spec) is True
|
||||
assert await absent.has_user_token(subject, spec) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_has_user_token_false_for_a_non_per_user_mode():
|
||||
# A none-mode server has no per-user token to check.
|
||||
provider = UpstreamCredentialProvider()
|
||||
spec = _spec(NoneConfig())
|
||||
assert (
|
||||
await provider.has_user_token(Subject(tenant_id="", subject_id="a"), spec)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
_STUBBED = [
|
||||
("api_key_byok", ApiKeyConfig(key_source=Byok())),
|
||||
("passthrough", PassthroughConfig()),
|
||||
("client_credentials", ClientCredentialsConfig()),
|
||||
("token_exchange", TokenExchangeConfig()),
|
||||
("authorization_code", AuthorizationCodeConfig()),
|
||||
("aws_sigv4", AwsSigV4Config(region="us-east-1")),
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,73 @@
|
|||
"""Tests for the v2-native per-user token read store: validate the decoded blob into a typed token."""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.v2_token_store import (
|
||||
V2PerUserTokenStore,
|
||||
)
|
||||
|
||||
|
||||
def _reader(payload):
|
||||
async def read(user_id: str, server_id: str):
|
||||
return payload
|
||||
|
||||
return read
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_builds_typed_token_with_epoch_expiry():
|
||||
store = V2PerUserTokenStore(
|
||||
_reader(
|
||||
{
|
||||
"access_token": "at",
|
||||
"expires_at": "2099-01-01T00:00:00+00:00",
|
||||
"refresh_token": "rt",
|
||||
}
|
||||
)
|
||||
)
|
||||
token = await store.fetch("alice", "srv")
|
||||
assert token is not None
|
||||
assert token.access_token == "at"
|
||||
assert token.refresh_token == "rt"
|
||||
# ISO string is converted to epoch seconds for RefreshingTokenStore to compare against.
|
||||
assert token.expires_at == 4070908800.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_optional_fields_absent_yield_none():
|
||||
store = V2PerUserTokenStore(_reader({"access_token": "at"}))
|
||||
token = await store.fetch("alice", "srv")
|
||||
assert token is not None
|
||||
assert token.expires_at is None and token.refresh_token is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unparseable_expiry_is_dropped_not_raised():
|
||||
store = V2PerUserTokenStore(
|
||||
_reader({"access_token": "at", "expires_at": "not-a-date"})
|
||||
)
|
||||
token = await store.fetch("alice", "srv")
|
||||
assert token is not None and token.expires_at is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_access_token_is_not_authorized():
|
||||
store = V2PerUserTokenStore(_reader({"expires_at": "2099-01-01T00:00:00+00:00"}))
|
||||
assert await store.fetch("alice", "srv") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_credential_is_not_authorized():
|
||||
assert await V2PerUserTokenStore(_reader(None)).fetch("alice", "srv") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_user_short_circuits_without_reading():
|
||||
calls = []
|
||||
|
||||
async def read(user_id: str, server_id: str):
|
||||
calls.append((user_id, server_id))
|
||||
return {"access_token": "at"}
|
||||
|
||||
assert await V2PerUserTokenStore(read).fetch("", "srv") is None
|
||||
assert calls == []
|
||||
|
|
@ -284,8 +284,11 @@ def test_prepare_mcp_server_headers_oauth2_m2m_omits_litellm_caller_authorizatio
|
|||
assert extra_headers is None
|
||||
|
||||
|
||||
def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers():
|
||||
"""Interactive OAuth still forwards the user's OAuth token in extra_headers."""
|
||||
def test_prepare_mcp_server_headers_oauth2_interactive_drops_caller_authorization():
|
||||
"""A v2-migrated interactive OAuth (authorization_code) server must NOT forward the
|
||||
caller's Authorization: the resolver injects the stored per-user token, so a
|
||||
caller-supplied bearer must not override another user's stored credential. Non-auth
|
||||
headers are still carried; only the credential is dropped."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_prepare_mcp_server_headers,
|
||||
|
|
@ -293,7 +296,7 @@ def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers():
|
|||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
user_oauth = {"Authorization": "Bearer upstream-user-token"}
|
||||
caller_oauth = {"Authorization": "Bearer caller-supplied-token"}
|
||||
|
||||
server = MCPServer(
|
||||
server_id="3lo-server",
|
||||
|
|
@ -307,12 +310,13 @@ def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers():
|
|||
server=server,
|
||||
mcp_server_auth_headers=None,
|
||||
mcp_auth_header=None,
|
||||
oauth2_headers=user_oauth,
|
||||
oauth2_headers=caller_oauth,
|
||||
raw_headers=None,
|
||||
)
|
||||
|
||||
assert server_auth_header is None
|
||||
assert extra_headers == user_oauth
|
||||
# Caller's Authorization is dropped (only key present) -> extra_headers is None.
|
||||
assert extra_headers is None
|
||||
|
||||
|
||||
def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers():
|
||||
|
|
@ -2813,8 +2817,10 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.no_parallel
|
||||
async def test_oauth2_headers_passed_to_mcp_client():
|
||||
"""Test that OAuth2 headers are properly passed through to the MCP client for OAuth2 servers like github_mcp"""
|
||||
async def test_oauth2_caller_headers_not_forwarded_for_migrated_server():
|
||||
"""A v2-migrated authorization_code server (like github_mcp) must NOT forward the
|
||||
caller's oauth2 Authorization to the MCP client — the resolver injects the stored
|
||||
per-user token, so a caller-supplied bearer cannot override another user's credential."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
|
|
@ -2928,20 +2934,13 @@ async def test_oauth2_headers_passed_to_mcp_client():
|
|||
assert captured_client_args["server"].server_id == oauth2_server.server_id
|
||||
assert captured_client_args["server"].auth_type == MCPAuth.oauth2
|
||||
|
||||
# Most importantly: verify that OAuth2 headers were passed as extra_headers
|
||||
assert (
|
||||
captured_client_args["extra_headers"] is not None
|
||||
), "Expected extra_headers to be passed for OAuth2 server"
|
||||
assert (
|
||||
captured_client_args["extra_headers"] == oauth2_headers
|
||||
), f"Expected OAuth2 headers to be passed as extra_headers, got {captured_client_args['extra_headers']}"
|
||||
|
||||
# Verify the Authorization header specifically
|
||||
assert "Authorization" in captured_client_args["extra_headers"]
|
||||
assert (
|
||||
captured_client_args["extra_headers"]["Authorization"]
|
||||
== "Bearer github_oauth_token_12345"
|
||||
)
|
||||
# Security: a v2-migrated authorization_code server must NOT forward the caller's
|
||||
# oauth2 Authorization upstream. The v2 resolver injects the stored per-user token,
|
||||
# so a caller-supplied bearer cannot override another user's stored credential.
|
||||
extra_headers = captured_client_args["extra_headers"]
|
||||
assert extra_headers is None or "Authorization" not in {
|
||||
k.lower() for k in extra_headers
|
||||
}, f"Caller Authorization must not be forwarded, got {extra_headers}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -31,6 +31,8 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
|||
_deserialize_json_dict,
|
||||
_deserialize_json_list,
|
||||
_normalize_mcp_server_cost_info,
|
||||
_should_strip_caller_authorization,
|
||||
_without_authorization,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_MCPServerTable,
|
||||
|
|
@ -137,6 +139,43 @@ class TestMCPServerManager:
|
|||
assert client.stdio_config["env"]["NODE_ENV"] == "test"
|
||||
assert client.stdio_config["env"]["NPM_CONFIG_CACHE"] == MCP_NPM_CACHE_DIR
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_auth_header_cannot_bypass_v2_for_authorization_code(self):
|
||||
"""A caller-supplied per-request override must not substitute the stored authorization_code
|
||||
token: _create_mcp_client keeps the v2 spec and resolves through the injected provider
|
||||
rather than deferring to the v1 caller-override path."""
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import (
|
||||
StaticHeaderAuth,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import (
|
||||
Ok,
|
||||
)
|
||||
from litellm.types.mcp import MCPAuth
|
||||
|
||||
calls = []
|
||||
|
||||
class _FakeProvider:
|
||||
async def resolve_credentials(self, subject, server):
|
||||
calls.append((subject.subject_id, server.server_id))
|
||||
return Ok(StaticHeaderAuth("stored-token"))
|
||||
|
||||
manager = MCPServerManager(cred_provider=_FakeProvider())
|
||||
server = MCPServer(
|
||||
server_id="authz-srv",
|
||||
name="authz",
|
||||
url="https://upstream.example/mcp",
|
||||
transport=MCPTransport.sse,
|
||||
auth_type=MCPAuth.oauth2, # oauth2 + no client creds + not delegate -> authorization_code
|
||||
)
|
||||
|
||||
client = await manager._create_mcp_client(
|
||||
server, mcp_auth_header="Bearer caller-supplied-token"
|
||||
)
|
||||
|
||||
# the v2 resolver ran (the caller override did NOT defer to v1); the stored token wins
|
||||
assert calls == [("", "authz-srv")]
|
||||
assert client is not None
|
||||
|
||||
async def test_create_mcp_client_stdio_injects_npm_config_cache(self):
|
||||
"""Test that _create_mcp_client injects NPM_CONFIG_CACHE when not already set,
|
||||
and preserves user-provided NPM_CONFIG_CACHE when present."""
|
||||
|
|
@ -600,6 +639,84 @@ class TestMCPServerManager:
|
|||
|
||||
assert captured_extra_headers == {"x-request-id": "req-123"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_regular_mcp_tool_v2_authz_code_drops_caller_authorization(
|
||||
self,
|
||||
):
|
||||
"""A v2-migrated per-user OAuth (authorization_code) server must NOT seed a
|
||||
caller-forwarded Authorization into extra_headers — the resolver injects the
|
||||
stored per-user token, and apply-if-absent would otherwise let the caller's
|
||||
header override another user's stored credential (matches v1's overwrite)."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
# oauth2, not M2M, not delegate => to_server_spec maps it to AuthorizationCodeConfig
|
||||
server = MCPServer(
|
||||
server_id="server-authz-code-call",
|
||||
name="authz-code-server",
|
||||
url="https://example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
# Migrated authorization_code => the centralized strip decision says drop the
|
||||
# caller's Authorization (the v2 resolver injects the stored token).
|
||||
assert (
|
||||
_should_strip_caller_authorization(
|
||||
mcp_server=server, raw_headers=None, user_api_key_auth=None
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
mock_client = AsyncMock()
|
||||
mock_client.call_tool = AsyncMock(
|
||||
return_value=CallToolResult(content=[], isError=False)
|
||||
)
|
||||
captured_extra_headers = "unset"
|
||||
|
||||
async def capture_create_mcp_client(
|
||||
server,
|
||||
mcp_auth_header,
|
||||
extra_headers,
|
||||
stdio_env,
|
||||
subject_token=None,
|
||||
**kwargs,
|
||||
): # pragma: no cover - helper
|
||||
nonlocal captured_extra_headers
|
||||
captured_extra_headers = extra_headers
|
||||
return mock_client
|
||||
|
||||
manager._create_mcp_client = AsyncMock(side_effect=capture_create_mcp_client)
|
||||
|
||||
await manager._call_regular_mcp_tool(
|
||||
mcp_server=server,
|
||||
original_tool_name="tool",
|
||||
arguments={},
|
||||
tasks=[],
|
||||
mcp_auth_header=None,
|
||||
mcp_server_auth_headers=None,
|
||||
oauth2_headers={"Authorization": "Bearer caller-supplied-token"},
|
||||
raw_headers={"authorization": "Bearer caller-supplied-token"},
|
||||
proxy_logging_obj=None,
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-key"),
|
||||
)
|
||||
|
||||
# The caller's Authorization must not reach extra_headers; the v2 resolver is
|
||||
# the sole Authorization source for this server.
|
||||
assert captured_extra_headers != "unset"
|
||||
if captured_extra_headers:
|
||||
assert "authorization" not in {k.lower() for k in captured_extra_headers}
|
||||
|
||||
def test_without_authorization_drops_only_the_credential(self):
|
||||
# None / empty -> None
|
||||
assert _without_authorization(None) is None
|
||||
assert _without_authorization({}) is None
|
||||
# Only Authorization present -> nothing left -> None (case-insensitive)
|
||||
assert _without_authorization({"authorization": "Bearer x"}) is None
|
||||
# Authorization dropped, other headers kept
|
||||
assert _without_authorization(
|
||||
{"Authorization": "Bearer x", "X-Trace-Id": "t"}
|
||||
) == {"X-Trace-Id": "t"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_regular_mcp_tool_passthrough_forwards_authorization_with_admission_header(
|
||||
self,
|
||||
|
|
@ -2221,7 +2338,7 @@ class TestMCPServerManager:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_looks_up_stored_token(self):
|
||||
"""Falls back to stored per-user OAuth headers when no token is supplied."""
|
||||
"""Falls back to stored per-user OAuth headers for a non-migrated (delegate) oauth2 server."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
|
|
@ -2230,6 +2347,7 @@ class TestMCPServerManager:
|
|||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
delegate_auth_to_upstream=True, # stays on v1, so v1 still builds the header
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
stored = {"Authorization": "Bearer stored-user-token"}
|
||||
|
|
@ -2245,9 +2363,36 @@ class TestMCPServerManager:
|
|||
assert result == stored
|
||||
mock_lookup.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_steps_aside_for_migrated_server(self):
|
||||
"""A migrated authorization_code server is owned by the v2 resolver, so v1 must not also
|
||||
build the token into extra_headers (which the v2 graft would defer to and shadow).
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="oauth-srv",
|
||||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2, # per-user, not delegate -> migrated to v2
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
new=AsyncMock(return_value={"Authorization": "Bearer should-not-be-used"}),
|
||||
) as mock_lookup:
|
||||
result = await manager._resolve_oauth2_headers_for_tool_call(
|
||||
server, oauth2_headers=None, user_api_key_auth=user_auth
|
||||
)
|
||||
|
||||
assert result is None # stepped aside; the v2 resolver handles the token
|
||||
mock_lookup.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_swallows_lookup_exception(self):
|
||||
"""Returns supplied headers (None) when the stored-token lookup raises."""
|
||||
"""Returns supplied headers (None) when the v1 stored-token lookup raises (delegate path)."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
|
|
@ -2256,6 +2401,7 @@ class TestMCPServerManager:
|
|||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
delegate_auth_to_upstream=True, # non-migrated, so it reaches the v1 lookup
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
|
||||
|
|
@ -2268,6 +2414,51 @@ class TestMCPServerManager:
|
|||
)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_has_user_oauth_token_delegates_to_provider(self):
|
||||
"""has_user_oauth_token maps the server and delegates the verdict to the v2 resolver."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
for verdict in (True, False):
|
||||
|
||||
class _Provider:
|
||||
async def has_user_token(self, subject, spec):
|
||||
return verdict
|
||||
|
||||
manager = MCPServerManager(cred_provider=_Provider())
|
||||
server = MCPServer(
|
||||
server_id="s",
|
||||
name="n",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk", user_id="alice")
|
||||
assert await manager.has_user_oauth_token(server, user_auth) is verdict
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_has_user_oauth_token_short_circuits_for_unmigrated_server(self):
|
||||
"""A server the resolver does not own (None spec, e.g. delegate) is False without a call."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
calls: list = []
|
||||
|
||||
class _Provider:
|
||||
async def has_user_token(self, subject, spec):
|
||||
calls.append(spec)
|
||||
return True
|
||||
|
||||
manager = MCPServerManager(cred_provider=_Provider())
|
||||
server = MCPServer(
|
||||
server_id="s",
|
||||
name="n",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
delegate_auth_to_upstream=True,
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk", user_id="alice")
|
||||
assert await manager.has_user_oauth_token(server, user_auth) is False
|
||||
assert calls == [] # short-circuited on the None spec, never hit the resolver
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_no_user_id(self):
|
||||
"""Skip lookup entirely when user_api_key_auth has no user_id."""
|
||||
|
|
@ -2356,7 +2547,6 @@ class TestMCPServerManager:
|
|||
|
||||
# Unprefixed resolution
|
||||
resolved_server_unpref = manager._get_mcp_server_from_tool_name("create_zap")
|
||||
print(resolved_server_unpref)
|
||||
assert resolved_server_unpref is not None
|
||||
assert resolved_server_unpref.server_id == server.server_id
|
||||
|
||||
|
|
@ -2970,14 +3160,17 @@ class TestMCPServerManager:
|
|||
object_permission_id="perm_no_mcp",
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
manager, "get_allow_all_keys_server_ids", return_value=["global-server"]
|
||||
), patch.object(
|
||||
MCPRequestHandler,
|
||||
"get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=["leaked-server"],
|
||||
) as mock_inner:
|
||||
with (
|
||||
patch.object(
|
||||
manager, "get_allow_all_keys_server_ids", return_value=["global-server"]
|
||||
),
|
||||
patch.object(
|
||||
MCPRequestHandler,
|
||||
"get_allowed_mcp_servers",
|
||||
new_callable=AsyncMock,
|
||||
return_value=["leaked-server"],
|
||||
) as mock_inner,
|
||||
):
|
||||
result = await manager.get_allowed_mcp_servers(user_api_key_auth)
|
||||
|
||||
assert result == []
|
||||
|
|
@ -4641,15 +4834,23 @@ class TestCreateMcpClientV2Graft:
|
|||
client._resolved_auth._header_value.get_secret_value() == f"Basic {encoded}"
|
||||
)
|
||||
|
||||
async def test_deferred_mode_uses_v1_auth_value(self):
|
||||
async def test_m2m_client_credentials_defers_to_v1(self):
|
||||
# M2M (oauth2 client_credentials) is not migrated: to_server_spec returns
|
||||
# None, so the graft sets no resolved auth and leaves v1 in charge (v1
|
||||
# performs the client_credentials grant itself - the static
|
||||
# authentication_token is never consumed for oauth2, so it does not flow
|
||||
# to _mcp_auth_value). Per-user oauth2 (authorization_code) is migrated to
|
||||
# v2 and is exercised separately.
|
||||
client = await MCPServerManager()._create_mcp_client(
|
||||
self._http_server(
|
||||
auth_type=MCPAuth.oauth2, authentication_token="legacy-token"
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
authentication_token="legacy-token",
|
||||
)
|
||||
)
|
||||
|
||||
assert client._resolved_auth is None
|
||||
assert client._mcp_auth_value == "legacy-token"
|
||||
assert client._mcp_auth_value is None
|
||||
|
||||
async def test_static_token_missing_defers_to_v1(self):
|
||||
client = await MCPServerManager()._create_mcp_client(
|
||||
|
|
|
|||
|
|
@ -628,6 +628,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
|
|||
oauth_server = MagicMock()
|
||||
oauth_server.auth_type = MCPAuth.oauth2
|
||||
oauth_server.needs_user_oauth_token = True
|
||||
oauth_server.delegate_auth_to_upstream = False
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
@ -648,10 +649,10 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
|
|||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
) as mock_get_stored_token,
|
||||
return_value=False,
|
||||
) as mock_has_token,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
|
||||
return_value=oauth_server,
|
||||
|
|
@ -666,7 +667,7 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
|
|||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
# Verify a 401 was raised
|
||||
assert mock_get_stored_token.await_count == 1
|
||||
assert mock_has_token.await_count == 1
|
||||
assert mock_handle_request.await_count == 0
|
||||
assert exc_info.value.status_code == 401
|
||||
assert "www-authenticate" in exc_info.value.headers
|
||||
|
|
@ -817,6 +818,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
|
|||
oauth_server = MagicMock()
|
||||
oauth_server.auth_type = MCPAuth.oauth2
|
||||
oauth_server.needs_user_oauth_token = True
|
||||
oauth_server.delegate_auth_to_upstream = False
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
@ -837,10 +839,10 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
|
|||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"Authorization": "Bearer cached-token"},
|
||||
) as mock_get_stored_token,
|
||||
return_value=True,
|
||||
) as mock_has_token,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
|
||||
return_value=oauth_server,
|
||||
|
|
@ -858,7 +860,7 @@ async def test_per_user_oauth_with_stored_token_skips_preemptive_401():
|
|||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert mock_get_stored_token.await_count == 1
|
||||
assert mock_has_token.await_count == 1
|
||||
assert mock_handle_request.await_count == 1
|
||||
|
||||
|
||||
|
|
@ -941,10 +943,9 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns
|
|||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
) as mock_get_stored_token,
|
||||
) as mock_has_token,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
|
||||
return_value=delegated_server,
|
||||
|
|
@ -958,7 +959,9 @@ async def test_handle_streamable_http_mcp_delegated_server_without_token_returns
|
|||
with pytest.raises(HTTPException) as exc_info:
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert mock_get_stored_token.await_count == 1
|
||||
# Delegate-auth servers raise the resource_metadata challenge before any
|
||||
# per-user existence check, so the v2 token store is never consulted.
|
||||
assert mock_has_token.await_count == 0
|
||||
assert mock_handle_request.await_count == 0
|
||||
assert exc_info.value.status_code == 401
|
||||
challenge = exc_info.value.headers["www-authenticate"]
|
||||
|
|
|
|||
|
|
@ -270,6 +270,155 @@ class TestExecuteWithMcpClient:
|
|||
or "Authorization" not in captured["extra_headers"]
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interactive_oauth_resolves_forwarded_token_via_presented_store(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""Interactive authorization_code preview (oauth2, no client credentials): the forwarded
|
||||
just-authorized token is resolved THROUGH the v2 resolver via a one-shot presented store
|
||||
(cred_provider), not the caller-override path. The bare token (Bearer stripped) is the
|
||||
upstream credential and is not also forwarded in extra_headers."""
|
||||
captured: dict = {}
|
||||
|
||||
def fake_build_stdio_env(server, raw_headers):
|
||||
return None
|
||||
|
||||
async def fake_create_client(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_build_stdio_env",
|
||||
fake_build_stdio_env,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_create_mcp_client",
|
||||
fake_create_client,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
async def ok_operation(client):
|
||||
return {"status": "ok"}
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="linear",
|
||||
url="https://mcp.linear.app/mcp",
|
||||
auth_type=MCPAuth.oauth2,
|
||||
authorization_url="https://mcp.linear.app/authorize",
|
||||
)
|
||||
|
||||
result = await rest_endpoints._execute_with_mcp_client(
|
||||
payload,
|
||||
ok_operation,
|
||||
oauth2_headers={"Authorization": "Bearer forwarded-user-token"},
|
||||
)
|
||||
|
||||
assert result["status"] == "ok"
|
||||
# Resolved via the v2 resolver, never the caller-override header
|
||||
assert captured["mcp_auth_header"] is None
|
||||
provider = captured["cred_provider"]
|
||||
assert provider is not None
|
||||
token = await provider._oauth_token_store.fetch("u", "s")
|
||||
assert token is not None and token.access_token == "forwarded-user-token"
|
||||
# The resolver supplies the bearer, so it is not also forwarded as a caller header
|
||||
extra_headers = captured.get("extra_headers") or {}
|
||||
assert not any(k.lower() == "authorization" for k in extra_headers)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_does_not_build_presented_store(self, monkeypatch):
|
||||
"""M2M (client_credentials): to_server_spec returns None, so no presented provider is built;
|
||||
the auto-fetch path is unchanged (no cred_provider, the incoming header dropped as before)."""
|
||||
captured: dict = {}
|
||||
|
||||
def fake_build_stdio_env(server, raw_headers):
|
||||
return None
|
||||
|
||||
async def fake_create_client(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_build_stdio_env",
|
||||
fake_build_stdio_env,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_create_mcp_client",
|
||||
fake_create_client,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
async def ok_operation(client):
|
||||
return {"status": "ok"}
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="m2m-server",
|
||||
url="https://example.com",
|
||||
auth_type=MCPAuth.oauth2,
|
||||
token_url="https://auth.example.com/token",
|
||||
credentials={"client_id": "my-id", "client_secret": "my-secret"},
|
||||
)
|
||||
|
||||
result = await rest_endpoints._execute_with_mcp_client(
|
||||
payload,
|
||||
ok_operation,
|
||||
oauth2_headers={"Authorization": "Bearer sk-litellm-api-key"},
|
||||
)
|
||||
|
||||
assert result["status"] == "ok"
|
||||
assert captured.get("cred_provider") is None
|
||||
assert captured["mcp_auth_header"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_exchange_does_not_build_presented_store(self, monkeypatch):
|
||||
"""OBO / token-exchange (auth_type oauth2_token_exchange, not oauth2): excluded by the
|
||||
auth_type == oauth2 guard, so no presented provider is built and the v1 exchange path runs."""
|
||||
captured: dict = {}
|
||||
|
||||
def fake_build_stdio_env(server, raw_headers):
|
||||
return None
|
||||
|
||||
async def fake_create_client(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_build_stdio_env",
|
||||
fake_build_stdio_env,
|
||||
raising=False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
rest_endpoints.global_mcp_server_manager,
|
||||
"_create_mcp_client",
|
||||
fake_create_client,
|
||||
raising=False,
|
||||
)
|
||||
|
||||
async def ok_operation(client):
|
||||
return {"status": "ok"}
|
||||
|
||||
payload = NewMCPServerRequest(
|
||||
server_name="obo-server",
|
||||
url="https://example.com",
|
||||
auth_type=MCPAuth.oauth2_token_exchange,
|
||||
token_url="https://auth.example.com/token",
|
||||
)
|
||||
|
||||
result = await rest_endpoints._execute_with_mcp_client(
|
||||
payload,
|
||||
ok_operation,
|
||||
oauth2_headers={"Authorization": "Bearer subject-jwt"},
|
||||
)
|
||||
|
||||
assert result["status"] == "ok"
|
||||
assert captured.get("cred_provider") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_catches_exception_group(self, monkeypatch):
|
||||
"""MCP SDK's anyio TaskGroup raises BaseExceptionGroup which does not
|
||||
|
|
@ -1660,9 +1809,9 @@ class TestPreviewOpenAPITools:
|
|||
names = [t["name"] for t in result["tools"]]
|
||||
anthropic_re = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")
|
||||
for name in names:
|
||||
assert anthropic_re.match(
|
||||
name
|
||||
), f"preview tool name {name!r} violates ^[a-zA-Z0-9_-]+$"
|
||||
assert anthropic_re.match(name), (
|
||||
f"preview tool name {name!r} violates ^[a-zA-Z0-9_-]+$"
|
||||
)
|
||||
assert "actions_download-job-logs-for-workflow-run" in names
|
||||
assert "pulls_list-files" in names
|
||||
|
||||
|
|
@ -1719,9 +1868,7 @@ class TestPreviewOpenAPITools:
|
|||
|
||||
registered_summary_to_name: dict = {}
|
||||
|
||||
def fake_create_tool_function(
|
||||
path, method, operation, base_url
|
||||
): # noqa: ANN001
|
||||
def fake_create_tool_function(path, method, operation, base_url): # noqa: ANN001
|
||||
def _f():
|
||||
return None
|
||||
|
||||
|
|
@ -1734,9 +1881,7 @@ class TestPreviewOpenAPITools:
|
|||
)
|
||||
|
||||
class _StubRegistry:
|
||||
def register_tool(
|
||||
self, name, description, input_schema, handler
|
||||
): # noqa: ANN001
|
||||
def register_tool(self, name, description, input_schema, handler): # noqa: ANN001
|
||||
registered_summary_to_name[description] = name
|
||||
|
||||
monkeypatch.setattr(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue