From e02662c309aafae70f997b97f18e37e6e31604ad Mon Sep 17 00:00:00 2001 From: Shine1i Date: Sun, 15 Feb 2026 18:48:02 +0100 Subject: [PATCH] feat: improve model loading/unloading UX and remove `_WarmupIndicator_` from thread UI - Refactored loading/unloading logic to provide detailed toast notifications with statuses (loading, success, error). - Removed unused `WarmupIndicator` component from thread UI to simplify interface. - Introduced better error handling for model refresh and inference tasks. --- .../src/components/assistant-ui/thread.tsx | 25 -------- .../src/features/chat/api/chat-adapter.ts | 61 +++++++++++++++++++ .../chat/hooks/use-chat-model-runtime.ts | 58 ++++++++++++------ 3 files changed, 101 insertions(+), 43 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 9c94e4eb22..c109c92ead 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -7,9 +7,7 @@ import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; -import { AnimatedShinyText } from "@/components/ui/animated-shiny-text"; import { Button } from "@/components/ui/button"; -import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store"; import { cn } from "@/lib/utils"; import { ActionBarMorePrimitive, @@ -73,7 +71,6 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ - !thread.isEmpty}> {!hideComposer && } @@ -83,28 +80,6 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ ); }; -const WarmupIndicator: FC = () => { - const threadId = useAuiState(({ threads }) => threads.mainThreadId); - const isRunning = useAuiState(({ thread }) => thread.isRunning); - const isWarmingUp = useChatRuntimeStore((state) => - Boolean(state.warmingByThreadId[threadId ?? "__default"]), - ); - - if (!isRunning || !isWarmingUp) { - return null; - } - - return ( -
-
- - Warming up model... - -
-
- ); -}; - const ThreadScrollToBottom: FC = () => { return ( diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index a408d53683..89cbfcaa47 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1,4 +1,5 @@ import type { ChatModelAdapter } from "@assistant-ui/react"; +import { toast } from "sonner"; import { streamChatCompletions } from "./chat-api"; import { db } from "../db"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; @@ -106,6 +107,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const { params } = state; if (!params.checkpoint) { + toast.error("No model loaded", { + description: "Pick model in top bar, then retry.", + }); throw new Error("Load a model first."); } @@ -126,6 +130,29 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const threadKey = unstable_threadId || "__default"; let waitingFirstChunk = true; + let hasResolvedFirstToken = false; + let resolveFirstToken: (() => void) | undefined; + let rejectFirstToken: ((err: unknown) => void) | undefined; + const firstTokenPromise = new Promise((resolve, reject) => { + resolveFirstToken = resolve; + rejectFirstToken = reject; + }); + // Avoid unhandled rejections if toast.promise never attached. + void firstTokenPromise.catch(() => {}); + let warmupToastShown = false; + const warmupDelayMs = 450; + const warmupTimer = setTimeout(() => { + if (!waitingFirstChunk || abortSignal.aborted) return; + warmupToastShown = true; + toast.promise(firstTokenPromise, { + loading: "Warming up model", + success: "Generating", + error: (err) => + err instanceof Error && err.message ? err.message : "Generation failed", + description: "Waiting for first token.", + duration: 900, + }); + }, warmupDelayMs); useChatRuntimeStore.getState().setThreadWarming(threadKey, true); useChatRuntimeStore.getState().setThreadRunning(threadKey, true); let cumulativeText = ""; @@ -157,6 +184,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (waitingFirstChunk) { waitingFirstChunk = false; useChatRuntimeStore.getState().setThreadWarming(threadKey, false); + if (!hasResolvedFirstToken) { + hasResolvedFirstToken = true; + resolveFirstToken?.(); + } } cumulativeText += delta; @@ -176,9 +207,39 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }; } } + if (!hasResolvedFirstToken) { + hasResolvedFirstToken = true; + resolveFirstToken?.(); + } + } catch (err) { + if (!hasResolvedFirstToken) { + hasResolvedFirstToken = true; + rejectFirstToken?.( + err instanceof Error ? err : new Error("Generation failed"), + ); + } + const isEarly = waitingFirstChunk; + if (!abortSignal.aborted && !(warmupToastShown && isEarly)) { + toast.error("Generation failed", { + description: err instanceof Error ? err.message : "Unknown error", + }); + } + throw err; } finally { + clearTimeout(warmupTimer); if (waitingFirstChunk) { useChatRuntimeStore.getState().setThreadWarming(threadKey, false); + if (warmupToastShown && !hasResolvedFirstToken) { + hasResolvedFirstToken = true; + rejectFirstToken?.( + abortSignal.aborted + ? new Error("Cancelled") + : new Error("No tokens received"), + ); + } else if (!hasResolvedFirstToken) { + hasResolvedFirstToken = true; + resolveFirstToken?.(); + } } useChatRuntimeStore.getState().setThreadRunning(threadKey, false); } 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 3eb9301302..febf6da367 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 @@ -105,6 +105,9 @@ export function useChatModelRuntime() { const message = error instanceof Error ? error.message : "Failed to load models"; setModelsError(message); + toast.error("Failed to refresh models", { + description: message, + }); } }, [setCheckpoint, setLoras, setModels, setModelsError]); @@ -122,30 +125,38 @@ export function useChatModelRuntime() { const isLora = explicitIsLora ?? model?.isLora ?? (lora ? true : false); const displayName = model?.name || lora?.name || modelId; - const loadingToastId = toast.loading(`Loading ${displayName}...`); setModelsError(null); try { - if (params.checkpoint) { - await unloadModel({ model_path: params.checkpoint }); - } + await toast.promise( + (async () => { + if (params.checkpoint) { + await unloadModel({ model_path: params.checkpoint }); + } - await loadModel({ - model_path: modelId, - hf_token: null, - max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, - load_in_4bit: true, - is_lora: isLora, - }); + await loadModel({ + model_path: modelId, + hf_token: null, + max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH, + load_in_4bit: true, + is_lora: isLora, + }); - setCheckpoint(modelId); - await refresh(); - toast.success(`${displayName} loaded`, { id: loadingToastId }); + setCheckpoint(modelId); + await refresh(); + })(), + { + loading: `Loading ${displayName}`, + success: `${displayName} loaded`, + error: (err) => + err instanceof Error ? err.message : "Failed to load model", + description: isLora ? "Fine-tuned (LoRA) selected." : "Base model selected.", + }, + ); } catch (error) { const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); - toast.error(message, { id: loadingToastId }); } }, [loras, models, params.checkpoint, refresh, setCheckpoint, setModelsError], @@ -157,9 +168,20 @@ export function useChatModelRuntime() { } setModelsError(null); try { - await unloadModel({ model_path: params.checkpoint }); - clearCheckpoint(); - await refresh(); + await toast.promise( + (async () => { + await unloadModel({ model_path: params.checkpoint }); + clearCheckpoint(); + await refresh(); + })(), + { + loading: "Unloading model", + success: "Model unloaded", + error: (err) => + err instanceof Error ? err.message : "Failed to unload model", + description: "Releases VRAM and resets inference state.", + }, + ); } catch (error) { const message = error instanceof Error ? error.message : "Failed to unload model";