diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index e1c1fcb232c..ceff1809b7b 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -94,6 +94,7 @@ export interface KeyResponse { object_permission_id: string; mcp_servers: string[]; mcp_access_groups?: string[]; + mcp_toolsets?: string[] | null; mcp_tool_permissions?: Record; vector_stores: string[]; agents?: string[]; diff --git a/ui/litellm-dashboard/src/components/object_permissions_view.tsx b/ui/litellm-dashboard/src/components/object_permissions_view.tsx index be021a5b59d..b0ee38bd834 100644 --- a/ui/litellm-dashboard/src/components/object_permissions_view.tsx +++ b/ui/litellm-dashboard/src/components/object_permissions_view.tsx @@ -9,7 +9,7 @@ interface ObjectPermission { mcp_servers: string[]; mcp_access_groups?: string[]; mcp_tool_permissions?: Record; - mcp_toolsets?: string[]; + mcp_toolsets?: string[] | null; vector_stores: string[]; agents?: string[]; agent_access_groups?: string[]; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index f84b95b8d4d..2f72fa37d86 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -443,6 +443,36 @@ describe("CreateKey", () => { }); }); + it("should include mcp_toolsets in keyCreateCall payload when only toolsets are selected", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /create key/i })).toBeInTheDocument(); + }); + + act(() => { + formMock.setFieldValue("key_alias", "Test Key"); + formMock.setFieldValue("allowed_mcp_servers_and_groups", { + servers: [], + accessGroups: [], + toolsets: ["ts-1"], + }); + }); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create key/i })); + }); + + await waitFor(() => { + expect(mockKeyCreateCall).toHaveBeenCalled(); + }); + expect(mockKeyCreateCall.mock.calls[0][2].object_permission?.mcp_toolsets).toEqual(["ts-1"]); + }); + it("should prefill models when provided without team_id", async () => { renderWithProviders( = ({ team, teams, data, addKey, autoOp if ( formValues.allowed_mcp_servers_and_groups && (formValues.allowed_mcp_servers_and_groups.servers?.length > 0 || - formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0) + formValues.allowed_mcp_servers_and_groups.accessGroups?.length > 0 || + formValues.allowed_mcp_servers_and_groups.toolsets?.length > 0) ) { if (!formValues.object_permission) { formValues.object_permission = {}; } - const { servers, accessGroups } = formValues.allowed_mcp_servers_and_groups; + const { servers, accessGroups, toolsets } = formValues.allowed_mcp_servers_and_groups; if (servers && servers.length > 0) { formValues.object_permission.mcp_servers = servers; } if (accessGroups && accessGroups.length > 0) { formValues.object_permission.mcp_access_groups = accessGroups; } + if (toolsets && toolsets.length > 0) { + formValues.object_permission.mcp_toolsets = toolsets; + } // Remove the original field as it's now part of object_permission delete formValues.allowed_mcp_servers_and_groups; } 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 e58e24d0bf7..fa8ffcd7366 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -453,6 +453,28 @@ describe("KeyInfoView handleKeyUpdate guardrails guard", () => { }); }); +describe("KeyInfoView handleKeyUpdate mcp_toolsets", () => { + it("should forward the toolsets the edit form supplies into object_permission", async () => { + renderView(true); + + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + (globalThis as any).__TEST_FORM_VALUES = { + token: "tok_123", + max_budget: 40000, + mcp_servers_and_groups: { servers: [], accessGroups: [], toolsets: ["ts-1"] }, + }; + + fireEvent.click(screen.getByText("Mock Submit")); + + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + + const [, sentPayload] = keyUpdateCallMock.mock.calls[0]; + expect(sentPayload.object_permission.mcp_toolsets).toEqual(["ts-1"]); + expect(sentPayload.max_budget).toBe(40000); + }); +}); + describe("KeyInfoView handleKeyUpdate budget_duration", () => { it("should send a canonical budget_duration through unchanged", async () => { renderView(true); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 693c50374c1..756481487e7 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -631,6 +631,40 @@ describe("KeyEditView", () => { }); }); + it("should keep mcp_toolsets when saving an edit that does not touch the MCP selector", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + const keyDataWithToolset = { + ...MOCK_KEY_DATA, + object_permission: { + ...MOCK_KEY_DATA.object_permission!, + mcp_toolsets: ["ts-1"], + }, + }; + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="test-token" + userID="test-user" + userRole="admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Save Changes")).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].mcp_servers_and_groups.toolsets).toEqual(["ts-1"]); + }); + it("should submit budget_limits: [] when the last budget window is deleted", async () => { const onSubmitMock = vi.fn().mockResolvedValue(undefined); const keyDataWithWindow = { diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 23b868d2b87..a9fa05d817d 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -200,6 +200,7 @@ export function KeyEditView({ mcp_servers_and_groups: { servers: keyData.object_permission?.mcp_servers || [], accessGroups: keyData.object_permission?.mcp_access_groups || [], + toolsets: keyData.object_permission?.mcp_toolsets || [], }, mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {}, agents_and_groups: { @@ -233,6 +234,7 @@ export function KeyEditView({ mcp_servers_and_groups: { servers: keyData.object_permission?.mcp_servers || [], accessGroups: keyData.object_permission?.mcp_access_groups || [], + toolsets: keyData.object_permission?.mcp_toolsets || [], }, mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {}, throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false,