diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index ad886c66de7..9c8ad2f4613 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -9,6 +9,7 @@ import contextlib import contextvars import hashlib import json +import os import time import traceback import types @@ -84,7 +85,13 @@ from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, get_chain_id_from_headers, ) -from litellm.types.mcp import MCPAuth, MCPSpecVersion +from litellm.types.mcp import ( + MCPAuth, + MCPGatewaySession, + MCPGatewaySessionGroupCount, + MCPGatewaySessionsResponse, + MCPSpecVersion, +) from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall from litellm.utils import Rules, client, function_setup @@ -454,6 +461,8 @@ if MCP_AVAILABLE: StreamableHTTPSessionManager = None from mcp.types import ( CallToolResult, + Implementation, + InitializeRequest, ListToolsResult, Prompt, TextContent, @@ -607,6 +616,7 @@ if MCP_AVAILABLE: # still reading the shared object. _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 class _TerminableTransport(Protocol): async def terminate(self) -> None: ... @@ -625,6 +635,7 @@ if MCP_AVAILABLE: _stateful_session_owners.pop(session_id, None) _stateful_session_locks.pop(session_id, None) _stateful_session_active_request_counts.pop(session_id, None) + _stateful_session_client_info.pop(session_id, None) # Keep this alias so existing references to session_manager still work session_manager: Final = session_manager_stateless @@ -3816,6 +3827,63 @@ if MCP_AVAILABLE: except (json.JSONDecodeError, TypeError): return False + def _extract_initialize_client_info(body: bytes) -> Implementation | None: + try: + return InitializeRequest.model_validate_json(body).params.clientInfo + except ValidationError: + return None + + def _group_session_counts( + sessions: Sequence[MCPGatewaySession], + label_for: Callable[[MCPGatewaySession], str | None], + ) -> tuple[MCPGatewaySessionGroupCount, ...]: + labels: Final = tuple(label_for(session) for session in sessions) + return tuple( + sorted( + (MCPGatewaySessionGroupCount(label=label, count=labels.count(label)) for label in frozenset(labels)), + key=lambda group: (-group.count, group.label is None, group.label or ""), + ) + ) + + def _gateway_session_for(session_id: str, auth_user: MCPAuthenticatedUser, now: float) -> MCPGatewaySession: + 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], + 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, + user_email=key_auth.user_email if key_auth is not None else None, + key_alias=key_auth.key_alias if key_auth is not None else None, + team_id=key_auth.team_id if key_auth is not None else None, + team_alias=key_auth.team_alias if key_auth is not None else None, + client_ip=auth_user.client_ip, + idle_seconds=max(0.0, now - _stateful_session_auth_context_last_seen.get(session_id, now)), + in_flight_requests=_stateful_session_active_request_counts.get(session_id, 0), + ) + + def get_mcp_gateway_sessions_report(now: float | None = None) -> MCPGatewaySessionsResponse: + """Live stateful Streamable HTTP sessions held by this worker process. + + Only sessions whose transport is still registered with the stateful + session manager are reported; SSE and stateless requests hold no + session and are never counted. + """ + report_time: Final = time.monotonic() if now is None else now + live_session_ids: Final = frozenset(_stateful_server_instances()) + sessions: Final = tuple( + _gateway_session_for(session_id, auth_user, report_time) + for session_id, auth_user in tuple(_stateful_session_auth_contexts.items()) + if session_id in live_session_ids + ) + return MCPGatewaySessionsResponse( + worker_pid=os.getpid(), + total_sessions=len(sessions), + by_client=_group_session_counts(sessions, lambda session: session.client_name), + by_user=_group_session_counts(sessions, lambda session: session.user_id), + sessions=sessions, + ) + async def _read_request_body_for_routing( receive: Receive, ) -> tuple[list[Message], bytes]: @@ -4652,6 +4720,7 @@ if MCP_AVAILABLE: auth_user, _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip), _track_initialized_stateful_session, + client_info=_extract_initialize_client_info(body), ) async with _gateway_initialize_instructions_request_scope( @@ -4965,6 +5034,7 @@ if MCP_AVAILABLE: auth_user: MCPAuthenticatedUser, owner_fingerprint: str, on_session_registered: Callable[[str], None] | None = None, + client_info: Implementation | None = None, ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": @@ -4979,6 +5049,8 @@ if MCP_AVAILABLE: _stateful_session_auth_contexts[session_id] = auth_user _stateful_session_auth_context_last_seen[session_id] = time.monotonic() _stateful_session_owners[session_id] = owner_fingerprint + if client_info is not None: + _stateful_session_client_info[session_id] = client_info break await send(message) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 216b9f6def6..82b709feb92 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -27473,6 +27473,181 @@ "title": "MCPEnvVarScope", "type": "string" }, + "MCPGatewaySession": { + "description": "One live stateful Streamable HTTP session held by this proxy worker.", + "properties": { + "client_ip": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Ip" + }, + "client_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Name" + }, + "client_version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Client Version" + }, + "idle_seconds": { + "title": "Idle Seconds", + "type": "number" + }, + "in_flight_requests": { + "title": "In Flight Requests", + "type": "integer" + }, + "key_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Alias" + }, + "session_id_prefix": { + "title": "Session Id Prefix", + "type": "string" + }, + "team_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Alias" + }, + "team_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" + }, + "user_email": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Email" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" + } + }, + "required": [ + "session_id_prefix", + "idle_seconds", + "in_flight_requests" + ], + "title": "MCPGatewaySession", + "type": "object" + }, + "MCPGatewaySessionGroupCount": { + "properties": { + "count": { + "title": "Count", + "type": "integer" + }, + "label": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Label" + } + }, + "required": [ + "count" + ], + "title": "MCPGatewaySessionGroupCount", + "type": "object" + }, + "MCPGatewaySessionsResponse": { + "properties": { + "by_client": { + "items": { + "$ref": "#/components/schemas/MCPGatewaySessionGroupCount" + }, + "title": "By Client", + "type": "array" + }, + "by_user": { + "items": { + "$ref": "#/components/schemas/MCPGatewaySessionGroupCount" + }, + "title": "By User", + "type": "array" + }, + "sessions": { + "items": { + "$ref": "#/components/schemas/MCPGatewaySession" + }, + "title": "Sessions", + "type": "array" + }, + "total_sessions": { + "title": "Total Sessions", + "type": "integer" + }, + "worker_pid": { + "title": "Worker Pid", + "type": "integer" + } + }, + "required": [ + "worker_pid", + "total_sessions" + ], + "title": "MCPGatewaySessionsResponse", + "type": "object" + }, "MCPOAuthUserCredentialRequest": { "description": "Stores a user's OAuth2 token for an OpenAPI MCP server.", "properties": { @@ -30207,6 +30382,33 @@ ] } }, + "/v1/mcp/sessions": { + "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", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/MCPGatewaySessionsResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Mcp Gateway Sessions", + "tags": [ + "mcp_management" + ] + } + }, "/v1/mcp/tools": { "get": { "description": "Get all MCP tools available for the current key, including those from access groups", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 680c63393e8..ff32d4784df 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -531,6 +531,7 @@ class LiteLLMRoutes(enum.Enum): mcp_management_routes = [ "/v1/mcp/server", "/v1/mcp/server/{path:path}", + "/v1/mcp/sessions", ] # Backwards-compat union — virtual keys may be configured with diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 6fa91c16eb2..82a1cdcdd00 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -22,7 +22,7 @@ import os from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Final, Literal, Protocol +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol from fastapi import ( APIRouter, @@ -220,6 +220,7 @@ if MCP_AVAILABLE: MCP_ADMIN_CONFIG_CREDENTIAL_KEYS, MCPAuth, MCPCredentials, + MCPGatewaySessionsResponse, normalize_upstream_header_name, ) from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -1346,6 +1347,32 @@ if MCP_AVAILABLE: # Do NOT add to runtime registry — pending servers are not active return _redact_mcp_credentials(new_mcp_server) + @router.get( + "/sessions", + description="Live stateful MCP gateway sessions on this proxy worker, grouped by AI client and by user.", + dependencies=(Depends(user_api_key_auth),), + response_model=MCPGatewaySessionsResponse, + ) + @management_endpoint_wrapper + async def get_mcp_gateway_sessions( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + ) -> MCPGatewaySessionsResponse: + 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: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape + "error": "Admin access required to view MCP gateway sessions." + }, + ) + from litellm.proxy._experimental.mcp_server.server import ( + get_mcp_gateway_sessions_report, + ) + + return get_mcp_gateway_sessions_report() + @router.get( "/server/submissions", description="Returns all MCP servers submitted by non-admin users (admin review queue). Mirrors GET /guardrails/submissions.", diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index a59fcb1bcb5..2d06bb9a009 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -435,3 +435,32 @@ class MCPPostCallResponseObject(BaseModel): mcp_tool_call_response: list[MCPTextContent | MCPImageContent | MCPEmbeddedResource] hidden_params: HiddenParams + + +class MCPGatewaySession(BaseModel): + """One live stateful Streamable HTTP session held by this proxy worker.""" + + session_id_prefix: str + client_name: str | None = None + client_version: str | None = None + user_id: str | None = None + user_email: str | None = None + key_alias: str | None = None + team_id: str | None = None + team_alias: str | None = None + client_ip: str | None = None + idle_seconds: float + in_flight_requests: int + + +class MCPGatewaySessionGroupCount(BaseModel): + label: str | None = None + count: int + + +class MCPGatewaySessionsResponse(BaseModel): + worker_pid: int + total_sessions: int + by_client: list[MCPGatewaySessionGroupCount] = Field(default_factory=list) + by_user: list[MCPGatewaySessionGroupCount] = Field(default_factory=list) + sessions: list[MCPGatewaySession] = Field(default_factory=list) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 02182ebbe60..bbc36991e21 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2609,6 +2609,246 @@ async def test_initialize_request_tracks_active_session_after_response_header(): mcp_server._remove_stateful_session_tracking(session_id) +_INITIALIZE_WITH_CLIENT_INFO: Final = ( + b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18",' + b'"capabilities":{},"clientInfo":{"name":"claude-code","version":"1.0.0"}}}' +) + + +@pytest.mark.parametrize( + ("body", "expected_name", "expected_version"), + [ + (_INITIALIZE_WITH_CLIENT_INFO, "claude-code", "1.0.0"), + ( + b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18",' + b'"capabilities":{},"clientInfo":{"name":"","version":"0"}}}', + "", + "0", + ), + ], +) +def test_extract_initialize_client_info_reads_client_name_and_version(body, expected_name, expected_version): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + client_info = mcp_server._extract_initialize_client_info(body) + + assert client_info is not None + assert client_info.name == expected_name + assert client_info.version == expected_version + + +@pytest.mark.parametrize( + "body", + [ + b"", + b"not json", + b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', + b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', + ], +) +def test_extract_initialize_client_info_returns_none_without_client_info(body): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + + assert mcp_server._extract_initialize_client_info(body) is None + + +@pytest.mark.asyncio +async def test_initialize_request_records_client_name_in_gateway_sessions_report(): + """The real initialize body's clientInfo is attributed to the session the + stateful manager creates, together with the authenticated user.""" + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + session_manager_stateful, + session_manager_stateless, + ) + except ImportError: + pytest.skip("MCP server not available") + + session_id = "initialize-client-info-session-1" + owner_auth = UserAPIKeyAuth( + api_key="initialize-key", + user_id="user-a", + user_email="a@example.com", + key_alias="alice-key", + team_id="team-1", + ) + scope = { + "type": "http", + "method": "POST", + "path": "/mcp", + "headers": [ + (b"content-type", b"application/json"), + (b"authorization", b"Bearer initialize-key"), + ], + } + receive = AsyncMock(return_value={"type": "http.request", "body": _INITIALIZE_WITH_CLIENT_INFO, "more_body": False}) + instances: dict[str, object] = {} + + async def stateful_handle(s, r, se): + instances[session_id] = MagicMock() + await se( + { + "type": "http.response.start", + "headers": [(b"mcp-session-id", session_id.encode())], + } + ) + + async def stateless_handle(s, r, se): + raise AssertionError("initialize request should use stateful manager") + + try: + with ( + patch( # test-quality-ok: admission auth is resolved by a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(owner_auth, None, None, None, None, None), + ), + patch( # test-quality-ok: registry is empty in unit tests; key owns one server + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new_callable=AsyncMock, + return_value=[MagicMock()], + ), + patch( # test-quality-ok: session manager init is a module-level flag; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", + True, + ), + patch.object( # test-quality-ok: the transports are module-level singletons; the suite's only seam + session_manager_stateful, "handle_request", side_effect=stateful_handle + ), + patch.object( # test-quality-ok: the transports are module-level singletons; the suite's only seam + session_manager_stateless, "handle_request", side_effect=stateless_handle + ), + patch.object( # test-quality-ok: the transport registry is a module-level singleton; the suite's only seam + session_manager_stateful, "_server_instances", instances + ), + patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam + mcp_server._stateful_session_auth_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 handle_streamable_http_mcp(scope, receive, AsyncMock()) + report = mcp_server.get_mcp_gateway_sessions_report() + + assert report.total_sessions == 1 + assert [session.model_dump() for session in report.sessions] == [ + { + "session_id_prefix": session_id[:8], + "client_name": "claude-code", + "client_version": "1.0.0", + "user_id": "user-a", + "user_email": "a@example.com", + "key_alias": "alice-key", + "team_id": "team-1", + "team_alias": None, + "client_ip": "", + "idle_seconds": report.sessions[0].idle_seconds, + "in_flight_requests": 0, + } + ] + assert [(group.label, group.count) for group in report.by_client] == [("claude-code", 1)] + assert [(group.label, group.count) for group in report.by_user] == [("user-a", 1)] + assert "initialize-key" not in report.model_dump_json() + finally: + mcp_server._remove_stateful_session_tracking(session_id) + + +def test_gateway_sessions_report_groups_live_sessions_by_client_and_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") + from mcp.types import Implementation + + def auth_user(user_id: str) -> object: + return mcp_server.MCPAuthenticatedUser( + user_api_key_auth=UserAPIKeyAuth(api_key=f"key-{user_id}", user_id=user_id), + client_ip="10.0.0.1", + ) + + contexts = { + "alice-1": auth_user("alice"), + "alice-2": auth_user("alice"), + "bob-1": auth_user("bob"), + "anon-1": mcp_server.MCPAuthenticatedUser(user_api_key_auth=None), + "gone-1": auth_user("alice"), + } + client_info = { + "alice-1": Implementation(name="claude-code", version="1.0.0"), + "alice-2": Implementation(name="claude-code", version="1.0.1"), + "bob-1": Implementation(name="cursor", version="0.50.0"), + "gone-1": Implementation(name="cursor", version="0.50.0"), + } + last_seen = {"alice-1": 90.0, "alice-2": 100.0, "bob-1": 70.0, "anon-1": 100.0, "gone-1": 100.0} + live_instances = {session_id: MagicMock() for session_id in ("alice-1", "alice-2", "bob-1", "anon-1")} + + 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_instances + ), + 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, client_info, clear=True + ), + patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam + mcp_server._stateful_session_auth_context_last_seen, 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_active_request_counts, {"bob-1": 2}, clear=True + ), + ): + report = mcp_server.get_mcp_gateway_sessions_report(now=100.0) + + assert report.total_sessions == 4 + assert [(group.label, group.count) for group in report.by_client] == [ + ("claude-code", 2), + ("cursor", 1), + (None, 1), + ] + assert [(group.label, group.count) for group in report.by_user] == [ + ("alice", 2), + ("bob", 1), + (None, 1), + ] + by_prefix = {session.session_id_prefix: session for session in report.sessions} + assert set(by_prefix) == {"alice-1", "alice-2", "bob-1", "anon-1"} + assert by_prefix["alice-1"].idle_seconds == 10.0 + assert by_prefix["bob-1"].in_flight_requests == 2 + assert by_prefix["bob-1"].client_ip == "10.0.0.1" + assert by_prefix["anon-1"].client_name is None + assert by_prefix["anon-1"].user_id is None + assert "key-alice" not in report.model_dump_json() + + +def test_remove_stateful_session_tracking_drops_client_info(): + try: + from litellm.proxy._experimental.mcp_server import server as mcp_server + except ImportError: + pytest.skip("MCP server not available") + from mcp.types import Implementation + + session_id = "client-info-cleanup-session" + with patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam + mcp_server._stateful_session_client_info, + {session_id: Implementation(name="cursor", version="1")}, + clear=True, + ): + mcp_server._remove_stateful_session_tracking(session_id) + assert session_id not in mcp_server._stateful_session_client_info + + @pytest.mark.asyncio async def test_initialize_request_with_existing_session_tracks_new_session(): try: diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 54b190f7195..7c874aff3df 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -2360,7 +2360,7 @@ class TestTemporaryMCPSessionEndpoints: "litellm.proxy.management_endpoints.mcp_management_endpoints.validate_and_normalize_mcp_server_payload", MagicMock(), ): - with pytest.raises(Exception, match='User does not have permission to create temporary mcp') as exc_info: + with pytest.raises(Exception, match="User does not have permission to create temporary mcp") as exc_info: await add_session_mcp_server( payload=payload, user_api_key_dict=non_admin, @@ -4093,8 +4093,11 @@ async def test_health_discovery_respects_route_restricted_key_grants( manager: Final = mcp_server_manager.MCPServerManager() manager.registry = { server_id: MCPServer( - server_id=server_id, name=server_id, transport=MCPTransport.http, - spec_path=f"https://93.184.216.34/{server_id}.json", auth_type=MCPAuth.none, + server_id=server_id, + name=server_id, + transport=MCPTransport.http, + spec_path=f"https://93.184.216.34/{server_id}.json", + auth_type=MCPAuth.none, ) for server_id in ("server-x", "server-y") } @@ -4107,18 +4110,24 @@ async def test_health_discovery_respects_route_restricted_key_grants( api_key="test-health-key", allowed_routes=["/v1/mcp/server", "/v1/mcp/server/health"] if restricted else [], object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id="health-permissions", mcp_servers=list(grants), + object_permission_id="health-permissions", + mcp_servers=list(grants), ), ) with ( patch.object( # test-quality-ok: TQ008 inject real registry into legacy route binding - mgmt_endpoints, "global_mcp_server_manager", manager, + mgmt_endpoints, + "global_mcp_server_manager", + manager, ), patch.object( # test-quality-ok: TQ008 inject shared registry without mocking permission policy - mcp_server_manager, "global_mcp_server_manager", manager, + mcp_server_manager, + "global_mcp_server_manager", + manager, ), patch( # test-quality-ok: TQ008 configure mode without mocking authorization - "litellm.proxy.proxy_server.general_settings", {"user_mcp_management_mode": mode}, + "litellm.proxy.proxy_server.general_settings", + {"user_mcp_management_mode": mode}, ), ): result: Final = await mgmt_endpoints.health_check_servers( @@ -7125,9 +7134,7 @@ class TestImportMCPServers: import_mcp_servers, ) - payload = MCPConnectorImportRequest.model_validate( - {"mcpServers": {"srv": {"url": "https://x.example/mcp"}}} - ) + payload = MCPConnectorImportRequest.model_validate({"mcpServers": {"srv": {"url": "https://x.example/mcp"}}}) caller = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.INTERNAL_USER) with patch( # test-quality-ok: endpoint takes collaborators from module scope, matching the suite's pattern @@ -7263,3 +7270,54 @@ class TestImportMCPServers: assert [entry.name for entry in result.imported] == ["new-server"] mock_manager.reload_servers_from_database.assert_awaited_once() + + +class TestGetMCPGatewaySessions: + @pytest.mark.asyncio + async def test_non_admin_forbidden(self): + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_mcp_gateway_sessions, + ) + + non_admin = generate_mock_user_api_key_auth(user_role=LitellmUserRoles.INTERNAL_USER) + with pytest.raises(HTTPException) as exc_info: + await get_mcp_gateway_sessions(user_api_key_dict=non_admin) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) + async def test_admin_roles_receive_live_session_report(self, role): + from mcp.types import Implementation + + from litellm.proxy._experimental.mcp_server import server as mcp_server + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_mcp_gateway_sessions, + ) + from litellm.types.mcp import MCPGatewaySessionsResponse + + session_id = "gateway-sessions-endpoint-1" + auth_user = mcp_server.MCPAuthenticatedUser( + user_api_key_auth=UserAPIKeyAuth(api_key="sk-live-secret", 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: MagicMock()} + ), + 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 + ), + patch.dict( # test-quality-ok: the session tables are module-level singletons; the suite's only seam + mcp_server._stateful_session_client_info, + {session_id: Implementation(name="cursor", version="0.50.0")}, + clear=True, + ), + ): + result = await get_mcp_gateway_sessions( + user_api_key_dict=generate_mock_user_api_key_auth(user_role=role), + ) + + assert isinstance(result, MCPGatewaySessionsResponse) + assert result.total_sessions == 1 + 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() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.integration.test.tsx new file mode 100644 index 00000000000..11328ff1a3d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.integration.test.tsx @@ -0,0 +1,138 @@ +import React from "react"; +import { render, screen, within } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { MCPGatewaySessionsTab, formatIdleSeconds } from "./MCPGatewaySessionsTab"; +import * as networking from "@/components/networking"; +import type { MCPGatewaySessionsResponse } from "@/components/mcp_tools/types"; + +vi.mock("@/components/networking", () => ({ + fetchMCPGatewaySessions: vi.fn(), +})); + +const REPORT: MCPGatewaySessionsResponse = { + worker_pid: 4242, + total_sessions: 3, + by_client: [ + { label: "claude-code", count: 2 }, + { label: "cursor", count: 1 }, + ], + by_user: [ + { label: "alice", count: 2 }, + { label: null, count: 1 }, + ], + sessions: [ + { + session_id_prefix: "aaaa1111", + client_name: "claude-code", + client_version: "1.0.0", + user_id: "alice", + user_email: "alice@example.com", + key_alias: "alice-key", + team_id: "team-1", + team_alias: "platform", + client_ip: "10.0.0.1", + idle_seconds: 75, + in_flight_requests: 0, + }, + { + session_id_prefix: "bbbb2222", + client_name: "claude-code", + client_version: "1.0.1", + user_id: "alice", + user_email: null, + key_alias: null, + team_id: null, + team_alias: null, + client_ip: "", + idle_seconds: 3, + in_flight_requests: 1, + }, + { + session_id_prefix: "cccc3333", + client_name: "cursor", + client_version: null, + user_id: null, + user_email: null, + key_alias: null, + team_id: null, + team_alias: null, + client_ip: null, + idle_seconds: 0, + in_flight_requests: 0, + }, + ], +}; + +const renderTab = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + return render( + + + , + ); +}; + +describe("formatIdleSeconds", () => { + it("renders seconds under a minute and minutes plus seconds above it", () => { + expect(formatIdleSeconds(0)).toBe("0s"); + expect(formatIdleSeconds(59.9)).toBe("59s"); + expect(formatIdleSeconds(60)).toBe("1m"); + expect(formatIdleSeconds(75)).toBe("1m 15s"); + expect(formatIdleSeconds(-4)).toBe("0s"); + }); +}); + +describe("MCPGatewaySessionsTab", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("shows grouped counts and session rows from /v1/mcp/sessions", async () => { + vi.mocked(networking.fetchMCPGatewaySessions).mockResolvedValue(REPORT); + renderTab(); + + const byClient = await screen.findByRole("region", { name: "Sessions by AI client" }); + expect(within(byClient).getByRole("row", { name: /claude-code 2/ })).toBeInTheDocument(); + expect(within(byClient).getByRole("row", { name: /cursor 1/ })).toBeInTheDocument(); + + const byUser = screen.getByRole("region", { name: "Sessions by user" }); + expect(within(byUser).getByRole("row", { name: /alice 2/ })).toBeInTheDocument(); + expect(within(byUser).getByRole("row", { name: /\(unknown\) 1/ })).toBeInTheDocument(); + + const sessions = screen.getByRole("region", { name: "Live sessions" }); + const firstRow = within(sessions).getByRole("row", { name: /aaaa1111/ }); + expect(firstRow).toHaveTextContent("claude-code"); + expect(firstRow).toHaveTextContent("v1.0.0"); + expect(firstRow).toHaveTextContent("alice@example.com"); + expect(firstRow).toHaveTextContent("platform"); + expect(firstRow).toHaveTextContent("1m 15s"); + expect(within(sessions).getByRole("row", { name: /cccc3333/ })).toHaveTextContent("(unknown)"); + expect(screen.getByText("Live sessions (worker pid 4242)")).toBeInTheDocument(); + expect(networking.fetchMCPGatewaySessions).toHaveBeenCalledWith("token"); + }); + + it("shows an empty state when the worker holds no live sessions", async () => { + const emptyReport: MCPGatewaySessionsResponse = { + worker_pid: 7, + total_sessions: 0, + by_client: [], + by_user: [], + sessions: [], + }; + vi.mocked(networking.fetchMCPGatewaySessions).mockResolvedValue(emptyReport); + renderTab(); + + expect(await screen.findByText(/No live MCP connections on this worker \(pid 7\)/)).toBeInTheDocument(); + expect(screen.queryByRole("region", { name: "Live sessions" })).not.toBeInTheDocument(); + }); + + it("shows the API error when the request fails", async () => { + vi.mocked(networking.fetchMCPGatewaySessions).mockRejectedValue(new Error("Admin access required")); + renderTab(); + + const alert = await screen.findByRole("alert"); + expect(alert).toHaveTextContent("Could not load live connections"); + expect(alert).toHaveTextContent("Admin access required"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx new file mode 100644 index 00000000000..18f44d5d090 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPGatewaySessionsTab.tsx @@ -0,0 +1,222 @@ +"use client"; + +import React from "react"; +import { useQuery } from "@tanstack/react-query"; +import { RefreshCw } from "lucide-react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert"; +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 { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory"; + +const mcpGatewaySessionKeys = createQueryKeys("mcpGatewaySessions"); +const REFETCH_INTERVAL_MS = 15000; +const UNKNOWN_LABEL = "(unknown)"; + +export function formatIdleSeconds(idleSeconds: number): string { + const total = Math.max(0, Math.floor(idleSeconds)); + if (total < 60) return `${total}s`; + const minutes = Math.floor(total / 60); + const seconds = total % 60; + return seconds === 0 ? `${minutes}m` : `${minutes}m ${seconds}s`; +} + +function groupLabel(label: string | null): string { + if (label === null) return UNKNOWN_LABEL; + return label === "" ? '""' : label; +} + +function StatCard({ label, value }: { label: string; value: number }) { + return ( +
+
{value}
+
{label}
+
+ ); +} + +function GroupCountTable({ + title, + groups, + labelHeader, +}: { + title: string; + groups: MCPGatewaySessionGroupCount[]; + labelHeader: string; +}) { + return ( +
+

{title}

+ + + + {labelHeader} + Sessions + + + + {groups.map((group) => ( + + {groupLabel(group.label)} + {group.count} + + ))} + +
+
+ ); +} + +function SessionsBody({ + data, + error, + isLoading, +}: { + data: MCPGatewaySessionsResponse | undefined; + error: Error | null; + isLoading: boolean; +}) { + if (isLoading) { + return ( +
+ +

Loading live connections...

+
+ ); + } + if (error) { + return ( + + Could not load live connections + {error.message} + + ); + } + if (!data) return null; + if (data.total_sessions === 0) { + return ( +
+

+ No live MCP connections on this worker (pid {data.worker_pid}). Connect an AI client to the gateway to see it + here. +

+
+ ); + } + return ( + <> +
+ + + +
+
+ + +
+
+

+ Live sessions (worker pid {data.worker_pid}) +

+ + + + Session + Client + User + Key alias + Team + Client IP + Idle + In flight + + + + {data.sessions.map((session) => ( + + {session.session_id_prefix} + + {session.client_name === null ? ( + {UNKNOWN_LABEL} + ) : ( + <> + {groupLabel(session.client_name)} + {session.client_version ? ( + v{session.client_version} + ) : null} + + )} + + + {session.user_id === null ? ( + {UNKNOWN_LABEL} + ) : ( + <> + {session.user_id} + {session.user_email ? ( + {session.user_email} + ) : null} + + )} + + {session.key_alias ?? "-"} + {session.team_alias ?? session.team_id ?? "-"} + {session.client_ip || "-"} + {formatIdleSeconds(session.idle_seconds)} + {session.in_flight_requests} + + ))} + +
+
+ + ); +} + +interface MCPGatewaySessionsTabProps { + accessToken: string | null; +} + +export function MCPGatewaySessionsTab({ accessToken }: MCPGatewaySessionsTabProps) { + const queryOptions = { + queryKey: mcpGatewaySessionKeys.lists(), + queryFn: () => fetchMCPGatewaySessions(accessToken!), + enabled: !!accessToken, + refetchInterval: REFETCH_INTERVAL_MS, + }; + const { data, error, isLoading, isFetching, refetch } = useQuery(queryOptions); + + return ( +
+
+
+

Live Connections

+

+ Stateful Streamable HTTP sessions currently open on this proxy worker, grouped by the AI client that sent + the MCP initialize request and by the authenticated LiteLLM user. Stateless requests and SSE connections are + not counted. +

+
+ +
+ + +
+ ); +} + +export default MCPGatewaySessionsTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index e6148d5d997..8a1f8aa9206 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -22,6 +22,7 @@ import { useMCPServerHealth } from "@/app/(dashboard)/hooks/mcpServers/useMCPSer import { toast } from "@/lib/toast"; import { deleteMCPServer } from "@/components/networking"; import { MCPSubmissionsTab } from "./MCPSubmissionsTab"; +import { MCPGatewaySessionsTab } from "./MCPGatewaySessionsTab"; import { MCPToolsetsTab } from "./MCPToolsetsTab"; import CreateMCPServer from "./CreateMCPServer"; import ImportMCPServers from "./ImportMCPServers"; @@ -560,6 +561,11 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) Submitted MCPs )} + {isAdminRole(userRole) && ( + + Live Connections + + )} {selectedServerId ? ( @@ -747,6 +753,11 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) )} + {isAdminRole(userRole) && ( + + + + )} {byokModalServer && ( diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index bc14e7a87ec..fe04eb4969b 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -560,3 +560,30 @@ export interface MCPSubmissionsSummary { rejected: number; items: MCPServer[]; } + +export interface MCPGatewaySession { + session_id_prefix: string; + client_name: string | null; + client_version: string | null; + user_id: string | null; + user_email: string | null; + key_alias: string | null; + team_id: string | null; + team_alias: string | null; + client_ip: string | null; + idle_seconds: number; + in_flight_requests: number; +} + +export interface MCPGatewaySessionGroupCount { + label: string | null; + count: number; +} + +export interface MCPGatewaySessionsResponse { + worker_pid: number; + total_sessions: number; + by_client: MCPGatewaySessionGroupCount[]; + by_user: MCPGatewaySessionGroupCount[]; + sessions: MCPGatewaySession[]; +} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index e77c8ba7e41..a8ea1b66488 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -97,7 +97,7 @@ 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 { MCPUserEnvVarsStatus } from "./mcp_tools/types"; +import type { MCPGatewaySessionsResponse, MCPUserEnvVarsStatus } from "./mcp_tools/types"; import type { CoordinationRedisSettings, CoordinationRedisSettingsResponse, @@ -5109,6 +5109,9 @@ export const fetchMCPSubmissions = async (accessToken: string) => { } }; +export const fetchMCPGatewaySessions = async (accessToken: string): Promise => + apiClient.get(`/v1/mcp/sessions`, { accessToken }); + export const approveMCPServer = async (accessToken: string, serverId: string) => { try { const url = (proxyBaseUrl ? `${proxyBaseUrl}` : "") + `/v1/mcp/server/${encodeURIComponent(serverId)}/approve`; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 872875cc535..f1779c7cb5a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -19002,6 +19002,26 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/mcp/sessions": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Mcp Gateway Sessions + * @description Live stateful MCP gateway sessions on this proxy worker, grouped by AI client and by user. + */ + get: operations["get_mcp_gateway_sessions_v1_mcp_sessions_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/mcp/tools": { parameters: { query?: never; @@ -32315,6 +32335,54 @@ export interface components { * @enum {string} */ MCPEnvVarScope: "global" | "user"; + /** + * MCPGatewaySession + * @description One live stateful Streamable HTTP session held by this proxy worker. + */ + MCPGatewaySession: { + /** Client Ip */ + client_ip?: string | null; + /** Client Name */ + client_name?: string | null; + /** Client Version */ + client_version?: string | null; + /** Idle Seconds */ + idle_seconds: number; + /** In Flight Requests */ + in_flight_requests: number; + /** Key Alias */ + key_alias?: string | null; + /** Session Id Prefix */ + session_id_prefix: string; + /** Team Alias */ + team_alias?: string | null; + /** Team Id */ + team_id?: string | null; + /** User Email */ + user_email?: string | null; + /** User Id */ + user_id?: string | null; + }; + /** MCPGatewaySessionGroupCount */ + MCPGatewaySessionGroupCount: { + /** Count */ + count: number; + /** Label */ + label?: string | null; + }; + /** MCPGatewaySessionsResponse */ + MCPGatewaySessionsResponse: { + /** By Client */ + by_client?: components["schemas"]["MCPGatewaySessionGroupCount"][]; + /** By User */ + by_user?: components["schemas"]["MCPGatewaySessionGroupCount"][]; + /** Sessions */ + sessions?: components["schemas"]["MCPGatewaySession"][]; + /** Total Sessions */ + total_sessions: number; + /** Worker Pid */ + worker_pid: number; + }; /** * MCPOAuthUserCredentialRequest * @description Stores a user's OAuth2 token for an OpenAPI MCP server. @@ -65299,6 +65367,26 @@ export interface operations { }; }; }; + get_mcp_gateway_sessions_v1_mcp_sessions_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["MCPGatewaySessionsResponse"]; + }; + }; + }; + }; get_mcp_tools_v1_mcp_tools_get: { parameters: { query?: never;