scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 69

refactor(ui): load auto-router models with useQuery instead of a hand-rolled effect

Browse source 
The model list drives preset availability and is scoped to accessToken. Loading it with a useState + useEffect + ignore-flag loader coupled the fetch lifecycle to manual resets, and each token-change edge (stale list, out-of-order resolution, config erased on reset) was handled by adding another line to that effect. This replaces the whole loader with useQuery keyed on accessToken, matching how EditFallbacks and the rest of the dashboard fetch model data.

react-query owns the race surface: a caller switch is a new query key, so the previous caller's list is never read for the new caller and out-of-order resolutions are discarded by key. There is no reset-on-token-change anymore, so a token change can no longer erase the user's in-progress configuration; only the model list is token-scoped, and preset availability recomputes from it. isLoading and isError replace the hand-rolled load-state union.

Net effect deletes the two model-related useState hooks, the loader effect, and the ignore guard. Tests updated to drive the query mock and assert re-gating on caller switch; mutation-checked (removing accessToken from the query key, skipping the loading gate, and swapping the nullish match_threshold prefill each turn a test red).
This commit is contained in:
Tin Chi Lo 2026-07-31 13:29:04 -07:00
parent abba8e0102
commit 85cf656827
2 changed files with 31 additions and 63 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
View file

