diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 7686cc05fa6..f4fffa4ffbc 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1137,14 +1137,11 @@ "count": 1 }, "no-nested-ternary": { - "count": 7 + "count": 6 }, "no-restricted-imports": { "count": 2 }, - "prefer-const": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 4 }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index 9da3e3a4a08..2972d7d4a17 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -1,7 +1,11 @@ -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import { renderWithProviders as render, testQueryClient } from "@/../tests/test-utils"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ChatUI from "./ChatUI"; import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; +import * as networking from "@/components/networking"; +import { PLAYGROUND_TEAM_KEY_DURATION, playgroundTeamKeyStorageKey } from "../../hooks/usePlaygroundTeamSession"; +import { getSecureItem } from "@/utils/secureStorage"; // Mock the fetchAvailableModels function vi.mock("@/components/llm_calls/fetch_models", () => ({ @@ -14,6 +18,8 @@ vi.mock("@/components/networking", () => ({ vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), getGuardrailsList: vi.fn().mockResolvedValue({ data: [] }), modelHubCall: vi.fn().mockResolvedValue({ data: [] }), + teamListCall: vi.fn().mockResolvedValue([]), + keyCreateCall: vi.fn().mockResolvedValue({ key: "sk-team-key" }), })); // Mock scrollIntoView which is not available in jsdom @@ -26,6 +32,7 @@ describe("ChatUI", () => { // Reset mocks before each test vi.clearAllMocks(); sessionStorage.clear(); + testQueryClient.clear(); // Mock scrollIntoView which is not available in JSDOM Element.prototype.scrollIntoView = vi.fn(); @@ -413,4 +420,112 @@ describe("ChatUI", () => { } } }); + + describe("team scoped session", () => { + const clickOption = async (label: string) => { + await waitFor(() => { + const option = Array.from(document.querySelectorAll(".ant-select-item-option")).find( + (el) => el.getAttribute("title") === label, + ); + expect(option).toBeTruthy(); + }); + const option = Array.from(document.querySelectorAll(".ant-select-item-option")).find( + (el) => el.getAttribute("title") === label, + ); + await act(async () => { + fireEvent.click(option!); + }); + }; + + const selectTeamKeySource = async () => { + const keySourceSelect = screen + .getByText("Virtual Key Source") + .parentElement?.querySelector(".ant-select-selector"); + expect(keySourceSelect).toBeTruthy(); + + await act(async () => { + fireEvent.mouseDown(keySourceSelect!); + }); + await clickOption("Team"); + }; + + const selectTeam = async (label: string) => { + const teamSelect = screen.getByTestId("playground-team-select").querySelector(".ant-select-selector"); + expect(teamSelect).toBeTruthy(); + + await act(async () => { + fireEvent.mouseDown(teamSelect!); + }); + await clickOption(label); + }; + + beforeEach(() => { + (networking.teamListCall as any).mockResolvedValue([ + { team_id: "team-alpha", team_alias: "Alpha Team" }, + { team_id: "team-beta", team_alias: "Beta Team" }, + ]); + (networking.keyCreateCall as any).mockResolvedValue({ key: "sk-team-alpha-key" }); + }); + + const renderChatUI = () => + render( + , + ); + + it("mints a team scoped key and lists the team's models with it", async () => { + renderChatUI(); + await waitFor(() => expect(screen.getByText("Test Key")).toBeInTheDocument()); + + await selectTeamKeySource(); + await waitFor(() => expect(networking.teamListCall).toHaveBeenCalledWith("ui-session-key", null, "user-1")); + + (fetchModelsModule.fetchAvailableModels as any).mockResolvedValue([{ model_group: "team-only-model" }]); + await selectTeam("Alpha Team"); + + await waitFor(() => + expect(networking.keyCreateCall).toHaveBeenCalledWith( + "ui-session-key", + "user-1", + expect.objectContaining({ team_id: "team-alpha", duration: PLAYGROUND_TEAM_KEY_DURATION }), + ), + ); + await waitFor(() => expect(fetchModelsModule.fetchAvailableModels).toHaveBeenCalledWith("sk-team-alpha-key")); + expect(getSecureItem(playgroundTeamKeyStorageKey("team-alpha"))).toBe("sk-team-alpha-key"); + }); + + it("reuses a cached team key instead of minting another one", async () => { + renderChatUI(); + await waitFor(() => expect(screen.getByText("Test Key")).toBeInTheDocument()); + + await selectTeamKeySource(); + await selectTeam("Alpha Team"); + await waitFor(() => expect(networking.keyCreateCall).toHaveBeenCalledTimes(1)); + + await selectTeam("Beta Team"); + await waitFor(() => expect(networking.keyCreateCall).toHaveBeenCalledTimes(2)); + + await selectTeam("Alpha Team"); + await waitFor(() => expect(fetchModelsModule.fetchAvailableModels).toHaveBeenCalledWith("sk-team-alpha-key")); + expect(networking.keyCreateCall).toHaveBeenCalledTimes(2); + }); + + it("surfaces an error and keeps the model list empty when the team key cannot be minted", async () => { + (networking.keyCreateCall as any).mockRejectedValue(new Error("not a member of team")); + + renderChatUI(); + await waitFor(() => expect(screen.getByText("Test Key")).toBeInTheDocument()); + + await selectTeamKeySource(); + await selectTeam("Alpha Team"); + + await waitFor(() => expect(screen.getByText("not a member of team")).toBeInTheDocument()); + expect(getSecureItem(playgroundTeamKeyStorageKey("team-alpha"))).toBeNull(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index d2cf27e0c8b..63f2a3d2909 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -23,7 +23,7 @@ import { } from "@ant-design/icons"; import { Card, Text, TextInput, Title, Button as TremorButton } from "@tremor/react"; import { Button, Input, Modal, Popover, Select, Spin, Tooltip, Typography, Upload } from "antd"; -import React, { useEffect, useRef, useState } from "react"; +import React, { useEffect, useMemo, useRef, useState } from "react"; import ReactMarkdown from "react-markdown"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; @@ -76,6 +76,7 @@ import RealtimePlayground from "./RealtimePlayground"; import { A2ATaskMetadata, MessageType } from "@/components/chat_ui/types"; import { useCodeInterpreter } from "../../hooks/useCodeInterpreter"; import { useChatHistory } from "../../hooks/useChatHistory"; +import { usePlaygroundTeamSession } from "../../hooks/usePlaygroundTeamSession"; import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; @@ -107,6 +108,8 @@ const MCP_SUPPORTED_ENDPOINTS = new Set([ const CUSTOM_MODEL_DEBOUNCE_WAIT_MS = 500; +type ApiKeySource = "session" | "team" | "custom"; + const ChatUI: React.FC = ({ accessToken, token, @@ -172,11 +175,11 @@ const ChatUI: React.FC = ({ clearMCPEvents, } = useChatHistory({ simplified }); // codeql[js/clear-text-storage-of-sensitive-data] - const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => { + const [apiKeySource, setApiKeySource] = useState(() => { const saved = getSecureItem("apiKeySource"); if (saved) { try { - return JSON.parse(saved) as "session" | "custom"; + return JSON.parse(saved) as ApiKeySource; } catch (error) { console.error("Error parsing apiKeySource from sessionStorage", error); } @@ -184,6 +187,9 @@ const ChatUI: React.FC = ({ return disabledPersonalKeyCreation ? "custom" : "session"; }); const [apiKey, setApiKey] = useState(() => getSecureItem("apiKey") || ""); + const [selectedTeamId, setSelectedTeamId] = useState( + () => sessionStorage.getItem("playgroundTeamId") || null, + ); const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState( () => sessionStorage.getItem("customProxyBaseUrl") || "", ); @@ -267,9 +273,29 @@ const ChatUI: React.FC = ({ const chatEndRef = useRef(null); + const { + teams, + teamKey, + isLoadingTeams, + isMintingKey, + error: teamSessionError, + } = usePlaygroundTeamSession({ + accessToken, + userID, + teamId: selectedTeamId, + enabled: !simplified && apiKeySource === "team", + }); + + const effectiveApiKey = useMemo(() => { + if (simplified) return accessToken; + if (apiKeySource === "team") return teamKey; + if (apiKeySource === "custom") return apiKey || null; + return accessToken; + }, [simplified, apiKeySource, accessToken, apiKey, teamKey]); + // Fetch MCP servers and toolsets const loadMCPServers = async () => { - const userApiKey = apiKeySource === "session" ? accessToken : apiKey; + const userApiKey = effectiveApiKey; if (!userApiKey) return; setIsLoadingMCPServers(true); @@ -297,7 +323,7 @@ const ChatUI: React.FC = ({ // Fetch tools for a specific server const loadServerTools = async (serverId: string) => { - const userApiKey = apiKeySource === "session" ? accessToken : apiKey; + const userApiKey = effectiveApiKey; if (!userApiKey || serverToolsMap[serverId]) return; try { @@ -314,9 +340,7 @@ const ChatUI: React.FC = ({ useEffect(() => { if (isGetCodeModalVisible) { const code = generateCodeSnippet({ - apiKeySource, - accessToken, - apiKey, + apiKey: effectiveApiKey, inputMessage, chatHistory, selectedTags, @@ -337,9 +361,7 @@ const ChatUI: React.FC = ({ }, [ isGetCodeModalVisible, selectedSdk, - apiKeySource, - accessToken, - apiKey, + effectiveApiKey, inputMessage, chatHistory, selectedTags, @@ -370,6 +392,11 @@ const ChatUI: React.FC = ({ sessionStorage.setItem("mcpServerToolRestrictions", JSON.stringify(mcpServerToolRestrictions)); sessionStorage.setItem("selectedVoice", selectedVoice); sessionStorage.removeItem("selectedMCPTools"); // Clean up old key + if (selectedTeamId) { + sessionStorage.setItem("playgroundTeamId", selectedTeamId); + } else { + sessionStorage.removeItem("playgroundTeamId"); + } if (!simplified) { if (selectedModel) { @@ -383,6 +410,7 @@ const ChatUI: React.FC = ({ simplified, apiKeySource, apiKey, + selectedTeamId, selectedModel, endpointType, selectedTags, @@ -395,8 +423,9 @@ const ChatUI: React.FC = ({ ]); useEffect(() => { - let userApiKey = apiKeySource === "session" ? accessToken : apiKey; + const userApiKey = effectiveApiKey; if (!userApiKey || !token || !userRole || !userID) { + setModelInfo([]); return; } @@ -426,7 +455,7 @@ const ChatUI: React.FC = ({ loadModels(); } loadMCPServers(); - }, [accessToken, userID, userRole, apiKeySource, apiKey, token, simplified]); + }, [effectiveApiKey, userID, userRole, token, simplified]); // Load tools when MCP direct mode has a server (or toolset) selected useEffect(() => { @@ -450,7 +479,7 @@ const ChatUI: React.FC = ({ // Fetch agents when A2A endpoint is selected useEffect(() => { - const userApiKey = apiKeySource === "session" ? accessToken : apiKey; + const userApiKey = effectiveApiKey; if (!userApiKey || endpointType !== EndpointType.A2A_AGENTS) { return; } @@ -469,7 +498,7 @@ const ChatUI: React.FC = ({ }; loadAgents(); - }, [accessToken, apiKeySource, apiKey, endpointType, customProxyBaseUrl, selectedAgent]); + }, [effectiveApiKey, endpointType, customProxyBaseUrl, selectedAgent]); useEffect(() => { // Scroll to the bottom of the chat whenever chatHistory updates @@ -650,10 +679,12 @@ const ChatUI: React.FC = ({ return; } - const effectiveApiKey = simplified ? accessToken : apiKeySource === "session" ? accessToken : apiKey; - if (!effectiveApiKey) { - NotificationsManager.fromBackend("Please provide a Virtual Key or select Current UI Session"); + NotificationsManager.fromBackend( + apiKeySource === "team" + ? "Please select a team to chat as" + : "Please provide a Virtual Key or select Current UI Session", + ); return; } @@ -1051,14 +1082,14 @@ const ChatUI: React.FC = ({ Virtual Key Source setSelectedTeamId(value)} + showSearch + optionFilterProp="label" + options={teams.map((team) => ({ + value: team.team_id, + label: team.team_alias || team.team_id, + }))} + notFoundContent={isLoadingTeams ? : "You are not a member of any team"} + /> + + Requests use a short lived key scoped to this team, so the model list shows the team's + models. + + {teamSessionError && {teamSessionError}} + + )}
@@ -1694,7 +1750,7 @@ const ChatUI: React.FC = ({ {endpointType === EndpointType.RESPONSES && (
= ({
{endpointType === EndpointType.REALTIME ? ( 0 ? selectedGuardrails : undefined} @@ -1755,7 +1811,7 @@ const ChatUI: React.FC = ({ endpointType={endpointType as EndpointType} mcpEvents={mcpEvents} codeInterpreterResult={codeInterpreter.result} - accessToken={apiKeySource === "session" ? accessToken || "" : apiKey} + accessToken={effectiveApiKey || ""} />
))} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/usePlaygroundTeamSession.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/usePlaygroundTeamSession.ts new file mode 100644 index 00000000000..d41c1c5e688 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/hooks/usePlaygroundTeamSession.ts @@ -0,0 +1,102 @@ +"use client"; + +import { useQuery } from "@tanstack/react-query"; +import { keyCreateCall, teamListCall } from "@/components/networking"; +import { Team } from "@/components/key_team_helpers/key_list"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; + +export const PLAYGROUND_TEAM_KEY_DURATION = "24h"; + +export const playgroundTeamKeyStorageKey = (teamId: string) => `playgroundTeamKey:${teamId}`; + +interface UsePlaygroundTeamSessionArgs { + accessToken: string | null; + userID: string | null; + teamId: string | null; + enabled: boolean; +} + +interface PlaygroundTeamSession { + teams: Team[]; + teamKey: string | null; + isLoadingTeams: boolean; + isMintingKey: boolean; + error: string | null; +} + +const readCachedKey = (teamId: string): string | null => { + try { + return getSecureItem(playgroundTeamKeyStorageKey(teamId)); + } catch { + return null; + } +}; + +const cacheKey = (teamId: string, key: string): void => { + try { + setSecureItem(playgroundTeamKeyStorageKey(teamId), key); + } catch { + return; + } +}; + +const errorMessage = (error: unknown, fallback: string): string => + error instanceof Error && error.message ? error.message : fallback; + +export const usePlaygroundTeamSession = ({ + accessToken, + userID, + teamId, + enabled, +}: UsePlaygroundTeamSessionArgs): PlaygroundTeamSession => { + const teamsQuery = useQuery({ + queryKey: ["playgroundTeams", userID], + queryFn: async (): Promise => { + if (!accessToken) return []; + const response: unknown = await teamListCall(accessToken, null, userID); + return Array.isArray(response) ? (response as Team[]) : []; + }, + enabled: enabled && !!accessToken, + }); + + const teamKeyQuery = useQuery({ + queryKey: ["playgroundTeamKey", teamId], + queryFn: async (): Promise => { + if (!accessToken || !userID || !teamId) { + throw new Error("Sign in and select a team to chat as that team"); + } + const cached = readCachedKey(teamId); + if (cached) { + return cached; + } + const response: { key?: string } = await keyCreateCall(accessToken, userID, { + team_id: teamId, + duration: PLAYGROUND_TEAM_KEY_DURATION, + key_alias: `playground-${teamId}-${Date.now()}`, + }); + if (!response?.key) { + throw new Error("Creating a key for this team returned no key"); + } + cacheKey(teamId, response.key); + return response.key; + }, + enabled: enabled && !!accessToken && !!userID && !!teamId, + retry: false, + staleTime: Infinity, + }); + + let error: string | null = null; + if (teamsQuery.error) { + error = errorMessage(teamsQuery.error, "Failed to load teams"); + } else if (teamKeyQuery.error) { + error = errorMessage(teamKeyQuery.error, "Failed to create a key for this team"); + } + + return { + teams: teamsQuery.data ?? [], + teamKey: teamKeyQuery.data ?? null, + isLoadingTeams: teamsQuery.isFetching, + isMintingKey: teamKeyQuery.isFetching, + error, + }; +}; diff --git a/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx index 0113f7b2832..a9b5abbaa06 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.test.tsx @@ -27,8 +27,6 @@ describe("CodeSnippets", () => { endpointType: EndpointType.EMBEDDINGS, inputMessage: "Hello, world!", selectedModel: "text-embedding-3-small", - apiKeySource: "session" as const, - accessToken: "1234567890", apiKey: "1234567890", chatHistory: [], selectedTags: [], diff --git a/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx index 576b094b5d3..5397ae7aafd 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx +++ b/ui/litellm-dashboard/src/components/chat_ui/CodeSnippets.tsx @@ -10,9 +10,7 @@ interface CodeGenMetadata { } interface GenerateCodeParams { - apiKeySource: "session" | "custom"; - accessToken: string | null; - apiKey: string; + apiKey: string | null; inputMessage: string; chatHistory: MessageType[]; selectedTags: string[]; @@ -34,9 +32,7 @@ interface GenerateCodeParams { export const generateCodeSnippet = (params: GenerateCodeParams): string => { const { - apiKeySource, - accessToken, - apiKey, + apiKey: effectiveApiKey, inputMessage, chatHistory, selectedTags, @@ -52,7 +48,6 @@ export const generateCodeSnippet = (params: GenerateCodeParams): string => { selectedSdk, proxySettings, } = params; - const effectiveApiKey = apiKeySource === "session" ? accessToken : apiKey; // Determine base URL with priority: LITELLM_UI_API_DOC_BASE_URL > PROXY_BASE_URL > window.location.origin let apiBase = window.location.origin; diff --git a/ui/litellm-dashboard/src/components/public_model_hub.tsx b/ui/litellm-dashboard/src/components/public_model_hub.tsx index 2bdb3055835..8701d5a647f 100644 --- a/ui/litellm-dashboard/src/components/public_model_hub.tsx +++ b/ui/litellm-dashboard/src/components/public_model_hub.tsx @@ -1020,8 +1020,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded 
                     {(() => {
                       const codeSnippet = generateCodeSnippet({
-                        apiKeySource: "custom",
-                        accessToken: null,
                         apiKey: "your_api_key",
                         inputMessage: "Hello, how are you?",
                         chatHistory: [{ role: "user", content: "Hello, how are you?", isImage: false } as MessageType],
@@ -1042,8 +1040,6 @@ const PublicModelHub: React.FC = ({ accessToken, isEmbedded