From a05a1eef94d608ca5226eb176ec1dfadaa963fb4 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 7 Jul 2026 18:54:19 -0700 Subject: [PATCH] fix(ui): scope key models dropdown options to the key's team (#32382) * fix(ui): scope key models dropdown options to the key's team A teamless key no longer offers the all-team-models option in the create and edit forms; the backend expands that sentinel to the full proxy model list when no team is attached, which is rarely what the user intended. A team key no longer surfaces the all-proxy-models sentinel that leaks in verbatim when the team's own model list carries it; the dropdown keeps All Team Models plus the team's individual models. Adds browser coverage to the management e2e suite: playwright (an optional dependency behind importorskip) drives the proxy-served dashboard at /ui, asserts the dropdown options a real user sees for teamless and team keys on both create and edit, and walks the create modal end to end, reading the persisted key back through /key/info. * fix(ui): offer all-proxy-models on teamless keys in the models dropdown A teamless key has no team allowlist to inherit, so the dropdown now offers All Proxy Models in place of All Team Models on both the create and edit forms, with the same exclusive-selection handling. Component and browser e2e tests updated to pin the swapped option pair; the teamless create case now also walks the modal end to end and reads the persisted key back through /key/info. * test(ui): update no-team key creation spec to pick All Proxy Models The create modal no longer offers All Team Models without a team; the teamless path now offers All Proxy Models, which is what this spec exercises * fix(ui): gate All Team Models on the team object being loaded When a key has a team_id but the teams prop does not yet include the matching team, availableModels stays empty and the models dropdown rendered All Team Models on its own with nothing to compare against. Gate the option on the team object being present so it only appears once team models are known, and add a regression test for the loading state Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): filter all-proxy-models from teamless model fetch in key edit form The teamless fetch path stored modelAvailableCall results without excludeProxyWideSentinel, so an all-proxy-models entry in the response rendered a second option colliding with the hardcoded All Proxy Models sentinel. Apply the same filter used on the team path and add a regression test asserting the sentinel option is not duplicated Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Mubashir Osmani Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/CLAUDE.md | 2 +- tests/e2e/CONTRIBUTING.md | 14 ++ tests/e2e/coverage_registry/mgmt.yaml | 1 + tests/e2e/docker-compose.yml | 2 + tests/e2e/e2e_config.py | 3 + tests/e2e/management/conftest.py | 42 +++- .../test_key_models_dropdown_e2e.py | 161 +++++++++++++ .../e2e_tests/tests/proxy-admin/keys.spec.ts | 7 +- .../fetch_available_models_team_key.tsx | 6 + .../organisms/create_key_button.test.tsx | 154 +++++++++--- .../organisms/create_key_button.tsx | 22 +- .../templates/key_edit_view.test.tsx | 219 +++++++++++++++++- .../components/templates/key_edit_view.tsx | 24 +- 13 files changed, 606 insertions(+), 51 deletions(-) create mode 100644 tests/e2e/management/test_key_models_dropdown_e2e.py 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} ))}