@ -1,4 +1,4 @@
import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
import { renderWithProviders, screen, waitFor, testQueryClient } from "../../../tests/test-utils";
import { fireEvent } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { vi } from "vitest";
@ -114,6 +114,10 @@ const Harness = () => <AddAutoRouterTab handleOk={vi.fn()} accessToken="token" u
describe("AddAutoRouterTab", () => {
beforeEach(() => {
vi.clearAllMocks();
// testQueryClient is a shared singleton with staleTime: Infinity, so cached model lists would
// otherwise bleed across tests (a later test reusing accessToken="token" would read an earlier
// test's data instead of its own mock).
testQueryClient.clear();
mockFetchAvailableModels.mockResolvedValue([]);
mockHandleAddAutoRouterSubmit.mockResolvedValue(undefined);
});
@ -158,8 +162,10 @@ describe("AddAutoRouterTab", () => {
renderWithProviders(<Harness />);
openTemplateDropdown();
await waitFor(() => expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(true));
expect(optionByLabel("Anthropic Family")).toHaveTextContent(/Missing:.*claude-opus-5/);
// Wait for the terminal (loaded) state; "disabled" alone is also true mid-load, so assert on
// the missing-model text that only appears once the list has resolved.
await waitFor(() => expect(optionByLabel("Anthropic Family")).toHaveTextContent(/Missing:.*claude-opus-5/));
expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(true);
// The other family, fully available, stays selectable.
expect(isOptionDisabled(optionByLabel("OpenAI Family")!)).toBe(false);
});
@ -244,46 +250,24 @@ describe("AddAutoRouterTab", () => {
expect(mockHandleAddAutoRouterSubmit).not.toHaveBeenCalled();
});
// A caller switch must clear the preset selection, or a preset selected on the previous caller
// stays selected (and its models available in the old list) for the new caller. Only the preset
// choice and model-verification state are token-scoped; user config survives.
it("clears the preset selection when the access token changes", async () => {
mockFetchAvailableModels.mockResolvedValueOnce(ALL_FAMILY_MODELS).mockResolvedValueOnce(ALL_FAMILY_MODELS);
// Availability is scoped to the caller. Because the model query is keyed on accessToken, a caller
// switch re-fetches and re-gates the presets against the NEW caller's models: a preset the first
// caller could select greys out for a second caller who lacks one of its models. Nothing carries
// the first caller's list forward, so no preset stays wrongly selectable across the switch.
it("re-gates presets against the new caller when the access token changes", async () => {
mockFetchAvailableModels
.mockResolvedValueOnce(ALL_FAMILY_MODELS)
.mockResolvedValueOnce(ALL_FAMILY_MODELS.filter((m) => m.model_group !== "o3"));
const { rerender } = renderWithProviders(
<AddAutoRouterTab handleOk={vi.fn()} accessToken="caller-a" userRole="Admin" />,
);
openTemplateDropdown();
await waitFor(() => expect(isOptionDisabled(optionByLabel("OpenAI Family")!)).toBe(false));
fireEvent.click(optionByLabel("OpenAI Family")!);
rerender(<AddAutoRouterTab handleOk={vi.fn()} accessToken="caller-b" userRole="Admin" />);
await waitFor(() => expect(mockFetchAvailableModels).toHaveBeenCalledTimes(2));
openTemplateDropdown();
expect(optionByLabel("OpenAI Family")).not.toHaveClass("ant-select-item-option-selected");
});
// A late-resolving fetch from the previous caller must not overwrite the current caller's list.
// Caller A's request is held open, caller B's resolves empty; when A finally resolves with the
// full family, the ignore guard drops it so the preset stays greyed out for B. Without the guard,
// A's response would land after B's and wrongly re-enable the preset.
it("ignores a stale in-flight model fetch that resolves after the token changed", async () => {
let resolveA: (models: ModelGroup[]) => void = () => undefined;
mockFetchAvailableModels
.mockReturnValueOnce(new Promise<ModelGroup[]>((resolve) => (resolveA = resolve)))
.mockResolvedValueOnce([]);
const { rerender } = renderWithProviders(
<AddAutoRouterTab handleOk={vi.fn()} accessToken="caller-a" userRole="Admin" />,
);
rerender(<AddAutoRouterTab handleOk={vi.fn()} accessToken="caller-b" userRole="Admin" />);
await waitFor(() => expect(mockFetchAvailableModels).toHaveBeenCalledTimes(2));
resolveA(ALL_FAMILY_MODELS);
openTemplateDropdown();
await waitFor(() => expect(optionByLabel("OpenAI Family")).toBeTruthy());
await waitFor(() => expect(optionByLabel("OpenAI Family")).toHaveTextContent(/Missing:.*o3/));
expect(isOptionDisabled(optionByLabel("OpenAI Family")!)).toBe(true);
});

42
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
View file

@ -1,4 +1,5 @@
import React, { useEffect, useState } from "react";
import { useQuery } from "@tanstack/react-query";
import { Card, Form, Button, Tooltip, Typography, Select as AntdSelect, Modal } from "antd";
import { TextInput } from "@tremor/react";
import { modelAvailableCall } from "../networking";
@ -6,7 +7,7 @@ import { all_admin_roles } from "@/utils/roles";
import { type ModelWriteScope } from "@/utils/modelPermissions";
import TeamDropdown from "../common_components/team_dropdown";
import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit";
import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models";
import { fetchAvailableModels } from "@/components/llm_calls/fetch_models";
import ComplexityRouterConfig, {
ComplexityRouterConfigValue,
DEFAULT_ADAPTIVE_WEIGHTS,
@ -75,8 +76,6 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
const requiresTeamScope = createScope === "team-required";
const [form] = Form.useForm();
const [modelAccessGroups, setModelAccessGroups] = useState<string[]>([]);
const [modelInfo, setModelInfo] = useState<ModelGroup[]>([]);
const [modelsLoadState, setModelsLoadState] = useState<"loading" | "loaded" | "error">("loading");
const [complexityRouterConfig, setComplexityRouterConfig] = useState<ComplexityRouterConfigValue>({
tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] },
@ -106,30 +105,15 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
fetchModelAccessGroups();
}, [accessToken]);
useEffect(() => {
let ignore = false;
const loadModels = async () => {
setModelsLoadState("loading");
setModelInfo([]);
setSelectedPreset(undefined);
try {
const uniqueModels = await fetchAvailableModels(accessToken);
if (ignore) return;
setModelInfo(uniqueModels);
setModelsLoadState("loaded");
} catch (error) {
console.error("Error fetching model info for auto router:", error);
if (ignore) return;
setModelsLoadState("error");
}
};
loadModels();
return () => {
ignore = true;
};
}, [accessToken]);
const {
data: modelInfo = [],
isLoading: modelsLoading,
isError: modelsError,
} = useQuery({
queryKey: ["availableModels", "autoRouter", accessToken],
queryFn: () => fetchAvailableModels(accessToken),
enabled: Boolean(accessToken),
});
const isAdmin = all_admin_roles.includes(userRole);
@ -147,8 +131,8 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
// whose models we cannot yet verify, and a failed fetch leaves every preset unverifiable. This
// makes the load-race (pick during loading, then discover a missing model) unrepresentable.
const presetAvailability = (preset: AutoRouterPreset): PresetAvailability => {
if (modelsLoadState === "loading") return { kind: "loading" };
if (modelsLoadState === "error") return { kind: "unverifiable" };
if (modelsLoading) return { kind: "loading" };
if (modelsError) return { kind: "unverifiable" };
const missing = getMissingModelsInPreset(preset, availableModelSet);
return missing.length > 0 ? { kind: "missing_models", models: missing } : { kind: "available" };
};