From a2653bcd5e6f368b652e17fbaafe5e6df5b7d356 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 3 Feb 2026 15:26:51 -0800 Subject: [PATCH] Adding Allowed Routes to Key Info and Edit Pages --- .../templates/key_edit_view.test.tsx | 353 ++++++++++++++++-- .../components/templates/key_edit_view.tsx | 67 +++- .../templates/key_info_view.test.tsx | 174 ++++++++- .../components/templates/key_info_view.tsx | 65 ++-- 4 files changed, 575 insertions(+), 84 deletions(-) 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 03e1085937..0123e22eb1 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,24 +1,61 @@ -import { fireEvent, waitFor } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; +import { 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 { KeyEditView } from "./key_edit_view"; -// Mock window.matchMedia -Object.defineProperty(window, "matchMedia", { - writable: true, - value: vi.fn().mockImplementation((query) => ({ - matches: false, - media: query, - onchange: null, - addListener: vi.fn(), - removeListener: vi.fn(), - addEventListener: vi.fn(), - removeEventListener: vi.fn(), - dispatchEvent: vi.fn(), - })), +vi.mock("../networking", async () => { + const actual = await vi.importActual("../networking"); + return { + ...actual, + getPromptsList: vi.fn().mockResolvedValue({ + prompts: [{ prompt_id: "prompt-1" }, { prompt_id: "prompt-2" }], + }), + modelAvailableCall: vi.fn().mockResolvedValue({ + data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }], + }), + tagListCall: vi.fn().mockResolvedValue({ + tag1: { name: "tag1", description: "Test tag 1" }, + tag2: { name: "tag2", description: "Test tag 2" }, + }), + getGuardrailsList: vi.fn().mockResolvedValue({ + guardrails: [{ guardrail_name: "guardrail-1" }], + }), + getPoliciesList: vi.fn().mockResolvedValue({ + policies: [{ policy_name: "policy-1" }], + }), + getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ + endpoints: [], + }), + vectorStoreListCall: vi.fn().mockResolvedValue({ + data: [], + }), + mcpToolsCall: vi.fn().mockResolvedValue({ + data: [], + }), + agentListCall: vi.fn().mockResolvedValue({ + data: [], + }), + fetchMCPServers: vi.fn().mockResolvedValue([]), + fetchMCPAccessGroups: vi.fn().mockResolvedValue([]), + listMCPTools: vi.fn().mockResolvedValue({ + tools: [], + error: null, + message: null, + stack_trace: null, + }), + getAgentsList: vi.fn().mockResolvedValue({ + agents: [], + }), + getAgentAccessGroups: vi.fn().mockResolvedValue([]), + }; }); +vi.mock("../organisms/create_key_button", () => ({ + fetchTeamModels: vi.fn().mockResolvedValue(["team-model-1", "team-model-2"]), +})); + describe("KeyEditView", () => { const MOCK_KEY_DATA: KeyResponse = { token: "test-token-123", @@ -93,8 +130,8 @@ describe("KeyEditView", () => { const { getByText } = renderWithProviders( {}} - onSubmit={async () => {}} + onCancel={() => { }} + onSubmit={async () => { }} accessToken={""} userID={""} userRole={""} @@ -111,8 +148,8 @@ describe("KeyEditView", () => { const { getByText } = renderWithProviders( {}} - onSubmit={async () => {}} + onCancel={() => { }} + onSubmit={async () => { }} accessToken={""} userID={""} userRole={""} @@ -129,8 +166,8 @@ describe("KeyEditView", () => { const { getByLabelText } = renderWithProviders( {}} - onSubmit={async () => {}} + onCancel={() => { }} + onSubmit={async () => { }} accessToken={""} userID={""} userRole={""} @@ -144,13 +181,17 @@ describe("KeyEditView", () => { }); }); + beforeEach(() => { + vi.clearAllMocks(); + }); + it("should call onCancel when cancel button is clicked", async () => { const onCancelMock = vi.fn(); - const { getByText } = renderWithProviders( + renderWithProviders( {}} + onSubmit={async () => { }} accessToken={""} userID={""} userRole={""} @@ -159,12 +200,272 @@ describe("KeyEditView", () => { ); await waitFor(() => { - expect(getByText("Cancel")).toBeInTheDocument(); + expect(screen.getByRole("button", { name: /cancel/i })).toBeInTheDocument(); }); - const cancelButton = getByText("Cancel"); - fireEvent.click(cancelButton); + const cancelButton = screen.getByRole("button", { name: /cancel/i }); + await userEvent.click(cancelButton); expect(onCancelMock).toHaveBeenCalledTimes(1); }); + + it("should display key alias input field", async () => { + renderWithProviders( + { }} + onSubmit={async () => { }} + accessToken={""} + userID={""} + userRole={""} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Key Alias")).toBeInTheDocument(); + }); + }); + + it("should display models select field", async () => { + renderWithProviders( + { }} + onSubmit={async () => { }} + accessToken={""} + userID={""} + userRole={""} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models")).toBeInTheDocument(); + }); + }); + + it("should display max budget input field", async () => { + renderWithProviders( + { }} + onSubmit={async () => { }} + accessToken={""} + userID={""} + userRole={""} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Max Budget (USD)")).toBeInTheDocument(); + }); + }); + + it("should display allowed routes input field", async () => { + renderWithProviders( + { }} + onSubmit={async () => { }} + accessToken={""} + userID={""} + userRole={""} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText(/allowed routes/i)).toBeInTheDocument(); + }); + }); + + it("should call onSubmit with form values when form is submitted", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + { }} + onSubmit={onSubmitMock} + accessToken={"test-token"} + userID={"test-user"} + userRole={"admin"} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + }); + + const submitButton = screen.getByRole("button", { name: /save changes/i }); + await userEvent.click(submitButton); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + }); + + it("should disable models field when management routes are selected", async () => { + const keyDataWithManagementRoutes = { + ...MOCK_KEY_DATA, + allowed_routes: ["management_routes"], + }; + + renderWithProviders( + { }} + onSubmit={async () => { }} + accessToken={""} + userID={""} + userRole={""} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models field is disabled for this key type")).toBeInTheDocument(); + }); + }); + + it("should disable models field when info routes are selected", async () => { + const keyDataWithInfoRoutes = { + ...MOCK_KEY_DATA, + allowed_routes: ["info_routes"], + }; + + renderWithProviders( + { }} + onSubmit={async () => { }} + accessToken={""} + userID={""} + userRole={""} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Models field is disabled for this key type")).toBeInTheDocument(); + }); + }); + + it("should disable guardrails selector when user is not premium", async () => { + renderWithProviders( + { }} + onSubmit={async () => { }} + accessToken={"test-token"} + userID={""} + userRole={""} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Guardrails")).toBeInTheDocument(); + }); + }); + + it("should parse comma-separated allowed routes on submit", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + { }} + onSubmit={onSubmitMock} + accessToken={"test-token"} + userID={"test-user"} + userRole={"admin"} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText(/allowed routes/i)).toBeInTheDocument(); + }); + + const allowedRoutesInput = screen.getByLabelText(/allowed routes/i); + await userEvent.clear(allowedRoutesInput); + await userEvent.type(allowedRoutesInput, "route1, route2, route3"); + + 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(Array.isArray(callArgs.allowed_routes)).toBe(true); + expect(callArgs.allowed_routes).toEqual(["route1", "route2", "route3"]); + }); + }); + + it("should handle empty allowed routes string on submit", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderWithProviders( + { }} + onSubmit={onSubmitMock} + accessToken={"test-token"} + userID={"test-user"} + userRole={"admin"} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByLabelText(/allowed routes/i)).toBeInTheDocument(); + }); + + const allowedRoutesInput = screen.getByLabelText(/allowed routes/i); + await userEvent.clear(allowedRoutesInput); + + 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.allowed_routes).toEqual([]); + }); + }); + + + it("should disable cancel button during submission", async () => { + const onSubmitMock = vi.fn( + () => + new Promise((resolve) => { + setTimeout(resolve, 100); + }), + ); + + renderWithProviders( + { }} + onSubmit={onSubmitMock} + accessToken={"test-token"} + userID={"test-user"} + userRole={"admin"} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /cancel/i })).toBeInTheDocument(); + }); + + const submitButton = screen.getByRole("button", { name: /save changes/i }); + await userEvent.click(submitButton); + + await waitFor(() => { + const cancelButton = screen.getByRole("button", { name: /cancel/i }); + expect(cancelButton).toBeDisabled(); + }); + }); }); 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 2d3eaf2bb7..64676e3b94 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -35,7 +35,6 @@ interface KeyEditViewProps { // Add this helper function const getAvailableModelsForKey = (keyData: KeyResponse, teams: any[] | null): string[] => { // If no teams data is available, return empty array - console.log("getAvailableModelsForKey:", teams); if (!teams || !keyData.team_id) { return []; } @@ -172,7 +171,9 @@ export function KeyEditView({ : [], auto_rotate: keyData.auto_rotate || false, ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }), - allowed_routes: keyData.allowed_routes, + allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 + ? keyData.allowed_routes.join(", ") + : "", }; useEffect(() => { @@ -197,7 +198,9 @@ export function KeyEditView({ : [], auto_rotate: keyData.auto_rotate || false, ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }), - allowed_routes: keyData.allowed_routes, + allowed_routes: Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 + ? keyData.allowed_routes.join(", ") + : "", }); }, [keyData, form]); @@ -226,11 +229,24 @@ export function KeyEditView({ fetchTags(); }, [accessToken]); - console.log("premiumUser:", premiumUser); - const handleSubmit = async (values: any) => { try { setIsKeySaving(true); + + // Parse allowed_routes from comma-separated string to array + if (typeof values.allowed_routes === "string") { + const trimmedInput = values.allowed_routes.trim(); + if (trimmedInput === "") { + values.allowed_routes = []; + } else { + values.allowed_routes = trimmedInput + .split(",") + .map((route: string) => route.trim()) + .filter((route: string) => route.length > 0); + } + } + // If it's already an array (shouldn't happen, but handle it), keep as is + await onSubmit(values); } finally { setIsKeySaving(false); @@ -251,7 +267,11 @@ export function KeyEditView({ } > {({ getFieldValue, setFieldValue }) => { - const allowedRoutes = getFieldValue("allowed_routes") || []; + const allowedRoutesValue = getFieldValue("allowed_routes") || ""; + // Convert string to array for checking + const allowedRoutes = typeof allowedRoutesValue === "string" && allowedRoutesValue.trim() !== "" + ? allowedRoutesValue.split(",").map((r: string) => r.trim()).filter((r: string) => r.length > 0) + : []; const isDisabled = allowedRoutes.includes("management_routes") || allowedRoutes.includes("info_routes"); const models = getFieldValue("models") || []; @@ -290,7 +310,11 @@ export function KeyEditView({ shouldUpdate={(prevValues, currentValues) => prevValues.allowed_routes !== currentValues.allowed_routes} > {({ getFieldValue, setFieldValue }) => { - const allowedRoutes = getFieldValue("allowed_routes"); + const allowedRoutesValue = getFieldValue("allowed_routes") || ""; + // Convert string to array for getKeyTypeFromRoutes + const allowedRoutes = typeof allowedRoutesValue === "string" && allowedRoutesValue.trim() !== "" + ? allowedRoutesValue.split(",").map((r: string) => r.trim()).filter((r: string) => r.length > 0) + : []; const keyTypeValue = getKeyTypeFromRoutes(allowedRoutes); return ( @@ -302,13 +326,13 @@ export function KeyEditView({ onChange={(value) => { switch (value) { case "default": - setFieldValue("allowed_routes", []); + setFieldValue("allowed_routes", ""); break; case "llm_api": - setFieldValue("allowed_routes", ["llm_api_routes"]); + setFieldValue("allowed_routes", "llm_api_routes"); break; case "management": - setFieldValue("allowed_routes", ["management_routes"]); + setFieldValue("allowed_routes", "management_routes"); setFieldValue("models", []); break; } @@ -344,6 +368,22 @@ export function KeyEditView({ + + Allowed Routes{" "} + + + + + } + name="allowed_routes" + > + + + @@ -473,7 +513,7 @@ export function KeyEditView({ !premiumUser ? "Premium feature - Upgrade to set allowed pass through routes by key" : Array.isArray(keyData.metadata?.allowed_passthrough_routes) && - keyData.metadata.allowed_passthrough_routes.length > 0 + keyData.metadata.allowed_passthrough_routes.length > 0 ? `Current: ${keyData.metadata.allowed_passthrough_routes.join(", ")}` : "Select or enter allowed pass through routes" } @@ -590,11 +630,6 @@ export function KeyEditView({ - {/* Hidden form field for allowed_routes */} - - {/* Hidden form field for disabled callbacks */}