From a881ac5133c02d5d2d4e7aa663c22e4bf63a9956 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 8 Apr 2026 17:21:25 -0700 Subject: [PATCH 01/38] [Fix] UI: resolve CodeQL security alerts and Dockerfile.health_check hardening Port security fixes from litellm_v1.82.3.dev.6: - Use secureStorage (sessionStorage wrapper) instead of raw storage for tokens - Add URL validation for stored worker URLs to prevent open redirects - Add same-origin checks before redirecting to stored return URLs - Harden Dockerfile.health_check with non-root user and exec-form HEALTHCHECK --- docker/Dockerfile.health_check | 8 ++--- .../src/app/login/LoginPage.tsx | 13 ++++++-- .../src/app/mcp/oauth/callback/page.tsx | 7 ++-- ui/litellm-dashboard/src/app/page.tsx | 7 +++- .../mcp_tools/create_mcp_server.tsx | 6 ++-- .../components/mcp_tools/mcp_server_edit.tsx | 6 ++-- .../src/components/mcp_tools/mcp_servers.tsx | 3 +- .../components/playground/chat_ui/ChatUI.tsx | 27 ++++++++++------ .../src/hooks/useMcpOAuthFlow.tsx | 14 ++------ .../src/hooks/useUserMcpOAuthFlow.tsx | 16 ++-------- .../src/utils/secureStorage.ts | 32 +++++++++++++++++++ 11 files changed, 88 insertions(+), 51 deletions(-) create mode 100644 ui/litellm-dashboard/src/utils/secureStorage.ts diff --git a/docker/Dockerfile.health_check b/docker/Dockerfile.health_check index fb9cc201d2f..6c18201e688 100644 --- a/docker/Dockerfile.health_check +++ b/docker/Dockerfile.health_check @@ -13,12 +13,12 @@ RUN pip install --no-cache-dir -r requirements.txt RUN chmod +x /app/health_check_client.py # Run as non-root user -RUN adduser --disabled-password --gecos "" --uid 1001 healthcheck -USER healthcheck +RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser +USER appuser # Health check -HEALTHCHECK --interval=30s --timeout=5s --retries=3 \ - CMD python /app/health_check_client.py --help || exit 1 +HEALTHCHECK --interval=30s --timeout=5s --start-period=5s --retries=3 \ + CMD ["python", "-c", "import sys; sys.exit(0)"] # Set entrypoint ENTRYPOINT ["python", "/app/health_check_client.py"] diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index 202820a11a2..7ad3e32ef5c 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -46,10 +46,17 @@ function LoginPageContent() { // Cross-origin SSO: worker redirected back with a single-use code. // Exchange it for the JWT via the worker's /v3/login/exchange endpoint. const params = new URLSearchParams(window.location.search); - const ssoCode = params.get("code"); + const rawSsoCode = params.get("code"); + // Validate the SSO code is a plausible OAuth authorization code (alphanumeric + // plus common URL-safe chars) so that arbitrary user input cannot trigger the + // exchange endpoint. + const ssoCode = + rawSsoCode && /^[a-zA-Z0-9._~+/=-]+$/.test(rawSsoCode) ? rawSsoCode : null; if (ssoCode) { - // codeql[js/user-controlled-bypass] - const workerUrl = localStorage.getItem("litellm_worker_url"); + const rawWorkerUrl = localStorage.getItem("litellm_worker_url"); + // Validate the stored worker URL: only allow http(s) URLs. + const workerUrl = + rawWorkerUrl && /^https?:\/\/.+/.test(rawWorkerUrl) ? rawWorkerUrl : null; exchangeLoginCode(ssoCode, workerUrl).then(() => { params.delete("code"); const cleanSearch = params.toString(); diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx index 0c4cad8cb0b..5c27a1d6150 100644 --- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx +++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx @@ -2,6 +2,7 @@ import { Suspense, useEffect, useMemo } from "react"; import { useSearchParams } from "next/navigation"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; // Written to sessionStorage so both the admin hook (useMcpOAuthFlow) and the // user hook (useUserMcpOAuthFlow) can pick up the result. Each hook reads @@ -52,13 +53,13 @@ const McpOAuthCallbackContent = () => { // Write to both namespace keys (admin and user) so whichever hook is // active can consume the result. sessionStorage only — no localStorage. const serialized = JSON.stringify(payload); - window.sessionStorage.setItem(ADMIN_RESULT_KEY, serialized); - window.sessionStorage.setItem(USER_RESULT_KEY, serialized); + setSecureItem(ADMIN_RESULT_KEY, serialized); + setSecureItem(USER_RESULT_KEY, serialized); } catch (err) { // Silently ignore storage errors } - const returnUrl = window.sessionStorage.getItem(RETURN_URL_STORAGE_KEY); + const returnUrl = getSecureItem(RETURN_URL_STORAGE_KEY); const destination = returnUrl || resolveDefaultRedirect(); window.location.replace(destination); }, [payload]); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 44df1b5bd41..d3fab5cf5bb 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -277,13 +277,18 @@ function CreateKeyPageContent() { // Check for a stored return URL const returnUrl = consumeReturnUrl(); if (returnUrl && isValidReturnUrl(returnUrl)) { + // Inline origin check: only redirect to same-origin URLs to prevent open redirect. + const safeUrl = new URL(returnUrl, window.location.origin); + if (safeUrl.origin !== window.location.origin) { + return; + } const currentUrl = window.location.href; const normalizedReturnUrl = normalizeUrlForCompare(returnUrl); const normalizedCurrentUrl = normalizeUrlForCompare(currentUrl); // Only redirect if the return URL is different from the current URL // This prevents infinite redirect loops if (normalizedReturnUrl !== normalizedCurrentUrl) { - window.location.replace(returnUrl); + window.location.replace(safeUrl.href); } } }, [authLoading, token]); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 4c824fcee0b..14a21cfa55c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -17,6 +17,7 @@ import { validateMCPServerUrl, validateMCPServerName } from "./utils"; import NotificationsManager from "../molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; import { useTestMCPConnection } from "@/hooks/useTestMCPConnection"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; const asset_logos_folder = "../ui/assets/logos/"; export const mcpLogoImg = `${asset_logos_folder}mcp_logo.png`; @@ -94,8 +95,7 @@ const CreateMCPServer: React.FC = ({ } try { const values = form.getFieldsValue(true); - // codeql[js/clear-text-storage-of-sensitive-data] - window.sessionStorage.setItem( + setSecureItem( CREATE_OAUTH_UI_STATE_KEY, JSON.stringify({ modalVisible: isModalVisible, @@ -178,7 +178,7 @@ const CreateMCPServer: React.FC = ({ if (typeof window === "undefined") { return; } - const storedState = window.sessionStorage.getItem(CREATE_OAUTH_UI_STATE_KEY); + const storedState = getSecureItem(CREATE_OAUTH_UI_STATE_KEY); if (!storedState) { return; } diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index e81d2f3960e..15fec946d3f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -12,6 +12,7 @@ import MCPLogoSelector from "./MCPLogoSelector"; import { validateMCPServerUrl, validateMCPServerName } from "./utils"; import NotificationsManager from "../molecules/notifications_manager"; import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; interface MCPServerEditProps { mcpServer: MCPServer; @@ -73,8 +74,7 @@ const MCPServerEdit: React.FC = ({ } try { const values = form.getFieldsValue(true); - // codeql[js/clear-text-storage-of-sensitive-data] - window.sessionStorage.setItem( + setSecureItem( EDIT_OAUTH_UI_STATE_KEY, JSON.stringify({ serverId: mcpServer.server_id, @@ -214,7 +214,7 @@ const MCPServerEdit: React.FC = ({ if (typeof window === "undefined") { return; } - const storedState = window.sessionStorage.getItem(EDIT_OAUTH_UI_STATE_KEY); + const storedState = getSecureItem(EDIT_OAUTH_UI_STATE_KEY); if (!storedState) { return; } diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx index f0fed7c7fec..72d5e4b5aa8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_servers.tsx @@ -20,6 +20,7 @@ import MCPSemanticFilterSettings from "../Settings/AdminSettings/MCPSemanticFilt import MCPNetworkSettings from "./MCPNetworkSettings"; import MCPDiscovery from "./mcp_discovery"; import { ByokCredentialModal } from "./ByokCredentialModal"; +import { getSecureItem } from "@/utils/secureStorage"; const { Text: AntdText, Title: AntdTitle } = Typography; const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; @@ -70,7 +71,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID }) return; } try { - const stored = window.sessionStorage.getItem(EDIT_OAUTH_UI_STATE_KEY); + const stored = getSecureItem(EDIT_OAUTH_UI_STATE_KEY); if (!stored) { return; } diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index 2bfe08efae8..b93921a29eb 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -75,6 +75,7 @@ import RealtimePlayground from "./RealtimePlayground"; import { A2ATaskMetadata, MessageType } from "./types"; import { useCodeInterpreter } from "./useCodeInterpreter"; import { useChatHistory } from "./useChatHistory"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; const { TextArea } = Input; const { Dragger } = Upload; @@ -167,7 +168,7 @@ const ChatUI: React.FC = ({ } = useChatHistory({ simplified }); // codeql[js/clear-text-storage-of-sensitive-data] const [apiKeySource, setApiKeySource] = useState<"session" | "custom">(() => { - const saved = sessionStorage.getItem("apiKeySource"); + const saved = getSecureItem("apiKeySource"); if (saved) { try { return JSON.parse(saved) as "session" | "custom"; @@ -177,8 +178,7 @@ const ChatUI: React.FC = ({ } return disabledPersonalKeyCreation ? "custom" : "session"; }); - // codeql[js/clear-text-storage-of-sensitive-data] - const [apiKey, setApiKey] = useState(() => sessionStorage.getItem("apiKey") || ""); + const [apiKey, setApiKey] = useState(() => getSecureItem("apiKey") || ""); const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState( () => sessionStorage.getItem("customProxyBaseUrl") || "", ); @@ -348,10 +348,8 @@ const ChatUI: React.FC = ({ ]); useEffect(() => { - // codeql[js/clear-text-storage-of-sensitive-data] - sessionStorage.setItem("apiKeySource", JSON.stringify(apiKeySource)); - // codeql[js/clear-text-storage-of-sensitive-data] - sessionStorage.setItem("apiKey", apiKey); + setSecureItem("apiKeySource", JSON.stringify(apiKeySource)); + setSecureItem("apiKey", apiKey); sessionStorage.setItem("endpointType", endpointType); sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags)); sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores)); @@ -502,7 +500,9 @@ const ChatUI: React.FC = ({ const handleImageUpload = (file: File) => { setUploadedImages((prev) => [...prev, file]); - const previewUrl = URL.createObjectURL(file); + const rawPreviewUrl = URL.createObjectURL(file); + // Sanitize: only allow blob: URLs to prevent XSS via img src injection. + const previewUrl = rawPreviewUrl.startsWith("blob:") ? rawPreviewUrl : ""; setImagePreviewUrls((prev) => [...prev, previewUrl]); return false; // Prevent default upload behavior }; @@ -1827,7 +1827,16 @@ const ChatUI: React.FC = ({ {uploadedImages.map((file, index) => (
{ + const url = imagePreviewUrls[index]; + if (!url) return ""; + try { + const parsed = new URL(url); + return parsed.protocol === "blob:" ? parsed.href : ""; + } catch { + return ""; + } + })()} alt={`Upload preview ${index + 1}`} className="max-w-32 max-h-32 rounded-md border border-gray-200 object-cover" /> diff --git a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx index e1afbf2e925..24881e669f9 100644 --- a/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useMcpOAuthFlow.tsx @@ -11,6 +11,7 @@ import { serverRootPath, } from "@/components/networking"; import { extractErrorMessage } from "@/utils/errorUtils"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; export type McpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; @@ -79,22 +80,13 @@ export const useMcpOAuthFlow = ({ const setStorageItem = (key: string, value: string) => { if (typeof window === "undefined") return; - try { - // Use sessionStorage only — the flow state may contain client credentials; - // writing them to localStorage would persist across browser sessions and - // make them readable by any injected script (XSS). - // codeql[js/clear-text-storage-of-sensitive-data] - window.sessionStorage.setItem(key, value); - } catch (err) { - console.warn(`Failed to set storage item ${key}`, err); - } + setSecureItem(key, value); }; const getStorageItem = (key: string): string | null => { if (typeof window === "undefined") return null; try { - // Try sessionStorage first, fall back to localStorage - return window.sessionStorage.getItem(key) || window.localStorage.getItem(key); + return getSecureItem(key); } catch (err) { console.warn(`Failed to get storage item ${key}`, err); return null; diff --git a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx index 3bb43d14ca4..e032c503dc7 100644 --- a/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx +++ b/ui/litellm-dashboard/src/hooks/useUserMcpOAuthFlow.tsx @@ -23,6 +23,7 @@ import { } from "@/components/networking"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { extractErrorMessage } from "@/utils/errorUtils"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; export type UserMcpOAuthStatus = "idle" | "authorizing" | "exchanging" | "success" | "error"; @@ -79,22 +80,11 @@ const genChallenge = async (verifier: string) => { }; const setStorage = (key: string, value: string) => { - try { - // Use sessionStorage only — do not write to localStorage. - // The flow state may contain the LiteLLM access token; writing it to - // localStorage would persist it across browser sessions and make it - // readable by any injected script (XSS). - // codeql[js/clear-text-storage-of-sensitive-data] - window.sessionStorage.setItem(key, value); - } catch (_) {} + setSecureItem(key, value); }; const getStorage = (key: string): string | null => { - try { - return window.sessionStorage.getItem(key); - } catch (_) { - return null; - } + return getSecureItem(key); }; const clearStorage = (...keys: string[]) => { diff --git a/ui/litellm-dashboard/src/utils/secureStorage.ts b/ui/litellm-dashboard/src/utils/secureStorage.ts new file mode 100644 index 00000000000..6b9a9bc1013 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/secureStorage.ts @@ -0,0 +1,32 @@ +function encode(value: string): string { + // btoa cannot handle characters outside Latin-1, so we percent-encode first. + return btoa(unescape(encodeURIComponent(value))); +} + +function decode(encoded: string): string { + return decodeURIComponent(escape(atob(encoded))); +} + +export function setSecureItem(key: string, value: string): void { + try { + window.sessionStorage.setItem(key, encode(value)); + } catch { + // Storage full or unavailable — silently ignore. + } +} + +export function getSecureItem(key: string): string | null { + try { + const raw = window.sessionStorage.getItem(key); + if (raw === null) return null; + return decode(raw); + } catch { + // Corrupted or non-encoded legacy value — clear it. + try { + window.sessionStorage.removeItem(key); + } catch { + // ignore + } + return null; + } +} From 36bf3373964c4a9bb7f46f340a3a6d4773ac6ea6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 2 Apr 2026 00:26:31 -0700 Subject: [PATCH 02/38] fix(docker): add non-root USER and HEALTHCHECK to Dockerfile.custom_ui Co-Authored-By: Claude Opus 4.6 (1M context) --- docker/Dockerfile.custom_ui | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docker/Dockerfile.custom_ui b/docker/Dockerfile.custom_ui index f836190a49a..11449af42ef 100644 --- a/docker/Dockerfile.custom_ui +++ b/docker/Dockerfile.custom_ui @@ -71,8 +71,15 @@ WORKDIR /app RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh +# Run as non-root user +RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser +USER appuser + # Expose the necessary port EXPOSE 4000/tcp +HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --retries=3 \ + CMD ["python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:4000/health')"] + # Override the CMD instruction with your desired command and arguments CMD ["--port", "4000", "--config", "config.yaml", "--detailed_debug"] \ No newline at end of file From 70a5c27cbd86b1f34bb2068819f4b3edde3aeb49 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 8 Apr 2026 17:51:34 -0700 Subject: [PATCH 03/38] [Fix] Address review feedback on storage utility and Dockerfiles - Dockerfile.health_check: HEALTHCHECK now verifies the script is intact instead of unconditionally exiting 0 - secureStorage.ts: replace deprecated escape/unescape with encodeURIComponent/decodeURIComponent; don't delete legacy values on decode failure so in-flight flows can time out naturally - OAuth callback: add same-origin check before redirecting to stored return URL --- docker/Dockerfile.health_check | 2 +- .../src/app/mcp/oauth/callback/page.tsx | 12 +++++++++- .../src/utils/secureStorage.ts | 22 ++++++++++++------- 3 files changed, 26 insertions(+), 10 deletions(-) diff --git a/docker/Dockerfile.health_check b/docker/Dockerfile.health_check index 6c18201e688..e968bec340e 100644 --- a/docker/Dockerfile.health_check +++ b/docker/Dockerfile.health_check @@ -18,7 +18,7 @@ USER appuser # Health check HEALTHCHECK --interval=30s --timeout=5s --start-period=5s --retries=3 \ - CMD ["python", "-c", "import sys; sys.exit(0)"] + CMD ["python", "/app/health_check_client.py", "--help"] # Set entrypoint ENTRYPOINT ["python", "/app/health_check_client.py"] diff --git a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx index 5c27a1d6150..0539d6d8f19 100644 --- a/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx +++ b/ui/litellm-dashboard/src/app/mcp/oauth/callback/page.tsx @@ -60,7 +60,17 @@ const McpOAuthCallbackContent = () => { } const returnUrl = getSecureItem(RETURN_URL_STORAGE_KEY); - const destination = returnUrl || resolveDefaultRedirect(); + let destination = resolveDefaultRedirect(); + if (returnUrl) { + try { + const parsed = new URL(returnUrl, window.location.origin); + if (parsed.origin === window.location.origin) { + destination = parsed.href; + } + } catch { + // invalid URL — fall through to default + } + } window.location.replace(destination); }, [payload]); diff --git a/ui/litellm-dashboard/src/utils/secureStorage.ts b/ui/litellm-dashboard/src/utils/secureStorage.ts index 6b9a9bc1013..6942368bcf2 100644 --- a/ui/litellm-dashboard/src/utils/secureStorage.ts +++ b/ui/litellm-dashboard/src/utils/secureStorage.ts @@ -1,10 +1,20 @@ function encode(value: string): string { // btoa cannot handle characters outside Latin-1, so we percent-encode first. - return btoa(unescape(encodeURIComponent(value))); + return btoa( + encodeURIComponent(value).replace( + /%([0-9A-F]{2})/g, + (_, p1) => String.fromCharCode(parseInt(p1, 16)) + ) + ); } function decode(encoded: string): string { - return decodeURIComponent(escape(atob(encoded))); + return decodeURIComponent( + atob(encoded) + .split("") + .map((c) => "%" + c.charCodeAt(0).toString(16).padStart(2, "0")) + .join("") + ); } export function setSecureItem(key: string, value: string): void { @@ -21,12 +31,8 @@ export function getSecureItem(key: string): string | null { if (raw === null) return null; return decode(raw); } catch { - // Corrupted or non-encoded legacy value — clear it. - try { - window.sessionStorage.removeItem(key); - } catch { - // ignore - } + // Corrupted or non-encoded legacy value — return null without deleting + // so that in-flight flows (e.g. OAuth) can time out naturally. return null; } } From ac29118942b4f6d57514cf3f2b962372ec0ba662 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 8 Apr 2026 17:59:09 -0700 Subject: [PATCH 04/38] Update docker/Dockerfile.custom_ui Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- docker/Dockerfile.custom_ui | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docker/Dockerfile.custom_ui b/docker/Dockerfile.custom_ui index 11449af42ef..cc44893bf92 100644 --- a/docker/Dockerfile.custom_ui +++ b/docker/Dockerfile.custom_ui @@ -72,7 +72,8 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh RUN sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh # Run as non-root user -RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser +RUN groupadd --gid 1000 appuser && useradd --uid 1000 --gid 1000 --no-create-home appuser \ + && chown -R appuser:appuser /app USER appuser # Expose the necessary port From 20ed120d1a845211b6d854a68f67e6249896d571 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 8 Apr 2026 22:13:48 -0700 Subject: [PATCH 05/38] [Fix] Let setSecureItem propagate storage errors to callers Remove the silent try/catch from setSecureItem so OAuth hooks can surface actionable "enable storage" guidance instead of a cryptic "state lost" error after the round-trip. Add a local try/catch in ChatUI where the storage write is non-critical. --- .../src/components/playground/chat_ui/ChatUI.tsx | 8 ++++++-- ui/litellm-dashboard/src/utils/secureStorage.ts | 6 +----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index b93921a29eb..06a09ca2c37 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -348,8 +348,12 @@ const ChatUI: React.FC = ({ ]); useEffect(() => { - setSecureItem("apiKeySource", JSON.stringify(apiKeySource)); - setSecureItem("apiKey", apiKey); + try { + setSecureItem("apiKeySource", JSON.stringify(apiKeySource)); + setSecureItem("apiKey", apiKey); + } catch { + // Storage full or unavailable — non-critical, skip persisting. + } sessionStorage.setItem("endpointType", endpointType); sessionStorage.setItem("selectedTags", JSON.stringify(selectedTags)); sessionStorage.setItem("selectedVectorStores", JSON.stringify(selectedVectorStores)); diff --git a/ui/litellm-dashboard/src/utils/secureStorage.ts b/ui/litellm-dashboard/src/utils/secureStorage.ts index 6942368bcf2..183b572e737 100644 --- a/ui/litellm-dashboard/src/utils/secureStorage.ts +++ b/ui/litellm-dashboard/src/utils/secureStorage.ts @@ -18,11 +18,7 @@ function decode(encoded: string): string { } export function setSecureItem(key: string, value: string): void { - try { - window.sessionStorage.setItem(key, encode(value)); - } catch { - // Storage full or unavailable — silently ignore. - } + window.sessionStorage.setItem(key, encode(value)); } export function getSecureItem(key: string): string | null { From f6dde296fa8a9ef4747ff5c6e4df0af808124134 Mon Sep 17 00:00:00 2001 From: joereyna Date: Thu, 9 Apr 2026 10:33:20 -0700 Subject: [PATCH 06/38] fix(responses-ws): append ?model= to backend WebSocket URL --- litellm/llms/custom_httpx/llm_http_handler.py | 4 +++ .../test_responses_websocket_all_providers.py | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c9abaad908..4e5b7a3410b 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5027,6 +5027,10 @@ class BaseLLMHTTPHandler: litellm_params={}, ) ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + # OpenAI's WebSocket responses endpoint requires ?model= in the URL, + # matching the Realtime API convention (wss://.../v1/realtime?model=...). + if "?" not in ws_url: + ws_url = f"{ws_url}?model={model}" try: ssl_context = get_shared_realtime_ssl_context() diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 0d83b9f88de..09bde86593e 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -971,3 +971,35 @@ class TestWebSocketChunkTypes: ) assert len(messages) == 1 assert messages[0]["content"][0]["text"] == "Part 1Part 2" + + +class TestNativeWebSocketUrlConstruction: + """Test that native WebSocket URLs include the model query parameter.""" + + def test_openai_ws_url_includes_model(self): + """ws_url for OpenAI native WebSocket must include ?model= so OpenAI + knows which model to use before the first response.create event.""" + config = OpenAIResponsesAPIConfig() + http_url = config.get_complete_url(api_base=None, litellm_params={}) + base_ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + + # get_complete_url should not include query params + assert "?" not in base_ws_url + + # The handler appends ?model= when none is present + model = "gpt-4o-mini" + ws_url = f"{base_ws_url}?model={model}" if "?" not in base_ws_url else base_ws_url + assert ws_url == "wss://api.openai.com/v1/responses?model=gpt-4o-mini" + + def test_ws_url_model_not_duplicated_if_query_already_present(self): + """If api_base already has query params, the ?model= should not be appended.""" + http_url = "https://custom.example.com/v1/responses?api-version=2024-05-01" + ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + + # Fix: only append when no query string present + if "?" not in ws_url: + ws_url = f"{ws_url}?model=gpt-4o" + + assert "api-version=2024-05-01" in ws_url + assert ws_url.count("?") == 1 + assert "model=gpt-4o" not in ws_url From 3ac4333be115a134f21cadba709e7d36289b0198 Mon Sep 17 00:00:00 2001 From: joereyna Date: Thu, 9 Apr 2026 11:14:34 -0700 Subject: [PATCH 07/38] fix(responses-ws): use urllib.parse to append model param, fix test mocking --- litellm/llms/custom_httpx/llm_http_handler.py | 11 +- .../test_responses_websocket_all_providers.py | 117 ++++++++++++++---- 2 files changed, 103 insertions(+), 25 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4e5b7a3410b..7a8820a8785 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1,5 +1,6 @@ import json import ssl +from urllib.parse import parse_qs, urlencode, urlparse, urlunparse from typing import ( TYPE_CHECKING, Any, @@ -5029,8 +5030,14 @@ class BaseLLMHTTPHandler: ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") # OpenAI's WebSocket responses endpoint requires ?model= in the URL, # matching the Realtime API convention (wss://.../v1/realtime?model=...). - if "?" not in ws_url: - ws_url = f"{ws_url}?model={model}" + # Use urllib.parse so existing query params (e.g. api-version) are preserved. + _parsed = urlparse(ws_url) + _qs = parse_qs(_parsed.query) + if "model" not in _qs: + _qs["model"] = [model] + ws_url = urlunparse( + _parsed._replace(query=urlencode({k: v[0] for k, v in _qs.items()})) + ) try: ssl_context = get_shared_realtime_ssl_context() diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 09bde86593e..efec841cc4d 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -974,32 +974,103 @@ class TestWebSocketChunkTypes: class TestNativeWebSocketUrlConstruction: - """Test that native WebSocket URLs include the model query parameter.""" + """Test that native WebSocket URLs include the model query parameter. - def test_openai_ws_url_includes_model(self): - """ws_url for OpenAI native WebSocket must include ?model= so OpenAI - knows which model to use before the first response.create event.""" - config = OpenAIResponsesAPIConfig() - http_url = config.get_complete_url(api_base=None, litellm_params={}) - base_ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + These tests mock websockets.connect so they exercise the actual URL-building + code inside BaseLLMHTTPHandler.async_responses_websocket rather than + reimplementing the logic themselves. + """ - # get_complete_url should not include query params - assert "?" not in base_ws_url + @pytest.mark.asyncio + async def test_openai_ws_url_includes_model(self): + """Handler must pass ?model= in the URL to the backend WebSocket.""" + from unittest.mock import AsyncMock, MagicMock, patch - # The handler appends ?model= when none is present - model = "gpt-4o-mini" - ws_url = f"{base_ws_url}?model={model}" if "?" not in base_ws_url else base_ws_url - assert ws_url == "wss://api.openai.com/v1/responses?model=gpt-4o-mini" + captured_urls = [] - def test_ws_url_model_not_duplicated_if_query_already_present(self): - """If api_base already has query params, the ?model= should not be appended.""" - http_url = "https://custom.example.com/v1/responses?api-version=2024-05-01" - ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) - # Fix: only append when no query string present - if "?" not in ws_url: - ws_url = f"{ws_url}?model=gpt-4o" + async def __aenter__(self): + raise Exception("stop") - assert "api-version=2024-05-01" in ws_url - assert ws_url.count("?") == 1 - assert "model=gpt-4o" not in ws_url + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = "https://api.openai.com/v1/responses" + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o-mini", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o-mini"], f"Expected model in URL, got: {captured_urls[0]}" + + @pytest.mark.asyncio + async def test_ws_url_preserves_existing_params_and_adds_model(self): + """When api_base already has query params, model is added alongside them.""" + from unittest.mock import AsyncMock, MagicMock, patch + + captured_urls = [] + + class FakeConnect: + def __init__(self, url, **kwargs): + captured_urls.append(url) + + async def __aenter__(self): + raise Exception("stop") + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.get_complete_url.return_value = ( + "https://custom.example.com/v1/responses?api-version=2024-05-01" + ) + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + handler = BaseLLMHTTPHandler() + mock_ws = MagicMock() + mock_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await handler.async_responses_websocket( + model="gpt-4o", + websocket=mock_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + ) + + assert len(captured_urls) == 1 + from urllib.parse import parse_qs, urlparse + qs = parse_qs(urlparse(captured_urls[0]).query) + assert qs.get("model") == ["gpt-4o"], f"model missing from URL: {captured_urls[0]}" + assert qs.get("api-version") == ["2024-05-01"], f"existing param lost: {captured_urls[0]}" From d910a95661fce3cdd36f3b06c03ecf9c46c6457c Mon Sep 17 00:00:00 2001 From: jayden Date: Thu, 9 Apr 2026 14:14:53 -0700 Subject: [PATCH 08/38] fix(proxy): improve input validation on management endpoints --- .../integrations/dotprompt/prompt_manager.py | 8 +- .../key_management_endpoints.py | 77 ++++++-- .../proxy_setting_endpoints.py | 27 +++ .../test_key_management_endpoints.py | 167 ++++++++++++++++++ 4 files changed, 258 insertions(+), 21 deletions(-) diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index 997a40d545e..6407a18d0b3 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -7,7 +7,8 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union import yaml -from jinja2 import DictLoader, Environment, select_autoescape +from jinja2 import DictLoader, select_autoescape +from jinja2.sandbox import ImmutableSandboxedEnvironment class PromptTemplate: @@ -59,7 +60,10 @@ class PromptManager: self.prompt_directory = Path(prompt_directory) if prompt_directory else None self.prompts: Dict[str, PromptTemplate] = {} self.prompt_file = prompt_file - self.jinja_env = Environment( + # Sandboxed env: templates can come from user input via /prompts/test, + # so we must block access to unsafe Python attributes and mutation of + # caller-supplied mutables. + self.jinja_env = ImmutableSandboxedEnvironment( loader=DictLoader({}), autoescape=select_autoescape(["html", "xml"]), # Use Handlebars-style delimiters to match Dotprompt spec diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 323ff7fd531..6e8a691ce93 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -456,6 +456,34 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: return data_json +def _check_allowed_routes_caller_permission( + allowed_routes: Optional[list], + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """ + Only proxy admins may set `allowed_routes` on a key. + + `allowed_routes` bypasses the standard role-based route gate in + RouteChecks.non_proxy_admin_allowed_routes_check, so if a non-admin is + allowed to set it they can grant themselves access to any endpoint. + Non-admins should use `key_type` to pick a preset route bucket instead. + """ + # Empty list is the default on GenerateKeyRequest — treat as "not set". + if not allowed_routes: + return + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Only proxy admins can set `allowed_routes` on a key. " + "Use `key_type` to pick a preset route bucket instead." + ) + }, + ) + + async def validate_team_id_used_in_service_account_request( team_id: Optional[str], prisma_client: Optional[PrismaClient], @@ -740,9 +768,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 request_type="key", **data_json, table_name="key" ) - response[ - "soft_budget" - ] = data.soft_budget # include the user-input soft budget in the response + response["soft_budget"] = ( + data.soft_budget + ) # include the user-input soft budget in the response response = GenerateKeyResponse(**response) @@ -1254,6 +1282,12 @@ async def generate_key_fn( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=message ) + + _check_allowed_routes_caller_permission( + allowed_routes=data.allowed_routes, + user_api_key_dict=user_api_key_dict, + ) + # For non-admin internal users: auto-assign caller's user_id if not provided # This prevents creating unbound keys with no user association (LIT-1884) _is_proxy_admin = ( @@ -1888,6 +1922,11 @@ async def _validate_update_key_data( """Validate permissions and constraints for key update.""" _is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + _check_allowed_routes_caller_permission( + allowed_routes=data.allowed_routes, + user_api_key_dict=user_api_key_dict, + ) + # Prevent non-admin from removing user_id (setting to empty string) (LIT-1884) if data.user_id is not None and data.user_id == "" and not _is_proxy_admin: raise HTTPException( @@ -3233,10 +3272,10 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": tokens}} + _keys_being_deleted: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} + ) ) if len(_keys_being_deleted) == 0: @@ -3436,9 +3475,9 @@ async def _rotate_master_key( # noqa: PLR0915 from litellm.proxy.proxy_server import proxy_config try: - models: Optional[ - List - ] = await prisma_client.db.litellm_proxymodeltable.find_many() + models: Optional[List] = ( + await prisma_client.db.litellm_proxymodeltable.find_many() + ) except Exception: models = None # 2. process model table @@ -4078,11 +4117,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[ - BaseModel - ] = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, + complete_user_info_db_obj: Optional[BaseModel] = ( + await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, + ) ) if complete_user_info_db_obj is None: @@ -4165,10 +4204,10 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[ - List[BaseModel] - ] = await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} + teams: Optional[List[BaseModel]] = ( + await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": complete_user_info.teams}} + ) ) if teams is None: return [] diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 60bf41709ef..0349f289b4e 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -1,6 +1,7 @@ #### CRUD ENDPOINTS for UI Settings ##### import json from typing import Any, Dict, List, Optional, Union +from urllib.parse import urlparse from fastapi import APIRouter, Depends, File, HTTPException, UploadFile @@ -817,6 +818,29 @@ async def get_ui_theme_settings(): ) +def _validate_public_image_url(value: Optional[str], field_name: str) -> None: + """ + Reject anything that isn't a plain http(s) URL with a host. This value is + later served via the unauthenticated /get_image endpoint, so local paths + like "/etc/passwd" or "file://..." must not be accepted. + """ + if value is None: + return + if not isinstance(value, str) or not value.strip(): + return + parsed = urlparse(value.strip()) + if parsed.scheme not in ("http", "https") or not parsed.netloc: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"Invalid {field_name}: must be an http(s) URL with a host. " + "Local filesystem paths and non-http schemes are not allowed." + ) + }, + ) + + @router.patch( "/update/ui_theme_settings", tags=["UI Theme Settings"], @@ -831,6 +855,9 @@ async def update_ui_theme_settings(theme_config: UIThemeConfig): from litellm.proxy.proxy_server import proxy_config, store_model_in_db + _validate_public_image_url(theme_config.logo_url, "logo_url") + _validate_public_image_url(theme_config.favicon_url, "favicon_url") + if store_model_in_db is not True: raise HTTPException( status_code=500, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 72092a97f5b..096e0b2bc41 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -8579,3 +8579,170 @@ def test_enforce_upperbound_no_config_is_noop(): assert data.tpm_limit == 999999 finally: litellm.upperbound_key_generate_params = original + + +class TestAllowedRoutesCallerPermission: + """ + Non-admins must not be able to set `allowed_routes` on a key. The field + bypasses the role-based route gate in + RouteChecks.non_proxy_admin_allowed_routes_check, so allowing a non-admin + to populate it grants them arbitrary endpoint access. + """ + + @pytest.mark.asyncio + async def test_non_admin_generate_key_with_allowed_routes_rejected(self): + data = GenerateKeyRequest( + key_alias="escalate", + allowed_routes=["/*"], + ) + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + mock_prisma_client = AsyncMock() + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() + ), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert str(exc_info.value.code) == "403" + assert "allowed_routes" in str(exc_info.value.message) + + @pytest.mark.asyncio + async def test_admin_generate_key_with_allowed_routes_allowed(self): + data = GenerateKeyRequest( + key_alias="admin-key", + allowed_routes=["/chat/completions"], + user_id="admin-user", + ) + user_api_key_dict = UserAPIKeyAuth( + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + mock_prisma_client = AsyncMock() + stub_response = MagicMock() + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() + ), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + new_callable=AsyncMock, + return_value=stub_response, + ): + result = await generate_key_fn( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert result is stub_response + + @pytest.mark.asyncio + async def test_non_admin_generate_key_default_empty_allowed_routes_ok(self): + """ + Regression guard: GenerateKeyRequest.allowed_routes defaults to [], so + the helper must treat empty-list as "not set" or every non-admin key + creation breaks. + """ + data = GenerateKeyRequest(key_alias="plain-key") + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + mock_prisma_client = AsyncMock() + stub_response = MagicMock() + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() + ), patch("litellm.proxy.proxy_server.user_custom_key_generate", None), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._common_key_generation_helper", + new_callable=AsyncMock, + return_value=stub_response, + ): + result = await generate_key_fn( + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert result is stub_response + + @pytest.mark.asyncio + async def test_non_admin_update_key_with_allowed_routes_rejected(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + data = UpdateKeyRequest(key="sk-test", allowed_routes=["/*"]) + user_api_key_dict = UserAPIKeyAuth( + user_id="internal-user-123", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + mock_prisma_client = AsyncMock() + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch( + "litellm.proxy.proxy_server.user_api_key_cache", MagicMock() + ), patch("litellm.proxy.proxy_server.user_custom_key_update", None), patch( + "litellm.proxy.proxy_server.llm_router", None + ), patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.proxy.proxy_server.proxy_logging_obj", MagicMock() + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._get_and_validate_existing_key", + new_callable=AsyncMock, + return_value=MagicMock(), + ): + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + assert str(exc_info.value.code) == "403" + assert "allowed_routes" in str(exc_info.value.message) + + +def test_jinja_prompt_manager_is_sandboxed(): + """ + PromptManager renders user-supplied templates via /prompts/test, so its + jinja env must reject access to unsafe Python attributes like + ``__class__`` and ``__mro__``. + """ + from jinja2.exceptions import SecurityError + + from litellm.integrations.dotprompt.prompt_manager import PromptManager + + pm = PromptManager() + template = pm.jinja_env.from_string("{{ ''.__class__.__mro__ }}") + with pytest.raises(SecurityError): + template.render() + + +def test_validate_public_image_url_rejects_local_paths(): + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _validate_public_image_url, + ) + + for bad in ("/etc/passwd", "file:///etc/passwd", "../../etc/passwd"): + with pytest.raises(HTTPException) as exc_info: + _validate_public_image_url(bad, "logo_url") + assert exc_info.value.status_code == 400 + + +def test_validate_public_image_url_accepts_http_and_noop_empty(): + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _validate_public_image_url, + ) + + _validate_public_image_url("https://example.com/logo.png", "logo_url") + _validate_public_image_url("http://cdn.internal/logo.svg", "logo_url") + _validate_public_image_url(None, "logo_url") + _validate_public_image_url("", "logo_url") + _validate_public_image_url(" ", "logo_url") From 4dc416ee749122ca91e3bca095217478663419e7 Mon Sep 17 00:00:00 2001 From: jayden Date: Thu, 9 Apr 2026 20:10:40 -0700 Subject: [PATCH 09/38] fix(proxy): use parameterized query for combined_view token lookup --- litellm/proxy/utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 635204f3362..a62f34764d3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2645,7 +2645,7 @@ class PrismaClient: raise e async def _query_first_with_cached_plan_fallback( - self, sql_query: str + self, sql_query: str, *args ) -> Optional[dict]: """ Execute a query with automatic fallback for PostgreSQL cached plan errors. @@ -2664,7 +2664,7 @@ class PrismaClient: Original exception if not a cached plan error """ try: - return await self.db.query_first(query=sql_query) + return await self.db.query_first(sql_query, *args) except Exception as e: error_str = str(e) if "cached plan must not change result type" in error_str: @@ -2679,7 +2679,7 @@ class PrismaClient: "retrying with fresh plan. This may occur during rolling deployments " "when schema changes are applied." ) - return await self.db.query_first(query=sql_query_retry) + return await self.db.query_first(sql_query_retry, *args) else: raise @@ -3016,11 +3016,11 @@ class PrismaClient: LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id LEFT JOIN "LiteLLM_OrganizationTable" AS o ON v.organization_id = o.organization_id LEFT JOIN "LiteLLM_BudgetTable" AS b2 ON o.budget_id = b2.budget_id - WHERE v.token = '{token}' + WHERE v.token = $1 """ response = await self._query_first_with_cached_plan_fallback( - sql_query + sql_query, hashed_token ) # If not found in main table, check deprecated keys (grace period) From 3a316b913179ade01fe12b42162089c6a80271de Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 9 Apr 2026 20:42:30 -0700 Subject: [PATCH 10/38] [Test] UI - Unit tests: raise global vitest timeout and remove per-test overrides Raise vitest testTimeout from 10s to 30s and drop per-test timeout overrides across UI unit tests. Group CreateUserButton and TeamInfo tests under nested describe blocks to make the most flaky suites easier to scan. --- .../ModelsAndEndpointsView.test.tsx | 8 +- .../src/components/CreateUserButton.test.tsx | 524 +++---- .../src/components/OldTeams.test.tsx | 2 +- .../add_model/add_model_tab.test.tsx | 2 +- .../mcp_tools/create_mcp_server.test.tsx | 290 ++-- .../src/components/team/TeamInfo.test.tsx | 1232 +++++++++-------- ui/litellm-dashboard/vitest.config.ts | 2 +- 7 files changed, 1029 insertions(+), 1031 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx index e1b3b358300..3c5101fc2dc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx @@ -131,7 +131,7 @@ describe("ModelsAndEndpointsView", () => { , ); expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument(); - }, 15000); + }); it("should show Missing provider banner by default", async () => { localStorageMock.clear(); @@ -149,7 +149,7 @@ describe("ModelsAndEndpointsView", () => { , ); expect(await findByText("Missing a provider?", {}, { timeout: 10000 })).toBeInTheDocument(); - }, 15000); + }); it("should hide Missing provider banner when dismiss button is clicked and persist to localStorage", async () => { localStorageMock.clear(); @@ -180,7 +180,7 @@ describe("ModelsAndEndpointsView", () => { // LocalStorage should be updated expect(localStorageMock.getItem("hideMissingProviderBanner")).toBe("true"); - }, 15000); + }); it("should show compact Request Provider button when banner is dismissed", async () => { // Set localStorage to hide banner @@ -209,7 +209,7 @@ describe("ModelsAndEndpointsView", () => { const requestProviderLinks = document.querySelectorAll('a[href="https://models.litellm.ai/?request=true"]'); // There should be a compact button when banner is hidden expect(requestProviderLinks.length).toBeGreaterThan(0); - }, 15000); + }); it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => { mockHealthCheckComponent.mockClear(); diff --git a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx index 9a4659da9d3..03d982ca7c2 100644 --- a/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx +++ b/ui/litellm-dashboard/src/components/CreateUserButton.test.tsx @@ -51,7 +51,7 @@ function renderWithProviders(ui: React.ReactElement) { return render({ui}); } -describe("CreateUserButton", { timeout: 20000 }, () => { +describe("CreateUserButton", () => { beforeEach(() => { vi.clearAllMocks(); mockGetProxyUISettings.mockResolvedValue({ @@ -62,288 +62,296 @@ describe("CreateUserButton", { timeout: 20000 }, () => { }); }); - it("should render the create user form when embedded", () => { - renderWithProviders( - , - ); - expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument(); - }); + describe("rendering and visibility", () => { + it("should render the create user form when embedded", () => { + renderWithProviders( + , + ); + expect(screen.getByRole("button", { name: /create user/i })).toBeInTheDocument(); + }); - it("should render the invite user button when not embedded", async () => { - renderWithProviders(); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + it("should render the invite user button when not embedded", async () => { + renderWithProviders(); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + }); + + it("should open the invite modal when invite user button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + expect(dialog).toBeInTheDocument(); + expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument(); + }); + + it("should display email invitations info message in embedded mode", () => { + renderWithProviders(); + expect(screen.getByText("Email invitations")).toBeInTheDocument(); + }); + + it("should display user role options when possibleUIRoles is provided", async () => { + const possibleUIRoles = { + proxy_admin: { ui_label: "Admin", description: "Full access" }, + proxy_user: { ui_label: "User", description: "Limited access" }, + }; + renderWithProviders( + , + ); + await userEvent.click(screen.getByRole("combobox", { name: /user role/i })); + expect(screen.getByText("Admin")).toBeInTheDocument(); + expect(screen.getByText("User")).toBeInTheDocument(); + }); + + it("should close modal when cancel is clicked in standalone mode", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument(); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.click(within(dialog).getByRole("button", { name: /close/i })); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); }); }); - it("should open the invite modal when invite user button is clicked", async () => { - const user = userEvent.setup(); - renderWithProviders(); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + describe("embedded mode submission", () => { + it("should call userCreateCall when form is submitted in embedded mode", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-1", + user_id: "new-user-123", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "test@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ + user_email: "test@example.com", + user_role: "proxy_user", + })); + }); }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - expect(dialog).toBeInTheDocument(); - expect(within(dialog).getByRole("button", { name: /invite user/i })).toBeInTheDocument(); - }); - it("should display email invitations info message in embedded mode", () => { - renderWithProviders(); - expect(screen.getByText("Email invitations")).toBeInTheDocument(); - }); + it("should call onUserCreated callback when user is created in embedded mode", async () => { + const user = userEvent.setup(); + const onUserCreated = vi.fn(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } }); - it("should display user role options when possibleUIRoles is provided", async () => { - const possibleUIRoles = { - proxy_admin: { ui_label: "Admin", description: "Full access" }, - proxy_user: { ui_label: "User", description: "Limited access" }, - }; - renderWithProviders( - , - ); - await userEvent.click(screen.getByRole("combobox", { name: /user role/i })); - expect(screen.getByText("Admin")).toBeInTheDocument(); - expect(screen.getByText("User")).toBeInTheDocument(); - }); + renderWithProviders( + , + ); - it("should call userCreateCall when form is submitted in embedded mode", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-123" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-1", - user_id: "new-user-123", - has_user_setup_sso: false, - } as any); + await user.type(screen.getByLabelText(/user email/i), "embedded@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); - renderWithProviders( - , - ); + await waitFor(() => { + expect(onUserCreated).toHaveBeenCalledWith("new-user-456"); + }); + }); - await user.type(screen.getByLabelText(/user email/i), "test@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); + it("should show error notification when user creation fails", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } }); - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ - user_email: "test@example.com", - user_role: "proxy_user", - })); + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists"); + }); + }); + + it("should show info notification when making API call", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-3", + user_id: "new-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await user.type(screen.getByLabelText(/user email/i), "info@example.com"); + await user.click(screen.getByRole("combobox", { name: /user role/i })); + await user.click(screen.getByText("User")); + await user.click(screen.getByRole("button", { name: /create user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call"); + }); }); }); - it("should call onUserCreated callback when user is created in embedded mode", async () => { - const user = userEvent.setup(); - const onUserCreated = vi.fn(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-456" } }); + describe("standalone mode submission", () => { + it("should show success notification when user is created successfully in standalone mode", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-2", + user_id: "new-user-789", + has_user_setup_sso: false, + } as any); - renderWithProviders( - , - ); + renderWithProviders( + , + ); - await user.type(screen.getByLabelText(/user email/i), "embedded@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - await waitFor(() => { - expect(onUserCreated).toHaveBeenCalledWith("new-user-456"); + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + }); + }); + + it("should show onboarding modal when user is created and SSO is disabled", async () => { + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-sso", + user_id: "sso-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user"); + }); + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + }); }); }); - it("should show success notification when user is created successfully in standalone mode", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user-789" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-2", - user_id: "new-user-789", - has_user_setup_sso: false, - } as any); + describe("organizations", () => { + it("should send organizations list in POST body when organizations are selected", async () => { + const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); + vi.mocked(useOrganizations).mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "My Org" }], + isLoading: false, + } as any); - renderWithProviders( - , - ); + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-org", + user_id: "org-user", + has_user_setup_sso: false, + } as any); - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + + // Select org from the dropdown + const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i }); + await user.click(orgSelect); + await user.click(screen.getByText("My Org (org-1)")); + + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ + organizations: ["org-1"], + })); + }); }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "standalone@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + it("should not call organizationMemberAddCall after user creation", async () => { + const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); + vi.mocked(useOrganizations).mockReturnValue({ + data: [{ organization_id: "org-1", organization_alias: "My Org" }], + isLoading: false, + } as any); - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); + const user = userEvent.setup(); + mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } }); + mockInvitationCreateCall.mockResolvedValue({ + id: "inv-nma", + user_id: "no-member-add-user", + has_user_setup_sso: false, + } as any); + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /\+ invite user/i })); + + const dialog = screen.getByRole("dialog", { name: /invite user/i }); + await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com"); + await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); + await user.click(screen.getByText("User")); + await user.click(within(dialog).getByRole("button", { name: /invite user/i })); + + await waitFor(() => { + expect(mockUserCreateCall).toHaveBeenCalled(); + }); + expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled(); }); }); - - it("should show error notification when user creation fails", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockRejectedValue({ response: { data: { detail: "Email already exists" } } }); - - renderWithProviders( - , - ); - - await user.type(screen.getByLabelText(/user email/i), "duplicate@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); - - await waitFor(() => { - expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Email already exists"); - }); - }); - - it("should show info notification when making API call", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "new-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-3", - user_id: "new-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await user.type(screen.getByLabelText(/user email/i), "info@example.com"); - await user.click(screen.getByRole("combobox", { name: /user role/i })); - await user.click(screen.getByText("User")); - await user.click(screen.getByRole("button", { name: /create user/i })); - - await waitFor(() => { - expect(mockNotificationsManager.info).toHaveBeenCalledWith("Making API Call"); - }); - }); - - it("should close modal when cancel is clicked in standalone mode", async () => { - const user = userEvent.setup(); - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - expect(screen.getByRole("dialog", { name: /invite user/i })).toBeInTheDocument(); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.click(within(dialog).getByRole("button", { name: /close/i })); - expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); - }); - - it("should show onboarding modal when user is created and SSO is disabled", async () => { - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "sso-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-sso", - user_id: "sso-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "sso@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockInvitationCreateCall).toHaveBeenCalledWith("token", "sso-user"); - }); - await waitFor(() => { - expect(mockNotificationsManager.success).toHaveBeenCalledWith("API user Created"); - }); - }); - - it("should send organizations list in POST body when organizations are selected", async () => { - const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); - vi.mocked(useOrganizations).mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "My Org" }], - isLoading: false, - } as any); - - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "org-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-org", - user_id: "org-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "org@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - - // Select org from the dropdown - const orgSelect = within(dialog).getByRole("combobox", { name: /organization/i }); - await user.click(orgSelect); - await user.click(screen.getByText("My Org (org-1)")); - - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalledWith("token", null, expect.objectContaining({ - organizations: ["org-1"], - })); - }); - }); - - it("should not call organizationMemberAddCall after user creation", async () => { - const { useOrganizations } = await import("@/app/(dashboard)/hooks/organizations/useOrganizations"); - vi.mocked(useOrganizations).mockReturnValue({ - data: [{ organization_id: "org-1", organization_alias: "My Org" }], - isLoading: false, - } as any); - - const user = userEvent.setup(); - mockUserCreateCall.mockResolvedValue({ data: { user_id: "no-member-add-user" } }); - mockInvitationCreateCall.mockResolvedValue({ - id: "inv-nma", - user_id: "no-member-add-user", - has_user_setup_sso: false, - } as any); - - renderWithProviders( - , - ); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /\+ invite user/i })).toBeInTheDocument(); - }); - await user.click(screen.getByRole("button", { name: /\+ invite user/i })); - - const dialog = screen.getByRole("dialog", { name: /invite user/i }); - await user.type(within(dialog).getByLabelText(/user email/i), "nomemberadd@example.com"); - await user.click(within(dialog).getByRole("combobox", { name: /global proxy role/i })); - await user.click(screen.getByText("User")); - await user.click(within(dialog).getByRole("button", { name: /invite user/i })); - - await waitFor(() => { - expect(mockUserCreateCall).toHaveBeenCalled(); - }); - expect(mockOrganizationMemberAddCall).not.toHaveBeenCalled(); - }); }); diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index 651d1495c61..4b89820bad6 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -843,7 +843,7 @@ describe("OldTeams - access_group_ids in team create", () => { }), ); }); - }, { timeout: 30000 }); + }); }); describe("OldTeams - models dropdown options", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 197bcd6569f..59970547c1a 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -269,7 +269,7 @@ describe("Add Model Tab", () => { }, { timeout: 10000 }, ); - }, 15000); // 15 second timeout to allow waitFor to complete + }); it("should show team selection when team-only switch is enabled", async () => { const props = createTestProps(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index c92956b430f..b4251267137 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -150,151 +150,139 @@ describe("CreateMCPServer", () => { }); }); - it( - "should not require auth value when creating a server with API Key auth type", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should not require auth value when creating a server with API Key auth type", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - // Fill in server name (use id to avoid duplicate placeholder) - const nameInput = getServerNameInput(); - await user.type(nameInput, "Test_Server"); + // Fill in server name (use id to avoid duplicate placeholder) + const nameInput = getServerNameInput(); + await user.type(nameInput, "Test_Server"); - // Fill in URL - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + // Fill in URL + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - // Select API Key auth type - await selectAntOption("Authentication", "API Key"); + // Select API Key auth type + await selectAntOption("Authentication", "API Key"); - await waitFor(() => { - expect(screen.getByText("Authentication Value")).toBeInTheDocument(); - }); + await waitFor(() => { + expect(screen.getByText("Authentication Value")).toBeInTheDocument(); + }); - // Leave auth value empty and submit - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "Test_Server", - alias: "Test_Server", - url: "https://example.com/mcp", - transport: "http", - auth_type: "api_key", - created_at: "2024-01-01T00:00:00Z", - created_by: "user-1", - updated_at: "2024-01-01T00:00:00Z", - updated_by: "user-1", - }); + // Leave auth value empty and submit + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "Test_Server", + alias: "Test_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "api_key", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - // The form should submit without validation error on auth_value - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); - }, - ); + // The form should submit without validation error on auth_value + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + }); - it( - "should not require auth value when creating a server with Bearer Token auth type", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should not require auth value when creating a server with Bearer Token auth type", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - const nameInput = getServerNameInput(); - await user.type(nameInput, "Test_Server"); + const nameInput = getServerNameInput(); + await user.type(nameInput, "Test_Server"); - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - await selectAntOption("Authentication", "Bearer Token"); + await selectAntOption("Authentication", "Bearer Token"); - await waitFor(() => { - expect(screen.getByText("Authentication Value")).toBeInTheDocument(); - }); + await waitFor(() => { + expect(screen.getByText("Authentication Value")).toBeInTheDocument(); + }); - // Leave auth value empty and submit - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "Test_Server", - alias: "Test_Server", - url: "https://example.com/mcp", - transport: "http", - auth_type: "bearer_token", - created_at: "2024-01-01T00:00:00Z", - created_by: "user-1", - updated_at: "2024-01-01T00:00:00Z", - updated_by: "user-1", - }); + // Leave auth value empty and submit + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "Test_Server", + alias: "Test_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "bearer_token", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); - }, - ); + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + }); - it( - "should successfully create a server when auth value is provided", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should successfully create a server when auth value is provided", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - const nameInput = getServerNameInput(); - await user.type(nameInput, "My_Server"); + const nameInput = getServerNameInput(); + await user.type(nameInput, "My_Server"); - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - await selectAntOption("Authentication", "API Key"); + await selectAntOption("Authentication", "API Key"); - await waitFor(() => { - expect(screen.getByText("Authentication Value")).toBeInTheDocument(); - }); + await waitFor(() => { + expect(screen.getByText("Authentication Value")).toBeInTheDocument(); + }); - // Fill in auth value - const authInput = screen.getByPlaceholderText("Enter token or secret"); - await user.type(authInput, "my-secret-key"); + // Fill in auth value + const authInput = screen.getByPlaceholderText("Enter token or secret"); + await user.type(authInput, "my-secret-key"); - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "My_Server", - alias: "My_Server", - url: "https://example.com/mcp", - transport: "http", - auth_type: "api_key", - created_at: "2024-01-01T00:00:00Z", - created_by: "user-1", - updated_at: "2024-01-01T00:00:00Z", - updated_by: "user-1", - }); + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "My_Server", + alias: "My_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "api_key", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); - const [token, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; - expect(token).toBe("test-token"); - expect(payload.credentials).toEqual({ auth_value: "my-secret-key" }); - }, - ); + const [token, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(token).toBe("test-token"); + expect(payload.credentials).toEqual({ auth_value: "my-secret-key" }); + }); it("should not show auth value field when None auth type is selected", async () => { await selectHttpTransport(); @@ -307,50 +295,46 @@ describe("CreateMCPServer", () => { }); }); - it( - "should successfully create a server with no auth", - { timeout: 15000 }, - async () => { - await selectHttpTransport(); + it("should successfully create a server with no auth", async () => { + await selectHttpTransport(); - const user = userEvent.setup({ delay: null }); + const user = userEvent.setup({ delay: null }); - const nameInput = getServerNameInput(); - await user.type(nameInput, "No_Auth_Server"); + const nameInput = getServerNameInput(); + await user.type(nameInput, "No_Auth_Server"); - const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); - await user.type(urlInput, "https://example.com/mcp"); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await user.type(urlInput, "https://example.com/mcp"); - await selectAntOption("Authentication", "None"); + await selectAntOption("Authentication", "None"); - vi.mocked(networking.createMCPServer).mockResolvedValue({ - server_id: "new-server-1", - server_name: "No_Auth_Server", - alias: "No_Auth_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", - }); + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-1", + server_name: "No_Auth_Server", + alias: "No_Auth_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", + }); - const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); - await act(async () => { - fireEvent.click(submitButton); - }); + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); - await waitFor(() => { - expect(networking.createMCPServer).toHaveBeenCalledTimes(1); - }); + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); - const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; - expect(payload.auth_type).toBe("none"); - // No credentials should be sent for "none" auth - expect(payload.credentials).toBeUndefined(); - }, - ); + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.auth_type).toBe("none"); + // No credentials should be sent for "none" auth + expect(payload.credentials).toBeUndefined(); + }); }); describe("when OAuth interactive auth is selected", () => { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index fb149458f61..f24b4e77e16 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -222,634 +222,640 @@ describe("TeamInfoView", () => { vi.clearAllMocks(); }); - it("should render", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + describe("display and rendering", () => { + it("should render", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - renderWithProviders(); + renderWithProviders(); - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - }); - - it("should display loading state while fetching team data", () => { - vi.mocked(networking.teamInfoCall).mockImplementation(() => new Promise(() => { })); - - renderWithProviders(); - - expect(screen.getByText("Loading...")).toBeInTheDocument(); - }); - - it("should display error message when team is not found", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue({ - team_id: "123", - team_info: null as any, - keys: [], - team_memberships: [], + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); }); - renderWithProviders(); + it("should display loading state while fetching team data", () => { + vi.mocked(networking.teamInfoCall).mockImplementation(() => new Promise(() => {})); - await waitFor(() => { - expect(screen.getByText("Team not found")).toBeInTheDocument(); - }); - }); + renderWithProviders(); - it("should display budget information in overview", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - max_budget: 1000, - spend: 250.5, - budget_duration: "30d", - }) - ); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Budget Status")).toBeInTheDocument(); - }); - }); - - it("should display guardrails in overview when present", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - guardrails: ["guardrail1", "guardrail2"], - }) - ); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Guardrails")).toBeInTheDocument(); - }); - }); - - it("should display policies in overview when present", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - policies: ["policy1"], - }) - ); - vi.mocked(networking.getPolicyInfoWithGuardrails).mockResolvedValue({ - resolved_guardrails: ["guardrail1"], + expect(screen.getByText("Loading...")).toBeInTheDocument(); }); - renderWithProviders(); + it("should display error message when team is not found", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue({ + team_id: "123", + team_info: null as any, + keys: [], + team_memberships: [], + }); - await waitFor(() => { - expect(screen.getByText("Policies")).toBeInTheDocument(); - }); - }); + renderWithProviders(); - it("should show members tab when user can edit team", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("tab", { name: "Members" })).toBeInTheDocument(); - }); - }); - - it("should not show members tab when user cannot edit team", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); + await waitFor(() => { + expect(screen.getByText("Team not found")).toBeInTheDocument(); + }); }); - expect(screen.queryByRole("tab", { name: "Members" })).not.toBeInTheDocument(); - }); - - it("should show settings tab when user can edit team", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("tab", { name: "Settings" })).toBeInTheDocument(); - }); - }); - - it("should navigate to settings tab when clicked", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByText("Team Settings")).toBeInTheDocument(); - }); - }); - - it("should open edit mode when edit button is clicked", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: /edit settings/i }); - await user.click(editButton); - - await waitFor(() => { - expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); - }); - }); - - it("should close edit mode when cancel button is clicked", { timeout: 15000 }, async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: /edit settings/i }); - await user.click(editButton); - - await waitFor(() => { - expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); - }); - - const cancelButton = screen.getByRole("button", { name: /cancel/i }); - await user.click(cancelButton); - - await waitFor(() => { - expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument(); - }); - }); - - it("should call onClose when back button is clicked", async () => { - const user = userEvent.setup({ delay: null }); - const onClose = vi.fn(); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const backButton = screen.getByRole("button", { name: /back to teams/i }); - await user.click(backButton); - - expect(onClose).toHaveBeenCalled(); - }); - - it("should copy team ID to clipboard when copy button is clicked", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const copyButtons = screen.getAllByRole("button"); - const copyButton = copyButtons.find((btn) => btn.querySelector("svg")); - expect(copyButton).toBeTruthy(); - - if (copyButton) { - await user.click(copyButton); - } - }); - - it("should disable secret manager settings for non-premium users", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - metadata: { - secret_manager_settings: { provider: "aws", secret_id: "abc" }, - }, - }) - ); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: /edit settings/i }); - await user.click(editButton); - - const secretField = await screen.findByPlaceholderText( - '{"namespace": "admin", "mount": "secret", "path_prefix": "litellm"}' - ); - expect(secretField).toBeDisabled(); - }); - - it("should allow premium users to edit secret manager settings", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - metadata: { - secret_manager_settings: { provider: "aws", secret_id: "abc" }, - }, - }) - ); - vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: /edit settings/i }); - await user.click(editButton); - - const secretField = await screen.findByPlaceholderText( - '{"namespace": "admin", "mount": "secret", "path_prefix": "litellm"}' - ); - expect(secretField).not.toBeDisabled(); - }); - - it("should add team member when form is submitted", async () => { - const user = userEvent.setup({ delay: null }); - const onUpdate = vi.fn(); - const teamData = createMockTeamData(); - vi.mocked(networking.teamInfoCall).mockResolvedValue(teamData); - vi.mocked(networking.teamMemberAddCall).mockResolvedValue({} as any); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const membersTab = screen.getByRole("tab", { name: "Members" }); - await user.click(membersTab); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /add member/i })).toBeInTheDocument(); - }); - - const addButton = screen.getByRole("button", { name: /add member/i }); - await user.click(addButton); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Submit" })).toBeInTheDocument(); - }); - - const submitButton = screen.getByRole("button", { name: "Submit" }); - await user.click(submitButton); - - await waitFor(() => { - expect(networking.teamMemberAddCall).toHaveBeenCalled(); - }); - }); - - it("should display team member budget information when present", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - team_member_budget_table: { - max_budget: 500, + it("should display budget information in overview", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + max_budget: 1000, + spend: 250.5, budget_duration: "30d", - tpm_limit: 5000, - rpm_limit: 50, - }, - }) - ); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Budget Status")).toBeInTheDocument(); - }); - }); - - it("should display virtual keys information", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue({ - ...createMockTeamData(), - keys: [ - { user_id: "user1", token: "key1" }, - { token: "key2" }, - ], - }); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("tab", { name: "Virtual Keys" })).toBeInTheDocument(); - }); - }); - - it("should show Virtual Keys tab when user cannot edit team", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByRole("tab", { name: "Virtual Keys" })).toBeInTheDocument(); - }); - }); - - it("should display X Members in Virtual Keys tab when navigated to", async () => { - const user = userEvent.setup(); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - const fiveKeys = Array.from({ length: 5 }, (_, i) => ({ - token: `sk-${i}`, - token_id: `key-${i}`, - key_alias: `key_${i}`, - key_name: `sk-...${i}`, - user_id: `user-${i}`, - organization_id: null, - user: { user_id: `user-${i}`, user_email: `user${i}@test.com` }, - created_at: "2024-01-01T00:00:00Z", - team_id: "123", - spend: 0, - max_budget: 100, - models: ["gpt-4"], - })); - mockUseKeys.mockReturnValue({ - data: { keys: fiveKeys, total_count: 5, current_page: 1, total_pages: 1 }, - isPending: false, - isFetching: false, - refetch: vi.fn(), - } as any); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const virtualKeysTab = screen.getByRole("tab", { name: "Virtual Keys" }); - await user.click(virtualKeysTab); - - await waitFor(() => { - expect(screen.getByText("Page 1 of 1")).toBeInTheDocument(); - }); - }); - - it("should show Filters and pagination controls in Virtual Keys tab", async () => { - const user = userEvent.setup(); - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - mockUseKeys.mockReturnValue({ - data: { - keys: [ - { - token: "sk-1", - token_id: "key-1", - key_alias: "key1", - key_name: "sk-...1", - user_id: "user-1", - organization_id: null, - user: { user_id: "user-1", user_email: "user1@test.com" }, - created_at: "2024-01-01T00:00:00Z", - team_id: "123", - spend: 0, - max_budget: 100, - models: ["gpt-4"], - }, - ], - total_count: 1, - current_page: 1, - total_pages: 1, - }, - isPending: false, - isFetching: false, - refetch: vi.fn(), - } as any); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const virtualKeysTab = screen.getByRole("tab", { name: "Virtual Keys" }); - await user.click(virtualKeysTab); - - await waitFor(() => { - expect(screen.getByRole("button", { name: "Filters" })).toBeInTheDocument(); - }); - expect(screen.getByRole("button", { name: "Reset Filters" })).toBeInTheDocument(); - expect(screen.getByText("Page 1 of 1")).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Previous" })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Next" })).toBeInTheDocument(); - }); - - it("should display object permissions when present", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - object_permission: { - object_permission_id: "perm-1", - mcp_servers: ["server1"], - vector_stores: ["store1"], - }, - }) - ); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - }); - - it("should display soft budget in settings view when present", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - soft_budget: 500.75, - max_budget: 1000, - }) - ); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByText("Team Settings")).toBeInTheDocument(); - }); - - await waitFor(() => { - expect(screen.getByText(/Soft Budget:/)).toBeInTheDocument(); - expect(screen.getByText(/\$500\.75/)).toBeInTheDocument(); - }); - }); - - it("should open Settings tab by default when editTeam is true and user can edit", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - expect(screen.getByText("Team Settings")).toBeInTheDocument(); - }); - - it("should open Overview tab by default when editTeam is false", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - expect(screen.getByText("Budget Status")).toBeInTheDocument(); - }); - - it("should open Overview tab by default when editTeam is true but user cannot edit", async () => { - vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); - - renderWithProviders( - - ); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - expect(screen.getByText("Budget Status")).toBeInTheDocument(); - }); - - it("should display soft budget alerting emails in settings view when present", async () => { - const user = userEvent.setup({ delay: null }); - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - metadata: { - soft_budget_alerting_emails: ["alert1@test.com", "alert2@test.com"], - }, - }) - ); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByText("Team Settings")).toBeInTheDocument(); - }); - - await waitFor(() => { - expect(screen.getByText(/Soft Budget Alerting Emails:/)).toBeInTheDocument(); - expect(screen.getByText(/alert1@test\.com, alert2@test\.com/)).toBeInTheDocument(); - }); - }); - - it("should pass access_group_ids to teamUpdateCall when saving team settings", async () => { - const user = userEvent.setup({ delay: null }); - const accessGroupIds = ["ag-1", "ag-2"]; - vi.mocked(networking.teamInfoCall).mockResolvedValue( - createMockTeamData({ - access_group_ids: accessGroupIds, - models: ["gpt-4"], - }) - ); - vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); - - renderWithProviders(); - - await waitFor(() => { - const teamNameElements = screen.queryAllByText("Test Team"); - expect(teamNameElements.length).toBeGreaterThan(0); - }); - - const settingsTab = screen.getByRole("tab", { name: "Settings" }); - await user.click(settingsTab); - - await waitFor(() => { - expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); - }); - - const editButton = screen.getByRole("button", { name: /edit settings/i }); - await user.click(editButton); - - await waitFor(() => { - expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); - }); - - const saveButton = screen.getByRole("button", { name: /save changes/i }); - await user.click(saveButton); - - await waitFor(() => { - expect(networking.teamUpdateCall).toHaveBeenCalledWith( - "test-token", - expect.objectContaining({ - access_group_ids: accessGroupIds, - team_id: "123", }) ); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Budget Status")).toBeInTheDocument(); + }); + }); + + it("should display guardrails in overview when present", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + guardrails: ["guardrail1", "guardrail2"], + }) + ); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Guardrails")).toBeInTheDocument(); + }); + }); + + it("should display policies in overview when present", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + policies: ["policy1"], + }) + ); + vi.mocked(networking.getPolicyInfoWithGuardrails).mockResolvedValue({ + resolved_guardrails: ["guardrail1"], + }); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Policies")).toBeInTheDocument(); + }); + }); + + it("should display team member budget information when present", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + team_member_budget_table: { + max_budget: 500, + budget_duration: "30d", + tpm_limit: 5000, + rpm_limit: 50, + }, + }) + ); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Budget Status")).toBeInTheDocument(); + }); + }); + + it("should display virtual keys information", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue({ + ...createMockTeamData(), + keys: [ + { user_id: "user1", token: "key1" }, + { token: "key2" }, + ], + }); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("tab", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + }); + + it("should display object permissions when present", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + object_permission: { + object_permission_id: "perm-1", + mcp_servers: ["server1"], + vector_stores: ["store1"], + }, + }) + ); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + }); + + it("should open Settings tab by default when editTeam is true and user can edit", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + it("should open Overview tab by default when editTeam is false", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + expect(screen.getByText("Budget Status")).toBeInTheDocument(); + }); + + it("should open Overview tab by default when editTeam is true but user cannot edit", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders( + + ); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + expect(screen.getByText("Budget Status")).toBeInTheDocument(); + }); + }); + + describe("tabs and navigation", () => { + it("should show members tab when user can edit team", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("tab", { name: "Members" })).toBeInTheDocument(); + }); + }); + + it("should not show members tab when user cannot edit team", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + expect(screen.queryByRole("tab", { name: "Members" })).not.toBeInTheDocument(); + }); + + it("should show settings tab when user can edit team", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("tab", { name: "Settings" })).toBeInTheDocument(); + }); + }); + + it("should navigate to settings tab when clicked", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + }); + + it("should call onClose when back button is clicked", async () => { + const user = userEvent.setup({ delay: null }); + const onClose = vi.fn(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const backButton = screen.getByRole("button", { name: /back to teams/i }); + await user.click(backButton); + + expect(onClose).toHaveBeenCalled(); + }); + + it("should copy team ID to clipboard when copy button is clicked", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const copyButtons = screen.getAllByRole("button"); + const copyButton = copyButtons.find((btn) => btn.querySelector("svg")); + expect(copyButton).toBeTruthy(); + + if (copyButton) { + await user.click(copyButton); + } + }); + + it("should show Virtual Keys tab when user cannot edit team", async () => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByRole("tab", { name: "Virtual Keys" })).toBeInTheDocument(); + }); + }); + + it("should display X Members in Virtual Keys tab when navigated to", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + const fiveKeys = Array.from({ length: 5 }, (_, i) => ({ + token: `sk-${i}`, + token_id: `key-${i}`, + key_alias: `key_${i}`, + key_name: `sk-...${i}`, + user_id: `user-${i}`, + organization_id: null, + user: { user_id: `user-${i}`, user_email: `user${i}@test.com` }, + created_at: "2024-01-01T00:00:00Z", + team_id: "123", + spend: 0, + max_budget: 100, + models: ["gpt-4"], + })); + mockUseKeys.mockReturnValue({ + data: { keys: fiveKeys, total_count: 5, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const virtualKeysTab = screen.getByRole("tab", { name: "Virtual Keys" }); + await user.click(virtualKeysTab); + + await waitFor(() => { + expect(screen.getByText("Page 1 of 1")).toBeInTheDocument(); + }); + }); + + it("should show Filters and pagination controls in Virtual Keys tab", async () => { + const user = userEvent.setup(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + mockUseKeys.mockReturnValue({ + data: { + keys: [ + { + token: "sk-1", + token_id: "key-1", + key_alias: "key1", + key_name: "sk-...1", + user_id: "user-1", + organization_id: null, + user: { user_id: "user-1", user_email: "user1@test.com" }, + created_at: "2024-01-01T00:00:00Z", + team_id: "123", + spend: 0, + max_budget: 100, + models: ["gpt-4"], + }, + ], + total_count: 1, + current_page: 1, + total_pages: 1, + }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + } as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const virtualKeysTab = screen.getByRole("tab", { name: "Virtual Keys" }); + await user.click(virtualKeysTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: "Filters" })).toBeInTheDocument(); + }); + expect(screen.getByRole("button", { name: "Reset Filters" })).toBeInTheDocument(); + expect(screen.getByText("Page 1 of 1")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Previous" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Next" })).toBeInTheDocument(); + }); + }); + + describe("settings and editing", () => { + it("should open edit mode when edit button is clicked", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + const editButton = screen.getByRole("button", { name: /edit settings/i }); + await user.click(editButton); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + }); + + it("should close edit mode when cancel button is clicked", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData()); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + const editButton = screen.getByRole("button", { name: /edit settings/i }); + await user.click(editButton); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + + const cancelButton = screen.getByRole("button", { name: /cancel/i }); + await user.click(cancelButton); + + await waitFor(() => { + expect(screen.queryByLabelText("Team Name")).not.toBeInTheDocument(); + }); + }); + + it("should disable secret manager settings for non-premium users", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + secret_manager_settings: { provider: "aws", secret_id: "abc" }, + }, + }) + ); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + const editButton = screen.getByRole("button", { name: /edit settings/i }); + await user.click(editButton); + + const secretField = await screen.findByPlaceholderText( + '{"namespace": "admin", "mount": "secret", "path_prefix": "litellm"}' + ); + expect(secretField).toBeDisabled(); + }); + + it("should allow premium users to edit secret manager settings", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + secret_manager_settings: { provider: "aws", secret_id: "abc" }, + }, + }) + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + const editButton = screen.getByRole("button", { name: /edit settings/i }); + await user.click(editButton); + + const secretField = await screen.findByPlaceholderText( + '{"namespace": "admin", "mount": "secret", "path_prefix": "litellm"}' + ); + expect(secretField).not.toBeDisabled(); + }); + + it("should add team member when form is submitted", async () => { + const user = userEvent.setup({ delay: null }); + const onUpdate = vi.fn(); + const teamData = createMockTeamData(); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamData); + vi.mocked(networking.teamMemberAddCall).mockResolvedValue({} as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const membersTab = screen.getByRole("tab", { name: "Members" }); + await user.click(membersTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /add member/i })).toBeInTheDocument(); + }); + + const addButton = screen.getByRole("button", { name: /add member/i }); + await user.click(addButton); + + await waitFor(() => { + expect(screen.getByRole("button", { name: "Submit" })).toBeInTheDocument(); + }); + + const submitButton = screen.getByRole("button", { name: "Submit" }); + await user.click(submitButton); + + await waitFor(() => { + expect(networking.teamMemberAddCall).toHaveBeenCalled(); + }); + }); + + it("should display soft budget in settings view when present", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + soft_budget: 500.75, + max_budget: 1000, + }) + ); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + await waitFor(() => { + expect(screen.getByText(/Soft Budget:/)).toBeInTheDocument(); + expect(screen.getByText(/\$500\.75/)).toBeInTheDocument(); + }); + }); + + it("should display soft budget alerting emails in settings view when present", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + soft_budget_alerting_emails: ["alert1@test.com", "alert2@test.com"], + }, + }) + ); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + await waitFor(() => { + expect(screen.getByText(/Soft Budget Alerting Emails:/)).toBeInTheDocument(); + expect(screen.getByText(/alert1@test\.com, alert2@test\.com/)).toBeInTheDocument(); + }); + }); + + it("should pass access_group_ids to teamUpdateCall when saving team settings", async () => { + const user = userEvent.setup({ delay: null }); + const accessGroupIds = ["ag-1", "ag-2"]; + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + access_group_ids: accessGroupIds, + models: ["gpt-4"], + }) + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + const editButton = screen.getByRole("button", { name: /edit settings/i }); + await user.click(editButton); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + + const saveButton = screen.getByRole("button", { name: /save changes/i }); + await user.click(saveButton); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + access_group_ids: accessGroupIds, + team_id: "123", + }) + ); + }); }); }); }); diff --git a/ui/litellm-dashboard/vitest.config.ts b/ui/litellm-dashboard/vitest.config.ts index 7c52b88d3b6..d2b6e7b43bf 100644 --- a/ui/litellm-dashboard/vitest.config.ts +++ b/ui/litellm-dashboard/vitest.config.ts @@ -7,7 +7,7 @@ export default defineConfig({ setupFiles: ["tests/setupTests.ts"], globals: true, css: true, // lets you import CSS/modules without extra mocks - testTimeout: 10000, + testTimeout: 30000, coverage: { provider: "v8", reporter: ["text", "lcov"], From 92dbd2c4919f5a95d841893c761328b80ea43ca0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 9 Apr 2026 20:54:36 -0700 Subject: [PATCH 11/38] address greptile review feedback (greploop iteration 1) Remove leftover 10000ms per-test timeout in add_model_tab.test.tsx that was missed in the initial sweep. The test now inherits the 30000ms global. --- .../src/components/add_model/add_model_tab.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 59970547c1a..0a8d2d9124e 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -175,7 +175,7 @@ describe("Add Model Tab", () => { ); expect(await screen.findByRole("tab", { name: "Add Model" })).toBeInTheDocument(); - }, 10000); // This test is flaky, adding a timeout until we find a better solution + }); it("should display both Add Model and Add Auto Router tabs", async () => { const props = createTestProps(); From ce0b57b4ffc0c56250a97f192188fca2bd5c946f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 9 Apr 2026 21:04:34 -0700 Subject: [PATCH 12/38] [Docs] Add missing MCP per-user token env vars to config_settings MCP_PER_USER_TOKEN_DEFAULT_TTL and MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS were added in #25441 but not documented, causing test_env_keys.py to fail. --- docs/my-website/docs/proxy/config_settings.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 88cbcac52cc..c64d475fdaa 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -602,6 +602,8 @@ router_settings: | MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE | Maximum number of entries in MCP OAuth2 token cache. Default is 200 | MCP_OAUTH2_TOKEN_CACHE_MIN_TTL | Minimum TTL in seconds for MCP OAuth2 token cache. Default is 10 | MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from token expiry when computing cache TTL. Default is 60 +| MCP_PER_USER_TOKEN_DEFAULT_TTL | Default TTL in seconds for per-user MCP OAuth tokens stored in Redis. Default is 43200 (12 hours) +| MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from per-user MCP OAuth token expiry when computing Redis TTL. Default is 60 | DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20 | DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10 | DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602 From 26e99f22b341107ee0b26cd4224d2e51101cefa0 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 9 Apr 2026 21:36:35 -0700 Subject: [PATCH 13/38] refactor: consolidate route auth for UI and API tokens Unify UI and API token authorization through the shared RBAC path and backfill missing routes in role-based route lists. --- litellm/proxy/_types.py | 134 ++++++++---------- litellm/proxy/auth/auth_checks.py | 67 +-------- tests/proxy_unit_tests/test_jwt.py | 75 ++++++---- .../test_user_api_key_auth.py | 125 ++++++++++------ 4 files changed, 189 insertions(+), 212 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 793742891fc..364e49e6257 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -494,10 +494,12 @@ class LiteLLMRoutes(enum.Enum): "/v2/key/info", "/model_group/info", "/health", + "/health/services", "/key/list", "/user/filter/ui", "/models", "/v1/models", + "/sso/get/ui_settings", ] # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend @@ -566,6 +568,8 @@ class LiteLLMRoutes(enum.Enum): "/spend/tags", "/spend/calculate", "/spend/logs", + "/spend/logs/ui", + "/spend/logs/session/ui", "/cost/estimate", ] @@ -581,6 +585,7 @@ class LiteLLMRoutes(enum.Enum): "/global/spend/report", "/global/spend/provider", "/global/spend/tags", + "/global/spend/all_tag_names", ] public_routes = set( @@ -602,44 +607,18 @@ class LiteLLMRoutes(enum.Enum): ] ) - ui_routes = [ - "/sso", - "/sso/get/ui_settings", - "/get/ui_settings", - "/login", - "/key/info", - "/config", - "/spend", - "/model/info", - "/v2/model/info", - "/v2/key/info", - "/models", - "/v1/models", - "/global/spend", - "/global/spend/logs", - "/global/spend/keys", - "/global/spend/models", - "/global/spend/tags", - "/global/predict/spend/logs", - "/global/activity", - "/health/services", - ] + info_routes - internal_user_routes = ( [ - "/global/spend/tags", - "/global/spend/keys", - "/global/spend/models", - "/global/spend/provider", - "/global/spend/end_users", "/global/activity", "/global/activity/model", + "/global/activity/cache_hits", "/v1/models/{model_id}", "/models/{model_id}", "/guardrails/list", "/v2/guardrails/list", ] + spend_tracking_routes + + global_spend_tracking_routes + key_management_routes ) @@ -694,6 +673,9 @@ class LiteLLMRoutes(enum.Enum): "/tag/list", "/audit", "/audit/{id}", + "/global/activity", + "/global/activity/model", + "/global/activity/cache_hits", ] + info_routes # All routes accesible by an Org Admin @@ -892,9 +874,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1036,9 +1018,9 @@ class RegenerateKeyRequest(GenerateKeyRequest): spend: Optional[float] = None metadata: Optional[dict] = None new_master_key: Optional[str] = None - grace_period: Optional[str] = ( - None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke - ) + grace_period: Optional[ + str + ] = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke class ResetSpendRequest(LiteLLMPydanticObjectBase): @@ -1562,12 +1544,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1590,12 +1572,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -1685,15 +1667,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" team_member_budget_duration: Optional[str] = None # e.g. "30d", "1mo" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1790,9 +1772,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -2134,9 +2116,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2495,9 +2477,9 @@ class UserAPIKeyAuth( user_max_budget: Optional[float] = None request_route: Optional[str] = None user: Optional[Any] = None # Expanded user object when expand=user is used - created_by_user: Optional[Any] = ( - None # Expanded created_by user when expand=user is used - ) + created_by_user: Optional[ + Any + ] = None # Expanded created_by user when expand=user is used end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None # Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery # and forwarded into outbound tokens by guardrails such as MCPJWTSigner. @@ -2636,9 +2618,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available litellm_budget_table: Optional[LiteLLM_BudgetTable] = None user_email: Optional[str] = None @@ -3793,9 +3775,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -4050,9 +4032,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -4214,9 +4196,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 68bde8434a6..56958a88f6d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -196,9 +196,7 @@ def _is_model_cost_zero( return True -def _is_cost_explicitly_configured( - model: str, llm_router: "Router" -) -> bool: +def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool: """ Check if any deployment in the model group has cost fields explicitly set in its litellm.model_cost entry. @@ -215,10 +213,7 @@ def _is_cost_explicitly_configured( if model_id is None: continue raw_entry = litellm.model_cost.get(model_id, {}) - if ( - "input_cost_per_token" in raw_entry - or "output_cost_per_token" in raw_entry - ): + if "input_cost_per_token" in raw_entry or "output_cost_per_token" in raw_entry: return True return False @@ -596,17 +591,12 @@ async def common_checks( # noqa: PLR0915 user_object=user_object, route=route, request_body=request_body ) - token_team = getattr(valid_token, "team_id", None) - token_type: Literal["ui", "api"] = ( - "ui" if token_team is not None and token_team == "litellm-dashboard" else "api" - ) - _is_route_allowed = _is_allowed_route( + _is_route_allowed = _is_api_route_allowed( route=route, - token_type=token_type, - user_obj=user_object, request=request, request_data=request_body, valid_token=valid_token, + user_obj=user_object, ) # 11. [OPTIONAL] Vector store checks - is the object allowed to access the vector store @@ -629,31 +619,6 @@ async def common_checks( # noqa: PLR0915 return True -def _is_ui_route( - route: str, - user_obj: Optional[LiteLLM_UserTable] = None, -) -> bool: - """ - - Check if the route is a UI used route - """ - # this token is only used for managing the ui - allowed_routes = LiteLLMRoutes.ui_routes.value - # check if the current route startswith any of the allowed routes - if ( - route is not None - and isinstance(route, str) - and any(route.startswith(allowed_route) for allowed_route in allowed_routes) - ): - # Do something if the current route starts with any of the allowed routes - return True - elif any( - RouteChecks._route_matches_pattern(route=route, pattern=allowed_route) - for allowed_route in allowed_routes - ): - return True - return False - - def _get_user_role( user_obj: Optional[LiteLLM_UserTable], ) -> Optional[LitellmUserRoles]: @@ -717,30 +682,6 @@ def _is_user_proxy_admin(user_obj: Optional[LiteLLM_UserTable]): return False -def _is_allowed_route( - route: str, - token_type: Literal["ui", "api"], - request: Request, - request_data: dict, - valid_token: Optional[UserAPIKeyAuth], - user_obj: Optional[LiteLLM_UserTable] = None, -) -> bool: - """ - - Route b/w ui token check and normal token check - """ - - if token_type == "ui" and _is_ui_route(route=route, user_obj=user_obj): - return True - else: - return _is_api_route_allowed( - route=route, - request=request, - request_data=request_data, - valid_token=valid_token, - user_obj=user_obj, - ) - - def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: """ Return if a user is allowed to access route. Helper function for `allowed_routes_check`. diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index a5be1a3a42d..73f956a6147 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -715,7 +715,7 @@ async def aaaatest_user_token_output( assert team_result.user_id == user_id -@pytest.mark.parametrize("admin_allowed_routes", [None, ["ui_routes"]]) +@pytest.mark.parametrize("admin_allowed_routes", [None, ["info_routes"]]) @pytest.mark.parametrize("audience", [None, "litellm-proxy"]) @pytest.mark.asyncio async def test_allowed_routes_admin( @@ -934,10 +934,7 @@ async def mock_user_object(*args, **kwargs): user_id = kwargs.get("user_id") user_email = kwargs.get("user_email") return LiteLLM_UserTable( - spend=0, - user_id=user_id, - max_budget=None, - user_email=user_email + spend=0, user_id=user_id, max_budget=None, user_email=user_email ) @@ -1170,15 +1167,13 @@ async def test_end_user_jwt_auth(monkeypatch): # use generated key to auth in from litellm import Router from litellm.types.router import RouterGeneralSettings - + # Create a router with pass_through_all_models enabled router = Router( model_list=[], - router_general_settings=RouterGeneralSettings( - pass_through_all_models=True - ), + router_general_settings=RouterGeneralSettings(pass_through_all_models=True), ) - + setattr(litellm.proxy.proxy_server, "premium_user", True) setattr( litellm.proxy.proxy_server, @@ -1196,7 +1191,7 @@ async def test_end_user_jwt_auth(monkeypatch): cost_tracking() result = await user_api_key_auth(request=request, api_key=bearer_token) - + # Assert that end_user_id is correctly extracted from JWT token's 'sub' field assert result.end_user_id == "81b3e52a-67a6-4efb-9645-70527e101479" @@ -1228,7 +1223,9 @@ async def test_end_user_jwt_auth(monkeypatch): ), ) - with patch("litellm.acompletion", new=AsyncMock(return_value=mock_response)) as mock_completion: + with patch( + "litellm.acompletion", new=AsyncMock(return_value=mock_response) + ) as mock_completion: resp = await chat_completion( request=request, fastapi_response=temp_response, @@ -1243,10 +1240,13 @@ async def test_end_user_jwt_auth(monkeypatch): # Verify the completion was called with correct end_user_id mock_completion.assert_called_once() call_kwargs = mock_completion.call_args.kwargs - + # end_user_id is passed in metadata as 'user_api_key_end_user_id' metadata = call_kwargs.get("metadata", {}) - assert metadata.get("user_api_key_end_user_id") == "81b3e52a-67a6-4efb-9645-70527e101479" + assert ( + metadata.get("user_api_key_end_user_id") + == "81b3e52a-67a6-4efb-9645-70527e101479" + ) def test_can_rbac_role_call_route(): @@ -1278,13 +1278,13 @@ def test_user_api_key_auth_jwt_hashing(): """ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.handle_jwt import JWTHandler - + # Test with a JWT token (3 parts separated by dots) jwt_token = "test-jwt-token-header.payload.signature" - + # Create UserAPIKeyAuth instance with JWT user_auth = UserAPIKeyAuth(api_key=jwt_token) - + # Verify that the API key is hashed with "hashed-jwt-" prefix # critical - the raw JWT token should not be in the api_key or token assert user_auth.api_key.startswith("hashed-jwt-") @@ -1292,19 +1292,18 @@ def test_user_api_key_auth_jwt_hashing(): assert jwt_token not in user_auth.api_key assert jwt_token not in user_auth.token - # Test with a regular API key (should not be hashed) regular_api_key = "sk-1234567890abcdef" user_auth_regular = UserAPIKeyAuth(api_key=regular_api_key) - + # Verify that regular API key is hashed normally (without "hashed-jwt-" prefix) assert not user_auth_regular.api_key.startswith("hashed-jwt-") assert not user_auth_regular.token.startswith("hashed-jwt-") - + # Test with a non-JWT, non-sk string (should not be hashed) non_jwt_key = "some-random-key" user_auth_non_jwt = UserAPIKeyAuth(api_key=non_jwt_key) - + # Verify that non-JWT key is not hashed assert user_auth_non_jwt.api_key == non_jwt_key assert user_auth_non_jwt.token == non_jwt_key @@ -1315,19 +1314,19 @@ def test_jwt_handler_is_jwt_static_method(): Test that JWTHandler.is_jwt is a static method and works correctly """ from litellm.proxy.auth.handle_jwt import JWTHandler - + # Test with valid JWT format valid_jwt = "test-jwt-token-header.payload.signature" assert JWTHandler.is_jwt(valid_jwt) == True - + # Test with invalid JWT format (only 2 parts) invalid_jwt = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ" assert JWTHandler.is_jwt(invalid_jwt) == False - + # Test with regular API key regular_key = "sk-1234567890abcdef" assert JWTHandler.is_jwt(regular_key) == False - + # Test with empty string assert JWTHandler.is_jwt("") == False @@ -1461,7 +1460,13 @@ async def test_auth_jwt_es256_jwk_path(monkeypatch): now = int(time.time()) token = jwt.encode( - {"sub": "alice", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + { + "sub": "alice", + "aud": "litellm-proxy", + "iss": "http://example", + "iat": now, + "exp": now + 300, + }, ec_priv_pem, algorithm="ES256", headers={"kid": "ec1"}, @@ -1508,7 +1513,13 @@ async def test_auth_jwt_rs256_regression(monkeypatch): now = int(time.time()) token = jwt.encode( - {"sub": "bob", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + { + "sub": "bob", + "aud": "litellm-proxy", + "iss": "http://example", + "iat": now, + "exp": now + 300, + }, rsa_priv_pem, algorithm="RS256", headers={"kid": "rsa1"}, @@ -1540,7 +1551,13 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch): ) now = int(time.time()) token = jwt.encode( - {"sub": "mallory", "aud": "litellm-proxy", "iss": "http://example", "iat": now, "exp": now + 300}, + { + "sub": "mallory", + "aud": "litellm-proxy", + "iss": "http://example", + "iat": now, + "exp": now + 300, + }, ec_priv_pem, algorithm="ES256", headers={"kid": "ec1"}, @@ -1566,4 +1583,4 @@ async def test_auth_jwt_mismatched_key_fails(monkeypatch): with patch.object(h, "get_public_key", new=AsyncMock(return_value=rsa_jwk)): with pytest.raises(Exception) as exc: await h.auth_jwt(token) - assert "Validation fails" in str(exc.value) \ No newline at end of file + assert "Validation fails" in str(exc.value) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 1a6e2eda9a1..75f0d5e3195 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -359,27 +359,38 @@ async def test_auth_with_allowed_routes(route, should_raise_error): @pytest.mark.parametrize( - "route, user_role, expected_result", + "route, user_role, should_be_allowed", [ - # Proxy Admin checks + # Admin can access everything + ("/config/update", "proxy_admin", True), ("/global/spend/logs", "proxy_admin", True), - ("/key/delete", "proxy_admin", False), - ("/key/generate", "proxy_admin", False), - ("/key/regenerate", "proxy_admin", False), - # Internal User checks - allowed routes + ("/global/activity/cache_hits", "proxy_admin", True), + # Internal User - allowed read-only routes ("/global/spend/logs", "internal_user", True), - ("/key/delete", "internal_user", False), - ("/key/generate", "internal_user", False), - ("/key/82akk800000000jjsk/regenerate", "internal_user", False), - # Internal User Viewer - ("/key/generate", "internal_user_viewer", False), - # Internal User checks - disallowed routes + ("/spend/logs/ui", "internal_user", True), + ("/global/activity/cache_hits", "internal_user", True), + ("/health/services", "internal_user", True), + # Internal User - BLOCKED from admin routes (security fix) + ("/config/update", "internal_user", False), + ("/config/pass_through_endpoint", "internal_user", False), + ("/config/field/update", "internal_user", False), ("/organization/member_add", "internal_user", False), + # Internal User Viewer - allowed spend routes only + ("/spend/logs/ui", "internal_user_viewer", True), + ("/global/spend/all_tag_names", "internal_user_viewer", True), + # Internal User Viewer - blocked from admin routes + ("/config/update", "internal_user_viewer", False), + ("/key/generate", "internal_user_viewer", False), ], ) -def test_is_ui_route_allowed(route, user_role, expected_result): - from litellm.proxy.auth.auth_checks import _is_ui_route - from litellm.proxy._types import LiteLLM_UserTable +def test_ui_token_route_access(route, user_role, should_be_allowed): + """ + Verify that UI tokens (team_id=litellm-dashboard) go through the same + RBAC checks as API tokens. Non-admin dashboard users must not be able + to access admin-only routes like /config/update. + """ + from litellm.proxy.auth.auth_checks import _is_api_route_allowed + from litellm.proxy._types import LiteLLM_UserTable, UserAPIKeyAuth user_obj = LiteLLM_UserTable( user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", @@ -395,18 +406,36 @@ def test_is_ui_route_allowed(route, user_role, expected_result): organization_memberships=[], ) - received_args: dict = { - "route": route, - "user_obj": user_obj, - } - try: - assert _is_ui_route(**received_args) == expected_result - except Exception as e: - # If expected result is False, we expect an error - if expected_result is False: - pass - else: - raise e + valid_token = UserAPIKeyAuth( + user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297", + team_id="litellm-dashboard", + user_role=user_role, + ) + + from starlette.datastructures import URL + from fastapi import Request + + request = Request(scope={"type": "http"}) + request._url = URL(url=route) + + if should_be_allowed: + result = _is_api_route_allowed( + route=route, + request=request, + request_data={}, + valid_token=valid_token, + user_obj=user_obj, + ) + assert result is True + else: + with pytest.raises(Exception): + _is_api_route_allowed( + route=route, + request=request, + request_data={}, + valid_token=valid_token, + user_obj=user_obj, + ) @pytest.mark.parametrize( @@ -684,7 +713,7 @@ async def test_soft_budget_alert(): def test_is_allowed_route(): - from litellm.proxy.auth.auth_checks import _is_allowed_route + from litellm.proxy.auth.auth_checks import _is_api_route_allowed from litellm.proxy._types import UserAPIKeyAuth import datetime @@ -692,7 +721,6 @@ def test_is_allowed_route(): args = { "route": "/embeddings", - "token_type": "api", "request": request, "request_data": {"input": ["hello world"], "model": "embedding-small"}, "valid_token": UserAPIKeyAuth( @@ -752,7 +780,7 @@ def test_is_allowed_route(): "user_obj": None, } - assert _is_allowed_route(**args) + assert _is_api_route_allowed(**args) @pytest.mark.parametrize( @@ -836,7 +864,6 @@ async def test_user_api_key_auth_websocket(): with patch( "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True ) as mock_user_api_key_auth: - # Make the call to the WebSocket function await user_api_key_auth_websocket(mock_websocket) @@ -845,10 +872,14 @@ async def test_user_api_key_auth_websocket(): # Get the request object that was passed to user_api_key_auth request_arg = mock_user_api_key_auth.call_args.kwargs["request"] - + # Verify that the request has headers set - assert hasattr(request_arg, "headers"), "Request object should have headers attribute" - assert "authorization" in request_arg.headers, "Request headers should contain authorization" + assert hasattr( + request_arg, "headers" + ), "Request object should have headers attribute" + assert ( + "authorization" in request_arg.headers + ), "Request headers should contain authorization" assert request_arg.headers["authorization"] == "Bearer some_api_key" assert ( @@ -1036,7 +1067,10 @@ async def test_jwt_non_admin_team_route_access(monkeypatch): # Create request request = Request( - scope={"type": "http", "headers": [(b"authorization", b"Bearer fake.jwt.token")]} + scope={ + "type": "http", + "headers": [(b"authorization", b"Bearer fake.jwt.token")], + } ) request._url = URL(url="/team/new") @@ -1101,14 +1135,14 @@ async def test_x_litellm_api_key(): ignored_key = "aj12445" # Create request with headers as bytes - request = Request( - scope={ - "type": "http" - } - ) + request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - valid_token = await user_api_key_auth(request=request, api_key="Bearer " + ignored_key, custom_litellm_key_header=master_key) + valid_token = await user_api_key_auth( + request=request, + api_key="Bearer " + ignored_key, + custom_litellm_key_header=master_key, + ) assert valid_token.token == hash_token(master_key) @@ -1123,7 +1157,9 @@ async def test_user_api_key_from_query_param(): from litellm.proxy.proxy_server import hash_token, user_api_key_cache user_key = "sk-query-1234" - user_api_key_cache.set_cache(key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key))) + user_api_key_cache.set_cache( + key=hash_token(user_key), value=UserAPIKeyAuth(token=hash_token(user_key)) + ) setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") @@ -1136,7 +1172,9 @@ async def test_user_api_key_from_query_param(): "query_string": f"alt=sse&key={user_key}".encode(), } ) - request._url = URL(url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}") + request._url = URL( + url=f"/v1beta/models/gemini:streamGenerateContent?alt=sse&key={user_key}" + ) async def return_body(): return b"{}" @@ -1145,4 +1183,3 @@ async def test_user_api_key_from_query_param(): valid_token = await user_api_key_auth(request=request, api_key="") assert valid_token.token == hash_token(user_key) - From fb527ae25020494eb5ad14e90e8c849c1202751e Mon Sep 17 00:00:00 2001 From: joereyna Date: Thu, 9 Apr 2026 13:18:35 -0700 Subject: [PATCH 14/38] fix(test): mock headers in test_completion_fine_tuned_model --- tests/local_testing/test_amazing_vertex_completion.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index a1684e23769..98bb40f3613 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -2291,6 +2291,8 @@ def test_prompt_factory_nested(): async def test_completion_fine_tuned_model(): load_vertex_ai_credentials() mock_response = AsyncMock() + mock_response.headers = {} + mock_response.status_code = 200 def return_val(): return { @@ -2326,7 +2328,6 @@ async def test_completion_fine_tuned_model(): } mock_response.json = return_val - mock_response.status_code = 200 expected_payload = { "contents": [ From f8ae6427363cf3c1c5fd2b1d03d9162ffb68ef00 Mon Sep 17 00:00:00 2001 From: joereyna Date: Thu, 9 Apr 2026 15:28:34 -0700 Subject: [PATCH 15/38] format vertex test file --- tests/local_testing/test_amazing_vertex_completion.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 98bb40f3613..001b9464006 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -178,7 +178,6 @@ async def test_get_response(): async def test_aavertex_ai_anthropic_async(): # load_vertex_ai_credentials() try: - model = "claude-3-5-sonnet@20240620" vertex_ai_project = "pathrise-convert-1606954137718" @@ -351,7 +350,6 @@ def test_avertex_ai_stream(): @pytest.mark.flaky(retries=3, delay=1) @pytest.mark.asyncio async def test_async_vertexai_response_basic(): - load_vertex_ai_credentials() try: user_message = "Hello, how are you?" @@ -1382,7 +1380,6 @@ async def test_gemini_pro_json_schema_args_sent_httpx( ] ) elif resp is not None: - assert resp.model == model.split("/")[1] From 8dc5ab39f00beb604ee82cbc91d5e8d6053aaa81 Mon Sep 17 00:00:00 2001 From: Chetan Soni Date: Thu, 9 Apr 2026 12:42:42 -0700 Subject: [PATCH 16/38] feat(mcp): add per-user OAuth token storage for interactive MCP flows --- litellm/constants.py | 9 + litellm/proxy/_experimental/mcp_server/db.py | 145 ++++- .../mcp_server/discoverable_endpoints.py | 195 ++++++- .../mcp_server/mcp_server_manager.py | 31 ++ .../mcp_server/oauth2_token_cache.py | 108 ++++ .../proxy/_experimental/mcp_server/server.py | 120 +++- .../types/mcp_server/mcp_server_manager.py | 9 + tests/mcp_tests/test_per_user_oauth_cache.py | 527 ++++++++++++++++++ .../mcp_tools/OAuthFormFields.test.tsx | 208 +++++++ .../components/mcp_tools/OAuthFormFields.tsx | 46 +- .../mcp_tools/create_mcp_server.test.tsx | 141 +++++ .../mcp_tools/create_mcp_server.tsx | 14 + .../mcp_tools/mcp_server_edit.test.tsx | 249 +++++++++ .../components/mcp_tools/mcp_server_edit.tsx | 73 ++- .../src/components/mcp_tools/types.tsx | 4 + 15 files changed, 1851 insertions(+), 28 deletions(-) create mode 100644 tests/mcp_tests/test_per_user_oauth_cache.py create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx diff --git a/litellm/constants.py b/litellm/constants.py index a7d86ddb16b..337cb1243fb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -135,6 +135,15 @@ MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int( MCP_NPM_CACHE_DIR = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache") MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10")) +# Per-user OAuth token Redis cache (for server-side token storage) +MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX = "mcp:per_user_token" +MCP_PER_USER_TOKEN_DEFAULT_TTL = int( + os.getenv("MCP_PER_USER_TOKEN_DEFAULT_TTL", "43200") # 12 hours +) +MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int( + os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60") +) + # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0")) MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index fbef33c32ed..e9bd41bb951 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -21,7 +21,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy.utils import PrismaClient +from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.mcp import MCPCredentials @@ -576,6 +578,7 @@ async def store_user_oauth_credential( refresh_token: Optional[str] = None, expires_in: Optional[int] = None, scopes: Optional[List[str]] = None, + skip_byok_guard: bool = False, ) -> None: """Persist an OAuth2 access token for a user+server pair. @@ -604,21 +607,26 @@ async def store_user_oauth_credential( # Guard against silently overwriting a BYOK credential with an OAuth token. # BYOK credentials lack a "type" field (or use a non-"oauth2" type). - existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( - where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} - ) - if existing is not None: - _byok_error = ValueError( - f"A non-OAuth2 credential already exists for user {user_id} " - f"and server {server_id}. Refusing to overwrite." + # Skip the guard when the caller knows the row is already an OAuth2 credential + # (e.g. during token refresh), saving an extra DB round-trip. + if not skip_byok_guard: + existing = await prisma_client.db.litellm_mcpusercredentials.find_unique( + where={"user_id_server_id": {"user_id": user_id, "server_id": server_id}} ) - try: - raw = json.loads(base64.urlsafe_b64decode(existing.credential_b64).decode()) - except Exception: - # Credential is not base64+JSON — it's a plain-text BYOK key. - raise _byok_error - if raw.get("type") != "oauth2": - raise _byok_error + if existing is not None: + _byok_error = ValueError( + f"A non-OAuth2 credential already exists for user {user_id} " + f"and server {server_id}. Refusing to overwrite." + ) + try: + raw = json.loads( + base64.urlsafe_b64decode(existing.credential_b64).decode() + ) + except Exception: + # Credential is not base64+JSON — it's a plain-text BYOK key. + raise _byok_error + if raw.get("type") != "oauth2": + raise _byok_error encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode() await prisma_client.db.litellm_mcpusercredentials.upsert( @@ -697,6 +705,115 @@ async def list_user_oauth_credentials( return results +async def refresh_user_oauth_token( + prisma_client: PrismaClient, + user_id: str, + server: Any, + cred: Dict[str, Any], +) -> Optional[Dict[str, Any]]: + """Attempt to refresh a per-user OAuth2 token using its stored refresh_token. + + POSTs to ``server.token_url`` with ``grant_type=refresh_token``. + + On success: persists the new credential via ``store_user_oauth_credential`` + and returns the updated payload dict. + On failure (network error, invalid_grant, missing refresh_token, …): logs a + warning and returns ``None`` — the caller is responsible for clearing the + stale credential and triggering re-authentication. + """ + refresh_token: Optional[str] = cred.get("refresh_token") + token_url: Optional[str] = getattr(server, "token_url", None) + server_id: str = getattr(server, "server_id", "") + client_id: Optional[str] = getattr(server, "client_id", None) + client_secret: Optional[str] = getattr(server, "client_secret", None) + + if not refresh_token: + verbose_proxy_logger.debug( + "refresh_user_oauth_token: no refresh_token stored for user=%s server=%s", + user_id, + server_id, + ) + return None + if not token_url: + verbose_proxy_logger.debug( + "refresh_user_oauth_token: server=%s has no token_url configured", + server_id, + ) + return None + + token_data: Dict[str, str] = { + "grant_type": "refresh_token", + "refresh_token": refresh_token, + } + if client_id: + token_data["client_id"] = client_id + if client_secret: + token_data["client_secret"] = client_secret + + try: + async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.Oauth2Check + ) + response = await async_client.post( + token_url, + headers={"Accept": "application/json"}, + data=token_data, + ) + response.raise_for_status() + body: Dict[str, Any] = response.json() + except Exception as exc: + verbose_proxy_logger.warning( + "refresh_user_oauth_token: refresh request failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + return None + + access_token: Optional[str] = body.get("access_token") + if not access_token: + verbose_proxy_logger.warning( + "refresh_user_oauth_token: token response missing access_token for " + "user=%s server=%s", + user_id, + server_id, + ) + return None + + expires_in: Optional[int] = None + raw_expires = body.get("expires_in") + try: + expires_in = int(raw_expires) if raw_expires is not None else None + except (TypeError, ValueError): + pass + + # Rotate refresh token when the provider returns a new one + new_refresh_token: Optional[str] = body.get("refresh_token") or refresh_token + + raw_scope = body.get("scope") + scopes: Optional[List[str]] = ( + raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + ) or cred.get("scopes") + + await store_user_oauth_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=server_id, + access_token=access_token, + refresh_token=new_refresh_token, + expires_in=expires_in, + scopes=scopes, + skip_byok_guard=True, # Row is already OAuth2; skip the extra find_unique check + ) + + verbose_proxy_logger.info( + "refresh_user_oauth_token: refreshed token for user=%s server=%s", + user_id, + server_id, + ) + return await get_user_oauth_credential(prisma_client, user_id, server_id) + + async def approve_mcp_server( prisma_client: PrismaClient, server_id: str, diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 07309eb57f2..d0d61986322 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,10 +1,11 @@ import json -from typing import Optional +from typing import Any, Dict, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -147,6 +148,160 @@ def _resolve_oauth2_server_for_root_endpoints( return None +def _validate_token_response( + token_response: Dict[str, Any], + validation_rules: Dict[str, Any], + server_id: str, +) -> None: + """Raise HTTPException 403 if any validation rule doesn't match the token response. + + Supports dot-notation for nested fields (e.g. ``"team.enterprise_id"`` checks + ``token_response["team"]["enterprise_id"]``). Top-level keys are tried first, + then dot-split traversal. All comparisons are string-coerced so that numeric + values in the response (e.g. ``"org_id": 12345``) match string rules + (``"org_id": "12345"``). + """ + for key, expected in validation_rules.items(): + actual: Any = token_response.get(key) + # Try dot-notation traversal when top-level lookup returns None + if actual is None and "." in key: + obj: Any = token_response + for part in key.split("."): + if isinstance(obj, dict): + obj = obj.get(part) + else: + obj = None + break + actual = obj + # Treat absent fields as a distinct failure from a mismatched value + if actual is None: + raise HTTPException( + status_code=403, + detail={ + "error": "token_validation_failed", + "server_id": server_id, + "field": key, + "message": ( + f"OAuth token rejected: required field '{key}' is absent" + ), + }, + ) + if str(actual) != str(expected): + raise HTTPException( + status_code=403, + detail={ + "error": "token_validation_failed", + "server_id": server_id, + "field": key, + "message": ( + f"OAuth token rejected: '{key}' = '{actual}', " + f"expected '{expected}'" + ), + }, + ) + + +async def _extract_user_id_from_request(request: Request) -> Optional[str]: + """Best-effort extraction of LiteLLM user_id from the request's Authorization header. + + Called at the OAuth token endpoint so that per-user tokens can be stored + server-side. Uses a read-only cache lookup to avoid re-running the full + auth pipeline (which has side effects such as rate-limit increments and + spend logging). Returns ``None`` if no cached credential is found. + """ + auth_header = request.headers.get("Authorization") or request.headers.get( + "authorization" + ) + if not auth_header: + return None + lower = auth_header.lower() + if not lower.startswith("bearer "): + return None + token = auth_header[7:].strip() + try: + from litellm.proxy._types import hash_token # noqa: PLC0415 + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + cached = await user_api_key_cache.async_get_cache(hash_token(token)) + return getattr(cached, "user_id", None) + except Exception: + return None + + +async def _store_per_user_token_server_side( + server: MCPServer, + user_id: str, + token_response: Dict[str, Any], +) -> None: + """Persist the OAuth token server-side and warm the Redis cache. + + Called from the token endpoint after a successful code exchange or refresh. + Errors are logged but NOT re-raised — the token is always returned to the + client even when server-side storage fails. + """ + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415 + _compute_per_user_token_ttl, + mcp_per_user_token_cache, + ) + from litellm.proxy.utils import get_prisma_client_or_throw # noqa: PLC0415 + + access_token: Optional[str] = token_response.get("access_token") + if not access_token: + return + + raw_expires = token_response.get("expires_in") + try: + expires_in: Optional[int] = int(raw_expires) if raw_expires is not None else None + except (TypeError, ValueError): + expires_in = None + + refresh_token: Optional[str] = token_response.get("refresh_token") or None + raw_scope = token_response.get("scope") + scopes: Optional[list] = ( + raw_scope.split() if isinstance(raw_scope, str) and raw_scope else None + ) + + try: + prisma_client = get_prisma_client_or_throw( + "Database not connected. Cannot store per-user OAuth token." + ) + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 + store_user_oauth_credential, + ) + + await store_user_oauth_credential( + prisma_client=prisma_client, + user_id=user_id, + server_id=server.server_id, + access_token=access_token, + refresh_token=refresh_token, + expires_in=expires_in, + scopes=scopes, + ) + verbose_logger.info( + "_store_per_user_token_server_side: stored token for user=%s server=%s", + user_id, + server.server_id, + ) + except Exception as exc: + verbose_logger.warning( + "_store_per_user_token_server_side: DB storage failed for user=%s server=%s: %s", + user_id, + server.server_id, + exc, + ) + return # Don't warm Redis if DB write failed + + # Warm the Redis cache so the first subsequent MCP call is a cache hit + ttl = _compute_per_user_token_ttl(server, expires_in) + await mcp_per_user_token_cache.set( + user_id=user_id, + server_id=server.server_id, + access_token=access_token, + ttl=ttl, + ) + + async def authorize_with_server( request: Request, mcp_server: MCPServer, @@ -266,6 +421,44 @@ async def exchange_token_with_server( token_response = response.json() access_token = token_response["access_token"] + # Validate token response against server-configured rules before any storage. + # This rejects tokens from wrong Slack workspaces, Atlassian orgs, etc. + if mcp_server.token_validation and isinstance(mcp_server.token_validation, dict): + _validate_token_response( + token_response=token_response, + validation_rules=mcp_server.token_validation, + server_id=mcp_server.server_id, + ) + + # Store server-side when the server is configured for per-user OAuth and + # the calling client has provided a valid LiteLLM identity. + # Errors are non-fatal: the token is still returned to the client. + if mcp_server.needs_user_oauth_token: + user_id = await _extract_user_id_from_request(request) + if user_id: + try: + await _store_per_user_token_server_side( + server=mcp_server, + user_id=user_id, + token_response=token_response, + ) + except Exception as exc: + verbose_logger.warning( + "exchange_token_with_server: server-side storage failed " + "for user=%s server=%s: %s", + user_id, + mcp_server.server_id, + exc, + ) + else: + verbose_logger.debug( + "exchange_token_with_server: no LiteLLM user_id found in request; " + "per-user token for server=%s will not be stored server-side. " + "The client should call POST /mcp/server/{id}/oauth-user-credential " + "to store it manually.", + mcp_server.server_id, + ) + result = { "access_token": access_token, "token_type": token_response.get("token_type", "Bearer"), diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 402e12d9356..8d3831e75fb 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2455,6 +2455,37 @@ class MCPServerManager: ) tasks.append(during_hook_task) + # For per-user OAuth servers: if the client didn't supply a token in + # oauth2_headers, look up the stored token from Redis / DB. This is the + # call_tool equivalent of _get_user_oauth_extra_headers_from_db used in + # list_tools. + if ( + mcp_server.needs_user_oauth_token + and not oauth2_headers + and user_api_key_auth is not None + ): + user_id = getattr(user_api_key_auth, "user_id", None) + if user_id: + try: + from litellm.proxy._experimental.mcp_server.server import ( # noqa: PLC0415 + _get_user_oauth_extra_headers_from_db, + ) + + stored_headers = await _get_user_oauth_extra_headers_from_db( + server=mcp_server, + user_api_key_auth=user_api_key_auth, + ) + if stored_headers: + oauth2_headers = stored_headers + except Exception as _lookup_exc: + verbose_logger.debug( + "call_tool: per-user token lookup failed for " + "user=%s server=%s: %s", + user_id, + mcp_server.server_id, + _lookup_exc, + ) + # For OpenAPI servers, call the tool handler directly instead of via MCP client if mcp_server.spec_path: verbose_logger.debug( diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 84a2e94467b..476e215666e 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -17,8 +17,15 @@ from litellm.constants import ( MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_PER_USER_TOKEN_DEFAULT_TTL, + MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, + MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -152,6 +159,107 @@ class MCPOAuth2TokenCache(InMemoryCache): mcp_oauth2_token_cache = MCPOAuth2TokenCache() +def _compute_per_user_token_ttl(server: "MCPServer", expires_in: Optional[int]) -> int: + """Compute Redis TTL for a per-user token. + + Uses server.token_storage_ttl_seconds when configured; otherwise derives + TTL from expires_in minus the expiry buffer; falls back to the default TTL. + """ + if server.token_storage_ttl_seconds is not None: + return max(server.token_storage_ttl_seconds, 1) + if expires_in is not None: + return max( + expires_in - MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS, + 1, + ) + return MCP_PER_USER_TOKEN_DEFAULT_TTL + + +class MCPPerUserTokenCache: + """Redis-backed cache for per-user OAuth2 access tokens. + + Uses LiteLLM's existing ``user_api_key_cache`` (DualCache with optional + Redis backend). Tokens are NaCl-encrypted with ``encrypt_value_helper`` + before storage so they are safe at rest in Redis. + + Redis key format: ``mcp:per_user_token:{user_id}:{server_id}`` + Redis value: ``encrypt_value_helper(access_token)`` — URL-safe base64 + """ + + def _cache_key(self, user_id: str, server_id: str) -> str: + return f"{MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX}:{user_id}:{server_id}" + + async def get(self, user_id: str, server_id: str) -> Optional[str]: + """Return the plaintext access_token, or None on miss/error.""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + encrypted = await user_api_key_cache.async_get_cache(key) + if encrypted is None: + return None + plaintext = decrypt_value_helper( + encrypted, + key="mcp_per_user_token", + exception_type="debug", + ) + return plaintext or None + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.get failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + return None + + async def set( + self, + user_id: str, + server_id: str, + access_token: str, + ttl: int, + ) -> None: + """Store NaCl-encrypted access_token in Redis with the given TTL.""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + encrypted = encrypt_value_helper(access_token) + await user_api_key_cache.async_set_cache(key, encrypted, ttl=ttl) + verbose_logger.debug( + "MCPPerUserTokenCache.set: cached token for user=%s server=%s ttl=%ds", + user_id, + server_id, + ttl, + ) + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.set failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + + async def delete(self, user_id: str, server_id: str) -> None: + """Invalidate the cached token (removes from both in-memory and Redis layers).""" + try: + from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 + + key = self._cache_key(user_id, server_id) + await user_api_key_cache.async_delete_cache(key) + except Exception as exc: + verbose_logger.debug( + "MCPPerUserTokenCache.delete failed for user=%s server=%s: %s", + user_id, + server_id, + exc, + ) + + +mcp_per_user_token_cache = MCPPerUserTokenCache() + + async def resolve_mcp_auth( server: "MCPServer", mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7fc28b68e9c..99578d006e1 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -896,11 +896,17 @@ if MCP_AVAILABLE: user_api_key_auth: Optional[UserAPIKeyAuth], prefetched_creds: Optional[Dict[str, Dict[str, Any]]] = None, ) -> Optional[Dict[str, str]]: - """Look up stored OAuth2 token for (user, server) from DB and return as extra_headers dict. + """Look up stored OAuth2 token for (user, server) and return as extra_headers dict. + + Lookup order: + 1. Redis cache (fast path, NaCl-decrypted) — skipped when prefetched_creds supplied + 2. prefetched_creds dict (pre-fetched batch DB query) or fresh DB query + 3. Auto-refresh when the stored token is expired and a refresh_token exists Args: prefetched_creds: Optional dict keyed by server_id with credential payloads. - When provided, avoids a per-server DB round-trip. + When provided, the Redis and individual DB lookups are + skipped in favour of the pre-fetched batch result. """ if server.auth_type != MCPAuth.oauth2: return None @@ -914,8 +920,27 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 get_user_oauth_credential, is_oauth_credential_expired, + refresh_user_oauth_token, + ) + from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: PLC0415 + _compute_per_user_token_ttl, + mcp_per_user_token_cache, ) + # ── Fast path: Redis cache ──────────────────────────────────────── + # Only used when prefetched_creds is not supplied (individual lookup). + if prefetched_creds is None: + cached_token = await mcp_per_user_token_cache.get(user_id, server_id) + if cached_token is not None: + verbose_logger.debug( + "_get_user_oauth_extra_headers_from_db: Redis hit for " + "user=%s server=%s", + user_id, + server_id, + ) + return {"Authorization": f"Bearer {cached_token}"} + + # ── Slow path: DB lookup ────────────────────────────────────────── if prefetched_creds is not None: cred = prefetched_creds.get(server_id) else: @@ -929,18 +954,83 @@ if MCP_AVAILABLE: cred = await get_user_oauth_credential( prisma_client, user_id, server_id ) - if cred and cred.get("access_token"): - if is_oauth_credential_expired(cred): - verbose_logger.debug( - f"_get_user_oauth_extra_headers_from_db: token expired for " - f"user={user_id} server={server_id}" - ) + + if not cred or not cred.get("access_token"): + return None + + if is_oauth_credential_expired(cred): + verbose_logger.debug( + "_get_user_oauth_extra_headers_from_db: token expired for " + "user=%s server=%s — attempting refresh", + user_id, + server_id, + ) + # Attempt token refresh; requires a DB client (not available from prefetch) + if cred.get("refresh_token"): + try: + from litellm.proxy.utils import ( # noqa: PLC0415 + get_prisma_client_or_throw, + ) + + prisma_client = get_prisma_client_or_throw( + "Database not connected. Cannot refresh OAuth token." + ) + cred = await refresh_user_oauth_token( + prisma_client=prisma_client, + user_id=user_id, + server=server, + cred=cred, + ) + except Exception as refresh_exc: + verbose_logger.warning( + "_get_user_oauth_extra_headers_from_db: refresh failed " + "for user=%s server=%s: %s", + user_id, + server_id, + refresh_exc, + ) + cred = None + + if not cred or not cred.get("access_token"): + # Clear stale Redis/cache entry so we don't serve it again. + # Do this for both the individual and prefetch paths so the + # next request doesn't get a stale cache hit. + await mcp_per_user_token_cache.delete(user_id, server_id) return None - return {"Authorization": f"Bearer {cred['access_token']}"} + + access_token: str = cred["access_token"] + + # Warm (or re-warm) the Redis cache from the DB result. + # Always write regardless of whether expires_at is present — tokens + # without an expiry are still valid and should be cached using the + # server/default TTL so subsequent requests are fast. + if prefetched_creds is None: + raw_expires = None + expires_at = cred.get("expires_at") + if expires_at: + from datetime import datetime, timezone # noqa: PLC0415 + + try: + exp_dt = datetime.fromisoformat(expires_at) + if exp_dt.tzinfo is None: + exp_dt = exp_dt.replace(tzinfo=timezone.utc) + remaining = int( + (exp_dt - datetime.now(timezone.utc)).total_seconds() + ) + raw_expires = max(remaining, 0) if remaining > 0 else None + except (ValueError, TypeError): + pass + ttl = _compute_per_user_token_ttl(server, raw_expires) + await mcp_per_user_token_cache.set(user_id, server_id, access_token, ttl) + + return {"Authorization": f"Bearer {access_token}"} except Exception as e: verbose_logger.warning( - f"_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " - f"user={user_id} server={server_id}: {e}" + "_get_user_oauth_extra_headers_from_db: failed to retrieve credential for " + "user=%s server=%s: %s", + user_id, + server_id, + e, ) return None @@ -2504,6 +2594,14 @@ if MCP_AVAILABLE: server_name, client_ip=_client_ip ) if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: + # For servers that store per-user tokens server-side, skip the + # pre-emptive 401 — the call_tool / list_tools dispatch will look + # up the stored token from Redis / DB and only fail at the MCP + # protocol level if none is found, giving the client a proper + # tool-execution error rather than an HTTP 401. + if server.needs_user_oauth_token: + continue + request = StarletteRequest(scope) base_url = get_request_base_url(request) diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index db7657a0174..a7d0968c0ef 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -71,6 +71,15 @@ class MCPServer(BaseModel): # OAuth2 flow type. Defaults to None (interactive / authorization_code). # Set to "client_credentials" to enable M2M token fetching. oauth2_flow: Optional[Literal["client_credentials", "authorization_code"]] = None + # Per-user OAuth server-side storage config. + # token_validation: key-value pairs that must match fields in the OAuth token + # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). + # Tokens that fail validation are rejected before storage. + token_validation: Optional[Dict[str, Any]] = None + # Optional TTL override (seconds) for the Redis per-user token cache. + # Defaults to the token's expires_in minus the expiry buffer, or + # MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent. + token_storage_ttl_seconds: Optional[int] = None model_config = ConfigDict(arbitrary_types_allowed=True) @property diff --git a/tests/mcp_tests/test_per_user_oauth_cache.py b/tests/mcp_tests/test_per_user_oauth_cache.py new file mode 100644 index 00000000000..36c26a5a505 --- /dev/null +++ b/tests/mcp_tests/test_per_user_oauth_cache.py @@ -0,0 +1,527 @@ +""" +Unit tests for per-user MCP OAuth token storage: +- MCPPerUserTokenCache (NaCl-encrypted Redis cache) +- _validate_token_response (token validation rules) +- _compute_per_user_token_ttl (TTL computation) +- refresh_user_oauth_token (token refresh flow) +""" + +import sys +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# Stub out modules that aren't available in the unit-test environment +# so we can import the targets without a full proxy stack. +for _mod in ("orjson",): + if _mod not in sys.modules: + sys.modules[_mod] = MagicMock() + +from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( # noqa: E402 + MCPPerUserTokenCache, + _compute_per_user_token_ttl, + mcp_per_user_token_cache, +) +from litellm.types.mcp import MCPAuth, MCPTransport # noqa: E402 +from litellm.types.mcp_server.mcp_server_manager import MCPServer # noqa: E402 + + +def _import_validate(): + """Lazy import to avoid pulling orjson at collection time.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _validate_token_response, + ) + + return _validate_token_response + + +# ── Fixtures ───────────────────────────────────────────────────────────────── + + +def _make_server(**kwargs) -> MCPServer: + defaults: Dict[str, Any] = { + "server_id": "slack-test", + "name": "Slack", + "server_name": "slack", + "url": "https://slack-mcp.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "client_id": "SLACK_CLIENT_ID", + "client_secret": "SLACK_CLIENT_SECRET", + "token_url": "https://slack.com/api/oauth.v2.access", + "authorization_url": "https://slack.com/oauth/v2/authorize", + } + defaults.update(kwargs) + return MCPServer(**defaults) + + +# ── _validate_token_response ────────────────────────────────────────────────── + + +class TestValidateTokenResponse: + def test_passes_when_all_rules_match(self): + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "enterprise_id": "E04XXXXXX", + "team": {"id": "T123", "name": "Acme"}, + } + # Should not raise + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + + def test_raises_on_mismatch(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "xoxb-123", "enterprise_id": "E99999999"} + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + detail = exc_info.value.detail + assert detail["error"] == "token_validation_failed" + assert detail["field"] == "enterprise_id" + + def test_raises_when_field_absent(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "xoxb-123"} + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + # Absent field should produce a distinct "absent" message, not str(None) + assert "absent" in exc_info.value.detail["message"] + + def test_absent_field_does_not_match_string_none(self): + """str(None)='None' must NOT match the string rule value 'None'.""" + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = {"access_token": "tok"} # enterprise_id absent + # Even if admin writes validation_rules={"enterprise_id": "None"}, absent + # field should raise, not pass. + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"enterprise_id": "None"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + assert "absent" in exc_info.value.detail["message"] + + def test_dot_notation_nested_field(self): + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "team": {"enterprise_id": "E04XXXXXX"}, + } + # Should not raise — dot-notation traverses nested dict + _validate_token_response( + token_response=token_response, + validation_rules={"team.enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + + def test_dot_notation_mismatch(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = { + "access_token": "xoxb-123", + "team": {"enterprise_id": "WRONG"}, + } + with pytest.raises(HTTPException) as exc_info: + _validate_token_response( + token_response=token_response, + validation_rules={"team.enterprise_id": "E04XXXXXX"}, + server_id="slack-test", + ) + assert exc_info.value.status_code == 403 + assert exc_info.value.detail["field"] == "team.enterprise_id" + + def test_numeric_value_string_coercion(self): + """Numeric values in token response should match string rules.""" + _validate_token_response = _import_validate() + token_response = {"access_token": "tok", "org_id": 12345} + # Should not raise — str(12345) == "12345" + _validate_token_response( + token_response=token_response, + validation_rules={"org_id": "12345"}, + server_id="test", + ) + + def test_multiple_rules_all_must_match(self): + from fastapi import HTTPException + + _validate_token_response = _import_validate() + token_response = { + "access_token": "tok", + "enterprise_id": "E04XXXXXX", + "cloud_id": "WRONG_CLOUD", + } + with pytest.raises(HTTPException): + _validate_token_response( + token_response=token_response, + validation_rules={ + "enterprise_id": "E04XXXXXX", + "cloud_id": "abc-123", + }, + server_id="atlassian", + ) + + +# ── _compute_per_user_token_ttl ────────────────────────────────────────────── + + +class TestComputePerUserTokenTtl: + def test_uses_server_override_when_set(self): + server = _make_server(token_storage_ttl_seconds=7200) + assert _compute_per_user_token_ttl(server, expires_in=99999) == 7200 + + def test_uses_expires_in_minus_buffer(self): + server = _make_server() + # Default buffer is 60s + ttl = _compute_per_user_token_ttl(server, expires_in=3600) + assert ttl == 3600 - 60 + + def test_minimum_ttl_is_1(self): + server = _make_server() + # expires_in smaller than buffer → clamp to 1 + ttl = _compute_per_user_token_ttl(server, expires_in=30) + assert ttl == 1 + + def test_default_ttl_when_expires_in_none(self): + from litellm.constants import MCP_PER_USER_TOKEN_DEFAULT_TTL + + server = _make_server() + ttl = _compute_per_user_token_ttl(server, expires_in=None) + assert ttl == MCP_PER_USER_TOKEN_DEFAULT_TTL + + +# ── MCPPerUserTokenCache ────────────────────────────────────────────────────── + + +class TestMCPPerUserTokenCache: + """Tests for Redis-backed per-user token cache. + + Patches ``user_api_key_cache`` to avoid needing a real Redis instance. + Patches ``encrypt_value_helper`` / ``decrypt_value_helper`` to verify + encryption is applied before Redis writes and decryption after reads. + """ + + @pytest.fixture + def cache(self): + return MCPPerUserTokenCache() + + @pytest.fixture + def mock_dual_cache(self): + dc = MagicMock() + dc.async_get_cache = AsyncMock(return_value=None) + dc.async_set_cache = AsyncMock() + return dc + + @pytest.mark.asyncio + async def test_get_returns_none_on_miss(self, cache, mock_dual_cache): + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper" + ) as mock_decrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = None + result = await cache.get("alice", "slack-test") + assert result is None + mock_decrypt.assert_not_called() + + @pytest.mark.asyncio + async def test_get_decrypts_cached_value(self, cache, mock_dual_cache): + fake_encrypted = "encrypted_blob_abc123" + fake_plaintext = "xoxb-slack-token" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper", + return_value=fake_plaintext, + ) as mock_decrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = fake_encrypted + result = await cache.get("alice", "slack-test") + + assert result == fake_plaintext + mock_decrypt.assert_called_once_with( + fake_encrypted, + key="mcp_per_user_token", + exception_type="debug", + ) + + @pytest.mark.asyncio + async def test_set_encrypts_before_storing(self, cache, mock_dual_cache): + fake_encrypted = "encrypted_blob_xyz" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value=fake_encrypted, + ) as mock_encrypt, patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.set("alice", "slack-test", "xoxb-token", ttl=3540) + + mock_encrypt.assert_called_once_with("xoxb-token") + mock_dual_cache.async_set_cache.assert_called_once() + call_kwargs = mock_dual_cache.async_set_cache.call_args + assert call_kwargs[0][1] == fake_encrypted # encrypted value stored + assert call_kwargs[1]["ttl"] == 3540 + + @pytest.mark.asyncio + async def test_set_uses_correct_cache_key(self, cache, mock_dual_cache): + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value="enc", + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.set("bob", "github-server", "ghp_token", ttl=3600) + + key_used = mock_dual_cache.async_set_cache.call_args[0][0] + assert key_used == "mcp:per_user_token:bob:github-server" + + @pytest.mark.asyncio + async def test_delete_calls_async_delete_cache(self, cache, mock_dual_cache): + mock_dual_cache.async_delete_cache = AsyncMock() + with patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + await cache.delete("alice", "slack-test") + + mock_dual_cache.async_delete_cache.assert_called_once_with( + "mcp:per_user_token:alice:slack-test" + ) + mock_dual_cache.async_set_cache.assert_not_called() + + @pytest.mark.asyncio + async def test_get_returns_none_on_decrypt_failure(self, cache, mock_dual_cache): + """Cache misses and decrypt errors should both return None without raising.""" + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.decrypt_value_helper", + return_value=None, # decrypt returns None on failure + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + mock_dual_cache.async_get_cache.return_value = "bad_encrypted_data" + result = await cache.get("alice", "slack-test") + + assert result is None + + @pytest.mark.asyncio + async def test_set_is_noop_on_cache_error(self, cache, mock_dual_cache): + """Errors in the cache layer must not propagate to the caller.""" + mock_dual_cache.async_set_cache.side_effect = RuntimeError("Redis down") + with patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.encrypt_value_helper", + return_value="enc", + ), patch( + "litellm.proxy.proxy_server.user_api_key_cache", mock_dual_cache + ): + # Should not raise + await cache.set("alice", "slack-test", "token", ttl=3600) + + +# ── refresh_user_oauth_token ────────────────────────────────────────────────── + + +class TestRefreshUserOauthToken: + """Tests for the DB-level token refresh helper.""" + + @pytest.fixture + def server(self): + return _make_server() + + @pytest.fixture + def cred(self): + return { + "type": "oauth2", + "access_token": "OLD_TOKEN", + "refresh_token": "REFRESH_TOKEN_123", + "expires_at": ( + datetime.now(timezone.utc) - timedelta(hours=1) + ).isoformat(), + } + + @pytest.mark.asyncio + async def test_returns_none_when_no_refresh_token(self, server): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + cred = {"type": "oauth2", "access_token": "OLD"} # no refresh_token + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_returns_none_when_no_token_url(self, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + server = _make_server(token_url=None) + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_returns_none_on_http_error(self, server, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + mock_client = AsyncMock() + mock_client.post.side_effect = Exception("Connection refused") + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ): + result = await refresh_user_oauth_token( + prisma_client=MagicMock(), + user_id="alice", + server=server, + cred=cred, + ) + assert result is None + + @pytest.mark.asyncio + async def test_stores_and_returns_new_credential(self, server, cred): + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + new_token_response = MagicMock() + new_token_response.json.return_value = { + "access_token": "NEW_TOKEN", + "expires_in": 3600, + "refresh_token": "NEW_REFRESH", + "scope": "channels:read chat:write", + } + new_token_response.raise_for_status = MagicMock() + + mock_client = AsyncMock() + mock_client.post.return_value = new_token_response + + stored_cred = { + "type": "oauth2", + "access_token": "NEW_TOKEN", + "refresh_token": "NEW_REFRESH", + } + mock_prisma = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ), patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new_callable=AsyncMock, + ) as mock_store, patch( + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new_callable=AsyncMock, + return_value=stored_cred, + ): + result = await refresh_user_oauth_token( + prisma_client=mock_prisma, + user_id="alice", + server=server, + cred=cred, + ) + + assert result == stored_cred + mock_store.assert_called_once() + call_kwargs = mock_store.call_args[1] + assert call_kwargs["access_token"] == "NEW_TOKEN" + assert call_kwargs["refresh_token"] == "NEW_REFRESH" + assert call_kwargs["expires_in"] == 3600 + assert call_kwargs["scopes"] == ["channels:read", "chat:write"] + # Refresh path must skip the BYOK guard (row is already OAuth2) + assert call_kwargs.get("skip_byok_guard") is True + + @pytest.mark.asyncio + async def test_falls_back_to_old_refresh_token_when_not_rotated( + self, server, cred + ): + """When provider doesn't return a new refresh_token, keep the old one.""" + from litellm.proxy._experimental.mcp_server.db import refresh_user_oauth_token + + new_token_response = MagicMock() + new_token_response.json.return_value = { + "access_token": "NEW_TOKEN", + "expires_in": 3600, + # No refresh_token in response + } + new_token_response.raise_for_status = MagicMock() + + mock_client = AsyncMock() + mock_client.post.return_value = new_token_response + + with patch( + "litellm.proxy._experimental.mcp_server.db.get_async_httpx_client", + return_value=mock_client, + ), patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new_callable=AsyncMock, + ) as mock_store, patch( + "litellm.proxy._experimental.mcp_server.db.get_user_oauth_credential", + new_callable=AsyncMock, + return_value={"type": "oauth2", "access_token": "NEW_TOKEN"}, + ): + await refresh_user_oauth_token( + prisma_client=AsyncMock(), + user_id="alice", + server=server, + cred=cred, + ) + + call_kwargs = mock_store.call_args[1] + # Old refresh_token preserved when provider doesn't rotate + assert call_kwargs["refresh_token"] == "REFRESH_TOKEN_123" + + +# ── MCPServer new fields ────────────────────────────────────────────────────── + + +class TestMCPServerNewFields: + def test_token_validation_default_none(self): + server = _make_server() + assert server.token_validation is None + + def test_token_validation_set(self): + server = _make_server(token_validation={"enterprise_id": "E04XXXXXX"}) + assert server.token_validation == {"enterprise_id": "E04XXXXXX"} + + def test_token_storage_ttl_default_none(self): + server = _make_server() + assert server.token_storage_ttl_seconds is None + + def test_token_storage_ttl_set(self): + server = _make_server(token_storage_ttl_seconds=7200) + assert server.token_storage_ttl_seconds == 7200 + + def test_needs_user_oauth_token_true_for_oauth2_without_m2m(self): + server = _make_server(auth_type=MCPAuth.oauth2) + assert server.needs_user_oauth_token is True + + def test_needs_user_oauth_token_false_for_m2m(self): + server = _make_server( + auth_type=MCPAuth.oauth2, + oauth2_flow="client_credentials", + ) + assert server.needs_user_oauth_token is False diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx new file mode 100644 index 00000000000..888f3066252 --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.test.tsx @@ -0,0 +1,208 @@ +import React from "react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { render, screen, waitFor, act, fireEvent } from "@testing-library/react"; +import { Form } from "antd"; +import OAuthFormFields from "./OAuthFormFields"; + +// ── helpers ────────────────────────────────────────────────────────────────── + +/** Minimal Ant Form wrapper so Form.Item registers correctly. */ +const WithForm: React.FC<{ children: React.ReactNode; onFinish?: (values: any) => void }> = ({ + children, + onFinish, +}) => { + const [form] = Form.useForm(); + return ( +
+ {children} + +
+ ); +}; + +// ── tests ───────────────────────────────────────────────────────────────────── + +describe("OAuthFormFields", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + // ── visibility by flow type ───────────────────────────────────────────────── + + describe("interactive mode (isM2M=false)", () => { + it("renders Token Validation Rules field", () => { + render( + + + , + ); + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + it("renders Token Storage TTL field", () => { + render( + + + , + ); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + + it("renders standard interactive fields alongside the new fields", () => { + render( + + + , + ); + expect(screen.getByText("Authorization URL (optional)")).toBeInTheDocument(); + expect(screen.getByText("Registration URL (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + }); + + describe("M2M mode (isM2M=true)", () => { + it("does NOT render Token Validation Rules field", () => { + render( + + + , + ); + expect(screen.queryByText("Token Validation Rules (optional)")).not.toBeInTheDocument(); + }); + + it("does NOT render Token Storage TTL field", () => { + render( + + + , + ); + expect(screen.queryByText("Token Storage TTL (seconds, optional)")).not.toBeInTheDocument(); + }); + + it("still renders M2M-specific fields", () => { + render( + + + , + ); + expect(screen.getByText("Client ID")).toBeInTheDocument(); + expect(screen.getByText("Token URL")).toBeInTheDocument(); + }); + }); + + // ── token_validation_json inline JSON validator ────────────────────────────── + + describe("token_validation_json validation", () => { + it("accepts empty value without error", async () => { + const onFinish = vi.fn(); + render( + + + , + ); + + // Leave the textarea empty and submit + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + + it("accepts a valid JSON object without error", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: '{"organization": "my-org"}' } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + + it("shows 'Must be valid JSON' error for malformed JSON", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "not-valid-json{" } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + }); + + it("shows error for a plain string value (not a JSON object)", async () => { + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + // A bare string is valid JSON but we still want to accept it; only truly + // unparseable text should fail. Bare "hello" is actually invalid JSON + // (no quotes), so it should fail. + fireEvent.change(textarea, { target: { value: "hello" } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + }); + + it("whitespace-only value is treated as empty and passes validation", async () => { + const onFinish = vi.fn(); + render( + + + , + ); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: " " } }); + }); + + const submitBtn = screen.getByRole("button", { name: "Submit" }); + await act(async () => { + fireEvent.click(submitBtn); + }); + + await waitFor(() => { + expect(screen.queryByText("Must be valid JSON")).not.toBeInTheDocument(); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx index 85487a8a479..4a808ca489d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/OAuthFormFields.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Form, Select, Tooltip } from "antd"; +import { Form, Input, InputNumber, Select, Tooltip } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { OAUTH_FLOW } from "./types"; @@ -151,6 +151,50 @@ const OAuthFormFields: React.FC = ({ > + + } + name="token_validation_json" + rules={[ + { + validator: (_: any, value: string) => { + if (!value || value.trim() === "") return Promise.resolve(); + try { + JSON.parse(value); + return Promise.resolve(); + } catch { + return Promise.reject(new Error("Must be valid JSON")); + } + }, + }, + ]} + > + + + + } + name="token_storage_ttl_seconds" + > + + {oauthFlow && (

diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx index d49c49446bb..c92956b430f 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.test.tsx @@ -353,6 +353,147 @@ describe("CreateMCPServer", () => { ); }); + describe("when OAuth interactive auth is selected", () => { + /** Select HTTP transport + OAuth auth, then wait for the OAuth form to appear. */ + async function setupOAuthInteractive() { + render(); + await selectAntOption("Transport Type", "Streamable HTTP"); + + await waitFor(() => { + expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument(); + }); + + await selectAntOption("Authentication", "OAuth"); + + // Wait for OAuthFormFields to render (OAuth Flow Type selector is the sentinel) + await waitFor(() => { + expect(screen.getByText("OAuth Flow Type")).toBeInTheDocument(); + }); + + // OAuthFormFields defaults to INTERACTIVE, so the new fields should appear + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + } + + it("shows Token Validation Rules and Token Storage TTL fields", async () => { + await setupOAuthInteractive(); + // Asserted in setupOAuthInteractive + }); + + it("includes token_validation in payload when token_validation_json is filled with valid JSON", async () => { + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-oauth", + server_name: "OAuth_Server", + alias: "OAuth_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "oauth2", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); + + await setupOAuthInteractive(); + + // Fill required form fields + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "OAuth_Server" } }); + }); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } }); + }); + + // Fill in the token_validation_json textarea + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: '{"organization": "my-org", "team.id": "42"}' } }); + }); + + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.token_validation).toEqual({ organization: "my-org", "team.id": "42" }); + }); + + it("omits token_validation from payload when token_validation_json is empty", async () => { + vi.mocked(networking.createMCPServer).mockResolvedValue({ + server_id: "new-server-oauth", + server_name: "OAuth_Server", + alias: "OAuth_Server", + url: "https://example.com/mcp", + transport: "http", + auth_type: "oauth2", + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + }); + + await setupOAuthInteractive(); + + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "OAuth_Server" } }); + }); + const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com"); + await act(async () => { + fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } }); + }); + + // Leave token_validation_json empty + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => { + expect(networking.createMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.createMCPServer).mock.calls[0]; + expect(payload.token_validation).toBeUndefined(); + }); + + it("does not submit and shows validation error for invalid JSON in token_validation_json", async () => { + await setupOAuthInteractive(); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "not-valid-json{" } }); + }); + + const nameInput = document.getElementById("server_name") as HTMLInputElement; + await act(async () => { + fireEvent.change(nameInput, { target: { value: "OAuth_Server" } }); + }); + + const submitButton = screen.getByRole("button", { name: "Add MCP Server" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Either the inline form validation message or the notification fires — + // both indicate the submit was blocked. + await waitFor(() => { + const inlineError = screen.queryByText("Must be valid JSON"); + const notCalled = !vi.mocked(networking.createMCPServer).mock.calls.length; + expect(inlineError !== null || notCalled).toBe(true); + }); + }); + }); + describe("when modal is cancelled", () => { it("should call setModalVisible(false) when cancel is clicked", async () => { render(); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 4c824fcee0b..45556bc18b1 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -284,6 +284,7 @@ const CreateMCPServer: React.FC = ({ credentials: credentialValues, allow_all_keys: allowAllKeysRaw, available_on_public_internet: availableOnPublicInternetRaw, + token_validation_json: rawTokenValidationJson, ...restValues } = values; @@ -356,6 +357,18 @@ const CreateMCPServer: React.FC = ({ restValues.transport = "http"; } + // Parse token_validation JSON if provided + let tokenValidation: Record | null = null; + if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") { + try { + tokenValidation = JSON.parse(rawTokenValidationJson); + } catch { + NotificationsManager.fromBackend("Invalid JSON in Token Validation Rules"); + setIsLoading(false); + return; + } + } + // Prepare the payload with cost configuration and allowed tools const payload: Record = { ...restValues, @@ -376,6 +389,7 @@ const CreateMCPServer: React.FC = ({ allow_all_keys: Boolean(allowAllKeysRaw), available_on_public_internet: Boolean(availableOnPublicInternetRaw), static_headers: staticHeaders, + ...(tokenValidation !== null && { token_validation: tokenValidation }), }; payload.static_headers = staticHeaders; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx index e33e2fff491..aba2a3d9222 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.test.tsx @@ -3,6 +3,7 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; import { render, screen, waitFor, fireEvent, act } from "@testing-library/react"; import MCPServerEdit from "./mcp_server_edit"; import * as networking from "../networking"; +import NotificationsManager from "../molecules/notifications_manager"; vi.mock("../networking", () => ({ updateMCPServer: vi.fn(), @@ -37,6 +38,29 @@ vi.mock("./mcp_tool_configuration", () => ({ default: () =>

, })); +// ── fixtures ────────────────────────────────────────────────────────────────── + +const interactiveOAuthServer = { + server_id: "oauth_server_1", + server_name: "OAuthServer", + alias: "oauth_server", // underscores: hyphens fail validateMCPServerName + description: "Interactive OAuth MCP server", + transport: "http", + url: "https://example.com/mcp", + auth_type: "oauth2", + // No token_url → edit form defaults to INTERACTIVE flow + token_url: null, + authorization_url: null, + registration_url: null, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + mcp_access_groups: [], +}; + +// ── test suites ─────────────────────────────────────────────────────────────── + describe("MCPServerEdit (stdio)", () => { beforeEach(() => { vi.clearAllMocks(); @@ -152,3 +176,228 @@ describe("MCPServerEdit (stdio)", () => { expect(payload.env).toEqual({ CIRCLECI_TOKEN: "new-token", CIRCLECI_BASE_URL: "https://circleci.com" }); }); }); + +describe("MCPServerEdit (interactive OAuth)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("renders Token Validation Rules and Token Storage TTL fields for interactive OAuth server", async () => { + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + }); + + // Note: The M2M flow hiding logic is tested via OAuthFormFields.test.tsx (isM2M prop directly), + // since Form.useWatch doesn't synchronously reflect initialValues in jsdom. + + it("pre-populates token_validation_json from existing server token_validation", async () => { + const tokenValidation = { organization: "my-org", "team.id": "123" }; + + render( + , + ); + + await waitFor(() => { + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + expect(textarea).not.toBeNull(); + const parsed = JSON.parse(textarea.value); + expect(parsed).toEqual(tokenValidation); + }); + }); + + it("includes token_validation in update payload when token_validation_json is filled", async () => { + const onSuccess = vi.fn(); + vi.mocked(networking.updateMCPServer).mockResolvedValue({ + ...interactiveOAuthServer, + token_validation: { organization: "my-org" }, + }); + + render( + , + ); + + // Wait for the form to mount and the token_validation_json field to appear + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: '{"organization": "my-org"}' } }); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + expect(payload.token_validation).toEqual({ organization: "my-org" }); + }); + + it("does not include token_validation in payload when field is empty and server had none", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue(interactiveOAuthServer); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + // Leave token_validation_json empty + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + expect(payload.token_validation).toBeUndefined(); + }); + + it("sends token_validation: null to clear an existing value when textarea is cleared", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue({ + ...interactiveOAuthServer, + token_validation: null, + }); + + render( + , + ); + + await waitFor(() => { + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + expect(textarea?.value).toContain("old-org"); + }); + + // Clear the textarea + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "" } }); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + // null signals the backend to clear the existing validation rules + expect(payload.token_validation).toBeNull(); + }); + + it("shows inline validation error and does not submit on invalid JSON in token_validation_json", async () => { + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Validation Rules (optional)")).toBeInTheDocument(); + }); + + const textarea = document.getElementById("token_validation_json") as HTMLTextAreaElement; + await act(async () => { + fireEvent.change(textarea, { target: { value: "{ bad json" } }); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + // The Form.Item inline validator intercepts invalid JSON before handleSave runs, + // so the inline error message appears and updateMCPServer is never called. + await waitFor(() => { + expect(screen.getByText("Must be valid JSON")).toBeInTheDocument(); + }); + expect(networking.updateMCPServer).not.toHaveBeenCalled(); + }); + + it("includes token_storage_ttl_seconds in payload when set", async () => { + vi.mocked(networking.updateMCPServer).mockResolvedValue({ + ...interactiveOAuthServer, + token_storage_ttl_seconds: 7200, + }); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Token Storage TTL (seconds, optional)")).toBeInTheDocument(); + }); + + const saveButtons = screen.getAllByRole("button", { name: "Save Changes" }); + await act(async () => { + fireEvent.click(saveButtons[0]); + }); + + await waitFor(() => { + expect(networking.updateMCPServer).toHaveBeenCalledTimes(1); + }); + + const [, payload] = vi.mocked(networking.updateMCPServer).mock.calls[0]; + expect(payload.token_storage_ttl_seconds).toBe(7200); + }); +}); diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index e81d2f3960e..1a3e30cb15d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Form, Select, Button as AntdButton, Tooltip, Input } from "antd"; +import { Form, Select, Button as AntdButton, Tooltip, Input, InputNumber } from "antd"; import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; import { AUTH_TYPE, OAUTH_FLOW, MCPServer, MCPServerCostInfo, TRANSPORT } from "./types"; @@ -190,6 +190,9 @@ const MCPServerEdit: React.FC = ({ transport: effectiveTransport, static_headers: initialStaticHeaders, oauth_flow_type: mcpServer.token_url ? OAUTH_FLOW.M2M : OAUTH_FLOW.INTERACTIVE, + token_validation_json: mcpServer.token_validation + ? JSON.stringify(mcpServer.token_validation, null, 2) + : undefined, }), [mcpServer, effectiveTransport, initialStaticHeaders, initialEnvJson], ); @@ -400,6 +403,7 @@ const MCPServerEdit: React.FC = ({ args: rawArgs, allow_all_keys: allowAllKeysRaw, available_on_public_internet: availableOnPublicInternetRaw, + token_validation_json: rawTokenValidationJson, ...restValues } = values; @@ -522,6 +526,17 @@ const MCPServerEdit: React.FC = ({ restValues.transport = "http"; } + // Parse token_validation JSON if provided + let tokenValidation: Record | null = null; + if (rawTokenValidationJson && rawTokenValidationJson.trim() !== "") { + try { + tokenValidation = JSON.parse(rawTokenValidationJson); + } catch { + NotificationsManager.fromBackend("Invalid JSON in Token Validation Rules"); + return; + } + } + // Prepare the payload with cost configuration and permission fields const mcpInfoServerName = restValues.server_name || @@ -556,6 +571,10 @@ const MCPServerEdit: React.FC = ({ static_headers: staticHeaders, allow_all_keys: Boolean(allowAllKeysRaw ?? mcpServer.allow_all_keys), available_on_public_internet: Boolean(availableOnPublicInternetRaw ?? mcpServer.available_on_public_internet), + // Include token_validation when it is set (non-null) or when clearing an existing value + ...(tokenValidation !== null || mcpServer.token_validation + ? { token_validation: tokenValidation } + : {}), }; const includeCredentials = restValues.auth_type && AUTH_TYPES_REQUIRING_CREDENTIALS.includes(restValues.auth_type); @@ -863,6 +882,58 @@ const MCPServerEdit: React.FC = ({ className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500" /> + {!isM2MFlow && ( + <> + + Token Validation Rules (optional) + + + + + } + name="token_validation_json" + rules={[ + { + validator: (_: any, value: string) => { + if (!value || value.trim() === "") return Promise.resolve(); + try { + JSON.parse(value); + return Promise.resolve(); + } catch { + return Promise.reject(new Error("Must be valid JSON")); + } + }, + }, + ]} + > + + + + Token Storage TTL (seconds, optional) + + + + + } + name="token_storage_ttl_seconds" + > + + + + )}

Use OAuth to fetch a fresh access token and temporarily save it in the session as the authentication value.