feat(ui): add budget fallbacks configuration to key create/edit forms

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-07 08:26:10 +00:00 · 2026-07-03 19:50:36 +00:00 · 2026-07-03 19:50:36 +00:00 · ed51c96d3f
commit ed51c96d3f
parent e06adb5588
6 changed files with 300 additions and 0 deletions
--- a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.test.tsx
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.test.tsx
@ -0,0 +1,82 @@
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, expect, it, vi } from "vitest";
+import { BudgetFallbacksEditor } from "./BudgetFallbacksEditor";
+
+const MODELS = ["gpt-4", "gpt-3.5-turbo", "claude-3", "claude-haiku"];
+
+describe("BudgetFallbacksEditor", () => {
+  it("renders empty state with add button", () => {
+    const onChange = vi.fn();
+    render(<BudgetFallbacksEditor value={{}} onChange={onChange} availableModels={MODELS} />);
+    expect(screen.getByText("Add Budget Fallback")).toBeTruthy();
+    expect(screen.getByText(/reroute to fallback models/)).toBeTruthy();
+  });
+
+  it("renders existing entries from value prop", () => {
+    const onChange = vi.fn();
+    render(
+      <BudgetFallbacksEditor
+        value={{ "gpt-4": ["gpt-3.5-turbo", "claude-3"] }}
+        onChange={onChange}
+        availableModels={MODELS}
+      />,
+    );
+    expect(screen.getByText("IF BUDGET EXCEEDED, TRY")).toBeTruthy();
+    expect(screen.getByText("Primary Model")).toBeTruthy();
+    expect(screen.getByText("Fallback Models")).toBeTruthy();
+  });
+
+  it("adds a new empty entry when clicking add button", async () => {
+    const user = userEvent.setup();
+    const onChange = vi.fn();
+    render(<BudgetFallbacksEditor value={{}} onChange={onChange} availableModels={MODELS} />);
+
+    await user.click(screen.getByText("Add Budget Fallback"));
+    expect(screen.getByText("Primary Model")).toBeTruthy();
+    expect(onChange).toHaveBeenCalledWith({});
+  });
+
+  it("removes an entry and emits updated dict", async () => {
+    const user = userEvent.setup();
+    const onChange = vi.fn();
+    const { container } = render(
+      <BudgetFallbacksEditor
+        value={{ "gpt-4": ["gpt-3.5-turbo"], "claude-3": ["claude-haiku"] }}
+        onChange={onChange}
+        availableModels={MODELS}
+      />,
+    );
+
+    const removeButtons = container.querySelectorAll<HTMLButtonElement>(".relative > button[type='button']");
+    expect(removeButtons.length).toBe(2);
+
+    await user.click(removeButtons[0]);
+    expect(onChange).toHaveBeenLastCalledWith({ "claude-3": ["claude-haiku"] });
+  });
+
+  it("renders multiple entries for multiple fallback groups", () => {
+    const onChange = vi.fn();
+    render(
+      <BudgetFallbacksEditor
+        value={{ "gpt-4": ["claude-3"], "gpt-3.5-turbo": ["claude-haiku"] }}
+        onChange={onChange}
+        availableModels={MODELS}
+      />,
+    );
+    const labels = screen.getAllByText("Primary Model");
+    expect(labels.length).toBe(2);
+  });
+
+  it("shows ordering hint when multiple fallback models are configured", () => {
+    const onChange = vi.fn();
+    render(
+      <BudgetFallbacksEditor
+        value={{ "gpt-4": ["gpt-3.5-turbo", "claude-3"] }}
+        onChange={onChange}
+        availableModels={MODELS}
+      />,
+    );
+    expect(screen.getByText(/first model still within its own budget/)).toBeTruthy();
+  });
+});
--- a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx
@ -0,0 +1,153 @@
+import { Button, Select, Tooltip } from "antd";
+import { ArrowDown, Plus, X } from "lucide-react";
+import React, { useState } from "react";
+
+interface FallbackEntry {
+  id: string;
+  primaryModel: string | null;
+  fallbackModels: string[];
+}
+
+interface BudgetFallbacksEditorProps {
+  value: Record<string, string[]>;
+  onChange: (v: Record<string, string[]>) => void;
+  availableModels: string[];
+}
+
+const entriesToDict = (entries: readonly FallbackEntry[]): Record<string, string[]> =>
+  Object.fromEntries(
+    entries
+      .filter(
+        (e): e is FallbackEntry & { primaryModel: string } => e.primaryModel !== null && e.fallbackModels.length > 0,
+      )
+      .map((e) => [e.primaryModel, e.fallbackModels]),
+  );
+
+const dictToEntries = (dict: Record<string, string[]>): FallbackEntry[] => {
+  const keys = Object.keys(dict);
+  if (keys.length === 0) return [];
+  return keys.map((model, i) => ({
+    id: String(i + 1),
+    primaryModel: model,
+    fallbackModels: dict[model],
+  }));
+};
+
+export function BudgetFallbacksEditor({ value, onChange, availableModels }: BudgetFallbacksEditorProps) {
+  const [entries, setEntries] = useState<FallbackEntry[]>(() => dictToEntries(value));
+
+  const emitChange = (updated: FallbackEntry[]) => {
+    setEntries(updated);
+    onChange(entriesToDict(updated));
+  };
+
+  const addEntry = () => {
+    emitChange([...entries, { id: Date.now().toString(), primaryModel: null, fallbackModels: [] }]);
+  };
+
+  const removeEntry = (id: string) => {
+    emitChange(entries.filter((e) => e.id !== id));
+  };
+
+  const updateEntry = (id: string, patch: Partial<FallbackEntry>) => {
+    emitChange(entries.map((e) => (e.id === id ? { ...e, ...patch } : e)));
+  };
+
+  const usedPrimaryModels = new Set(entries.map((e) => e.primaryModel).filter(Boolean));
+
+  if (entries.length === 0) {
+    return (
+      <div>
+        <div className="text-xs text-gray-500 mb-2">
+          When a model exceeds its per-model budget, requests automatically reroute to fallback models
+        </div>
+        <Button size="small" onClick={addEntry} icon={<Plus className="w-3 h-3" />}>
+          Add Budget Fallback
+        </Button>
+      </div>
+    );
+  }
+
+  return (
+    <div className="space-y-4">
+      <div className="text-xs text-gray-500">
+        When a model exceeds its per-model budget, requests automatically reroute to fallback models
+      </div>
+      {entries.map((entry) => {
+        const availablePrimaryOptions = availableModels.filter(
+          (m) => m === entry.primaryModel || !usedPrimaryModels.has(m),
+        );
+        const availableFallbackOptions = availableModels.filter((m) => m !== entry.primaryModel);
+
+        return (
+          <div key={entry.id} className="relative rounded-lg border border-gray-200 bg-gray-50 p-4">
+            <button
+              type="button"
+              onClick={() => removeEntry(entry.id)}
+              className="absolute top-2 right-2 text-gray-400 hover:text-red-500 transition-colors p-1"
+            >
+              <X className="w-4 h-4" />
+            </button>
+
+            <div className="mb-3">
+              <label className="block text-xs font-medium text-gray-600 mb-1">Primary Model</label>
+              <Select
+                className="w-full"
+                placeholder="Select model"
+                value={entry.primaryModel}
+                onChange={(v) => {
+                  const newFallbacks = entry.fallbackModels.filter((m) => m !== v);
+                  updateEntry(entry.id, { primaryModel: v, fallbackModels: newFallbacks });
+                }}
+                showSearch
+                filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())}
+                options={availablePrimaryOptions.map((m) => ({ label: m, value: m }))}
+                getPopupContainer={(trigger) => trigger.parentElement || document.body}
+              />
+            </div>
+
+            <div className="flex items-center justify-center -my-1 mb-2">
+              <div className="bg-amber-50 text-amber-600 px-3 py-0.5 rounded-full text-[10px] font-bold border border-amber-100 flex items-center gap-1">
+                <ArrowDown className="w-3 h-3" />
+                IF BUDGET EXCEEDED, TRY
+              </div>
+            </div>
+
+            <div>
+              <label className="block text-xs font-medium text-gray-600 mb-1">Fallback Models</label>
+              <Select
+                mode="multiple"
+                className="w-full"
+                placeholder={entry.primaryModel ? "Select fallback models" : "Select a primary model first"}
+                value={entry.fallbackModels}
+                onChange={(values) => updateEntry(entry.id, { fallbackModels: values })}
+                disabled={!entry.primaryModel}
+                showSearch
+                filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())}
+                options={availableFallbackOptions.map((m) => ({ label: m, value: m }))}
+                getPopupContainer={(trigger) => trigger.parentElement || document.body}
+                maxTagCount="responsive"
+                maxTagPlaceholder={(omittedValues) => (
+                  <Tooltip
+                    styles={{ root: { pointerEvents: "none" } }}
+                    title={omittedValues.map(({ value: v }) => v).join(", ")}
+                  >
+                    <span>+{omittedValues.length} more</span>
+                  </Tooltip>
+                )}
+              />
+              {entry.fallbackModels.length > 1 && (
+                <div className="text-[10px] text-gray-400 mt-1 ml-1">
+                  Tried in order; first model still within its own budget is used
+                </div>
+              )}
+            </div>
+          </div>
+        );
+      })}
+      <Button size="small" onClick={addEntry} icon={<Plus className="w-3 h-3" />}>
+        Add Budget Fallback
+      </Button>
+    </div>
+  );
+}
--- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
+++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
@ -97,6 +97,7 @@ export interface KeyResponse {
    agent_access_groups?: string[];
  };
  access_group_ids?: string[];
