scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

adding access groups on keys and teams

Browse source
This commit is contained in:
yuneng-jiang 2026-02-14 17:01:29 -08:00
parent 27fe16be4b
commit b3fd76a713
8 changed files with 344 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

107
ui/litellm-dashboard/src/components/OldTeams.test.tsx
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();

16
ui/litellm-dashboard/src/components/OldTeams.tsx
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>

1
ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
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;

80
ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx
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 },
);
});
});

57
ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx
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",
})
);
});
});
});

22
ui/litellm-dashboard/src/components/team/TeamInfo.tsx
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)}

58
ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx
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) => {

17
ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
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" : ""}