diff --git a/litellm/constants.py b/litellm/constants.py index 78aba30f9c0..5a70623e10b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -115,6 +115,10 @@ MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACH MCP_NPM_CACHE_DIR: Final = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache") MCP_OAUTH2_TOKEN_CACHE_MIN_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10")) +# Broadcast-only (never stored) key prefix telling every worker to drop its process-local minted +# client_credentials tokens for a server. +MCP_OAUTH2_MINT_INVALIDATION_KEY_PREFIX: Final = "mcp:oauth2_mint_invalidation" + # Per-user OAuth token Redis cache (for server-side token storage) MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX: Final = "mcp:per_user_token" MCP_PER_USER_TOKEN_DEFAULT_TTL: Final = int( diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 28638ed9c77..a8e0f299a70 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -3,12 +3,14 @@ import binascii import hashlib import json from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, TypeVar, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request from litellm.proxy._types import ( @@ -40,7 +42,7 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.types.llms.custom_http import httpxSpecialProvider -from litellm.types.mcp import MCPCredentials +from litellm.types.mcp import MCPAuth, MCPCredentials if TYPE_CHECKING: from prisma import models as prisma_db_models @@ -133,6 +135,10 @@ _CLIENT_FORWARDED_AUTH_TYPES: Final["frozenset[str]"] = frozenset({"true_passthr # Minted token material that must never survive a client rotation on a persisted row. _MINTED_TOKEN_CREDENTIAL_FIELDS: Final["frozenset[str]"] = frozenset({"access_token", "refresh_token", "expires_in"}) +# For an M2M server the grant IS the standing authorization: the gateway mints on demand, so a +# disconnect that kept it would just force a remint on the next call instead of ending access. +_CLIENT_CREDENTIALS_GRANT_FIELDS: Final["frozenset[str]"] = frozenset({"client_id", "client_secret"}) + class _OAuthCredentialAccessToken(TypedDict): access_token: str @@ -334,7 +340,6 @@ def _prepare_mcp_server_data( Returns: Dict with properly serialized JSON fields """ - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # Convert model to dict. # - Partial update (exclude_unset): only caller-provided keys are emitted, so @@ -986,7 +991,6 @@ async def update_mcp_server( """ Update a new mcp server record in the db """ - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # Use helper to prepare data with proper JSON serialization. # exclude_unset=True makes this a true partial update: fields the caller did @@ -1102,6 +1106,92 @@ async def update_mcp_server( return updated_mcp_server +@dataclass(frozen=True, slots=True) +class ClearedMCPServerOAuthToken: + """Outcome of clearing a server's stored OAuth authorization: the row as persisted, whether + minted token material was present to remove, and whether an M2M client_credentials grant was + dropped with it (so the caller can report a no-op, and tell the admin what to re-enter).""" + + server: LiteLLM_MCPServerTable + had_token: bool + cleared_client_credentials: bool + + @property + def changed_row(self) -> bool: + return self.had_token or self.cleared_client_credentials + + +def _credential_str(credentials: Mapping[str, object], key: str) -> str | None: + value: Final = credentials.get(key) + return value if isinstance(value, str) else None + + +def _oauth_fields_to_clear( + row: "prisma_db_models.LiteLLM_MCPServerTable", credentials: Mapping[str, object] +) -> "frozenset[str]": + """Which credential keys a disconnect must remove. An interactive (authorization_code) server + keeps its declared client, because its authorization lives in the per-user rows the caller purges + separately. A client_credentials server loses the grant too, since keeping it would let the very + next request mint a replacement token and continue upstream access. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # avoids circular import + MCPServerManager, + ) + + flow: Final = MCPServerManager.resolve_oauth2_flow_from_fields( + auth_type=MCPAuth.oauth2 if row.auth_type == MCPAuth.oauth2 else None, + oauth2_flow=row.oauth2_flow, + token_url=row.token_url or _credential_str(credentials, "token_url"), + authorization_url=row.authorization_url, + client_id=_credential_str(credentials, "client_id"), + client_secret=_credential_str(credentials, "client_secret"), + ) + if flow != "client_credentials": + return _MINTED_TOKEN_CREDENTIAL_FIELDS + return _MINTED_TOKEN_CREDENTIAL_FIELDS | _CLIENT_CREDENTIALS_GRANT_FIELDS + + +async def clear_mcp_server_oauth_token( + prisma_client: PrismaClient, server_id: str, touched_by: str +) -> ClearedMCPServerOAuthToken | None: + """Drop a server's stored OAuth authorization: the gateway-minted token material + (``access_token``/``refresh_token``/``expires_in``) always, plus the ``client_id``/``client_secret`` + grant when the row is a client_credentials (M2M) server, whose grant would otherwise mint again on + the next request. Everything else the admin declared (url, auth_type, issuer, authorization/token + urls, token-exchange columns) stays, so reauthorizing an interactive server needs no reconfiguration. + Returns ``None`` when the server row does not exist. + """ + + existing: Final = await _db_find_mcp_server_row(prisma_client, server_id) + if existing is None: + return None + + existing_creds: Final = _credentials_blob_to_mutable_dict(existing.credentials) if existing.credentials else {} + fields_to_clear: Final = _oauth_fields_to_clear(existing, existing_creds) + cleared_fields: Final = fields_to_clear & existing_creds.keys() + if not cleared_fields: + _decrypt_env_vars_on_returned_row(existing) + return ClearedMCPServerOAuthToken(server=existing, had_token=False, cleared_client_credentials=False) + + remaining: Final = { # mutable-ok: safe_dumps serializes a plain dict and str()s a mapping view + key: value for key, value in existing_creds.items() if key not in fields_to_clear + } + updated: Final = await _db_update_mcp_server_row( + prisma_client, + server_id, + { # mutable-ok: prisma update-inputs must be plain dicts + "credentials": safe_dumps(remaining), + "updated_by": touched_by, + }, + ) + _decrypt_env_vars_on_returned_row(updated) + return ClearedMCPServerOAuthToken( + server=updated, + had_token=bool(_MINTED_TOKEN_CREDENTIAL_FIELDS & cleared_fields), + cleared_client_credentials=bool(_CLIENT_CREDENTIALS_GRANT_FIELDS & cleared_fields), + ) + + async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None: """Read the persisted (encrypted) DCR OAuth client blob for a server from the server-scoped store, or None. Config.yaml-declared servers have no @@ -1124,7 +1214,6 @@ async def upsert_mcp_server_oauth_client_credentials( client_id/client_secret are encrypted at rest with the same salt key used for the server row's credentials blob, so ``_apply_persisted_dcr_credentials`` decrypts them the same way regardless of which store a server's client came from.""" - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps encrypted: Final = encrypt_credentials(credentials=MCPCredentials(**credentials), encryption_key=_get_salt_key()) blob: Final = safe_dumps(encrypted) @@ -1146,7 +1235,6 @@ def _reencrypt_mcp_credentials_blob( uniformly and cannot silently skip one.""" if not credentials: return None - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import creds_dict: Final = _credentials_blob_to_mutable_dict(credentials) decrypted: Final = decrypt_credentials(credentials=cast(MCPCredentials, creds_dict)) @@ -1155,7 +1243,6 @@ def _reencrypt_mcp_credentials_blob( async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, touched_by: str, new_master_key: str): - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps # noqa: PLC0415 # avoids circular import mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client) @@ -1589,12 +1676,9 @@ async def purge_user_oauth_credentials_for_server( invalidate_token_cache is injectable for tests; it defaults to the manager's shared invalidate_user_oauth_token_cache, the single invalidation point for per-user tokens.""" rows: Final = await _db_find_user_credential_rows(prisma_client, {"server_id": server_id}) - oauth_rows: Final = [row for row in rows if _decode_oauth_payload(row.credential_b64) is not None] + oauth_rows: Final = tuple(row for row in rows if _decode_oauth_payload(row.credential_b64) is not None) if not oauth_rows: return 0 - deleted_count: Final = await _user_credential_actions(prisma_client).delete_many( - where={"server_id": server_id, "user_id": {"in": [row.user_id for row in oauth_rows]}} - ) if invalidate_token_cache is None: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -1602,8 +1686,20 @@ async def purge_user_oauth_credentials_for_server( invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache - for row in oauth_rows: - await invalidate_token_cache(row.user_id, server_id) + # Resolved above and dropped in a finally so a delete that succeeds is never followed by a skipped + # invalidation: the rows would be gone while their cached tokens stayed servable to the end of + # their TTL, and a retry would find nothing left to enumerate and so could never drop them. An + # invalidation for a row that survived the delete only costs that user's next resolve a DB read. + try: + deleted_count: Final = await _user_credential_actions(prisma_client).delete_many( + where={ # mutable-ok: prisma where-inputs must be plain dicts + "server_id": server_id, + "user_id": {"in": [row.user_id for row in oauth_rows]}, # mutable-ok: prisma `in` takes a plain list + } + ) + finally: + for row in oauth_rows: + await invalidate_token_cache(row.user_id, server_id) if deleted_count != len(oauth_rows): verbose_proxy_logger.warning( "MCP server %s: purge removed %d OAuth credential row(s) but %d were enumerated; " diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7ab26db0f3e..fcb00aaa71c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -13,7 +13,7 @@ import json import os import re import time -from collections.abc import AsyncIterator, Callable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypedDict, cast @@ -70,8 +70,12 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ( ) from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPPerUserTokenCache, + mcp_oauth2_mint_invalidation_key, + mcp_oauth2_token_cache, mcp_per_user_token_cache, + per_user_token_cache_key, resolve_mcp_auth, + server_id_from_mint_invalidation_key, ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, @@ -143,6 +147,9 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + publish_auth_cache_invalidation, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path @@ -1399,7 +1406,7 @@ class MCPServerManager: return None @staticmethod - def _resolve_oauth2_flow( + def resolve_oauth2_flow_from_fields( *, auth_type: MCPAuthType | None, oauth2_flow: str | None, @@ -1410,8 +1417,9 @@ class MCPServerManager: ) -> Literal["client_credentials", "authorization_code"] | None: """Infer oauth2_flow from field shape when the value is omitted. - SECURITY-SENSITIVE: this is the shape-inference engine both request-time security - helpers delegate to, so it is what decides M2M-vs-interactive for an unstamped row. + SECURITY-SENSITIVE: this is the shape-inference engine the request-time security + helpers and the disconnect write path (which must know whether a row's stored client + can mint again) delegate to, so it is what decides M2M-vs-interactive for an unstamped row. Always access it through ``effective_oauth2_flow`` (boolean/enum decisions) or ``resolve_oauth2_flow_for_request`` (the egress object backstop), which are the single choke points for request-time resolution; do not call it directly from security sites @@ -1445,7 +1453,7 @@ class MCPServerManager: resolution) goes through this one helper rather than reading the bare ``has_client_credentials`` column, which is unreliable for null rows. """ - return MCPServerManager._resolve_oauth2_flow( + return MCPServerManager.resolve_oauth2_flow_from_fields( auth_type=server.auth_type, oauth2_flow=server.oauth2_flow, token_url=server.token_url, @@ -1496,11 +1504,13 @@ class MCPServerManager: cred_provider: UpstreamCredentialProvider | None = None, per_user_oauth_token_store: InvalidatableOAuthTokenStore | None = None, per_user_token_cache: MCPPerUserTokenCache | None = None, + publish_cache_invalidation: Callable[[str], Awaitable[None]] = publish_auth_cache_invalidation, ): self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore( self.get_mcp_server_by_id ) self._per_user_token_cache = per_user_token_cache or mcp_per_user_token_cache + self._publish_cache_invalidation = publish_cache_invalidation self._cred_provider = cred_provider or UpstreamCredentialProvider( oauth_token_store=self._per_user_oauth_token_store, token_exchanger=build_token_exchanger(), @@ -5541,9 +5551,10 @@ class MCPServerManager: (re-auth, revoke, config-change purge): the v2 chain's cache and the legacy per-user token cache, so the next resolve reads the new row instead of serving the replaced token until its cache TTL, whichever path resolves it. This is the single invalidation point for per-user - OAuth tokens; callers must not evict individual caches directly. Best-effort: a cache-drop - failure is logged, never raised, because the DB write already succeeded and the TTL remains - the backstop. + OAuth tokens; callers must not evict individual caches directly. Both caches are local-first + DualCaches, so the shared key is also broadcast on the auth cache invalidation channel to + evict the in-memory copy other workers hold. Best-effort: a cache-drop failure is logged, + never raised, because the DB write already succeeded and the TTL remains the backstop. """ try: await self._per_user_oauth_token_store.invalidate(user_id, server_id) @@ -5557,6 +5568,46 @@ class MCPServerManager: verbose_logger.warning( "Failed to drop legacy cached MCP OAuth token for user=%s server=%s: %s", user_id, server_id, exc ) + await self._publish_cache_invalidation(per_user_token_cache_key(user_id, server_id)) + + async def invalidate_local_upstream_m2m_tokens(self, server_id: str) -> None: + """Drop this worker's cached gateway-minted (``client_credentials``) tokens for a server. + + Both mint caches are process-local, so this is per worker; ``revoke_upstream_m2m_tokens`` + is what makes a revoke reach the others. + """ + mcp_oauth2_token_cache.invalidate(server_id) + await self._cred_provider.invalidate_server_m2m_credentials(server_id) + + async def refresh_local_server_from_db(self, server_id: str) -> None: + """Re-read one server's persisted config into this worker's registry. + + Dropping the cached token is not enough on a worker that did not serve the revoke: its + registry entry still holds the pre-revoke ``client_id``/``client_secret``, and mints a + replacement token on the next request. Best-effort: the DB write has already committed. + """ + from litellm.proxy._experimental.mcp_server.db import get_mcp_server # noqa: PLC0415 # avoids circular import + + try: + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # avoids circular import + + if prisma_client is None: + return + server: Final = await get_mcp_server(prisma_client, server_id) + if server is not None: + await self.update_server(server) + except Exception as exc: # noqa: BLE001 - registry refresh is best-effort; the periodic reload is the backstop + verbose_logger.warning("Failed to refresh MCP server %s from the database: %s", server_id, exc) + + async def revoke_upstream_m2m_tokens(self, server_id: str) -> None: + """Revoke a server's gateway-minted tokens on every worker, not just this one. + + A worker that already minted keeps a usable upstream token in memory, so invalidating only + the worker that served the admin request left the revoked authorization in use until its + TTL. Best-effort broadcast: the DB write has already committed. + """ + await self.invalidate_local_upstream_m2m_tokens(server_id) + await self._publish_cache_invalidation(mcp_oauth2_mint_invalidation_key(server_id)) async def _resolve_oauth2_headers_for_tool_call( self, @@ -6535,3 +6586,17 @@ class MCPServerManager: global_mcp_server_manager: Final[MCPServerManager] = MCPServerManager() + + +async def apply_mcp_upstream_m2m_invalidation(cache_key: str) -> None: + """Handle a broadcast mint-invalidation key on the worker that receives it. + + The auth cache invalidation subscriber evicts the shared ``user_api_key_cache`` only, which + cannot reach the process-local mint caches holding a usable upstream token, so it routes + broadcast keys here. Any other key is ignored. + """ + server_id: Final = server_id_from_mint_invalidation_key(cache_key) + if server_id is None: + return + await global_mcp_server_manager.invalidate_local_upstream_m2m_tokens(server_id) + await global_mcp_server_manager.refresh_local_server_from_db(server_id) diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index c76c933c5b5..96032ece483 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -14,6 +14,7 @@ import httpx from litellm._logging import verbose_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( + MCP_OAUTH2_MINT_INVALIDATION_KEY_PREFIX, MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, @@ -199,6 +200,24 @@ class MCPOAuth2TokenCache(InMemoryCache): mcp_oauth2_token_cache: Final = MCPOAuth2TokenCache() +def mcp_oauth2_mint_invalidation_key(server_id: str) -> str: + return f"{MCP_OAUTH2_MINT_INVALIDATION_KEY_PREFIX}:{server_id}" + + +def server_id_from_mint_invalidation_key(cache_key: str) -> str | None: + """The server a broadcast invalidation key names, or ``None`` for any other key.""" + prefix: Final = f"{MCP_OAUTH2_MINT_INVALIDATION_KEY_PREFIX}:" + if not cache_key.startswith(prefix): + return None + return cache_key.removeprefix(prefix) or None + + +def per_user_token_cache_key(user_id: str, server_id: str) -> str: + """The ``user_api_key_cache`` key holding a user's token for a server. Shared by the legacy + per-user cache and the v2 store backend, so broadcasting this one key evicts both.""" + return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + + def _compute_per_user_token_ttl(server: "MCPServer", expires_in: int | None) -> int: """Compute Redis TTL for a per-user token. @@ -229,7 +248,7 @@ class MCPPerUserTokenCache: """ def _cache_key(self, user_id: str, server_id: str) -> str: - return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + return per_user_token_cache_key(user_id, server_id) async def get(self, user_id: str, server_id: str) -> str | None: """Return the plaintext access_token, or None on miss/error.""" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py index d0053fbe0a8..93bc22f507f 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py @@ -179,6 +179,7 @@ class ClientCredentialsTokenSource: self._max_locks = max_locks self._clock = clock self._locks: dict[str, asyncio.Lock] = {} + self._identities: dict[str, frozenset[str]] = {} # mutable-ok: process-lifetime index, like _locks def _lock(self, server_id: str) -> asyncio.Lock: """Per-server single-flight lock, bounded so ephemeral server ids (e.g. the REST tools @@ -190,6 +191,20 @@ class ClientCredentialsTokenSource: self._locks.pop(next(iter(self._locks)), None) return self._locks.setdefault(server_id, asyncio.Lock()) + def _remember_identity(self, server_id: str, identity_key: str) -> None: + """Record the identities a server minted under, bounded like the locks. ``invalidate`` runs + after the stored client credentials were cleared, so the identity key is no longer + derivable from config and the cached token would be unreachable without this. + """ + if server_id not in self._identities and len(self._identities) >= self._max_locks: + self._identities.pop(next(iter(self._identities)), None) + self._identities[server_id] = self._identities.get(server_id, frozenset()) | {identity_key} + + async def invalidate(self, server_id: str) -> None: + """Drop every cached token for a server, whichever client identity minted it.""" + for identity_key in self._identities.pop(server_id, frozenset()): + await self._backend.delete(identity_key, server_id) + async def get(self, server_id: str, config: ClientCredentialsConfig) -> Result[OAuthToken, CredError]: match _prepare_grant(config): case Error(err): @@ -262,6 +277,7 @@ class ClientCredentialsTokenSource: ) if ttl > 0: await self._backend.set(grant.identity_key, server_id, token, ttl) + self._remember_identity(server_id, grant.identity_key) return Ok(token) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py index 94c59962b70..84cb78cb741 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py @@ -356,6 +356,14 @@ class UpstreamCredentialProvider: subject.inbound_token.get_secret_value(), server, server.config, tenant_id=subject.tenant_id ) + async def invalidate_server_m2m_credentials(self, server_id: str) -> None: + """Drop the gateway's cached `client_credentials` tokens for a server, for every caller. + + The M2M token belongs to the gateway rather than a subject, so a revoke has to clear the + whole server, unlike the per-subject `invalidate_credentials`. + """ + await self._client_credentials_source.invalidate(server_id) + def _invalidate_id_jag(self, subject: Subject, server: ServerSpec) -> None: """Evict the bearer this `(subject, server)` last resolved, without depending on the store. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 628e569e1b8..e597b54d5de 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1579,6 +1579,19 @@ class MCPOAuthUserCredentialStatus(LiteLLMPydanticObjectBase): connected_at: str | None = None # ISO-8601 +class MCPServerOAuthTokenStatus(LiteLLMPydanticObjectBase): + """Outcome of the admin-wide OAuth token clear for one MCP server. + + Distinct from MCPOAuthUserCredentialStatus, which is scoped to the calling user alone. + """ + + server_id: str + has_token: bool + cleared: bool = False + cleared_user_tokens: int = 0 + cleared_client_credentials: bool = False + + class MCPUserCredentialListItem(LiteLLMPydanticObjectBase): """One entry in the /user-credentials list.""" diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index fb2ca6372c0..9ca087de9ee 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import Sequence +from collections.abc import Awaitable, Callable, Sequence from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Final @@ -122,17 +122,25 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache") + __slots__ = ( + "_additional_in_memory_caches", + "_local_invalidators", + "_redis_cache", + "_task", + "_user_api_key_cache", + ) def __init__( self, redis_cache: "RedisCache", user_api_key_cache: "UserApiKeyCache", additional_in_memory_caches: Sequence["InMemoryCache"] = (), + local_invalidators: Sequence[Callable[[str], Awaitable[None]]] = (), ) -> None: self._redis_cache = redis_cache self._user_api_key_cache = user_api_key_cache self._additional_in_memory_caches = tuple(additional_in_memory_caches) + self._local_invalidators = tuple(local_invalidators) self._task: asyncio.Task[None] | None = None def start(self) -> None: @@ -185,9 +193,9 @@ class AuthCacheInvalidationSubscriber: message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) if message is None: continue - self._apply_message(message) + await self._apply_message(message) - def _apply_message(self, message: object) -> None: + async def _apply_message(self, message: object) -> None: data: Final = message.get("data") if isinstance(message, dict) else None parsed: Final = _message_from_data(data) if parsed is None: @@ -201,6 +209,13 @@ class AuthCacheInvalidationSubscriber: in_memory_cache.delete_cache(parsed.cache_key) for additional_cache in self._additional_in_memory_caches: additional_cache.delete_cache(parsed.cache_key) + for invalidate_local in self._local_invalidators: + try: + await invalidate_local(parsed.cache_key) + except Exception as e: # noqa: BLE001 # one bad invalidator must not stop the subscriber + verbose_proxy_logger.warning( + "auth cache invalidation local handler failed for %s: %s", parsed.cache_key, e + ) @staticmethod async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 54a591a5e1a..7267d0b2274 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -22,7 +22,7 @@ import os from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final, Literal, Protocol +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol from fastapi import ( APIRouter, @@ -128,6 +128,7 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.db import ( approve_mcp_server, + clear_mcp_server_oauth_token, create_draft_mcp_server, create_mcp_server, delete_mcp_server, @@ -175,6 +176,7 @@ if MCP_AVAILABLE: MCPApprovalStatus, MCPOAuthUserCredentialRequest, MCPOAuthUserCredentialStatus, + MCPServerOAuthTokenStatus, MCPSubmissionsSummary, MCPTransport, MCPUserCredentialListItem, @@ -2032,6 +2034,88 @@ if MCP_AVAILABLE: return Response(status_code=status.HTTP_202_ACCEPTED) + @router.delete( + "/server/{server_id}/oauth-token", + description=( + "Clear every OAuth token stored for an MCP server (admin only), so it can be reauthorized " + "with a different set of upstream grants. An interactive (authorization_code) server keeps " + "its declared OAuth app (url, auth_type, issuer, client_id/secret) so users just reconnect; " + "a client_credentials server also loses the client_id/secret grant, which is the only way " + "to stop it minting a replacement token, so an admin must re-enter the client to resume." + ), + response_model=MCPServerOAuthTokenStatus, + ) + @management_endpoint_wrapper + async def clear_mcp_server_oauth_token_endpoint( + server_id: str, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + ): + """ + Clear every stored OAuth token for a server so the next call reauthorizes from scratch. + + A token minted by 'Authorize & Fetch Token' can be served from three places, and a clear that + misses any one of them leaves the old upstream grants live: + - the per-user credential rows the interactive (authorization_code) flow writes, plus the two + caches that hydrate from them, + - the in-memory client_credentials mint cache an M2M server serves from, + - the minted fields on the server row's own credential blob. + This clears all three. A user's own BYOK API key is never touched. + + A client_credentials server mints on demand from its stored client, so the grant itself is the + standing authorization and gets dropped with the token; the rest of the declared app stays, and + an admin re-enters the client to authorize again. The mint cache is process-local, so the + server-level invalidation is broadcast on the auth cache invalidation channel and every worker + drops its copy. Without redis pub/sub there is nothing to broadcast over, and a worker that + already minted keeps its token until the TTL, but the DB no longer backs a fresh mint. + + ``` + curl -X "DELETE" --location 'http://localhost:4000/v1/mcp/server/server_id/oauth-token' \ + --header 'Authorization: Bearer your_api_key_here' + ``` + """ + prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy") + + if LitellmUserRoles.PROXY_ADMIN != user_api_key_dict.user_role: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ # mutable-ok: HTTPException details are plain dicts across every management route + "error": ( + "Call not allowed to clear the MCP server OAuth token. User is not a proxy admin. " + "route=DELETE /v1/mcp/server/{server_id}/oauth-token" + ) + }, + ) + + cleared: Final = await clear_mcp_server_oauth_token( + prisma_client, + server_id, + touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + ) + if cleared is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"MCP Server not found, passed server_id={server_id}"}, # mutable-ok: plain dict + ) + + # The in-memory stores are reconciled even when the purge raises. The server row has already + # been written by then, so skipping them on the error path would leave the registry and the + # mint cache serving material the DB no longer backs, and the admin's retry would read as a + # no-op while the old grants stayed live. Both are idempotent, so running them twice is free. + try: + cleared_user_tokens: Final = await purge_user_oauth_credentials_for_server(prisma_client, server_id) + finally: + await global_mcp_server_manager.revoke_upstream_m2m_tokens(server_id) + if cleared.changed_row: + await global_mcp_server_manager.update_server(cleared.server) + + return MCPServerOAuthTokenStatus( + server_id=server_id, + has_token=False, + cleared=cleared.changed_row or cleared_user_tokens > 0, + cleared_user_tokens=cleared_user_tokens, + cleared_client_credentials=cleared.cleared_client_credentials, + ) + @router.post( "/server/{server_id}/user-credential", description="Store or update the calling user's API key for a BYOK MCP server", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0abcdeaf3f6..029e4d8bbf7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16,7 +16,7 @@ import threading import time import traceback import warnings -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping, MutableMapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -4198,6 +4198,20 @@ def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool: return any(str(obj) == object_type_str for obj in supported_db_objects) +def _local_cache_invalidators() -> tuple[Callable[[str], Awaitable[None]], ...]: + """Handlers for broadcast invalidation keys that name a process-local cache the shared + ``user_api_key_cache`` eviction cannot reach, such as the MCP client_credentials mint caches.""" + from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + + if not is_mcp_available(): + return () + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + apply_mcp_upstream_m2m_invalidation, + ) + + return (apply_mcp_upstream_m2m_invalidation,) + + class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. @@ -6805,6 +6819,7 @@ class ProxyConfig: redis_cache=redis_cache, user_api_key_cache=user_api_key_cache, additional_in_memory_caches=(spend_counter_cache.in_memory_cache,), + local_invalidators=_local_cache_invalidators(), ) self.auth_cache_invalidation_subscriber = subscriber subscriber.start() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py index a5d17428b37..456c8e7567f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_client_credentials.py @@ -425,3 +425,50 @@ def test_bearer_auth_rejects_sync_clients(): with httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200)), auth=auth) as client: with pytest.raises(RuntimeError): client.get("https://upstream.example.com/mcp") + + +@pytest.mark.asyncio +async def test_invalidate_stops_serving_the_cached_token_for_a_server(): + """An admin revoke has to reach this source, not just the row: otherwise the worker that already + minted keeps serving the revoked authorization from cache until its TTL.""" + poster = _FakePoster([_success(access_token="tok-before", expires_in=3600), _success(access_token="tok-after", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + first = await source.get("s", _config()) + + await source.invalidate("s") + second = await source.get("s", _config()) + + assert isinstance(first, Ok) and isinstance(second, Ok) + assert first.ok.access_token == "tok-before" + assert second.ok.access_token == "tok-after" + assert len(poster.calls) == 2 + + +@pytest.mark.asyncio +async def test_invalidate_drops_tokens_minted_under_credentials_that_are_already_gone(): + """The revoke clears the stored client credentials first, so the identity key is no longer + derivable from config by the time invalidation runs.""" + poster = _FakePoster([_success(access_token="tok-old", expires_in=3600), _success(access_token="tok-new", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + _ = await source.get("s", _config(client_secret=SecretStr("old-secret"))) + + await source.invalidate("s") + reused = await source.get("s", _config(client_secret=SecretStr("old-secret"))) + + assert isinstance(reused, Ok) + assert reused.ok.access_token == "tok-new" + + +@pytest.mark.asyncio +async def test_invalidate_leaves_other_servers_alone(): + poster = _FakePoster([_success(access_token="tok-a", expires_in=3600), _success(access_token="tok-b", expires_in=3600)]) + source = ClientCredentialsTokenSource(poster) + _ = await source.get("srv-a", _config()) + _ = await source.get("srv-b", _config()) + + await source.invalidate("srv-a") + kept = await source.get("srv-b", _config()) + + assert isinstance(kept, Ok) + assert kept.ok.access_token == "tok-b" + assert len(poster.calls) == 2 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 50248e95ffa..b7ef05db6ed 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -12,7 +12,7 @@ import base64 import json from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -200,6 +200,28 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_each_user(): assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")} +@pytest.mark.asyncio +async def test_purge_user_oauth_credentials_for_server_invalidates_even_when_the_delete_raises(): + """A delete that raises after touching rows must not skip the invalidation: the cached tokens would + stay servable to the end of their TTL, and a retry could not fix it because the enumeration it + depends on no longer finds those rows. Invalidating a row that survived only costs a DB read.""" + from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server + + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[_oauth_row("alice"), _oauth_row("bob")]) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(side_effect=RuntimeError("connection reset")) + + invalidations = [] + + async def record_invalidation(user_id: str, server_id: str) -> None: + invalidations.append((user_id, server_id)) + + with pytest.raises(RuntimeError, match="connection reset"): + await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation) + + assert set(invalidations) == {("alice", "srv-1"), ("bob", "srv-1")} + + @pytest.mark.asyncio async def test_purge_user_oauth_credentials_for_server_spares_byok_rows(): """Regression: the purge used to delete_many on server_id alone, wiping BYOK API keys that share @@ -1337,3 +1359,243 @@ def test_mcp_oauth_token_identity_changes_when_only_upstream_resource_is_edited( assert mcp_oauth_token_identity(set_to_explicit) == mcp_oauth_token_identity( _identity_server(credentials={**creds, "upstream_resource": "api://audience-one"}) ) + + +def _server_row_with_credentials(credentials, **overrides): + """A LiteLLM_MCPServerTable-shaped row as prisma hands it back, for the server-level + minted-token clear tests.""" + base = { + "server_id": "srv-clear", + "url": "https://upstream.example.com/mcp", + "auth_type": MCPAuth.oauth2, + "issuer": "https://issuer.example.com", + "transport": MCPTransport.http, + "credentials": credentials, + "env_vars": None, + "oauth2_flow": None, + "token_url": None, + "authorization_url": None, + } + base.update(overrides) + return SimpleNamespace(**base) + + +def _prisma_with_server_row(row): + prisma = MagicMock() + prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=row) + prisma.db.litellm_mcpservertable.update = AsyncMock( + side_effect=lambda where, data: _server_row_with_credentials( + data.get("credentials"), server_id=where["server_id"] + ) + ) + return prisma + + +def _written_credentials(prisma) -> dict: + return json.loads(prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"]["credentials"]) + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_drops_only_minted_fields(): + """The server-level clear must remove exactly the minted token material and leave every + declared-app credential key behind, so an admin can reauthorize without reconfiguring the + OAuth app. A clear that dropped the whole blob would still pass a test that only checked the + access token is gone, so the surviving keys are asserted explicitly.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + stored = { + "access_token": "at-live", + "refresh_token": "rt-live", + "expires_in": 3600, + "client_id": "declared-client", + "client_secret": "declared-secret", + "scopes": ["read", "write"], + "token_endpoint_auth_method": "client_secret_post", + "upstream_resource": "api://audience-one", + } + prisma = _prisma_with_server_row(_server_row_with_credentials(json.dumps(stored))) + + result = await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + + assert result is not None + assert result.had_token is True + + written = _written_credentials(prisma) + assert "access_token" not in written + assert "refresh_token" not in written + assert "expires_in" not in written + assert written == { + "client_id": "declared-client", + "client_secret": "declared-secret", + "scopes": ["read", "write"], + "token_endpoint_auth_method": "client_secret_post", + "upstream_resource": "api://audience-one", + } + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_leaves_declared_app_columns_untouched(): + """Clearing the token must not write url, auth_type, issuer or any other declared column. + The pre-existing clear paths (client rotation, auth_type/url/issuer change) are destructive + side effects; this one has to be surgical or it re-creates the workaround it replaces.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + prisma = _prisma_with_server_row( + _server_row_with_credentials(json.dumps({"access_token": "at-live", "client_id": "declared-client"})) + ) + + await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + + written_data = prisma.db.litellm_mcpservertable.update.call_args.kwargs["data"] + assert set(written_data) == {"credentials", "updated_by"} + assert written_data["updated_by"] == "admin-1" + for untouched in ("url", "auth_type", "issuer", "client_id", "transport", "env_vars"): + assert untouched not in written_data + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_is_a_noop_without_a_minted_token(): + """A server that only carries a declared app must not be written at all, so a repeated + Disconnect does not churn updated_by or bump the row.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + prisma = _prisma_with_server_row( + _server_row_with_credentials(json.dumps({"client_id": "declared-client", "client_secret": "declared-secret"})) + ) + + result = await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + + assert result is not None + assert result.had_token is False + prisma.db.litellm_mcpservertable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_handles_dict_shaped_credentials(): + """Prisma hands JSONB back as a dict on some read paths; the clear must work there too.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + prisma = _prisma_with_server_row( + _server_row_with_credentials({"access_token": "at-live", "client_id": "declared-client"}) + ) + + result = await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + + assert result is not None and result.had_token is True + assert _written_credentials(prisma) == {"client_id": "declared-client"} + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_returns_none_for_missing_server(): + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + prisma = _prisma_with_server_row(None) + + assert await clear_mcp_server_oauth_token(prisma, "srv-missing", touched_by="admin-1") is None + prisma.db.litellm_mcpservertable.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_drops_the_client_credentials_grant_for_m2m(): + """An M2M server mints on demand from its stored client, so a clear that kept the grant left the + next request free to mint a replacement token and keep calling upstream.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + stored = { + "access_token": "at-live", + "client_id": "m2m-client", + "client_secret": "m2m-secret", + "scopes": ["read"], + } + prisma = _prisma_with_server_row( + _server_row_with_credentials(json.dumps(stored), oauth2_flow="client_credentials") + ) + + result = await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + + assert result is not None + assert result.had_token is True + assert result.cleared_client_credentials is True + assert _written_credentials(prisma) == {"scopes": ["read"]} + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_drops_the_grant_for_an_unstamped_m2m_shaped_row(): + """Rows written before oauth2_flow was stamped carry the M2M shape instead, and the same + inference the request path uses has to classify them here or they keep minting.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + prisma = _prisma_with_server_row( + _server_row_with_credentials( + json.dumps({"client_id": "m2m-client", "client_secret": "m2m-secret"}), + token_url="https://auth.example.com/token", + ) + ) + + result = await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + + assert result is not None + assert result.had_token is False + assert result.cleared_client_credentials is True + assert result.changed_row is True + assert _written_credentials(prisma) == {} + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_keeps_the_client_for_an_interactive_server(): + """An authorization_code server's authorization lives in the per-user rows the caller purges, so + its OAuth app must survive for users to reconnect without an admin reconfiguring it.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + + prisma = _prisma_with_server_row( + _server_row_with_credentials( + json.dumps({"access_token": "at-live", "client_id": "declared-client", "client_secret": "declared-secret"}), + oauth2_flow="authorization_code", + authorization_url="https://auth.example.com/authorize", + ) + ) + + result = await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + + assert result is not None + assert result.cleared_client_credentials is False + assert _written_credentials(prisma) == {"client_id": "declared-client", "client_secret": "declared-secret"} + + +@pytest.mark.asyncio +async def test_cleared_m2m_row_can_no_longer_mint_a_replacement_token(): + """End state check: the row the clear persists must fail the mint path's precondition, so no + request can quietly re-authorize the server.""" + from litellm.proxy._experimental.mcp_server.db import clear_mcp_server_oauth_token + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import MCPOAuth2TokenCache + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + prisma = _prisma_with_server_row( + _server_row_with_credentials( + json.dumps({"access_token": "at-live", "client_id": "m2m-client", "client_secret": "m2m-secret"}), + oauth2_flow="client_credentials", + ) + ) + + await clear_mcp_server_oauth_token(prisma, "srv-clear", touched_by="admin-1") + remaining = _written_credentials(prisma) + server = MCPServer( + server_id="srv-clear", + name="srv-clear", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + token_url="https://auth.example.com/token", + client_id=remaining.get("client_id"), + client_secret=remaining.get("client_secret"), + ) + token_endpoint = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client", + return_value=token_endpoint, + ): + assert await MCPOAuth2TokenCache().async_get_token(server) is None + + token_endpoint.post.assert_not_awaited() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index cdea803ebf3..7dc788637d6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -5170,6 +5170,138 @@ class TestMCPServerManager: manager = MCPServerManager(per_user_oauth_token_store=_Store(), per_user_token_cache=_RaisingLegacyCache()) await manager.invalidate_user_oauth_token_cache("alice", "srv-1") + @pytest.mark.asyncio + async def test_invalidate_user_oauth_token_cache_broadcasts_the_shared_key(self): + """Both per-user caches are local-first, so evicting them here leaves every other worker + holding the revoked token until its TTL unless the shared key is broadcast.""" + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + per_user_token_cache_key, + ) + + class _Store: + async def fetch(self, user_id: str, server_id: str): + return None + + async def invalidate(self, user_id: str, server_id: str) -> None: + return None + + class _LegacyCache: + async def delete(self, user_id: str, server_id: str) -> None: + return None + + published: list[str] = [] + + async def _publish(cache_key: str) -> None: + published.append(cache_key) + + manager = MCPServerManager( + per_user_oauth_token_store=_Store(), + per_user_token_cache=_LegacyCache(), + publish_cache_invalidation=_publish, + ) + await manager.invalidate_user_oauth_token_cache("alice", "srv-1") + + assert published == [per_user_token_cache_key("alice", "srv-1")] + + @pytest.mark.asyncio + async def test_revoke_upstream_m2m_tokens_clears_both_mint_caches_and_broadcasts(self): + """A revoke has to reach the v2 token source that actually serves M2M traffic, the legacy + mint cache, and every other worker, or the gateway keeps minting or replaying upstream + access the admin just revoked.""" + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + mcp_oauth2_mint_invalidation_key, + mcp_oauth2_token_cache, + ) + + invalidated: list[str] = [] + + class _FakeProvider: + async def invalidate_server_m2m_credentials(self, server_id: str) -> None: + invalidated.append(server_id) + + published: list[str] = [] + + async def _publish(cache_key: str) -> None: + published.append(cache_key) + + mcp_oauth2_token_cache.set_cache("srv-1:identity", "tok-revoked") + mcp_oauth2_token_cache.set_cache("srv-2:identity", "tok-untouched") + + manager = MCPServerManager(cred_provider=_FakeProvider(), publish_cache_invalidation=_publish) + await manager.revoke_upstream_m2m_tokens("srv-1") + + assert invalidated == ["srv-1"] + assert published == [mcp_oauth2_mint_invalidation_key("srv-1")] + assert mcp_oauth2_token_cache.get_cache("srv-1:identity") is None + assert mcp_oauth2_token_cache.get_cache("srv-2:identity") == "tok-untouched" + + @pytest.mark.asyncio + async def test_broadcast_mint_invalidation_only_fires_for_its_own_keys(self): + """The channel carries every management-object eviction, so a receiving worker must not read + a per-user key as a server id and wipe an unrelated server's minted tokens.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + apply_mcp_upstream_m2m_invalidation, + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + mcp_oauth2_mint_invalidation_key, + ) + + invalidated: list[str] = [] + + with ( + patch.object( + global_mcp_server_manager, + "invalidate_local_upstream_m2m_tokens", + new=AsyncMock(side_effect=lambda server_id: invalidated.append(server_id)), + ), + patch.object(global_mcp_server_manager, "refresh_local_server_from_db", new=AsyncMock()), + ): + await apply_mcp_upstream_m2m_invalidation("mcp:per_user_token:alice:srv-1") + assert invalidated == [] + await apply_mcp_upstream_m2m_invalidation(mcp_oauth2_mint_invalidation_key("srv-1")) + + assert invalidated == ["srv-1"] + + @pytest.mark.asyncio + async def test_broadcast_mint_invalidation_rereads_the_revoked_server_from_the_database(self): + """Dropping the cached token is not enough: a worker that did not serve the revoke still holds + the pre-revoke client_id/client_secret in its registry and would mint a replacement token.""" + from litellm.proxy._experimental.mcp_server import db as mcp_db + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + apply_mcp_upstream_m2m_invalidation, + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( + mcp_oauth2_mint_invalidation_key, + ) + + revoked_row = LiteLLM_MCPServerTable( + server_id="srv-1", + server_name="srv-1", + url="http://upstream/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + created_at=None, + updated_at=None, + ) + applied: list[LiteLLM_MCPServerTable] = [] + + with ( + patch.object(global_mcp_server_manager, "invalidate_local_upstream_m2m_tokens", new=AsyncMock()), + patch.object( + global_mcp_server_manager, + "update_server", + new=AsyncMock(side_effect=lambda server: applied.append(server)), + ), + patch.object(mcp_db, "get_mcp_server", new=AsyncMock(return_value=revoked_row)), + patch("litellm.proxy.proxy_server.prisma_client", new=MagicMock()), + ): + await apply_mcp_upstream_m2m_invalidation(mcp_oauth2_mint_invalidation_key("srv-1")) + + assert applied == [revoked_row] + @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.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py index 72589fd8b3e..01758d44aae 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_oauth2_token_cache.py @@ -12,7 +12,10 @@ import pytest from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( MCPOAuth2TokenCache, + mcp_oauth2_mint_invalidation_key, + mcp_oauth2_token_cache, resolve_mcp_auth, + server_id_from_mint_invalidation_key, ) from litellm.proxy._types import MCPTransport from litellm.types.mcp import MCPAuth @@ -392,3 +395,14 @@ async def test_invalidate_clears_every_identity_for_a_server(): assert refetched == "tok-after-invalidate" assert mock_client.post.call_count == 3 + + +def test_mint_invalidation_key_round_trips_to_its_server_id(): + assert server_id_from_mint_invalidation_key(mcp_oauth2_mint_invalidation_key("srv-1")) == "srv-1" + + +def test_unrelated_cache_keys_do_not_name_a_server(): + """The channel carries every management-object eviction, so a user key must not be read as a + server id and wipe an unrelated server's minted tokens.""" + assert server_id_from_mint_invalidation_key("mcp:per_user_token:user-1:srv-1") is None + assert server_id_from_mint_invalidation_key("mcp:oauth2_mint_invalidation:") is None diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 7d5fc1a3544..8152a341f4e 100644 --- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -185,9 +185,55 @@ async def test_subscriber_ignores_malformed_messages() -> None: redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), user_api_key_cache=cache, ) - subscriber._apply_message({"type": "message", "data": b"not json"}) - subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()}) - subscriber._apply_message("raw string") - subscriber._apply_message(None) + await subscriber._apply_message({"type": "message", "data": b"not json"}) + await subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()}) + await subscriber._apply_message("raw string") + await subscriber._apply_message(None) assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + + +@pytest.mark.asyncio +async def test_subscriber_routes_broadcast_keys_to_injected_local_invalidators() -> None: + """Some caches a broadcast has to reach are process-local rather than entries in + user_api_key_cache (the MCP client_credentials mint cache), so the receiving worker hands every + key to the injected handlers as well.""" + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + seen: List[str] = [] + + async def _record(cache_key: str) -> None: + seen.append(cache_key) + + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), + user_api_key_cache=cache, + local_invalidators=(_record,), + ) + await subscriber._apply_message(_invalidation_message("project_id:p-1")) + + assert seen == ["project_id:p-1"] + assert cache.in_memory_cache.get_cache("project_id:p-1") is None + + +@pytest.mark.asyncio +async def test_subscriber_survives_a_failing_local_invalidator() -> None: + """One handler raising must not stop the others or kill the subscriber loop, which would leave + the worker serving stale authorization for every later broadcast.""" + cache = UserApiKeyCache() + seen: List[str] = [] + + async def _boom(cache_key: str) -> None: + raise RuntimeError("handler exploded") + + async def _record(cache_key: str) -> None: + seen.append(cache_key) + + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), + user_api_key_cache=cache, + local_invalidators=(_boom, _record), + ) + await subscriber._apply_message(_invalidation_message("project_id:p-1")) + + assert seen == ["project_id:p-1"] diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 0d639e1cb6a..87a7a04cfe1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1,8 +1,9 @@ +import base64 import os import sys import types import json -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from datetime import datetime, timedelta from types import SimpleNamespace from typing import List, Optional @@ -6766,3 +6767,347 @@ class TestConnectedAppViewAnnotation: assert all(server.connected_app_reachable is None for server in result) reload_mock.assert_not_awaited() + + +def _make_admin_auth(user_id: str = "admin-abc") -> "UserAPIKeyAuth": + return UserAPIKeyAuth( + api_key="sk-test", + user_id=user_id, + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + +@contextmanager +def _patched_clear_oauth_token_endpoint(clear_mock, purge_mock, manager, revoke_mock): + """Patch every collaborator the clear endpoint reaches, so a test can assert which stores it + actually emptied rather than only what it returned.""" + manager.revoke_upstream_m2m_tokens = revoke_mock + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token", + new=clear_mock, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.purge_user_oauth_credentials_for_server", + new=purge_mock, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + manager, + ), + ): + yield + + +def _clear_oauth_token_manager(): + manager = MagicMock() + manager.update_server = AsyncMock() + return manager + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_empties_every_store_a_token_is_served_from(): + """A token authorized for a server can be served from the per-user credential rows, from the + in-memory client_credentials mint cache, or from the server row's own minted fields. The clear + has to empty all three: leaving any one of them behind keeps the old upstream grants live, which + is exactly the state the admin asked to get out of.""" + from litellm.proxy._types import MCPServerOAuthTokenStatus + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server.db import ClearedMCPServerOAuthToken + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + server = generate_mock_mcp_server_db_record(server_id="srv-oauth") + clear_mock = AsyncMock(return_value=ClearedMCPServerOAuthToken(server=server, had_token=True, cleared_client_credentials=False)) + purge_mock = AsyncMock(return_value=3) + manager = _clear_oauth_token_manager() + revoke_mock = AsyncMock() + + with _patched_clear_oauth_token_endpoint(clear_mock, purge_mock, manager, revoke_mock): + result = await clear_mcp_server_oauth_token_endpoint( + server_id="srv-oauth", + user_api_key_dict=_make_admin_auth("admin-1"), + ) + + assert isinstance(result, MCPServerOAuthTokenStatus) + assert result.server_id == "srv-oauth" + assert result.has_token is False + assert result.cleared is True + assert result.cleared_user_tokens == 3 + assert clear_mock.await_args.kwargs["touched_by"] == "admin-1" + purge_mock.assert_awaited_once() + assert purge_mock.await_args.args[1] == "srv-oauth" + revoke_mock.assert_awaited_once_with("srv-oauth") + manager.update_server.assert_awaited_once_with(server) + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_clears_per_user_tokens_with_no_minted_server_row(): + """The interactive flow stores its token per user and writes nothing to the server row, so a clear + that keyed 'did anything happen' off the server row alone would report a no-op while revoking + three users. cleared must follow the per-user purge too.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server.db import ClearedMCPServerOAuthToken + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + server = generate_mock_mcp_server_db_record(server_id="srv-oauth") + purge_mock = AsyncMock(return_value=3) + manager = _clear_oauth_token_manager() + revoke_mock = AsyncMock() + + with _patched_clear_oauth_token_endpoint( + AsyncMock(return_value=ClearedMCPServerOAuthToken(server=server, had_token=False, cleared_client_credentials=False)), + purge_mock, + manager, + revoke_mock, + ): + result = await clear_mcp_server_oauth_token_endpoint( + server_id="srv-oauth", + user_api_key_dict=_make_admin_auth(), + ) + + assert result.cleared is True + assert result.cleared_user_tokens == 3 + purge_mock.assert_awaited_once() + manager.update_server.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_drives_a_real_purge_for_the_authorize_and_fetch_state(): + """End to end through the endpoint with the real purge, in the state 'Authorize & Fetch Token' + actually leaves behind: nothing minted on the server row, one stored per-user token. + + The sibling tests either mock the purge (so they pin the wiring, not the delete) or call the purge + helper directly (so they never prove the endpoint reaches it). Between those two a gate reinstated + on the endpoint would pass both while no-opping on the ticket's own scenario, so this one drives + the endpoint against a prisma double and asserts the delete really issued.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server.db import ClearedMCPServerOAuthToken + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + stored_token = base64.urlsafe_b64encode( + json.dumps({"type": "oauth2", "access_token": "upstream-token-a"}).encode() + ).decode() + prisma = _make_prisma_client() + prisma.db.litellm_mcpusercredentials.find_many = AsyncMock( + return_value=[SimpleNamespace(user_id="alice", server_id="srv-oauth", credential_b64=stored_token)] + ) + prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=1) + + server = generate_mock_mcp_server_db_record(server_id="srv-oauth") + manager = _clear_oauth_token_manager() + manager.invalidate_user_oauth_token_cache = AsyncMock() + manager.revoke_upstream_m2m_tokens = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=prisma, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token", + new=AsyncMock(return_value=ClearedMCPServerOAuthToken(server=server, had_token=False, cleared_client_credentials=False)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + manager, + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + manager, + ), + ): + result = await clear_mcp_server_oauth_token_endpoint( + server_id="srv-oauth", + user_api_key_dict=_make_admin_auth(), + ) + + assert result.cleared_user_tokens == 1 + assert result.cleared is True + prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once_with( + where={"server_id": "srv-oauth", "user_id": {"in": ["alice"]}} + ) + manager.invalidate_user_oauth_token_cache.assert_awaited_once_with("alice", "srv-oauth") + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_reconciles_caches_when_the_purge_fails(): + """The server row is written before the per-user purge runs, so a purge that raises must not skip + the in-memory reconciliation: the registry and the mint cache would go on serving material the DB + no longer backs, and the admin's retry would look like a no-op while the old grants stayed live. + The error still surfaces, because the clear did not finish.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server.db import ClearedMCPServerOAuthToken + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + server = generate_mock_mcp_server_db_record(server_id="srv-oauth") + manager = _clear_oauth_token_manager() + revoke_mock = AsyncMock() + + with _patched_clear_oauth_token_endpoint( + AsyncMock(return_value=ClearedMCPServerOAuthToken(server=server, had_token=True, cleared_client_credentials=False)), + AsyncMock(side_effect=RuntimeError("database went away mid-purge")), + manager, + revoke_mock, + ): + with pytest.raises(RuntimeError, match="database went away mid-purge"): + await clear_mcp_server_oauth_token_endpoint( + server_id="srv-oauth", + user_api_key_dict=_make_admin_auth(), + ) + + revoke_mock.assert_awaited_once_with("srv-oauth") + manager.update_server.assert_awaited_once_with(server) + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_reports_a_true_noop(): + """Nothing stored anywhere means nothing to clear: the response says so and the registry is left + alone. The mint cache is still flushed, because it is the one store the endpoint cannot count.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server.db import ClearedMCPServerOAuthToken + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + server = generate_mock_mcp_server_db_record(server_id="srv-oauth") + manager = _clear_oauth_token_manager() + revoke_mock = AsyncMock() + + with _patched_clear_oauth_token_endpoint( + AsyncMock(return_value=ClearedMCPServerOAuthToken(server=server, had_token=False, cleared_client_credentials=False)), + AsyncMock(return_value=0), + manager, + revoke_mock, + ): + result = await clear_mcp_server_oauth_token_endpoint( + server_id="srv-oauth", + user_api_key_dict=_make_admin_auth(), + ) + + assert result.cleared is False + assert result.has_token is False + assert result.cleared_user_tokens == 0 + revoke_mock.assert_awaited_once_with("srv-oauth") + manager.update_server.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_reports_a_dropped_m2m_grant(): + """An M2M server has no per-user rows and may have nothing minted on the row yet, so dropping its + client grant is the whole disconnect: the registry has to be reloaded and the caller told what to + re-enter, or the dashboard reports a no-op on the one change that revoked access.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server.db import ClearedMCPServerOAuthToken + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + server = generate_mock_mcp_server_db_record(server_id="srv-oauth") + manager = _clear_oauth_token_manager() + revoke_mock = AsyncMock() + + with _patched_clear_oauth_token_endpoint( + AsyncMock( + return_value=ClearedMCPServerOAuthToken( + server=server, had_token=False, cleared_client_credentials=True + ) + ), + AsyncMock(return_value=0), + manager, + revoke_mock, + ): + result = await clear_mcp_server_oauth_token_endpoint( + server_id="srv-oauth", + user_api_key_dict=_make_admin_auth(), + ) + + assert result.cleared is True + assert result.cleared_client_credentials is True + revoke_mock.assert_awaited_once_with("srv-oauth") + manager.update_server.assert_awaited_once_with(server) + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_rejects_non_admin(): + """Clearing the shared token logs every user of the server out of the upstream, so it is + admin-only. A non-admin who wants to drop their own connection uses the per-user revoke.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + clear_mock = AsyncMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token", + new=clear_mock, + ), + pytest.raises(HTTPException) as exc_info, + ): + await clear_mcp_server_oauth_token_endpoint( + server_id="srv-oauth", + user_api_key_dict=_make_user_auth("bob"), + ) + + assert exc_info.value.status_code == 403 + clear_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_clear_mcp_server_oauth_token_endpoint_404s_for_unknown_server(): + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + clear_mcp_server_oauth_token_endpoint, + ) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.clear_mcp_server_oauth_token", + new=AsyncMock(return_value=None), + ), + pytest.raises(HTTPException) as exc_info, + ): + await clear_mcp_server_oauth_token_endpoint( + server_id="srv-missing", + user_api_key_dict=_make_admin_auth(), + ) + + assert exc_info.value.status_code == 404 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPDisconnectDialog.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPDisconnectDialog.test.tsx new file mode 100644 index 00000000000..cae0c7653df --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPDisconnectDialog.test.tsx @@ -0,0 +1,139 @@ +import React from "react"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import MCPDisconnectDialog from "./MCPDisconnectDialog"; +import { toast } from "@/lib/toast"; + +const deleteMCPServerOAuthToken = vi.fn(); +const deleteMCPOAuthUserCredential = vi.fn(); + +vi.mock("@/components/networking", () => ({ + deleteMCPServerOAuthToken: (...args: unknown[]) => deleteMCPServerOAuthToken(...args), + deleteMCPOAuthUserCredential: (...args: unknown[]) => deleteMCPOAuthUserCredential(...args), +})); + +vi.mock("@/lib/toast", () => ({ + toast: { success: vi.fn(), fromError: vi.fn(), error: vi.fn() }, +})); + +const SHARED_BLAST_RADIUS = /Clears every OAuth token LiteLLM holds for this server, for every user/i; +const OWN_BLAST_RADIUS = /Clears only the token stored against your user/i; + +function renderDialog(overrides: Partial> = {}) { + const onCleared = vi.fn(); + const onOpenChange = vi.fn(); + render( + , + ); + return { onCleared, onOpenChange }; +} + +describe("MCPDisconnectDialog blast radius", () => { + beforeEach(() => { + vi.clearAllMocks(); + deleteMCPServerOAuthToken.mockResolvedValue({ + server_id: "srv-1", + has_token: false, + cleared: true, + cleared_user_tokens: 2, + cleared_client_credentials: false, + }); + deleteMCPOAuthUserCredential.mockResolvedValue({ server_id: "srv-1", has_credential: false, is_expired: false }); + }); + + it("states both blast radii, distinguishing the all-users clear from the caller's own connection", () => { + renderDialog(); + + const shared = screen.getByText(SHARED_BLAST_RADIUS); + const own = screen.getByText(OWN_BLAST_RADIUS); + + expect(shared).toBeInTheDocument(); + expect(own).toBeInTheDocument(); + expect(shared).not.toHaveTextContent(OWN_BLAST_RADIUS); + expect(own).not.toHaveTextContent(SHARED_BLAST_RADIUS); + expect(screen.getByText(/Every stored token for this server \(affects all users\)/i)).toBeInTheDocument(); + expect(screen.getByText(/Only your own connection \(affects just you\)/i)).toBeInTheDocument(); + }); + + it("warns that the all-users clear revokes everyone while the per-user clear leaves everyone else connected", () => { + renderDialog(); + + expect(screen.getByText(SHARED_BLAST_RADIUS)).toHaveTextContent( + /Anyone who authorized interactively loses upstream access until they authorize again/i, + ); + expect(screen.getByText(SHARED_BLAST_RADIUS)).toHaveTextContent( + /machine-to-machine server also loses the stored client credentials .* re-enter the client secret/i, + ); + expect(screen.getByText(SHARED_BLAST_RADIUS)).toHaveTextContent(/BYOK API keys are left alone/i); + expect(screen.getByText(OWN_BLAST_RADIUS)).toHaveTextContent(/Every other user keeps their connection/i); + }); + + it("defaults an admin to the server-level clear and calls the server-level endpoint", async () => { + const { onCleared } = renderDialog(); + + await userEvent.click(screen.getByRole("button", { name: "Disconnect" })); + + await waitFor(() => expect(deleteMCPServerOAuthToken).toHaveBeenCalledWith("sk-test", "srv-1")); + expect(deleteMCPOAuthUserCredential).not.toHaveBeenCalled(); + expect(onCleared).toHaveBeenCalledWith("server"); + }); + + it("calls the per-user revoke when the caller picks their own connection", async () => { + const { onCleared } = renderDialog(); + + await userEvent.click(screen.getByText(/Only your own connection \(affects just you\)/i)); + await userEvent.click(screen.getByRole("button", { name: "Disconnect" })); + + await waitFor(() => expect(deleteMCPOAuthUserCredential).toHaveBeenCalledWith("sk-test", "srv-1")); + expect(deleteMCPServerOAuthToken).not.toHaveBeenCalled(); + expect(onCleared).toHaveBeenCalledWith("self"); + }); + + it("defaults a non-admin to their own connection and blocks the shared clear", async () => { + const { onCleared } = renderDialog({ isProxyAdmin: false }); + + expect(screen.getByText(/Requires proxy admin/i)).toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: "Disconnect" })); + + await waitFor(() => expect(onCleared).toHaveBeenCalledWith("self")); + expect(deleteMCPServerOAuthToken).not.toHaveBeenCalled(); + }); + + it("tells the admin to re-enter the client when the M2M grant was dropped", async () => { + deleteMCPServerOAuthToken.mockResolvedValue({ + server_id: "srv-1", + has_token: false, + cleared: true, + cleared_user_tokens: 0, + cleared_client_credentials: true, + }); + renderDialog(); + + await userEvent.click(screen.getByRole("button", { name: "Disconnect" })); + + await waitFor(() => + expect(toast.success).toHaveBeenCalledWith(expect.stringMatching(/Re-enter the client credentials/i)), + ); + }); + + it("labels the reauthorize entry point and still shows both blast radii", () => { + renderDialog({ mode: "reauthorize" }); + + expect(screen.getByText("Reauthorize MCP Server?")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Clear & Reauthorize" })).toBeInTheDocument(); + expect(screen.getByText(SHARED_BLAST_RADIUS)).toBeInTheDocument(); + expect(screen.getByText(OWN_BLAST_RADIUS)).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPDisconnectDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPDisconnectDialog.tsx new file mode 100644 index 00000000000..7ffcd9265ef --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPDisconnectDialog.tsx @@ -0,0 +1,150 @@ +import { type FC, useState } from "react"; +import { + AlertDialog, + AlertDialogCancel, + AlertDialogContent, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Button } from "@/components/ui/button"; +import { Label } from "@/components/ui/label"; +import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group"; +import { cn } from "@/lib/cva.config"; +import { toast } from "@/lib/toast"; +import { deleteMCPOAuthUserCredential, deleteMCPServerOAuthToken } from "@/components/networking"; + +export type MCPDisconnectScope = "server" | "self"; + +export type MCPDisconnectMode = "disconnect" | "reauthorize"; + +export const MCP_DISCONNECT_SCOPE_COPY: Record = { + server: { + label: "Every stored token for this server (affects all users)", + blastRadius: + "Clears every OAuth token LiteLLM holds for this server, for every user. Anyone who authorized interactively loses upstream access until they authorize again, and their OAuth app stays configured. A machine-to-machine server also loses the stored client credentials it mints from, so an admin has to re-enter the client secret before it can call upstream again. Stored BYOK API keys are left alone.", + }, + self: { + label: "Only your own connection (affects just you)", + blastRadius: + "Clears only the token stored against your user. Every other user keeps their connection and the server stays authorized for them.", + }, +}; + +interface MCPDisconnectDialogProps { + open: boolean; + mode: MCPDisconnectMode; + serverId: string; + serverName?: string; + accessToken: string; + isProxyAdmin: boolean; + onOpenChange: (open: boolean) => void; + onCleared: (scope: MCPDisconnectScope) => void; +} + +const MCPDisconnectDialog: FC = ({ + open, + mode, + serverId, + serverName, + accessToken, + isProxyAdmin, + onOpenChange, + onCleared, +}) => { + const [scope, setScope] = useState(isProxyAdmin ? "server" : "self"); + const [isSubmitting, setIsSubmitting] = useState(false); + + const confirmLabel = mode === "reauthorize" ? "Clear & Reauthorize" : "Disconnect"; + + const handleConfirm = async () => { + setIsSubmitting(true); + try { + if (scope === "server") { + const result = await deleteMCPServerOAuthToken(accessToken, serverId); + toast.success( + result.cleared + ? `Cleared every stored OAuth token for this server (${result.cleared_user_tokens} user token${ + result.cleared_user_tokens === 1 ? "" : "s" + })${ + result.cleared_client_credentials + ? ". Re-enter the client credentials to let this server mint tokens again" + : "" + }` + : "This server had no stored OAuth token", + ); + } else { + await deleteMCPOAuthUserCredential(accessToken, serverId); + toast.success("Cleared your connection to this MCP server"); + } + onOpenChange(false); + onCleared(scope); + } catch (error) { + toast.fromError(error); + } finally { + setIsSubmitting(false); + } + }; + + return ( + + + + + {mode === "reauthorize" ? "Reauthorize MCP Server?" : "Disconnect MCP Server?"} + + +
+

