Merge pull request #19468 from BerriAI/litellm_ui_team_select_001

[Fix] UI - MCP Servers: MCP Tools Tab Resetting to Overview

ui/litellm-dashboard/src/components/mcp_tools/mcp_server_view.tsx

@@ -36,6 +36,8 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
   const [editing, setEditing] = useState(isEditing);
   const [showFullUrl, setShowFullUrl] = useState(false);
   const [copiedStates, setCopiedStates] = useState<Record<string, boolean>>({});
+  const [selectedTabIndex, setSelectedTabIndex] = useState(0);
+
   const handleSuccess = (updated: MCPServer) => {
     setEditing(false);
     onBack();
@@ -72,11 +74,10 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
               size="small"
               icon={copiedStates["mcp-server_name"] ? <CheckIcon size={12} /> : <CopyIcon size={12} />}
               onClick={() => copyToClipboard(mcpServer.server_name, "mcp-server_name")}
-              className={`left-2 z-10 transition-all duration-200 ${
-                copiedStates["mcp-server_name"]
-                  ? "text-green-600 bg-green-50 border-green-200"
-                  : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
-              }`}
+              className={`left-2 z-10 transition-all duration-200 ${copiedStates["mcp-server_name"]
+                ? "text-green-600 bg-green-50 border-green-200"
+                : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
+                }`}
             />
             {mcpServer.alias && (
               <>
@@ -87,11 +88,10 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
                   size="small"
                   icon={copiedStates["mcp-alias"] ? <CheckIcon size={12} /> : <CopyIcon size={12} />}
                   onClick={() => copyToClipboard(mcpServer.alias, "mcp-alias")}
-                  className={`left-2 z-10 transition-all duration-200 ${
-                    copiedStates["mcp-alias"]
-                      ? "text-green-600 bg-green-50 border-green-200"
-                      : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
-                  }`}
+                  className={`left-2 z-10 transition-all duration-200 ${copiedStates["mcp-alias"]
+                    ? "text-green-600 bg-green-50 border-green-200"
+                    : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
+                    }`}
                 />
               </>
             )}
@@ -103,18 +103,17 @@ export const MCPServerView: React.FC<MCPServerViewProps> = ({
               size="small"
               icon={copiedStates["mcp-server-id"] ? <CheckIcon size={12} /> : <CopyIcon size={12} />}
               onClick={() => copyToClipboard(mcpServer.server_id, "mcp-server-id")}
-              className={`left-2 z-10 transition-all duration-200 ${
-                copiedStates["mcp-server-id"]
-                  ? "text-green-600 bg-green-50 border-green-200"
-                  : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
-              }`}
+              className={`left-2 z-10 transition-all duration-200 ${copiedStates["mcp-server-id"]
+                ? "text-green-600 bg-green-50 border-green-200"
+                : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
+                }`}
             />
           </div>
         </div>
       </div>
 
       {/* TODO: magic number for index */}
-      <TabGroup defaultIndex={editing ? 2 : 0}>
+      <TabGroup index={selectedTabIndex} onIndexChange={setSelectedTabIndex}>
         <TabList className="mb-4">
           {[
             <Tab key="overview">Overview</Tab>,

ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.test.tsx

@@ -1,5 +1,5 @@
 import React from "react";
-import { render, waitFor } from "@testing-library/react";
+import { render, waitFor, screen, fireEvent, act } from "@testing-library/react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import MCPServers from "./mcp_servers";
@@ -208,7 +208,7 @@ describe("MCPServers", () => {
     vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers);
     // Mock health check to never resolve (to test loading state)
     vi.mocked(networking.fetchMCPServerHealth).mockImplementation(
-      () => new Promise(() => {}), // Never resolves
+      () => new Promise(() => { }), // Never resolves
     );
 
     const queryClient = createQueryClient();
@@ -228,4 +228,120 @@ describe("MCPServers", () => {
       expect(networking.fetchMCPServerHealth).toHaveBeenCalled();
     });
   });
+
+  it("should filter servers by team when a team is selected", async () => {
+    // Mock MCP servers with different teams
+    const mockServers = [
+      {
+        server_id: "server-1",
+        server_name: "Team A Server",
+        alias: "team-a-server",
+        url: "https://example.com/mcp",
+        transport: "http",
+        auth_type: "none",
+        created_at: "2024-01-01T00:00:00Z",
+        created_by: "user-1",
+        updated_at: "2024-01-01T00:00:00Z",
+        updated_by: "user-1",
+        teams: [{ team_id: "team-a", team_alias: "Team A" }],
+        mcp_access_groups: [],
+      },
+      {
+        server_id: "server-2",
+        server_name: "Team B Server",
+        alias: "team-b-server",
+        url: "https://example2.com/mcp",
+        transport: "sse",
+        auth_type: "api_key",
+        created_at: "2024-01-02T00:00:00Z",
+        created_by: "user-2",
+        updated_at: "2024-01-02T00:00:00Z",
+        updated_by: "user-2",
+        teams: [{ team_id: "team-b", team_alias: "Team B" }],
+        mcp_access_groups: [],
+      },
+      {
+        server_id: "server-3",
+        server_name: "Team A Server 2",
+        alias: "team-a-server-2",
+        url: "https://example3.com/mcp",
+        transport: "http",
+        auth_type: "none",
+        created_at: "2024-01-03T00:00:00Z",
+        created_by: "user-1",
+        updated_at: "2024-01-03T00:00:00Z",
+        updated_by: "user-1",
+        teams: [{ team_id: "team-a", team_alias: "Team A" }],
+        mcp_access_groups: [],
+      },
+    ];
+
+    vi.mocked(networking.fetchMCPServers).mockResolvedValue(mockServers);
+    vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]);
+
+    const queryClient = createQueryClient();
+    render(
+      <QueryClientProvider client={queryClient}>
+        <MCPServers {...defaultProps} />
+      </QueryClientProvider>,
+    );
+
+    // Wait for the component to load
+    await waitFor(() => {
+      expect(screen.getByText("MCP Servers")).toBeInTheDocument();
+    });
+
+    // Wait for servers to be rendered
+    await waitFor(() => {
+      expect(screen.getByText("Team A Server")).toBeInTheDocument();
+    });
+
+    // Verify all servers are initially displayed
+    expect(screen.getByText("Team A Server")).toBeInTheDocument();
+    expect(screen.getByText("Team B Server")).toBeInTheDocument();
+    expect(screen.getByText("Team A Server 2")).toBeInTheDocument();
+
+    // Find the team select dropdown by looking for the "Current Team:" label
+    const teamLabel = screen.getByText("Current Team:");
+    const teamSelectContainer = teamLabel.closest("div")?.querySelector(".ant-select");
+    expect(teamSelectContainer).toBeTruthy();
+
+    // Open the dropdown by clicking on the selector
+    const selectSelector = teamSelectContainer?.querySelector(".ant-select-selector");
+    expect(selectSelector).toBeTruthy();
+
+    act(() => {
+      fireEvent.mouseDown(selectSelector!);
+    });
+
+    // Wait for dropdown to open
+    await waitFor(
+      () => {
+        const dropdownOptions = document.querySelectorAll(".ant-select-item-option");
+        expect(dropdownOptions.length).toBeGreaterThan(0);
+      },
+      { timeout: 5000 },
+    );
+
+    // Find and click on "Team A" option
+    const dropdownOptions = document.querySelectorAll(".ant-select-item-option");
+    const teamAOption = Array.from(dropdownOptions).find((option) =>
+      option.textContent?.includes("Team A"),
+    );
+    expect(teamAOption).toBeTruthy();
+
+    act(() => {
+      fireEvent.click(teamAOption!);
+    });
+
+    // Wait for filtering to complete
+    await waitFor(() => {
+      // Team A servers should still be visible
+      expect(screen.getByText("Team A Server")).toBeInTheDocument();
+      expect(screen.getByText("Team A Server 2")).toBeInTheDocument();
+    });
+
+    // Team B server should not be visible
+    expect(screen.queryByText("Team B Server")).not.toBeInTheDocument();
+  });
 });

ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx

