From 49d4c61623d55ffd36a5c94979ab651c6fe5c68c Mon Sep 17 00:00:00 2001 From: oobabooga Date: Mon, 22 Jun 2026 02:19:09 -0300 Subject: [PATCH] studio/frontend: Prevent staged model's load settings from disappearing while it loads (#6458) * Studio: keep staged model's load settings visible while it loads * Studio: pin staged load settings at click time, match loading pick by variant * Studio: drop any stale staged pick after a successful load * Studio: disable staged Load button while a different model loads When a model is staged with "Load on selection" off and a different model (or a different GGUF variant of the same repo) is already loading, selectModel's in-flight-load guard matches on id + native path token only, so the click was silently deduped to a no-op. The staged Load button stayed enabled but did nothing, and the stale stage was then cleared once the other load finished. Disable the staged Load button (showing "Another model loading...") whenever a different model is loading, so it is no longer an enabled no-op. The stage stays put and the user can retry once the in-flight load settles. * Studio: refuse loading a different model while one is in flight The in-flight-load guard in selectModel matched on id + native path token only, so a different GGUF variant of the same repo fell through and was silently deduped to a no-op. The earlier commit disabled the staged Load button for this case, but other entry points (e.g. selecting a different quant from the model picker with "Load on selection" on) still hit selectModel directly and no-op'd. Make the guard variant-aware (id + GGUF variant + native path token) and, for a genuinely different model while a load is in flight, surface a "Another model is already loading" toast instead of silently returning. The load path has no clean supersession, so a second concurrent load is not started; the user is told to wait or cancel. Centralized in selectModel so every entry point is covered. * Studio: lock staged model's load settings while it loads The staged Load button snapshots context length, KV cache dtype, speculative decoding, draft tokens, and tensor parallelism at click time, but the settings sheet stays mounted during the load, so edits made while "Loading..." shows were silently ignored and then overwritten by the load response. Disable those controls while the staged pick is loading so the visible state matches what the run actually uses. * Studio: refuse to stage a model while another load is in flight With Load on selection off, selecting a model mid-load staged it into pendingSelection and showed a disabled "Another model loading..." button, implying the user could retry once the load settled. The post-load cleanup then treated that queued pick as stale and silently cancelled and cleared it. Refuse to stage while a load (or a cancel's background unload) is in flight: stageModel no-ops in that state so every entry point (the chat selector and the Hub) is covered, and stageOrLoad surfaces a toast on the common path. The immediate-load path is unchanged; selectModel already rejects a concurrent load. * Tighten staged-load guard comments --------- Co-authored-by: Daniel Han --- .../frontend/src/features/chat/chat-page.tsx | 26 ++-- .../src/features/chat/chat-settings-sheet.tsx | 96 ++++++++++--- .../chat/hooks/use-chat-model-runtime.ts | 135 ++++++++++++------ .../hooks/use-staged-model-preparation.ts | 12 +- .../chat/stores/chat-runtime-store.ts | 24 +++- 5 files changed, 211 insertions(+), 82 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 758aa00e1a..83452c2ea8 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1607,6 +1607,14 @@ export function ChatPage({ await selectModel(selection); return; } + // Refuse staging while a load is in flight (it would be silently dropped); + // the immediate-load branch above is already guarded in selectModel. + if (store.modelLoading) { + toast.info("Another model is already loading", { + description: "Wait for it to finish or cancel it first.", + }); + return; + } // Tear down any existing staged pick first so its in-flight download is // cancelled, not left running after we rebind to the new pick. abandonStaged(); @@ -2495,6 +2503,7 @@ export function ChatPage({ ); }} externalProviderType={activeExternalProviderType} + loadingModel={loadingModel} onReloadModel={() => { const state = useChatRuntimeStore.getState(); if (state.params.checkpoint) { @@ -2512,22 +2521,23 @@ export function ChatPage({ if (!pending) return; const keyAtLoad = chatContextKey; // forceReload: the staged model isn't loaded yet, so bypass the - // same-checkpoint dedupe (and selectModel clears pendingSelection). - // keepSpeculative: honor the speculative mode set on the sidebar. + // same-checkpoint dedupe. keepSpeculative: honor the speculative mode + // set on the sidebar. void selectModel({ ...pending, forceReload: true, keepSpeculative: true, throwOnError: true, }).catch(() => { - // Recoverable failure (expired token, gated repo, OOM…): selectModel - // cleared the pick but left the edited knobs intact. + // Recoverable failure (expired token, gated repo, OOM…): the pick is + // cleared only on success, so it normally stays staged with edited + // knobs intact — nothing to restore. const store = useChatRuntimeStore.getState(); - // A pick staged meanwhile owns the knobs now; leave it untouched. + // Still staged (this pick, or a newer one queued meanwhile): leave it. if (store.pendingSelection) return; - // Restore (not re-stage, which would reset the knobs) only if the - // staged-load is still wanted: same chat context, sheet still open, - // page still mounted. + // Cleared mid-load (sheet closed / switched chats). Re-stage only if + // the staged-load is still wanted: same chat context, sheet still + // open, page still mounted. const stillWanted = mountedRef.current && store.settingsPanelOpen && diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index bf0cfcb277..77d7cf4c6d 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -52,6 +52,7 @@ import { SheetTitle, } from "@/components/ui/sheet"; import { Slider } from "@/components/ui/slider"; +import { Spinner } from "@/components/ui/spinner"; import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { InfoHint } from "@/components/ui/info-hint"; @@ -99,6 +100,7 @@ import { } from "./provider-capabilities"; import { isPendingGguf, + pendingSelectionMatches, useChatRuntimeStore, } from "./stores/chat-runtime-store"; import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section"; @@ -145,6 +147,7 @@ function NumericValueInput({ className, ariaLabel, size: sizeAttr, + disabled = false, }: { value: number; min?: number; @@ -155,6 +158,7 @@ function NumericValueInput({ className?: string; ariaLabel?: string; size?: number; + disabled?: boolean; }) { const [focused, setFocused] = useState(false); const [draft, setDraft] = useState(""); @@ -177,6 +181,7 @@ function NumericValueInput({ void; + /** The in-flight load (id + GGUF variant + native path token), or null when + * idle. Used to show a loading state for the staged pick only — not for an + * unrelated load or a cancel's background unload. */ + loadingModel?: { + id: string; + ggufVariant?: string | null; + nativePathToken?: string | null; + } | null; /** Loads the staged `pendingSelection` (deferred "Load on selection" flow). */ onLoadPendingModel?: () => void; /** Download progress (0–1) for a staged GGUF being fetched, or null when idle. */ @@ -457,6 +470,7 @@ export function ChatSettingsPanel({ onExternalProviderChange, externalProviderType = null, onReloadModel, + loadingModel = null, onLoadPendingModel, stagedDownloadFraction, onCancelStagedDownload, @@ -475,6 +489,19 @@ export function ChatSettingsPanel({ !isExternalModel || Boolean(providerCapabilities?.presencePenalty); const isMobile = useIsMobile(); const pendingSelection = useChatRuntimeStore((s) => s.pendingSelection); + // "Loading" only when the in-flight load IS this staged pick (full id + GGUF + // variant + native token match), not an unrelated load or a cancel's + // background unload. The variant matters: a different quant of the same repo + // staged mid-load must not read as this one loading. + const stagedLoading = + loadingModel != null && + pendingSelectionMatches(pendingSelection, { + id: loadingModel.id, + ggufVariant: loadingModel.ggufVariant, + nativePathToken: loadingModel.nativePathToken, + }); + // Load settings are snapshotted at click time; lock them while loading. + const modelControlsDisabled = stagedLoading; const abandonStagedModel = useChatRuntimeStore((s) => s.abandonStagedModel); const resetModelSettingsToLoaded = useChatRuntimeStore( (s) => s.resetModelSettingsToLoaded, @@ -867,10 +894,14 @@ export function ChatSettingsPanel({ {pendingSelection && ( - {stagedLabel} is staged, not loaded yet + {stagedLoading + ? `Loading ${stagedLabel}…` + : `${stagedLabel} is staged, not loaded yet`} - Set the options below, then choose Load model to load it. + {stagedLoading + ? "Applying your settings." + : "Set the options below, then choose Load model to load it."} )} @@ -898,6 +929,7 @@ export function ChatSettingsPanel({ }} ariaLabel="Context Length" size={8} + disabled={modelControlsDisabled} /> {ggufMaxContextLength != null && typeof ctxDisplayValue === "number" && @@ -944,6 +977,7 @@ export function ChatSettingsPanel({
{ setSpeculativeType(v); @@ -1057,6 +1092,7 @@ export function ChatSettingsPanel({
@@ -1118,31 +1155,46 @@ export function ChatSettingsPanel({ {Math.round((stagedDownloadFraction ?? 0) * 100)}%

)} -
+ {stagedLoading ? ( + // Mid-load: nothing to load or abandon until it settles, so disable. - -
+ ) : ( +
+ + +
+ )} ) : modelSettingsDirty ? (
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index f8a0eb579b..134878bdac 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -25,6 +25,7 @@ import { } from "../api/chat-api"; import { formatEta, formatRate } from "../utils/format-transfer"; import { + pendingSelectionMatches, readPersistedSpeculativeType, resolveToolsEnabledOnLoad, saveSpeculativeType, @@ -253,6 +254,7 @@ export function useChatModelRuntime() { displayName: string; isDownloaded?: boolean; isCachedLora?: boolean; + ggufVariant?: string | null; nativePathToken?: string | null; } | null>(null); const [loadToastDismissed, setLoadToastDismissed] = useState(false); @@ -399,29 +401,51 @@ export function useChatModelRuntime() { typeof selection === "string" ? false : selection.keepSpeculative ?? false; // Picking/loading any model abandons a staged (deferred) selection. // Before the early-returns below so even a no-op re-select clears the - // stage, and so the Load button unmounts on first click (no double-load). + // stage. const staged = useChatRuntimeStore.getState().pendingSelection; if (staged) { - // Loading a DIFFERENT model abandons this stage, so cancel its in-flight - // download. Loading the staged pick itself keeps it (that download feeds - // this load). - const loadingStagedPick = - staged.id === modelId && - (staged.ggufVariant ?? null) === (ggufVariant ?? null) && - (staged.nativePathToken ?? null) === (nativePathToken ?? null); - if (!loadingStagedPick) cancelStagedModelDownload(staged); - useChatRuntimeStore.getState().setPendingSelection(null); + // Loading a DIFFERENT model abandons this stage. Loading the staged pick + // ITSELF keeps it so the sidebar can show its load settings (context, KV + // cache, …) during the load. Cleared on success below; on failure it's + // left staged so the user can retry (see onLoadPendingModel's catch). + const loadingStagedPick = pendingSelectionMatches(staged, { + id: modelId, + ggufVariant, + nativePathToken, + }); + if (!loadingStagedPick) { + cancelStagedModelDownload(staged); + useChatRuntimeStore.getState().setPendingSelection(null); + } } const currentVariant = useChatRuntimeStore.getState().activeGgufVariant; if (!forceReload && (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null)))) { return; } - // Prevent duplicate loads if already loading this model - if ( - loadingModelRef.current?.id === modelId && - (loadingModelRef.current?.nativePathToken ?? null) === (nativePathToken ?? null) - ) + // A load is already in flight. If it's this exact pick (id + GGUF variant + + // native path token), ignore the duplicate click. If it's a DIFFERENT model + // -- crucially including a different GGUF variant of the same repo, which the + // old id+token-only guard wrongly treated as a duplicate and silently + // no-op'd -- don't start a second concurrent load (the load path has no clean + // supersession) and don't silently swallow the request: surface it so the + // user knows to wait for, or cancel, the in-flight load. Centralized here so + // every entry point is covered, not just the staged Load button. + const inFlightLoad = loadingModelRef.current; + if (inFlightLoad) { + const loadingSamePick = + inFlightLoad.id === modelId && + (inFlightLoad.ggufVariant ?? null) === (ggufVariant ?? null) && + (inFlightLoad.nativePathToken ?? null) === (nativePathToken ?? null); + if (loadingSamePick) return; + const message = + "Another model is already loading. Wait for it to finish or cancel it first."; + setModelsError(message); + if (throwOnError) throw new Error(message); + toast.info("Another model is already loading", { + description: "Wait for it to finish or cancel it first.", + }); return; + } const explicitIsLora = typeof selection === "string" ? undefined : selection.isLora; @@ -475,6 +499,7 @@ export function useChatModelRuntime() { displayName, isDownloaded, isCachedLora, + ggufVariant: ggufVariant ?? null, nativePathToken: nativePathToken ?? null, }; setLoadingModel(loadInfo); @@ -509,6 +534,21 @@ export function useChatModelRuntime() { stateBeforeUnload.modelRequiresTrustRemoteCode; const previousActiveNativePathToken = stateBeforeUnload.activeNativePathToken; + // Snapshot the load settings at click time, before the awaits below + // (validation, the trust dialog, unload). For a staged Load these knobs + // stay editable and a sheet-close revert (abandonStagedModel) can fire + // mid-load; reading them live just before loadModel would let the load + // use post-click values. The model-switch speculative reset below + // updates this snapshot in lock-step so non-staged loads are unchanged. + const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride; + const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype; + const loadCustomContextLength = stateBeforeUnload.customContextLength; + const loadGgufContextLength = stateBeforeUnload.ggufContextLength; + const loadTensorParallel = stateBeforeUnload.tensorParallel; + const loadActivePresetSource = stateBeforeUnload.activePresetSource; + const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant; + let loadSpeculativeType = stateBeforeUnload.speculativeType; + let loadSpecDraftNMax = stateBeforeUnload.specDraftNMax; try { // Lightweight pre-flight validation: avoid unloading a working model // if the new identifier is clearly invalid (e.g. bad HF id / path). @@ -517,18 +557,17 @@ export function useChatModelRuntime() { : undefined; // Validate with the same effective context /load uses: a GGUF native // context can exceed maxSeqLength, so sizing on raw maxSeqLength could - // pass, unload, then have /load refuse it. Read pre-unload state; load - // recomputes its own value, so this leaves loading untouched. - const preUnloadState = useChatRuntimeStore.getState(); + // pass, unload, then have /load refuse it. Uses the click-time + // snapshot (same values loadModel uses below), so the two agree. const validateMaxSeqLength = resolveLoadMaxSeqLength({ modelId, ggufVariant, - customContextLength: preUnloadState.customContextLength, - ggufContextLength: preUnloadState.ggufContextLength, + customContextLength: loadCustomContextLength, + ggufContextLength: loadGgufContextLength, currentCheckpoint, - activeGgufVariant: preUnloadState.activeGgufVariant, + activeGgufVariant: loadActiveGgufVariant, maxSeqLength, - presetSource: preUnloadState.activePresetSource, + presetSource: loadActivePresetSource, }); const validation = await validateModel({ model_path: modelId, @@ -586,32 +625,23 @@ export function useChatModelRuntime() { specDraftNMax: null, loadedSpecDraftNMax: null, }); + loadSpeculativeType = persistedSpeculativeType; + loadSpecDraftNMax = null; } - const { - chatTemplateOverride, - kvCacheDtype, - customContextLength, - ggufContextLength, - speculativeType, - specDraftNMax, - tensorParallel, - activePresetSource, - activeGgufVariant, - } = useChatRuntimeStore.getState(); const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ modelId, ggufVariant, isGguf, - customContextLength, - ggufContextLength, + customContextLength: loadCustomContextLength, + ggufContextLength: loadGgufContextLength, currentCheckpoint, - activeGgufVariant, + activeGgufVariant: loadActiveGgufVariant, maxSeqLength, - presetSource: activePresetSource, + presetSource: loadActivePresetSource, }); const effectiveChatTemplateOverride = - chatTemplateOverride?.trim() ? chatTemplateOverride : null; + loadChatTemplateOverride?.trim() ? loadChatTemplateOverride : null; const loadResponse = await loadModel({ model_path: modelId, nativePathLease: loadNativePathLease, @@ -623,10 +653,10 @@ export function useChatModelRuntime() { trust_remote_code: trustRemoteCode, approved_remote_code_fingerprint: approvedRemoteCodeFingerprint, chat_template_override: effectiveChatTemplateOverride, - cache_type_kv: kvCacheDtype, - speculative_type: speculativeType, - spec_draft_n_max: specDraftNMax, - tensor_parallel: tensorParallel, + cache_type_kv: loadKvCacheDtype, + speculative_type: loadSpeculativeType, + spec_draft_n_max: loadSpecDraftNMax, + tensor_parallel: loadTensorParallel, }); // If cancelled while loading, don't update UI to show @@ -636,7 +666,7 @@ export function useChatModelRuntime() { // The load applied this spec mode, so persist the user's standing // preference now (the requested intent, not the resolved echo; // saveSpeculativeType keeps only the universal auto/ngram/off). - saveSpeculativeType(speculativeType); + saveSpeculativeType(loadSpeculativeType); const currentParams = useChatRuntimeStore.getState().params; setParams( @@ -773,6 +803,25 @@ export function useChatModelRuntime() { } } await refresh({ signal: abortCtrl.signal }); + // A successful load owns the shared (pick-unscoped) settings fields, + // so any surviving stage is stale: the just-loaded pick itself, or a + // pick queued for a different model mid-load whose knobs this load + // overwrote. Drop it. Only a DIFFERENT pick's download needs + // cancelling; the loaded pick's is already consumed, and cancelling + // it inside its post-complete linger window would flicker its card. + const staleStage = useChatRuntimeStore.getState().pendingSelection; + if (staleStage) { + if ( + !pendingSelectionMatches(staleStage, { + id: modelId, + ggufVariant, + nativePathToken, + }) + ) { + cancelStagedModelDownload(staleStage); + } + useChatRuntimeStore.getState().setPendingSelection(null); + } } catch (error) { // Skip rollback if user cancelled -- model is already being unloaded. if (abortCtrl.signal.aborted) throw error; diff --git a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts index d154b931d4..b11f496aa5 100644 --- a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts +++ b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts @@ -10,6 +10,7 @@ import type { DownloadJob } from "@/features/hub/download-manager/use-repo-downl import { fetchGgufContextLength } from "../api/chat-api"; import { isPendingGguf, + pendingSelectionMatches, useChatRuntimeStore, } from "../stores/chat-runtime-store"; @@ -54,15 +55,12 @@ export function useStagedModelPreparation(): DownloadJob { nativePathToken, }); // Apply only if the same model is still staged (the user may have switched - // picks or loaded/cancelled while the request was in flight). Native ids - // are display labels, not paths, so two files can share an id -- compare - // the path token too, or a stale response could land on the wrong pick. + // picks or loaded/cancelled while the request was in flight). const latest = useChatRuntimeStore.getState().pendingSelection; if ( - latest?.id === id && - (latest.ggufVariant ?? null) === (ggufVariant ?? null) && - (latest.nativePathToken ?? null) === (nativePathToken ?? null) && - contextLength != null + latest && + contextLength != null && + pendingSelectionMatches(latest, { id, ggufVariant, nativePathToken }) ) { setPendingSelection({ ...latest, contextLength }); } diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index e2d8701695..d8e578e790 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -453,6 +453,22 @@ export function isPendingGguf(pending: PendingModelSelection | null): boolean { return pending != null && hasGgufSource(pending); } +/** Whether `pending` refers to the same model as `pick` (id + GGUF variant + + * native path token, optionals null-normalized). Native ids are display labels + * that can collide, so the token must match too — id alone can land on the + * wrong file. */ +export function pendingSelectionMatches( + pending: PendingModelSelection | null, + pick: { id: string; ggufVariant?: string | null; nativePathToken?: string | null }, +): boolean { + return ( + pending != null && + pending.id === pick.id && + (pending.ggufVariant ?? null) === (pick.ggufVariant ?? null) && + (pending.nativePathToken ?? null) === (pick.nativePathToken ?? null) + ); +} + type ChatRuntimeStore = { settingsHydrated: boolean; params: InferenceParams; @@ -1455,7 +1471,10 @@ export const useChatRuntimeStore = create((set, get) => ({ set({ loadOnSelection }); }, setPendingSelection: (pendingSelection) => set({ pendingSelection }), - stageModel: (selection) => + stageModel: (selection) => { + // Refuse staging mid-load: post-load cleanup would silently drop the queued + // pick. stageOrLoad toasts first for callers that can. + if (get().modelLoading) return; set((s) => { if ( s.pendingSelection && @@ -1475,7 +1494,8 @@ export const useChatRuntimeStore = create((set, get) => ({ speculativeType: readPersistedSpeculativeType(), specDraftNMax: null, }; - }), + }); + }, abandonStagedModel: () => { const { pendingSelection } = get(); if (!pendingSelection) return;