From 85cf6568279b145d08df0c57c7abb566223b97b2 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 31 Jul 2026 13:29:04 -0700 Subject: [PATCH] refactor(ui): load auto-router models with useQuery instead of a hand-rolled effect 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). --- .../add_model/add_auto_router_tab.test.tsx | 52 +++++++------------ .../add_model/add_auto_router_tab.tsx | 42 +++++---------- 2 files changed, 31 insertions(+), 63 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index b095a858eee..5fc39bbc45a 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -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 = () => { 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(); 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( , ); openTemplateDropdown(); await waitFor(() => expect(isOptionDisabled(optionByLabel("OpenAI Family")!)).toBe(false)); - fireEvent.click(optionByLabel("OpenAI Family")!); rerender(); - 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((resolve) => (resolveA = resolve))) - .mockResolvedValueOnce([]); - - const { rerender } = renderWithProviders( - , - ); - rerender(); - 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); }); 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 c66cfb9c4e4..9abc1c17ca1 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 @@ -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 = ({ const requiresTeamScope = createScope === "team-required"; const [form] = Form.useForm(); const [modelAccessGroups, setModelAccessGroups] = useState([]); - const [modelInfo, setModelInfo] = useState([]); - const [modelsLoadState, setModelsLoadState] = useState<"loading" | "loaded" | "error">("loading"); const [complexityRouterConfig, setComplexityRouterConfig] = useState({ tiers: { SIMPLE: [], MEDIUM: [], COMPLEX: [], REASONING: [] }, @@ -106,30 +105,15 @@ const AddAutoRouterTab: React.FC = ({ 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 = ({ // 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" }; };