Merge pull request #21234 from BerriAI/litellm_ui_access_groups

[Feature] UI - Keys/Teams: Add Access Group Selector to Create and Edit Flow
This commit is contained in:
yuneng-jiang 2026-02-14 17:08:13 -08:00 • committed by GitHub
commit 801fc22d6e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 474 additions and 14 deletions

View file

@ -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 }) => (
<input
data-testid="access-group-selector"
value={Array.isArray(value) ? value.join(",") : ""}
onChange={(e) => 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(
<OldTeams
teams={[]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="Admin"
organizations={[{ organization_id: "org-1", organization_alias: "Org 1", models: [], members: [] }]}
/>,
);
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();

View file

@ -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<TeamProps> = ({
}))}
/>
</Form.Item>
<Form.Item
label={
<span>
Access Groups{" "}
<Tooltip title="Assign access groups to this team. Access groups control which models, MCP servers, and agents this team can use">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
}
name="access_group_ids"
className="mt-8"
help="Select access groups to assign to this team"
>
<AccessGroupSelector placeholder="Select access groups (optional)" />
</Form.Item>
<Form.Item
label={
<span>

View file

@ -0,0 +1,112 @@
import React from "react";
import { Select, Skeleton } from "antd";
import { TeamOutlined } from "@ant-design/icons";
import { Text } from "@tremor/react";
import {
useAccessGroups,
AccessGroupResponse,
} from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
export interface AccessGroupSelectorProps {
value?: string[];
onChange?: (value: string[]) => void;
placeholder?: string;
disabled?: boolean;
style?: React.CSSProperties;
className?: string;
showLabel?: boolean;
labelText?: string;
/** Allow clearing the selection */
allowClear?: boolean;
}
/**
* Reusable multi-select selector for access groups.
*
* - Displays the **access_group_name** in the dropdown.
* - Returns an array of **access_group_id** values.
* - Always multi-select since users can assign multiple access groups.
* - Integrates with Ant Design `<Form.Item>` out of the box via `value` / `onChange`.
*/
const AccessGroupSelector: React.FC<AccessGroupSelectorProps> = ({
value,
onChange,
placeholder = "Select access groups",
disabled = false,
style,
className,
showLabel = false,
labelText = "Access Group",
allowClear = true,
}) => {
const { data: accessGroups, isLoading, isError } = useAccessGroups();
// ── Loading skeleton ─────────────────────────────────────────────────────
if (isLoading) {
return (
<div>
{showLabel && (
<Text className="font-medium block mb-2 text-gray-700 flex items-center">
<TeamOutlined className="mr-2" /> {labelText}
</Text>
)}
<Skeleton.Input active block style={{ height: 32, ...style }} />
</div>
);
}
// ── Build options ────────────────────────────────────────────────────────
const options = (accessGroups ?? []).map((group: AccessGroupResponse) => ({
label: (
<span>
<span className="font-medium">{group.access_group_name}</span>{" "}
<span className="text-gray-400 text-xs">({group.access_group_id})</span>
</span>
),
value: group.access_group_id,
selectedLabel: group.access_group_name,
searchText: `${group.access_group_name} ${group.access_group_id}`,
}));
// ── Render ───────────────────────────────────────────────────────────────
return (
<div>
{showLabel && (
<Text className="font-medium block mb-2 text-gray-700 flex items-center">
<TeamOutlined className="mr-2" /> {labelText}
</Text>
)}
<Select
mode="multiple"
value={value}
placeholder={placeholder}
onChange={onChange}
disabled={disabled}
allowClear={allowClear}
showSearch
style={{ width: "100%", ...style }}
className={`rounded-md ${className ?? ""}`}
notFoundContent={
isError ? (
<span className="text-red-500">Failed to load access groups</span>
) : (
"No access groups found"
)
}
filterOption={(input, option) => {
const searchText =
options.find((opt) => opt.value === option?.value)?.searchText ?? "";
return searchText.toLowerCase().includes(input.toLowerCase());
}}
optionLabelProp="selectedLabel"
options={options.map((opt) => ({
label: opt.label,
value: opt.value,
selectedLabel: opt.selectedLabel,
}))}
/>
</div>
);
};
export default AccessGroupSelector;

View file

@ -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;

View file

@ -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 }) => (
<input
data-testid="access-group-selector"
value={Array.isArray(value) ? value.join(",") : ""}
onChange={(e) => 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(<CreateKey {...defaultProps} />);
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(<CreateKey {...defaultProps} />);
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 },
);
});
});

View file

@ -20,6 +20,7 @@ import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSel
import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings";
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion";
import AccessGroupSelector from "../common_components/AccessGroupSelector";
import TeamDropdown from "../common_components/team_dropdown";
import { CreateUserButton } from "../CreateUserButton";
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
@ -998,6 +999,23 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
options={promptsList.map((name) => ({ value: name, label: name }))}
/>
</Form.Item>
<Form.Item
label={
<span>
Access Groups{" "}
<Tooltip title="Assign access groups to this key. Access groups control which models, MCP servers, and agents this key can use">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
}
name="access_group_ids"
className="mt-4"
help="Select access groups to assign to this key"
>
<AccessGroupSelector
placeholder="Select access groups (optional)"
/>
</Form.Item>
<Form.Item
label={
<span>

View file

@ -89,6 +89,17 @@ vi.mock("@/components/team/member_permissions", () => ({
default: vi.fn(() => <div>Member Permissions</div>),
}));
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(<TeamInfoView {...defaultProps} />);
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",
})
);
});
});
});

View file

@ -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<string, string>;
} | null;
created_at: string;
access_group_ids?: string[];
guardrails?: string[];
policies?: string[];
object_permission?: {
@ -535,6 +537,11 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
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<TeamInfoProps> = ({
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<TeamInfoProps> = ({
/>
</Form.Item>
<Form.Item
label={
<span>
Access Groups{" "}
<Tooltip title="Assign access groups to this team. Access groups control which models, MCP servers, and agents this team can use">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
}
name="access_group_ids"
>
<AccessGroupSelector placeholder="Select access groups (optional)" />
</Form.Item>
<Form.Item label="Vector Stores" name="vector_stores" aria-label="Vector Stores">
<VectorStoreSelector
onChange={(values: string[]) => form.setFieldValue("vector_stores", values)}

View file

@ -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 }) => (
<input
data-testid="access-group-selector"
value={Array.isArray(value) ? value.join(",") : ""}
onChange={(e) => 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(
<KeyEditView
keyData={keyDataWithAccessGroups}
onCancel={() => {}}
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<void>((resolve) => {

View file

@ -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({
</Tooltip>
</Form.Item>
<Form.Item
label={
<span>
Access Groups{" "}
<Tooltip title="Assign access groups to this key. Access groups control which models, MCP servers, and agents this key can use">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
}
name="access_group_ids"
>
<AccessGroupSelector placeholder="Select access groups (optional)" />
</Form.Item>
<Form.Item label="Allowed Pass Through Routes" name="allowed_passthrough_routes">
<Tooltip
title={!premiumUser ? "Setting allowed pass through routes by key is a premium feature" : ""}