@@ -2,7 +2,7 @@ import { isAdminRole } from "@/utils/roles";
 import { QuestionCircleOutlined } from "@ant-design/icons";
 import { Button, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react";
 import { Descriptions, Modal, Select, Tooltip, Typography } from "antd";
-import React, { useEffect, useState, useMemo } from "react";
+import React, { useEffect, useState, useMemo, useCallback } from "react";
 import { useMCPServers } from "../../app/(dashboard)/hooks/mcpServers/useMCPServers";
 import { useMCPServerHealth } from "../../app/(dashboard)/hooks/mcpServers/useMCPServerHealth";
 import NotificationsManager from "../molecules/notifications_manager";
@@ -115,20 +115,8 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
     );
   }, [serversWithHealth]);
 
-  // Handle team filter change
-  const handleTeamChange = (teamId: string) => {
-    setSelectedTeam(teamId);
-    filterServers(teamId, selectedMcpAccessGroup);
-  };
-
-  // Handle MCP access group filter change
-  const handleMcpAccessGroupChange = (group: string) => {
-    setSelectedMcpAccessGroup(group);
-    filterServers(selectedTeam, group);
-  };
-
   // Filtering logic for both team and access group
-  const filterServers = (teamId: string, group: string) => {
+  const filterServers = useCallback((teamId: string, group: string) => {
     if (!serversWithHealth) return setFilteredServers([]);
     let filtered = serversWithHealth;
     if (teamId === "personal") {
@@ -144,12 +132,24 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
       );
     }
     setFilteredServers(filtered);
+  }, [serversWithHealth]);
+
+  // Handle team filter change
+  const handleTeamChange = (teamId: string) => {
+    setSelectedTeam(teamId);
+    filterServers(teamId, selectedMcpAccessGroup);
+  };
+
+  // Handle MCP access group filter change
+  const handleMcpAccessGroupChange = (group: string) => {
+    setSelectedMcpAccessGroup(group);
+    filterServers(selectedTeam, group);
   };
 
   // Initial and effect-based filtering (trigger on query data updates and health data updates)
   useEffect(() => {
     filterServers(selectedTeam, selectedMcpAccessGroup);
-  }, [serversWithHealth, selectedTeam, selectedMcpAccessGroup]);
+  }, [serversWithHealth, selectedTeam, selectedMcpAccessGroup, filterServers]);
 
   const columns = React.useMemo(
     () =>
@@ -207,109 +207,34 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
     setModalVisible(false);
   };
 
