diff --git a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts index cb04a8fa1ec..18b50906e65 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerPayload.ts @@ -1,13 +1,11 @@ import type { EndUser } from "@/app/(dashboard)/hooks/customers/useCustomers"; import type { UpdateCustomerRequest } from "@/app/(dashboard)/hooks/customers/customerApi"; import type { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; +import { mcpAllowedToolsFor, mcpServersForIdentifier } from "@/components/mcp_server_management/effectiveMcpServers"; import { - mcpAllowedToolsFor, - mcpServersForIdentifier, - mcpToolPermissionKeysFor, - resolveEffectiveMcpServers, -} from "@/components/mcp_server_management/effectiveMcpServers"; -import { extractMcpEntitlement, type McpEntitlementUpdate } from "@/components/mcp_server_management/mcpEntitlement"; + extractMcpEntitlementUpdate, + type McpEntitlementUpdate, +} from "@/components/mcp_server_management/mcpEntitlement"; export const getCustomerBudgetCounts = (customers: readonly EndUser[]): ReadonlyMap => customers.reduce( @@ -57,85 +55,6 @@ export const toCustomerFormValues = (customer: EndUser | null): CustomerFormValu }; }; -const normalizedStrings = (values: string[]): string[] => Array.from(new Set(values)).toSorted(); - -const grantIdentity = (grant: McpEntitlementUpdate): string => { - const normalized = { - mcp_servers: normalizedStrings(grant.mcp_servers), - mcp_access_groups: normalizedStrings(grant.mcp_access_groups), - mcp_toolsets: normalizedStrings(grant.mcp_toolsets), - mcp_tool_permissions: Object.fromEntries( - Object.keys(grant.mcp_tool_permissions) - .toSorted() - .map((server) => [server, normalizedStrings(grant.mcp_tool_permissions[server])]), - ), - }; - return JSON.stringify(normalized); -}; - -const extractEditedMcpGrant = ( - values: CustomerFormValues, - storedGrant: McpEntitlementUpdate, - servers: MCPServer[], - toolsets: MCPToolset[], -): McpEntitlementUpdate | null => { - const grant = extractMcpEntitlement(values, servers, toolsets); - if (!grant) return null; - const storedSelection = { - allServers: servers, - selectedServers: storedGrant.mcp_servers, - selectedAccessGroups: storedGrant.mcp_access_groups, - selectedToolsets: storedGrant.mcp_toolsets, - toolsets, - toolPermissions: {}, - }; - const selectedServerIds = new Set(resolveEffectiveMcpServers(storedSelection).map(({ server }) => server.server_id)); - // A stored tool-permission key grants every server it names, so servers that only a key reaches survive edits to - // the selection, including the other servers behind an alias a selected server shares - const keyOnlyServers = Object.keys(storedGrant.mcp_tool_permissions) - .flatMap((key) => mcpServersForIdentifier(servers, key)) - .filter((server) => !selectedServerIds.has(server.server_id)); - const keyOnlyServerIds = new Set(keyOnlyServers.map((server) => server.server_id)); - const namesOnlyKeyOnlyServers = (key: string): boolean => { - const named = mcpServersForIdentifier(servers, key); - return named.length > 0 && named.every((server) => keyOnlyServerIds.has(server.server_id)); - }; - const independentPermissions = Object.fromEntries( - Object.entries(values.mcp_tool_permissions).filter(([key]) => namesOnlyKeyOnlyServers(key)), - ); - const keyOnlyPermissions = Object.fromEntries( - keyOnlyServers - .filter((server) => mcpToolPermissionKeysFor(server, independentPermissions, servers).length === 0) - .flatMap((server) => { - const allowed = mcpAllowedToolsFor(server, values.mcp_tool_permissions, servers); - return allowed ? [[server.server_id, [...allowed]]] : []; - }), - ); - return { - ...grant, - mcp_tool_permissions: { ...grant.mcp_tool_permissions, ...independentPermissions, ...keyOnlyPermissions }, - }; -}; - -const changedMcpGrant = ( - values: CustomerFormValues, - existing: EndUser | null, - servers: MCPServer[], - toolsets: MCPToolset[], -): McpEntitlementUpdate | null => { - const storedGrant = storedMcpGrant(existing); - const storedIdentity = grantIdentity(storedGrant); - const formGrant = { - mcp_servers: values.mcp_servers_and_groups.servers, - mcp_access_groups: values.mcp_servers_and_groups.accessGroups, - mcp_toolsets: values.mcp_servers_and_groups.toolsets, - mcp_tool_permissions: values.mcp_tool_permissions, - }; - if (grantIdentity(formGrant) === storedIdentity) return null; - const grant = extractEditedMcpGrant(values, storedGrant, servers, toolsets); - return grant && grantIdentity(grant) !== storedIdentity ? grant : null; -}; - const optionalCustomerFields = ( values: CustomerFormValues, existing: EndUser | null, @@ -153,7 +72,7 @@ export const buildCustomerPayload = ( servers: MCPServer[], toolsets: MCPToolset[], ): UpdateCustomerRequest => { - const grant = changedMcpGrant(values, existing, servers, toolsets); + const grant = extractMcpEntitlementUpdate(values, existing?.object_permission, servers, toolsets); return { user_id: existing?.user_id ?? values.user_id.trim(), ...optionalCustomerFields(values, existing), @@ -187,7 +106,7 @@ export const serversLosingToolLimits = ( servers: MCPServer[], toolsets: MCPToolset[], ): readonly MCPServer[] => { - const grant = changedMcpGrant(values, existing, servers, toolsets); + const grant = extractMcpEntitlementUpdate(values, existing?.object_permission, servers, toolsets); if (!grant) return []; const unrestricted = unrestrictedDirectServerIds(grant, servers, toolsets); const alreadyUnrestricted = unrestrictedDirectServerIds(storedMcpGrant(existing), servers, toolsets); 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 94fd2cec8ac..223ae2416b6 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 @@ -13,6 +13,7 @@ const mockTeamInfoCall = vi.fn(); const mockUserUpdateUserCall = vi.fn(); const mockFetchMCPServers = vi.fn(); const mockListMCPTools = vi.fn(); +const mockMCPServerCatalog = vi.fn(); const MCP_SERVER = { server_id: "srv-1", server_name: "GitHub MCP", alias: "GitHub MCP" }; @@ -70,7 +71,7 @@ vi.mock("@/components/networking", () => { }); vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ - useMCPServers: () => ({ data: [MCP_SERVER], isLoading: false }), + useMCPServers: () => mockMCPServerCatalog(), })); vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPAccessGroups", () => ({ useMCPAccessGroups: () => ({ data: ["dev-group"], isLoading: false }), @@ -108,6 +109,7 @@ describe("UserInfoView", () => { mockTeamMemberDeleteCall.mockResolvedValue({}); mockUserUpdateUserCall.mockResolvedValue({}); mockFetchMCPServers.mockResolvedValue([MCP_SERVER]); + mockMCPServerCatalog.mockReturnValue({ data: [MCP_SERVER], isLoading: false }); mockListMCPTools.mockResolvedValue({ tools: [{ name: "list_issues", description: "List issues" }] }); }); @@ -331,6 +333,43 @@ describe("UserInfoView", () => { }); describe("MCP permissions", () => { + it.each([ + { field: "user_email", label: "Email", value: "renamed@example.com", catalogLoaded: true }, + { field: "user_alias", label: "User Alias", value: "Renamed user", catalogLoaded: false }, + ])( + "preserves shared-alias grants when only $field changes with catalogLoaded=$catalogLoaded", + async ({ field, label, value, catalogLoaded }) => { + const servers = [ + { server_id: "server-a", server_name: "Server A", alias: "shared" }, + { server_id: "server-b", server_name: "Server B", alias: "shared" }, + ]; + const permission = { + mcp_servers: ["server-a"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { shared: ["list_issues"] }, + }; + mockUserGetInfoV2.mockResolvedValue({ ...MOCK_USER_DATA, object_permission: permission }); + mockMCPServerCatalog.mockReturnValue({ data: catalogLoaded ? servers : [], isLoading: !catalogLoaded }); + if (catalogLoaded) mockFetchMCPServers.mockResolvedValue(servers); + else mockFetchMCPServers.mockReturnValue(new Promise(() => {})); + render(); + + fireEvent.change(await screen.findByRole("textbox", { name: label, exact: true }), { target: { value } }); + await userEvent.click(screen.getByRole("button", { name: "Save Changes", exact: true })); + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalledOnce()); + + const [, payload] = mockUserUpdateUserCall.mock.calls[0]; + expect(payload[field]).toBe(value); + expect(payload).not.toHaveProperty("object_permission"); + await userEvent.click(await screen.findByRole("button", { name: "Edit Settings", exact: true })); + expect(await screen.findByRole("textbox", { name: label, exact: true })).toHaveValue(value); + await userEvent.click(screen.getByRole("button", { name: "Save Changes", exact: true })); + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalledTimes(2)); + expect(mockUserUpdateUserCall.mock.calls[1][1]).not.toHaveProperty("object_permission"); + }, + ); + it("should render the user's MCP entitlements in read mode", async () => { const user = userEvent.setup(); render(); @@ -348,7 +387,7 @@ describe("UserInfoView", () => { expect(await screen.findByText("list_issues")).toBeInTheDocument(); }); - it("should nest MCP entitlements under object_permission when an admin saves", async () => { + it("omits unchanged MCP entitlements when an admin saves", async () => { const user = userEvent.setup(); render(); @@ -363,13 +402,7 @@ describe("UserInfoView", () => { expect(token).toBe("test-token"); expect(roleArg).toBeNull(); expect(payload.user_id).toBe("user-123"); - const expectedObjectPermission = { - mcp_servers: ["srv-1"], - mcp_access_groups: ["dev-group"], - mcp_toolsets: [], - mcp_tool_permissions: { "srv-1": ["list_issues"] }, - }; - expect(payload.object_permission).toEqual(expectedObjectPermission); + expect(payload).not.toHaveProperty("object_permission"); expect(payload).not.toHaveProperty("mcp_servers_and_groups"); expect(payload).not.toHaveProperty("mcp_tool_permissions"); expect(payload).not.toHaveProperty("mcp_servers"); @@ -418,9 +451,35 @@ describe("UserInfoView", () => { }); const [, payload] = mockUserUpdateUserCall.mock.calls[0]; - expect(payload.object_permission.mcp_tool_permissions).toEqual({ - "srv-1": ["list_issues"], - "srv-via-group": ["read_only"], + expect(payload).not.toHaveProperty("object_permission"); + }); + + it("retains an independent tool-only server grant when another server's tools are edited", async () => { + const user = userEvent.setup(); + const servers = [MCP_SERVER, { server_id: "server-b", server_name: "Server B" }]; + const permission = { + mcp_servers: ["srv-1"], + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: { "srv-1": ["read_direct"], "server-b": ["read_independent"] }, + }; + mockUserGetInfoV2.mockResolvedValue({ ...MOCK_USER_DATA, object_permission: permission }); + mockFetchMCPServers.mockResolvedValue(servers); + mockMCPServerCatalog.mockReturnValue({ data: servers, isLoading: false }); + mockListMCPTools.mockImplementation((_token: string, serverId: string) => + Promise.resolve({ tools: [{ name: serverId === "srv-1" ? "read_direct" : "read_independent" }] }), + ); + render(); + + await user.click((await screen.findAllByRole("radio", { name: "Flat List" }))[0]); + await user.click(await screen.findByRole("checkbox", { name: "read_direct", exact: true })); + await user.click(screen.getByRole("button", { name: "Save Changes", exact: true })); + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalledOnce()); + + const [, payload] = mockUserUpdateUserCall.mock.calls[0]; + expect(payload.object_permission).toEqual({ + ...permission, + mcp_tool_permissions: { "srv-1": [], "server-b": ["read_independent"] }, }); }); 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 2056142b50a..2839cae82cd 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,7 +42,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import MCPServerPermissions from "@/components/permissions/MCPServerPermissions"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; -import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; +import { extractMcpEntitlementUpdate } from "@/components/mcp_server_management/mcpEntitlement"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; import { AutoRouterUsageView } from "@/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab"; @@ -321,7 +321,12 @@ export default function UserInfoView({ try { if (!accessToken || !userData) return; - const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers, allMcpToolsets); + const mcpEntitlement = extractMcpEntitlementUpdate( + formValues, + userData.object_permission, + 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 index 0288e85f170..3678c8e61f3 100644 --- a/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts +++ b/ui/litellm-dashboard/src/components/mcp_server_management/mcpEntitlement.ts @@ -1,6 +1,12 @@ -import { ALL_PROXY_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; +import { ALL_PROXY_MCP_SERVERS_SENTINEL, NO_MCP_SERVERS_SENTINEL } from "@/components/mcp_tools/constants"; import { MCPServer, MCPToolset } from "@/components/mcp_tools/types"; -import { mcpAllowedToolsFor, mcpServersForIdentifier } from "./effectiveMcpServers"; +import type { ObjectPermission } from "@/components/object_permission_types"; +import { + mcpAllowedToolsFor, + mcpServersForIdentifier, + mcpToolPermissionKeysFor, + resolveEffectiveMcpServers, +} from "./effectiveMcpServers"; export interface McpEntitlementUpdate { mcp_servers: string[]; @@ -19,18 +25,31 @@ const asToolPermissions = (value: unknown): Record => { ); }; +const formMcpGrant = (formValues: Record): McpEntitlementUpdate | null => { + const selection = formValues.mcp_servers_and_groups; + if (selection === null || typeof selection !== "object") return null; + const { servers, accessGroups, toolsets } = selection as Record; + return { + mcp_servers: asStringArray(servers), + mcp_access_groups: asStringArray(accessGroups), + mcp_toolsets: asStringArray(toolsets), + mcp_tool_permissions: asToolPermissions(formValues.mcp_tool_permissions), + }; +}; + 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 formGrant = formMcpGrant(formValues); + if (!formGrant) return null; + const { + mcp_servers: mcpServers, + mcp_access_groups: mcpAccessGroups, + mcp_toolsets: mcpToolsets, + mcp_tool_permissions: toolPermissions, + } = formGrant; const grantsEveryServer = mcpServers.includes(ALL_PROXY_MCP_SERVERS_SENTINEL) || mcpToolsets.some((toolsetId) => !allToolsets.some((toolset) => toolset.toolset_id === toolsetId)); @@ -57,7 +76,6 @@ export const extractMcpEntitlement = ( return named.some(grants); }; - const toolPermissions = asToolPermissions(formValues.mcp_tool_permissions); const splitKeys = new Set( Object.keys(toolPermissions).filter((key) => { const named = mcpServersForIdentifier(allServers, key); @@ -93,3 +111,87 @@ export const extractMcpEntitlement = ( mcp_tool_permissions: { ...retainedPermissions, ...canonicalPermissions }, }; }; + +const normalizedStrings = (values: string[]): string[] => Array.from(new Set(values)).toSorted(); + +const grantIdentity = (grant: McpEntitlementUpdate): string => { + const normalized = { + mcp_servers: normalizedStrings(grant.mcp_servers), + mcp_access_groups: normalizedStrings(grant.mcp_access_groups), + mcp_toolsets: normalizedStrings(grant.mcp_toolsets), + mcp_tool_permissions: Object.fromEntries( + Object.keys(grant.mcp_tool_permissions) + .toSorted() + .map((server) => [server, normalizedStrings(grant.mcp_tool_permissions[server])]), + ), + }; + return JSON.stringify(normalized); +}; + +const independentToolPermissions = ( + storedGrant: McpEntitlementUpdate, + toolPermissions: Record, + allServers: MCPServer[], + allToolsets: MCPToolset[], +): Record => { + const storedSelection = { + allServers, + selectedServers: storedGrant.mcp_servers, + selectedAccessGroups: storedGrant.mcp_access_groups, + selectedToolsets: storedGrant.mcp_toolsets, + toolsets: allToolsets, + toolPermissions: {}, + }; + const selectedServers = resolveEffectiveMcpServers(storedSelection); + const selectedServerIds = new Set(selectedServers.map(({ server }) => server.server_id)); + const keyOnlyServers = Object.keys(storedGrant.mcp_tool_permissions) + .flatMap((key) => mcpServersForIdentifier(allServers, key)) + .filter((server) => !selectedServerIds.has(server.server_id)); + const keyOnlyServerIds = new Set(keyOnlyServers.map((server) => server.server_id)); + const namesOnlyKeyOnlyServers = (key: string): boolean => { + const named = mcpServersForIdentifier(allServers, key); + return named.length > 0 && named.every((server) => keyOnlyServerIds.has(server.server_id)); + }; + const independentPermissions = Object.fromEntries( + Object.entries(toolPermissions).filter(([key]) => namesOnlyKeyOnlyServers(key)), + ); + const canonicalPermissions = Object.fromEntries( + keyOnlyServers + .filter((server) => mcpToolPermissionKeysFor(server, independentPermissions, allServers).length === 0) + .flatMap((server) => { + const allowed = mcpAllowedToolsFor(server, toolPermissions, allServers); + return allowed ? [[server.server_id, [...allowed]]] : []; + }), + ); + return { ...independentPermissions, ...canonicalPermissions }; +}; + +export const extractMcpEntitlementUpdate = ( + formValues: Record, + existingPermission: ObjectPermission | null | undefined, + allServers: MCPServer[], + allToolsets: MCPToolset[] = [], +): McpEntitlementUpdate | null => { + const formGrant = formMcpGrant(formValues); + if (!formGrant) return null; + const storedGrant = { + mcp_servers: existingPermission?.mcp_servers ?? [], + mcp_access_groups: existingPermission?.mcp_access_groups ?? [], + mcp_toolsets: existingPermission?.mcp_toolsets ?? [], + mcp_tool_permissions: existingPermission?.mcp_tool_permissions ?? {}, + }; + const storedIdentity = grantIdentity(storedGrant); + if (grantIdentity(formGrant) === storedIdentity) return null; + const extracted = extractMcpEntitlement(formValues, allServers, allToolsets); + if (!extracted) return null; + const grant = { + ...extracted, + mcp_tool_permissions: formGrant.mcp_servers.includes(NO_MCP_SERVERS_SENTINEL) + ? {} + : { + ...extracted.mcp_tool_permissions, + ...independentToolPermissions(storedGrant, formGrant.mcp_tool_permissions, allServers, allToolsets), + }, + }; + return grantIdentity(grant) === storedIdentity ? null : grant; +}; 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 c82a4823dfd..b271737fe07 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -1,6 +1,11 @@ import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { toast } from "@/lib/toast"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import type { KeyResponse } from "../key_team_helpers/key_list"; +import type { MCPServer } from "../mcp_tools/types"; +import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; +import { toKeyEditFormValues, toSubmittedValues } from "./keyEditFormValues"; vi.mock("@/lib/toast", () => ({ toast: { error: vi.fn(), success: vi.fn() }, @@ -228,7 +233,7 @@ const baseKeyData = { next_rotation_at: null as any, }; -const renderView = (premiumUser: boolean) => { +const renderView = (premiumUser: boolean, keyData: KeyResponse = baseKeyData as KeyResponse) => { // Configure the mock for this test mockUseAuthorized.mockReturnValue({ accessToken: "access_abc", @@ -242,18 +247,13 @@ const renderView = (premiumUser: boolean) => { }); return render( - {}} - keyData={baseKeyData as any} - onKeyDataUpdate={() => {}} - teams={[]} - />, + {}} keyData={keyData} onKeyDataUpdate={() => {}} teams={[]} />, ); }; beforeEach(() => { vi.clearAllMocks(); + vi.mocked(useMCPServers).mockReturnValue({ data: [] } as ReturnType); (globalThis as any).__TEST_FORM_VALUES = undefined; }); @@ -375,6 +375,85 @@ describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { }); }); +describe("KeyInfoView handleKeyUpdate shared MCP aliases", () => { + it.each([ + { + edit: "only renaming the key", + storedServers: ["server-a"], + storedToolPermissions: { shared: ["read"] }, + selectedServers: ["server-a"], + expectedToolPermissions: { shared: ["read"] }, + }, + { + edit: "removing an explicitly selected server", + storedServers: ["server-a", "server-b"], + storedToolPermissions: { shared: ["read"] }, + selectedServers: ["server-a"], + expectedToolPermissions: { "server-a": ["read"] }, + }, + { + edit: "choosing No MCP Servers", + storedServers: ["server-a"], + storedToolPermissions: { "server-b": ["read"] }, + selectedServers: [NO_MCP_SERVERS_SENTINEL], + expectedToolPermissions: {}, + }, + ])( + "preserves unrelated permissions when $edit", + async ({ storedServers, storedToolPermissions, selectedServers, expectedToolPermissions }) => { + const servers: MCPServer[] = ["server-a", "server-b"].map((server_id) => ({ + server_id, + alias: "shared", + created_at: "", + created_by: "", + updated_at: "", + updated_by: "", + })); + vi.mocked(useMCPServers).mockReturnValue({ data: servers } as ReturnType); + const permission = { + mcp_servers: storedServers, + mcp_access_groups: [], + mcp_toolsets: [], + mcp_tool_permissions: storedToolPermissions, + vector_stores: ["vector-store"], + agents: ["agent"], + agent_access_groups: ["agent-group"], + skills: ["skill"], + search_tools: ["search-tool"], + }; + const keyData = { ...baseKeyData, object_permission: permission } as KeyResponse; + renderView(true, keyData); + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as { __TEST_FORM_VALUES?: Record }).__TEST_FORM_VALUES = toSubmittedValues( + { + ...toKeyEditFormValues(keyData), + key_alias: "Renamed key", + mcp_servers_and_groups: { servers: selectedServers, accessGroups: [], toolsets: [] }, + }, + { canViewPolicies: true, canViewPrompts: true }, + ); + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => + expect(keyUpdateCallMock).toHaveBeenCalledExactlyOnceWith( + "access_abc", + expect.objectContaining({ + key: keyData.token, + key_alias: "Renamed key", + object_permission: { + ...permission, + mcp_servers: selectedServers, + mcp_tool_permissions: expectedToolPermissions, + }, + }), + ), + ); + }, + ); +}); + describe("KeyInfoView handleKeyUpdate cache sync", () => { it("should invalidate every cached key query so the list and detail views re-read the saved key", async () => { keyUpdateCallMock.mockResolvedValueOnce({ 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 7b924243fb5..148eedffd04 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -45,7 +45,7 @@ import { useQueryClient } from "@tanstack/react-query"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useTeamMemberBudgets } from "@/app/(dashboard)/hooks/teams/useTeamMemberBudgets"; -import { extractMcpEntitlement } from "../mcp_server_management/mcpEntitlement"; +import { extractMcpEntitlementUpdate } from "../mcp_server_management/mcpEntitlement"; import ObjectPermissionsView from "../object_permissions_view"; import { RegenerateKeyModal } from "../organisms/RegenerateKeyModal"; import { parseErrorMessage } from "../shared/errorUtils"; @@ -269,7 +269,12 @@ export default function KeyInfoView({ delete formValues.vector_stores; } - const mcpEntitlement = extractMcpEntitlement(formValues, allMcpServers ?? [], allMcpToolsets ?? []); + const mcpEntitlement = extractMcpEntitlementUpdate( + formValues, + currentKeyData.object_permission, + 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.