diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 20ab0212854..4a4c31c4708 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -13,7 +13,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `realtime/` - realtime websocket sessions, including the pipecat audio path - `budgets/` - budget definition, enforcement, and reset windows (key, team, tag, soft, multi-window) - `spend_tracking/` - spend logging and cost attribution on `/spend/*` -- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials +- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials; also the dashboard UI behavior on top of them, driven through the proxy-served UI at /ui with playwright (optional dep behind importorskip) - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (rate limits, fallbacks, cooldowns) diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 4037f324393..4395d299c5b 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -34,6 +34,20 @@ The suites run against a live proxy, so bring one up first. `docker-compose.yml` uv run pytest tests/e2e/llm_translation/ -v ``` + The browser tests in the `management/` suite drive the dashboard the proxy serves at `/ui` through playwright, an optional dependency behind `importorskip` (the suite's API tests run without it). Install it once into your environment along with its browser: + + ```bash + uv pip install playwright + uv run playwright install chromium + ``` + + They also need a proxy whose bundled UI contains the change under test. The published `main-latest` image ships the UI from the last release; to test local UI changes, build the image from your branch and point the compose stack at it: + + ```bash + docker build -t litellm-local . + LITELLM_E2E_IMAGE=litellm-local docker compose up -d + ``` + 4. Tear it down when you're done: ```bash diff --git a/tests/e2e/coverage_registry/mgmt.yaml b/tests/e2e/coverage_registry/mgmt.yaml index 8971a0cb42c..4fbaac0a205 100644 --- a/tests/e2e/coverage_registry/mgmt.yaml +++ b/tests/e2e/coverage_registry/mgmt.yaml @@ -4,6 +4,7 @@ - {id: mgmt.key.generate.happy_path, module: mgmt, tier: P0, surface: ui, assertions: [happy_path], source: "ui_sso.py:420", rationale: "SSO-driven key gen (UI path)"} - {id: mgmt.key.update.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:2462", rationale: "Budget/model changes persist"} - {id: mgmt.key.update.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:2462", rationale: "Non-admin cannot escalate perms"} +- {id: mgmt.key.update.happy_path, module: mgmt, tier: P1, surface: ui, assertions: [happy_path], source: "key_management_endpoints.py:2462", rationale: "Key edit through the dashboard"} - {id: mgmt.key.delete.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3122", rationale: "Deletion revokes future calls"} - {id: mgmt.key.delete.admin_only, module: mgmt, tier: P0, surface: api, assertions: [admin_only], source: "key_management_endpoints.py:3122", rationale: "Non-owner cannot delete"} - {id: mgmt.key.info.persists, module: mgmt, tier: P0, surface: api, assertions: [persists], source: "key_management_endpoints.py:3380", rationale: "Info reflects all writes"} diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml index 2de38ca5c91..cdf5d6cbf6f 100644 --- a/tests/e2e/docker-compose.yml +++ b/tests/e2e/docker-compose.yml @@ -58,6 +58,8 @@ services: environment: LITELLM_MASTER_KEY: sk-1234 DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm + UI_USERNAME: admin + UI_PASSWORD: sk-1234 ports: - "4000:4000" configs: diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 3865804b08f..75bd715a23a 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -21,6 +21,9 @@ CONTROL_PLANE_BASE_URL = os.environ.get( "LITELLM_CONTROL_PLANE_URL", PROXY_BASE_URL ).rstrip("/") +UI_USERNAME = os.environ.get("E2E_UI_USERNAME", "admin") +UI_PASSWORD = os.environ.get("E2E_UI_PASSWORD", MASTER_KEY) + # Writes on the proxy are eventually consistent (e.g. spend rows flush on # proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once. POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) diff --git a/tests/e2e/management/conftest.py b/tests/e2e/management/conftest.py index 1a1a740cc0d..47c1d34baae 100644 --- a/tests/e2e/management/conftest.py +++ b/tests/e2e/management/conftest.py @@ -1,9 +1,23 @@ -"""Management suite client fixture; lifecycle/skip/marker live in the parent conftest.""" +"""Management suite fixtures: the client plus a logged-in dashboard page. + +Lifecycle/skip/marker live in the parent conftest. The browser fixtures drive +the dashboard the proxy serves at /ui, so browser tests exercise exactly what an +end user sees. playwright is an optional dependency loaded behind importorskip +inside the fixture, so the API tests in this suite collect and run without it: + + uv pip install playwright && uv run playwright install chromium +""" + +from typing import TYPE_CHECKING, Iterator import pytest +from e2e_config import PROXY_BASE_URL, UI_PASSWORD, UI_USERNAME from management_client import ManagementClient, build_client +if TYPE_CHECKING: + from playwright.sync_api import Browser, Page + def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line( @@ -15,3 +29,29 @@ def pytest_configure(config: pytest.Config) -> None: @pytest.fixture(scope="session") def client() -> ManagementClient: return build_client() + + +@pytest.fixture(scope="session") +def browser() -> "Iterator[Browser]": + pytest.importorskip("playwright.sync_api", reason="playwright not installed") + from playwright.sync_api import sync_playwright + + with sync_playwright() as playwright: + launched = playwright.chromium.launch() + yield launched + launched.close() + + +@pytest.fixture +def ui_page(browser: "Browser") -> "Iterator[Page]": + context = browser.new_context() + try: + page = context.new_page() + page.goto(f"{PROXY_BASE_URL}/ui/") + page.fill("#username", UI_USERNAME) + page.fill("#password", UI_PASSWORD) + page.click('input[type="submit"]') + page.wait_for_url("**/ui/**") + yield page + finally: + context.close() diff --git a/tests/e2e/management/test_key_models_dropdown_e2e.py b/tests/e2e/management/test_key_models_dropdown_e2e.py new file mode 100644 index 00000000000..f0ba21699e0 --- /dev/null +++ b/tests/e2e/management/test_key_models_dropdown_e2e.py @@ -0,0 +1,161 @@ +"""The dashboard's key create/edit Models dropdown scopes its options to the key's team. + +A teamless key offers All Proxy Models but not the all-team-models sentinel (the +backend expands the latter to the full proxy model list when no team is attached), +and a team key offers all-team-models plus the team's own models but never the +all-proxy-models sentinel, even when the team's model list carries it. The create +cases also walk the full product path: submit the modal with the offered sentinel +and read the persisted key back through /key/info. + +The tests drive gpt-5.5, one of the example models prewired in the proxy config in +tests/e2e/docker-compose.yml; the dropdown wait fails with a pointer there when the +proxy under test does not serve it. +""" + +import pytest + +from e2e_config import PROXY_BASE_URL, unique_marker +from lifecycle import ResourceManager +from management_client import ManagementClient +from models import KeyGenerateBody, TeamNewBody + +pytest.importorskip("playwright.sync_api", reason="playwright not installed") + +from playwright.sync_api import Locator, Page, expect # noqa: E402 # import must follow the importorskip guard above + + +def _form_item(page: Page, label: str) -> Locator: + return page.locator(".ant-form-item").filter(has=page.get_by_text(label, exact=True)).first + + +def _open_dropdown(page: Page, label: str) -> Locator: + _form_item(page, label).locator(".ant-select-selector").first.click() + dropdown = page.locator(".ant-select-dropdown:not(.ant-select-dropdown-hidden)").last + expect(dropdown).to_be_visible() + return dropdown + + +def _models_dropdown_texts(page: Page, must_contain: str) -> list[str]: + dropdown = _open_dropdown(page, "Models") + expect( + dropdown.locator(".ant-select-item-option-content", has_text=must_contain).first, + f"{must_contain!r} never appeared in the Models dropdown; the proxy must serve it " + f"(see the model_list in tests/e2e/docker-compose.yml)", + ).to_be_visible() + return dropdown.locator(".ant-select-item-option-content").all_inner_texts() + + +def _open_create_key_modal(page: Page) -> None: + page.goto(f"{PROXY_BASE_URL}/ui/api-keys/?create=true") + expect(page.locator(".ant-modal").first).to_be_visible() + + +def _select_team(page: Page, alias: str) -> None: + dropdown = _open_dropdown(page, "Team") + dropdown.get_by_text(alias).first.click() + + +def _submit_create_modal(page: Page, sentinel_label: str) -> str: + dropdown = page.locator(".ant-select-dropdown:not(.ant-select-dropdown-hidden)").last + dropdown.locator(".ant-select-item-option-content", has_text=sentinel_label).first.click() + page.keyboard.press("Escape") + _form_item(page, "Key Name").locator("input").first.fill(f"e2e-ui-key-{unique_marker()}") + page.get_by_role("button", name="Create Key", exact=True).click() + + expect(page.get_by_text("Save your Key")).to_be_visible() + key = page.locator(".ant-modal pre").last.inner_text().strip() + assert key.startswith("sk-"), f"expected the created key in the success modal, got {key!r}" + return key + + +def _open_key_edit_form(page: Page, key_alias: str) -> None: + page.goto(f"{PROXY_BASE_URL}/ui/api-keys/") + page.get_by_text(key_alias).first.click() + page.get_by_role("tab", name="Settings").click() + page.get_by_role("button", name="Edit Settings").click() + expect(_form_item(page, "Models")).to_be_visible() + + +def _provision_team(client: ManagementClient, resources: ResourceManager, alias: str) -> str: + team_id = client.create_team(TeamNewBody(team_alias=alias, models=["all-proxy-models", "gpt-5.5"])) + resources.defer(lambda: client.delete_team(team_id)) + return team_id + + +def _provision_key( + client: ManagementClient, resources: ResourceManager, alias: str, team_id: str | None = None +) -> str: + key = client.gateway.generate_key(KeyGenerateBody(key_alias=alias, models=["gpt-5.5"], team_id=team_id)) + resources.defer(lambda: client.gateway.delete_key(key)) + return key + + +@pytest.mark.e2e +class TestKeyModelsDropdownUI: + @pytest.mark.covers("mgmt.key.generate.happy_path", exercised_on=[]) + def test_create_teamless_key_offers_proxy_scope_and_persists( + self, ui_page: Page, client: ManagementClient, resources: ResourceManager + ) -> None: + _open_create_key_modal(ui_page) + + options = _models_dropdown_texts(ui_page, must_contain="gpt-5.5") + assert "All Proxy Models" in options, f"teamless create lost 'All Proxy Models': {options}" + assert "All Team Models" not in options, f"teamless create offered 'All Team Models': {options}" + + key = _submit_create_modal(ui_page, sentinel_label="All Proxy Models") + resources.defer(lambda: client.gateway.delete_key(key)) + + info = client.gateway.key_info(key) + assert info.models == ["all-proxy-models"], f"persisted models {info.models}" + assert info.team_id is None, f"teamless key persisted with team {info.team_id}" + + @pytest.mark.covers("mgmt.key.generate.happy_path", exercised_on=[]) + def test_create_team_key_offers_team_scope_and_persists( + self, ui_page: Page, client: ManagementClient, resources: ResourceManager + ) -> None: + team_alias = f"e2e-ui-team-{unique_marker()}" + team_id = _provision_team(client, resources, team_alias) + + _open_create_key_modal(ui_page) + _select_team(ui_page, team_alias) + + options = _models_dropdown_texts(ui_page, must_contain="All Team Models") + assert "gpt-5.5" in options, f"team key create lost the team's own model: {options}" + assert "All Proxy Models" not in options, f"team key create offered 'All Proxy Models': {options}" + assert "all-proxy-models" not in options, f"team key create offered the raw sentinel: {options}" + + key = _submit_create_modal(ui_page, sentinel_label="All Team Models") + resources.defer(lambda: client.gateway.delete_key(key)) + + info = client.gateway.key_info(key) + assert info.models == ["all-team-models"], f"persisted models {info.models}" + assert info.team_id == team_id, f"persisted team {info.team_id}, expected {team_id}" + + @pytest.mark.covers("mgmt.key.update.happy_path", exercised_on=[]) + def test_edit_teamless_key_offers_proxy_scope( + self, ui_page: Page, client: ManagementClient, resources: ResourceManager + ) -> None: + key_alias = f"e2e-ui-teamless-{unique_marker()}" + _provision_key(client, resources, key_alias) + + _open_key_edit_form(ui_page, key_alias) + + options = _models_dropdown_texts(ui_page, must_contain="gpt-5.5") + assert "All Proxy Models" in options, f"teamless edit lost 'All Proxy Models': {options}" + assert "All Team Models" not in options, f"teamless edit offered 'All Team Models': {options}" + + @pytest.mark.covers("mgmt.key.update.happy_path", exercised_on=[]) + def test_edit_team_key_offers_team_scope_only( + self, ui_page: Page, client: ManagementClient, resources: ResourceManager + ) -> None: + team_alias = f"e2e-ui-team-{unique_marker()}" + team_id = _provision_team(client, resources, team_alias) + key_alias = f"e2e-ui-teamkey-{unique_marker()}" + _provision_key(client, resources, key_alias, team_id=team_id) + + _open_key_edit_form(ui_page, key_alias) + + options = _models_dropdown_texts(ui_page, must_contain="All Team Models") + assert "gpt-5.5" in options, f"team key edit lost the team's own model: {options}" + assert "All Proxy Models" not in options, f"team key edit offered 'All Proxy Models': {options}" + assert "all-proxy-models" not in options, f"team key edit offered the raw sentinel: {options}" diff --git a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts index 644228c5ff9..a55c19a53de 100644 --- a/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts +++ b/ui/litellm-dashboard/e2e_tests/tests/proxy-admin/keys.spec.ts @@ -137,11 +137,10 @@ test.describe("Proxy Admin - Keys", () => { // No team selection — leave team dropdown empty so the key is owned by the admin user // Select models — open the multi-select and pick the all-models meta-option. - // The Create Key modal labels this "All Team Models" even when no team is selected - // (see src/components/organisms/create_key_button.tsx:944), unlike the team/user - // settings screens which use "All Proxy Models". + // With no team selected the modal offers "All Proxy Models"; the team-scoped + // "All Team Models" option only appears once a team is picked. await page.locator(".ant-select-selection-overflow").click(); - await page.locator(".ant-select-dropdown:visible").getByText("All Team Models").click(); + await page.locator(".ant-select-dropdown:visible").getByText("All Proxy Models").click(); await page.keyboard.press("Escape"); await page.getByRole("button", { name: "Create Key", exact: true }).click(); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx index d4bb7b56bf7..3a3226336b2 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/fetch_available_models_team_key.tsx @@ -35,6 +35,12 @@ export const fetchAvailableModelsForTeamOrKey = async ( } }; +export const excludeProxyWideSentinel = (models: string[]): string[] => + models.filter((model) => model !== "all-proxy-models"); + +export const hasAllModelsSentinel = (models: string[]): boolean => + models.includes("all-proxy-models") || models.includes("all-team-models"); + export const getModelDisplayName = (model: string) => { if (model === "all-proxy-models") { return "All Proxy Models"; 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 355e870e848..0fe8adb70e1 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,35 +1,39 @@ -import { act, fireEvent } from "@testing-library/react"; +import { act, fireEvent, within } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; +import { Team } from "../key_team_helpers/key_list"; import CreateKey from "./create_key_button"; -const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall } = vi.hoisted(() => { - const formStateRef = { current: {} as Record }; - const mockKeyCreateCall = vi.fn().mockResolvedValue({ - key: "test-api-key", - soft_budget: null, +const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall, teamDropdownTeamsRef } = + vi.hoisted(() => { + const formStateRef = { current: {} as Record }; + const teamDropdownTeamsRef = { current: [] as Array<{ team_id: string; team_alias: string; models: string[] }> }; + const mockKeyCreateCall = vi.fn().mockResolvedValue({ + key: "test-api-key", + soft_budget: null, + }); + const formMock = { + setFieldsValue: vi.fn((values: Record) => { + Object.assign(formStateRef.current, values); + }), + setFieldValue: vi.fn((name: string, value: any) => { + formStateRef.current[name] = value; + }), + getFieldValue: vi.fn((name: string) => formStateRef.current[name]), + resetFields: vi.fn(() => { + formStateRef.current = {}; + }), + }; + const radioGroupValueRef = { current: null as string | null }; + return { + formMock, + setFieldsValueMock: formMock.setFieldsValue, + radioGroupValueRef, + formStateRef, + mockKeyCreateCall, + teamDropdownTeamsRef, + }; }); - const formMock = { - setFieldsValue: vi.fn((values: Record) => { - Object.assign(formStateRef.current, values); - }), - setFieldValue: vi.fn((name: string, value: any) => { - formStateRef.current[name] = value; - }), - getFieldValue: vi.fn((name: string) => formStateRef.current[name]), - resetFields: vi.fn(() => { - formStateRef.current = {}; - }), - }; - const radioGroupValueRef = { current: null as string | null }; - return { - formMock, - setFieldsValueMock: formMock.setFieldsValue, - radioGroupValueRef, - formStateRef, - mockKeyCreateCall, - }; -}); const defaultAuthorizedState = { accessToken: "test-token", @@ -125,6 +129,8 @@ vi.mock("antd", () => { Form.useForm = () => [formMock]; + Form.useWatch = (name: string) => formStateRef.current[name]; + const Select = ({ children, onChange, @@ -256,11 +262,26 @@ vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ }), })); vi.mock("../common_components/team_dropdown", () => ({ - default: ({ onChange, disabled }: { onChange?: (v: string) => void; disabled?: boolean }) => ( - + onTeamSelect?.(teamDropdownTeamsRef.current.find((team) => team.team_id === e.target.value) ?? null) + } + > - - + {teamDropdownTeamsRef.current.map((team) => ( + + ))} ), })); @@ -269,9 +290,13 @@ vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => nu vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null })); vi.mock("../shared/numerical_input", () => ({ default: () => null })); vi.mock("../vector_store_management/VectorStoreSelector", () => ({ default: () => null })); -vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({ - getModelDisplayName: (model: string) => model, -})); +vi.mock("../key_team_helpers/fetch_available_models_team_key", async () => { + const actual = await vi.importActual("../key_team_helpers/fetch_available_models_team_key"); + return { + ...actual, + getModelDisplayName: (model: string) => model, + }; +}); vi.mock("@/app/(dashboard)/hooks/tags/useTags", () => ({ useTags: vi.fn().mockReturnValue({ @@ -344,6 +369,10 @@ describe("CreateKey", () => { authorizedState = { ...defaultAuthorizedState }; radioGroupValueRef.current = null; formStateRef.current = {}; + teamDropdownTeamsRef.current = [ + { team_id: "team-1", team_alias: "Team One", models: [] }, + { team_id: "team-2", team_alias: "Team Two", models: [] }, + ]; mockKeyCreateCall.mockResolvedValue({ key: "test-api-key", soft_budget: null, @@ -555,6 +584,63 @@ describe("CreateKey", () => { }); }); + describe("models dropdown team gating", () => { + const getModelsSelect = async (): Promise => { + return waitFor(() => { + const element = document.querySelector('select[placeholder="Select models"]'); + expect(element).toBeTruthy(); + return element as HTMLElement; + }); + }; + + it("should offer all-proxy-models but not all-team-models when no team is selected", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + const modelsSelect = await getModelsSelect(); + + await waitFor(() => { + expect(within(modelsSelect).getByText("gpt-4")).toBeInTheDocument(); + }); + + expect(within(modelsSelect).getByText("All Proxy Models")).toBeInTheDocument(); + expect(within(modelsSelect).queryByText("All Team Models")).not.toBeInTheDocument(); + }); + + it("should offer all-team-models but hide all-proxy-models when a team is selected", async () => { + teamDropdownTeamsRef.current = [ + { team_id: "team-1", team_alias: "Team One", models: ["all-proxy-models", "team-model-1"] }, + ]; + + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("team-dropdown")).toBeInTheDocument(); + }); + + act(() => { + fireEvent.change(screen.getByTestId("team-dropdown"), { target: { value: "team-1" } }); + }); + + const modelsSelect = await getModelsSelect(); + + await waitFor(() => { + expect(within(modelsSelect).getByText("team-model-1")).toBeInTheDocument(); + }); + + expect(within(modelsSelect).getByText("All Team Models")).toBeInTheDocument(); + expect(within(modelsSelect).queryByText("All Proxy Models")).not.toBeInTheDocument(); + expect(within(modelsSelect).queryByText("all-proxy-models")).not.toBeInTheDocument(); + }); + }); + describe("tags dropdown", () => { it("should populate tags dropdown with options from useTags hook", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index bf0f0cc3fae..0f371b72efe 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -30,7 +30,11 @@ import ProjectDropdown from "../common_components/ProjectDropdown"; import { CreateUserButton } from "../CreateUserButton"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; -import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { + excludeProxyWideSentinel, + getModelDisplayName, + hasAllModelsSentinel, +} from "../key_team_helpers/fetch_available_models_team_key"; import { Team } from "../key_team_helpers/key_list"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; @@ -203,6 +207,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const [routerSettingsKey, setRouterSettingsKey] = useState(0); const [agentsList, setAgentsList] = useState<{ agent_id: string; agent_name: string }[]>([]); const [selectedAgentId, setSelectedAgentId] = useState(null); + const selectedModels: string[] = Form.useWatch("models", form) ?? []; const handleOk = () => { setIsModalVisible(false); form.resetFields(); @@ -589,7 +594,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp } if (userID && userRole && accessToken) { fetchTeamModels(userID, userRole, accessToken, selectedCreateKeyTeam?.team_id ?? null).then((models) => { - let allModels = Array.from(new Set([...(selectedCreateKeyTeam?.models ?? []), ...models])); + const allModels = excludeProxyWideSentinel( + Array.from(new Set([...(selectedCreateKeyTeam?.models ?? []), ...models])), + ); setModelsToPick(allModels); }); } @@ -948,16 +955,23 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp onChange={(values) => { if (values.includes("all-team-models")) { form.setFieldsValue({ models: ["all-team-models"] }); + } else if (values.includes("all-proxy-models")) { + form.setFieldsValue({ models: ["all-proxy-models"] }); } }} > - {!selectedProjectId && ( + {!selectedProjectId && selectedCreateKeyTeam && ( )} + {!selectedProjectId && !selectedCreateKeyTeam && ( + + )} {modelsToPick.map((model: string) => ( - ))} 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 ba68124beee..f8b8e5b6de3 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 @@ -1,8 +1,9 @@ -import { screen, waitFor } from "@testing-library/react"; +import { fireEvent, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import { KeyResponse } from "../key_team_helpers/key_list"; +import { modelAvailableCall } from "../networking"; import { KeyEditView } from "./key_edit_view"; vi.mock("../networking", async () => { @@ -895,4 +896,220 @@ describe("KeyEditView", () => { }); }); }); + + describe("models dropdown team gating", () => { + const openModelsDropdown = () => { + const modelsFormItem = screen.getByText("Models", { selector: "label" }).closest(".ant-form-item"); + const selector = modelsFormItem?.querySelector(".ant-select-selector"); + expect(selector).toBeTruthy(); + fireEvent.mouseDown(selector as Element); + }; + + it("should offer all-proxy-models but not all-team-models for a teamless key", async () => { + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="test-token" + userID="user-123" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models", { selector: "label" })).toBeInTheDocument(); + }); + + openModelsDropdown(); + + await waitFor(() => { + expect(screen.getAllByText("gpt-4").length).toBeGreaterThan(0); + }); + + expect(screen.getAllByText("All Proxy Models").length).toBeGreaterThan(0); + expect(screen.queryAllByText("All Team Models")).toHaveLength(0); + }); + + it("should offer all-team-models but hide all-proxy-models for a team key", async () => { + const teamKeyData = { ...MOCK_KEY_DATA, team_id: "team-1" }; + const teams = [{ team_id: "team-1", models: ["all-proxy-models", "team-model-1"] }]; + + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="test-token" + userID="user-123" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models", { selector: "label" })).toBeInTheDocument(); + }); + + openModelsDropdown(); + + await waitFor(() => { + expect(screen.getAllByText("team-model-1").length).toBeGreaterThan(0); + }); + + expect(screen.getAllByText("All Team Models").length).toBeGreaterThan(0); + expect(screen.queryAllByText("All Proxy Models")).toHaveLength(0); + expect(screen.queryAllByText("all-proxy-models")).toHaveLength(0); + }); + + it("should not offer all-team-models for a team key whose team has not loaded yet", async () => { + const teamKeyData = { ...MOCK_KEY_DATA, team_id: "team-1" }; + + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="test-token" + userID="user-123" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models", { selector: "label" })).toBeInTheDocument(); + }); + + openModelsDropdown(); + + expect(screen.queryAllByText("All Team Models")).toHaveLength(0); + expect(screen.queryAllByText("All Proxy Models")).toHaveLength(0); + }); + + it("should not duplicate the all-proxy-models option when the teamless model list already carries the sentinel", async () => { + vi.mocked(modelAvailableCall).mockResolvedValueOnce({ + data: [{ id: "all-proxy-models" }, { id: "gpt-4" }], + }); + + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="test-token" + userID="user-123" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models", { selector: "label" })).toBeInTheDocument(); + }); + + openModelsDropdown(); + + const proxyOptionLabels = () => + Array.from(document.querySelectorAll('[role="option"]')).map((option) => option.getAttribute("aria-label")); + + await waitFor(() => { + expect(proxyOptionLabels()).toContain("gpt-4"); + }); + + const labels = proxyOptionLabels(); + expect(labels.filter((label) => label === "All Proxy Models")).toHaveLength(1); + expect(labels).not.toContain("all-proxy-models"); + }); + + it("should collapse the selection to all-proxy-models when the sentinel is picked alongside a model", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken="test-token" + userID="user-123" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models", { selector: "label" })).toBeInTheDocument(); + }); + + openModelsDropdown(); + + const clickOption = async (label: string) => { + const option = await waitFor(() => { + const match = Array.from(document.querySelectorAll(".ant-select-item-option")).find( + (el) => el.querySelector(".ant-select-item-option-content")?.textContent === label, + ); + expect(match).toBeTruthy(); + return match as HTMLElement; + }); + fireEvent.click(option); + }; + + await clickOption("gpt-4"); + await clickOption("All Proxy Models"); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + expect(onSubmitMock.mock.calls[0][0].models).toEqual(["all-proxy-models"]); + }); + + it("should disable the individual model options once all-proxy-models is selected", async () => { + renderWithProviders( + {}} + onSubmit={async () => {}} + accessToken="test-token" + userID="user-123" + userRole="Admin" + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models", { selector: "label" })).toBeInTheDocument(); + }); + + openModelsDropdown(); + + const findOption = (label: string) => + Array.from(document.querySelectorAll(".ant-select-item-option")).find( + (el) => el.querySelector(".ant-select-item-option-content")?.textContent === label, + ) as HTMLElement | undefined; + + const gpt4Before = await waitFor(() => { + const match = findOption("gpt-4"); + expect(match).toBeTruthy(); + return match!; + }); + expect(gpt4Before.classList.contains("ant-select-item-option-disabled")).toBe(false); + + fireEvent.click( + await waitFor(() => { + const match = findOption("All Proxy Models"); + expect(match).toBeTruthy(); + return match!; + }), + ); + + await waitFor(() => { + expect(findOption("gpt-4")?.classList.contains("ant-select-item-option-disabled")).toBe(true); + }); + }); + }); }); 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 4821ea86b87..02bc039cf74 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -18,6 +18,7 @@ import OrganizationDropdown from "../common_components/OrganizationDropdown"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; +import { excludeProxyWideSentinel, hasAllModelsSentinel } from "../key_team_helpers/fetch_available_models_team_key"; import { KeyResponse } from "../key_team_helpers/key_list"; import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import { NO_MCP_SERVERS_SENTINEL } from "../mcp_tools/constants"; @@ -132,11 +133,11 @@ export function KeyEditView({ // Fetch user models if no team const model_available = await modelAvailableCall(accessToken, userID, userRole); const available_model_names = model_available["data"].map((element: { id: string }) => element.id); - setAvailableModels(available_model_names); + setAvailableModels(excludeProxyWideSentinel(available_model_names)); } else if (team?.team_id) { // Fetch team models if team exists const models = await fetchTeamModels(userID, userRole, accessToken, team.team_id); - setAvailableModels(Array.from(new Set([...team.models, ...models]))); + setAvailableModels(excludeProxyWideSentinel(Array.from(new Set([...team.models, ...models])))); } } catch (error) { console.error("Error fetching models:", error); @@ -357,12 +358,23 @@ export function KeyEditView({ style={{ width: "100%" }} disabled={isDisabled} value={isDisabled ? [] : models} - onChange={(value) => setFieldValue("models", value)} + onChange={(value) => { + if (value.includes("all-team-models")) { + setFieldValue("models", ["all-team-models"]); + } else if (value.includes("all-proxy-models")) { + setFieldValue("models", ["all-proxy-models"]); + } else { + setFieldValue("models", value); + } + }} > - {/* Only show All Team Models if team has models */} - {availableModels.length > 0 && All Team Models} + {keyData.team_id != null ? ( + team != null && All Team Models + ) : ( + All Proxy Models + )} {availableModels.map((model) => ( - + {model} ))}