Merge pull request #41725 from BerriAI/litellm_mcp_admin_terminate_sessions_revoke_credentials

feat(mcp): let proxy admins force-close live MCP sessions and revoke stored user credentials
This commit is contained in:
Yassin Kortam 2026-09-18 13:44:37 -07:00 committed by GitHub
commit 6e750e3d24
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
31 changed files with 2305 additions and 127 deletions

View file

@ -183,6 +183,9 @@ MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIME
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE: Final = 4096
# Allowlist of commands permitted for MCP stdio transport.
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.

View file

@ -0,0 +1,38 @@
"""Per-worker cache of stored BYOK credentials, keyed so peer workers can evict it over the auth cache pub/sub."""
from dataclasses import dataclass
from typing import Final
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE, MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS
_CACHE_KEY_PREFIX: Final = "mcp_byok_credential"
@dataclass(frozen=True, slots=True)
class CachedByokCredential:
credential: str | None
byok_credential_cache: Final = InMemoryCache(
max_size_in_memory=MCP_BYOK_CREDENTIAL_CACHE_MAX_SIZE,
default_ttl=MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS,
)
def byok_credential_cache_key(user_id: str, server_id: str) -> str:
return f"{_CACHE_KEY_PREFIX}:{user_id}:{server_id}"
def get_cached_byok_credential(user_id: str, server_id: str) -> CachedByokCredential | None:
cached: Final = byok_credential_cache.get_cache( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # InMemoryCache is untyped
byok_credential_cache_key(user_id, server_id)
)
return cached if isinstance(cached, CachedByokCredential) else None
def cache_byok_credential(user_id: str, server_id: str, credential: str | None) -> None:
byok_credential_cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped
byok_credential_cache_key(user_id, server_id),
CachedByokCredential(credential=credential),
)

View file

