From 56f009400f4c4c05f020212a63bf5a2ae5963f5d Mon Sep 17 00:00:00 2001
From: Tin <tin@berri.ai>
Date: Wed, 2 Sep 2026 15:54:54 -0700
Subject: [PATCH] fix(ui): restrict heuristic v2 to proxy admins

---
 .../add_model/ClassificationMethodConfig.tsx  | 20 +++++++++---
 .../add_model/ComplexityRouterConfig.test.tsx | 31 +++++++++++++++++--
 .../add_model/ComplexityRouterConfig.tsx      |  3 ++
 .../add_model/add_auto_router_tab.tsx         |  1 +
 4 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx
index a29527a20fa..6e93e64d760 100644
--- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx
@@ -159,15 +159,21 @@ interface ClassificationMethodConfigProps {
   showValidationErrors?: boolean;
   /** The resolved default model - see resolveComplexityDefaultModel. Names and gates the radio. */
   defaultModel?: string;
+  /** Mirrors the backend's proxy-admin-only heuristic_v2 write gate. */
+  allowHeuristicV2?: boolean;
 }
 
 const ClassifierTypeRadios: React.FC<{
   value: ComplexityRouterConfigValue;
   classifierType: ClassifierType;
   onTypeChange: (classifierType: ClassifierType) => void;
-}> = ({ value, classifierType, onTypeChange }) => {
+  allowHeuristicV2: boolean;
+}> = ({ value, classifierType, onTypeChange, allowHeuristicV2 }) => {
   const scorerLocked = Boolean(value.custom_tier_set);
   const scorerLockedReason = restrictedBy(value, "heuristicClassifier")?.reason;
+  const heuristicV2Locked = scorerLocked || !allowHeuristicV2;
+  const heuristicV2LockedReason =
+    scorerLockedReason ?? (!allowHeuristicV2 ? "Only proxy admins can configure Heuristic v2." : undefined);
   return (
     <RadioGroup
       value={classifierType}
@@ -186,9 +192,9 @@ const ClassifierTypeRadios: React.FC<{
             </span>
           </Label>
         </SimpleTooltip>
-        <SimpleTooltip content={scorerLockedReason}>
+        <SimpleTooltip content={heuristicV2LockedReason}>
           <Label className="items-start font-normal leading-normal has-data-disabled:cursor-not-allowed has-data-disabled:opacity-50">
-            <RadioGroupItem value="heuristic_v2" className="mt-0.5" disabled={scorerLocked} />
+            <RadioGroupItem value="heuristic_v2" className="mt-0.5" disabled={heuristicV2Locked} />
             <span>
               <strong className="font-semibold">Heuristic v2</strong>{" "}
               <span className="text-muted-foreground">
@@ -240,6 +246,7 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
   onCustomTechnicalKeywordsChange,
   showValidationErrors = false,
   defaultModel,
+  allowHeuristicV2 = false,
 }) => {
   const [draft, setDraft] = React.useState<{ id: string; raw: string } | null>(null);
   const hasDefaultModel = Boolean(defaultModel);
@@ -387,7 +394,12 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
 
   return (
     <>
-      <ClassifierTypeRadios value={value} classifierType={classifierType} onTypeChange={handleClassifierTypeChange} />
+      <ClassifierTypeRadios
+        value={value}
+        classifierType={classifierType}
+        onTypeChange={handleClassifierTypeChange}
+        allowHeuristicV2={allowHeuristicV2}
+      />
 
       {classifierType === "heuristic_first" && (
         <div className="mt-4 space-y-2">
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
index 349d03356f8..fe4a5d24cf6 100644
--- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
@@ -130,7 +130,12 @@ describe("ComplexityRouterConfig", () => {
   it("selects heuristic v2 without requiring a classifier model or showing weighted scoring", () => {
     const onChange = vi.fn();
     const { rerender } = renderWithProviders(
-      <ComplexityRouterConfig modelInfo={mockModelInfo} value={defaultValue} onChange={onChange} />,
+      <ComplexityRouterConfig
+        modelInfo={mockModelInfo}
+        value={defaultValue}
+        onChange={onChange}
+        allowHeuristicV2={true}
+      />,
     );
 
     fireEvent.click(screen.getByText("Advanced: Classification Method"));
@@ -144,7 +149,14 @@ describe("ComplexityRouterConfig", () => {
     );
 
     const heuristicV2Value: ComplexityRouterConfigValue = { ...defaultValue, classifier_type: "heuristic_v2" };
-    rerender(<ComplexityRouterConfig modelInfo={mockModelInfo} value={heuristicV2Value} onChange={onChange} />);
+    rerender(
+      <ComplexityRouterConfig
+        modelInfo={mockModelInfo}
+        value={heuristicV2Value}
+        onChange={onChange}
+        allowHeuristicV2={true}
+      />,
+    );
 
     expect(screen.queryByText("Classifier Model")).not.toBeInTheDocument();
     expect(screen.queryByText("Advanced scoring")).not.toBeInTheDocument();
@@ -152,6 +164,21 @@ describe("ComplexityRouterConfig", () => {
     expect(screen.queryByText(/Score < 0.15/)).not.toBeInTheDocument();
   });
 
+  it("reserves heuristic v2 selection for proxy admins", () => {
+    const onChange = vi.fn();
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} allowHeuristicV2={false} />);
+
+    fireEvent.click(screen.getByText("Advanced: Classification Method"));
+
+    const heuristicV2 = screen.getByRole("radio", { name: /Heuristic v2/ });
+    expect(heuristicV2).toHaveAttribute("aria-disabled", "true");
+    expect(screen.getByRole("radio", { name: /^Heuristic \(/ })).not.toHaveAttribute("aria-disabled", "true");
+    expect(screen.getByRole("radio", { name: /LLM Classifier/ })).not.toHaveAttribute("aria-disabled", "true");
+
+    fireEvent.click(heuristicV2);
+    expect(onChange).not.toHaveBeenCalled();
+  });
+
   it("should show classifier fields and use the configured values when classifier_type is llm", () => {
     const llmValue: ComplexityRouterConfigValue = {
       ...defaultValue,
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
index 73ab4f0abef..30076ac154e 100644
--- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
@@ -484,6 +484,7 @@ interface ComplexityRouterConfigProps {
   escalationKeywords?: string[];
   onEscalationKeywordsChange?: (keywords: string[]) => void;
   showValidationErrors?: boolean;
+  allowHeuristicV2?: boolean;
 }
 
 export const TIER_DESCRIPTIONS: Record<
@@ -605,6 +606,7 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
   escalationKeywords = [],
   onEscalationKeywordsChange,
   showValidationErrors = false,
+  allowHeuristicV2 = false,
 }) => {
   const customTierSet = value.custom_tier_set;
   const tierRows = activeTierRows(value);
@@ -821,6 +823,7 @@ const ComplexityRouterConfig: React.FC<ComplexityRouterConfigProps> = ({
                 onCustomTechnicalKeywordsChange={onCustomTechnicalKeywordsChange}
                 showValidationErrors={showValidationErrors}
                 defaultModel={defaultModel}
+                allowHeuristicV2={allowHeuristicV2}
               />
             ),
           },
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
index ab5ea4cbfc9..2178545b8ad 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
@@ -615,6 +615,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
                       escalationKeywords={escalationKeywords}
                       onEscalationKeywordsChange={setEscalationKeywords}
                       showValidationErrors={showValidationErrors}
+                      allowHeuristicV2={createScope === "unscoped-ok"}
                     />
                   </div>
                 )}