From b3fd76a7131dfac7ec49a5fb5daa82e6b7a509a6 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sat, 14 Feb 2026 17:01:29 -0800 Subject: [PATCH] adding access groups on keys and teams --- .../src/components/OldTeams.test.tsx | 107 +++++++++++++++++- .../src/components/OldTeams.tsx | 16 +++ .../components/key_team_helpers/key_list.tsx | 1 + .../organisms/create_key_button.test.tsx | 80 +++++++++++-- .../src/components/team/TeamInfo.test.tsx | 57 ++++++++++ .../src/components/team/TeamInfo.tsx | 22 ++++ .../templates/key_edit_view.test.tsx | 58 ++++++++++ .../components/templates/key_edit_view.tsx | 17 +++ 8 files changed, 344 insertions(+), 14 deletions(-) diff --git a/ui/litellm-dashboard/src/components/OldTeams.test.tsx b/ui/litellm-dashboard/src/components/OldTeams.test.tsx index 0ca37cc82d7..f0f08e3907a 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.test.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.test.tsx @@ -14,7 +14,8 @@ vi.mock("./networking", () => ({ teamDeleteCall: vi.fn(), fetchMCPAccessGroups: vi.fn(), v2TeamListCall: vi.fn(), - getGuardrailsList: vi.fn(), + getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), + getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), })); vi.mock("./common_components/fetch_teams", () => ({ @@ -70,12 +71,16 @@ vi.mock("./ModelSelect/ModelSelect", () => { data-testid={dataTestId || "model-select"} value={Array.isArray(value) ? value.join(", ") : ""} onChange={(e) => { - // Mock onChange - in real usage this would be handled by Ant Design Select if (onChange) { - onChange(value || []); + const newVal = e.target.value + ? e.target.value + .split(",") + .map((s: string) => s.trim()) + .filter(Boolean) + : []; + onChange(newVal); } }} - readOnly /> ); }); @@ -89,6 +94,27 @@ vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ useOrganizations: () => mockUseOrganizations(), })); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1" }, + { access_group_id: "ag-2", access_group_name: "Group 2" }, + ], + isLoading: false, + isError: false, + }), +})); + +vi.mock("./common_components/AccessGroupSelector", () => ({ + default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( + onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} + /> + ), +})); + const createQueryClient = () => { return new QueryClient({ defaultOptions: { @@ -740,6 +766,79 @@ describe("OldTeams - Default Team Settings tab visibility", () => { }); }); +describe("OldTeams - access_group_ids in team create", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "Test Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + } as any); + mockUseOrganizations.mockReturnValue({ data: [{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }] }); + }); + + it("should pass access_group_ids to teamCreateCall when creating team", async () => { + renderWithQueryClient( + , + ); + + const createButton = screen.getByRole("button", { name: /create new team/i }); + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + const teamNameInput = screen.getByLabelText(/team name/i); + fireEvent.change(teamNameInput, { target: { value: "Test Team" } }); + + const modelsInput = screen.getByTestId("create-team-models-select"); + fireEvent.change(modelsInput, { target: { value: "gpt-4" } }); + + const additionalSettingsAccordion = screen.getByText("Additional Settings"); + fireEvent.click(additionalSettingsAccordion); + + await waitFor(() => { + expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); + }); + + const accessGroupInput = screen.getByTestId("access-group-selector"); + fireEvent.change(accessGroupInput, { target: { value: "ag-1,ag-2" } }); + + const createTeamSubmitButton = screen.getByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButton); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_alias: "Test Team", + models: ["gpt-4"], + access_group_ids: ["ag-1", "ag-2"], + }), + ); + }); + }); +}); + describe("OldTeams - models dropdown options", () => { beforeEach(() => { vi.clearAllMocks(); diff --git a/ui/litellm-dashboard/src/components/OldTeams.tsx b/ui/litellm-dashboard/src/components/OldTeams.tsx index 7b906505759..702724c9729 100644 --- a/ui/litellm-dashboard/src/components/OldTeams.tsx +++ b/ui/litellm-dashboard/src/components/OldTeams.tsx @@ -38,6 +38,7 @@ import { import { Button as Button2, Form, Input, Modal, Select as Select2, Switch, Tooltip, Typography } from "antd"; import React, { useEffect, useState } from "react"; import { formatNumberWithCommas } from "../utils/dataUtils"; +import AccessGroupSelector from "./common_components/AccessGroupSelector"; import AgentSelector from "./agent_management/AgentSelector"; import { fetchTeams } from "./common_components/fetch_teams"; import ModelAliasManager from "./common_components/ModelAliasManager"; @@ -1376,6 +1377,21 @@ const Teams: React.FC = ({ }))} /> + + Access Groups{" "} + + + + + } + name="access_group_ids" + className="mt-8" + help="Select access groups to assign to this team" + > + + 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 5511a671db7..b54eb21a0ae 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 @@ -87,6 +87,7 @@ export interface KeyResponse { agents?: string[]; agent_access_groups?: string[]; }; + access_group_ids?: string[]; auto_rotate?: boolean; rotation_interval?: string; last_rotation_at?: 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 bd5e0ee1cbc..b5198e4f861 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 @@ -1,15 +1,19 @@ +import { act, fireEvent, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen } from "../../../tests/test-utils"; import CreateKey from "./create_key_button"; -const mockKeyCreateCall = vi.fn().mockResolvedValue({ - key: "test-api-key", - soft_budget: null, +const { mockKeyCreateCall } = vi.hoisted(() => { + const fn = vi.fn().mockResolvedValue({ + key: "test-api-key", + soft_budget: null, + }); + return { mockKeyCreateCall: fn }; }); -vi.mock("./networking", () => ({ +vi.mock("../networking", () => ({ keyCreateCall: mockKeyCreateCall, - modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }] }), getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }), proxyBaseUrl: "http://localhost:4000", @@ -25,7 +29,7 @@ vi.mock("./networking", () => ({ fetchMCPAccessGroups: vi.fn().mockResolvedValue([]), })); -vi.mock("./molecules/notifications_manager", () => ({ +vi.mock("../molecules/notifications_manager", () => ({ default: { success: vi.fn(), fromBackend: vi.fn(), @@ -36,16 +40,22 @@ vi.mock("./molecules/notifications_manager", () => ({ }, })); +vi.mock("../common_components/AccessGroupSelector", () => ({ + default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( + onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} + /> + ), +})); + describe("CreateKey", () => { const defaultProps = { - userID: "test-user-id", team: null, - userRole: "Admin", - accessToken: "test-token", data: [], teams: [], addKey: vi.fn(), - premiumUser: false, }; beforeEach(() => { @@ -61,4 +71,54 @@ describe("CreateKey", () => { renderWithProviders(); expect(screen.getByRole("button", { name: /create new key/i })).toBeInTheDocument(); }); + + it("should include access_group_ids in keyCreateCall payload when access groups are selected", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/key name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/key name/i), { target: { value: "Test Key" } }); + + const optionalSettingsAccordion = screen.getByText("Optional Settings"); + act(() => { + fireEvent.click(optionalSettingsAccordion); + }); + + await waitFor(() => { + expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); + + const modelsCombobox = screen.getAllByRole("combobox").find((el) => el.closest('[class*="ant-form-item"]')?.textContent?.includes("Models")) || + screen.getAllByRole("combobox")[1]; + if (modelsCombobox) { + act(() => fireEvent.mouseDown(modelsCombobox)); + await waitFor(() => { + const allTeamModels = [...document.body.querySelectorAll(".ant-select-item")].find( + (el) => el.textContent?.includes("All Team Models"), + ); + if (allTeamModels) fireEvent.click(allTeamModels); + }); + } + + const createButton = screen.getByRole("button", { name: /create key/i }); + act(() => fireEvent.click(createButton)); + + await waitFor( + () => { + expect(mockKeyCreateCall).toHaveBeenCalled(); + const formValues = mockKeyCreateCall.mock.calls[0][2]; + expect(formValues).toHaveProperty("access_group_ids"); + expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]); + }, + { timeout: 5000 }, + ); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 7a15a6e3f22..cc0483aafd6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -89,6 +89,17 @@ vi.mock("@/components/team/member_permissions", () => ({ default: vi.fn(() =>
Member Permissions
), })); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1" }, + { access_group_id: "ag-2", access_group_name: "Group 2" }, + ], + isLoading: false, + isError: false, + }), +})); + import { useAllProxyModels } from "@/app/(dashboard)/hooks/models/useModels"; import { useOrganization } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useTeam } from "@/app/(dashboard)/hooks/teams/useTeams"; @@ -669,4 +680,50 @@ describe("TeamInfoView", () => { expect(screen.getByText(/alert1@test\.com, alert2@test\.com/)).toBeInTheDocument(); }); }); + + it("should pass access_group_ids to teamUpdateCall when saving team settings", async () => { + const user = userEvent.setup(); + const accessGroupIds = ["ag-1", "ag-2"]; + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + access_group_ids: accessGroupIds, + models: ["gpt-4"], + }) + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + const settingsTab = screen.getByRole("tab", { name: "Settings" }); + await user.click(settingsTab); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + + const editButton = screen.getByRole("button", { name: /edit settings/i }); + await user.click(editButton); + + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + + const saveButton = screen.getByRole("button", { name: /save changes/i }); + await user.click(saveButton); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + access_group_ids: accessGroupIds, + team_id: "123", + }) + ); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index d081538fb64..ea7a9a1c460 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -23,6 +23,7 @@ import { Button, Form, Input, message, Select, Switch, Tabs, Tooltip } from "ant import { CheckIcon, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; +import AccessGroupSelector from "../common_components/AccessGroupSelector"; import AgentSelector from "../agent_management/AgentSelector"; import DeleteResourceModal from "../common_components/DeleteResourceModal"; import DurationSelect from "../common_components/DurationSelect"; @@ -90,6 +91,7 @@ export interface TeamData { model_aliases: Record; } | null; created_at: string; + access_group_ids?: string[]; guardrails?: string[]; policies?: string[]; object_permission?: { @@ -535,6 +537,11 @@ const TeamInfoView: React.FC = ({ updateData.object_permission.vector_stores = values.vector_stores; } + // Pass access_group_ids to the update request + if (values.access_group_ids !== undefined) { + updateData.access_group_ids = values.access_group_ids; + } + const response = await teamUpdateCall(accessToken, updateData); NotificationsManager.success("Team settings updated successfully"); @@ -800,6 +807,7 @@ const TeamInfoView: React.FC = ({ agents: info.object_permission?.agents || [], accessGroups: info.object_permission?.agent_access_groups || [], }, + access_group_ids: info.access_group_ids || [], }} layout="vertical" > @@ -970,6 +978,20 @@ const TeamInfoView: React.FC = ({ /> + + Access Groups{" "} + + + + + } + name="access_group_ids" + > + + + form.setFieldValue("vector_stores", values)} 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 f923dfeadd3..bbe3c206cee 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 @@ -56,6 +56,27 @@ vi.mock("../organisms/create_key_button", () => ({ fetchTeamModels: vi.fn().mockResolvedValue(["team-model-1", "team-model-2"]), })); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ + data: [ + { access_group_id: "ag-1", access_group_name: "Group 1" }, + { access_group_id: "ag-2", access_group_name: "Group 2" }, + ], + isLoading: false, + isError: false, + }), +})); + +vi.mock("../common_components/AccessGroupSelector", () => ({ + default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( + onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} + /> + ), +})); + describe("KeyEditView", () => { const MOCK_KEY_DATA: KeyResponse = { token: "test-token-123", @@ -436,6 +457,43 @@ describe("KeyEditView", () => { }); + it("should pass access_group_ids to onSubmit when saving key with access groups", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + const keyDataWithAccessGroups = { + ...MOCK_KEY_DATA, + access_group_ids: ["ag-1"], + }; + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="test-token" + userID="test-user" + userRole="admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); + }); + + const accessGroupInput = screen.getByTestId("access-group-selector"); + await userEvent.clear(accessGroupInput); + await userEvent.type(accessGroupInput, "ag-1,ag-2"); + + const submitButton = screen.getByRole("button", { name: /save changes/i }); + await userEvent.click(submitButton); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + const callArgs = onSubmitMock.mock.calls[0][0]; + expect(callArgs.access_group_ids).toEqual(["ag-1", "ag-2"]); + }); + }); + it("should disable cancel button during submission", async () => { let resolveSubmit: (() => void) | undefined; const submitPromise = new Promise((resolve) => { 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 64676e3b94a..1ba3cc7f7b4 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -5,6 +5,7 @@ import { TextInput, Button as TremorButton } from "@tremor/react"; import { Form, Input, Select, Switch, Tooltip } from "antd"; import { useEffect, useState } from "react"; import AgentSelector from "../agent_management/AgentSelector"; +import AccessGroupSelector from "../common_components/AccessGroupSelector"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; @@ -169,6 +170,7 @@ export function KeyEditView({ disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], + access_group_ids: keyData.access_group_ids || [], auto_rotate: keyData.auto_rotate || false, ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }), allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 @@ -196,6 +198,7 @@ export function KeyEditView({ disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) : [], + access_group_ids: keyData.access_group_ids || [], auto_rotate: keyData.auto_rotate || false, ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }), allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 @@ -500,6 +503,20 @@ export function KeyEditView({ + + Access Groups{" "} + + + + + } + name="access_group_ids" + > + + +