+  // Memoize the selected server to prevent unnecessary re-renders
+  const selectedServer = React.useMemo(() => {
+    return filteredServers.find((server: MCPServer) => server.server_id === selectedServerId) || {
+      server_id: "",
+      server_name: "",
+      alias: "",
+      url: "",
+      transport: "",
+      auth_type: "",
+      created_at: "",
+      created_by: "",
+      updated_at: "",
+      updated_by: "",
+    };
+  }, [filteredServers, selectedServerId]);
+
+  // Memoize the onBack callback to prevent unnecessary re-renders
+  const handleBack = React.useCallback(() => {
+    setEditServer(false);
+    setSelectedServerId(null);
+    refetch();
+  }, [refetch]);
+
   if (!accessToken || !userRole || !userID) {
     console.log("Missing required authentication parameters", { accessToken, userRole, userID });
     return <div className="p-6 text-center text-gray-500">Missing required authentication parameters.</div>;
   }
 
-  const ServersTab = () =>
-    selectedServerId ? (
-      <MCPServerView
-        mcpServer={
-          filteredServers.find((server: MCPServer) => server.server_id === selectedServerId) || {
-            server_id: "",
-            server_name: "",
-            alias: "",
-            url: "",
-            transport: "",
-            auth_type: "",
-            created_at: "",
-            created_by: "",
-            updated_at: "",
-            updated_by: "",
-          }
-        }
-        onBack={() => {
-          setEditServer(false);
-          setSelectedServerId(null);
-          refetch();
-        }}
-        isProxyAdmin={isAdminRole(userRole)}
-        isEditing={editServer}
-        accessToken={accessToken}
-        userID={userID}
-        userRole={userRole}
-        availableAccessGroups={uniqueMcpAccessGroups}
-      />
-    ) : (
-      <div className="w-full h-full">
-        <div className="w-full px-6">
-          <div className="flex flex-col space-y-4">
-            <div className="flex items-center justify-between bg-gray-50 rounded-lg p-4 border-2 border-gray-200">
-              <div className="flex items-center gap-4">
-                <Text className="text-lg font-semibold text-gray-900">Current Team:</Text>
-                <Select value={selectedTeam} onChange={handleTeamChange} style={{ width: 300 }}>
-                  <Option value="all">
-                    <div className="flex items-center gap-2">
-                      <div className="w-2 h-2 bg-blue-500 rounded-full"></div>
-                      <span className="font-medium">{isInternalUser ? "All Available Servers" : "All Servers"}</span>
-                    </div>
-                  </Option>
-                  <Option value="personal">
-                    <div className="flex items-center gap-2">
-                      <div className="w-2 h-2 bg-green-500 rounded-full"></div>
-                      <span className="font-medium">Personal</span>
-                    </div>
-                  </Option>
-                  {uniqueTeams.map((team) => (
-                    <Option key={team.team_id} value={team.team_id}>
-                      <div className="flex items-center gap-2">
-                        <div className="w-2 h-2 bg-green-500 rounded-full"></div>
-                        <span className="font-medium">{team.team_alias || team.team_id}</span>
-                      </div>
-                    </Option>
-                  ))}
-                </Select>
-                <Text className="text-lg font-semibold text-gray-900 ml-6">
-                  Access Group:
-                  <Tooltip title="An MCP Access Group is a set of users or teams that have permission to access specific MCP servers. Use access groups to control and organize who can connect to which servers.">
-                    <QuestionCircleOutlined style={{ marginLeft: 4, color: "#888" }} />
-                  </Tooltip>
-                </Text>
-                <Select value={selectedMcpAccessGroup} onChange={handleMcpAccessGroupChange} style={{ width: 300 }}>
-                  <Option value="all">
-                    <div className="flex items-center gap-2">
-                      <div className="w-2 h-2 bg-blue-500 rounded-full"></div>
-                      <span className="font-medium">All Access Groups</span>
-                    </div>
-                  </Option>
-                  {uniqueMcpAccessGroups.map((group) => (
-                    <Option key={group} value={group}>
-                      <div className="flex items-center gap-2">
-                        <div className="w-2 h-2 bg-green-500 rounded-full"></div>
-                        <span className="font-medium">{group}</span>
-                      </div>
-                    </Option>
-                  ))}
-                </Select>
-              </div>
-            </div>
-          </div>
-        </div>
-        <div className="w-full px-6 mt-6">
-          <DataTable
-            data={filteredServers}
-            columns={columns}
-            renderSubComponent={() => <div></div>}
-            getRowCanExpand={() => false}
-            isLoading={isLoadingServers}
-            noDataMessage="No MCP servers configured"
-            loadingMessage="🚅 Loading MCP servers..."
-          />
-        </div>
-      </div>
-    );
-
   return (
     <div className="w-full h-full p-6">
       <Modal
@@ -381,7 +306,86 @@ const MCPServers: React.FC<MCPServerProps> = ({ accessToken, userRole, userID })
         </TabList>
         <TabPanels>
           <TabPanel>
-            <ServersTab />
+            {selectedServerId ? (
+              <MCPServerView
+                key={selectedServerId}
+                mcpServer={selectedServer}
+                onBack={handleBack}
+                isProxyAdmin={isAdminRole(userRole)}
+                isEditing={editServer}
+                accessToken={accessToken}
+                userID={userID}
+                userRole={userRole}
+                availableAccessGroups={uniqueMcpAccessGroups}
+              />
+            ) : (
+              <div className="w-full h-full">
+                <div className="w-full px-6">
+                  <div className="flex flex-col space-y-4">
+                    <div className="flex items-center justify-between bg-gray-50 rounded-lg p-4 border-2 border-gray-200">
+                      <div className="flex items-center gap-4">
+                        <Text className="text-lg font-semibold text-gray-900">Current Team:</Text>
+                        <Select value={selectedTeam} onChange={handleTeamChange} style={{ width: 300 }}>
+                          <Option value="all">
+                            <div className="flex items-center gap-2">
+                              <div className="w-2 h-2 bg-blue-500 rounded-full"></div>
+                              <span className="font-medium">{isInternalUser ? "All Available Servers" : "All Servers"}</span>
+                            </div>
+                          </Option>
+                          <Option value="personal">
+                            <div className="flex items-center gap-2">
+                              <div className="w-2 h-2 bg-green-500 rounded-full"></div>
+                              <span className="font-medium">Personal</span>
+                            </div>
+                          </Option>
+                          {uniqueTeams.map((team) => (
+                            <Option key={team.team_id} value={team.team_id}>
+                              <div className="flex items-center gap-2">
+                                <div className="w-2 h-2 bg-green-500 rounded-full"></div>
+                                <span className="font-medium">{team.team_alias || team.team_id}</span>
+                              </div>
+                            </Option>
+                          ))}
+                        </Select>
+                        <Text className="text-lg font-semibold text-gray-900 ml-6">
+                          Access Group:
+                          <Tooltip title="An MCP Access Group is a set of users or teams that have permission to access specific MCP servers. Use access groups to control and organize who can connect to which servers.">
+                            <QuestionCircleOutlined style={{ marginLeft: 4, color: "#888" }} />
+                          </Tooltip>
+                        </Text>
+                        <Select value={selectedMcpAccessGroup} onChange={handleMcpAccessGroupChange} style={{ width: 300 }}>
+                          <Option value="all">
+                            <div className="flex items-center gap-2">
+                              <div className="w-2 h-2 bg-blue-500 rounded-full"></div>
+                              <span className="font-medium">All Access Groups</span>
+                            </div>
+                          </Option>
+                          {uniqueMcpAccessGroups.map((group) => (
+                            <Option key={group} value={group}>
+                              <div className="flex items-center gap-2">
+                                <div className="w-2 h-2 bg-green-500 rounded-full"></div>
+                                <span className="font-medium">{group}</span>
+                              </div>
+                            </Option>
+                          ))}
+                        </Select>
+                      </div>
+                    </div>
+                  </div>
+                </div>
+                <div className="w-full px-6 mt-6">
+                  <DataTable
+                    data={filteredServers}
+                    columns={columns}
+                    renderSubComponent={() => <div></div>}
+                    getRowCanExpand={() => false}
+                    isLoading={isLoadingServers}
+                    noDataMessage="No MCP servers configured"
+                    loadingMessage="🚅 Loading MCP servers..."
+                  />
+                </div>
+              </div>
+            )}
           </TabPanel>
           <TabPanel>
             <MCPConnect />

ui/litellm-dashboard/src/components/mcp_tools/mcp_tools.tsx

@@ -127,11 +127,10 @@ const MCPToolsViewer = ({
                     {toolsData.map((tool: MCPTool) => (
                       <div
                         key={tool.name}
-                        className={`border rounded-lg p-3 cursor-pointer transition-all hover:shadow-sm ${
-                          selectedTool?.name === tool.name
+                        className={`border rounded-lg p-3 cursor-pointer transition-all hover:shadow-sm ${selectedTool?.name === tool.name
                             ? "border-blue-500 bg-blue-50 ring-1 ring-blue-200"
                             : "border-gray-200 bg-white hover:border-gray-300"
-                        }`}
+                          }`}
                         onClick={() => {
                           setSelectedTool(tool);
                           setToolResult(null);