From 79aea5ddcfacd2281a6f94074ca7cd2bb5cd6fdf Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 20 Mar 2026 16:25:40 -0700 Subject: [PATCH] [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}