+ {mode === "reauthorize" + ? "Clearing the stored token lets you run Authorize & Fetch Token again and grant a different set of upstream scopes. The server's URL, auth type and issuer are kept, so nothing has to be recreated." + : "Choose what to clear. The server's URL, auth type and issuer are kept, so the server can be authorized again without being recreated."} +

+
+ {serverName && ( +
+
Name
+
{serverName}
+
+ )} +
+
ID
+
{serverId}
+
+
+ setScope(value as MCPDisconnectScope)}> + {(["server", "self"] as const).map((option) => { + const disabled = option === "server" && !isProxyAdmin; + return ( + + ); + })} + +
+ + Cancel + + + + + ); +}; + +export default MCPDisconnectDialog; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.test.tsx index 8c7fe46a1a2..505e1855a9e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.test.tsx @@ -1,5 +1,6 @@ import React from "react"; import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, afterEach } from "vitest"; import MCPServerCard from "./MCPServerCard"; import type { MCPServer } from "@/components/mcp_tools/types"; @@ -67,3 +68,60 @@ describe("MCPServerCard logo", () => { expect(screen.getByText("DE")).toBeInTheDocument(); }); }); + +describe("MCPServerCard OAuth token actions", () => { + async function openMenu(overrides: Partial, handlers: Record void>) { + render( + , + ); + await userEvent.click(screen.getByLabelText("Server actions")); + } + + it("offers Disconnect and Reauthorize for an oauth2 server", async () => { + await openMenu({ auth_type: "oauth2" }, { onDisconnectOAuth: vi.fn(), onReauthorizeOAuth: vi.fn() }); + + expect(await screen.findByText("Disconnect / Clear token")).toBeInTheDocument(); + expect(screen.getByText("Reauthorize")).toBeInTheDocument(); + expect(screen.getByText("Test Connection")).toBeInTheDocument(); + expect(screen.getByText("Delete")).toBeInTheDocument(); + }); + + it("invokes the disconnect handler without triggering the card's own click", async () => { + const onDisconnectOAuth = vi.fn(); + const onClick = vi.fn(); + render( + , + ); + await userEvent.click(screen.getByLabelText("Server actions")); + await userEvent.click(await screen.findByText("Disconnect / Clear token")); + + expect(onDisconnectOAuth).toHaveBeenCalledTimes(1); + expect(onClick).not.toHaveBeenCalled(); + }); + + it("hides both token actions for a non-oauth2 server", async () => { + await openMenu({ auth_type: "api_key" }, { onDisconnectOAuth: vi.fn(), onReauthorizeOAuth: vi.fn() }); + + expect(await screen.findByText("Test Connection")).toBeInTheDocument(); + expect(screen.queryByText("Disconnect / Clear token")).not.toBeInTheDocument(); + expect(screen.queryByText("Reauthorize")).not.toBeInTheDocument(); + }); + + it("hides both token actions when the caller passes no handlers (non-admin)", async () => { + await openMenu({ auth_type: "oauth2" }, {}); + + expect(await screen.findByText("Test Connection")).toBeInTheDocument(); + expect(screen.queryByText("Disconnect / Clear token")).not.toBeInTheDocument(); + expect(screen.queryByText("Reauthorize")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx index cf3c2863e83..d18078d06d5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPServerCard.tsx @@ -1,5 +1,5 @@ import { type FC, type KeyboardEvent, type MouseEvent } from "react"; -import { Check, CircleAlert, Ellipsis, Trash2, Zap } from "lucide-react"; +import { Check, CircleAlert, Ellipsis, Trash2, Unplug, RefreshCw, Zap } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { @@ -27,6 +27,8 @@ interface MCPServerCardProps { onRecheckHealth?: () => void; onByokConnect?: () => void; onOpenFillFields?: () => void; + onDisconnectOAuth?: () => void; + onReauthorizeOAuth?: () => void; onDelete?: () => void; } @@ -48,6 +50,8 @@ const MCPServerCard: FC = ({ onRecheckHealth, onByokConnect, onOpenFillFields, + onDisconnectOAuth, + onReauthorizeOAuth, onDelete, }) => { const alias = server.alias || server.server_name || ""; @@ -106,7 +110,9 @@ const MCPServerCard: FC = ({ } }; - const hasMenu = !!onRecheckHealth || !!onDelete; + const isOAuthServer = server.auth_type === AUTH_TYPE.OAUTH2; + const showOAuthTokenActions = isOAuthServer && (!!onDisconnectOAuth || !!onReauthorizeOAuth); + const hasMenu = !!onRecheckHealth || !!onDelete || showOAuthTokenActions; // Card uses role="button" + nested + + + + )}

Cost Configuration

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 8dca87e24ef..8aab4af2637 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -17,6 +17,24 @@ vi.mock("@/components/networking", () => ({ updateConfigFieldSetting: vi.fn().mockResolvedValue(undefined), deleteConfigFieldSetting: vi.fn().mockResolvedValue(undefined), listMCPUserEnvVarStatus: vi.fn().mockResolvedValue([]), + deleteMCPServerOAuthToken: vi + .fn() + .mockResolvedValue({ server_id: "server-1", has_token: false, cleared: true, cleared_user_tokens: 1 }), + deleteMCPOAuthUserCredential: vi + .fn() + .mockResolvedValue({ server_id: "server-1", has_credential: false, is_expired: false }), +})); + +// The detail view is stubbed so this file can assert which tab the list page hands it without +// mounting the edit form and tools viewer underneath. MCP_TAB stays real: the point of the +// assertion is the shared index, so a stubbed constant would let the two sides drift apart. +const mcpServerViewProps = vi.fn(); +vi.mock("./mcp_server_view", async (importOriginal) => ({ + MCP_TAB: ((await importOriginal()) as { MCP_TAB: unknown }).MCP_TAB, + MCPServerView: (props: Record) => { + mcpServerViewProps(props); + return
; + }, })); const createQueryClient = () => @@ -321,6 +339,48 @@ describe("MCPServers", () => { expect(screen.queryByText("Team B Server")).not.toBeInTheDocument(); }); + it("opens Settings after an all-users Reauthorize from a list card, where Authorize & Fetch Token lives", async () => { + // Regression: the list page hardcoded "Tools tab or Overview", so an all-users Reauthorize + // dropped the admin on Overview with no authorize button in sight, while the same action on the + // detail page correctly opened Settings. Reauthorize that lands nowhere useful is the whole bug. + const { MCP_TAB } = await import("./mcp_server_view"); + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { + server_id: "server-1", + server_name: "OAuth Server", + alias: "oauth-server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "oauth2", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + teams: [], + mcp_access_groups: [], + }, + ]); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + + render( + + + , + ); + + expect(await screen.findByText("OAuth Server")).toBeInTheDocument(); + + await userEvent.click(screen.getByLabelText("Server actions")); + await userEvent.click(await screen.findByText("Reauthorize")); + await userEvent.click(await screen.findByRole("button", { name: "Clear & Reauthorize" })); + + expect(await screen.findByTestId("mcp-server-view")).toBeInTheDocument(); + expect(networking.deleteMCPServerOAuthToken).toHaveBeenCalledWith("123", "server-1"); + expect(mcpServerViewProps).toHaveBeenLastCalledWith( + expect.objectContaining({ initialTabIndex: MCP_TAB.SETTINGS, isEditing: true }), + ); + }); + it("should not trigger an extra health check when the server list changes after deletion", async () => { // Regression test: previously useMCPServerHealth received serverIds derived from the // server list. Deleting a server changed serverIds, which changed the React Query key, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index c6418a13897..dcb871e3adc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -27,7 +27,8 @@ import { MCPToolsetsTab } from "./MCPToolsetsTab"; import CreateMCPServer from "./CreateMCPServer"; import MCPConnect from "./mcp_connect"; import MCPServerCard from "./MCPServerCard"; -import { MCPServerView } from "./mcp_server_view"; +import MCPDisconnectDialog, { type MCPDisconnectMode } from "./MCPDisconnectDialog"; +import { MCP_TAB, MCPServerView } from "./mcp_server_view"; import type { DiscoverableMCPServer, MCPServer, @@ -152,6 +153,11 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const [prefillData, setPrefillData] = useState(null); const [isDeletingServer, setIsDeletingServer] = useState(false); const [byokModalServer, setByokModalServer] = useState(null); + const [disconnectServer, setDisconnectServer] = useState(null); + const [disconnectMode, setDisconnectMode] = useState("disconnect"); + // Set when an all-users Reauthorize is confirmed from a list card, so the detail view opens on + // Settings where "Authorize & Fetch Token" lives rather than on Overview. + const [settingsTabServerId, setSettingsTabServerId] = useState(null); // Per-user env-var fill modal target + deep-link source captured once from the URL. const [envVarsModalServer, setEnvVarsModalServer] = useState(null); const [deepLinkServerId, setDeepLinkServerId] = useState(() => @@ -338,6 +344,30 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) return [...matches].sort((a, b) => compareServers(a, b, sortKey)); }, [filteredServers, searchQuery, sortKey]); + const openDisconnectDialog = (server: MCPServer, mode: MCPDisconnectMode) => { + setDisconnectMode(mode); + setDisconnectServer(server); + }; + + const handleDisconnectCleared = (scope: "server" | "self") => { + const server = disconnectServer; + setDisconnectServer(null); + refetch(); + if (!server || disconnectMode !== "reauthorize") { + return; + } + setToolsTabServerId(scope === "self" ? server.server_id : null); + setSettingsTabServerId(scope === "server" ? server.server_id : null); + setEditServer(scope === "server"); + setSelectedServerId(server.server_id); + }; + + const initialTabIndexForSelectedServer = () => { + if (selectedServerId === toolsTabServerId) return MCP_TAB.TOOLS; + if (selectedServerId === settingsTabServerId) return MCP_TAB.SETTINGS; + return MCP_TAB.OVERVIEW; + }; + function handleDelete(server_id: string) { setServerToDelete(server_id); setIsDeleteModalOpen(true); @@ -407,8 +437,9 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) const handleBack = React.useCallback(() => { setEditServer(false); setSelectedServerId(null); - // Drop the post-redirect one-shot so re-selecting that server opens Overview. + // Drop the post-redirect one-shots so re-selecting that server opens Overview. setToolsTabServerId(null); + setSettingsTabServerId(null); refetch(); }, [refetch]); @@ -458,6 +489,18 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) + {disconnectServer && ( + !open && setDisconnectServer(null)} + onCleared={handleDisconnectCleared} + /> + )} = ({ accessToken, userRole, userID }) userID={userID} userRole={userRole} availableAccessGroups={uniqueMcpAccessGroups} - initialTabIndex={selectedServerId === toolsTabServerId ? 1 : 0} + initialTabIndex={initialTabIndexForSelectedServer()} /> ) : (
@@ -694,6 +737,8 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) } onByokConnect={server.is_byok ? () => setByokModalServer(server) : undefined} onOpenFillFields={() => setEnvVarsModalServer(server)} + onDisconnectOAuth={() => openDisconnectDialog(server, "disconnect")} + onReauthorizeOAuth={() => openDisconnectDialog(server, "reauthorize")} onDelete={isAdminRole(userRole) ? () => handleDelete(server.server_id) : undefined} /> ))} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 5b6d70b4771..6090a5c64bc 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -7803,6 +7803,14 @@ export interface MCPOAuthUserCredentialStatus { connected_at?: string | null; } +export interface MCPServerOAuthTokenStatus { + server_id: string; + has_token: boolean; + cleared: boolean; + cleared_user_tokens: number; + cleared_client_credentials: boolean; +} + export interface MCPUserCredentialListItem { server_id: string; server_name?: string | null; @@ -7880,6 +7888,13 @@ export const deleteMCPOAuthUserCredential = async ( return response.json(); }; +export const deleteMCPServerOAuthToken = async ( + accessToken: string, + serverId: string, +): Promise => { + return apiClient.delete(`/v1/mcp/server/${serverId}/oauth-token`, { accessToken }); +}; + export const getMCPOAuthUserCredentialStatus = async ( accessToken: string, serverId: string,