@ -865,7 +865,7 @@ async def byok_token(
_invalidate_byok_cred_cache,
)
_invalidate_byok_cred_cache(user_id, server_id)
await _invalidate_byok_cred_cache(user_id, server_id)
except Exception as exc:
verbose_proxy_logger.error(
"byok_token: failed to store user credential for user=%s server=%s: %s",

View file

@ -24,6 +24,7 @@ from litellm.proxy._types import (
MCPApprovalStatus,
MCPEnvVar,
MCPEnvVarScope,
MCPServerUserCredentialListItem,
MCPSubmissionsSummary,
NewMCPServerRequest,
SpecialMCPServerName,
@ -1504,6 +1505,37 @@ async def get_user_oauth_credential(
return _parse_oauth_payload(decoded)
def _server_user_credential_item(
row: "prisma_db_models.LiteLLM_MCPUserCredentials",
) -> MCPServerUserCredentialListItem:
oauth_payload: Final = _decode_oauth_payload(row.credential_b64)
if oauth_payload is None:
return MCPServerUserCredentialListItem(
user_id=row.user_id,
credential_type="byok",
updated_at=row.updated_at.isoformat(),
)
return MCPServerUserCredentialListItem(
user_id=row.user_id,
credential_type="oauth2",
expires_at=oauth_payload.get("expires_at"),
connected_at=oauth_payload.get("connected_at"),
updated_at=row.updated_at.isoformat(),
)
async def list_server_user_credentials(
prisma_client: PrismaClient,
server_id: str,
) -> tuple[MCPServerUserCredentialListItem, ...]:
"""Every user's stored credential for one server, typed but without the secret, for admins."""
rows: Final = await _db_find_user_credential_rows(
prisma_client,
{"server_id": server_id}, # mutable-ok: prisma where-inputs must be plain dicts
)
return tuple(_server_user_credential_item(row) for row in rows)
async def list_user_oauth_credentials(
prisma_client: PrismaClient,
user_id: str,

View file

@ -295,12 +295,15 @@ class MCPPerUserTokenCache:
)
async def delete(self, user_id: str, server_id: str) -> None:
"""Invalidate the cached token (removes from both in-memory and Redis layers)."""
"""Invalidate the cached token in Redis, here, and in every peer worker's in-memory layer."""
try:
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( # noqa: PLC0415 # proxy import cycle
evict_and_broadcast,
)
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
key: Final = self._cache_key(user_id, server_id)
await user_api_key_cache.async_delete_cache(key)
await evict_and_broadcast((key,), user_api_key_cache)
except Exception as exc:
verbose_logger.debug(
"MCPPerUserTokenCache.delete failed for user=%s server=%s: %s",

View file

@ -28,7 +28,10 @@ from starlette.types import Message, Receive, Scope, Send
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
from litellm.constants import (
MAXIMUM_TRACEBACK_LINES_TO_LOG,
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -38,6 +41,12 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
_is_mcp_admitted_user_subject,
)
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
byok_credential_cache,
byok_credential_cache_key,
cache_byok_credential,
get_cached_byok_credential,
)
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
get_request_base_url,
)
@ -82,6 +91,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.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
get_chain_id_from_headers,
@ -91,6 +103,7 @@ from litellm.types.mcp import (
MCPGatewaySession,
MCPGatewaySessionGroupCount,
MCPGatewaySessionsResponse,
MCPGatewaySessionsTerminateResponse,
MCPSpecVersion,
)
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
@ -102,13 +115,6 @@ if TYPE_CHECKING:
from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
# Short-lived in-memory cache for BYOK credentials.
# Keyed by (user_id, server_id); value is (credential_or_None, monotonic_timestamp).
# Storing the credential value (not just a bool) means _get_byok_credential and
# _check_byok_credential share a single DB round-trip per TTL window.
_byok_cred_cache: Final[dict[tuple[str, str], tuple[str | None, float]]] = {}
_BYOK_CRED_CACHE_TTL: Final = 60 # seconds
_BYOK_CRED_CACHE_MAX_SIZE: Final = 4096 # cap to prevent unbounded growth
_STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS: Final = 30 * 60
# Upper bound on concurrent stateful sessions a single caller may hold. Each
# `initialize` creates a session that survives until the idle timeout, so
@ -127,20 +133,11 @@ _MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span"
_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations"
def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
"""Remove a (user_id, server_id) entry from the BYOK credential cache.
Call this after storing or deleting a credential so subsequent calls
see the fresh value rather than a stale cached result.
"""
_byok_cred_cache.pop((user_id, server_id), None)
def _write_byok_cred_cache(user_id: str, server_id: str, credential: str | None) -> None:
"""Write a credential value to the cache, evicting all entries if at capacity."""
if len(_byok_cred_cache) >= _BYOK_CRED_CACHE_MAX_SIZE:
_byok_cred_cache.clear()
_byok_cred_cache[(user_id, server_id)] = (credential, time.monotonic())
async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
"""Drop a stored-or-deleted BYOK credential from this worker's cache and from every peer worker's."""
cache_key: Final = byok_credential_cache_key(user_id, server_id)
byok_credential_cache.delete_cache(cache_key)
await publish_auth_cache_invalidation(cache_key=cache_key)
# Check if MCP is available
@ -618,6 +615,7 @@ if MCP_AVAILABLE:
_stateful_session_locks: Final[dict[str, asyncio.Lock]] = {}
_stateful_session_active_request_counts: Final[dict[str, int]] = {}
_stateful_session_client_info: Final[dict[str, Implementation]] = {} # mutable-ok: cleared on session teardown
_admin_terminated_session_ids: Final[dict[str, float]] = {} # mutable-ok: admin-closed id -> last replay
class _TerminableTransport(Protocol):
async def terminate(self) -> None: ...
@ -689,6 +687,7 @@ if MCP_AVAILABLE:
for session_id in list(_stateful_session_auth_context_last_seen):
if session_id not in _stateful_session_auth_contexts:
_remove_stateful_session_tracking(session_id)
_forget_expired_admin_terminated_session_ids(now)
async def _enforce_stateful_session_cap_for_owner(owner: str) -> bool:
"""
@ -2811,35 +2810,28 @@ if MCP_AVAILABLE:
mcp_server: MCPServer,
user_api_key_auth: UserAPIKeyAuth | None,
) -> str | None:
"""Retrieve the stored BYOK credential for a user+server pair.
Uses the shared _byok_cred_cache to avoid a DB round-trip on every
tool call within the TTL window.
"""
"""Retrieve the stored BYOK credential for a user+server pair, served from the worker cache within its TTL."""
if not mcp_server.is_byok:
return None
user_id: Final = (user_api_key_auth.user_id if user_api_key_auth else None) or ""
if not user_id:
return None
cache_key: Final = (user_id, mcp_server.server_id)
cached: Final = _byok_cred_cache.get(cache_key)
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
if cached is not None:
credential, ts = cached
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
return credential
return cached.credential
from litellm.proxy._experimental.mcp_server.db import get_user_credential
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return None
credential = await get_user_credential(
credential: Final = await get_user_credential(
prisma_client=prisma_client,
user_id=user_id,
server_id=mcp_server.server_id,
)
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
cache_byok_credential(user_id, mcp_server.server_id, credential)
return credential
async def _check_byok_credential(
@ -2868,27 +2860,23 @@ if MCP_AVAILABLE:
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
# Check shared credential cache before hitting the DB.
cache_key: Final = (user_id, mcp_server.server_id)
cached: Final = _byok_cred_cache.get(cache_key)
cached: Final = get_cached_byok_credential(user_id, mcp_server.server_id)
if cached is not None:
cached_cred, ts = cached
if time.monotonic() - ts < _BYOK_CRED_CACHE_TTL:
if cached_cred is None:
raise HTTPException(
status_code=401,
detail={
"error": "byok_auth_required",
"server_id": mcp_server.server_id,
"server_name": mcp_server.server_name or mcp_server.name,
"message": (
"No stored credential found for this BYOK server. "
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
return
if cached.credential is None:
raise HTTPException(
status_code=401,
detail={
"error": "byok_auth_required",
"server_id": mcp_server.server_id,
"server_name": mcp_server.server_name or mcp_server.name,
"message": (
"No stored credential found for this BYOK server. "
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={"WWW-Authenticate": get_byok_www_authenticate()},
)
return
from litellm.proxy._experimental.mcp_server.db import get_user_credential
from litellm.proxy.proxy_server import prisma_client
@ -2912,7 +2900,7 @@ if MCP_AVAILABLE:
user_id=user_id,
server_id=mcp_server.server_id,
)
_write_byok_cred_cache(user_id, mcp_server.server_id, credential)
cache_byok_credential(user_id, mcp_server.server_id, credential)
if credential is None:
raise HTTPException(
status_code=401,
@ -3850,7 +3838,7 @@ if MCP_AVAILABLE:
client_info: Final = _stateful_session_client_info.get(session_id)
key_auth: Final = auth_user.user_api_key_auth
return MCPGatewaySession(
session_id_prefix=session_id[:8],
session_id_prefix=session_id[:MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH],
client_name=client_info.name if client_info is not None else None,
client_version=client_info.version if client_info is not None else None,
user_id=key_auth.user_id if key_auth is not None else None,
@ -3885,6 +3873,72 @@ if MCP_AVAILABLE:
sessions=sessions,
)
def _session_matches_admin_selector(
session_id: str,
auth_user: MCPAuthenticatedUser,
session_id_prefix: str | None,
user_id: str | None,
) -> bool:
if session_id_prefix is not None and not session_id.startswith(session_id_prefix):
return False
if user_id is None:
return True
key_auth: Final = auth_user.user_api_key_auth
return key_auth is not None and key_auth.user_id == user_id
def _forget_expired_admin_terminated_session_ids(now: float) -> None:
for session_id in [
session_id
for session_id, last_replayed in _admin_terminated_session_ids.items()
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
]:
del _admin_terminated_session_ids[session_id]
def _is_admin_terminated_session_id(session_id: str, now: float) -> bool:
last_replayed: Final = _admin_terminated_session_ids.get(session_id)
if last_replayed is None:
return False
if now - last_replayed >= _STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS:
del _admin_terminated_session_ids[session_id]
return False
_admin_terminated_session_ids[session_id] = now
return True
async def terminate_mcp_gateway_sessions(
*,
session_id_prefix: str | None = None,
user_id: str | None = None,
) -> MCPGatewaySessionsTerminateResponse:
"""Force-close every live stateful session on this worker matching the selector.
The transport is terminated (open streams close), all per-session
tracking is dropped, and the id is remembered so a client that keeps
sending it receives 404 and has to ``initialize`` again, which re-runs
admission. Only sessions held by this worker process are affected.
"""
now: Final = time.monotonic()
_forget_expired_admin_terminated_session_ids(now)
server_instances: Final = _stateful_server_instances()
targets: Final = tuple(
(session_id, auth_user)
for session_id, auth_user in tuple(_stateful_session_auth_contexts.items())
if session_id in server_instances
and _session_matches_admin_selector(session_id, auth_user, session_id_prefix, user_id)
)
terminated: Final = tuple(_gateway_session_for(session_id, auth_user, now) for session_id, auth_user in targets)
for session_id, _ in targets:
_admin_terminated_session_ids[session_id] = now
transport = server_instances.pop(session_id, None)
_remove_stateful_session_tracking(session_id)
if transport is not None:
await transport.terminate()
verbose_logger.warning("MCP session '%s' terminated by an administrator.", session_id)
return MCPGatewaySessionsTerminateResponse(
worker_pid=os.getpid(),
terminated_sessions=len(terminated),
sessions=terminated,
)
async def _read_request_body_for_routing(
receive: Receive,
) -> tuple[list[Message], bytes]:
@ -4009,6 +4063,17 @@ if MCP_AVAILABLE:
await success_response(scope, receive, send)
return True
if _is_admin_terminated_session_id(_session_id, time.monotonic()):
terminated_response: Final = JSONResponse(
status_code=404,
content={ # mutable-ok: JSONResponse content must be a plain dict
"error": "Not Found",
"details": "mcp-session-id was terminated by an administrator. Send initialize to start a new session.",
},
)
await terminated_response(scope, receive, send)
return True
# Non-DELETE: strip stale session ID to allow new session creation
verbose_logger.warning(
"MCP session ID '%s' not found in this worker's memory. "

View file

@ -27989,6 +27989,32 @@
"title": "MCPGatewaySessionsResponse",
"type": "object"
},
"MCPGatewaySessionsTerminateResponse": {
"description": "Stateful sessions an administrator force-closed on this proxy worker.",
"properties": {
"sessions": {
"items": {
"$ref": "#/components/schemas/MCPGatewaySession"
},
"title": "Sessions",
"type": "array"
},
"terminated_sessions": {
"title": "Terminated Sessions",
"type": "integer"
},
"worker_pid": {
"title": "Worker Pid",
"type": "integer"
}
},
"required": [
"worker_pid",
"terminated_sessions"
],
"title": "MCPGatewaySessionsTerminateResponse",
"type": "object"
},
"MCPOAuthUserCredentialRequest": {
"description": "Stores a user's OAuth2 token for an OpenAPI MCP server.",
"properties": {
@ -28085,6 +28111,56 @@
"title": "MCPOAuthUserCredentialStatus",
"type": "object"
},
"MCPServerUserCredentialListItem": {
"description": "One user's stored credential for an MCP server, as an admin sees it. Never carries the secret.",
"properties": {
"connected_at": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Connected At"
},
"credential_type": {
"enum": [
"oauth2",
"byok"
],
"title": "Credential Type",
"type": "string"
},
"expires_at": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Expires At"
},
"updated_at": {
"title": "Updated At",
"type": "string"
},
"user_id": {
"title": "User Id",
"type": "string"
}
},
"required": [
"user_id",
"credential_type",
"updated_at"
],
"title": "MCPServerUserCredentialListItem",
"type": "object"
},
"MCPSubmissionsSummary": {
"properties": {
"active": {
@ -30261,7 +30337,7 @@
},
"/v1/mcp/server/{server_id}/oauth-user-credential": {
"delete": {
"description": "Revoke the calling user's stored OAuth2 token for an MCP server",
"description": "Revoke the calling user's stored OAuth2 token for an MCP server. A proxy admin may pass user_id to revoke another user's stored token.",
"operationId": "delete_mcp_oauth_user_credential_v1_mcp_server__server_id__oauth_user_credential_delete",
"parameters": [
{
@ -30272,6 +30348,23 @@
"title": "Server Id",
"type": "string"
}
},
{
"in": "query",
"name": "user_id",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"title": "User Id"
}
}
],
"responses": {
@ -30471,7 +30564,7 @@
},
"/v1/mcp/server/{server_id}/user-credential": {
"delete": {
"description": "Delete the calling user's stored API key for a BYOK MCP server",
"description": "Delete the calling user's stored API key for a BYOK MCP server. A proxy admin may pass user_id to revoke another user's stored key.",
"operationId": "delete_mcp_user_credential_v1_mcp_server__server_id__user_credential_delete",
"parameters": [
{
@ -30482,6 +30575,23 @@
"title": "Server Id",
"type": "string"
}
},
{
"in": "query",
"name": "user_id",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"title": "User Id"
}
}
],
"responses": {
@ -30573,6 +30683,58 @@
]
}
},
"/v1/mcp/server/{server_id}/user-credentials": {
"get": {
"description": "List every user's stored BYOK or OAuth2 credential for an MCP server (admin only, no secrets)",
"operationId": "list_mcp_server_user_credentials_v1_mcp_server__server_id__user_credentials_get",
"parameters": [
{
"in": "path",
"name": "server_id",
"required": true,
"schema": {
"title": "Server Id",
"type": "string"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"items": {
"$ref": "#/components/schemas/MCPServerUserCredentialListItem"
},
"title": "Response List Mcp Server User Credentials V1 Mcp Server Server Id User Credentials Get",
"type": "array"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "List Mcp Server User Credentials",
"tags": [
"mcp_management"
]
}
},
"/v1/mcp/server/{server_id}/user-env-vars": {
"delete": {
"description": "Clear the calling user's per-user MCP env var values for this server.",
@ -30724,6 +30886,77 @@
}
},
"/v1/mcp/sessions": {
"delete": {
"description": "Force-close live stateful MCP gateway sessions on this proxy worker, selected by session id prefix and/or by the LiteLLM user that opened them (proxy admin only).",
"operationId": "delete_mcp_gateway_sessions_v1_mcp_sessions_delete",
"parameters": [
{
"in": "query",
"name": "session_id_prefix",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 8,
"type": "string"
},
{
"type": "null"
}
],
"title": "Session Id Prefix"
}
},
{
"in": "query",
"name": "user_id",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"title": "User Id"
}
}
],
"responses": {
"200": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/MCPGatewaySessionsTerminateResponse"
}
}
},
"description": "Successful Response"
},
"422": {
"content": {
"application/json": {
"schema": {
"$ref": "#/components/schemas/HTTPValidationError"
}
}
},
"description": "Validation Error"
}
},
"security": [
{
"APIKeyHeader": []
}
],
"summary": "Delete Mcp Gateway Sessions",
"tags": [
"mcp_management"
]
},
"get": {
"description": "Live stateful MCP gateway sessions on this proxy worker, grouped by AI client and by user.",
"operationId": "get_mcp_gateway_sessions_v1_mcp_sessions_get",

View file

@ -1726,6 +1726,16 @@ class MCPUserCredentialListItem(LiteLLMPydanticObjectBase):
connected_at: str | None = None # ISO-8601
class MCPServerUserCredentialListItem(LiteLLMPydanticObjectBase):
"""One user's stored credential for an MCP server, as an admin sees it. Never carries the secret."""
user_id: str
credential_type: Literal["oauth2", "byok"]
expires_at: str | None = None
connected_at: str | None = None
updated_at: str
class MCPUserEnvVarsRequest(LiteLLMPydanticObjectBase):
"""Payload for storing the calling user's per-user env var values."""

View file

@ -47,7 +47,7 @@ except ImportError:
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.constants import LITELLM_PROXY_ADMIN_NAME, MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH
from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_DESCRIPTION,
LITELLM_MCP_SERVER_NAME,
@ -145,6 +145,7 @@ if MCP_AVAILABLE:
get_user_env_vars,
get_user_env_vars_bulk,
get_user_oauth_credential,
list_server_user_credentials,
list_user_oauth_credentials,
mcp_oauth_token_identity,
merge_user_env_vars,
@ -180,6 +181,7 @@ if MCP_AVAILABLE:
MCPApprovalStatus,
MCPOAuthUserCredentialRequest,
MCPOAuthUserCredentialStatus,
MCPServerUserCredentialListItem,
MCPSubmissionsSummary,
MCPTransport,
MCPUserCredentialListItem,
@ -221,6 +223,7 @@ if MCP_AVAILABLE:
MCPAuth,
MCPCredentials,
MCPGatewaySessionsResponse,
MCPGatewaySessionsTerminateResponse,
normalize_upstream_header_name,
)
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -662,6 +665,31 @@ if MCP_AVAILABLE:
"""
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
def _resolve_credential_target_user_id(user_api_key_dict: UserAPIKeyAuth, requested_user_id: str | None) -> str:
"""The user whose stored MCP credential a request acts on.
Defaults to the caller. Naming another user is a revocation and needs
``PROXY_ADMIN``; a read-only admin or a regular user gets 403.
"""
caller_user_id: Final = user_api_key_dict.user_id or ""
if requested_user_id is not None and requested_user_id != caller_user_id:
if not _user_is_full_admin(user_api_key_dict):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Proxy admin access required to revoke another user's MCP credential.",
},
)
return requested_user_id
if not caller_user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": "User ID not found in token"
}, # mutable-ok: FastAPI HTTPException detail requires a plain dict
)
return caller_user_id
def _is_restricted_virtual_key_request(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""Best-effort detection for route-restricted virtual keys.
@ -1373,6 +1401,41 @@ if MCP_AVAILABLE:
return get_mcp_gateway_sessions_report()
@router.delete(
"/sessions",
description=(
"Force-close live stateful MCP gateway sessions on this proxy worker, selected by session id prefix "
"and/or by the LiteLLM user that opened them (proxy admin only)."
),
dependencies=(Depends(user_api_key_auth),),
response_model=MCPGatewaySessionsTerminateResponse,
)
@management_endpoint_wrapper
async def delete_mcp_gateway_sessions(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
session_id_prefix: Annotated[str | None, Query(min_length=MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH)] = None,
user_id: Annotated[str | None, Query(min_length=1)] = None,
) -> MCPGatewaySessionsTerminateResponse:
if not _user_is_full_admin(user_api_key_dict):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Proxy admin access required to terminate MCP gateway sessions.",
},
)
if session_id_prefix is None and user_id is None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Provide session_id_prefix and/or user_id to select the sessions to terminate.",
},
)
from litellm.proxy._experimental.mcp_server.server import (
terminate_mcp_gateway_sessions,
)
return await terminate_mcp_gateway_sessions(session_id_prefix=session_id_prefix, user_id=user_id)
@router.get(
"/server/submissions",
description="Returns all MCP servers submitted by non-admin users (admin review queue). Mirrors GET /guardrails/submissions.",
@ -2254,14 +2317,17 @@ if MCP_AVAILABLE:
_invalidate_byok_cred_cache,
)
_invalidate_byok_cred_cache(user_id, server_id)
await _invalidate_byok_cred_cache(user_id, server_id)
return MCPUserCredentialResponse(server_id=server_id, has_credential=True)
# save=False: credential not persisted
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
@router.delete(
"/server/{server_id}/user-credential",
description="Delete the calling user's stored API key for a BYOK MCP server",
description=(
"Delete the calling user's stored API key for a BYOK MCP server. "
"A proxy admin may pass user_id to revoke another user's stored key."
),
dependencies=[Depends(user_api_key_auth)],
response_model=MCPUserCredentialResponse,
)
@ -2269,24 +2335,20 @@ if MCP_AVAILABLE:
async def delete_mcp_user_credential(
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_id: Annotated[str | None, Query(min_length=1)] = None,
):
"""Remove the calling user's BYOK credential."""
"""Remove the target user's BYOK credential (the caller unless an admin names another user)."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
user_id: Final = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
target_user_id: Final = _resolve_credential_target_user_id(user_api_key_dict, user_id)
try:
await delete_user_credential(prisma_client, user_id, server_id)
await delete_user_credential(prisma_client, target_user_id, server_id)
except RecordNotFoundError:
pass # Already deleted or didn't exist
from litellm.proxy._experimental.mcp_server.server import (
_invalidate_byok_cred_cache,
)
_invalidate_byok_cred_cache(user_id, server_id)
await _invalidate_byok_cred_cache(target_user_id, server_id)
return MCPUserCredentialResponse(server_id=server_id, has_credential=False)
# ── OAuth2 user-credential endpoints ──────────────────────────────────────
@ -2362,7 +2424,10 @@ if MCP_AVAILABLE:
@router.delete(
"/server/{server_id}/oauth-user-credential",
description="Revoke the calling user's stored OAuth2 token for an MCP server",
description=(
"Revoke the calling user's stored OAuth2 token for an MCP server. "
"A proxy admin may pass user_id to revoke another user's stored token."
),
dependencies=[Depends(user_api_key_auth)],
response_model=MCPOAuthUserCredentialStatus,
)
@ -2370,29 +2435,25 @@ if MCP_AVAILABLE:
async def delete_mcp_oauth_user_credential(
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_id: Annotated[str | None, Query(min_length=1)] = None,
):
"""Revoke/delete the user's OAuth2 credential."""
"""Revoke the target user's OAuth2 credential (the caller unless an admin names another user)."""
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
user_id: Final = user_api_key_dict.user_id or ""
if not user_id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={"error": "User ID not found in token"},
)
target_user_id: Final = _resolve_credential_target_user_id(user_api_key_dict, user_id)
# Only delete if the stored credential is actually an OAuth2 token.
# This prevents accidentally deleting a BYOK credential if one exists
# for the same (user_id, server_id) pair.
cred_to_delete: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
cred_to_delete: Final = await get_user_oauth_credential(prisma_client, target_user_id, server_id)
if cred_to_delete is not None:
try:
await delete_user_credential(prisma_client, user_id, server_id)
await delete_user_credential(prisma_client, target_user_id, server_id)
except RecordNotFoundError:
pass # Already gone — treat as a successful delete
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415
global_mcp_server_manager,
)
await global_mcp_server_manager.invalidate_user_oauth_token_cache(user_id, server_id)
await global_mcp_server_manager.invalidate_user_oauth_token_cache(target_user_id, server_id)
return MCPOAuthUserCredentialStatus(
server_id=server_id,
has_credential=False,
@ -2481,6 +2542,30 @@ if MCP_AVAILABLE:
)
return items
@router.get(
"/server/{server_id}/user-credentials",
description="List every user's stored BYOK or OAuth2 credential for an MCP server (admin only, no secrets)",
dependencies=(Depends(user_api_key_auth),),
response_model=list[MCPServerUserCredentialListItem],
)
@management_endpoint_wrapper
async def list_mcp_server_user_credentials(
server_id: str,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
) -> tuple[MCPServerUserCredentialListItem, ...]:
if user_api_key_dict.user_role not in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict
"error": "Admin access required to view MCP server user credentials.",
},
)
prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
return await list_server_user_credentials(prisma_client, server_id)
# ── Per-user MCP env var endpoints ────────────────────────────────────────
async def _authorize_and_fetch_mcp_server(

View file

@ -309,6 +309,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import (
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache
from litellm.proxy._lazy_features import attach_lazy_features, reserve_lazy_slot
from litellm.proxy._types import *
from litellm.proxy.analytics_endpoints.analytics_endpoints import (
@ -7550,7 +7551,7 @@ class ProxyConfig:
subscriber: Final = AuthCacheInvalidationSubscriber(
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
additional_in_memory_caches=(spend_counter_cache.in_memory_cache,),
additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache),
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()

View file

@ -464,3 +464,11 @@ class MCPGatewaySessionsResponse(BaseModel):
by_client: list[MCPGatewaySessionGroupCount] = Field(default_factory=list)
by_user: list[MCPGatewaySessionGroupCount] = Field(default_factory=list)
sessions: list[MCPGatewaySession] = Field(default_factory=list)
class MCPGatewaySessionsTerminateResponse(BaseModel):
"""Stateful sessions an administrator force-closed on this proxy worker."""
worker_pid: int
terminated_sessions: int
sessions: list[MCPGatewaySession] = Field(default_factory=list)

View file

@ -331,9 +331,7 @@ class TestMCPPerUserTokenCache:
with patch("litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache):
await cache.delete("alice", "slack-test")
mock_dual_cache.async_delete_cache.assert_called_once_with(
"mcp:per_user_token:alice:slack-test"
)
mock_dual_cache.async_delete_cache.assert_called_once_with(key="mcp:per_user_token:alice:slack-test")
mock_dual_cache.async_set_cache.assert_not_called()
@pytest.mark.asyncio

View file

@ -0,0 +1,57 @@
import json
import pytest
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
CachedByokCredential,
byok_credential_cache,
byok_credential_cache_key,
cache_byok_credential,
get_cached_byok_credential,
)
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
class _FakeRedisCache:
namespace = None
def init_async_client(self) -> object:
return object()
@pytest.fixture(autouse=True)
def _empty_cache():
byok_credential_cache.flush_cache()
yield
byok_credential_cache.flush_cache()
def test_a_cached_negative_lookup_is_distinguishable_from_a_miss():
assert get_cached_byok_credential("u-1", "srv-1") is None
cache_byok_credential("u-1", "srv-1", None)
assert get_cached_byok_credential("u-1", "srv-1") == CachedByokCredential(credential=None)
cache_byok_credential("u-1", "srv-1", "sk-stored")
assert get_cached_byok_credential("u-1", "srv-1") == CachedByokCredential(credential="sk-stored")
assert get_cached_byok_credential("u-1", "srv-2") is None
def test_peer_worker_invalidation_message_evicts_the_cached_credential():
"""The key a mutating worker broadcasts must be the key every other worker caches under."""
cache_byok_credential("mallory", "srv-byok", "sk-revoked")
cache_byok_credential("alice", "srv-byok", "sk-kept")
subscriber = AuthCacheInvalidationSubscriber(
redis_cache=_FakeRedisCache(), # pyright: ignore[reportArgumentType] # subscriber is never started; only its message handler runs
user_api_key_cache=UserApiKeyCache(),
additional_in_memory_caches=(byok_credential_cache,),
)
subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler
{
"type": "message",
"data": json.dumps({"cache_key": byok_credential_cache_key("mallory", "srv-byok")}).encode(),
}
)
assert get_cached_byok_credential("mallory", "srv-byok") is None
assert get_cached_byok_credential("alice", "srv-byok") == CachedByokCredential(credential="sk-kept")

View file

@ -592,7 +592,7 @@ async def test_check_byok_credential_missing_credential(monkeypatch):
monkeypatch.delenv("PROXY_BASE_URL", raising=False)
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
monkeypatch.setattr(server_module, "_byok_cred_cache", {})
server_module.byok_credential_cache.flush_cache()
mock_prisma = MagicMock()
with (
@ -628,7 +628,7 @@ async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monk
from litellm.types.mcp_server.mcp_server_manager import MCPServer
monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy")
monkeypatch.setattr(mcp_module, "_byok_cred_cache", {})
mcp_module.byok_credential_cache.flush_cache()
server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True)
prisma = MagicMock()
prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None)
@ -677,6 +677,40 @@ async def test_check_byok_credential_has_credential():
await _check_byok_credential(server, user_auth)
@pytest.mark.asyncio
async def test_invalidate_byok_cred_cache_evicts_locally_and_broadcasts_the_same_key():
"""A revoked credential must stop being served here and on every peer worker within the TTL."""
from litellm.proxy._experimental.mcp_server import server as server_module
from litellm.proxy._experimental.mcp_server.byok_credential_cache import byok_credential_cache_key
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(server_id="byok-revoke", name="byok-server", transport=MCPTransport.http, is_byok=True)
user_auth = UserAPIKeyAuth(user_id="mallory", api_key="sk-test")
server_module.byok_credential_cache.flush_cache()
db_lookup = AsyncMock(side_effect=["sk-before-revoke", None])
publish = AsyncMock()
with (
patch( # test-quality-ok: the DB row lookup is the only seam below the credential resolver; no Prisma fake exists
"litellm.proxy._experimental.mcp_server.db.get_user_credential", new=db_lookup
),
patch( # test-quality-ok: the resolver reads the module-level prisma_client singleton; the suite's only seam
"litellm.proxy.proxy_server.prisma_client", MagicMock()
),
patch.object( # test-quality-ok: the redis publisher is module-level; asserting the broadcast without a redis
server_module, "publish_auth_cache_invalidation", new=publish
),
):
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
assert await server_module._get_byok_credential(server, user_auth) == "sk-before-revoke"
await server_module._invalidate_byok_cred_cache("mallory", "byok-revoke")
assert await server_module._get_byok_credential(server, user_auth) is None
assert db_lookup.await_count == 2
publish.assert_awaited_once_with(cache_key=byok_credential_cache_key("mallory", "byok-revoke"))
@pytest.mark.asyncio
async def test_check_byok_credential_db_unavailable_fails_closed():
"""BYOK server with no prisma_client → 503, not silent pass.

View file

@ -212,6 +212,53 @@ 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_list_server_user_credentials_types_each_row_without_leaking_the_secret():
"""The admin view of one server's stored credentials names the user and the kind of
credential (OAuth2 vs BYOK) and echoes OAuth expiry, but never the token or key itself."""
from litellm.proxy._experimental.mcp_server.db import list_server_user_credentials
oauth_row = _legacy_row(
json.dumps(
{
"type": "oauth2",
"access_token": "tok-alice",
"expires_at": "2026-12-31T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
}
)
)
oauth_row.user_id = "alice"
oauth_row.updated_at = datetime(2026, 1, 1, tzinfo=timezone.utc)
byok_row = _byok_row("carol")
byok_row.updated_at = datetime(2026, 2, 1, tzinfo=timezone.utc)
prisma = MagicMock()
prisma.db.litellm_mcpusercredentials.find_many = AsyncMock(return_value=[oauth_row, byok_row])
items = await list_server_user_credentials(prisma, "srv-1")
prisma.db.litellm_mcpusercredentials.find_many.assert_awaited_once_with(where={"server_id": "srv-1"})
assert [item.model_dump() for item in items] == [
{
"user_id": "alice",
"credential_type": "oauth2",
"expires_at": "2026-12-31T00:00:00+00:00",
"connected_at": "2026-01-01T00:00:00+00:00",
"updated_at": "2026-01-01T00:00:00+00:00",
},
{
"user_id": "carol",
"credential_type": "byok",
"expires_at": None,
"connected_at": None,
"updated_at": "2026-02-01T00:00:00+00:00",
},
]
serialized = "".join(item.model_dump_json() for item in items)
assert "tok-alice" not in serialized
assert "sk-byok-carol" not in serialized
@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

View file

@ -2871,6 +2871,255 @@ def test_remove_stateful_session_tracking_drops_client_info():
assert session_id not in mcp_server._stateful_session_client_info
def _admin_terminate_fixture(mcp_server):
def auth_user(user_id: str):
return mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key=f"key-{user_id}", user_id=user_id),
)
contexts = {
"alice-session-1": auth_user("alice"),
"alice-session-2": auth_user("alice"),
"bob-session-1": auth_user("bob"),
"anon-session-1": mcp_server.MCPAuthenticatedUser(user_api_key_auth=None),
"gone-session-1": auth_user("alice"),
}
transports = {
session_id: MagicMock(terminate=AsyncMock())
for session_id in ("alice-session-1", "alice-session-2", "bob-session-1", "anon-session-1")
}
return contexts, transports
@pytest.mark.asyncio
async def test_terminate_mcp_gateway_sessions_by_user_closes_every_live_session_of_that_user():
try:
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
contexts, transports = _admin_terminate_fixture(mcp_server)
live_transports = dict(transports)
last_seen = {session_id: 100.0 for session_id in contexts}
locks = {session_id: asyncio.Lock() for session_id in contexts}
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_context_last_seen, last_seen, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_locks, locks, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_owners, {session_id: "owner" for session_id in contexts}, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_active_request_counts, {}, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
):
result = await mcp_server.terminate_mcp_gateway_sessions(user_id="alice")
assert set(live_transports) == {"bob-session-1", "anon-session-1"}
assert set(mcp_server._stateful_session_auth_contexts) == {"bob-session-1", "anon-session-1", "gone-session-1"}
assert set(mcp_server._stateful_session_locks) == {"bob-session-1", "anon-session-1", "gone-session-1"}
assert set(mcp_server._stateful_session_owners) == {"bob-session-1", "anon-session-1", "gone-session-1"}
assert set(mcp_server._stateful_session_auth_context_last_seen) == {
"bob-session-1",
"anon-session-1",
"gone-session-1",
}
transports["alice-session-1"].terminate.assert_awaited_once()
transports["alice-session-2"].terminate.assert_awaited_once()
transports["bob-session-1"].terminate.assert_not_awaited()
transports["anon-session-1"].terminate.assert_not_awaited()
assert result.terminated_sessions == 2
assert sorted(session.session_id_prefix for session in result.sessions) == ["alice-se", "alice-se"]
assert {session.user_id for session in result.sessions} == {"alice"}
assert "key-alice" not in result.model_dump_json()
assert "alice-session-1" not in result.model_dump_json()
@pytest.mark.asyncio
async def test_terminate_mcp_gateway_sessions_prefix_and_user_must_both_match():
try:
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
contexts, transports = _admin_terminate_fixture(mcp_server)
live_transports = dict(transports)
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
):
mismatch = await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix="alice-session-1", user_id="bob")
assert mismatch.terminated_sessions == 0
assert set(live_transports) == set(transports)
stale = await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix="gone-session-1")
assert stale.terminated_sessions == 0
exact = await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix="alice-session-1", user_id="alice")
assert exact.terminated_sessions == 1
assert set(live_transports) == {"alice-session-2", "bob-session-1", "anon-session-1"}
@pytest.mark.asyncio
async def test_admin_terminated_session_id_gets_404_instead_of_a_fresh_stateless_session():
"""Once an admin closes a session, a client replaying its id must not be silently upgraded to a
new stateless session by the stale-header path; it gets 404 and has to initialize again."""
try:
from starlette.types import Scope
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
session_id = "admin-closed-session-1"
live_transports = {session_id: MagicMock(terminate=AsyncMock())}
contexts = {
session_id: mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key="key-alice", user_id="alice"),
)
}
def scope_with_session_header() -> Scope:
return {
"type": "http",
"method": "POST",
"headers": [(b"content-type", b"application/json"), (b"mcp-session-id", session_id.encode())],
}
try:
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
):
await mcp_server.terminate_mcp_gateway_sessions(session_id_prefix=session_id)
terminated_scope = scope_with_session_header()
send = AsyncMock()
handled = await mcp_server._handle_stale_mcp_session(
terminated_scope, AsyncMock(), send, session_manager_stateful
)
assert handled is True
statuses = [m["status"] for (m,), _ in send.await_args_list if m["type"] == "http.response.start"]
assert statuses == [404]
assert [k for k, _ in terminated_scope["headers"]] == [b"content-type", b"mcp-session-id"]
unknown_scope = scope_with_session_header()
unknown_scope["headers"][1] = (b"mcp-session-id", b"never-seen-session")
assert (
await mcp_server._handle_stale_mcp_session(
unknown_scope, AsyncMock(), AsyncMock(), session_manager_stateful
)
is False
)
assert [k for k, _ in unknown_scope["headers"]] == [b"content-type"]
finally:
mcp_server._admin_terminated_session_ids.clear()
@pytest.mark.asyncio
async def test_admin_terminated_session_id_stays_refused_while_replayed_and_is_forgotten_like_an_idle_session():
"""The refusal window slides on every replay, so a client that keeps retrying is never silently
upgraded to a stateless session no matter how many other sessions an admin closes later; an id
nobody has replayed for a full idle timeout is dropped from the table by the idle sweep."""
try:
from starlette.types import Scope
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
except ImportError:
pytest.skip("MCP server not available")
idle_timeout = mcp_server._STATEFUL_SESSION_IDLE_TIMEOUT_SECONDS
retrying_id, silent_id = "admin-closed-retrying", "admin-closed-silent"
contexts = {
session_id: mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key="key-alice", user_id="alice"),
)
for session_id in (retrying_id, silent_id)
}
live_transports = {session_id: MagicMock(terminate=AsyncMock()) for session_id in contexts}
async def replay(session_id: str, now: float) -> tuple[bool, list[bytes]]:
scope: Scope = {
"type": "http",
"method": "POST",
"headers": [(b"content-type", b"application/json"), (b"mcp-session-id", session_id.encode())],
}
with patch.object( # test-quality-ok: the stale-session handler reads the clock directly; no injectable now
mcp_server.time, "monotonic", return_value=now
):
handled = await mcp_server._handle_stale_mcp_session(
scope, AsyncMock(), AsyncMock(), session_manager_stateful
)
return handled, [k for k, _ in scope["headers"]]
try:
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
session_manager_stateful, "_server_instances", live_transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_client_info, {}, clear=True
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_context_last_seen, {}, clear=True
),
):
with patch.object( # test-quality-ok: termination stamps the tombstone from the clock directly; no injectable now
mcp_server.time, "monotonic", return_value=1000.0
):
closed = await mcp_server.terminate_mcp_gateway_sessions(user_id="alice")
assert closed.terminated_sessions == 2
for elapsed in (idle_timeout - 1, 2 * idle_timeout - 2, 3 * idle_timeout - 3):
assert await replay(retrying_id, 1000.0 + elapsed) == (True, [b"content-type", b"mcp-session-id"])
await mcp_server._purge_expired_stateful_session_auth_contexts(now=1000.0 + idle_timeout)
assert set(mcp_server._admin_terminated_session_ids) == {retrying_id}
assert await replay(silent_id, 1000.0 + idle_timeout) == (False, [b"content-type"])
assert await replay(retrying_id, 1000.0 + 4 * idle_timeout) == (False, [b"content-type"])
assert mcp_server._admin_terminated_session_ids == {}
finally:
mcp_server._admin_terminated_session_ids.clear()
@pytest.mark.asyncio
async def test_initialize_request_with_existing_session_tracks_new_session():
try:

View file

@ -395,6 +395,34 @@ async def test_invalidate_clears_every_identity_for_a_server():
assert mock_client.post.call_count == 3
@pytest.mark.asyncio
async def test_per_user_token_delete_evicts_locally_and_broadcasts_to_peer_workers():
"""Revoking a user's OAuth token must not leave peer workers serving it from their in-memory layer."""
from litellm.proxy import proxy_server
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import MCPPerUserTokenCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
local_cache = UserApiKeyCache()
publish = AsyncMock()
token_cache = MCPPerUserTokenCache()
key = token_cache._cache_key("mallory", "srv-oauth") # pyright: ignore[reportPrivateUsage] # asserting the broadcast names the stored key
local_cache.in_memory_cache.set_cache(key, "encrypted-token")
with (
patch.object( # test-quality-ok: the token cache reads the module-level user_api_key_cache singleton; the suite's only seam
proxy_server, "user_api_key_cache", local_cache
),
patch( # test-quality-ok: the redis publisher is module-level; asserting the broadcast without a redis
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
new=publish,
),
):
await token_cache.delete("mallory", "srv-oauth")
assert local_cache.in_memory_cache.get_cache(key) is None
publish.assert_awaited_once_with(cache_key=key)
@pytest.mark.asyncio
async def test_m2m_mint_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the

View file

@ -24,6 +24,7 @@ from litellm.proxy._types import (
LiteLLM_MCPServerTable,
LitellmUserRoles,
MCPTransport,
MCPUserCredentialResponse,
NewMCPServerRequest,
UpdateMCPServerRequest,
UserAPIKeyAuth,
@ -5136,6 +5137,266 @@ async def test_delete_mcp_oauth_user_credential_invalidates_when_record_already_
assert result.has_credential is False
def _make_admin_auth(role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN) -> "UserAPIKeyAuth":
return UserAPIKeyAuth(api_key="sk-admin", user_id="admin-user", user_role=role)
@pytest.mark.asyncio
async def test_admin_revokes_another_users_byok_credential():
"""A proxy admin naming user_id deletes and cache-invalidates that user's stored key, not their own."""
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_user_credential,
)
delete_mock = AsyncMock(return_value=None)
invalidate_mock = AsyncMock()
with (
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: endpoint test stubs the credential row delete
"litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential",
new=delete_mock,
),
patch.object( # test-quality-ok: the cache invalidator is module scoped; the suite's only seam
mcp_server, "_invalidate_byok_cred_cache", new=invalidate_mock
),
):
result = await delete_mcp_user_credential(
server_id="srv-byok-admin",
user_api_key_dict=_make_admin_auth(),
user_id="mallory",
)
delete_mock.assert_awaited_once()
assert delete_mock.await_args.args[1:] == ("mallory", "srv-byok-admin")
invalidate_mock.assert_awaited_once_with("mallory", "srv-byok-admin")
assert result.has_credential is False
@pytest.mark.asyncio
@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
async def test_non_full_admin_cannot_revoke_another_users_byok_credential(role):
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_user_credential,
)
delete_mock = AsyncMock(return_value=None)
with (
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: endpoint test stubs the credential row delete
"litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential",
new=delete_mock,
),
):
with pytest.raises(HTTPException) as exc_info:
await delete_mcp_user_credential(
server_id="srv-byok-forbidden",
user_api_key_dict=_make_admin_auth(role),
user_id="mallory",
)
assert exc_info.value.status_code == 403
delete_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_user_naming_themselves_still_deletes_own_byok_credential():
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_user_credential,
)
deleted_rows: list[tuple[str, str]] = [] # mutable-ok: test-local recorder for the fake delete boundary
async def _fake_delete_user_credential(_prisma_client: object, user_id: str, server_id: str) -> None:
deleted_rows.append((user_id, server_id))
with (
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: endpoint test stubs the credential row delete
"litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential",
new=_fake_delete_user_credential,
),
patch.object( # test-quality-ok: the cache invalidator is module scoped; the suite's only seam
mcp_server, "_invalidate_byok_cred_cache", new=AsyncMock()
),
):
result = await delete_mcp_user_credential(
server_id="srv-byok-self",
user_api_key_dict=_make_user_auth("user-self"),
user_id="user-self",
)
assert deleted_rows == [("user-self", "srv-byok-self")]
assert result == MCPUserCredentialResponse(server_id="srv-byok-self", has_credential=False)
@pytest.mark.asyncio
async def test_admin_revokes_another_users_oauth_credential():
"""A proxy admin naming user_id reads, deletes, and cache-invalidates that user's OAuth token."""
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_oauth_user_credential,
)
get_mock = AsyncMock(return_value={"type": "oauth2", "access_token": "mallory-tok"})
delete_mock = AsyncMock(return_value=None)
invalidate_mock = AsyncMock(return_value=None)
with (
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: endpoint test stubs the stored OAuth token read
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential",
new=get_mock,
),
patch( # test-quality-ok: endpoint test stubs the credential row delete
"litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential",
new=delete_mock,
),
patch.object( # test-quality-ok: the OAuth cache lives on the global manager; the suite's only seam
manager_module.global_mcp_server_manager,
"invalidate_user_oauth_token_cache",
new=invalidate_mock,
),
):
result = await delete_mcp_oauth_user_credential(
server_id="srv-oauth-admin",
user_api_key_dict=_make_admin_auth(),
user_id="mallory",
)
assert get_mock.await_args.args[1:] == ("mallory", "srv-oauth-admin")
assert delete_mock.await_args.args[1:] == ("mallory", "srv-oauth-admin")
invalidate_mock.assert_awaited_once_with("mallory", "srv-oauth-admin")
assert result.has_credential is False
@pytest.mark.asyncio
@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
async def test_non_full_admin_cannot_revoke_another_users_oauth_credential(role):
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_oauth_user_credential,
)
get_mock = AsyncMock(return_value={"type": "oauth2", "access_token": "mallory-tok"})
delete_mock = AsyncMock(return_value=None)
with (
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: endpoint test stubs the stored OAuth token read
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential",
new=get_mock,
),
patch( # test-quality-ok: endpoint test stubs the credential row delete
"litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential",
new=delete_mock,
),
):
with pytest.raises(HTTPException) as exc_info:
await delete_mcp_oauth_user_credential(
server_id="srv-oauth-forbidden",
user_api_key_dict=_make_admin_auth(role),
user_id="mallory",
)
assert exc_info.value.status_code == 403
get_mock.assert_not_awaited()
delete_mock.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
async def test_admin_lists_every_users_credential_for_a_server(role):
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy._types import MCPServerUserCredentialListItem
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
list_mcp_server_user_credentials,
)
items = (
MCPServerUserCredentialListItem(user_id="alice", credential_type="byok", updated_at="2026-01-01T00:00:00"),
MCPServerUserCredentialListItem(user_id="bob", credential_type="oauth2", updated_at="2026-01-02T00:00:00"),
)
list_mock = AsyncMock(return_value=items)
with (
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: endpoint test stubs the credential row listing
"litellm.proxy.management_endpoints.mcp_management_endpoints.list_server_user_credentials",
new=list_mock,
),
):
result = await list_mcp_server_user_credentials(
server_id="srv-list-admin",
user_api_key_dict=_make_admin_auth(role),
)
assert list_mock.await_args.args[1:] == ("srv-list-admin",)
assert [(item.user_id, item.credential_type) for item in result] == [("alice", "byok"), ("bob", "oauth2")]
@pytest.mark.asyncio
async def test_non_admin_cannot_list_a_servers_user_credentials():
if not mgmt_endpoints.MCP_AVAILABLE:
pytest.skip("MCP module not installed")
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
list_mcp_server_user_credentials,
)
list_mock = AsyncMock(return_value=())
with (
patch( # test-quality-ok: endpoint test stubs the Prisma client lookup
"litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw",
return_value=_make_prisma_client(),
),
patch( # test-quality-ok: endpoint test stubs the credential row listing
"litellm.proxy.management_endpoints.mcp_management_endpoints.list_server_user_credentials",
new=list_mock,
),
):
with pytest.raises(HTTPException) as exc_info:
await list_mcp_server_user_credentials(
server_id="srv-list-forbidden",
user_api_key_dict=_make_user_auth("user-plain"),
)
assert exc_info.value.status_code == 403
list_mock.assert_not_awaited()
@pytest.mark.asyncio
async def test_list_mcp_user_credentials_batch_server_fetch():
"""list_mcp_user_credentials uses a single batch DB call, not N+1 queries."""
@ -7321,3 +7582,146 @@ class TestGetMCPGatewaySessions:
assert [(group.label, group.count) for group in result.by_client] == [("cursor", 1)]
assert [(group.label, group.count) for group in result.by_user] == [("alice", 1)]
assert "sk-live-secret" not in result.model_dump_json()
class TestDeleteMCPGatewaySessions:
@pytest.fixture(autouse=True)
def _forget_admin_terminated_ids(self):
from litellm.proxy._experimental.mcp_server import server as mcp_server
yield
mcp_server._admin_terminated_session_ids.clear()
@pytest.mark.asyncio
@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
async def test_non_full_admin_forbidden_before_any_session_is_touched(self, role):
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_gateway_sessions,
)
session_id = "gateway-terminate-forbidden-1"
transport = MagicMock(terminate=AsyncMock())
auth_user = mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key="sk-live", user_id="alice"),
)
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
mcp_server.session_manager_stateful, "_server_instances", {session_id: transport}
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, {session_id: auth_user}, clear=True
),
):
with pytest.raises(HTTPException) as exc_info:
await delete_mcp_gateway_sessions(
user_api_key_dict=generate_mock_user_api_key_auth(user_role=role),
session_id_prefix=session_id,
user_id=None,
)
assert exc_info.value.status_code == 403
transport.terminate.assert_not_awaited()
assert session_id in mcp_server._stateful_session_auth_contexts
@pytest.mark.asyncio
async def test_requires_a_selector(self):
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_gateway_sessions,
)
with pytest.raises(HTTPException) as exc_info:
await delete_mcp_gateway_sessions(
user_api_key_dict=generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN),
session_id_prefix=None,
user_id=None,
)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_admin_terminates_only_the_selected_session(self):
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_gateway_sessions,
)
from litellm.types.mcp import MCPGatewaySessionsTerminateResponse
target_id = "11111111-target-session"
other_id = "22222222-other-session"
target_transport = MagicMock(terminate=AsyncMock())
other_transport = MagicMock(terminate=AsyncMock())
transports = {target_id: target_transport, other_id: other_transport}
contexts = {
target_id: mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key="sk-live-target", user_id="alice"),
),
other_id: mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key="sk-live-other", user_id="bob"),
),
}
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
mcp_server.session_manager_stateful, "_server_instances", transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
):
result = await delete_mcp_gateway_sessions(
user_api_key_dict=generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN),
session_id_prefix=target_id[:8],
user_id=None,
)
assert target_id not in transports
assert other_id in transports
assert target_id not in mcp_server._stateful_session_auth_contexts
assert other_id in mcp_server._stateful_session_auth_contexts
target_transport.terminate.assert_awaited_once()
other_transport.terminate.assert_not_awaited()
assert isinstance(result, MCPGatewaySessionsTerminateResponse)
assert result.terminated_sessions == 1
assert [(s.session_id_prefix, s.user_id) for s in result.sessions] == [(target_id[:8], "alice")]
assert target_id not in result.model_dump_json()
assert "sk-live-target" not in result.model_dump_json()
@pytest.mark.asyncio
async def test_admin_terminates_every_session_of_the_selected_user(self):
from litellm.proxy._experimental.mcp_server import server as mcp_server
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
delete_mcp_gateway_sessions,
)
def auth_user(user_id: str):
return mcp_server.MCPAuthenticatedUser(
user_api_key_auth=UserAPIKeyAuth(api_key=f"sk-live-{user_id}", user_id=user_id),
)
transports = {
"bob-session-1": MagicMock(terminate=AsyncMock()),
"bob-session-2": MagicMock(terminate=AsyncMock()),
"alice-session-1": MagicMock(terminate=AsyncMock()),
}
contexts = {
"bob-session-1": auth_user("bob"),
"bob-session-2": auth_user("bob"),
"alice-session-1": auth_user("alice"),
}
with (
patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam
mcp_server.session_manager_stateful, "_server_instances", transports
),
patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam
mcp_server._stateful_session_auth_contexts, contexts, clear=True
),
):
result = await delete_mcp_gateway_sessions(
user_api_key_dict=generate_mock_user_api_key_auth(user_role=LitellmUserRoles.PROXY_ADMIN),
session_id_prefix=None,
user_id="bob",
)
assert set(transports) == {"alice-session-1"}
assert set(mcp_server._stateful_session_auth_contexts) == {"alice-session-1"}
assert result.terminated_sessions == 2
assert {s.user_id for s in result.sessions} == {"bob"}
assert "sk-live-bob" not in result.model_dump_json()

