From fe8cc100de83eb8b8064eca79c0e14cd6950812f Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 2 Apr 2026 17:51:17 -0700 Subject: [PATCH] feat: wire submit guardrail form to POST /guardrails/register Backend: - Add required team_id to RegisterGuardrailRequest body - Validate team membership before allowing submission Frontend: - Add useRegisterGuardrail React Query mutation hook - Wire form onFinish to mutation, pass team_id from dropdown - Show notification and refresh submissions on success --- .../proxy/guardrails/guardrail_endpoints.py | 25 ++++++- .../hooks/guardrails/useRegisterGuardrail.ts | 74 +++++++++++++++++++ .../guardrails/TeamGuardrailsTab.tsx | 27 ++++--- 3 files changed, 113 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 2b20876ba2..2b10537d8e 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -542,6 +542,7 @@ class RegisterGuardrailRequest(BaseModel): str, Any ] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: Optional[Dict[str, Any]] = None + team_id: str def get_litellm_params_dict(self) -> Dict[str, Any]: return dict(self.litellm_params) @@ -603,11 +604,29 @@ async def register_guardrail( if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") - if not user_api_key_dict.team_id: + if not request.team_id: raise HTTPException( status_code=400, - detail="Registration requires an API key associated with a team. Use a team-scoped key.", + detail="team_id is required.", ) + team_id = request.team_id + + # Validate the user is a member of the specified team + if team_id != user_api_key_dict.team_id: + from litellm.proxy.auth.auth_checks import get_team_membership + from litellm.proxy.proxy_server import user_api_key_cache + + membership = await get_team_membership( + user_id=user_api_key_dict.user_id or "", + team_id=request.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if membership is None: + raise HTTPException( + status_code=403, + detail=f"You are not a member of team {request.team_id!r}", + ) params = request.get_litellm_params_dict() if params.get("guardrail") != GENERIC_GUARDRAIL_API: @@ -673,7 +692,7 @@ async def register_guardrail( "litellm_params": litellm_params_str, "guardrail_info": guardrail_info_str, "status": "pending_review", - "team_id": user_api_key_dict.team_id, + "team_id": team_id, "submitted_at": now, "created_at": now, "updated_at": now, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts new file mode 100644 index 0000000000..ccc26c4b17 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useRegisterGuardrail.ts @@ -0,0 +1,74 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { + getProxyBaseUrl, + getGlobalLitellmHeaderName, + deriveErrorMessage, + handleError, +} from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +// ── Types ──────────────────────────────────────────────────────────────────── + +export interface RegisterGuardrailParams { + guardrail_name: string; + litellm_params: Record; + guardrail_info?: Record; + team_id: string; +} + +export interface RegisterGuardrailResponse { + guardrail_id: string; + guardrail_name: string; + status: string; + submitted_at?: string | null; +} + +// ── Fetch function ─────────────────────────────────────────────────────────── + +const registerGuardrail = async ( + accessToken: string, + params: RegisterGuardrailParams, +): Promise => { + const baseUrl = getProxyBaseUrl(); + const url = `${baseUrl}/guardrails/register`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(params), + }); + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})); + const errorMessage = deriveErrorMessage(errorData); + handleError(errorMessage); + throw new Error(errorMessage); + } + + return response.json(); +}; + +// ── Hook ───────────────────────────────────────────────────────────────────── + +const guardrailKeys = createQueryKeys("guardrails"); + +export const useRegisterGuardrail = () => { + const { accessToken } = useAuthorized(); + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (params) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return registerGuardrail(accessToken, params); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: guardrailKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx index c98ad18dca..fd3f8b1561 100644 --- a/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/TeamGuardrailsTab.tsx @@ -24,6 +24,7 @@ import { } from "@/components/networking"; import NotificationsManager from "@/components/molecules/notifications_manager"; import TeamDropdown from "@/components/common_components/team_dropdown"; +import { useRegisterGuardrail } from "@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail"; type GuardrailStatus = "active" | "pending" | "rejected"; @@ -824,6 +825,7 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { const [searchDebounced, setSearchDebounced] = useState(""); const [isSubmitModalOpen, setIsSubmitModalOpen] = useState(false); const [submitForm] = Form.useForm(); + const registerGuardrail = useRegisterGuardrail(); useEffect(() => { const t = setTimeout(() => setSearchDebounced(search), 300); @@ -1099,22 +1101,27 @@ export function TeamGuardrailsTab({ accessToken }: TeamGuardrailsTabProps) { form={submitForm} layout="vertical" initialValues={{ mode: "pre_call" }} - onFinish={(values) => { + onFinish={async (values) => { const litellm_params: Record = { guardrail: "generic_guardrail_api", mode: values.mode, api_base: values.api_base, ...(values.extra_litellm_params ? JSON.parse(values.extra_litellm_params) : {}), }; - const payload = { - guardrail_name: values.guardrail_name, - litellm_params, - guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined, - }; - // TODO: call registerGuardrailCall once backend is wired - console.log("Submit guardrail:", payload); - setIsSubmitModalOpen(false); - submitForm.resetFields(); + try { + await registerGuardrail.mutateAsync({ + team_id: values.team_id, + guardrail_name: values.guardrail_name, + litellm_params, + guardrail_info: values.guardrail_info ? JSON.parse(values.guardrail_info) : undefined, + }); + NotificationsManager.success("Guardrail submitted for review"); + setIsSubmitModalOpen(false); + submitForm.resetFields(); + fetchSubmissions(); + } catch { + // error already handled by networking layer + } }} >