From 9cbeecc16aaf75b28c17b499bcd5b9606e58f2f8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 16 Mar 2026 06:51:52 +0000 Subject: [PATCH] Incorporate PR #4304 toast UX improvements Merge the toast UX refactor from PR #4304 (by @Shine1i): - Toast duration 5s default with close button (X) for manual dismiss - Inline progress bar component (ModelLoadInlineStatus) shown in the header after toast is dismissed - Model switch warning only for image compatibility (not generic) - activeThreadId tracked in store via ActiveThreadSync - Loading state cleanup via resetLoadingUi helper - Toast uses Infinity duration during loading with onDismiss handler Re-applied non-GGUF download progress additions on top: - getDownloadProgress for all models (not just GGUF) - hasShownProgress flag, loadingModelRef race condition checks - First poll at 500ms, bytes-only fallback when expected size unknown --- studio/frontend/src/components/ui/sonner.tsx | 10 +- .../frontend/src/features/chat/chat-page.tsx | 103 +++---- .../chat/components/model-load-status.tsx | 106 +++++++ .../chat/hooks/use-chat-model-runtime.ts | 286 ++++++++++++------ .../src/features/chat/runtime-provider.tsx | 17 ++ .../chat/stores/chat-runtime-store.ts | 4 + 6 files changed, 359 insertions(+), 167 deletions(-) create mode 100644 studio/frontend/src/features/chat/components/model-load-status.tsx diff --git a/studio/frontend/src/components/ui/sonner.tsx b/studio/frontend/src/components/ui/sonner.tsx index a20f106d42..f211c4acd0 100644 --- a/studio/frontend/src/components/ui/sonner.tsx +++ b/studio/frontend/src/components/ui/sonner.tsx @@ -16,11 +16,11 @@ const Toaster = ({ ...props }: ToasterProps) => { const { theme = "system" } = useTheme(); return ( - { - if (view.mode !== "single") { - return undefined; - } - if (view.threadId) { - return view.threadId; - } - - // New-thread flow keeps threadId undefined in local view state. - // Fall back to most recent regular base thread. - const candidates = await db.threads.where("modelType").equals("base").toArray(); - const latest = candidates - .filter((thread) => !thread.archived && !thread.pairId) - .sort((a, b) => b.createdAt - a.createdAt)[0]; - return latest?.id; -} - const SingleContent = memo(function SingleContent({ threadId, newThreadNonce, @@ -321,7 +304,16 @@ export function ChatPage(): ReactElement { const modelsFromStore = useChatRuntimeStore((state) => state.models); const lorasFromStore = useChatRuntimeStore((state) => state.loras); const modelsError = useChatRuntimeStore((state) => state.modelsError); - const { refresh, selectModel, ejectModel, cancelLoading, loadingModel } = + const activeThreadId = useChatRuntimeStore((state) => state.activeThreadId); + const { + refresh, + selectModel, + ejectModel, + cancelLoading, + loadingModel, + loadProgress, + loadToastDismissed, + } = useChatModelRuntime(); const refreshRef = useRef(refresh); const selectModelRef = useRef(selectModel); @@ -343,42 +335,27 @@ export function ChatPage(): ReactElement { const currentVariant = store.activeGgufVariant; if (!value || (value === currentCheckpoint && (meta?.ggufVariant ?? null) === (currentVariant ?? null))) return; void (async () => { - let switchNote: string | undefined; - const activeThreadId = await resolveActiveSingleThreadId(view); - if (activeThreadId) { + let showImageCompatibilityWarning = false; + if (view.mode === "single" && activeThreadId) { const thread = await db.threads.get(activeThreadId); if (thread?.modelId && thread.modelId !== value) { const messages = await db.messages .where("threadId") .equals(activeThreadId) .toArray(); - if (messages.length === 0) { - // No history -- just switch silently - await db.threads.update(activeThreadId, { modelId: value }); - await selectModel({ - id: value, - isLora: meta?.isLora, - ggufVariant: meta?.ggufVariant, - isDownloaded: meta?.isDownloaded, - expectedBytes: meta?.expectedBytes, - }); - return; + if (messages.length > 0) { + const hasImage = messages.some(messageHasImage); + const targetModel = modelsFromStore.find((model) => model.id === value); + showImageCompatibilityWarning = + hasImage && targetModel?.isVision === false; } - const hasImage = messages.some(messageHasImage); - const targetModel = modelsFromStore.find((model) => model.id === value); - const nonVisionWithImages = hasImage && targetModel?.isVision === false; - - switchNote = nonVisionWithImages - ? "Full chat history will be sent to the new model. This chat has images; text-only models may fail." - : hasImage - ? "Full chat history will be sent to the new model. This chat includes images." - : "Full chat history will be sent to the new model."; } } - if (switchNote) { - toast.warning("Model changed for this chat", { - description: switchNote, + if (showImageCompatibilityWarning) { + toast.warning("Selected model may not handle earlier images", { + description: + "This chat already includes images. Text-only models can ignore them or fail on follow-up replies.", duration: 6000, }); } @@ -391,13 +368,16 @@ export function ChatPage(): ReactElement { }); })(); }, - [modelsFromStore, selectModel, view], + [activeThreadId, modelsFromStore, selectModel, view], ); const handleEject = useCallback(() => { void ejectModel(); }, [ejectModel]); const handleNewThread = useCallback( - () => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }), + () => { + useChatRuntimeStore.getState().setActiveThreadId(null); + setView({ mode: "single", newThreadNonce: crypto.randomUUID() }); + }, [], ); const handleNewCompare = useCallback( @@ -618,25 +598,22 @@ export function ChatPage(): ReactElement { contentDataTour="chat-model-selector-popover" className="max-w-[62vw] sm:max-w-none" /> - {loadingModel ? ( -
- - - {loadingModel.isDownloaded ? "Loading model…" : "Downloading model…"} - - -
+ progressPercent={loadProgress?.percent} + progressLabel={loadProgress?.label} + onStop={cancelLoading} + /> ) : null} {modelsError && ( diff --git a/studio/frontend/src/features/chat/components/model-load-status.tsx b/studio/frontend/src/features/chat/components/model-load-status.tsx new file mode 100644 index 0000000000..9686011aca --- /dev/null +++ b/studio/frontend/src/features/chat/components/model-load-status.tsx @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { Progress } from "@/components/ui/progress"; +import { Spinner } from "@/components/ui/spinner"; +import { Button } from "@/components/ui/button"; + +type ModelLoadDescriptionProps = { + message?: string | null; + progressPercent?: number | null; + progressLabel?: string | null; + onStop?: () => void; + stopLabel?: string; +}; + +function clampProgress(value: number): number { + return Math.max(0, Math.min(100, value)); +} + +export function ModelLoadDescription({ + message, + progressPercent, + progressLabel, + onStop, + stopLabel = "Stop loading", +}: ModelLoadDescriptionProps) { + const hasProgress = typeof progressPercent === "number"; + + return ( +
+
+ {hasProgress ? ( +
+
+ {progressLabel} + {Math.round(clampProgress(progressPercent))}% +
+ +
+ ) : message ? ( +

{message}

+ ) : null} +
+ {onStop ? ( + + ) : null} +
+ ); +} + +type ModelLoadInlineStatusProps = { + label: string; + title: string; + progressPercent?: number | null; + progressLabel?: string | null; + onStop?: () => void; +}; + +export function ModelLoadInlineStatus({ + label, + title, + progressPercent, + progressLabel, + onStop, +}: ModelLoadInlineStatusProps) { + const hasProgress = typeof progressPercent === "number"; + + return ( +
+
+ + {label} +
+ {hasProgress ? ( +
+
+ +
+
+ {progressLabel} + {Math.round(clampProgress(progressPercent))}% +
+
+ ) : null} + {onStop ? ( + + ) : null} +
+ ); +} 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 f1d0e16d0c..e5c58c7be7 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 @@ -1,8 +1,10 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { useCallback, useRef, useState } from "react"; +import { createElement, useCallback, useRef, useState } from "react"; import { toast } from "sonner"; +import { Spinner } from "@/components/ui/spinner"; +import { ModelLoadDescription } from "../components/model-load-status"; import { getDownloadProgress, getGgufDownloadProgress, @@ -30,6 +32,15 @@ type SelectedModelInput = { expectedBytes?: number; }; +const MODEL_LOAD_TOAST_CLASSNAMES = { + toast: "items-start gap-2.5 pr-8", + content: "gap-0.5", + title: "leading-5", + description: "mt-0", + closeButton: + "!left-auto !right-1 !top-2 !translate-x-0 !translate-y-0 !border-transparent !bg-transparent !shadow-none hover:!bg-transparent hover:opacity-70", +} as const; + const LORA_SUFFIX_RE = /_(\d{9,})$/; function parseTrailingEpoch(input: string): number | undefined { @@ -152,11 +163,46 @@ export function useChatModelRuntime() { displayName: string; isDownloaded?: boolean; } | null>(null); - const [_loadAbortController, setLoadAbortController] = - useState(null); + const [loadToastDismissed, setLoadToastDismissed] = useState(false); + const [loadProgress, setLoadProgress] = useState<{ + percent: number | null; + label: string | null; + phase: "downloading" | "starting"; + } | null>(null); const loadAbortRef = useRef(null); const loadingModelRef = useRef(null); const loadToastIdRef = useRef(null); + const loadToastDismissedRef = useRef(false); + + const setLoadToastDismissedState = useCallback((dismissed: boolean) => { + loadToastDismissedRef.current = dismissed; + setLoadToastDismissed(dismissed); + }, []); + + const resetLoadingUi = useCallback(() => { + setLoadingModel(null); + setLoadProgress(null); + loadingModelRef.current = null; + loadAbortRef.current = null; + loadToastIdRef.current = null; + setLoadToastDismissedState(false); + }, [setLoadToastDismissedState]); + + const renderLoadDescription = useCallback( + ( + message: string, + progressPercent?: number | null, + progressLabel?: string | null, + onStop?: () => void, + ) => + createElement(ModelLoadDescription, { + message, + progressPercent, + progressLabel, + onStop, + }), + [], + ); const refresh = useCallback(async () => { setModelsError(null); @@ -183,6 +229,26 @@ export function useChatModelRuntime() { } }, [setCheckpoint, setLoras, setModels, setModelsError]); + const cancelLoading = useCallback(() => { + const model = loadingModelRef.current; + if (!model) return; + loadAbortRef.current?.abort(); + loadAbortRef.current = null; + loadingModelRef.current = null; + const tid = loadToastIdRef.current; + loadToastIdRef.current = null; + setLoadingModel(null); + setLoadProgress(null); + setLoadToastDismissedState(false); + clearCheckpoint(); + if (tid != null) toast.dismiss(tid); + toast.info("Stopped loading model", { + description: "The current download may still finish in the background.", + }); + // Fire-and-forget: tell backend to stop, don't block UI + unloadModel({ model_path: model.id }).catch(() => {}); + }, [clearCheckpoint, setLoadToastDismissedState]); + const selectModel = useCallback( async (selection: string | SelectedModelInput) => { const modelId = typeof selection === "string" ? selection : selection.id; @@ -218,21 +284,23 @@ export function useChatModelRuntime() { const previousIsLora = previousModel?.isLora ?? (previousLora ? true : false); const loadingDescription = [ - currentCheckpoint ? "Unloading previous model first." : null, + currentCheckpoint ? "Switching models." : null, extraLoadingDescription ?? null, - isDownloaded - ? "Loading cached model into memory." - : "Downloading and loading model. Large models can take a while.", + isDownloaded ? "Loading cached model into memory." : null, ] .filter(Boolean) .join(" "); - setModelsError(null); + setLoadToastDismissedState(false); const loadInfo = { id: modelId, displayName, isDownloaded }; setLoadingModel(loadInfo); + setLoadProgress( + isDownloaded + ? { percent: null, label: null, phase: "starting" } + : { percent: 0, label: "Preparing download", phase: "downloading" }, + ); loadingModelRef.current = loadInfo; const abortCtrl = new AbortController(); - setLoadAbortController(abortCtrl); loadAbortRef.current = abortCtrl; try { async function performLoad(): Promise { @@ -301,56 +369,37 @@ export function useChatModelRuntime() { } } - const toastId = toast.loading( - isDownloaded ? "Loading model…" : "Downloading model…", + const toastId = toast( + isDownloaded ? "Starting model…" : "Downloading model…", { - description: loadingDescription, - duration: 10000, - action: { - label: "Cancel", - onClick: () => { - abortCtrl.abort(); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; - unloadModel({ model_path: modelId }).catch(() => {}); - clearCheckpoint(); - toast.dismiss(toastId); - toast.info("Model loading cancelled"); - }, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + loadingDescription, + isDownloaded ? null : 0, + isDownloaded ? null : "Preparing download", + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) { + return; + } + setLoadToastDismissedState(true); }, }, ); loadToastIdRef.current = toastId; - // Poll download progress for non-cached models + // Poll download progress for non-cached models (GGUF and non-GGUF) let progressInterval: ReturnType | null = null; if (!isDownloaded) { const expectedBytes = typeof selection !== "string" ? selection.expectedBytes ?? 0 : 0; - - const cancelAction = { - label: "Cancel", - onClick: () => { - abortCtrl.abort(); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; - unloadModel({ model_path: modelId }).catch(() => {}); - clearCheckpoint(); - toast.dismiss(toastId); - toast.info("Model loading cancelled"); - }, - }; - let hasShownProgress = false; const pollProgress = async () => { - // Stop if cancelled or if loading already finished if (abortCtrl.signal.aborted || !loadingModelRef.current) { if (progressInterval) clearInterval(progressInterval); return; @@ -360,7 +409,6 @@ export function useChatModelRuntime() { ? await getGgufDownloadProgress(modelId, ggufVariant, expectedBytes) : await getDownloadProgress(modelId); - // Re-check after await -- load may have finished while polling if (!loadingModelRef.current) return; if (prog.progress > 0 && prog.progress < 1) { @@ -368,36 +416,69 @@ export function useChatModelRuntime() { const dlGb = prog.downloaded_bytes / (1024 ** 3); const totalGb = prog.expected_bytes / (1024 ** 3); const pct = Math.round(prog.progress * 100); - toast.loading( - `Downloading model... ${pct}%`, + const progressLabel = totalGb > 0 + ? `${dlGb.toFixed(1)} of ${totalGb.toFixed(1)} GB` + : `${dlGb.toFixed(1)} GB downloaded`; + setLoadProgress({ + percent: pct, + label: progressLabel, + phase: "downloading", + }); + if (loadToastDismissedRef.current) return; + toast( + "Downloading model…", { id: toastId, - description: totalGb > 0 - ? `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB` - : `${dlGb.toFixed(1)} GB downloaded`, - duration: 10000, - action: cancelAction, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + loadingDescription, + pct, + progressLabel, + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) return; + setLoadToastDismissedState(true); + }, }, ); } else if (prog.downloaded_bytes > 0 && prog.expected_bytes === 0 && prog.progress === 0) { - // Have bytes but no total size -- show bytes only hasShownProgress = true; const dlGb = prog.downloaded_bytes / (1024 ** 3); - toast.loading( - "Downloading model...", - { - id: toastId, - description: `${dlGb.toFixed(1)} GB downloaded`, - duration: 10000, - action: cancelAction, - }, - ); + setLoadProgress({ + percent: null, + label: `${dlGb.toFixed(1)} GB downloaded`, + phase: "downloading", + }); } else if (prog.progress >= 1 && hasShownProgress) { - // Only show "download complete" if we actually showed progress - toast.loading("Loading model...", { + setLoadProgress({ + percent: 100, + label: "Download complete", + phase: "starting", + }); + if (loadToastDismissedRef.current) { + if (progressInterval) clearInterval(progressInterval); + return; + } + toast("Starting model…", { id: toastId, - description: "Download complete. Loading into memory...", - duration: 10000, + icon: createElement(Spinner, { className: "size-4" }), + description: renderLoadDescription( + "Download complete. Loading the model into memory.", + 100, + "Download complete", + cancelLoading, + ), + duration: Infinity, + closeButton: true, + classNames: MODEL_LOAD_TOAST_CLASSNAMES, + onDismiss: (dismissedToast) => { + if (loadToastIdRef.current !== dismissedToast.id) return; + setLoadToastDismissedState(true); + }, }); if (progressInterval) clearInterval(progressInterval); } @@ -406,40 +487,62 @@ export function useChatModelRuntime() { } }; - // First poll after 500ms, then every 2s setTimeout(pollProgress, 500); progressInterval = setInterval(pollProgress, 2000); } try { await performLoad(); - toast.success(`${displayName} loaded`, { id: toastId }); + if (loadToastDismissedRef.current) { + toast.success(`${displayName} loaded`); + } else { + toast.success(`${displayName} loaded`, { + id: toastId, + description: undefined, + closeButton: false, + duration: 2000, + }); + } } catch (err) { if (!abortCtrl.signal.aborted) { - toast.error( - err instanceof Error ? err.message : "Failed to load model", - { id: toastId }, - ); + const message = + err instanceof Error ? err.message : "Failed to load model"; + if (loadToastDismissedRef.current) { + toast.error(message); + } else { + toast.error(message, { + id: toastId, + description: undefined, + closeButton: false, + duration: 5000, + }); + } } throw err; } finally { if (progressInterval) clearInterval(progressInterval); - setLoadingModel(null); - setLoadAbortController(null); - loadingModelRef.current = null; - loadAbortRef.current = null; - loadToastIdRef.current = null; + resetLoadingUi(); } } catch (error) { if (abortCtrl.signal.aborted) return; // User cancelled, nothing to report - setLoadingModel(null); - loadingModelRef.current = null; + resetLoadingUi(); const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); } }, - [loras, models, params.checkpoint, refresh, setModelsError, setParams], + [ + cancelLoading, + loras, + models, + params.checkpoint, + refresh, + renderLoadDescription, + resetLoadingUi, + setLoadToastDismissedState, + setModelsError, + setParams, + ], ); const ejectModel = useCallback(async () => { @@ -468,28 +571,13 @@ export function useChatModelRuntime() { } }, [clearCheckpoint, params.checkpoint, refresh, setModelsError]); - const cancelLoading = useCallback(() => { - const model = loadingModelRef.current; - if (!model) return; - loadAbortRef.current?.abort(); - loadAbortRef.current = null; - loadingModelRef.current = null; - const tid = loadToastIdRef.current; - loadToastIdRef.current = null; - setLoadingModel(null); - setLoadAbortController(null); - clearCheckpoint(); - if (tid != null) toast.dismiss(tid); - toast.info("Model loading cancelled"); - // Fire-and-forget: tell backend to stop, don't block UI - unloadModel({ model_path: model.id }).catch(() => {}); - }, [clearCheckpoint]); - return { refresh, selectModel, ejectModel, cancelLoading, loadingModel, + loadProgress, + loadToastDismissed, }; } diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index dbc631a1d5..0ad6005509 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -559,6 +559,22 @@ function ThreadNewChatSwitch({ return null; } +function ActiveThreadSync({ + enabled, +}: { enabled: boolean }): ReactElement | null { + const mainThreadId = useAuiState(({ threads }) => threads.mainThreadId); + const setActiveThreadId = useChatRuntimeStore((state) => state.setActiveThreadId); + + useEffect(() => { + if (!enabled) { + return; + } + setActiveThreadId(mainThreadId ?? null); + }, [enabled, mainThreadId, setActiveThreadId]); + + return null; +} + export function ChatRuntimeProvider({ children, modelType = "base", @@ -586,6 +602,7 @@ export function ChatRuntimeProvider({ return ( + {initialThreadId && } {!initialThreadId && newThreadNonce && ( 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 87e162ed56..07d3aecc36 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -43,6 +43,7 @@ type ChatRuntimeStore = { autoTitle: boolean; modelsError: string | null; activeGgufVariant: string | null; + activeThreadId: string | null; pendingAudioBase64: string | null; pendingAudioName: string | null; setParams: (params: InferenceParams) => void; @@ -52,6 +53,7 @@ type ChatRuntimeStore = { setAutoTitle: (enabled: boolean) => void; setModelsError: (error: string | null) => void; setCheckpoint: (modelId: string, ggufVariant?: string | null) => void; + setActiveThreadId: (threadId: string | null) => void; clearCheckpoint: () => void; setPendingAudio: (base64: string, name: string) => void; clearPendingAudio: () => void; @@ -65,6 +67,7 @@ export const useChatRuntimeStore = create((set) => ({ autoTitle: loadBool(AUTO_TITLE_KEY, false), modelsError: null, activeGgufVariant: null, + activeThreadId: null, pendingAudioBase64: null, pendingAudioName: null, setParams: (params) => set({ params }), @@ -94,6 +97,7 @@ export const useChatRuntimeStore = create((set) => ({ }, activeGgufVariant: ggufVariant ?? null, })), + setActiveThreadId: (activeThreadId) => set({ activeThreadId }), clearCheckpoint: () => set((state) => ({ params: {