+  budget_fallbacks?: Record<string, string[]>;
  budget_limits?: Array<{ budget_duration: string; max_budget: number; reset_at?: string }>;
  auto_rotate?: boolean;
  rotation_interval?: string;
--- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
+++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx
@ -28,6 +28,7 @@ import TeamDropdown from "../common_components/team_dropdown";
 import OrganizationDropdown from "../common_components/OrganizationDropdown";
 import ProjectDropdown from "../common_components/ProjectDropdown";
 import { CreateUserButton } from "../CreateUserButton";
+import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor";
 import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
 import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
 import { Team } from "../key_team_helpers/key_list";
@ -202,6 +203,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
  const [rotationInterval, setRotationInterval] = useState<string>("30d");
  const [routerSettings, setRouterSettings] = useState<RouterSettingsAccordionValue | null>(null);
  const [budgetLimits, setBudgetLimits] = useState<BudgetWindowEntry[]>([]);
+  const [budgetFallbacks, setBudgetFallbacks] = useState<Record<string, string[]>>({});
  const [routerSettingsKey, setRouterSettingsKey] = useState<number>(0);
  const [agentsList, setAgentsList] = useState<{ agent_id: string; agent_name: string }[]>([]);
  const [selectedAgentId, setSelectedAgentId] = useState<string | null>(null);
