fix(ui): show newly added custom PII entity when a category filter is active

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.integration.test.tsx

@@ -1,7 +1,35 @@
+import { useState } from "react";
 import { fireEvent, render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import { describe, expect, it, vi } from "vitest";
 import PiiConfiguration from "./pii_configuration";
 
+const StatefulPiiConfiguration = () => {
+  const [selectedEntities, setSelectedEntities] = useState<string[]>([]);
+  const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({});
+
+  return (
+    <PiiConfiguration
+      entities={["PERSON", "EMAIL"]}
+      actions={["MASK", "BLOCK"]}
+      selectedEntities={selectedEntities}
+      selectedActions={selectedActions}
+      onEntitySelect={(entity) => {
+        setSelectedEntities((previous) =>
+          previous.includes(entity) ? previous.filter((selected) => selected !== entity) : [...previous, entity],
+        );
+      }}
+      onActionSelect={(entity, action) => {
+        setSelectedActions((previous) => ({ ...previous, [entity]: action }));
+      }}
+      entityCategories={[
+        { category: "Personal", entities: ["PERSON"] },
+        { category: "Contact", entities: ["EMAIL"] },
+      ]}
+    />
+  );
+};
+
 describe("PiiConfiguration", () => {
   it("should render", () => {
     render(
@@ -56,6 +84,31 @@ describe("PiiConfiguration", () => {
     expect(onActionSelect).toHaveBeenCalledWith("SE_PERSONNUMMER", "MASK");
   });
 
+  it("shows a newly added custom entity while a category filter is active", async () => {
+    const user = userEvent.setup();
+    render(<StatefulPiiConfiguration />);
+
+    const categoryInput = screen.getByPlaceholderText("Select categories to filter by");
+    await user.click(categoryInput);
+    if (categoryInput.getAttribute("aria-expanded") !== "true") {
+      categoryInput.focus();
+      await user.keyboard("{Enter}");
+    }
+    await user.click(await screen.findByRole("option", { name: "Personal" }));
+    await user.keyboard("{Escape}");
+
+    expect(screen.getByText("PERSON")).toBeVisible();
+    expect(screen.queryByText("EMAIL")).not.toBeInTheDocument();
+
+    fireEvent.change(screen.getByRole("textbox", { name: "Custom entity name" }), {
+      target: { value: "se personnummer" },
+    });
+    fireEvent.click(screen.getByRole("button", { name: "Add entity" }));
+
+    expect(await screen.findByText("SE PERSONNUMMER")).toBeVisible();
+    expect(screen.getByText("PERSON")).toBeVisible();
+  });
+
   it("rejects invalid custom entity names", () => {
     const onEntitySelect = vi.fn();
     const onActionSelect = vi.fn();

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx

@@ -77,6 +77,9 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
       return;
     }
 
+    if (selectedCategories.length > 0 && !selectedCategories.includes("Custom")) {
+      setSelectedCategories([...selectedCategories, "Custom"]);
+    }
     onEntitySelect(normalizedEntityName);
     onActionSelect(normalizedEntityName, "MASK");
   };