From a9d64f8620d76e3ce0e7391c300bd0a37f92575f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Sat, 11 Apr 2026 15:19:50 -0700 Subject: [PATCH] refactor(ui/guardrails): expose full data from useGuardrails hook Extend useGuardrails to return the full guardrail objects plus derived globalGuardrailNames / optionalGuardrailNames sets via React Query's select option, instead of just an array of names. Update its existing consumer (AddModelForm) to extract names from the new shape. The previous shape was tailored to AddModelForm's single use case (populate a Select with names). The team info per-guardrail opt-out work needs default_on per guardrail to split globals from non-globals, which the old shape couldn't provide. Consolidating into the existing hook gives both consumers one source of truth and one React Query cache entry instead of two parallel fetches. - useGuardrails.ts: rewrite return type, derive global/optional sets in select(); preserve the existing query key and auth-gate semantics - AddModelForm.tsx: extract names from data?.guardrails.map(...) - AddModelForm.test.tsx: update mock to return the new shape (also fixes a pre-existing shape mismatch in the mock) - useGuardrails.test.ts: update 3 assertions to read names via data?.guardrails.map(...) instead of asserting against the flat array --- .../hooks/guardrails/useGuardrails.test.ts | 13 ++--- .../hooks/guardrails/useGuardrails.ts | 48 ++++++++++++++++--- .../add_model/AddModelForm.test.tsx | 6 ++- .../src/components/add_model/AddModelForm.tsx | 3 +- 4 files changed, 56 insertions(+), 14 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts index d9e96a5308c..d5b788ab6a5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts @@ -74,7 +74,7 @@ describe("useGuardrails", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual(expectedGuardrailNames); + expect(result.current.data?.guardrails.map((g) => g.guardrail_name)).toEqual(expectedGuardrailNames); expect(result.current.error).toBeNull(); expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); expect(getGuardrailsList).toHaveBeenCalledTimes(1); @@ -228,7 +228,7 @@ describe("useGuardrails", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual([]); + expect(result.current.data?.guardrails).toEqual([]); expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); }); @@ -265,9 +265,10 @@ describe("useGuardrails", () => { expect(result.current.isSuccess).toBe(true); }); - expect(result.current.data).toEqual(expectedNames); - expect(result.current.data).toHaveLength(2); - expect(result.current.data).toContain("custom-guardrail-1"); - expect(result.current.data).toContain("custom-guardrail-2"); + const names = result.current.data?.guardrails.map((g) => g.guardrail_name); + expect(names).toEqual(expectedNames); + expect(names).toHaveLength(2); + expect(names).toContain("custom-guardrail-1"); + expect(names).toContain("custom-guardrail-2"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts index 9786b7fa359..5c7a8df050b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts @@ -3,16 +3,52 @@ import { createQueryKeys } from "../common/queryKeysFactory"; import { getGuardrailsList } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface GuardrailListItem { + guardrail_name: string; + litellm_params?: { + default_on?: boolean; + mode?: string | string[]; + [key: string]: unknown; + }; + guardrail_info?: Record | null; + guardrail_id?: string | null; + [key: string]: unknown; +} + +interface GuardrailsListResponse { + guardrails: GuardrailListItem[]; +} + +export interface GuardrailsListData { + guardrails: GuardrailListItem[]; + globalGuardrailNames: Set; + optionalGuardrailNames: Set; +} + +// ── Hook ───────────────────────────────────────────────────────────────────── + const guardrailKeys = createQueryKeys("guardrails"); -export const useGuardrails = (): UseQueryResult => { +export const useGuardrails = (): UseQueryResult => { const { accessToken, userId, userRole } = useAuthorized(); - return useQuery({ + return useQuery({ queryKey: guardrailKeys.list({}), - queryFn: async () => { - const response = await getGuardrailsList(accessToken!); - return response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name); - }, + queryFn: async () => getGuardrailsList(accessToken!), enabled: Boolean(accessToken && userId && userRole), + select: (data) => { + const guardrails: GuardrailListItem[] = data?.guardrails ?? []; + const globalGuardrailNames = new Set(); + const optionalGuardrailNames = new Set(); + for (const g of guardrails) { + if (g.litellm_params?.default_on) { + globalGuardrailNames.add(g.guardrail_name); + } else { + optionalGuardrailNames.add(g.guardrail_name); + } + } + return { guardrails, globalGuardrailNames, optionalGuardrailNames }; + }, }); }; diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx index de5e1387396..4a3dbaedf74 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx @@ -82,7 +82,11 @@ vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ vi.mock("@/app/(dashboard)/hooks/guardrails/useGuardrails", () => ({ useGuardrails: vi.fn().mockReturnValue({ - data: [{ guardrail_name: "test-guardrail" }], + data: { + guardrails: [{ guardrail_name: "test-guardrail" }], + globalGuardrailNames: new Set(), + optionalGuardrailNames: new Set(["test-guardrail"]), + }, isLoading: false, error: null, }), diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 65e239d58ea..cb35a6f4142 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -63,7 +63,8 @@ const AddModelForm: React.FC = ({ isLoading: isProviderMetadataLoading, error: providerMetadataError, } = useProviderFields(); - const { data: guardrailsList, isLoading: isGuardrailsLoading, error: guardrailsError } = useGuardrails(); + const { data: guardrailsData, isLoading: isGuardrailsLoading, error: guardrailsError } = useGuardrails(); + const guardrailsList = guardrailsData?.guardrails.map((g) => g.guardrail_name); const { data: tagsList, isLoading: isTagsLoading, error: tagsError } = useTags(); const handleTestConnection = async () => {