From 9058354ceda7a99ea6fa2a9007a59df27e2b24d8 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 10 Jun 2025 15:46:43 -0700
Subject: [PATCH] [Feat] UI - Add controls for MCP Permission Management
 (#11598)

* init commit to render MCP servers

* allow adding MCP servers to keys

* fix allow adding MCP permissions per team

* ui mcp permissions

* updates based on review
---
 .../src/components/key_edit_view.tsx          |  13 +-
 .../src/components/key_info_view.tsx          |   9 ++
 .../MCPServerSelector.tsx                     |  74 ++++++++++++
 .../components/object_permissions_view.tsx    | 111 ++----------------
 .../permissions/MCPServerPermissions.tsx      |  86 ++++++++++++++
 .../permissions/VectorStorePermissions.tsx    |  83 +++++++++++++
 .../src/components/team/team_info.tsx         |  18 ++-
 7 files changed, 292 insertions(+), 102 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx
 create mode 100644 ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
 create mode 100644 ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx

diff --git a/ui/litellm-dashboard/src/components/key_edit_view.tsx b/ui/litellm-dashboard/src/components/key_edit_view.tsx
index 959ce0362f2..cd4f32567f3 100644
--- a/ui/litellm-dashboard/src/components/key_edit_view.tsx
+++ b/ui/litellm-dashboard/src/components/key_edit_view.tsx
@@ -6,6 +6,7 @@ import { fetchTeamModels } from "../components/create_key_button";
 import { modelAvailableCall } from "./networking";
 import NumericalInput from "./shared/numerical_input";
 import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
+import MCPServerSelector from "./mcp_server_management/MCPServerSelector";
 
 interface KeyEditViewProps {
   keyData: KeyResponse;
@@ -98,7 +99,8 @@ export function KeyEditView({
     budget_duration: getBudgetDuration(keyData.budget_duration),
     metadata: keyData.metadata ? JSON.stringify(keyData.metadata, null, 2) : "",
     guardrails: keyData.metadata?.guardrails || [],
-    vector_stores: keyData.object_permission?.vector_stores || []
+    vector_stores: keyData.object_permission?.vector_stores || [],
+    mcp_servers: keyData.object_permission?.mcp_servers || []
   };
 
   return (
@@ -180,6 +182,15 @@ export function KeyEditView({
         />
       </Form.Item>
 
+      <Form.Item label="MCP Servers" name="mcp_servers">
+        <MCPServerSelector
+          onChange={(values) => form.setFieldValue('mcp_servers', values)}
+          value={form.getFieldValue('mcp_servers')}
+          accessToken={accessToken || ""}
+          placeholder="Select MCP servers"
+        />
+      </Form.Item>
+
       <Form.Item label="Metadata" name="metadata">
         <Input.TextArea rows={10} />
       </Form.Item>
diff --git a/ui/litellm-dashboard/src/components/key_info_view.tsx b/ui/litellm-dashboard/src/components/key_info_view.tsx
index 087d60d20ec..1326d37d708 100644
--- a/ui/litellm-dashboard/src/components/key_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/key_info_view.tsx
@@ -76,6 +76,15 @@ export default function KeyInfoView({ keyId, onClose, keyData, accessToken, user
         delete formValues.vector_stores;
       }
 
+      if (formValues.mcp_servers !== undefined) {
+        formValues.object_permission = {
+          ...keyData.object_permission,
+          mcp_servers: formValues.mcp_servers || []
+        };
+        // Remove mcp_servers from the top level as it should be in object_permission
+        delete formValues.mcp_servers;
+      }
+
       // Convert metadata back to an object if it exists and is a string
       if (formValues.metadata && typeof formValues.metadata === "string") {
         try {
diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx
new file mode 100644
index 00000000000..a8ddda2d257
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx
@@ -0,0 +1,74 @@
+import React, { useEffect, useState } from 'react';
+import { Select } from 'antd';
+import { fetchMCPServers } from '../networking';
+import { MCPServer } from '../mcp_tools/types';
+
+
+interface MCPServerSelectorProps {
+  onChange: (selectedMCPServers: string[]) => void;
+  value?: string[];
+  className?: string;
+  accessToken: string;
+  placeholder?: string;
+  disabled?: boolean;
+}
+
+const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({ 
+  onChange, 
+  value, 
+  className, 
+  accessToken,
+  placeholder = "Select MCP servers",
+  disabled = false
+}) => {
+  const [mcpServers, setMCPServers] = useState<MCPServer[]>([]);
+  const [loading, setLoading] = useState(false);
+
+  useEffect(() => {
+    const fetchMCPServerList = async () => {
+      if (!accessToken) return;
+      
+      setLoading(true);
+      try {
+        const response = await fetchMCPServers(accessToken);
+        if (response && Array.isArray(response)) {
+          // Direct array response
+          setMCPServers(response);
+        } else if (response.data && Array.isArray(response.data)) {
+          // Response with data wrapper
+          setMCPServers(response.data);
+        }
+      } catch (error) {
+        console.error("Error fetching MCP servers:", error);
+      } finally {
+        setLoading(false);
+      }
+    };
+
+    fetchMCPServerList();
+  }, [accessToken]);
+
+  return (
+    <div>
+      <Select
+        mode="multiple"
+        placeholder={placeholder}
+        onChange={onChange}
+        value={value}
+        loading={loading}
+        className={className}
+        options={mcpServers.map(server => ({
+          label: `${server.alias} (${server.server_id})`,
+          value: server.server_id,
+          title: server.description || server.alias,
+        }))}
+        optionFilterProp="label"
+        showSearch
+        style={{ width: '100%' }}
+        disabled={disabled}
+      />
+    </div>
+  );
+};
+
+export default MCPServerSelector; 
\ No newline at end of file
diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx
index e8143af5aa4..16ff13d0e63 100644
--- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx
+++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx
@@ -1,12 +1,7 @@
-import React, { useState, useEffect } from "react";
-import { Card, Text, Badge } from "@tremor/react";
-import { ServerIcon, DatabaseIcon } from "@heroicons/react/outline";
-import { vectorStoreListCall } from "./networking";
-
-interface VectorStoreDetails {
-  vector_store_id: string;
-  vector_store_name?: string;
-}
+import React from "react";
+import { Text } from "@tremor/react";
+import VectorStorePermissions from "./permissions/VectorStorePermissions";
+import MCPServerPermissions from "./permissions/MCPServerPermissions";
 
 interface ObjectPermission {
   object_permission_id: string;
@@ -29,97 +24,17 @@ export function ObjectPermissionsView({
 }: ObjectPermissionsViewProps) {
   const vectorStores = objectPermission?.vector_stores || [];
   const mcpServers = objectPermission?.mcp_servers || [];
-  const [vectorStoreDetails, setVectorStoreDetails] = useState<VectorStoreDetails[]>([]);
-
-  // Fetch vector store details when component mounts
-  useEffect(() => {
-    const fetchVectorStores = async () => {
-      if (!accessToken || vectorStores.length === 0) return;
-      
-      try {
-        const response = await vectorStoreListCall(accessToken);
-        if (response.data) {
-          setVectorStoreDetails(response.data.map((store: any) => ({
-            vector_store_id: store.vector_store_id,
-            vector_store_name: store.vector_store_name
-          })));
-        }
-      } catch (error) {
-        console.error("Error fetching vector stores:", error);
-      }
-    };
-
-    fetchVectorStores();
-  }, [accessToken, vectorStores.length]);
-
-  // Function to get display name for vector store
-  const getVectorStoreDisplayName = (storeId: string) => {
-    const storeDetail = vectorStoreDetails.find(store => store.vector_store_id === storeId);
-    if (storeDetail) {
-      return `${storeDetail.vector_store_name || storeDetail.vector_store_id} (${storeDetail.vector_store_id})`;
-    }
-    return storeId;
-  };
 
   const content = (
     <div className={variant === "card" ? "grid grid-cols-1 md:grid-cols-2 gap-6" : "space-y-4"}>
-      {/* Vector Stores Section */}
-      <div className="space-y-3">
-        <div className="flex items-center gap-2">
-          <DatabaseIcon className="h-4 w-4 text-blue-600" />
-          <Text className="font-semibold text-gray-900">Vector Stores</Text>
-          <Badge color="blue" size="xs">
-            {vectorStores.length}
-          </Badge>
-        </div>
-        
-        {vectorStores.length > 0 ? (
-          <div className="flex flex-wrap gap-2">
-            {vectorStores.map((store, index) => (
-              <div
-                key={index}
-                className="inline-flex items-center px-3 py-1.5 rounded-lg bg-blue-50 border border-blue-200 text-blue-800 text-sm font-medium"
-              >
-                {getVectorStoreDisplayName(store)}
-              </div>
-            ))}
-          </div>
-        ) : (
-          <div className="flex items-center gap-2 px-3 py-2 rounded-lg bg-gray-50 border border-gray-200">
-            <DatabaseIcon className="h-4 w-4 text-gray-400" />
-            <Text className="text-gray-500 text-sm">No vector stores configured</Text>
-          </div>
-        )}
-      </div>
-
-      {/* MCP Servers Section */}
-      {/* <div className="space-y-3">
-        <div className="flex items-center gap-2">
-          <ServerIcon className="h-4 w-4 text-blue-600" />
-          <Text className="font-semibold text-gray-900">MCP Servers</Text>
-          <Badge color="blue" size="xs">
-            {mcpServers.length}
-          </Badge>
-        </div>
-        
-        {mcpServers.length > 0 ? (
-          <div className="flex flex-wrap gap-2">
-            {mcpServers.map((server, index) => (
-              <div
-                key={index}
-                className="inline-flex items-center px-3 py-1.5 rounded-lg bg-blue-50 border border-blue-200 text-blue-800 text-sm font-medium"
-              >
-                {server}
-              </div>
-            ))}
-          </div>
-        ) : (
-          <div className="flex items-center gap-2 px-3 py-2 rounded-lg bg-gray-50 border border-gray-200">
-            <ServerIcon className="h-4 w-4 text-gray-400" />
-            <Text className="text-gray-500 text-sm">No MCP servers configured</Text>
-          </div>
-        )}
-      </div> */}
+      <VectorStorePermissions 
+        vectorStores={vectorStores} 
+        accessToken={accessToken} 
+      />
+      <MCPServerPermissions 
+        mcpServers={mcpServers} 
+        accessToken={accessToken} 
+      />
     </div>
   );
 
@@ -130,7 +45,7 @@ export function ObjectPermissionsView({
           <div>
             <Text className="font-semibold text-gray-900">Object Permissions</Text>
             <Text className="text-xs text-gray-500">
-              Access control for Vector Stores
+              Access control for Vector Stores and MCP Servers
             </Text>
           </div>
         </div>
diff --git a/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
new file mode 100644
index 00000000000..2d85a9b3f11
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/permissions/MCPServerPermissions.tsx
@@ -0,0 +1,86 @@
+import React, { useState, useEffect } from "react";
+import { Text, Badge } from "@tremor/react";
+import { ServerIcon } from "@heroicons/react/outline";
+import { Tooltip } from "antd";
+import { fetchMCPServers } from "../networking";
+import { MCPServer } from '../mcp_tools/types';
+
+interface MCPServerPermissionsProps {
+  mcpServers: string[];
+  accessToken?: string | null;
+}
+
+export function MCPServerPermissions({ 
+  mcpServers, 
+  accessToken 
+}: MCPServerPermissionsProps) {
+  const [mcpServerDetails, setMCPServerDetails] = useState<MCPServer[]>([]);
+
+  // Fetch MCP server details when component mounts
+  useEffect(() => {
+    const fetchMCPServerDetails = async () => {
+      if (!accessToken || mcpServers.length === 0) return;
+      
+      try {
+        const response = await fetchMCPServers(accessToken);
+        if (response && Array.isArray(response)) {
+          // Direct array response
+          setMCPServerDetails(response);
+        } else if (response.data && Array.isArray(response.data)) {
+          // Response with data wrapper
+          setMCPServerDetails(response.data);
+        }
+      } catch (error) {
+        console.error("Error fetching MCP servers:", error);
+      }
+    };
+
+    fetchMCPServerDetails();
+  }, [accessToken, mcpServers.length]);
+
+  // Function to get display name for MCP server
+  const getMCPServerDisplayName = (serverId: string) => {
+    const serverDetail = mcpServerDetails.find(server => server.server_id === serverId);
+    if (serverDetail) {
+      // Truncate server ID: first 3 chars + "..." + last 4 chars
+      const truncatedId = serverId.length > 7 
+        ? `${serverId.slice(0, 3)}...${serverId.slice(-4)}`
+        : serverId;
+      return `${serverDetail.alias} (${truncatedId})`;
+    }
+    return serverId;
+  };
+
+  return (
+    <div className="space-y-3">
+      <div className="flex items-center gap-2">
+        <ServerIcon className="h-4 w-4 text-green-600" />
+        <Text className="font-semibold text-gray-900">MCP Servers</Text>
+        <Badge color="green" size="xs">
+          {mcpServers.length}
+        </Badge>
+      </div>
+      
+      {mcpServers.length > 0 ? (
+        <div className="flex flex-wrap gap-2">
+          {mcpServers.map((server, index) => (
+            <Tooltip key={index} title={`Full ID: ${server}`} placement="top">
+              <div
+                className="inline-flex items-center px-3 py-1.5 rounded-lg bg-green-50 border border-green-200 text-green-800 text-sm font-medium cursor-help"
+              >
+                {getMCPServerDisplayName(server)}
+              </div>
+            </Tooltip>
+          ))}
+        </div>
+      ) : (
+        <div className="flex items-center gap-2 px-3 py-2 rounded-lg bg-gray-50 border border-gray-200">
+          <ServerIcon className="h-4 w-4 text-gray-400" />
+          <Text className="text-gray-500 text-sm">No MCP servers configured</Text>
+        </div>
+      )}
+    </div>
+  );
+}
+
+export default MCPServerPermissions; 
\ No newline at end of file
diff --git a/ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx b/ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx
new file mode 100644
index 00000000000..faa9067a68b
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/permissions/VectorStorePermissions.tsx
@@ -0,0 +1,83 @@
+import React, { useState, useEffect } from "react";
+import { Text, Badge } from "@tremor/react";
+import { DatabaseIcon } from "@heroicons/react/outline";
+import { vectorStoreListCall } from "../networking";
+
+interface VectorStoreDetails {
+  vector_store_id: string;
+  vector_store_name?: string;
+}
+
+interface VectorStorePermissionsProps {
+  vectorStores: string[];
+  accessToken?: string | null;
+}
+
+export function VectorStorePermissions({ 
+  vectorStores, 
+  accessToken 
+}: VectorStorePermissionsProps) {
+  const [vectorStoreDetails, setVectorStoreDetails] = useState<VectorStoreDetails[]>([]);
+
+  // Fetch vector store details when component mounts
+  useEffect(() => {
+    const fetchVectorStores = async () => {
+      if (!accessToken || vectorStores.length === 0) return;
+      
+      try {
+        const response = await vectorStoreListCall(accessToken);
+        if (response.data) {
+          setVectorStoreDetails(response.data.map((store: any) => ({
+            vector_store_id: store.vector_store_id,
+            vector_store_name: store.vector_store_name
+          })));
+        }
+      } catch (error) {
+        console.error("Error fetching vector stores:", error);
+      }
+    };
+
+    fetchVectorStores();
+  }, [accessToken, vectorStores.length]);
+
+  // Function to get display name for vector store
+  const getVectorStoreDisplayName = (storeId: string) => {
+    const storeDetail = vectorStoreDetails.find(store => store.vector_store_id === storeId);
+    if (storeDetail) {
+      return `${storeDetail.vector_store_name || storeDetail.vector_store_id} (${storeDetail.vector_store_id})`;
+    }
+    return storeId;
+  };
+
+  return (
+    <div className="space-y-3">
+      <div className="flex items-center gap-2">
+        <DatabaseIcon className="h-4 w-4 text-blue-600" />
+        <Text className="font-semibold text-gray-900">Vector Stores</Text>
+        <Badge color="blue" size="xs">
+          {vectorStores.length}
+        </Badge>
+      </div>
+      
+      {vectorStores.length > 0 ? (
+        <div className="flex flex-wrap gap-2">
+          {vectorStores.map((store, index) => (
+            <div
+              key={index}
+              className="inline-flex items-center px-3 py-1.5 rounded-lg bg-blue-50 border border-blue-200 text-blue-800 text-sm font-medium"
+            >
+              {getVectorStoreDisplayName(store)}
+            </div>
+          ))}
+        </div>
+      ) : (
+        <div className="flex items-center gap-2 px-3 py-2 rounded-lg bg-gray-50 border border-gray-200">
+          <DatabaseIcon className="h-4 w-4 text-gray-400" />
+          <Text className="text-gray-500 text-sm">No vector stores configured</Text>
+        </div>
+      )}
+    </div>
+  );
+}
+
+export default VectorStorePermissions; 
\ No newline at end of file
diff --git a/ui/litellm-dashboard/src/components/team/team_info.tsx b/ui/litellm-dashboard/src/components/team/team_info.tsx
index 084f58065ee..1c6bfd46832 100644
--- a/ui/litellm-dashboard/src/components/team/team_info.tsx
+++ b/ui/litellm-dashboard/src/components/team/team_info.tsx
@@ -35,6 +35,7 @@ import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_
 import { isAdminRole } from "@/utils/roles";
 import ObjectPermissionsView from "../object_permissions_view";
 import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
+import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
 import PremiumVectorStoreSelector from "../common_components/PremiumVectorStoreSelector";
 
 export interface TeamData {
@@ -235,10 +236,11 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
       };
 
       // Handle object_permission updates
-      if (values.vector_stores !== undefined) {
+      if (values.vector_stores !== undefined || values.mcp_servers !== undefined) {
         updateData.object_permission = {
           ...teamData?.team_info.object_permission,
-          vector_stores: values.vector_stores || []
+          vector_stores: values.vector_stores || [],
+          mcp_servers: values.mcp_servers || []
         };
       }
       
@@ -385,7 +387,8 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                     guardrails: info.metadata?.guardrails || [],
                     metadata: info.metadata ? JSON.stringify(info.metadata, null, 2) : "",
                     organization_id: info.organization_id,
-                    vector_stores: info.object_permission?.vector_stores || []
+                    vector_stores: info.object_permission?.vector_stores || [],
+                    mcp_servers: info.object_permission?.mcp_servers || []
                   }}
                   layout="vertical"
                 >
@@ -466,6 +469,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
                       placeholder="Select vector stores"
                     />
                   </Form.Item>
+
+                  <Form.Item label="MCP Servers" name="mcp_servers">
+                    <MCPServerSelector
+                      onChange={(values) => form.setFieldValue('mcp_servers', values)}
+                      value={form.getFieldValue('mcp_servers')}
+                      accessToken={accessToken || ""}
+                      placeholder="Select MCP servers"
+                    />
+                  </Form.Item>
                   
                   <Form.Item label="Organization ID" name="organization_id">
                     <Input type=""/>