From 716aaacad156532ae8994ffcd271d93099c8c9b2 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Sat, 21 Mar 2026 19:00:49 -0700
Subject: [PATCH] feat(mcp): add toolsets (purple) as third option type in
 MCPServerSelector

---
 .../MCPServerSelector.tsx                     | 58 ++++++++++++++-----
 1 file changed, 45 insertions(+), 13 deletions(-)

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 dc4ed25786b..fc4b20517cf 100644
--- a/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx
+++ b/ui/litellm-dashboard/src/components/mcp_server_management/MCPServerSelector.tsx
@@ -1,13 +1,15 @@
 import { useMCPAccessGroups } from "@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups";
 import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
+import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets";
 import { Select } from "antd";
 import React from "react";
 
 interface MCPServerSelectorProps {
-  onChange: (selected: { servers: string[]; accessGroups: string[] }) => void;
+  onChange: (selected: { servers: string[]; accessGroups: string[]; toolsets: string[] }) => void;
   value?: {
     servers: string[];
     accessGroups: string[];
+    toolsets?: string[];
   };
   className?: string;
   accessToken: string;
@@ -16,6 +18,8 @@ interface MCPServerSelectorProps {
   teamId?: string | null;
 }
 
+const TOOLSET_PREFIX = "toolset:";
+
 const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
   onChange,
   value,
@@ -27,33 +31,61 @@ const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
 }) => {
   const { data: mcpServers = [], isLoading: serversLoading } = useMCPServers(teamId);
   const { data: accessGroups = [], isLoading: groupsLoading } = useMCPAccessGroups();
+  const { data: toolsets = [], isLoading: toolsetsLoading } = useMCPToolsets();
 
-  const loading = serversLoading || groupsLoading;
+  const loading = serversLoading || groupsLoading || toolsetsLoading;
 
-  // Combine options, access groups first
+  const accessGroupSet = new Set(accessGroups);
+
+  // Combine options: access groups (green) + servers (blue) + toolsets (purple)
   const options = [
     ...accessGroups.map((group) => ({
       label: group,
       value: group,
-      isAccessGroup: true,
+      type: "accessGroup" as const,
       searchText: `${group} Access Group`,
     })),
     ...mcpServers.map((server) => ({
       label: `${server.server_name || server.server_id} (${server.server_id})`,
       value: server.server_id,
-      isAccessGroup: false,
+      type: "server" as const,
       searchText: `${server.server_name || server.server_id} ${server.server_id} MCP Server`,
     })),
+    ...toolsets.map((toolset) => ({
+      label: toolset.toolset_name,
+      value: `${TOOLSET_PREFIX}${toolset.toolset_id}`,
+      type: "toolset" as const,
+      searchText: `${toolset.toolset_name} ${toolset.toolset_id} Toolset`,
+    })),
   ];
 
-  // Flatten value for Select
-  const selectedValues = [...(value?.servers || []), ...(value?.accessGroups || [])];
+  const colorByType: Record<string, string> = {
+    accessGroup: "#52c41a",
+    server: "#1890ff",
+    toolset: "#722ed1",
+  };
+  const labelByType: Record<string, string> = {
+    accessGroup: "Access Group",
+    server: "MCP Server",
+    toolset: "Toolset",
+  };
+
+  // Flatten value for Select — prefix toolset IDs
+  const selectedValues = [
+    ...(value?.servers || []),
+    ...(value?.accessGroups || []),
+    ...(value?.toolsets || []).map((id) => `${TOOLSET_PREFIX}${id}`),
+  ];
 
   // Handle selection
   const handleChange = (selected: string[]) => {
-    const servers = selected.filter((v) => !accessGroups.includes(v));
-    const accessGroupsSelected = selected.filter((v) => accessGroups.includes(v));
-    onChange({ servers, accessGroups: accessGroupsSelected });
+    const toolsetsSelected = selected
+      .filter((v) => v.startsWith(TOOLSET_PREFIX))
+      .map((v) => v.slice(TOOLSET_PREFIX.length));
+    const rest = selected.filter((v) => !v.startsWith(TOOLSET_PREFIX));
+    const servers = rest.filter((v) => !accessGroupSet.has(v));
+    const accessGroupsSelected = rest.filter((v) => accessGroupSet.has(v));
+    onChange({ servers, accessGroups: accessGroupsSelected, toolsets: toolsetsSelected });
   };
 
   return (
@@ -83,20 +115,20 @@ const MCPServerSelector: React.FC<MCPServerSelectorProps> = ({
                   width: 8,
                   height: 8,
                   borderRadius: "50%",
-                  background: opt.isAccessGroup ? "#52c41a" : "#1890ff",
+                  background: colorByType[opt.type],
                   flexShrink: 0,
                 }}
               />
               <span style={{ flex: 1 }}>{opt.label}</span>
               <span
                 style={{
-                  color: opt.isAccessGroup ? "#52c41a" : "#1890ff",
+                  color: colorByType[opt.type],
                   fontSize: "12px",
                   fontWeight: 500,
                   opacity: 0.8,
                 }}
               >
-                {opt.isAccessGroup ? "Access Group" : "MCP Server"}
+                {labelByType[opt.type]}
               </span>
             </div>
           </Select.Option>