@ -220,6 +222,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
    setSelectedOrganizationId(null);
    setSelectedProjectId(null);
    setBudgetLimits([]);
+    setBudgetFallbacks({});
  };

  const handleCancel = () => {
@ -239,6 +242,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
    setSelectedOrganizationId(null);
    setSelectedProjectId(null);
    setBudgetLimits([]);
+    setBudgetFallbacks({});
  };

  useEffect(() => {
@ -536,6 +540,10 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
        formValues.budget_limits = validWindows;
      }

+      if (Object.keys(budgetFallbacks).length > 0) {
+        formValues.budget_fallbacks = budgetFallbacks;
+      }
+
      let response;
      if (keyOwner === "service_account") {
        response = await keyCreateServiceAccountCall(accessToken, formValues);
@ -558,6 +566,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
      NotificationsManager.success("Virtual Key Created");
      form.resetFields();
      setBudgetLimits([]);
+      setBudgetFallbacks({});
      localStorage.removeItem("userData" + userID);
    } catch (error) {
      console.log("error in create key:", error);
@ -1076,6 +1085,23 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
                  >
                    <BudgetWindowsEditor value={budgetLimits} onChange={setBudgetLimits} />
                  </Form.Item>
+                  <Form.Item
+                    className="mt-4"
+                    label={
+                      <span>
+                        Budget Fallbacks{" "}
+                        <Tooltip title="When a model exceeds its per-model budget (model_max_budget), requests automatically reroute to fallback models instead of failing. Configure per-model budgets in Advanced Settings.">
+                          <InfoCircleOutlined style={{ marginLeft: "4px" }} />
+                        </Tooltip>
+                      </span>
+                    }
+                  >
+                    <BudgetFallbacksEditor
+                      value={budgetFallbacks}
+                      onChange={setBudgetFallbacks}
+                      availableModels={modelsToPick}
+                    />
+                  </Form.Item>
                  <Form.Item
                    className="mt-4"
                    label={
--- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx
@ -16,6 +16,7 @@ import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSel
 import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
 import OrganizationDropdown from "../common_components/OrganizationDropdown";
 import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils";
+import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor";
 import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor";
 import { KeyResponse } from "../key_team_helpers/key_list";
 import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
@ -108,6 +109,9 @@ export function KeyEditView({
  const [budgetLimits, setBudgetLimits] = useState<BudgetWindowEntry[]>(
    Array.isArray(keyData.budget_limits) ? keyData.budget_limits : [],
  );
+  const [budgetFallbacks, setBudgetFallbacks] = useState<Record<string, string[]>>(
+    keyData.budget_fallbacks && typeof keyData.budget_fallbacks === "object" ? keyData.budget_fallbacks : {},
+  );
  const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations();
  const { data: projects } = useProjects();
  const { data: uiSettingsData } = useUISettings();
@ -304,6 +308,8 @@ export function KeyEditView({
        values.budget_limits = [];
      }

+      values.budget_fallbacks = budgetFallbacks;
+
      await onSubmit(values);
    } finally {
      setIsKeySaving(false);
@ -472,6 +478,23 @@ export function KeyEditView({
        <BudgetWindowsEditor value={budgetLimits} onChange={setBudgetLimits} />
      </Form.Item>

+      <Form.Item
+        label={
+          <span>
+            Budget Fallbacks{" "}
+            <Tooltip title="When a model exceeds its per-model budget, requests automatically reroute to fallback models instead of failing">
+              <InfoCircleOutlined style={{ marginLeft: "4px" }} />
+            </Tooltip>
+          </span>
+        }
+      >
+        <BudgetFallbacksEditor
+          value={budgetFallbacks}
+          onChange={setBudgetFallbacks}
+          availableModels={availableModels}
+        />
+      </Form.Item>
+
      <Form.Item label="TPM Limit" name="tpm_limit">
        <NumericalInput min={0} />
      </Form.Item>
--- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx
@ -748,6 +748,21 @@ export default function KeyInfoView({
                    </Text>
                  </div>

+                  {currentKeyData.budget_fallbacks && Object.keys(currentKeyData.budget_fallbacks).length > 0 && (
+                    <div>
+                      <Text className="font-medium">Budget Fallbacks</Text>
+                      <div className="mt-1 space-y-1">
+                        {Object.entries(currentKeyData.budget_fallbacks).map(([model, fallbacks]) => (
+                          <div key={model} className="text-xs text-gray-600">
+                            <span className="font-medium">{model}</span>
+                            <span className="mx-1 text-gray-400">-&gt;</span>
+                            {fallbacks.join(", ")}
+                          </div>
+                        ))}
+                      </div>
+                    </div>
+                  )}
+
                  <div>
                    <Text className="font-medium">Tags</Text>
                    <div className="flex flex-wrap gap-2 mt-1">