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
This commit is contained in:
Ryan Crabbe 2026-04-11 15:19:50 -07:00
parent 88beed905e
commit a9d64f8620
No known key found for this signature in database
4 changed files with 56 additions and 14 deletions

View file

@ -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");
});
});

View file

@ -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<string, unknown> | null;
guardrail_id?: string | null;
[key: string]: unknown;
}
interface GuardrailsListResponse {
guardrails: GuardrailListItem[];
}
export interface GuardrailsListData {
guardrails: GuardrailListItem[];
globalGuardrailNames: Set<string>;
optionalGuardrailNames: Set<string>;
}
// ── Hook ─────────────────────────────────────────────────────────────────────
const guardrailKeys = createQueryKeys("guardrails");
export const useGuardrails = (): UseQueryResult<string[]> => {
export const useGuardrails = (): UseQueryResult<GuardrailsListData> => {
const { accessToken, userId, userRole } = useAuthorized();
return useQuery<string[]>({
return useQuery<GuardrailsListResponse, Error, GuardrailsListData>({
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<string>();
const optionalGuardrailNames = new Set<string>();
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 };
},
});
};

View file

@ -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<string>(),
optionalGuardrailNames: new Set<string>(["test-guardrail"]),
},
isLoading: false,
error: null,
}),

View file

@ -63,7 +63,8 @@ const AddModelForm: React.FC<AddModelFormProps> = ({
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 () => {