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
This commit is contained in:
Ryan Crabbe
2026-04-02 17:51:17 -07:00
parent 7a27434f89
commit fe8cc100de
3 changed files with 113 additions and 13 deletions
@@ -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,
@@ -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<string, unknown>;
guardrail_info?: Record<string, unknown>;
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<RegisterGuardrailResponse> => {
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<RegisterGuardrailResponse, Error, RegisterGuardrailParams>({
mutationFn: async (params) => {
if (!accessToken) {
throw new Error("Access token is required");
}
return registerGuardrail(accessToken, params);
},
onSuccess: () => {
queryClient.invalidateQueries({ queryKey: guardrailKeys.all });
},
});
};
@@ -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<string, unknown> = {
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
}
}}
>
<Form.Item