diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b9e0379e445..1b86551972a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -4144,10 +4144,16 @@ class MCPServerManager: emitted so admins can diagnose stale/typo permission entries — the downstream access-check denies them when compared against the concrete request server_id. + + The ``all-proxy-mcps`` sentinel expands to every server in the live + registry, so an entity granted it stays in sync as servers are added + or removed without re-editing its permission list. """ if not identifiers: return [] registry = self.get_registry() + if SpecialMCPServerNames.all_proxy_mcp_servers.value in identifiers: + return list(registry.keys()) expanded: Set[str] = set() for identifier in identifiers: if identifier in registry: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7db0156bbe0..d3279b29a06 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2985,6 +2985,7 @@ class SpecialModelNames(enum.Enum): class SpecialMCPServerNames(enum.Enum): no_mcp_servers = "no-mcp-servers" + all_proxy_mcp_servers = "all-proxy-mcps" class SpecialProxyStrings(enum.Enum): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 20fd5c1d86a..5b7c4948c56 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1,6 +1,7 @@ import json import os import sys +from datetime import datetime from unittest.mock import AsyncMock, MagicMock, call as mock_call, patch import pytest @@ -16,6 +17,8 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) from litellm.proxy._types import ( + LiteLLM_MCPServerTable, + MCPTransport, SpecialHeaders, SpecialMCPServerNames, UserAPIKeyAuth, @@ -100,6 +103,66 @@ class TestMCPRequestHandler: mock_key_servers.assert_called_once_with(mock_user_auth) mock_team_servers.assert_called_once_with(mock_user_auth) + async def test_get_allowed_mcp_servers_for_team_all_proxy_mcps(self): + """A team scoped to the all-proxy-mcps sentinel resolves to every + server in the live registry — the smoothest path requested in the + ticket, exercised through the real expansion logic.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + ) + + manager = MCPServerManager() + for server_id in ("srv-a", "srv-b", "srv-c"): + await manager.add_server( + LiteLLM_MCPServerTable( + server_id=server_id, + alias=server_id, + description="", + url=None, + transport=MCPTransport.stdio, + command="python", + args=["-m", "server"], + env={}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + ) + + object_permission = MagicMock() + object_permission.mcp_servers = [ + SpecialMCPServerNames.all_proxy_mcp_servers.value + ] + object_permission.mcp_access_groups = [] + object_permission.mcp_tool_permissions = None + + team_obj = MagicMock() + team_obj.object_permission = object_permission + team_obj.access_group_ids = [] + + user_auth = UserAPIKeyAuth(api_key="k", user_id="u", team_id="t") + + with patch( + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), patch( + "litellm.proxy.auth.auth_checks.get_team_object", + AsyncMock(return_value=team_obj), + ), patch( + "litellm.proxy.auth.auth_checks._get_mcp_server_ids_from_access_groups", + AsyncMock(return_value=[]), + ), patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + manager, + ), patch.object( + MCPRequestHandler, + "_get_mcp_servers_from_access_groups", + AsyncMock(return_value=[]), + ): + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( + user_auth + ) + + assert sorted(result) == ["srv-a", "srv-b", "srv-c"] + @pytest.mark.parametrize( "team_servers,key_servers,expected_servers,scenario", [ diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 35a67391315..dc5a2783a91 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -38,6 +38,7 @@ from litellm.proxy._types import ( MCPEnvVar, MCPEnvVarScope, MCPTransport, + SpecialMCPServerNames, ) from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer @@ -111,6 +112,54 @@ class TestMCPServerManager: assert added_server.args == ["-m", "server"] assert added_server.env == {"DEBUG": "1", "TEST": "1"} + async def test_expand_permission_list_all_proxy_mcps_sentinel(self): + """The all-proxy-mcps sentinel expands to every server in the live + registry and re-expands against servers registered later, so a grant + stays current without the permission list being re-edited.""" + manager = MCPServerManager() + + async def _register(server_id: str): + await manager.add_server( + LiteLLM_MCPServerTable( + server_id=server_id, + alias=server_id, + description="", + url=None, + transport=MCPTransport.stdio, + command="python", + args=["-m", "server"], + env={}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + ) + + await _register("srv-a") + await _register("srv-b") + + sentinel = SpecialMCPServerNames.all_proxy_mcp_servers.value + + assert set(manager.expand_permission_list([sentinel])) == {"srv-a", "srv-b"} + + # The sentinel dominates: mixing it with a concrete id still grants all. + assert set(manager.expand_permission_list([sentinel, "srv-a"])) == { + "srv-a", + "srv-b", + } + + # A server registered after the grant is included with no re-edit. + await _register("srv-c") + assert set(manager.expand_permission_list([sentinel])) == { + "srv-a", + "srv-b", + "srv-c", + } + + # Without the sentinel, only the named server resolves — proving the + # full-registry result above comes from the sentinel branch, not a + # blanket "return everything". + assert manager.expand_permission_list(["srv-a"]) == ["srv-a"] + async def test_create_mcp_client_stdio(self): """Test creating MCP client for stdio transport""" manager = MCPServerManager() diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx index be9015e3730..1335a037c02 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.tsx @@ -1479,6 +1479,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser value={form.getFieldValue("allowed_mcp_servers_and_groups")} accessToken={accessToken || ""} placeholder="Select MCP servers or access groups (optional)" + allowAllProxyMcps /> diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx index 1517b2dfa27..f1b33a3534d 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.test.tsx @@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import MCPServerSelector from "./MCPServerSelector"; -import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { ALL_PROXY_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: vi.fn(), @@ -98,3 +98,57 @@ describe("MCPServerSelector no-mcp-servers option", () => { expect(optionByValue(NO_MCP_SERVERS_SENTINEL)?.disabled).toBe(false); }); }); + +describe("MCPServerSelector all-proxy-mcps option", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseMCPServers.mockReturnValue({ + data: [{ server_id: "srv-1", server_name: "Server One" }], + isLoading: false, + } as any); + mockUseMCPAccessGroups.mockReturnValue({ data: [], isLoading: false } as any); + mockUseMCPToolsets.mockReturnValue({ data: [], isLoading: false } as any); + }); + + const optionByValue = (value: string) => + Array.from(screen.getByTestId("mcp-select").querySelectorAll("option")).find( + (o) => (o as HTMLOptionElement).value === value, + ) as HTMLOptionElement | undefined; + + it("hides the All Proxy MCPs option by default", () => { + renderWithProviders( + , + ); + expect(optionByValue(ALL_PROXY_MCPS_SENTINEL)).toBeUndefined(); + }); + + it("emits an exclusive sentinel when All Proxy MCPs is selected", async () => { + const onChange = vi.fn(); + renderWithProviders( + , + ); + expect(optionByValue(ALL_PROXY_MCPS_SENTINEL)).toBeDefined(); + + await userEvent.selectOptions(screen.getByTestId("mcp-select"), [ALL_PROXY_MCPS_SENTINEL]); + + expect(onChange).toHaveBeenCalledWith({ servers: [ALL_PROXY_MCPS_SENTINEL], accessGroups: [], toolsets: [] }); + }); + + it("disables real server options while the sentinel is selected", () => { + renderWithProviders( + , + ); + expect(optionByValue("srv-1")?.disabled).toBe(true); + expect(optionByValue(ALL_PROXY_MCPS_SENTINEL)?.disabled).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx index bbda761938e..c966fd61e28 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx +++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx @@ -3,7 +3,7 @@ import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers" import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { Select } from "antd"; import React from "react"; -import { NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { ALL_PROXY_MCPS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; interface MCPServerSelectorProps { onChange: (selected: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void; @@ -18,10 +18,30 @@ interface MCPServerSelectorProps { disabled?: boolean; teamId?: string | null; allowNoMcpServers?: boolean; + allowAllProxyMcps?: boolean; } const TOOLSET_PREFIX = "toolset:"; +type MCPSelection = { servers: string[]; accessGroups: string[]; toolsets: string[] }; + +// "No MCP Servers" and "All Proxy MCPs" are mutually exclusive with everything +// else. The most restrictive ("block all") wins, mirroring the backend +// precedence where the no-mcp-servers sentinel overrides every grant. +const resolveExclusiveSelection = ( + selected: string[], + allowNoMcpServers: boolean, + allowAllProxyMcps: boolean, +): MCPSelection | null => { + if (allowNoMcpServers && selected.includes(NO_MCP_SERVERS_SENTINEL)) { + return { servers: [NO_MCP_SERVERS_SENTINEL], accessGroups: [], toolsets: [] }; + } + if (allowAllProxyMcps && selected.includes(ALL_PROXY_MCPS_SENTINEL)) { + return { servers: [ALL_PROXY_MCPS_SENTINEL], accessGroups: [], toolsets: [] }; + } + return null; +}; + const MCPServerSelector: React.FC = ({ onChange, value, @@ -31,6 +51,7 @@ const MCPServerSelector: React.FC = ({ disabled = false, teamId, allowNoMcpServers = false, + allowAllProxyMcps = false, }) => { const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(teamId); const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups(); @@ -81,12 +102,14 @@ const MCPServerSelector: React.FC = ({ ]; const hasNoMcpServersSelected = allowNoMcpServers && selectedValues.includes(NO_MCP_SERVERS_SENTINEL); + const hasAllProxyMcpsSelected = allowAllProxyMcps && selectedValues.includes(ALL_PROXY_MCPS_SENTINEL); + const hasExclusiveSelected = hasNoMcpServersSelected || hasAllProxyMcpsSelected; // Handle selection const handleChange = (selected: string[]) => { - // "No MCP Servers" is exclusive: picking it clears everything else. - if (allowNoMcpServers && selected.includes(NO_MCP_SERVERS_SENTINEL)) { - onChange({ servers: [NO_MCP_SERVERS_SENTINEL], accessGroups: [], toolsets: [] }); + const exclusive = resolveExclusiveSelection(selected, allowNoMcpServers, allowAllProxyMcps); + if (exclusive) { + onChange(exclusive); return; } const toolsetsSelected = selected @@ -98,6 +121,39 @@ const MCPServerSelector: React.FC = ({ onChange({ servers, accessGroups: accessGroupsSelected, toolsets: toolsetsSelected }); }; + const renderExclusiveOptions = () => [ + ...(allowAllProxyMcps + ? [ + +
+ All Proxy MCPs + Every server +
+ , + ] + : []), + ...(allowNoMcpServers + ? [ + +
+ No MCP Servers + Block all +
+ , + ] + : []), + ]; + return (