From b344043eaf3d0e758af8433865ad5b8f5e0c20ba Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 13 Aug 2026 18:00:42 -0700 Subject: [PATCH 1/3] fix(ui): stop a deselected MCP server keeping its grant on a virtual key (#36840) * fix(ui): stop a deselected MCP server keeping its grant on a virtual key The key editor sent `mcp_tool_permissions` unfiltered, and the MCP resolver counts a server named only under `mcp_tool_permissions` as entitled, unioning `tool_perm_servers` into `all_servers` at four sites in `user_api_key_auth_mcp.py`. Deselecting a server, or removing the access group that supplied it, therefore left a stale entry that kept the key reaching that server with its old tool allowlist attached. Reuse `extractMcpEntitlement`, which already landed for the internal-user surface, so the key surface drops an entry only once the server is known and no longer granted, and keeps it whenever a retained access group or toolset could still supply it. The helper moves to a shared module so the key template does not import a users page component. Setting the map unconditionally is part of the same fix: the old `Object.keys(...).length > 0` guard let the previous map ride through the `object_permission` spread, which filtering to an empty map would otherwise hit in exactly the case the fix is for. * fix(ui): resolve retained MCP groups and toolsets per server when pruning tool permissions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): mock the MCP toolsets hook in the key update suite Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../view_users/user_info_view.test.tsx | 48 ++++- .../_components/view_users/user_info_view.tsx | 75 +------ .../mcp_server_management/mcpEntitlement.ts | 91 ++++++++ .../KeyInfoView.handleKeyUpdate.test.tsx | 9 + .../templates/key_info_view.test.tsx | 196 ++++++++++++++++++ .../components/templates/key_info_view.tsx | 50 ++--- 6 files changed, 367 insertions(+), 102 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx index c704da301ad..237e3b2c842 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx @@ -1,7 +1,8 @@ import { render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; -import UserInfoView, { extractMcpEntitlement } from "./user_info_view"; +import UserInfoView from "./user_info_view"; +import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; const mockTeamMemberAddCall = vi.fn(); const mockTeamMemberDeleteCall = vi.fn(); @@ -371,7 +372,11 @@ describe("extractMcpEntitlement", () => { const CATALOG = [ { server_id: "srv-1", server_name: "deploy_tracker", alias: "deploy" }, { server_id: "srv-2", server_name: "issue_tracker", alias: null }, - { server_id: "srv-via-group", server_name: "audit_log", alias: null }, + { server_id: "srv-via-group", server_name: "audit_log", alias: null, mcp_access_groups: ["ops_readonly"] }, + ] as any; + + const TOOLSETS = [ + { toolset_id: "ts-1", toolset_name: "audit", tools: [{ server_id: "srv-via-group", tool_name: "read" }] }, ] as any; const form = ( @@ -465,12 +470,43 @@ describe("extractMcpEntitlement", () => { }); it("keeps the allowlist of a deselected server that a retained access group still supplies", () => { - const result = extractMcpEntitlement(form({ accessGroups: ["ops_readonly"] }, { "srv-1": ["read"] }), CATALOG); - expect(result?.mcp_tool_permissions).toEqual({ "srv-1": ["read"] }); + const result = extractMcpEntitlement( + form({ accessGroups: ["ops_readonly"] }, { "srv-via-group": ["read"] }), + CATALOG, + TOOLSETS, + ); + expect(result?.mcp_tool_permissions).toEqual({ "srv-via-group": ["read"] }); }); - it("keeps the allowlist of a deselected server when a toolset is retained", () => { - const result = extractMcpEntitlement(form({ toolsets: ["ts-1"] }, { "srv-1": ["read"] }), CATALOG); + it("drops the allowlist of a deselected server that the retained access group does not contain", () => { + // The gateway grants each allowlist key as its own server, so retaining a group covering only + // srv-via-group must not keep srv-1 callable after the admin removed it. + const result = extractMcpEntitlement( + form({ accessGroups: ["ops_readonly"] }, { "srv-1": ["read"] }), + CATALOG, + TOOLSETS, + ); + expect(result?.mcp_tool_permissions).toEqual({}); + }); + + it("keeps the allowlist of a deselected server that a retained toolset still supplies", () => { + const result = extractMcpEntitlement( + form({ toolsets: ["ts-1"] }, { "srv-via-group": ["read"] }), + CATALOG, + TOOLSETS, + ); + expect(result?.mcp_tool_permissions).toEqual({ "srv-via-group": ["read"] }); + }); + + it("drops the allowlist of a deselected server that the retained toolset does not cover", () => { + const result = extractMcpEntitlement(form({ toolsets: ["ts-1"] }, { "srv-1": ["read"] }), CATALOG, TOOLSETS); + expect(result?.mcp_tool_permissions).toEqual({}); + }); + + it("prunes nothing when a selected toolset is missing from the toolset catalog", () => { + // An unresolvable toolset could supply any server, so pruning against it would be a guess in + // the widening direction. + const result = extractMcpEntitlement(form({ toolsets: ["ts-unknown"] }, { "srv-1": ["read"] }), CATALOG, TOOLSETS); expect(result?.mcp_tool_permissions).toEqual({ "srv-1": ["read"] }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index 7c7b6b51ef4..34b34066f14 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -42,77 +42,9 @@ import NotificationsManager from "@/components/molecules/notifications_manager"; import { getBudgetDurationLabel } from "@/components/common_components/budget_duration_dropdown"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; import MCPServerPermissions from "@/components/permissions/MCPServerPermissions"; -import { MCPServer } from "@/components/mcp_tools/types"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; - -interface McpEntitlementUpdate { - mcp_servers: string[]; - mcp_access_groups: string[]; - mcp_toolsets: string[]; - mcp_tool_permissions: Record; -} - -const asStringArray = (value: unknown): string[] => - Array.isArray(value) ? value.filter((entry): entry is string => typeof entry === "string") : []; - -const asToolPermissions = (value: unknown): Record => { - if (value === null || typeof value !== "object" || Array.isArray(value)) return {}; - return Object.fromEntries( - Object.entries(value as Record).map(([serverId, tools]) => [serverId, asStringArray(tools)]), - ); -}; - -const mcpServerMatchesIdentifier = (server: MCPServer, identifier: string): boolean => - server.server_id === identifier || server.server_name === identifier || server.alias === identifier; - -/** - * The `object_permission` a save sends, derived from what the editor currently shows. - * - * A tool allowlist is what narrows a grant and an absent one reads as no restriction, so dropping - * an entry is the direction that widens. An entry is kept when an access group or toolset the admin - * retained could still supply its server, and dropped once nothing indirect survives, which is what - * makes removing a grant actually remove it. - * - * A tool-permission key may be a server id, a name or an alias: the gateway normalizes all three - * before looking up the allowlist, so an entry written by the API or by config can use any of them. - * `allServers` is what resolves a key to its servers, plural: names and aliases are not unique, and - * the gateway unions such a key into EVERY server answering to it, so the entry is kept while any - * one of them is still granted. Resolving to the first match instead would make the outcome depend - * on catalog order and could drop a restriction that was also covering a server still granted. A key - * that resolves to nothing is kept too, since a server we cannot identify is one we cannot confirm - * was deselected; that also covers a catalog that has not loaded or failed to load, where every key - * is unresolvable and nothing is pruned. - */ -export const extractMcpEntitlement = ( - formValues: Record, - allServers: MCPServer[], -): McpEntitlementUpdate | null => { - const selection = formValues.mcp_servers_and_groups; - if (selection === null || typeof selection !== "object") return null; - - const { servers, accessGroups, toolsets } = selection as Record; - const mcpServers = asStringArray(servers); - const mcpAccessGroups = asStringArray(accessGroups); - const mcpToolsets = asStringArray(toolsets); - const retainsIndirectGrant = mcpAccessGroups.length > 0 || mcpToolsets.length > 0; - - const grantsServerNamedBy = (permissionKey: string): boolean => { - const named = allServers.filter((candidate) => mcpServerMatchesIdentifier(candidate, permissionKey)); - if (named.length === 0) return true; - return named.some((server) => mcpServers.some((identifier) => mcpServerMatchesIdentifier(server, identifier))); - }; - - return { - mcp_servers: mcpServers, - mcp_access_groups: mcpAccessGroups, - mcp_toolsets: mcpToolsets, - mcp_tool_permissions: Object.fromEntries( - Object.entries(asToolPermissions(formValues.mcp_tool_permissions)).filter( - ([permissionKey]) => retainsIndirectGrant || grantsServerNamedBy(permissionKey), - ), - ), - }; -}; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; interface UserInfoViewProps { userId: string; @@ -164,6 +96,7 @@ export default function UserInfoView({ const [selectedRole, setSelectedRole] = useState("user"); const [isLoadingTeams, setIsLoadingTeams] = useState(false); const { data: allMcpServers = [] } = useMCPServers(); + const { data: allMcpToolsets = [] } = useMCPToolsets(); React.useEffect(() => { setBaseUrl(getProxyBaseUrl()); @@ -365,7 +298,7 @@ export default function UserInfoView({ try { if (!accessToken || !userData) return; - const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers); + const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers, allMcpToolsets); const userFields = Object.fromEntries( Object.entries(formValues).filter( ([field]) => field !== "mcp_servers_and_groups" && field !== "mcp_tool_permissions", diff --git a/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts new file mode 100644 index 00000000000..66f59f5b71d --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts @@ -0,0 +1,91 @@ +import { ALL_PROXY_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; + +export interface McpEntitlementUpdate { + mcp_servers: string[]; + mcp_access_groups: string[]; + mcp_toolsets: string[]; + mcp_tool_permissions: Record; +} + +const asStringArray = (value: unknown): string[] => + Array.isArray(value) ? value.filter((entry): entry is string => typeof entry === "string") : []; + +const asToolPermissions = (value: unknown): Record => { + if (value === null || typeof value !== "object" || Array.isArray(value)) return {}; + return Object.fromEntries( + Object.entries(value as Record).map(([serverId, tools]) => [serverId, asStringArray(tools)]), + ); +}; + +const mcpServerMatchesIdentifier = (server: MCPServer, identifier: string): boolean => + server.server_id === identifier || server.server_name === identifier || server.alias === identifier; + +/** + * The `object_permission` a save sends, derived from what the editor currently shows. + * + * A tool allowlist is what narrows a grant and an absent one reads as no restriction, so dropping + * an entry is the direction that widens. An entry is kept while its own server is still reachable, + * directly or through a retained access group or toolset, and dropped once nothing reaches it, which + * is what makes removing a grant actually remove it. Reachability is resolved per server rather than + * per selection: the gateway treats every allowlist key as an independent server grant, so keeping + * every key because some unrelated group survived would leave a deselected server callable. + * + * The catalog carries `mcp_access_groups` on each server and `tools[].server_id` on each toolset, + * which is the same membership the gateway resolves against. A selected toolset missing from the + * catalog is unresolvable, so nothing is pruned in that save. + * + * A tool-permission key may be a server id, a name or an alias: the gateway normalizes all three + * before looking up the allowlist, so an entry written by the API or by config can use any of them. + * `allServers` is what resolves a key to its servers, plural: names and aliases are not unique, and + * the gateway unions such a key into EVERY server answering to it, so the entry is kept while any + * one of them is still granted. Resolving to the first match instead would make the outcome depend + * on catalog order and could drop a restriction that was also covering a server still granted. A key + * that resolves to nothing is kept too, since a server we cannot identify is one we cannot confirm + * was deselected; that also covers a catalog that has not loaded or failed to load, where every key + * is unresolvable and nothing is pruned. + */ +export const extractMcpEntitlement = ( + formValues: Record, + allServers: MCPServer[], + allToolsets: MCPToolset[] = [], +): McpEntitlementUpdate | null => { + const selection = formValues.mcp_servers_and_groups; + if (selection === null || typeof selection !== "object") return null; + + const { servers, accessGroups, toolsets } = selection as Record; + const mcpServers = asStringArray(servers); + const mcpAccessGroups = asStringArray(accessGroups); + const mcpToolsets = asStringArray(toolsets); + const grantsEveryServer = + mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL) || + mcpToolsets.some((toolsetId) => !allToolsets.some((toolset) => toolset.toolset_id === toolsetId)); + + const toolsetServerIds = new Set( + allToolsets + .filter((toolset) => mcpToolsets.includes(toolset.toolset_id)) + .flatMap((toolset) => toolset.tools.map((tool) => tool.server_id)), + ); + + const grants = (server: MCPServer): boolean => + mcpServers.some((identifier) => mcpServerMatchesIdentifier(server, identifier)) || + (server.mcp_access_groups ?? []).some((group) => mcpAccessGroups.includes(group)) || + toolsetServerIds.has(server.server_id); + + const grantsServerNamedBy = (permissionKey: string): boolean => { + const named = allServers.filter((candidate) => mcpServerMatchesIdentifier(candidate, permissionKey)); + if (named.length === 0) return true; + return named.some(grants); + }; + + return { + mcp_servers: mcpServers, + mcp_access_groups: mcpAccessGroups, + mcp_toolsets: mcpToolsets, + mcp_tool_permissions: Object.fromEntries( + Object.entries(asToolPermissions(formValues.mcp_tool_permissions)).filter( + ([permissionKey]) => grantsEveryServer || grantsServerNamedBy(permissionKey), + ), + ), + }; +}; diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 16b1be0d74a..374e36029a0 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -255,6 +255,15 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: vi.fn().mockReturnValue({ data: { values: {} }, isLoading: false }), })); +// Mock useMCPServers hook (requires QueryClientProvider which is not available in this test) +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: vi.fn().mockReturnValue({ data: [] }), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: vi.fn().mockReturnValue({ data: [] }), +})); + // Mock useResetKeySpend hook (requires QueryClientProvider which is not available in this test) vi.mock("@/app/(dashboard)/hooks/keys/useResetKeySpend", () => ({ useResetKeySpend: vi.fn().mockReturnValue({ diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index d2390a633d5..1437689958c 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -32,6 +32,26 @@ vi.mock("@/app/(dashboard)/hooks/projects/useProjects", () => ({ useProjects: vi.fn().mockReturnValue({ data: [], isLoading: false }), })); +const MCP_CATALOG = [ + { server_id: "srv-1", server_name: "deploy_tracker", alias: "deploy" }, + { server_id: "srv-2", server_name: "incident_log", alias: "incidents", mcp_access_groups: ["ops_readonly"] }, +]; + +const MCP_TOOLSETS = [ + { toolset_id: "ts-1", toolset_name: "incidents", tools: [{ server_id: "srv-2", tool_name: "write" }] }, +]; + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: vi.fn(() => ({ data: MCP_CATALOG })), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: vi.fn(() => ({ data: MCP_TOOLSETS })), +})); + +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; + vi.mock("../networking", () => ({ keyDeleteCall: vi.fn().mockResolvedValue({}), keyUpdateCall: vi.fn().mockResolvedValue({}), @@ -881,6 +901,182 @@ describe("KeyInfoView", () => { }); }); + describe("MCP tool permissions on save", () => { + const KEY_WITH_TOOL_PERMISSIONS: KeyResponse = { + ...MOCK_KEY_DATA, + user_id: "proxy-admin-user", + object_permission: { + ...MOCK_KEY_DATA.object_permission, + mcp_servers: ["srv-1", "srv-2"], + mcp_access_groups: [], + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }, + } as KeyResponse; + + const enterEditMode = async (keyData: KeyResponse) => { + vi.mocked(useAuthorized).mockReturnValue({ + ...baseUseAuthorizedMock, + userId: "proxy-admin-user", + userRole: "proxy_admin", + }); + renderWithProviders( + {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />, + ); + await userEvent.click(screen.getByRole("tab", { name: /settings/i })); + await userEvent.click(screen.getByRole("button", { name: /edit settings/i })); + await waitFor(() => expect(editViewMocks.onSubmit).toBeDefined()); + }; + + const submittedToolPermissions = () => { + const payload = vi.mocked(keyUpdateCall).mock.calls.at(-1)?.[1] as Record; + return payload.object_permission.mcp_tool_permissions; + }; + + beforeEach(() => { + editViewMocks.onSubmit = undefined; + vi.mocked(keyUpdateCall).mockClear(); + vi.mocked(keyUpdateCall).mockResolvedValue({}); + vi.mocked(useMCPServers).mockReturnValue({ data: MCP_CATALOG } as unknown as ReturnType); + vi.mocked(useMCPToolsets).mockReturnValue({ data: MCP_TOOLSETS } as unknown as ReturnType); + }); + + it("drops the allowlist of every deselected server instead of leaving it entitled", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({}); + }); + + it("drops only the deselected server and keeps the one still granted", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["srv-1"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-1": ["read"] }); + }); + + it("keeps an allowlist whose server is reachable through a retained access group", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: ["ops_readonly"], toolsets: [] }, + mcp_tool_permissions: { "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-2": ["write"] }); + }); + + it("drops an allowlist the retained access group does not reach", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: ["ops_readonly"], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-2": ["write"] }); + }); + + it("drops an allowlist the retained toolset does not cover", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["ts-1"] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-2": ["write"] }); + }); + + it("refuses to save a permission change while a selected toolset is unresolvable", async () => { + vi.mocked(useMCPToolsets).mockReturnValue({ data: undefined } as unknown as ReturnType); + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["ts-1"] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(keyUpdateCall).not.toHaveBeenCalled(); + }); + + it("resolves a name-keyed allowlist against the server catalog", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["srv-1"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { deploy_tracker: ["read"], incident_log: ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ deploy_tracker: ["read"] }); + }); + + it("clears every allowlist when the admin picks the no-MCP-servers sentinel", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["no-mcp-servers"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({}); + }); + + it("keeps every allowlist when the admin grants all proxy servers", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: ["all-proxy-mcpservers"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(submittedToolPermissions()).toEqual({ "srv-1": ["read"], "srv-2": ["write"] }); + }); + + it("refuses to save a permission change it cannot compute without the server catalog", async () => { + vi.mocked(useMCPServers).mockReturnValue({ data: undefined } as ReturnType); + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"], "srv-2": ["write"] }, + }); + + expect(keyUpdateCall).not.toHaveBeenCalled(); + }); + + it("preserves a vector-store edit made in the same save", async () => { + await enterEditMode(KEY_WITH_TOOL_PERMISSIONS); + await editViewMocks.onSubmit!({ + key: KEY_WITH_TOOL_PERMISSIONS.token, + token: KEY_WITH_TOOL_PERMISSIONS.token, + vector_stores: ["vs-1"], + mcp_servers_and_groups: { servers: ["srv-1"], accessGroups: [], toolsets: [] }, + mcp_tool_permissions: { "srv-1": ["read"] }, + }); + + const payload = vi.mocked(keyUpdateCall).mock.calls.at(-1)?.[1] as Record; + expect(payload.object_permission.vector_stores).toEqual(["vs-1"]); + }); + }); + describe("delete flow", () => { it("invalidates the keys list query after a successful delete so active filters survive (LIT-4080)", async () => { const invalidateSpy = vi.spyOn(QueryClient.prototype, "invalidateQueries"); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index f7720874968..15d14d5abf5 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -24,6 +24,9 @@ import { useResetKeySpend } from "@/app/(dashboard)/hooks/keys/useResetKeySpend" import { useSetKeyBlockedState } from "@/app/(dashboard)/hooks/keys/useSetKeyBlockedState"; import { keyKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useQueryClient } from "@tanstack/react-query"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; +import { extractMcpEntitlement } from "../mcp_server_management/mcpEntitlement"; import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; @@ -74,6 +77,8 @@ export default function KeyInfoView({ const { teams: teamsData } = useTeams(); const { data: projects } = useProjects(); const { data: uiSettingsData } = useUISettings(); + const { data: allMcpServers } = useMCPServers(); + const { data: allMcpToolsets } = useMCPToolsets(); const enableProjectsUI = Boolean(uiSettingsData?.values?.enable_projects_ui); const [isEditing, setIsEditing] = useState(false); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); @@ -201,33 +206,28 @@ export default function KeyInfoView({ delete formValues.vector_stores; } - if (formValues.mcp_servers_and_groups !== undefined) { - const { servers, accessGroups, toolsets } = formValues.mcp_servers_and_groups || { - servers: [], - accessGroups: [], - toolsets: [], - }; - formValues.object_permission = { - ...currentKeyData.object_permission, - mcp_servers: servers || [], - mcp_access_groups: accessGroups || [], - mcp_toolsets: toolsets || [], - }; - // Remove mcp_servers_and_groups from the top level as it should be in object_permission - delete formValues.mcp_servers_and_groups; - } - - // Handle MCP tool permissions - if (formValues.mcp_tool_permissions !== undefined) { - const mcpToolPermissions = formValues.mcp_tool_permissions || {}; - if (Object.keys(mcpToolPermissions).length > 0) { - formValues.object_permission = { - ...formValues.object_permission, - mcp_tool_permissions: mcpToolPermissions, - }; + const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers ?? [], allMcpToolsets ?? []); + if (mcpEntitlement) { + // Without a catalog the grants an allowlist key still has are unresolvable, so nothing is + // pruned and a revocation would save as a no-op while reporting success. Refuse instead. + const unresolvableSelection = + allMcpServers === undefined || + mcpEntitlement.mcp_toolsets.some( + (toolsetId) => !(allMcpToolsets ?? []).some((toolset) => toolset.toolset_id === toolsetId), + ); + if (unresolvableSelection && Object.keys(mcpEntitlement.mcp_tool_permissions).length > 0) { + NotificationManager.error( + "MCP server or toolset list is unavailable, so MCP permissions cannot be saved yet. Retry.", + ); + return; } - delete formValues.mcp_tool_permissions; + formValues.object_permission = { + ...(formValues.object_permission ?? currentKeyData.object_permission), + ...mcpEntitlement, + }; } + delete formValues.mcp_servers_and_groups; + delete formValues.mcp_tool_permissions; // Handle agent permissions if (formValues.agents_and_groups !== undefined) { From 3615cccfef56fde230ef59297c9a7a5ae097f8d1 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 13 Aug 2026 18:01:38 -0700 Subject: [PATCH 2/3] fix(team): sweep dangling team references and cache on team delete (#36819) * fix(team): sweep dangling team references and cache on team delete delete_team drove all of its cleanup off the team's members_with_roles roster, so any user row referencing the team by another route kept a dangling team id forever and the deleted team stayed visible on /user/info. Nothing swept LiteLLM_UserTable.teams or LiteLLM_TeamMembership by team id, schema.prisma declares no relation between the membership table and the team table so there is no cascade to fall back on, and the cached team object was never invalidated on delete. Adds a sweep that runs before the team rows are dropped: it strips the deleted ids from every user row that still lists them and removes every membership row for those teams. Adds _delete_cache_team_object in auth_checks and calls it per deleted team so the team_id:{team_id} entry cannot outlive the team. The sweep is targeted, not indiscriminate: only the deleted ids are removed and the other teams on a user record are left intact. * fix(team): fail member_add when the team is deleted under the row lock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(team): correct the post-delete sweep note for the member_add lock path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 77 ++++ .../management_endpoints/team_endpoints.py | 140 +++++- litellm/repositories/team_repository.py | 10 +- .../test_team_endpoints.py | 421 +++++++++++++++++- .../repositories/test_repositories.py | 9 +- 5 files changed, 647 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4d9f0a58bd7..0239368ad0e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -65,6 +65,7 @@ from litellm.proxy.auth.budget_throttle import ( should_throttle_budget_exceeded, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, @@ -2016,6 +2017,44 @@ async def _cache_team_object( ) +async def delete_cache_team_object( + team_id: str, + team_alias: str | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> None: + """ + Evict both keys `_cache_team_object` writes. + + `get_team_object` reads the id key and the JWT `team_alias_jwt_field` path reads the alias key, + so leaving either behind keeps a deleted team resolvable for auth until its TTL expires. + + Mirrors `delete_cached_project_object`: evicting locally only reaches the worker handling the + delete, so every key is also broadcast to drop the other workers' in-memory copies. + + Eviction is best-effort, matching `_cache_team_object`. `delete_team` calls this after the team + rows are already gone, so letting an unreachable cache backend raise here would fail a request + whose delete has committed. + """ + keys: Final = (f"team_id:{team_id}", *((f"team_alias:{team_alias}",) if team_alias else ())) + + for key in keys: + try: + user_api_key_cache.delete_cache(key=key) + + ## UPDATE REDIS CACHE ## + if proxy_logging_obj is not None: + await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key) + except Exception as e: # noqa: BLE001 # best-effort invalidation: any cache backend error must not abort the delete + verbose_proxy_logger.warning( + "Failed to invalidate cached team entry %s on delete; " + "a deleted team may be served until its TTL expires: %s", + key, + e, + ) + await publish_auth_cache_invalidation(cache_key=key) + + async def _cache_key_object( hashed_token: str, user_api_key_obj: UserAPIKeyAuth, @@ -2068,6 +2107,44 @@ class TeamNotFoundError(HTTPException): ) +async def delete_cache_key_objects( + hashed_tokens: Sequence[str], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging | None, +) -> None: + """ + Evict a batch of key objects, for callers that delete keys in bulk rather than through + `/key/delete`. Auth resolves a cached key object without re-reading its team, so a key left + cached after its row is gone keeps buying access until its TTL expires. + + Evicting locally only reaches this worker, so each token is also broadcast: a deleted key left + in a peer worker's in-memory cache still authenticates there until its TTL expires. + + Best-effort per key: the rows are already deleted by the time this runs, so an unreachable + cache backend must not abort the caller partway through its own cascade. + """ + results: Final = await asyncio.gather( + *( + _delete_cache_key_object( + hashed_token=hashed_token, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for hashed_token in hashed_tokens + ), + return_exceptions=True, + ) + + for hashed_token, result in zip(hashed_tokens, results): + if isinstance(result, BaseException): + verbose_proxy_logger.warning( + "Failed to evict cached key entry for %s; a deleted key may authenticate until its TTL expires: %s", + hashed_token, + result, + ) + await publish_auth_cache_invalidation(cache_key=hashed_token) + + @log_db_metrics async def _get_team_db_check( team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 452828edc9e..b72b5d218b2 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -77,6 +77,8 @@ from litellm.proxy.auth.auth_checks import ( _cache_team_object, allowed_route_check_inside_route, can_org_access_model, + delete_cache_key_objects, + delete_cache_team_object, get_org_object, get_team_membership, get_team_object, @@ -313,6 +315,11 @@ class _TeamIdInFilter(TypedDict, total=False): team_id: Mapping[str, Sequence[str]] +_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """ +UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams) +""" + + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) @@ -2657,6 +2664,11 @@ async def _add_team_members_to_team( serialize on the row lock and each appends onto the other's committed result, instead of both rewriting the whole JSON array from a stale snapshot (which silently drops one member on the losing write). + + The same lock serializes this against /team/delete: the delete cannot remove + the row while the reconcile holds it, and a reconcile that finds the row + already gone cleans up after itself rather than leaving the member pointing + at a deleted team id. """ # Process and add new members updated_users, updated_team_memberships = await _process_team_members( @@ -2667,11 +2679,42 @@ async def _add_team_members_to_team( litellm_proxy_admin_name=litellm_proxy_admin_name, ) - async with prisma_client.tx() as tx: - complete_team_data.members_with_roles = await TeamRepository(prisma_client).get_members_with_roles_locked( - tx, data.team_id + updated_team: Final = await _write_members_with_roles_locked( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + updated_users=updated_users, + ) + if updated_team is None: + await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client) + raise HTTPException( + status_code=404, + detail={"error": f"Team={data.team_id} was deleted while this member add was running"}, ) + return updated_team, updated_users, updated_team_memberships + + +async def _write_members_with_roles_locked( + data: TeamMemberAddRequest, + complete_team_data: LiteLLM_TeamTable, + prisma_client: PrismaClient, + updated_users: list[LiteLLM_UserTable], +) -> LiteLLM_TeamTable | None: + """Reconcile members_with_roles under the team row lock. None when the team row is gone. + + That read is at least as recent as the user and membership writes the caller + already made, so a missing row means /team/delete committed after them. Its + post-delete sweep can have run before those writes landed, which is why the + caller sweeps this team id again rather than only reporting the 404. + """ + async with prisma_client.tx() as tx: + locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id) + if locked_members is None: + return None + + complete_team_data.members_with_roles = locked_members + await _update_team_members_list( data=data, complete_team_data=complete_team_data, @@ -2679,13 +2722,11 @@ async def _add_team_members_to_team( ) _db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles] - updated_team: Final = await tx.litellm_teamtable.update( + return await tx.litellm_teamtable.update( where={"team_id": data.team_id}, data={"members_with_roles": json.dumps(_db_team_members)}, ) - return updated_team, updated_users, updated_team_memberships - def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: """Update the Prometheus team members gauge after a membership change. @@ -3653,6 +3694,8 @@ async def delete_team( create_audit_log_for_update, litellm_proxy_admin_name, prisma_client, + proxy_logging_obj, + user_api_key_cache, ) if prisma_client is None: @@ -3746,6 +3789,12 @@ async def delete_team( await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key") + await _invalidate_deleted_key_cache( + keys=keys_to_delete, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ## DELETE ASSOCIATED BYOK MODELS # Runs before the team rows are deleted so a mid-flight failure never leaves # the team gone with its models orphaned. @@ -3779,11 +3828,90 @@ async def delete_team( ) await asyncio.gather(*tasks) + await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) + ## DELETE TEAMS deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team") + + # Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and + # `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a + # concurrent auth lookup re-caches the still-present team and the delete looks like it never + # invalidated anything. Nothing fallible runs between the delete and this, or a failure there + # would strand the deleted team in cache. + await _invalidate_deleted_team_cache( + teams=team_rows, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + # Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep + # and the delete would have re-appended the reference; an add still in flight sees the row + # missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and + # keeping the first one means a failure here still leaves a team the admin can retry deleting. + await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client) + return deleted_teams +async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client: PrismaClient) -> None: + """ + Strip the deleted team ids from every user row and team-membership row that still references them. + + The per-member `team_member_delete` pass above only reaches users listed in the team's + `members_with_roles`, so a user row that outlived its roster entry is invisible to it and keeps + surfacing the team on `/user/info` after the team is gone. + + #36839 closed the route that created that drift, by resolving member removal off the roster + entry's `user_id` rather than the identifier the caller happened to pass. It does not backfill + rows that already drifted, which is the state this was reported against, so the sweep still has + to run on delete. + + `array_remove` rather than read-filter-write: rewriting the whole array from a snapshot read + outside a transaction drops any team a concurrent `/team/member_add` appended in between. + """ + for team_id in team_ids: + _ = await prisma_client.db.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id) + + _ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)})) + + +async def _invalidate_deleted_key_cache( + keys: Sequence[LiteLLM_VerificationToken], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Evict the auth cache entry for every key deleted along with the team. + + `/key/delete` evicts as it goes, but the bulk delete above writes straight to the db. Auth + resolves a cached key object without re-reading the team, so a key belonging to a deleted team + keeps buying access until its TTL expires. + """ + await delete_cache_key_objects( + hashed_tokens=tuple(key.token for key in keys), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def _invalidate_deleted_team_cache( + teams: Sequence[LiteLLM_TeamTable], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + _ = await asyncio.gather( + *( + delete_cache_team_object( + team_id=team.team_id, + team_alias=team.team_alias, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for team in teams + ) + ) + + def _transform_teams_to_deleted_records( teams: list[LiteLLM_TeamTable], user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/repositories/team_repository.py b/litellm/repositories/team_repository.py index a4908af561a..7efd32288e4 100644 --- a/litellm/repositories/team_repository.py +++ b/litellm/repositories/team_repository.py @@ -57,9 +57,13 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): return LiteLLM_TeamTable.model_validate(data) - async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member]: + async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> list[Member] | None: """Return the team's members_with_roles, locking the row FOR UPDATE. + ``None`` when the team row is gone, which a caller holding the lock can + only see if a delete committed under it, as opposed to ``[]`` for a team + that simply has no members. + Must be called inside a transaction so the row lock is held until commit. This serializes concurrent membership writers on the team row so the losing writer appends onto the winner's committed result instead @@ -69,7 +73,9 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]): 'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = $1 FOR UPDATE', team_id, ) - raw_value: Final = rows[0]["members_with_roles"] if rows else None + if not rows: + return None + raw_value: Final = rows[0]["members_with_roles"] parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value if not parsed: return [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index fe844db9865..f038af4c7a1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -4,7 +4,7 @@ import os import sys from datetime import datetime, timezone from typing import Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from fastapi import HTTPException @@ -39,6 +39,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( from litellm.proxy.management_endpoints.team_endpoints import ( GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, + _STRIP_DELETED_TEAM_FROM_USERS_SQL, _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, @@ -1829,6 +1830,63 @@ async def test_add_team_members_reconciles_against_freshly_locked_row(): assert [m.user_id for m in updated_team.members_with_roles] == ["zed", "alice", "bob"] +@pytest.mark.asyncio +async def test_add_team_members_cleans_up_when_the_team_is_deleted_mid_request(): + """ + Regression pin for the /team/member_add vs /team/delete race. + + The user row and membership writes land before the reconcile takes the team + row lock, so a /team/delete that commits in between has already run its own + reference sweep and cannot see them. The empty locked SELECT is the only + signal that happened, and leaving it at that would strand the member on a + deleted team id, which authorization paths that trust `user.teams` would + treat as membership if the id were ever recreated. So the request must sweep + the references it just wrote and fail, not report success. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + _add_team_members_to_team, + ) + + tx = MagicMock() + tx.query_raw = AsyncMock(return_value=[]) + tx.litellm_teamtable.update = AsyncMock() + + tx_cm = MagicMock() + tx_cm.__aenter__ = AsyncMock(return_value=tx) + tx_cm.__aexit__ = AsyncMock(return_value=None) + + prisma_client = MagicMock() + prisma_client.tx = MagicMock(return_value=tx_cm) + prisma_client.db.execute_raw = AsyncMock() + prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + with patch( + "litellm.proxy.management_endpoints.team_endpoints._process_team_members", + new=AsyncMock(return_value=([], [])), + ): + with pytest.raises(HTTPException) as exc_info: + await _add_team_members_to_team( + data=TeamMemberAddRequest( + team_id="team-deleted-mid-add", + member=Member(user_id="bob", role="user"), + ), + complete_team_data=LiteLLM_TeamTable(team_id="team-deleted-mid-add", members_with_roles=[]), + prisma_client=cast(object, prisma_client), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + litellm_proxy_admin_name="admin", + ) + + assert exc_info.value.status_code == 404 + tx.litellm_teamtable.update.assert_not_awaited() + + assert prisma_client.db.execute_raw.await_args_list == [ + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-deleted-mid-add") + ] + prisma_client.db.litellm_teammembership.delete_many.assert_awaited_once_with( + where={"team_id": {"in": ("team-deleted-mid-add",)}} + ) + + def test_add_new_models_to_team_with_existing_models(): """ Test add_new_models_to_team function with existing models @@ -7238,6 +7296,367 @@ async def test_delete_team_persists_deleted_teams(monkeypatch): assert records[0]["litellm_changed_by"] == "admin-user" +@pytest.mark.asyncio +async def test_delete_team_sweeps_references_outside_members_with_roles(monkeypatch): + """ + Regression pin for LIT-5511: a deleted team stayed visible on user records. + + `delete_team` drove all of its cleanup off `team.members_with_roles`, so a user row that + referenced the team by any other route (`/user/update`, SSO sync, a membership row written + without a matching roster entry) kept the dangling team id forever and `/user/info` kept + listing the deleted team. The roster here is deliberately EMPTY, so nothing the per-member + `team_member_delete` path does can make this test pass. + + Both cache keys `_cache_team_object` writes are asserted in the same delete: the id key feeds + `get_team_object` and the alias key feeds the JWT `team_alias_jwt_field` path, so either one + surviving keeps the deleted team resolvable for auth until its TTL expires. + """ + from litellm.proxy._types import DeleteTeamRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + doomed_team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + cache_state_when_rows_deleted = {} + + async def record_cache_state_then_delete(*args, **kwargs): + if kwargs.get("table_name") == "team": + cache_state_when_rows_deleted["doomed_still_cached"] = ( + fresh_cache.get_cache(key="team_id:team-doomed") is not None + ) + return {"deleted_teams": ["team-doomed"]} + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=doomed_team) + mock_prisma_client.delete_data = AsyncMock(side_effect=record_cache_state_then_delete) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + mock_execute_raw = AsyncMock() + mock_prisma_client.db.execute_raw = mock_execute_raw + mock_membership_delete_many = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = mock_membership_delete_many + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + for cached_team_id, cached_alias in ( + ("team-doomed", "doomed-team"), + ("team-kept", "kept-team"), + ): + cached_obj = LiteLLM_TeamTableCachedObj( + team_id=cached_team_id, team_alias=cached_alias + ) + fresh_cache.set_cache(key=f"team_id:{cached_team_id}", value=cached_obj) + fresh_cache.set_cache(key=f"team_alias:{cached_alias}", value=cached_obj) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # array_remove strips just the deleted id in one statement; a read-filter-write of the whole + # array would drop any team a concurrent /team/member_add appended between read and write + assert "array_remove" in _STRIP_DELETED_TEAM_FROM_USERS_SQL + assert mock_execute_raw.await_args_list == [ + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), + call(_STRIP_DELETED_TEAM_FROM_USERS_SQL, "team-doomed"), + ], "the sweep must run once before the team row is deleted and again after, so a member_add racing the delete cannot leave the reference behind" + + # same two passes: the second one reaps a membership row inserted while the delete was running + assert mock_membership_delete_many.await_args_list == [ + call(where={"team_id": {"in": ("team-doomed",)}}), + call(where={"team_id": {"in": ("team-doomed",)}}), + ] + + assert fresh_cache.get_cache(key="team_id:team-doomed") is None + assert fresh_cache.get_cache(key="team_alias:doomed-team") is None + assert fresh_cache.get_cache(key="team_id:team-kept") is not None + assert fresh_cache.get_cache(key="team_alias:kept-team") is not None + + # Eviction must run AFTER the rows are gone: both writers of these keys hydrate from the db, + # so evicting first lets a concurrent auth lookup re-cache the still-present team. + assert cache_state_when_rows_deleted["doomed_still_cached"] is True + + +@pytest.mark.asyncio +async def test_delete_team_evicts_the_auth_cache_of_the_keys_it_deletes(monkeypatch): + """ + A virtual key scoped to the team is deleted from the db with the team, but auth resolves a + cached key object without re-reading the team, so leaving the cache entry behind lets that key + keep buying access until its TTL expires. Verified live: without this eviction the same key + still returns HTTP 200 on /v1/chat/completions right after /team/delete. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + team_key = LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed") + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[team_key]) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + fresh_cache.set_cache(key="hashed-doomed-key", value=UserAPIKeyAuth(token="hashed-doomed-key", team_id="team-doomed")) + fresh_cache.set_cache(key="hashed-unrelated-key", value=UserAPIKeyAuth(token="hashed-unrelated-key")) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + assert fresh_cache.get_cache(key="hashed-doomed-key") is None + # a key that had nothing to do with the deleted team must survive + assert fresh_cache.get_cache(key="hashed-unrelated-key") is not None + + +@pytest.mark.asyncio +async def test_delete_team_failing_reconcile_sweep_cannot_strand_the_team_in_cache(monkeypatch): + """ + The reconcile sweep runs after the team row is committed deleted. If it ran before cache + eviction, a sweep failure would return an error with the team gone from the db but still + served from cache, which is the exact bug this PR exists to fix. + """ + from litellm.proxy._types import DeleteTeamRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + # the first sweep succeeds, the post-delete reconcile sweep blows up + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[None, ConnectionError("db went away")]) + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + fresh_cache = UserApiKeyCache() + cached_obj = LiteLLM_TeamTableCachedObj(team_id="team-doomed", team_alias="doomed-team") + fresh_cache.set_cache(key="team_id:team-doomed", value=cached_obj) + fresh_cache.set_cache(key="team_alias:doomed-team", value=cached_obj) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", fresh_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + with pytest.raises(ConnectionError): + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # the delete committed, so the cache must not still be serving the team + assert fresh_cache.get_cache(key="team_id:team-doomed") is None + assert fresh_cache.get_cache(key="team_alias:doomed-team") is None + + +@pytest.mark.asyncio +async def test_delete_team_broadcasts_cache_invalidation_to_other_workers(monkeypatch): + """ + Evicting locally only reaches the worker that handled the delete. Without the broadcast, every + other worker keeps serving the deleted team, and the deleted team's keys, out of its own + in-memory cache until the TTL, so both stay usable for auth cluster-wide. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")] + ) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + published = [] + + async def record_publish(cache_key): + published.append(cache_key) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.publish_auth_cache_invalidation", record_publish) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + # the deleted key first, then both keys `_cache_team_object` writes: miss the alias one and the + # JWT-by-alias path keeps resolving the team, miss the token and the key still authenticates + assert published == ["hashed-doomed-key", "team_id:team-doomed", "team_alias:doomed-team"] + + +@pytest.mark.asyncio +async def test_delete_team_survives_a_failing_cache_backend(monkeypatch): + """ + Cache eviction runs after the reference sweep has already committed, so a cache backend that + is unreachable must not abort the delete. If it did, `/team/delete` would fail with the team + row still present but its user references and membership rows already gone. + """ + from litellm.proxy._types import DeleteTeamRequest, LiteLLM_VerificationToken + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + team = LiteLLM_TeamTable( + team_id="team-doomed", + team_alias="doomed-team", + members_with_roles=[], + metadata={}, + model_max_budget={}, + model_spend={}, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team) + mock_delete_data = AsyncMock(return_value={"deleted_teams": ["team-doomed"]}) + mock_prisma_client.delete_data = mock_delete_data + mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock() + mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock() + # a key to evict: its eviction runs after the key rows are already deleted, so it must not + # raise either + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[LiteLLM_VerificationToken(token="hashed-doomed-key", team_id="team-doomed")] + ) + mock_prisma_client.db.execute_raw = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + + mock_tx = AsyncMock() + mock_tx.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + mock_tx_cm = MagicMock() + mock_tx_cm.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_tx_cm) + + exploding_logging_obj = MagicMock() + exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( + side_effect=ConnectionError("redis is down") + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", exploding_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + result = await delete_team( + data=DeleteTeamRequest(team_ids=["team-doomed"]), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_id="admin-user", + api_key="sk-admin", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + ), + litellm_changed_by="admin-user", + ) + + assert result == {"deleted_teams": ["team-doomed"]} + mock_delete_data.assert_any_await(team_id_list=["team-doomed"], table_name="team") + assert exploding_logging_obj.internal_usage_cache.dual_cache.async_delete_cache.await_count > 0 + + @pytest.mark.asyncio async def test_team_member_delete_persists_deleted_keys(monkeypatch): from litellm.proxy._types import TeamMemberDeleteRequest diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index 3f567397e1a..38af52f165c 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -546,12 +546,19 @@ class TestTeamRepository: @pytest.mark.asyncio async def test_get_members_with_roles_locked_missing_row(self, repo): + """None, not [], so a caller can tell a deleted team from an empty one. + + /team/member_add reconciles membership under this lock and has to fail, + and clean up the references it already wrote, when a /team/delete + committed underneath it. An empty list would look like a live team with + no members and it would carry on writing. + """ tx = MagicMock() tx.query_raw = AsyncMock(return_value=[]) members = await repo.get_members_with_roles_locked(tx, "missing") - assert members == [] + assert members is None @pytest.mark.asyncio async def test_create_team_all_fields(self, repo): From 8841cbc10f2c27f2b054a0f1df324b59da90eeed Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 13 Aug 2026 18:03:12 -0700 Subject: [PATCH 3/3] fix(mcp): resolve admin OAuth sessions from any worker via DB-backed drafts (#36844) The Admin UI's Authorize & Fetch Token flow stored its pending server in a module-level dict, so /register, /authorize and /token only succeeded when every leg happened to land on the process that served /session. On a proxy with NUM_WORKERS greater than 1, or more than one replica, each click was an independent draw and failed with a bare 404, which reads as intermittent. Persist the pending server as a short-lived draft row instead, so any worker resolves it. The in-memory cache is kept as the fallback for proxies with no database configured, which keeps single-process deployments working as before. A session runs under a caller-supplied id only when that id names a server that really exists, which is the edit form re-authorizing a saved server. Anything else gets a fresh id, so two concurrent sessions can never share one draft and silently adopt each other's URL or client credentials. Drafts past their lifetime are swept on each write so abandoned sessions do not accumulate, and a lost create race adopts the winner rather than failing a caller whose session is ready. Drafts are excluded from listings and never enter the runtime registry. The exclusion keeps rows whose approval status is NULL, which both short spellings of the filter drop, silently hiding every server predating the approval workflow. Measured on a two-worker proxy against the live GitHub MCP server, 120 concurrent authorize calls per leg: staging 56/120 failures, this branch 0/120, staging again 65/120 as a positive control. --- litellm/proxy/_experimental/mcp_server/db.py | 107 ++++- litellm/proxy/_types.py | 3 + .../mcp_management_endpoints.py | 104 ++++- .../test_mcp_management_endpoints.py | 368 +++++++++++++++++- 4 files changed, 565 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 711119b5ab5..08a8b1bc7b3 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -552,13 +552,20 @@ async def get_all_mcp_servers( ) -> list[LiteLLM_MCPServerTable]: """ Returns mcp servers from the db, optionally filtered by approval_status. - Pass approval_status=None to return all servers regardless of approval state. + Pass approval_status=None to return every server except drafts, which back the admin OAuth + session flow, are addressable only by their own server_id, and must never appear in a listing. + NULL approval_status predates the approval workflow, so those rows are kept explicitly rather + than dropped by a bare inequality, which SQL evaluates as NULL and would silently hide them. """ try: - where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = {} - if approval_status is not None: - where["approval_status"] = approval_status - mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where if where else {}) + where: Final[prisma_db_types.LiteLLM_MCPServerTableWhereInput] = ( + {"approval_status": approval_status} + if approval_status is not None + # mutable-ok: prisma where-inputs must be plain dicts, and both `NOT` and `not` drop + # NULL rows (measured), so the OR is the only NULL-preserving way to exclude drafts + else {"OR": [{"approval_status": None}, {"approval_status": {"not": MCPApprovalStatus.draft}}]} + ) + mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where) tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers] for table in tables: @@ -814,6 +821,96 @@ async def create_mcp_server( return new_mcp_server +async def create_draft_mcp_server( + prisma_client: PrismaClient, + data: NewMCPServerRequest, + touched_by: str, + ttl_seconds: int, + server_id: str | None = None, +) -> LiteLLM_MCPServerTable: + """ + Persist a short-lived draft row backing the admin OAuth "Authorize & Fetch Token" flow. + + The draft lives in the database rather than in process memory so that the /register, + /authorize and /token legs resolve it whichever worker or replica accepts each request. + + Writing is strictly create-if-absent. Any existing row for the id is returned untouched, which + covers both a live draft for this same session and a real server the edit form is + re-authorizing against its own id, where writing a draft would collide on the primary key. + Each click of Authorize mints a fresh id, so nothing is lost by never overwriting, and it is + what makes concurrent callers sharing one id safe rather than mutually destructive. + """ + draft_id: Final = server_id or data.server_id or str(uuid.uuid4()) + await _prune_expired_draft_mcp_servers(prisma_client, ttl_seconds) + + existing: Final = await _db_find_mcp_server_row(prisma_client, draft_id) + if existing is not None: + # Already usable by every worker, whether it is a live draft for this same session or a + # real server the edit form is re-authorizing. Either way there is nothing to write, and + # not writing is what keeps concurrent callers for one server_id from racing each other. + return LiteLLM_MCPServerTable.model_validate(existing.model_dump()) + + draft_payload: Final = data.model_copy(update={"server_id": draft_id, "approval_status": MCPApprovalStatus.draft}) + try: + return await create_mcp_server(prisma_client, draft_payload, touched_by) + except Exception: + # Lost the create race: the read above and this create are two statements, not one. The + # winner wrote a draft for this same session, so adopt it rather than failing a caller + # whose session is in fact ready. Anything else still raises. + raced: Final = await _db_find_mcp_server_row(prisma_client, draft_id) + if raced is None or raced.approval_status != MCPApprovalStatus.draft: + raise + return LiteLLM_MCPServerTable.model_validate(raced.model_dump()) + + +async def _prune_expired_draft_mcp_servers(prisma_client: PrismaClient, ttl_seconds: int) -> None: + """Drop drafts already past ``ttl_seconds``, so abandoned OAuth sessions do not accumulate. + + Runs on each draft write rather than on a schedule, mirroring the in-memory cache this + replaces, which pruned on every store. Expired drafts are unreadable by then anyway, so the + only thing at stake is row count, and the work is bounded by how often admins authorize. + """ + cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds)) + # Age is filtered here rather than in the query: the draft set is bounded by how many OAuth + # authorizations are in flight, so it is a handful of rows even on a busy proxy. + drafts: Final = await _db_find_mcp_server_rows( + prisma_client, + where={"approval_status": MCPApprovalStatus.draft}, + ) + for row in drafts: + # A row without a timestamp has no age to judge, so leave it rather than guess it is stale. + # Two workers sweeping the same row is harmless: prisma's delete returns None for a row + # that is already gone rather than raising, so the loser of that race is a no-op. + if row.updated_at is not None and row.updated_at < cutoff: + await delete_mcp_server(prisma_client, row.server_id) + + +async def get_draft_mcp_server( + prisma_client: PrismaClient, server_id: str, ttl_seconds: int +) -> LiteLLM_MCPServerTable | None: + """ + Return the draft row for ``server_id`` if it has not yet aged past ``ttl_seconds``, else None. + + Age is enforced in the query rather than by a sweeper so an expired draft is unreadable the + moment it lapses, regardless of which process last ran a cleanup. + """ + cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=max(1, ttl_seconds)) + draft_rows: Final = await _db_find_mcp_server_rows( + prisma_client, + where={ + "server_id": server_id, + "approval_status": MCPApprovalStatus.draft, + "updated_at": {"gte": cutoff}, + }, + ) + if not draft_rows: + return None + + table: Final = LiteLLM_MCPServerTable.model_validate(draft_rows[0].model_dump()) + decrypt_global_env_var_values(table.env_vars) + return table + + async def update_mcp_server( prisma_client: PrismaClient, data: UpdateMCPServerRequest, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bb330d00756..6c01575e271 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1283,6 +1283,9 @@ class MCPApprovalStatus(str, enum.Enum): pending_review = "pending_review" active = "active" rejected = "rejected" + # Short-lived row backing the admin OAuth "Authorize & Fetch Token" flow. Never served: the + # registry loader and every listing exclude it, so it is reachable only by its own server_id. + draft = "draft" from litellm.models.mcp_server import ( # noqa: E402 diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c4f5ceaba7a..997012dbc65 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -22,7 +22,7 @@ import os from collections.abc import Iterable from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal from fastapi import ( APIRouter, @@ -92,6 +92,9 @@ def does_mcp_server_exist(mcp_server_records: Iterable[Any], mcp_server_id: str) DEFAULT_MCP_REGISTRY_VERSION: Final = "1.0.0" +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + try: importlib.import_module("mcp") except ImportError as e: @@ -115,11 +118,13 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.db import ( approve_mcp_server, + create_draft_mcp_server, create_mcp_server, delete_mcp_server, delete_user_credential, delete_user_env_vars, get_all_mcp_servers_for_user, + get_draft_mcp_server, get_mcp_server, get_mcp_servers, get_mcp_submissions, @@ -467,19 +472,68 @@ if MCP_AVAILABLE: verbose_proxy_logger.debug("Invalid temporary MCP server payload in Redis cache: %s", e) return None + def _get_prisma_client_or_none() -> "PrismaClient | None": + """Non-throwing counterpart to ``get_prisma_client_or_throw`` for paths that degrade + gracefully: a proxy configured without a database keeps the in-memory OAuth session.""" + from litellm.proxy.proxy_server import prisma_client + + return prisma_client + + async def _persist_draft_mcp_server( + payload: NewMCPServerRequest, + server_id: str, + created_by: str, + ) -> None: + """Write the draft row that makes the OAuth session resolvable from any worker. + + A failure here is raised, not swallowed: without the shared row the flow degrades to + the per-process cache and fails intermittently, which is the defect being fixed. + """ + prisma_client: Final = _get_prisma_client_or_none() + if prisma_client is None: + return + await create_draft_mcp_server( + prisma_client, + payload, + created_by, + ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, + server_id=server_id, + ) + + async def _get_draft_mcp_server_as_mcp_server(server_id: str) -> MCPServer | None: + """Resolve a database-backed draft, which is the only lookup that works across workers.""" + prisma_client: Final = _get_prisma_client_or_none() + if prisma_client is None: + return None + draft: Final = await get_draft_mcp_server( + prisma_client, server_id, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS + ) + if draft is None: + return None + return await global_mcp_server_manager.build_mcp_server_from_table(draft) + async def get_cached_temporary_mcp_server( server_id: str, ) -> MCPServer | None: _prune_expired_temporary_mcp_servers() entry: Final = _temporary_mcp_servers.get(server_id) - if entry is None: - redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id) - if redis_server is None: - return None - # Intentionally avoid repopulating local cache from Redis to prevent - # extending effective lifetime beyond the remaining Redis TTL. - return redis_server - return entry.server + if entry is not None: + return entry.server + + # A miss here means either an expired session or, on a multi-worker or multi-replica + # proxy, that a different process served /session. The draft row is shared, so it + # resolves the second case; the in-memory hit above still serves single-process + # deployments with no database configured. + draft_server: Final = await _get_draft_mcp_server_as_mcp_server(server_id) + if draft_server is not None: + return draft_server + + redis_server: Final = await _get_temporary_mcp_server_from_redis(server_id) + if redis_server is None: + return None + # Intentionally avoid repopulating local cache from Redis to prevent + # extending effective lifetime beyond the remaining Redis TTL. + return redis_server def _redact_mcp_credentials( mcp_server: LiteLLM_MCPServerTable, @@ -709,12 +763,36 @@ if MCP_AVAILABLE: payload_dict["credentials"] = inherited_credentials return NewMCPServerRequest.model_validate(payload_dict) + async def _resolve_session_server_id(payload: NewMCPServerRequest) -> str: + """Decide the id an OAuth session runs under. + + A caller-supplied id is honoured only when it names a server that really exists, which is + the edit form re-authorizing a saved server against its own id. Anything else gets a fresh + id, so two concurrent sessions can never land on one id and silently adopt each other's + URL or client credentials. Without a database there is nothing shared to collide over, so + the supplied id is kept and behaviour is unchanged. + """ + supplied: Final = payload.server_id + if not supplied: + return str(uuid.uuid4()) + if global_mcp_server_manager.get_mcp_server_by_id(supplied) is not None: + return supplied + prisma_client: Final = _get_prisma_client_or_none() + if prisma_client is None: + return supplied + # A draft is another session's row, not a saved server, so re-supplying an id this + # endpoint previously handed back must not let a later session adopt its configuration. + existing: Final = await get_mcp_server(prisma_client, supplied) + if existing is None or existing.approval_status == MCPApprovalStatus.draft: + return str(uuid.uuid4()) + return supplied + def _build_temporary_mcp_server_record( payload: NewMCPServerRequest, created_by: str | None, + server_id: str, ) -> LiteLLM_MCPServerTable: now: Final = datetime.utcnow() - server_id: Final = payload.server_id or str(uuid.uuid4()) server_name: Final = payload.server_name or payload.alias or server_id return LiteLLM_MCPServerTable( server_id=server_id, @@ -1544,6 +1622,7 @@ if MCP_AVAILABLE: temp_record: Final = _build_temporary_mcp_server_record( payload_with_credentials, created_by, + await _resolve_session_server_id(payload_with_credentials), ) try: @@ -1555,6 +1634,11 @@ if MCP_AVAILABLE: temporary_server, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, ) + await _persist_draft_mcp_server( + payload_with_credentials, + temp_record.server_id, + created_by, + ) await _cache_temporary_mcp_server_in_redis( temporary_server, ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS, diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index bf119c4fb2f..b6c2fd4c8e3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1697,6 +1697,362 @@ class TestTemporaryMCPSessionEndpoints: assert result is None assert "expired" not in cache + @pytest.mark.asyncio + async def test_get_cached_temporary_mcp_server_resolves_draft_written_by_another_worker(self): + """Regression: the OAuth session must resolve on a worker that did not serve /session. + + `_temporary_mcp_servers` is per-process, so on a multi-worker or multi-replica proxy the + /authorize and /token legs land on a process whose dict is empty and 404. An empty dict + here IS that other worker. Before the DB-backed draft this returned None. + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + get_cached_temporary_mcp_server, + ) + + draft_row = generate_mock_mcp_server_db_record(server_id="drafted-elsewhere") + rebuilt_server = generate_mock_mcp_server_config_record(server_id="drafted-elsewhere") + mock_manager = MagicMock() + mock_manager.build_mcp_server_from_table = AsyncMock(return_value=rebuilt_server) + get_draft = AsyncMock(return_value=draft_row) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers", + {}, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_draft_mcp_server", + get_draft, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + ): + result = await get_cached_temporary_mcp_server("drafted-elsewhere") + + assert result is rebuilt_server + # The shared row, not the empty per-process dict, is what answered. + assert get_draft.await_count == 1 + assert get_draft.await_args.args[1] == "drafted-elsewhere" + + @pytest.mark.asyncio + async def test_get_cached_temporary_mcp_server_still_works_without_a_database(self): + """A proxy configured with no database keeps the in-memory session, rather than 404ing. + + Pins the deliberate divergence from a DB-only design: single-process deployments with no + DATABASE_URL must keep working exactly as before. + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _TemporaryMCPServerEntry, + get_cached_temporary_mcp_server, + ) + + server = generate_mock_mcp_server_config_record(server_id="no-db") + entry = _TemporaryMCPServerEntry( + server=server, + expires_at=datetime.utcnow() + timedelta(seconds=300), + ) + get_draft = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers", + {"no-db": entry}, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=None, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_draft_mcp_server", + get_draft, + ), + ): + result = await get_cached_temporary_mcp_server("no-db") + + assert result is server + # No database means no draft lookup is even attempted. + assert get_draft.await_count == 0 + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_never_overwrites_a_real_server(self): + """The edit form authorizes against a saved server's own id, so a draft write would + collide on the primary key. That row is already visible to every worker, so it is + returned untouched and no draft is created.""" + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + real_row = generate_mock_mcp_server_db_record(server_id="already-saved") + real_row.approval_status = "active" + create_call = AsyncMock() + delete_call = AsyncMock() + + with ( + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", + AsyncMock(return_value=real_row), + ), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + AsyncMock(return_value=[]), + ), + patch("litellm.proxy._experimental.mcp_server.db.create_mcp_server", create_call), + patch("litellm.proxy._experimental.mcp_server.db.delete_mcp_server", delete_call), + ): + result = await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="already-saved", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + assert result.server_id == "already-saved" + assert create_call.await_count == 0 + assert delete_call.await_count == 0 + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_adopts_the_winner_when_it_loses_a_create_race(self): + """Regression: the read, delete and create are three statements, not one. + + Two concurrent sessions for the same server_id raced and 13 of 20 returned 500 against a + live two-worker proxy. The loser's session is in fact ready, because the winner wrote a + draft for it, so it adopts that row instead of failing the caller. + """ + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + winner_draft = generate_mock_mcp_server_db_record(server_id="raced") + winner_draft.approval_status = "draft" + # First lookup: nothing yet. After the losing create blows up: the winner's row. + lookups = AsyncMock(side_effect=[None, winner_draft]) + + with ( + patch("litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", lookups), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server", + AsyncMock(side_effect=Exception("duplicate key value violates unique constraint")), + ), + ): + result = await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="raced", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + assert result.server_id == "raced" + assert lookups.await_count == 2 + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_reraises_when_the_create_failure_was_not_a_race(self): + """A genuine database error must not be swallowed by the race-adoption path.""" + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + with ( + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", + AsyncMock(side_effect=[None, None]), + ), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + AsyncMock(return_value=[]), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server", + AsyncMock(side_effect=Exception("connection refused")), + ), + pytest.raises(Exception, match="connection refused"), + ): + await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="broken", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + @pytest.mark.asyncio + async def test_create_draft_mcp_server_prunes_drafts_past_their_lifetime(self): + """Regression: abandoned OAuth sessions accumulated forever. Verified against a live + proxy, where 12 drafts aged past the lifetime were still present and a 13th was added.""" + from litellm.proxy._experimental.mcp_server.db import create_draft_mcp_server + + from datetime import timezone + + now = datetime.now(timezone.utc) + stale_one = generate_mock_mcp_server_db_record(server_id="stale-1") + stale_one.updated_at = now - timedelta(hours=1) + stale_two = generate_mock_mcp_server_db_record(server_id="stale-2") + stale_two.updated_at = now - timedelta(hours=1) + # A draft still inside its lifetime must survive the sweep. + fresh_draft = generate_mock_mcp_server_db_record(server_id="still-live") + fresh_draft.updated_at = now + find_rows = AsyncMock(return_value=[stale_one, stale_two, fresh_draft]) + delete_call = AsyncMock() + + with ( + patch("litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", find_rows), + patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_row", + AsyncMock(return_value=None), + ), + patch("litellm.proxy._experimental.mcp_server.db.delete_mcp_server", delete_call), + patch( + "litellm.proxy._experimental.mcp_server.db.create_mcp_server", + AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id="fresh")), + ), + ): + await create_draft_mcp_server( + MagicMock(), + NewMCPServerRequest(server_id="fresh", url="https://x.example.com/mcp"), + "tester", + ttl_seconds=300, + ) + + # Only drafts are considered, only the expired ones are removed, and the live one stays. + assert find_rows.await_args.kwargs["where"]["approval_status"] == "draft" + assert sorted(c.args[1] for c in delete_call.await_args_list) == ["stale-1", "stale-2"] + + @pytest.mark.asyncio + async def test_get_all_mcp_servers_hides_drafts_without_hiding_legacy_null_rows(self): + """Drafts are addressable only by their own id and must never appear in a listing, but a + bare inequality would also drop pre-approval-workflow rows, since SQL evaluates + NULL != 'draft' as NULL.""" + from litellm.proxy._experimental.mcp_server.db import get_all_mcp_servers + + find_rows = AsyncMock(return_value=[]) + with patch( + "litellm.proxy._experimental.mcp_server.db._db_find_mcp_server_rows", + find_rows, + ): + await get_all_mcp_servers(MagicMock()) + + where = find_rows.await_args.args[1] + assert where == {"OR": [{"approval_status": None}, {"approval_status": {"not": "draft"}}]} + + @pytest.mark.asyncio + async def test_resolve_session_server_id_refuses_an_unknown_caller_supplied_id(self): + """Regression: two concurrent sessions must never land on one id. + + Honouring an arbitrary caller-supplied id lets a second session adopt the first's draft and + run OAuth against its URL and client credentials, silently. An id naming no real server is + therefore replaced with a fresh one. + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=None), + ), + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="someone-elses-id", url="https://x.example.com/mcp") + ) + + assert resolved != "someone-elses-id" + uuid.UUID(resolved) + + @pytest.mark.asyncio + async def test_resolve_session_server_id_refuses_an_id_that_names_another_sessions_draft(self): + """A draft row is another session's, not a saved server. Replaying an id this endpoint + previously returned must not let a later session inherit the earlier one's config.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + someone_elses_draft = generate_mock_mcp_server_db_record(server_id="earlier-session") + someone_elses_draft.approval_status = "draft" + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + AsyncMock(return_value=someone_elses_draft), + ), + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="earlier-session", url="https://x.example.com/mcp") + ) + + assert resolved != "earlier-session" + uuid.UUID(resolved) + + @pytest.mark.asyncio + async def test_resolve_session_server_id_keeps_a_real_servers_id_for_the_edit_flow(self): + """The edit form re-authorizes a saved server against its own id, which must be preserved + or the flow would authorize a throwaway id instead of the server being edited.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = generate_mock_mcp_server_config_record(server_id="saved") + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="saved", url="https://x.example.com/mcp") + ) + + assert resolved == "saved" + + @pytest.mark.asyncio + async def test_resolve_session_server_id_keeps_the_supplied_id_without_a_database(self): + """No database means nothing shared to collide over, so behaviour stays as it is today.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _resolve_session_server_id, + ) + + mock_manager = MagicMock() + mock_manager.get_mcp_server_by_id.return_value = None + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none", + return_value=None, + ), + ): + resolved = await _resolve_session_server_id( + NewMCPServerRequest(server_id="no-db-id", url="https://x.example.com/mcp") + ) + + assert resolved == "no-db-id" + @pytest.mark.asyncio async def test_get_cached_temporary_mcp_server_or_404(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -5705,7 +6061,9 @@ def _edit_endpoint_patches(old_record, update_mock): ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(side_effect=old_record) if isinstance(old_record, Exception) else AsyncMock(return_value=old_record), + AsyncMock(side_effect=old_record) + if isinstance(old_record, Exception) + else AsyncMock(return_value=old_record), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", @@ -6114,7 +6472,13 @@ def test_bundled_openapi_registry_parses_and_entries_are_well_formed(): registry_path = os.path.join( os.path.dirname(os.path.abspath(__file__)), - "..", "..", "..", "..", "litellm", "proxy", "openapi_registry.json", + "..", + "..", + "..", + "..", + "litellm", + "proxy", + "openapi_registry.json", ) with open(registry_path) as f: registry = json.load(f)