mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
feat(mcp): let proxy admins force-close live MCP sessions and revoke stored user credentials
Adds an admin-only DELETE /v1/mcp/sessions that terminates stateful MCP gateway sessions on the current worker by session id prefix and/or by the LiteLLM user that opened them, tombstones the terminated ids so a client reusing one gets 404 instead of a silently recreated stateless session, and lets PROXY_ADMIN name a user_id on the BYOK and OAuth credential delete routes. Full and view-only admins can list every user's stored credential metadata for a server (never the secret). The dashboard gains Disconnect controls on the Live Connections tab and a User Credentials tab with Revoke controls, both hidden from read-only admins. Resolves LIT-8001 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
72ef7033c2
commit
ea37596b88
19 changed files with 1826 additions and 49 deletions
|
|
@ -183,6 +183,8 @@ 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_ADMIN_TERMINATED_SESSION_IDS_MAX: Final = 1024
|
||||
|
||||
# Allowlist of commands permitted for MCP stdio transport.
|
||||
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import time
|
|||
import traceback
|
||||
import types
|
||||
import uuid
|
||||
from collections import Counter
|
||||
from collections import Counter, deque
|
||||
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol
|
||||
|
|
@ -28,7 +28,11 @@ 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_ADMIN_TERMINATED_SESSION_IDS_MAX,
|
||||
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,
|
||||
|
|
@ -91,6 +95,7 @@ from litellm.types.mcp import (
|
|||
MCPGatewaySession,
|
||||
MCPGatewaySessionGroupCount,
|
||||
MCPGatewaySessionsResponse,
|
||||
MCPGatewaySessionsTerminateResponse,
|
||||
MCPSpecVersion,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
|
||||
|
|
@ -618,6 +623,9 @@ 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[deque[str]] = deque( # mutable-ok: bounded ring, appended on admin termination
|
||||
maxlen=MCP_ADMIN_TERMINATED_SESSION_IDS_MAX
|
||||
)
|
||||
|
||||
class _TerminableTransport(Protocol):
|
||||
async def terminate(self) -> None: ...
|
||||
|
|
@ -3850,7 +3858,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 +3893,53 @@ 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
|
||||
|
||||
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()
|
||||
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.append(session_id)
|
||||
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 +4064,17 @@ if MCP_AVAILABLE:
|
|||
await success_response(scope, receive, send)
|
||||
return True
|
||||
|
||||
if _session_id in _admin_terminated_session_ids:
|
||||
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. "
|
||||
|
|
|
|||
|
|
@ -27648,6 +27648,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": {
|
||||
|
|
@ -27744,6 +27770,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": {
|
||||
|
|
@ -29920,7 +29996,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": [
|
||||
{
|
||||
|
|
@ -29931,6 +30007,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": {
|
||||
|
|
@ -30130,7 +30223,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": [
|
||||
{
|
||||
|
|
@ -30141,6 +30234,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": {
|
||||
|
|
@ -30232,6 +30342,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.",
|
||||
|
|
@ -30383,6 +30545,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",
|
||||
|
|
|
|||
|
|
@ -1716,6 +1716,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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
@ -2261,7 +2324,10 @@ if MCP_AVAILABLE:
|
|||
|
||||
@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)
|
||||
_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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -2871,6 +2871,183 @@ 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 litellm.proxy._experimental.mcp_server import server as mcp_server
|
||||
from litellm.proxy._experimental.mcp_server.server import session_manager_stateful
|
||||
from starlette.types import Scope
|
||||
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_initialize_request_with_existing_session_tracks_new_session():
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -5136,6 +5136,261 @@ 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 = MagicMock()
|
||||
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_called_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,
|
||||
)
|
||||
|
||||
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,
|
||||
),
|
||||
patch.object( # test-quality-ok: the cache invalidator is module scoped; the suite's only seam
|
||||
mcp_server, "_invalidate_byok_cred_cache", new=MagicMock()
|
||||
),
|
||||
):
|
||||
await delete_mcp_user_credential(
|
||||
server_id="srv-byok-self",
|
||||
user_api_key_dict=_make_user_auth("user-self"),
|
||||
user_id="user-self",
|
||||
)
|
||||
|
||||
assert delete_mock.await_args.args[1:] == ("user-self", "srv-byok-self")
|
||||
|
||||
|
||||
@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 +7576,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()
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
@ -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;
|
||||
|
|
@ -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";
|
||||
|
|
@ -63,6 +65,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);
|
||||
|
||||
const handleSuccess = (updated: MCPServer) => {
|
||||
setEditing(false);
|
||||
|
|
@ -142,6 +146,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 +396,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>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
@ -755,7 +755,7 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
|
|||
)}
|
||||
{isProxyAdminTierRole(userRole) && (
|
||||
<TabsContent value="connections">
|
||||
<MCPGatewaySessionsTab accessToken={accessToken} />
|
||||
<MCPGatewaySessionsTab accessToken={accessToken} canTerminate={isProxyAdminRole(userRole)} />
|
||||
</TabsContent>
|
||||
)}
|
||||
</Tabs>
|
||||
|
|
|
|||
|
|
@ -587,3 +587,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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -5112,6 +5119,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`;
|
||||
|
|
|
|||
132
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
132
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -18902,7 +18902,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;
|
||||
|
|
@ -18966,7 +18966,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;
|
||||
|
|
@ -18974,6 +18974,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;
|
||||
|
|
@ -19016,7 +19036,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;
|
||||
|
|
@ -32383,6 +32407,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.
|
||||
|
|
@ -32487,6 +32523,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 */
|
||||
|
|
@ -65109,7 +65164,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;
|
||||
|
|
@ -65241,7 +65298,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;
|
||||
|
|
@ -65270,6 +65329,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;
|
||||
|
|
@ -65387,6 +65477,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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue