From 625eb49c852de8b8be5e4fc641c0e5cdf3bdca9a Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 2 Oct 2026 07:36:59 +0000 Subject: [PATCH] feat(ui): edit logging_only_scope in the custom code guardrail modal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../custom_code/CustomCodeModal.test.tsx | 123 +++++++++++++++++- .../custom_code/CustomCodeModal.tsx | 49 ++++++- 2 files changed, 169 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx index fc19ca2a88b..a8b07a30329 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import CustomCodeModal from "./CustomCodeModal"; import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking"; @@ -130,6 +130,127 @@ describe("CustomCodeModal", () => { expect(screen.queryByText("logging_only")).not.toBeInTheDocument(); }); + it("should update the logging-only scope in edit mode", async () => { + const user = userEvent.setup(); + renderModal({ + editData: { + guardrail_id: "g-1", + guardrail_name: "existing-guardrail", + litellm_params: { + mode: ["logging_only"], + logging_only_scope: "input", + custom_code: "def apply_guardrail(): pass", + }, + }, + }); + + const scopeSelect = await screen.findByRole("combobox", { name: "Logging only scope" }); + expect(scopeSelect).toHaveTextContent("Input only (request)"); + + await user.click(scopeSelect); + await user.click(await screen.findByRole("option", { name: "Output only (response)" })); + await user.click(screen.getByRole("button", { name: /update guardrail/i })); + + await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1)); + expect(mockUpdate.mock.calls[0][2]).toMatchObject({ + litellm_params: { logging_only_scope: "output" }, + }); + }); + + it("should clear the logging-only scope when Default is selected in edit mode", async () => { + const user = userEvent.setup(); + renderModal({ + editData: { + guardrail_id: "g-1", + guardrail_name: "existing-guardrail", + litellm_params: { + mode: ["logging_only"], + logging_only_scope: "input", + custom_code: "def apply_guardrail(): pass", + }, + }, + }); + + await user.click(await screen.findByRole("combobox", { name: "Logging only scope" })); + await user.click(await screen.findByRole("option", { name: "Default (request and response)" })); + await user.click(screen.getByRole("button", { name: /update guardrail/i })); + + await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1)); + expect(mockUpdate.mock.calls[0][2]).toHaveProperty("litellm_params.logging_only_scope", null); + }); + + it("should omit an unchanged logging-only scope from the edit payload", async () => { + renderModal({ + editData: { + guardrail_id: "g-1", + guardrail_name: "existing-guardrail", + litellm_params: { + mode: ["logging_only"], + logging_only_scope: "input", + custom_code: "def apply_guardrail(): pass", + }, + }, + }); + + fireEvent.change(await screen.findByPlaceholderText("e.g., block-pii-custom"), { + target: { value: "renamed-guardrail" }, + }); + await userEvent.setup().click(screen.getByRole("button", { name: /update guardrail/i })); + + await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1)); + expect(mockUpdate.mock.calls[0][2]).not.toHaveProperty("litellm_params.logging_only_scope"); + }); + + it("should not show logging-only scope for other modes in edit mode", async () => { + renderModal({ + editData: { + guardrail_id: "g-1", + guardrail_name: "existing-guardrail", + litellm_params: { mode: "pre_call", custom_code: "def apply_guardrail(): pass" }, + }, + }); + + expect(await screen.findByText("Edit Custom Guardrail")).toBeInTheDocument(); + expect(screen.queryByRole("combobox", { name: "Logging only scope" })).not.toBeInTheDocument(); + }); + + it("should create a logging-only guardrail with its selected scope", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.click(screen.getAllByRole("combobox")[0]); + await user.keyboard("logging_only"); + await user.click(await screen.findByRole("option", { name: "logging_only" })); + await user.click(await screen.findByRole("combobox", { name: "Logging only scope" })); + await user.click(await screen.findByRole("option", { name: "Input only (request)" })); + fireEvent.change(screen.getByPlaceholderText("e.g., block-pii-custom"), { + target: { value: "logging-only-guardrail" }, + }); + await user.click(screen.getByRole("button", { name: /save guardrail/i })); + + await waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); + expect(mockCreate.mock.calls[0][1]).toMatchObject({ + litellm_params: { logging_only_scope: "input" }, + }); + }); + + it("should omit the default logging-only scope when creating a guardrail", async () => { + const user = userEvent.setup(); + renderModal(); + + await user.click(screen.getAllByRole("combobox")[0]); + await user.keyboard("logging_only"); + await user.click(await screen.findByRole("option", { name: "logging_only" })); + await screen.findByRole("combobox", { name: "Logging only scope" }); + fireEvent.change(screen.getByPlaceholderText("e.g., block-pii-custom"), { + target: { value: "logging-only-guardrail" }, + }); + await user.click(screen.getByRole("button", { name: /save guardrail/i })); + + await waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); + expect(mockCreate.mock.calls[0][1]).not.toHaveProperty("litellm_params.logging_only_scope"); + }); + it("should expand the test section and run a test against the backend", async () => { const user = userEvent.setup(); mockTest.mockResolvedValue({ success: true, result: { action: "allow" } } as never); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx index ccaf1918e1d..29523e9e294 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx @@ -2,7 +2,14 @@ import React, { useState, useRef, useEffect } from "react"; import { CheckCircle2, ChevronRight, Code, ExternalLink, PlayCircle, Save, Users, XCircle } from "lucide-react"; import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking"; import { toast } from "@/lib/toast"; -import { type LoggingOnlyScope } from "../guardrail_info_helpers"; +import { + choiceToLoggingOnlyScope, + getLoggingOnlyScopeOptions, + getLoggingOnlyScopeUpdate, + loggingOnlyScopeToChoice, + type LoggingOnlyScope, + type LoggingOnlyScopeChoice, +} from "../guardrail_info_helpers"; import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { @@ -165,6 +172,8 @@ const TEMPLATE_ITEMS = Object.entries(CODE_TEMPLATES).map(([key, template]) => ( label: template.name, })); +const LOGGING_ONLY_SCOPE_ITEMS = getLoggingOnlyScopeOptions(true); + type ModeOption = (typeof MODE_OPTIONS)[number]; const MODE_OPTION_BY_VALUE: Record = Object.fromEntries( @@ -193,11 +202,13 @@ interface CustomCodeModalProps { editData?: EditGuardrailData | null; } +/* eslint-disable max-lines -- The existing modal remains the requested location for this scope control */ const CustomCodeModal: React.FC = ({ visible, onClose, onSuccess, accessToken, editData }) => { const anchor = useComboboxAnchor(); const isEditMode = !!editData; const [guardrailName, setGuardrailName] = useState(""); const [mode, setMode] = useState(["pre_call"]); + const [loggingOnlyScopeChoice, setLoggingOnlyScopeChoice] = useState("default"); const [defaultOn, setDefaultOn] = useState(false); const [selectedTemplate, setSelectedTemplate] = useState("empty"); const [code, setCode] = useState(CODE_TEMPLATES.empty.code); @@ -322,6 +333,7 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS // Edit mode: populate with existing data setGuardrailName(editData.guardrail_name || ""); setMode(normalizeMode(editData.litellm_params?.mode)); + setLoggingOnlyScopeChoice(loggingOnlyScopeToChoice(editData.litellm_params?.logging_only_scope)); setDefaultOn(editData.litellm_params?.default_on || false); setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code); setSelectedTemplate(""); // No template selected in edit mode @@ -329,6 +341,7 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS // Create mode: reset to defaults setGuardrailName(""); setMode(["pre_call"]); + setLoggingOnlyScopeChoice("default"); setDefaultOn(false); setSelectedTemplate("empty"); setCode(CODE_TEMPLATES.empty.code); @@ -401,11 +414,17 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS if (defaultOn !== editData.litellm_params?.default_on) { updateData.litellm_params.default_on = defaultOn; } - + if (mode.includes("logging_only")) { + Object.assign( + updateData.litellm_params, + getLoggingOnlyScopeUpdate(editData.litellm_params, loggingOnlyScopeChoice), + ); + } await updateGuardrailCall(accessToken, editData.guardrail_id, updateData); toast.success("Custom code guardrail updated successfully"); } else { // Create new guardrail + const loggingOnlyScope = choiceToLoggingOnlyScope(loggingOnlyScopeChoice); const guardrailData = { guardrail_name: guardrailName, litellm_params: { @@ -413,6 +432,9 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS mode: mode, default_on: defaultOn, custom_code: code, + ...(mode.includes("logging_only") && loggingOnlyScope !== null + ? { logging_only_scope: loggingOnlyScope } + : {}), }, guardrail_info: {}, }; @@ -549,6 +571,29 @@ const CustomCodeModal: React.FC = ({ visible, onClose, onS + {mode.includes("logging_only") && ( +
+ + +
+ )}