From 79aea5ddcfacd2281a6f94074ca7cd2bb5cd6fdf Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 16:25:40 -0700 Subject: [PATCH 1/5] [Refactor] UI - Policies: Extract usePolicyVersions hook from FlowBuilderPage Replace useEffect + useState fetch pattern for policy version management with React Query hooks (useQuery + useMutation), following established codebase conventions. Co-Authored-By: Claude Opus 4.6 --- .../hooks/policies/usePolicyVersions.test.ts | 303 ++++++++++++++++++ .../hooks/policies/usePolicyVersions.ts | 123 +++++++ .../policies/pipeline_flow_builder.tsx | 111 ++----- 3 files changed, 461 insertions(+), 76 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts new file mode 100644 index 00000000000..438246eab07 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts @@ -0,0 +1,303 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + usePolicyVersions, + useCreatePolicyVersion, + useUpdatePolicyVersionStatus, +} from "./usePolicyVersions"; + +// ── Mocks ─────────────────────────────────────────────────────────────────── + +const mockListPolicyVersions = vi.fn(); +const mockCreatePolicyVersion = vi.fn(); +const mockUpdatePolicyVersionStatus = vi.fn(); + +vi.mock("@/components/networking", () => ({ + listPolicyVersions: (...args: unknown[]) => mockListPolicyVersions(...args), + createPolicyVersion: (...args: unknown[]) => mockCreatePolicyVersion(...args), + updatePolicyVersionStatus: (...args: unknown[]) => + mockUpdatePolicyVersionStatus(...args), +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +// Import the mocked module to assert on it +import NotificationsManager from "@/components/molecules/notifications_manager"; + +const mockUseAuthorized = vi.fn(); +vi.mock("../useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// ── Setup ─────────────────────────────────────────────────────────────────── + +describe("usePolicyVersions", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + mutations: { retry: false }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + // ── Query tests ───────────────────────────────────────────────────────── + + it("fetches versions when policyName is provided", async () => { + const mockResponse = { + policy_name: "my-policy", + versions: [ + { policy_id: "v1", policy_name: "my-policy", version_number: 1, version_status: "production" }, + { policy_id: "v2", policy_name: "my-policy", version_number: 2, version_status: "draft" }, + ], + total_count: 2, + }; + mockListPolicyVersions.mockResolvedValue(mockResponse); + + const { result } = renderHook( + () => usePolicyVersions({ policyName: "my-policy" }), + { wrapper } + ); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + expect(mockListPolicyVersions).toHaveBeenCalledWith("test-access-token", "my-policy"); + expect(result.current.data?.versions).toHaveLength(2); + expect(result.current.data?.versions[0].policy_id).toBe("v1"); + }); + + it("does not fetch when policyName is null", () => { + const { result } = renderHook( + () => usePolicyVersions({ policyName: null }), + { wrapper } + ); + + expect(result.current.fetchStatus).toBe("idle"); + expect(mockListPolicyVersions).not.toHaveBeenCalled(); + }); + + it("does not fetch when enabled is false", () => { + const { result } = renderHook( + () => usePolicyVersions({ policyName: "my-policy", enabled: false }), + { wrapper } + ); + + expect(result.current.fetchStatus).toBe("idle"); + expect(mockListPolicyVersions).not.toHaveBeenCalled(); + }); + + it("defaults versions to empty array when response has undefined versions", async () => { + mockListPolicyVersions.mockResolvedValue({ + policy_name: "my-policy", + versions: undefined, + total_count: 0, + }); + + const { result } = renderHook( + () => usePolicyVersions({ policyName: "my-policy" }), + { wrapper } + ); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(result.current.data?.versions).toEqual([]); + }); +}); + +describe("useCreatePolicyVersion", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + mutations: { retry: false }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("calls createPolicyVersion and shows success notification", async () => { + const newPolicy = { policy_id: "v3", policy_name: "my-policy", version_number: 3 }; + mockCreatePolicyVersion.mockResolvedValue(newPolicy); + + const { result } = renderHook( + () => useCreatePolicyVersion("my-policy"), + { wrapper } + ); + + const returned = await result.current.mutateAsync(); + + expect(mockCreatePolicyVersion).toHaveBeenCalledWith("test-access-token", "my-policy"); + expect(returned).toEqual(newPolicy); + expect(NotificationsManager.success).toHaveBeenCalledWith("New draft version created"); + }); + + it("shows error notification on failure", async () => { + mockCreatePolicyVersion.mockRejectedValue(new Error("Server error")); + + const { result } = renderHook( + () => useCreatePolicyVersion("my-policy"), + { wrapper } + ); + + await expect(result.current.mutateAsync()).rejects.toThrow("Server error"); + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + "Failed to create version: Server error" + ); + }); + + it("throws when policyName is null", async () => { + const { result } = renderHook( + () => useCreatePolicyVersion(null), + { wrapper } + ); + + await expect(result.current.mutateAsync()).rejects.toThrow( + "Missing access token or policy name" + ); + }); +}); + +describe("useUpdatePolicyVersionStatus", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { retry: false }, + mutations: { retry: false }, + }, + }); + + vi.clearAllMocks(); + + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("publishes a version and shows success notification", async () => { + const updatedPolicy = { policy_id: "v2", version_status: "published" }; + mockUpdatePolicyVersionStatus.mockResolvedValue(updatedPolicy); + + const { result } = renderHook( + () => useUpdatePolicyVersionStatus("my-policy"), + { wrapper } + ); + + const returned = await result.current.mutateAsync({ + policyId: "v2", + status: "published", + }); + + expect(mockUpdatePolicyVersionStatus).toHaveBeenCalledWith( + "test-access-token", + "v2", + "published" + ); + expect(returned).toEqual(updatedPolicy); + expect(NotificationsManager.success).toHaveBeenCalledWith( + expect.stringContaining("Version published") + ); + }); + + it("promotes to production and shows success notification", async () => { + const updatedPolicy = { policy_id: "v2", version_status: "production" }; + mockUpdatePolicyVersionStatus.mockResolvedValue(updatedPolicy); + + const { result } = renderHook( + () => useUpdatePolicyVersionStatus("my-policy"), + { wrapper } + ); + + await result.current.mutateAsync({ + policyId: "v2", + status: "production", + }); + + expect(mockUpdatePolicyVersionStatus).toHaveBeenCalledWith( + "test-access-token", + "v2", + "production" + ); + expect(NotificationsManager.success).toHaveBeenCalledWith( + "Version promoted to production" + ); + }); + + it("shows error notification on publish failure", async () => { + mockUpdatePolicyVersionStatus.mockRejectedValue(new Error("Forbidden")); + + const { result } = renderHook( + () => useUpdatePolicyVersionStatus("my-policy"), + { wrapper } + ); + + await expect( + result.current.mutateAsync({ policyId: "v2", status: "published" }) + ).rejects.toThrow("Forbidden"); + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + "Failed to publish: Forbidden" + ); + }); + + it("shows error notification on promote failure", async () => { + mockUpdatePolicyVersionStatus.mockRejectedValue(new Error("Not found")); + + const { result } = renderHook( + () => useUpdatePolicyVersionStatus("my-policy"), + { wrapper } + ); + + await expect( + result.current.mutateAsync({ policyId: "v2", status: "production" }) + ).rejects.toThrow("Not found"); + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + "Failed to promote to production: Not found" + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts new file mode 100644 index 00000000000..3eb55acafbe --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts @@ -0,0 +1,123 @@ +import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query"; +import useAuthorized from "../useAuthorized"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { + listPolicyVersions, + createPolicyVersion, + updatePolicyVersionStatus, +} from "@/components/networking"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { Policy } from "@/components/policies/types"; + +// ── Query keys ────────────────────────────────────────────────────────────── + +export const policyVersionKeys = createQueryKeys("policyVersions"); + +// ── Types ─────────────────────────────────────────────────────────────────── + +export interface PolicyVersionsResponse { + policy_name: string; + versions: Policy[]; + total_count: number; +} + +// ── Fetch function ────────────────────────────────────────────────────────── + +const fetchPolicyVersions = async ( + accessToken: string, + policyName: string +): Promise => { + return await listPolicyVersions(accessToken, policyName); +}; + +// ── Hook ──────────────────────────────────────────────────────────────────── + +export interface UsePolicyVersionsOptions { + policyName: string | null | undefined; + enabled?: boolean; +} + +export const usePolicyVersions = ({ + policyName, + enabled = true, +}: UsePolicyVersionsOptions) => { + const { accessToken } = useAuthorized(); + + return useQuery({ + queryKey: policyVersionKeys.detail(policyName ?? ""), + queryFn: async () => await fetchPolicyVersions(accessToken!, policyName!), + enabled: Boolean(accessToken && policyName && enabled), + select: (data) => ({ + ...data, + versions: data.versions ?? [], + }), + }); +}; + +// ── Mutations ─────────────────────────────────────────────────────────────── + +export const useCreatePolicyVersion = (policyName: string | null | undefined) => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async () => { + if (!accessToken || !policyName) { + throw new Error("Missing access token or policy name"); + } + return await createPolicyVersion(accessToken, policyName); + }, + onSuccess: () => { + NotificationsManager.success("New draft version created"); + if (policyName) { + queryClient.invalidateQueries({ + queryKey: policyVersionKeys.detail(policyName), + }); + } + }, + onError: (error) => { + NotificationsManager.fromBackend( + "Failed to create version: " + error.message + ); + }, + }); +}; + +export const useUpdatePolicyVersionStatus = ( + policyName: string | null | undefined +) => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation< + Policy, + Error, + { policyId: string; status: "published" | "production" } + >({ + mutationFn: async ({ policyId, status }) => { + if (!accessToken) { + throw new Error("Missing access token"); + } + return await updatePolicyVersionStatus(accessToken, policyId, status); + }, + onSuccess: (_data, variables) => { + const label = + variables.status === "published" + ? "Version published. You can test it in the Playground by selecting this version in the Policies dropdown." + : "Version promoted to production"; + NotificationsManager.success(label); + if (policyName) { + queryClient.invalidateQueries({ + queryKey: policyVersionKeys.detail(policyName), + }); + } + }, + onError: (error, variables) => { + const action = + variables.status === "published" ? "publish" : "promote to production"; + NotificationsManager.fromBackend( + `Failed to ${action}: ${error.message}` + ); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index 89665774eee..6dc235e459e 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -5,8 +5,13 @@ import { ArrowLeftIcon, PlusIcon } from "@heroicons/react/outline"; import { DotsVerticalIcon } from "@heroicons/react/solid"; import { GuardrailPipeline, PipelineStep, PipelineTestResult, PolicyCreateRequest, PolicyUpdateRequest, Policy } from "./types"; import { Guardrail } from "../guardrails/types"; -import { testPipelineCall, listPolicyVersions, createPolicyVersion, updatePolicyVersionStatus } from "../networking"; +import { testPipelineCall } from "../networking"; import NotificationsManager from "../molecules/notifications_manager"; +import { + usePolicyVersions, + useCreatePolicyVersion, + useUpdatePolicyVersionStatus, +} from "@/app/(dashboard)/hooks/policies/usePolicyVersions"; import { getComplianceDatasetPrompts, getFrameworks, @@ -1288,10 +1293,6 @@ export const FlowBuilderPage: React.FC = ({ const [pipeline, setPipeline] = useState( () => derivePipelineFromPolicy(editingPolicy) ); - const [versions, setVersions] = useState([]); - const [isVersionsLoading, setIsVersionsLoading] = useState(false); - const [isCreatingVersion, setIsCreatingVersion] = useState(false); - const [isUpdatingStatus, setIsUpdatingStatus] = useState(false); // Sync local state when editingPolicy changes (e.g. user switched version) React.useEffect(() => { @@ -1300,45 +1301,23 @@ export const FlowBuilderPage: React.FC = ({ setPipeline(derivePipelineFromPolicy(editingPolicy)); }, [editingPolicy?.policy_id, editingPolicy?.policy_name, editingPolicy?.description, editingPolicy?.pipeline, editingPolicy?.guardrails_add]); - // Fetch versions when editing an existing policy by name - React.useEffect(() => { - if (!showVersionsSidebar || !editingPolicy?.policy_name || !accessToken) { - setVersions([]); - return; - } - let cancelled = false; - setIsVersionsLoading(true); - listPolicyVersions(accessToken, editingPolicy.policy_name) - .then((res) => { - if (!cancelled) setVersions(res.versions || []); - }) - .catch(() => { - if (!cancelled) setVersions([]); - }) - .finally(() => { - if (!cancelled) setIsVersionsLoading(false); - }); - return () => { - cancelled = true; - }; - }, [showVersionsSidebar, editingPolicy?.policy_name, accessToken]); + // ── Version management via React Query hooks ────────────────────────────── + + const { + data: versionsData, + isPending: isVersionsLoading, + } = usePolicyVersions({ + policyName: editingPolicy?.policy_name, + enabled: showVersionsSidebar, + }); + const versions = versionsData?.versions ?? []; + + const createVersionMutation = useCreatePolicyVersion(editingPolicy?.policy_name); + const updateStatusMutation = useUpdatePolicyVersionStatus(editingPolicy?.policy_name); const handleNewVersion = async () => { - if (!accessToken || !editingPolicy?.policy_name) return; - setIsCreatingVersion(true); - try { - const newPolicy = await createPolicyVersion(accessToken, editingPolicy.policy_name); - NotificationsManager.success("New draft version created"); - onVersionCreated?.(newPolicy); - const list = await listPolicyVersions(accessToken, editingPolicy.policy_name); - setVersions(list.versions ?? []); - } catch (error) { - NotificationsManager.fromBackend( - "Failed to create version: " + (error instanceof Error ? error.message : String(error)) - ); - } finally { - setIsCreatingVersion(false); - } + const newPolicy = await createVersionMutation.mutateAsync(); + onVersionCreated?.(newPolicy); }; const handleSelectVersion = (policy: Policy) => { @@ -1346,41 +1325,21 @@ export const FlowBuilderPage: React.FC = ({ }; const handlePublishVersion = async () => { - if (!accessToken || !editingPolicy?.policy_id) return; - setIsUpdatingStatus(true); - try { - const updated = await updatePolicyVersionStatus(accessToken, editingPolicy.policy_id, "published"); - NotificationsManager.success( - "Version published. You can test it in the Playground by selecting this version in the Policies dropdown." - ); - const list = await listPolicyVersions(accessToken, editingPolicy.policy_name ?? ""); - setVersions(list.versions ?? []); - onVersionStatusUpdated?.(updated); - } catch (error) { - NotificationsManager.fromBackend( - "Failed to publish: " + (error instanceof Error ? error.message : String(error)) - ); - } finally { - setIsUpdatingStatus(false); - } + if (!editingPolicy?.policy_id) return; + const updated = await updateStatusMutation.mutateAsync({ + policyId: editingPolicy.policy_id, + status: "published", + }); + onVersionStatusUpdated?.(updated); }; const handlePromoteToProduction = async () => { - if (!accessToken || !editingPolicy?.policy_id) return; - setIsUpdatingStatus(true); - try { - const updated = await updatePolicyVersionStatus(accessToken, editingPolicy.policy_id, "production"); - NotificationsManager.success("Version promoted to production"); - const list = await listPolicyVersions(accessToken, editingPolicy.policy_name ?? ""); - setVersions(list.versions ?? []); - onVersionStatusUpdated?.(updated); - } catch (error) { - NotificationsManager.fromBackend( - "Failed to promote to production: " + (error instanceof Error ? error.message : String(error)) - ); - } finally { - setIsUpdatingStatus(false); - } + if (!editingPolicy?.policy_id) return; + const updated = await updateStatusMutation.mutateAsync({ + policyId: editingPolicy.policy_id, + status: "production", + }); + onVersionStatusUpdated?.(updated); }; const handleSave = async () => { @@ -1540,8 +1499,8 @@ export const FlowBuilderPage: React.FC = ({ accessToken={accessToken} versions={versions} isLoading={isVersionsLoading} - isCreatingVersion={isCreatingVersion} - isUpdatingStatus={isUpdatingStatus} + isCreatingVersion={createVersionMutation.isPending} + isUpdatingStatus={updateStatusMutation.isPending} onNewVersion={handleNewVersion} onSelectVersion={handleSelectVersion} onPublish={handlePublishVersion} From 8aa1ebfb07ca82da4ef38debe2d0739d491a1b6d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 23:04:25 -0700 Subject: [PATCH 2/5] [Fix] UI - Policies: Fix unhandled promise rejections and isPending semantics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Wrap mutateAsync calls in try/catch to swallow re-thrown errors (notifications already handled by onError in mutation hooks) - Use isLoading instead of isPending for version loading state — isPending is true when query is disabled with no cache, isLoading is only true during active fetches (matches original behavior) - Add isLoading assertions to disabled-state tests Co-Authored-By: Claude Opus 4.6 --- .../hooks/policies/usePolicyVersions.test.ts | 4 ++ .../policies/pipeline_flow_builder.tsx | 38 ++++++++++++------- 2 files changed, 29 insertions(+), 13 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts index 438246eab07..f773614388f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts @@ -96,6 +96,8 @@ describe("usePolicyVersions", () => { ); expect(result.current.fetchStatus).toBe("idle"); + // isLoading (not isPending) must be false when query is disabled + expect(result.current.isLoading).toBe(false); expect(mockListPolicyVersions).not.toHaveBeenCalled(); }); @@ -106,6 +108,8 @@ describe("usePolicyVersions", () => { ); expect(result.current.fetchStatus).toBe("idle"); + // isLoading (not isPending) must be false when query is disabled + expect(result.current.isLoading).toBe(false); expect(mockListPolicyVersions).not.toHaveBeenCalled(); }); diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index 6dc235e459e..6984948271e 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -1305,7 +1305,7 @@ export const FlowBuilderPage: React.FC = ({ const { data: versionsData, - isPending: isVersionsLoading, + isLoading: isVersionsLoading, } = usePolicyVersions({ policyName: editingPolicy?.policy_name, enabled: showVersionsSidebar, @@ -1316,8 +1316,12 @@ export const FlowBuilderPage: React.FC = ({ const updateStatusMutation = useUpdatePolicyVersionStatus(editingPolicy?.policy_name); const handleNewVersion = async () => { - const newPolicy = await createVersionMutation.mutateAsync(); - onVersionCreated?.(newPolicy); + try { + const newPolicy = await createVersionMutation.mutateAsync(); + onVersionCreated?.(newPolicy); + } catch { + // Notification already shown by onError in the mutation hook + } }; const handleSelectVersion = (policy: Policy) => { @@ -1326,20 +1330,28 @@ export const FlowBuilderPage: React.FC = ({ const handlePublishVersion = async () => { if (!editingPolicy?.policy_id) return; - const updated = await updateStatusMutation.mutateAsync({ - policyId: editingPolicy.policy_id, - status: "published", - }); - onVersionStatusUpdated?.(updated); + try { + const updated = await updateStatusMutation.mutateAsync({ + policyId: editingPolicy.policy_id, + status: "published", + }); + onVersionStatusUpdated?.(updated); + } catch { + // Notification already shown by onError in the mutation hook + } }; const handlePromoteToProduction = async () => { if (!editingPolicy?.policy_id) return; - const updated = await updateStatusMutation.mutateAsync({ - policyId: editingPolicy.policy_id, - status: "production", - }); - onVersionStatusUpdated?.(updated); + try { + const updated = await updateStatusMutation.mutateAsync({ + policyId: editingPolicy.policy_id, + status: "production", + }); + onVersionStatusUpdated?.(updated); + } catch { + // Notification already shown by onError in the mutation hook + } }; const handleSave = async () => { From 80e55804af0fadb92b2c71a1a1fc907138429bc2 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 23:27:02 -0700 Subject: [PATCH 3/5] [Fix] UI - Policies: Type accuracy and mutation guard for usePolicyVersions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Make PolicyVersionsResponse.versions optional (Policy[] | undefined) to match real API shape — select fallback handles normalization - Add policyName guard to useUpdatePolicyVersionStatus mutationFn to fail loudly instead of silently skipping cache invalidation - Add test for null policyName in updateStatus mutation Co-Authored-By: Claude Opus 4.6 --- .../hooks/policies/usePolicyVersions.test.ts | 11 +++++++++++ .../(dashboard)/hooks/policies/usePolicyVersions.ts | 6 +++--- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts index f773614388f..beb40742269 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts @@ -304,4 +304,15 @@ describe("useUpdatePolicyVersionStatus", () => { "Failed to promote to production: Not found" ); }); + + it("throws when policyName is null", async () => { + const { result } = renderHook( + () => useUpdatePolicyVersionStatus(null), + { wrapper } + ); + + await expect( + result.current.mutateAsync({ policyId: "v2", status: "published" }) + ).rejects.toThrow("Missing access token or policy name"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts index 3eb55acafbe..9dc78180cdc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts @@ -17,7 +17,7 @@ export const policyVersionKeys = createQueryKeys("policyVersions"); export interface PolicyVersionsResponse { policy_name: string; - versions: Policy[]; + versions: Policy[] | undefined; total_count: number; } @@ -95,8 +95,8 @@ export const useUpdatePolicyVersionStatus = ( { policyId: string; status: "published" | "production" } >({ mutationFn: async ({ policyId, status }) => { - if (!accessToken) { - throw new Error("Missing access token"); + if (!accessToken || !policyName) { + throw new Error("Missing access token or policy name"); } return await updatePolicyVersionStatus(accessToken, policyId, status); }, From 41d12ed10690a4316f3e6cc50fda3ce5e0338501 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 23:33:50 -0700 Subject: [PATCH 4/5] address greptile review feedback (greploop iteration 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add PolicyVersionsData type for select output; specify TData generic so consumers get Policy[] (not Policy[] | undefined) for versions - Remove empty-string queryKey fallback — use policyName! since enabled:false prevents fetch when policyName is null - Add cache invalidation tests for both mutation hooks - Add explanatory comment for ?? [] fallback in component Co-Authored-By: Claude Opus 4.6 --- .../hooks/policies/usePolicyVersions.test.ts | 32 +++++++++++++++++++ .../hooks/policies/usePolicyVersions.ts | 11 +++++-- .../policies/pipeline_flow_builder.tsx | 2 +- 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts index beb40742269..ba5792644c0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.test.ts @@ -172,6 +172,22 @@ describe("useCreatePolicyVersion", () => { expect(NotificationsManager.success).toHaveBeenCalledWith("New draft version created"); }); + it("invalidates the versions cache on success", async () => { + mockCreatePolicyVersion.mockResolvedValue({ policy_id: "v3" }); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + const { result } = renderHook( + () => useCreatePolicyVersion("my-policy"), + { wrapper } + ); + + await result.current.mutateAsync(); + + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: ["policyVersions", "detail", "my-policy"], + }); + }); + it("shows error notification on failure", async () => { mockCreatePolicyVersion.mockRejectedValue(new Error("Server error")); @@ -249,6 +265,22 @@ describe("useUpdatePolicyVersionStatus", () => { ); }); + it("invalidates the versions cache on success", async () => { + mockUpdatePolicyVersionStatus.mockResolvedValue({ policy_id: "v2" }); + const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); + + const { result } = renderHook( + () => useUpdatePolicyVersionStatus("my-policy"), + { wrapper } + ); + + await result.current.mutateAsync({ policyId: "v2", status: "published" }); + + expect(invalidateSpy).toHaveBeenCalledWith({ + queryKey: ["policyVersions", "detail", "my-policy"], + }); + }); + it("promotes to production and shows success notification", async () => { const updatedPolicy = { policy_id: "v2", version_status: "production" }; mockUpdatePolicyVersionStatus.mockResolvedValue(updatedPolicy); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts index 9dc78180cdc..95f8fa54eed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts @@ -21,6 +21,13 @@ export interface PolicyVersionsResponse { total_count: number; } +/** Output type after `select` normalizes the response — versions is always defined. */ +export interface PolicyVersionsData { + policy_name: string; + versions: Policy[]; + total_count: number; +} + // ── Fetch function ────────────────────────────────────────────────────────── const fetchPolicyVersions = async ( @@ -43,8 +50,8 @@ export const usePolicyVersions = ({ }: UsePolicyVersionsOptions) => { const { accessToken } = useAuthorized(); - return useQuery({ - queryKey: policyVersionKeys.detail(policyName ?? ""), + return useQuery({ + queryKey: policyVersionKeys.detail(policyName!), queryFn: async () => await fetchPolicyVersions(accessToken!, policyName!), enabled: Boolean(accessToken && policyName && enabled), select: (data) => ({ diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index 6984948271e..3e2d55207be 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -1310,7 +1310,7 @@ export const FlowBuilderPage: React.FC = ({ policyName: editingPolicy?.policy_name, enabled: showVersionsSidebar, }); - const versions = versionsData?.versions ?? []; + const versions = versionsData?.versions ?? []; // versionsData?.versions is Policy[] after select, fallback covers undefined data const createVersionMutation = useCreatePolicyVersion(editingPolicy?.policy_name); const updateStatusMutation = useUpdatePolicyVersionStatus(editingPolicy?.policy_name); From 492c0cd3ba6c056af5ab68096cf19639915438f3 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 23:39:01 -0700 Subject: [PATCH 5/5] address greptile review feedback (greploop iteration 3) - Move callbacks outside try/catch so only mutation errors are caught, not errors from onVersionCreated/onVersionStatusUpdated callbacks - Replace policyName! non-null assertion with DISABLED_POLICY_KEY sentinel to avoid undefined in cache keys when query is disabled Co-Authored-By: Claude Opus 4.6 --- .../hooks/policies/usePolicyVersions.ts | 8 +++++-- .../policies/pipeline_flow_builder.tsx | 21 +++++++++++-------- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts index 95f8fa54eed..1bcf3a1e856 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/policies/usePolicyVersions.ts @@ -44,16 +44,20 @@ export interface UsePolicyVersionsOptions { enabled?: boolean; } +/** Stable key used when the query is disabled to avoid undefined in cache keys. */ +const DISABLED_POLICY_KEY = "__disabled__"; + export const usePolicyVersions = ({ policyName, enabled = true, }: UsePolicyVersionsOptions) => { const { accessToken } = useAuthorized(); + const isEnabled = Boolean(accessToken && policyName && enabled); return useQuery({ - queryKey: policyVersionKeys.detail(policyName!), + queryKey: policyVersionKeys.detail(policyName ?? DISABLED_POLICY_KEY), queryFn: async () => await fetchPolicyVersions(accessToken!, policyName!), - enabled: Boolean(accessToken && policyName && enabled), + enabled: isEnabled, select: (data) => ({ ...data, versions: data.versions ?? [], diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index 3e2d55207be..7d1802749c6 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -1316,12 +1316,13 @@ export const FlowBuilderPage: React.FC = ({ const updateStatusMutation = useUpdatePolicyVersionStatus(editingPolicy?.policy_name); const handleNewVersion = async () => { + let newPolicy: Policy; try { - const newPolicy = await createVersionMutation.mutateAsync(); - onVersionCreated?.(newPolicy); + newPolicy = await createVersionMutation.mutateAsync(); } catch { - // Notification already shown by onError in the mutation hook + return; // Notification already shown by onError in the mutation hook } + onVersionCreated?.(newPolicy); }; const handleSelectVersion = (policy: Policy) => { @@ -1330,28 +1331,30 @@ export const FlowBuilderPage: React.FC = ({ const handlePublishVersion = async () => { if (!editingPolicy?.policy_id) return; + let updated: Policy; try { - const updated = await updateStatusMutation.mutateAsync({ + updated = await updateStatusMutation.mutateAsync({ policyId: editingPolicy.policy_id, status: "published", }); - onVersionStatusUpdated?.(updated); } catch { - // Notification already shown by onError in the mutation hook + return; // Notification already shown by onError in the mutation hook } + onVersionStatusUpdated?.(updated); }; const handlePromoteToProduction = async () => { if (!editingPolicy?.policy_id) return; + let updated: Policy; try { - const updated = await updateStatusMutation.mutateAsync({ + updated = await updateStatusMutation.mutateAsync({ policyId: editingPolicy.policy_id, status: "production", }); - onVersionStatusUpdated?.(updated); } catch { - // Notification already shown by onError in the mutation hook + return; // Notification already shown by onError in the mutation hook } + onVersionStatusUpdated?.(updated); }; const handleSave = async () => {