From 4b107e0c705acc30fe10faad4a71f4d4d8cdca43 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 13 Jun 2025 10:10:01 -0700 Subject: [PATCH] [Feat] UI - Allow setting MCP servers when creating keys, teams (#11711) * Add PremiumMCPSelector * add PremiumMCPSelector * Fix CTA enterprise --- .../common_components/PremiumMCPSelector.tsx | 50 +++++++++++++++++++ .../PremiumVectorStoreSelector.tsx | 2 +- .../src/components/create_key_button.tsx | 35 ++++++++++++- .../src/components/organizations.tsx | 2 +- ui/litellm-dashboard/src/components/teams.tsx | 33 +++++++++++- .../src/components/view_logs/audit_logs.tsx | 2 +- 6 files changed, 119 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/common_components/PremiumMCPSelector.tsx diff --git a/ui/litellm-dashboard/src/components/common_components/PremiumMCPSelector.tsx b/ui/litellm-dashboard/src/components/common_components/PremiumMCPSelector.tsx new file mode 100644 index 0000000000..f54d4d31df --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/PremiumMCPSelector.tsx @@ -0,0 +1,50 @@ +import React from 'react'; +import { Text } from "@tremor/react"; +import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; + +interface PremiumMCPSelectorProps { + onChange: (values: string[]) => void; + value: string[]; + accessToken: string; + placeholder?: string; + premiumUser?: boolean; +} + +export function PremiumMCPSelector({ + onChange, + value, + accessToken, + placeholder = "Select MCP servers", + premiumUser = false +}: PremiumMCPSelectorProps) { + if (!premiumUser) { + return ( +
+
+
+ ✨ premium-mcp-server-1 +
+
+ ✨ premium-mcp-server-2 +
+
+
+ + MCP server access control is a LiteLLM Enterprise feature. Get a trial key here. + +
+
+ ); + } + + return ( + + ); +} + +export default PremiumMCPSelector; \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/common_components/PremiumVectorStoreSelector.tsx b/ui/litellm-dashboard/src/components/common_components/PremiumVectorStoreSelector.tsx index 9134838342..7f7fbfd626 100644 --- a/ui/litellm-dashboard/src/components/common_components/PremiumVectorStoreSelector.tsx +++ b/ui/litellm-dashboard/src/components/common_components/PremiumVectorStoreSelector.tsx @@ -30,7 +30,7 @@ export function PremiumVectorStoreSelector({
- Vector store access control is a LiteLLM Enterprise feature. Get a trial key here. + Vector store access control is a LiteLLM Enterprise feature. Get a trial key here.
diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 2d05450298..5f9bbf62ea 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/create_key_button.tsx @@ -35,6 +35,7 @@ import { } from "./networking"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import PremiumVectorStoreSelector from "./common_components/PremiumVectorStoreSelector"; +import PremiumMCPSelector from "./common_components/PremiumMCPSelector"; import { Team } from "./key_team_helpers/key_list"; import TeamDropdown from "./common_components/team_dropdown"; import { InfoCircleOutlined } from '@ant-design/icons'; @@ -268,7 +269,7 @@ const CreateKey: React.FC = ({ formValues.metadata = JSON.stringify(metadata); } - // Transform allowed_vector_store_ids into object_permission format + // Transform allowed_vector_store_ids and allowed_mcp_server_ids into object_permission format if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { formValues.object_permission = { vector_stores: formValues.allowed_vector_store_ids @@ -277,6 +278,16 @@ const CreateKey: React.FC = ({ delete formValues.allowed_vector_store_ids; } + // Transform allowed_mcp_server_ids into object_permission format + if (formValues.allowed_mcp_server_ids && formValues.allowed_mcp_server_ids.length > 0) { + if (!formValues.object_permission) { + formValues.object_permission = {}; + } + formValues.object_permission.mcp_servers = formValues.allowed_mcp_server_ids; + // Remove the original field as it's now part of object_permission + delete formValues.allowed_mcp_server_ids; + } + const response = await keyCreateCall(accessToken, userID, formValues); console.log("key create Response:", response); @@ -711,6 +722,28 @@ const CreateKey: React.FC = ({ /> + + Allowed MCP Servers{' '} + + + + + } + name="allowed_mcp_server_ids" + className="mt-4" + help="Select MCP servers this key can access. Leave empty for access to all MCP servers" + > + form.setFieldValue('allowed_mcp_server_ids', values)} + value={form.getFieldValue('allowed_mcp_server_ids')} + accessToken={accessToken} + placeholder="Select MCP servers (optional)" + premiumUser={premiumUser} + /> + + diff --git a/ui/litellm-dashboard/src/components/organizations.tsx b/ui/litellm-dashboard/src/components/organizations.tsx index 6bb78ccee0..8a6557d7fd 100644 --- a/ui/litellm-dashboard/src/components/organizations.tsx +++ b/ui/litellm-dashboard/src/components/organizations.tsx @@ -135,7 +135,7 @@ const OrganizationsTable: React.FC = ({ if (!premiumUser) { return (
- This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here. + This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here.
); } diff --git a/ui/litellm-dashboard/src/components/teams.tsx b/ui/litellm-dashboard/src/components/teams.tsx index cc8b4220cc..44ffa09b9a 100644 --- a/ui/litellm-dashboard/src/components/teams.tsx +++ b/ui/litellm-dashboard/src/components/teams.tsx @@ -59,6 +59,7 @@ import { CogIcon } from "@heroicons/react/outline"; import AvailableTeamsPanel from "@/components/team/available_teams"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import PremiumVectorStoreSelector from "./common_components/PremiumVectorStoreSelector"; +import PremiumMCPSelector from "./common_components/PremiumMCPSelector"; import type { KeyResponse, Team } from "./key_team_helpers/key_list"; interface TeamProps { @@ -324,13 +325,22 @@ const Teams: React.FC = ({ } message.info("Creating Team"); - // Transform allowed_vector_store_ids into object_permission + // Transform allowed_vector_store_ids and allowed_mcp_server_ids into object_permission if (formValues.allowed_vector_store_ids && formValues.allowed_vector_store_ids.length > 0) { formValues.object_permission = { vector_stores: formValues.allowed_vector_store_ids }; delete formValues.allowed_vector_store_ids; } + + // Transform allowed_mcp_server_ids into object_permission + if (formValues.allowed_mcp_server_ids && formValues.allowed_mcp_server_ids.length > 0) { + if (!formValues.object_permission) { + formValues.object_permission = {}; + } + formValues.object_permission.mcp_servers = formValues.allowed_mcp_server_ids; + delete formValues.allowed_mcp_server_ids; + } const response: any = await teamCreateCall(accessToken, formValues); if (teams !== null) { setTeams([...teams, response]); @@ -1093,6 +1103,27 @@ const Teams: React.FC = ({ premiumUser={premiumUser} /> + + Allowed MCP Servers{' '} + + + + + } + name="allowed_mcp_server_ids" + className="mt-8" + help="Select MCP servers this team can access. Leave empty for access to all MCP servers" + > + form.setFieldValue('allowed_mcp_server_ids', values)} + value={form.getFieldValue('allowed_mcp_server_ids')} + accessToken={accessToken || ''} + placeholder="Select MCP servers (optional)" + premiumUser={premiumUser} + /> + diff --git a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx index 71032ba4b3..bea6a58739 100644 --- a/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/audit_logs.tsx @@ -390,7 +390,7 @@ export default function AuditLogs({ if (!premiumUser) { return (
- This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here. + This is a LiteLLM Enterprise feature, and requires a valid key to use. Get a trial key here.
); }