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.

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 }> = ({
 
         <ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 mt-auto flex w-full flex-col gap-4 overflow-visible bg-background pb-4 md:pb-4 before:pointer-events-none before:absolute before:inset-x-0 before:bottom-full before:h-20 before:bg-gradient-to-t before:from-background before:to-transparent">
           <ThreadScrollToBottom />
-          <WarmupIndicator />
           <AuiIf condition={({ thread }) => !thread.isEmpty}>
             {!hideComposer && <ComposerAnimated />}
           </AuiIf>
@@ -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 (
-    <div className="mx-auto -mb-2 w-full max-w-(--thread-max-width) px-2">
-      <div className="inline-flex items-center rounded-full border border-border/60 bg-background/90 px-3 py-1.5 text-xs text-muted-foreground shadow-sm">
-        <AnimatedShinyText className="text-xs">
-          Warming up model...
-        </AnimatedShinyText>
-      </div>
-    </div>
-  );
-};
-
 const ThreadScrollToBottom: FC = () => {
   return (
     <ThreadPrimitive.ScrollToBottom asChild={true}>

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<void>((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);
       }

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";