View file

@ -14368,3 +14368,74 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi
]
finally:
litellm.utils._select_custom_tokenizer_helper.cache_clear()
@pytest.mark.asyncio
async def test_auth_cache_invalidation_subscriber_evicts_byok_credentials_cached_by_this_worker():
"""A peer worker's BYOK revocation broadcast must reach this worker's BYOK credential cache."""
from redis.asyncio import Redis
from litellm.proxy._experimental.mcp_server.byok_credential_cache import (
byok_credential_cache,
byok_credential_cache_key,
cache_byok_credential,
get_cached_byok_credential,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
class _QueuePubSub:
def __init__(self, messages: list[object]) -> None:
self.queue: asyncio.Queue[object] = asyncio.Queue()
for message in messages:
self.queue.put_nowait(message)
async def subscribe(self, *channels: str) -> None:
return None
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> object | None:
try:
return await asyncio.wait_for(self.queue.get(), timeout)
except asyncio.TimeoutError:
return None
async def aclose(self) -> None:
return None
class _PubSubRedisClient(Redis):
def __init__(self, pubsub: _QueuePubSub) -> None:
self._scripted_pubsub = pubsub
def pubsub(self) -> _QueuePubSub:
return self._scripted_pubsub
class _FakeRedisCache:
namespace = None
def __init__(self, client: object) -> None:
self._client = client
def init_async_client(self) -> object:
return self._client
byok_credential_cache.flush_cache()
cache_byok_credential("mallory", "srv-byok", "sk-revoked-elsewhere")
message: Final = {
"type": "message",
"data": json.dumps({"cache_key": byok_credential_cache_key("mallory", "srv-byok")}).encode(),
}
proxy_config: Final = proxy_server_module.ProxyConfig()
proxy_config.start_auth_cache_invalidation_subscriber(
redis_cache=_FakeRedisCache(_PubSubRedisClient(_QueuePubSub([message]))), # pyright: ignore[reportArgumentType] # fake pub/sub capable redis; no live redis in this unit test
user_api_key_cache=UserApiKeyCache(),
)
try:
for _ in range(200):
if get_cached_byok_credential("mallory", "srv-byok") is None:
break
await asyncio.sleep(0.01)
evicted: Final = get_cached_byok_credential("mallory", "srv-byok") is None
finally:
await proxy_config.stop_auth_cache_invalidation_subscriber()
byok_credential_cache.flush_cache()
assert evicted, "the subscriber does not evict the BYOK credential cache on a peer worker's broadcast"

View file

@ -1,13 +1,15 @@
import React from "react";
import { render, screen, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { describe, it, expect, vi, beforeEach } from "vitest";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { MCPGatewaySessionsTab, formatIdleSeconds } from "./MCPGatewaySessionsTab";
import { MCPGatewaySessionsTab, describeTerminateResult, formatIdleSeconds } from "./MCPGatewaySessionsTab";
import * as networking from "@/components/networking";
import type { MCPGatewaySessionsResponse } from "@/components/mcp_tools/types";
import type { MCPGatewaySessionsResponse, MCPGatewaySessionsTerminateResponse } from "@/components/mcp_tools/types";
vi.mock("@/components/networking", () => ({
fetchMCPGatewaySessions: vi.fn(),
terminateMCPGatewaySessions: vi.fn(),
}));
const REPORT: MCPGatewaySessionsResponse = {
@ -64,11 +66,11 @@ const REPORT: MCPGatewaySessionsResponse = {
],
};
const renderTab = () => {
const renderTab = ({ canTerminate = false }: { canTerminate?: boolean } = {}) => {
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
return render(
<QueryClientProvider client={queryClient}>
<MCPGatewaySessionsTab accessToken="token" />
<MCPGatewaySessionsTab accessToken="token" canTerminate={canTerminate} />
</QueryClientProvider>,
);
};
@ -83,6 +85,17 @@ describe("formatIdleSeconds", () => {
});
});
describe("describeTerminateResult", () => {
it("pluralizes the session count and names the worker", () => {
expect(describeTerminateResult({ worker_pid: 9, terminated_sessions: 1, sessions: [] })).toBe(
"Disconnected 1 session on worker pid 9.",
);
expect(describeTerminateResult({ worker_pid: 9, terminated_sessions: 0, sessions: [] })).toBe(
"Disconnected 0 sessions on worker pid 9.",
);
});
});
describe("MCPGatewaySessionsTab", () => {
beforeEach(() => {
vi.clearAllMocks();
@ -135,4 +148,73 @@ describe("MCPGatewaySessionsTab", () => {
expect(alert).toHaveTextContent("Could not load live connections");
expect(alert).toHaveTextContent("Admin access required");
});
it("hides every disconnect control from a read-only admin", async () => {
vi.mocked(networking.fetchMCPGatewaySessions).mockResolvedValue(REPORT);
renderTab({ canTerminate: false });
await screen.findByRole("region", { name: "Live sessions" });
expect(screen.queryByRole("button", { name: /^Disconnect/ })).not.toBeInTheDocument();
});
it("disconnects one session by its displayed prefix after confirmation and refetches", async () => {
const user = userEvent.setup();
const terminated: MCPGatewaySessionsTerminateResponse = {
worker_pid: 4242,
terminated_sessions: 1,
sessions: [REPORT.sessions[1]],
};
vi.mocked(networking.fetchMCPGatewaySessions).mockResolvedValue(REPORT);
vi.mocked(networking.terminateMCPGatewaySessions).mockResolvedValue(terminated);
renderTab({ canTerminate: true });
await user.click(await screen.findByRole("button", { name: "Disconnect session bbbb2222" }));
expect(networking.terminateMCPGatewaySessions).not.toHaveBeenCalled();
const dialog = await screen.findByRole("alertdialog");
expect(dialog).toHaveTextContent("session bbbb2222");
await user.click(within(dialog).getByRole("button", { name: "Disconnect" }));
const status = await screen.findByText("Disconnected 1 session on worker pid 4242.", { exact: false });
expect(status).toBeInTheDocument();
expect(networking.terminateMCPGatewaySessions).toHaveBeenCalledWith("token", { session_id_prefix: "bbbb2222" });
expect(networking.fetchMCPGatewaySessions).toHaveBeenCalledTimes(2);
});
it("disconnects every session of a user from the by-user table", async () => {
const user = userEvent.setup();
vi.mocked(networking.fetchMCPGatewaySessions).mockResolvedValue(REPORT);
vi.mocked(networking.terminateMCPGatewaySessions).mockResolvedValue({
worker_pid: 4242,
terminated_sessions: 2,
sessions: [REPORT.sessions[0], REPORT.sessions[1]],
});
renderTab({ canTerminate: true });
const byUser = await screen.findByRole("region", { name: "Sessions by user" });
expect(within(byUser).queryByRole("button", { name: /\(unknown\)/ })).not.toBeInTheDocument();
await user.click(within(byUser).getByRole("button", { name: "Disconnect all sessions for user alice" }));
const dialog = await screen.findByRole("alertdialog");
expect(dialog).toHaveTextContent("every live session opened by user alice");
await user.click(within(dialog).getByRole("button", { name: "Disconnect" }));
expect(await screen.findByText(/Disconnected 2 sessions on worker pid 4242\./)).toBeInTheDocument();
expect(networking.terminateMCPGatewaySessions).toHaveBeenCalledWith("token", { user_id: "alice" });
});
it("shows the API error when a disconnect is refused", async () => {
const user = userEvent.setup();
vi.mocked(networking.fetchMCPGatewaySessions).mockResolvedValue(REPORT);
vi.mocked(networking.terminateMCPGatewaySessions).mockRejectedValue(
new Error("Proxy admin access required to terminate MCP gateway sessions."),
);
renderTab({ canTerminate: true });
await user.click(await screen.findByRole("button", { name: "Disconnect session aaaa1111" }));
await user.click(within(await screen.findByRole("alertdialog")).getByRole("button", { name: "Disconnect" }));
const alert = await screen.findByRole("alert");
expect(alert).toHaveTextContent("Could not disconnect");
expect(alert).toHaveTextContent("Proxy admin access required to terminate MCP gateway sessions.");
expect(screen.getByRole("region", { name: "Live sessions" })).toBeInTheDocument();
});
});

View file

@ -1,14 +1,27 @@
"use client";
import React from "react";
import { useQuery } from "@tanstack/react-query";
import { RefreshCw } from "lucide-react";
import React, { useState } from "react";
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { RefreshCw, Unplug } from "lucide-react";
import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
import {
AlertDialog,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { fetchMCPGatewaySessions } from "@/components/networking";
import type { MCPGatewaySessionGroupCount, MCPGatewaySessionsResponse } from "@/components/mcp_tools/types";
import { fetchMCPGatewaySessions, terminateMCPGatewaySessions } from "@/components/networking";
import type {
MCPGatewaySessionGroupCount,
MCPGatewaySessionSelector,
MCPGatewaySessionsResponse,
MCPGatewaySessionsTerminateResponse,
} from "@/components/mcp_tools/types";
import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory";
const mcpGatewaySessionKeys = createQueryKeys("mcpGatewaySessions");
@ -28,6 +41,16 @@ function groupLabel(label: string | null): string {
return label === "" ? '""' : label;
}
export function describeSelector(selector: MCPGatewaySessionSelector): string {
if (selector.user_id !== undefined) return `every live session opened by user ${groupLabel(selector.user_id)}`;
return `session ${selector.session_id_prefix}`;
}
export function describeTerminateResult(result: MCPGatewaySessionsTerminateResponse): string {
const noun = result.terminated_sessions === 1 ? "session" : "sessions";
return `Disconnected ${result.terminated_sessions} ${noun} on worker pid ${result.worker_pid}.`;
}
function StatCard({ label, value }: { label: string; value: number }) {
return (
<div className="bg-card border border-border rounded-lg px-4 py-3">
@ -37,14 +60,37 @@ function StatCard({ label, value }: { label: string; value: number }) {
);
}
function DisconnectUserButton({
userId,
onDisconnectUser,
}: {
userId: string | null;
onDisconnectUser: (userId: string) => void;
}) {
if (userId === null || userId === "") return null;
return (
<Button
variant="outline"
size="sm"
onClick={() => onDisconnectUser(userId)}
aria-label={`Disconnect all sessions for user ${groupLabel(userId)}`}
>
<Unplug className="size-4" />
Disconnect all
</Button>
);
}
function GroupCountTable({
title,
groups,
labelHeader,
onDisconnectUser,
}: {
title: string;
groups: MCPGatewaySessionGroupCount[];
labelHeader: string;
onDisconnectUser?: (userId: string) => void;
}) {
return (
<section aria-label={title} className="rounded-lg border border-border bg-card">
@ -54,6 +100,7 @@ function GroupCountTable({
<TableRow>
<TableHead>{labelHeader}</TableHead>
<TableHead className="text-right">Sessions</TableHead>
{onDisconnectUser ? <TableHead className="text-right">Actions</TableHead> : null}
</TableRow>
</TableHeader>
<TableBody>
@ -61,6 +108,11 @@ function GroupCountTable({
<TableRow key={group.label ?? "__unknown__"}>
<TableCell className="font-mono text-xs">{groupLabel(group.label)}</TableCell>
<TableCell className="text-right">{group.count}</TableCell>
{onDisconnectUser ? (
<TableCell className="text-right">
<DisconnectUserButton userId={group.label} onDisconnectUser={onDisconnectUser} />
</TableCell>
) : null}
</TableRow>
))}
</TableBody>
@ -73,10 +125,12 @@ function SessionsBody({
data,
error,
isLoading,
onDisconnect,
}: {
data: MCPGatewaySessionsResponse | undefined;
error: Error | null;
isLoading: boolean;
onDisconnect: ((selector: MCPGatewaySessionSelector) => void) | null;
}) {
if (isLoading) {
return (
@ -117,7 +171,12 @@ function SessionsBody({
</div>
<div className="grid grid-cols-1 gap-4 lg:grid-cols-2">
<GroupCountTable title="Sessions by AI client" labelHeader="Client" groups={data.by_client} />
<GroupCountTable title="Sessions by user" labelHeader="User" groups={data.by_user} />
<GroupCountTable
title="Sessions by user"
labelHeader="User"
groups={data.by_user}
onDisconnectUser={onDisconnect ? (userId) => onDisconnect({ user_id: userId }) : undefined}
/>
</div>
<section aria-label="Live sessions" className="rounded-lg border border-border bg-card">
<h3 className="border-b border-border px-4 py-2 text-sm font-semibold text-foreground">
@ -134,11 +193,12 @@ function SessionsBody({
<TableHead>Client IP</TableHead>
<TableHead className="text-right">Idle</TableHead>
<TableHead className="text-right">In flight</TableHead>
{onDisconnect ? <TableHead className="text-right">Actions</TableHead> : null}
</TableRow>
</TableHeader>
<TableBody>
{data.sessions.map((session) => (
<TableRow key={session.session_id_prefix}>
{data.sessions.map((session, index) => (
<TableRow key={`${session.session_id_prefix}-${index}`}>
<TableCell className="font-mono text-xs">{session.session_id_prefix}</TableCell>
<TableCell>
{session.client_name === null ? (
@ -169,6 +229,19 @@ function SessionsBody({
<TableCell className="font-mono text-xs">{session.client_ip || "-"}</TableCell>
<TableCell className="text-right text-xs">{formatIdleSeconds(session.idle_seconds)}</TableCell>
<TableCell className="text-right text-xs">{session.in_flight_requests}</TableCell>
{onDisconnect ? (
<TableCell className="text-right">
<Button
variant="outline"
size="sm"
onClick={() => onDisconnect({ session_id_prefix: session.session_id_prefix })}
aria-label={`Disconnect session ${session.session_id_prefix}`}
>
<Unplug className="size-4" />
Disconnect
</Button>
</TableCell>
) : null}
</TableRow>
))}
</TableBody>
@ -180,9 +253,12 @@ function SessionsBody({
interface MCPGatewaySessionsTabProps {
accessToken: string | null;
canTerminate: boolean;
}
export function MCPGatewaySessionsTab({ accessToken }: MCPGatewaySessionsTabProps) {
export function MCPGatewaySessionsTab({ accessToken, canTerminate }: MCPGatewaySessionsTabProps) {
const queryClient = useQueryClient();
const [pendingSelector, setPendingSelector] = useState<MCPGatewaySessionSelector | null>(null);
const queryOptions = {
queryKey: mcpGatewaySessionKeys.lists(),
queryFn: () => fetchMCPGatewaySessions(accessToken!),
@ -190,6 +266,15 @@ export function MCPGatewaySessionsTab({ accessToken }: MCPGatewaySessionsTabProp
refetchInterval: REFETCH_INTERVAL_MS,
};
const { data, error, isLoading, isFetching, refetch } = useQuery<MCPGatewaySessionsResponse, Error>(queryOptions);
const terminate = useMutation<MCPGatewaySessionsTerminateResponse, Error, MCPGatewaySessionSelector>({
mutationFn: (selector) => terminateMCPGatewaySessions(accessToken!, selector),
onSettled: () => queryClient.invalidateQueries({ queryKey: mcpGatewaySessionKeys.lists() }),
});
const confirmDisconnect = () => {
if (pendingSelector === null) return;
terminate.mutate(pendingSelector);
setPendingSelector(null);
};
return (
<div className="mt-4 space-y-4" data-testid="mcp-gateway-sessions-tab">
@ -214,7 +299,48 @@ export function MCPGatewaySessionsTab({ accessToken }: MCPGatewaySessionsTabProp
</Button>
</div>
<SessionsBody data={data} error={error} isLoading={isLoading} />
{terminate.isError ? (
<Alert variant="destructive">
<AlertTitle>Could not disconnect</AlertTitle>
<AlertDescription>{terminate.error.message}</AlertDescription>
</Alert>
) : null}
{terminate.isSuccess ? (
<Alert>
<AlertTitle>Disconnected</AlertTitle>
<AlertDescription>
{describeTerminateResult(terminate.data)} Clients holding those sessions must send a new initialize request,
which re-runs authentication. Sessions on other proxy workers are not affected.
</AlertDescription>
</Alert>
) : null}
<SessionsBody
data={data}
error={error}
isLoading={isLoading}
onDisconnect={canTerminate ? setPendingSelector : null}
/>
<AlertDialog open={pendingSelector !== null} onOpenChange={(open) => !open && setPendingSelector(null)}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Disconnect MCP session</AlertDialogTitle>
<AlertDialogDescription>
{pendingSelector ? `This force-closes ${describeSelector(pendingSelector)} on this proxy worker. ` : ""}
In-flight requests fail and the client must initialize again before it can call tools.
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<Button variant="outline" onClick={() => setPendingSelector(null)}>
Cancel
</Button>
<Button variant="destructive" onClick={confirmDisconnect} disabled={terminate.isPending}>
Disconnect
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}

View file

@ -0,0 +1,102 @@
import React from "react";
import { render, screen, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { describe, it, expect, vi, beforeEach } from "vitest";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { MCPServerUserCredentialsPanel } from "./MCPServerUserCredentialsPanel";
import * as networking from "@/components/networking";
import type { MCPServerUserCredentialListItem } from "@/components/mcp_tools/types";
vi.mock("@/components/networking", () => ({
fetchMCPServerUserCredentials: vi.fn(),
revokeMCPServerUserCredential: vi.fn(),
}));
const ITEMS: MCPServerUserCredentialListItem[] = [
{
user_id: "alice",
credential_type: "oauth2",
expires_at: "2026-12-31T00:00:00+00:00",
connected_at: "2026-01-01T00:00:00+00:00",
updated_at: "2026-01-01T00:00:00+00:00",
},
{
user_id: "carol",
credential_type: "byok",
expires_at: null,
connected_at: null,
updated_at: "2026-02-01T00:00:00+00:00",
},
];
const renderPanel = ({ canRevoke = false }: { canRevoke?: boolean } = {}) => {
const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } });
return render(
<QueryClientProvider client={queryClient}>
<MCPServerUserCredentialsPanel serverId="srv-1" accessToken="token" canRevoke={canRevoke} />
</QueryClientProvider>,
);
};
describe("MCPServerUserCredentialsPanel", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("lists each user's credential type without a revoke control for a read-only admin", async () => {
vi.mocked(networking.fetchMCPServerUserCredentials).mockResolvedValue(ITEMS);
renderPanel({ canRevoke: false });
const table = await screen.findByRole("region", { name: "Stored user credentials" });
expect(within(table).getByRole("row", { name: /alice/ })).toHaveTextContent("OAuth2");
expect(within(table).getByRole("row", { name: /carol/ })).toHaveTextContent("BYOK API key");
expect(screen.queryByRole("button", { name: /^Revoke credential/ })).not.toBeInTheDocument();
expect(networking.fetchMCPServerUserCredentials).toHaveBeenCalledWith("token", "srv-1");
});
it("revokes the selected user's credential through the route for its type and refetches", async () => {
const user = userEvent.setup();
vi.mocked(networking.fetchMCPServerUserCredentials).mockResolvedValueOnce(ITEMS).mockResolvedValueOnce([ITEMS[1]]);
vi.mocked(networking.revokeMCPServerUserCredential).mockResolvedValue(undefined);
renderPanel({ canRevoke: true });
await user.click(await screen.findByRole("button", { name: "Revoke credential for user alice" }));
expect(networking.revokeMCPServerUserCredential).not.toHaveBeenCalled();
const dialog = await screen.findByRole("alertdialog");
expect(dialog).toHaveTextContent("OAuth2 credential stored for user alice");
await user.click(within(dialog).getByRole("button", { name: "Revoke" }));
expect(await screen.findByText(/OAuth2 credential for user alice was deleted/)).toBeInTheDocument();
expect(networking.revokeMCPServerUserCredential).toHaveBeenCalledWith("token", "srv-1", "alice", "oauth2");
const table = await screen.findByRole("region", { name: "Stored user credentials" });
expect(within(table).queryByRole("row", { name: /alice/ })).not.toBeInTheDocument();
expect(within(table).getByRole("row", { name: /carol/ })).toBeInTheDocument();
});
it("shows the API error when a revoke is refused and keeps the list", async () => {
const user = userEvent.setup();
vi.mocked(networking.fetchMCPServerUserCredentials).mockResolvedValue(ITEMS);
vi.mocked(networking.revokeMCPServerUserCredential).mockRejectedValue(
new Error("Proxy admin access required to revoke another user's MCP credential."),
);
renderPanel({ canRevoke: true });
await user.click(await screen.findByRole("button", { name: "Revoke credential for user carol" }));
await user.click(within(await screen.findByRole("alertdialog")).getByRole("button", { name: "Revoke" }));
const alert = await screen.findByRole("alert");
expect(alert).toHaveTextContent("Could not revoke credential");
expect(alert).toHaveTextContent("Proxy admin access required to revoke another user's MCP credential.");
expect(networking.revokeMCPServerUserCredential).toHaveBeenCalledWith("token", "srv-1", "carol", "byok");
expect(screen.getByRole("region", { name: "Stored user credentials" })).toBeInTheDocument();
});
it("shows the API error when the list cannot be loaded", async () => {
vi.mocked(networking.fetchMCPServerUserCredentials).mockRejectedValue(new Error("Admin access required"));
renderPanel();
const alert = await screen.findByRole("alert");
expect(alert).toHaveTextContent("Could not load user credentials");
expect(alert).toHaveTextContent("Admin access required");
});
});

View file

@ -0,0 +1,212 @@
"use client";
import React, { useState } from "react";
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { RefreshCw, ShieldOff } from "lucide-react";
import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert";
import {
AlertDialog,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { fetchMCPServerUserCredentials, revokeMCPServerUserCredential } from "@/components/networking";
import type { MCPServerUserCredentialListItem } from "@/components/mcp_tools/types";
import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory";
const mcpServerUserCredentialKeys = createQueryKeys("mcpServerUserCredentials");
export function credentialTypeLabel(credentialType: MCPServerUserCredentialListItem["credential_type"]): string {
return credentialType === "oauth2" ? "OAuth2" : "BYOK API key";
}
export function formatTimestamp(value: string | null): string {
if (value === null) return "-";
const parsed = new Date(value);
return Number.isNaN(parsed.getTime()) ? value : parsed.toLocaleString();
}
function CredentialsBody({
items,
error,
isLoading,
onRevoke,
}: {
items: MCPServerUserCredentialListItem[] | undefined;
error: Error | null;
isLoading: boolean;
onRevoke: ((item: MCPServerUserCredentialListItem) => void) | null;
}) {
if (isLoading) {
return (
<div
role="status"
className="flex items-center justify-center gap-3 rounded-lg border border-dashed border-border bg-card p-12"
>
<UiLoadingSpinner className="size-6 text-muted-foreground" />
<p className="text-sm text-muted-foreground">Loading user credentials...</p>
</div>
);
}
if (error) {
return (
<Alert variant="destructive">
<AlertTitle>Could not load user credentials</AlertTitle>
<AlertDescription>{error.message}</AlertDescription>
</Alert>
);
}
if (!items) return null;
if (items.length === 0) {
return (
<div className="rounded-lg border border-dashed border-border bg-card p-12 text-center">
<p className="text-sm text-muted-foreground">No user has a stored credential for this server.</p>
</div>
);
}
return (
<section aria-label="Stored user credentials" className="rounded-lg border border-border bg-card">
<Table>
<TableHeader>
<TableRow>
<TableHead>User</TableHead>
<TableHead>Type</TableHead>
<TableHead>Connected</TableHead>
<TableHead>Expires</TableHead>
<TableHead>Updated</TableHead>
{onRevoke ? <TableHead className="text-right">Actions</TableHead> : null}
</TableRow>
</TableHeader>
<TableBody>
{items.map((item) => (
<TableRow key={item.user_id}>
<TableCell className="font-mono text-xs">{item.user_id}</TableCell>
<TableCell>
<Badge variant="secondary">{credentialTypeLabel(item.credential_type)}</Badge>
</TableCell>
<TableCell className="text-xs">{formatTimestamp(item.connected_at)}</TableCell>
<TableCell className="text-xs">{formatTimestamp(item.expires_at)}</TableCell>
<TableCell className="text-xs">{formatTimestamp(item.updated_at)}</TableCell>
{onRevoke ? (
<TableCell className="text-right">
<Button
variant="outline"
size="sm"
onClick={() => onRevoke(item)}
aria-label={`Revoke credential for user ${item.user_id}`}
>
<ShieldOff className="size-4" />
Revoke
</Button>
</TableCell>
) : null}
</TableRow>
))}
</TableBody>
</Table>
</section>
);
}
interface MCPServerUserCredentialsPanelProps {
serverId: string;
accessToken: string | null;
canRevoke: boolean;
}
export function MCPServerUserCredentialsPanel({
serverId,
accessToken,
canRevoke,
}: MCPServerUserCredentialsPanelProps) {
const queryClient = useQueryClient();
const [pendingItem, setPendingItem] = useState<MCPServerUserCredentialListItem | null>(null);
const queryKey = mcpServerUserCredentialKeys.detail(serverId);
const { data, error, isLoading, isFetching, refetch } = useQuery<MCPServerUserCredentialListItem[], Error>({
queryKey,
queryFn: () => fetchMCPServerUserCredentials(accessToken!, serverId),
enabled: !!accessToken,
});
const revoke = useMutation<void, Error, MCPServerUserCredentialListItem>({
mutationFn: (item) => revokeMCPServerUserCredential(accessToken!, serverId, item.user_id, item.credential_type),
onSettled: () => queryClient.invalidateQueries({ queryKey }),
});
const confirmRevoke = () => {
if (pendingItem === null) return;
revoke.mutate(pendingItem);
setPendingItem(null);
};
return (
<div className="space-y-4" data-testid="mcp-server-user-credentials-panel">
<div className="flex flex-wrap items-start justify-between gap-3">
<div>
<h2 className="text-lg font-medium">User Credentials</h2>
<p className="text-sm text-muted-foreground">
Per-user OAuth2 tokens and BYOK API keys stored for this server. Revoking one deletes it from the database
and clears the cached copy, so the user must connect again before the gateway will call this server for
them.
</p>
</div>
<Button
variant="outline"
size="sm"
onClick={() => refetch()}
disabled={isFetching}
aria-label="Refresh user credentials"
>
<RefreshCw className={`size-4 ${isFetching ? "animate-spin" : ""}`} />
Refresh
</Button>
</div>
{revoke.isError ? (
<Alert variant="destructive">
<AlertTitle>Could not revoke credential</AlertTitle>
<AlertDescription>{revoke.error.message}</AlertDescription>
</Alert>
) : null}
{revoke.isSuccess ? (
<Alert>
<AlertTitle>Credential revoked</AlertTitle>
<AlertDescription>
The stored {credentialTypeLabel(revoke.variables.credential_type)} credential for user{" "}
{revoke.variables.user_id} was deleted.
</AlertDescription>
</Alert>
) : null}
<CredentialsBody items={data} error={error} isLoading={isLoading} onRevoke={canRevoke ? setPendingItem : null} />
<AlertDialog open={pendingItem !== null} onOpenChange={(open) => !open && setPendingItem(null)}>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Revoke stored credential</AlertDialogTitle>
<AlertDialogDescription>
{pendingItem
? `This deletes the ${credentialTypeLabel(pendingItem.credential_type)} credential stored for user ${pendingItem.user_id}. `
: ""}
Their next MCP request to this server fails until they connect again.
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<Button variant="outline" onClick={() => setPendingItem(null)}>
Cancel
</Button>
<Button variant="destructive" onClick={confirmRevoke} disabled={revoke.isPending}>
Revoke
</Button>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}
export default MCPServerUserCredentialsPanel;

View file

@ -1,7 +1,9 @@
import { render, screen } from "@testing-library/react";
import { render, screen, within } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { describe, it, expect, vi, beforeEach } from "vitest";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { MCPServerView } from "./mcp_server_view";
import * as networking from "@/components/networking";
import type { MCPServer } from "@/components/mcp_tools/types";
vi.mock(".", () => ({
@ -13,6 +15,12 @@ vi.mock("./mcp_server_edit", () => ({
EDIT_OAUTH_UI_STATE_KEY: "litellm-mcp-oauth-edit-state",
}));
vi.mock("@/components/networking", async (importOriginal) => ({
...(await importOriginal<typeof import("@/components/networking")>()),
fetchMCPServerUserCredentials: vi.fn(),
revokeMCPServerUserCredential: vi.fn(),
}));
const baseServer = {
server_id: "srv-1",
server_name: "demo server",
@ -25,19 +33,38 @@ const baseServer = {
const renderView = (overrides: Partial<MCPServer> = {}, props: Record<string, unknown> = {}) =>
render(
<MCPServerView
mcpServer={{ ...baseServer, ...overrides } as MCPServer}
onBack={vi.fn()}
isProxyAdmin
isEditing={false}
accessToken="tok"
userRole="Admin"
userID="u1"
availableAccessGroups={[]}
{...props}
/>,
<QueryClientProvider client={new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } })}>
<MCPServerView
mcpServer={{ ...baseServer, ...overrides } as MCPServer}
onBack={vi.fn()}
isProxyAdmin
isEditing={false}
accessToken="tok"
userRole="Admin"
userID="u1"
availableAccessGroups={[]}
{...props}
/>
</QueryClientProvider>,
);
const openUserCredentials = async (props: Record<string, unknown>) => {
vi.mocked(networking.fetchMCPServerUserCredentials).mockResolvedValue([
{
user_id: "alice",
credential_type: "byok",
expires_at: null,
connected_at: null,
updated_at: "2026-01-01T00:00:00+00:00",
},
]);
renderView({}, props);
await userEvent.click(screen.getByRole("tab", { name: "User Credentials" }));
return within(await screen.findByRole("region", { name: "Stored user credentials" })).getByRole("row", {
name: /alice/,
});
};
describe("MCPServerView", () => {
beforeEach(() => {
vi.clearAllMocks();
@ -149,4 +176,15 @@ describe("MCPServerView", () => {
expect(await screen.findByText("All tools enabled")).toBeInTheDocument();
});
it("lets a full admin revoke a stored user credential", async () => {
const row = await openUserCredentials({});
expect(within(row).getByRole("button", { name: "Revoke credential for user alice" })).toBeInTheDocument();
});
it("shows stored credentials to a view-only admin session without a revoke control", async () => {
const row = await openUserCredentials({ isViewOnly: true });
expect(row).toHaveTextContent("BYOK API key");
expect(within(row).queryByRole("button", { name: /^Revoke credential/ })).not.toBeInTheDocument();
});
});

View file

@ -9,7 +9,9 @@ import { MCPServer, handleTransport, handleAuth } from "@/components/mcp_tools/t
// TODO: Move Tools viewer from index file
import { MCPToolsViewer } from ".";
import MCPServerEdit, { EDIT_OAUTH_UI_STATE_KEY } from "./mcp_server_edit";
import { MCPServerUserCredentialsPanel } from "./MCPServerUserCredentialsPanel";
import { getSecureItem } from "@/utils/secureStorage";
import { isProxyAdminRole, isProxyAdminTierRole } from "@/utils/roles";
import MCPServerCostDisplay from "./mcp_server_cost_display";
import { getMaskedAndFullUrl } from "./utils";
import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils";
@ -23,6 +25,7 @@ interface MCPServerViewProps {
accessToken: string | null;
userRole: string | null;
userID: string | null;
isViewOnly?: boolean;
availableAccessGroups: string[];
initialTabIndex?: number;
}
@ -53,6 +56,7 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
accessToken,
userRole,
userID,
isViewOnly = false,
availableAccessGroups,
initialTabIndex = 0,
}) => {
@ -63,6 +67,8 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
const [showFullUrl, setShowFullUrl] = useState(false);
const [copiedStates, setCopiedStates] = useState<Record<string, boolean>>({});
const [selectedTabIndex, setSelectedTabIndex] = useState(returningFromEditOAuth ? 2 : initialTabIndex);
const canViewUserCredentials = userRole !== null && isProxyAdminTierRole(userRole);
const canRevokeUserCredentials = userRole !== null && isProxyAdminRole(userRole) && !isViewOnly;
const handleSuccess = (updated: MCPServer) => {
setEditing(false);
@ -142,6 +148,11 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
Settings
</TabsTrigger>
)}
{canViewUserCredentials && (
<TabsTrigger value="3" className="flex-none rounded-none px-4 py-2">
User Credentials
</TabsTrigger>
)}
</TabsList>
{/* Overview Panel */}
@ -387,6 +398,18 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
)}
</Card>
</TabsContent>
{canViewUserCredentials && (
<TabsContent value="3">
<Card className="p-6">
<MCPServerUserCredentialsPanel
serverId={mcpServer.server_id}
accessToken={accessToken}
canRevoke={canRevokeUserCredentials}
/>
</Card>
</TabsContent>
)}
</Tabs>
</div>
);

View file

@ -17,6 +17,8 @@ vi.mock("@/components/networking", () => ({
updateConfigFieldSetting: vi.fn().mockResolvedValue(undefined),
deleteConfigFieldSetting: vi.fn().mockResolvedValue(undefined),
listMCPUserEnvVarStatus: vi.fn().mockResolvedValue([]),
fetchMCPGatewaySessions: vi.fn(),
terminateMCPGatewaySessions: vi.fn(),
}));
const createQueryClient = () =>
@ -400,4 +402,50 @@ describe("MCPServers", () => {
// The server list refresh must NOT trigger a second health check
expect(networking.fetchMCPServerHealth).toHaveBeenCalledTimes(1);
});
const liveSessionsReport = {
worker_pid: 4242,
total_sessions: 1,
by_client: [{ label: "claude-code", count: 1 }],
by_user: [{ label: "alice", count: 1 }],
sessions: [
{
session_id_prefix: "aaaa1111",
client_name: "claude-code",
client_version: "1.0.0",
user_id: "alice",
user_email: "alice@example.com",
key_alias: "alice-key",
team_id: null,
team_alias: null,
client_ip: "10.0.0.1",
idle_seconds: 5,
in_flight_requests: 0,
},
],
};
const openLiveConnections = async (props: { isViewOnly?: boolean }) => {
vi.mocked(networking.fetchMCPServers).mockResolvedValue([]);
vi.mocked(networking.fetchMCPGatewaySessions).mockResolvedValue(liveSessionsReport);
render(
<QueryClientProvider client={createQueryClient()}>
<MCPServers {...defaultProps} {...props} />
</QueryClientProvider>,
);
await userEvent.click(await screen.findByRole("tab", { name: "Live Connections" }));
return within(await screen.findByRole("region", { name: "Live sessions" })).getByRole("row", { name: /aaaa1111/ });
};
it("lets a full admin disconnect a live session", async () => {
const row = await openLiveConnections({ isViewOnly: false });
expect(within(row).getByRole("button", { name: "Disconnect session aaaa1111" })).toBeInTheDocument();
});
it("shows live sessions to a view-only admin session without any disconnect control", async () => {
const row = await openLiveConnections({ isViewOnly: true });
expect(row).toHaveTextContent("alice@example.com");
expect(within(row).queryByRole("button", { name: /^Disconnect/ })).not.toBeInTheDocument();
expect(screen.queryByRole("button", { name: /^Disconnect all/ })).not.toBeInTheDocument();
});
});

View file

@ -1,4 +1,4 @@
import { isAdminRole, isProxyAdminTierRole } from "@/utils/roles";
import { isAdminRole, isProxyAdminRole, isProxyAdminTierRole } from "@/utils/roles";
import { CircleHelp, Search } from "lucide-react";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
@ -109,7 +109,7 @@ const readToolsOAuthServerId = (): string | null => {
}
};
const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID }) => {
const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID, isViewOnly = false }) => {
const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers();
// Fetch health status for all servers
@ -578,6 +578,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
accessToken={accessToken}
userID={userID}
userRole={userRole}
isViewOnly={isViewOnly}
availableAccessGroups={uniqueMcpAccessGroups}
initialTabIndex={selectedServerId === toolsTabServerId ? 1 : 0}
/>
@ -755,7 +756,10 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
)}
{isProxyAdminTierRole(userRole) && (
<TabsContent value="connections">
<MCPGatewaySessionsTab accessToken={accessToken} />
<MCPGatewaySessionsTab
accessToken={accessToken}
canTerminate={isProxyAdminRole(userRole) && !isViewOnly}
/>
</TabsContent>
)}
</Tabs>

View file

@ -4,6 +4,6 @@ import { MCPServers } from "./_components";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
export default function McpServers() {
const { accessToken, userRole, userId } = useAuthorized();
return <MCPServers accessToken={accessToken} userRole={userRole} userID={userId} />;
const { accessToken, userRole, userId, isViewOnly } = useAuthorized();
return <MCPServers accessToken={accessToken} userRole={userRole} userID={userId} isViewOnly={isViewOnly} />;
}

View file

@ -517,6 +517,7 @@ export interface MCPServerProps {
accessToken: string | null;
userRole: string | null;
userID: string | null;
isViewOnly?: boolean;
}
export interface MCPToolsetTool {
@ -587,3 +588,23 @@ export interface MCPGatewaySessionsResponse {
by_user: MCPGatewaySessionGroupCount[];
sessions: MCPGatewaySession[];
}
export interface MCPGatewaySessionsTerminateResponse {
worker_pid: number;
terminated_sessions: number;
sessions: MCPGatewaySession[];
}
export type MCPGatewaySessionSelector =
| { session_id_prefix: string; user_id?: undefined }
| { user_id: string; session_id_prefix?: undefined };
export type MCPServerUserCredentialType = "oauth2" | "byok";
export interface MCPServerUserCredentialListItem {
user_id: string;
credential_type: MCPServerUserCredentialType;
expires_at: string | null;
connected_at: string | null;
updated_at: string;
}

View file

@ -97,7 +97,14 @@ import type { ModelBudgetUsage, ModelMaxBudget } from "./key_team_helpers/ModelM
import type { ObjectPermission } from "./object_permission_types";
import type { components } from "@/lib/http/schema";
import { jsonFields } from "./common_components/check_openapi_schema";
import type { MCPGatewaySessionsResponse, MCPUserEnvVarsStatus } from "./mcp_tools/types";
import type {
MCPGatewaySessionSelector,
MCPGatewaySessionsResponse,
MCPGatewaySessionsTerminateResponse,
MCPServerUserCredentialListItem,
MCPServerUserCredentialType,
MCPUserEnvVarsStatus,
} from "./mcp_tools/types";
import type {
CoordinationRedisSettings,
CoordinationRedisSettingsResponse,
@ -4976,6 +4983,33 @@ export const fetchMCPSubmissions = async (accessToken: string) => {
export const fetchMCPGatewaySessions = async (accessToken: string): Promise<MCPGatewaySessionsResponse> =>
apiClient.get<MCPGatewaySessionsResponse>(`/v1/mcp/sessions`, { accessToken });
export const terminateMCPGatewaySessions = async (
accessToken: string,
selector: MCPGatewaySessionSelector,
): Promise<MCPGatewaySessionsTerminateResponse> =>
apiClient.delete<MCPGatewaySessionsTerminateResponse>(`/v1/mcp/sessions`, { accessToken, query: { ...selector } });
export const fetchMCPServerUserCredentials = async (
accessToken: string,
serverId: string,
): Promise<MCPServerUserCredentialListItem[]> =>
apiClient.get<MCPServerUserCredentialListItem[]>(`/v1/mcp/server/${encodeURIComponent(serverId)}/user-credentials`, {
accessToken,
});
export const revokeMCPServerUserCredential = async (
accessToken: string,
serverId: string,
userId: string,
credentialType: MCPServerUserCredentialType,
): Promise<void> => {
const route = credentialType === "oauth2" ? "oauth-user-credential" : "user-credential";
await apiClient.delete(`/v1/mcp/server/${encodeURIComponent(serverId)}/${route}`, {
accessToken,
query: { user_id: userId },
});
};
export const approveMCPServer = async (accessToken: string, serverId: string) => {
try {
const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/server/${encodeURIComponent(serverId)}/approve`;

View file

@ -19037,7 +19037,7 @@ export interface paths {
post: operations["store_mcp_oauth_user_credential_v1_mcp_server__server_id__oauth_user_credential_post"];
/**
* Delete Mcp Oauth User Credential
* @description Revoke the calling user's stored OAuth2 token for an MCP server
* @description Revoke the calling user's stored OAuth2 token for an MCP server. A proxy admin may pass user_id to revoke another user's stored token.
*/
delete: operations["delete_mcp_oauth_user_credential_v1_mcp_server__server_id__oauth_user_credential_delete"];
options?: never;
@ -19101,7 +19101,7 @@ export interface paths {
post: operations["store_mcp_user_credential_v1_mcp_server__server_id__user_credential_post"];
/**
* Delete Mcp User Credential
* @description Delete the calling user's stored API key for a BYOK MCP server
* @description Delete the calling user's stored API key for a BYOK MCP server. A proxy admin may pass user_id to revoke another user's stored key.
*/
delete: operations["delete_mcp_user_credential_v1_mcp_server__server_id__user_credential_delete"];
options?: never;
@ -19109,6 +19109,26 @@ export interface paths {
patch?: never;
trace?: never;
};
"/v1/mcp/server/{server_id}/user-credentials": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
/**
* List Mcp Server User Credentials
* @description List every user's stored BYOK or OAuth2 credential for an MCP server (admin only, no secrets)
*/
get: operations["list_mcp_server_user_credentials_v1_mcp_server__server_id__user_credentials_get"];
put?: never;
post?: never;
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/v1/mcp/server/{server_id}/user-env-vars": {
parameters: {
query?: never;
@ -19151,7 +19171,11 @@ export interface paths {
get: operations["get_mcp_gateway_sessions_v1_mcp_sessions_get"];
put?: never;
post?: never;
delete?: never;
/**
* Delete Mcp Gateway Sessions
* @description Force-close live stateful MCP gateway sessions on this proxy worker, selected by session id prefix and/or by the LiteLLM user that opened them (proxy admin only).
*/
delete: operations["delete_mcp_gateway_sessions_v1_mcp_sessions_delete"];
options?: never;
head?: never;
patch?: never;
@ -32653,6 +32677,18 @@ export interface components {
/** Worker Pid */
worker_pid: number;
};
/**
* MCPGatewaySessionsTerminateResponse
* @description Stateful sessions an administrator force-closed on this proxy worker.
*/
MCPGatewaySessionsTerminateResponse: {
/** Sessions */
sessions?: components["schemas"]["MCPGatewaySession"][];
/** Terminated Sessions */
terminated_sessions: number;
/** Worker Pid */
worker_pid: number;
};
/**
* MCPOAuthUserCredentialRequest
* @description Stores a user's OAuth2 token for an OpenAPI MCP server.
@ -32757,6 +32793,25 @@ export interface components {
[key: string]: unknown;
};
};
/**
* MCPServerUserCredentialListItem
* @description One user's stored credential for an MCP server, as an admin sees it. Never carries the secret.
*/
MCPServerUserCredentialListItem: {
/** Connected At */
connected_at?: string | null;
/**
* Credential Type
* @enum {string}
*/
credential_type: "oauth2" | "byok";
/** Expires At */
expires_at?: string | null;
/** Updated At */
updated_at: string;
/** User Id */
user_id: string;
};
/** MCPSubmissionsSummary */
MCPSubmissionsSummary: {
/** Active */
@ -65725,7 +65780,9 @@ export interface operations {
};
delete_mcp_oauth_user_credential_v1_mcp_server__server_id__oauth_user_credential_delete: {
parameters: {
query?: never;
query?: {
user_id?: string | null;
};
header?: never;
path: {
server_id: string;
@ -65857,7 +65914,9 @@ export interface operations {
};
delete_mcp_user_credential_v1_mcp_server__server_id__user_credential_delete: {
parameters: {
query?: never;
query?: {
user_id?: string | null;
};
header?: never;
path: {
server_id: string;
@ -65886,6 +65945,37 @@ export interface operations {
};
};
};
list_mcp_server_user_credentials_v1_mcp_server__server_id__user_credentials_get: {
parameters: {
query?: never;
header?: never;
path: {
server_id: string;
};
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["MCPServerUserCredentialListItem"][];
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
get_mcp_user_env_vars_v1_mcp_server__server_id__user_env_vars_get: {
parameters: {
query?: never;
@ -66003,6 +66093,38 @@ export interface operations {
};
};
};
delete_mcp_gateway_sessions_v1_mcp_sessions_delete: {
parameters: {
query?: {
session_id_prefix?: string | null;
user_id?: string | null;
};
header?: never;
path?: never;
cookie?: never;
};
requestBody?: never;
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["MCPGatewaySessionsTerminateResponse"];
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
get_mcp_tools_v1_mcp_tools_get: {
parameters: {
query?: never;