From 475ba417dc6194ca6667be7b2fe570140aff1f54 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 15 Mar 2026 07:47:43 +0000
Subject: [PATCH] studio: context-aware loading text + download progress bar

1. Loading text: shows "Loading model..." for cached models,
   "Downloading model..." for new downloads. Toast description
   adapts accordingly.

2. Download progress: polls /api/models/gguf-download-progress every
   2s during downloads, updating the toast with percentage and GB
   downloaded. Progress is estimated by checking the HF cache folder
   size against the expected total bytes.

3. Passes isDownloaded and expectedBytes through the full chain from
   variant click to selectModel for accurate UI state.
---
 studio/backend/routes/models.py               | 40 ++++++++
 .../assistant-ui/model-selector/pickers.tsx   |  6 +-
 .../assistant-ui/model-selector/types.ts      |  2 +
 .../src/features/chat/api/chat-api.ts         | 12 +++
 .../frontend/src/features/chat/chat-page.tsx  | 10 +-
 .../chat/hooks/use-chat-model-runtime.ts      | 95 ++++++++++++++++---
 6 files changed, 145 insertions(+), 20 deletions(-)

diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py
index 1dee85ed82..9d92280274 100644
--- a/studio/backend/routes/models.py
+++ b/studio/backend/routes/models.py
@@ -604,6 +604,46 @@ async def get_gguf_variants(
         )
 
 
+@router.get("/gguf-download-progress")
+async def get_gguf_download_progress(
+    repo_id: str = Query(..., description = "HuggingFace repo ID"),
+    expected_bytes: int = Query(0, description = "Expected total download size in bytes"),
+    current_subject: str = Depends(get_current_subject),
+):
+    """Return download progress by checking current size of cached GGUF files."""
+    import re as _re
+    try:
+        if not _re.fullmatch(r"[A-Za-z0-9._-]+/[A-Za-z0-9._-]+", repo_id):
+            return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0}
+
+        from huggingface_hub import constants as hf_constants
+
+        cache_dir = Path(hf_constants.HF_HUB_CACHE)
+        target = f"models--{repo_id.replace('/', '--')}".lower()
+        downloaded_bytes = 0
+        for entry in cache_dir.iterdir():
+            if entry.name.lower() == target:
+                # Sum .gguf files in snapshots + incomplete downloads in blobs
+                for f in entry.rglob("*.gguf"):
+                    downloaded_bytes += f.stat().st_size
+                # Also check incomplete downloads (blobs without extension)
+                blobs_dir = entry / "blobs"
+                if blobs_dir.is_dir():
+                    for f in blobs_dir.iterdir():
+                        if f.is_file():
+                            downloaded_bytes += f.stat().st_size
+                break
+
+        progress = min(downloaded_bytes / expected_bytes, 1.0) if expected_bytes > 0 else 0
+        return {
+            "downloaded_bytes": downloaded_bytes,
+            "expected_bytes": expected_bytes,
+            "progress": round(progress, 3),
+        }
+    except Exception:
+        return {"downloaded_bytes": 0, "expected_bytes": expected_bytes, "progress": 0}
+
+
 @router.get("/cached-gguf")
 async def list_cached_gguf(
     current_subject: str = Depends(get_current_subject),
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index a37532b8e2..1c52d53915 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -184,11 +184,13 @@ function GgufVariantExpander({
   }, [repoId]);
 
   const handleVariantClick = useCallback(
-    (quant: string) => {
+    (quant: string, downloaded?: boolean, sizeBytes?: number) => {
       onSelect(repoId, {
         source: "hub",
         isLora: false,
         ggufVariant: quant,
+        isDownloaded: downloaded,
+        expectedBytes: sizeBytes,
       });
     },
     [repoId, onSelect],
@@ -294,7 +296,7 @@ function GgufVariantExpander({
           <button
             key={v.filename}
             type="button"
-            onClick={() => handleVariantClick(v.quant)}
+            onClick={() => handleVariantClick(v.quant, v.downloaded, v.size_bytes)}
             className={cn(
               "flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1 text-left text-sm transition-colors hover:bg-accent",
             )}
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts
index 75efc67830..9da7fd975f 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts
+++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts
@@ -21,5 +21,7 @@ export interface ModelSelectorChangeMeta {
   source: "hub" | "lora" | "exported";
   isLora: boolean;
   ggufVariant?: string;
+  isDownloaded?: boolean;
+  expectedBytes?: number;
 }
 
diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts
index a526366782..5831642606 100644
--- a/studio/frontend/src/features/chat/api/chat-api.ts
+++ b/studio/frontend/src/features/chat/api/chat-api.ts
@@ -103,6 +103,18 @@ export interface CachedGgufRepo {
   cache_path: string;
 }
 
+export async function getGgufDownloadProgress(
+  repoId: string,
+  expectedBytes: number,
+): Promise<{ downloaded_bytes: number; expected_bytes: number; progress: number }> {
+  const params = new URLSearchParams({
+    repo_id: repoId,
+    expected_bytes: String(expectedBytes),
+  });
+  const response = await authFetch(`/api/models/gguf-download-progress?${params}`);
+  return parseJsonOrThrow(response);
+}
+
 export async function listCachedGguf(): Promise<CachedGgufRepo[]> {
   const response = await authFetch("/api/models/cached-gguf");
   const data = await parseJsonOrThrow<{ cached: CachedGgufRepo[] }>(response);
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 9a086dd1d4..0e06d89b28 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -337,7 +337,7 @@ export function ChatPage(): ReactElement {
   }, [inferenceParams.checkpoint, lorasFromStore]);
 
   const handleCheckpointChange = useCallback(
-    (value: string, meta?: { isLora: boolean; ggufVariant?: string }) => {
+    (value: string, meta?: { isLora: boolean; ggufVariant?: string; isDownloaded?: boolean; expectedBytes?: number }) => {
       const store = useChatRuntimeStore.getState();
       const currentCheckpoint = store.params.checkpoint;
       const currentVariant = store.activeGgufVariant;
@@ -374,6 +374,8 @@ export function ChatPage(): ReactElement {
           id: value,
           isLora: meta?.isLora,
           ggufVariant: meta?.ggufVariant,
+          isDownloaded: meta?.isDownloaded,
+          expectedBytes: meta?.expectedBytes,
         });
       })();
     },
@@ -607,11 +609,13 @@ export function ChatPage(): ReactElement {
               {loadingModel ? (
                 <div
                   className="flex items-center gap-1.5 text-muted-foreground"
-                  title={`Loading ${loadingModel.displayName}. This may include downloading.`}
+                  title={loadingModel.isDownloaded
+                    ? `Loading ${loadingModel.displayName} from cache.`
+                    : `Loading ${loadingModel.displayName}. This may include downloading.`}
                 >
                   <Spinner className="size-3.5 shrink-0" />
                   <span className="text-xs">
-                    Downloading model…
+                    {loadingModel.isDownloaded ? "Loading model…" : "Downloading model…"}
                   </span>
                   <button
                     type="button"
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 c775ace084..07a1208287 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
@@ -4,6 +4,7 @@
 import { useCallback, useState } from "react";
 import { toast } from "sonner";
 import {
+  getGgufDownloadProgress,
   getInferenceStatus,
   listLoras,
   listModels,
@@ -24,6 +25,8 @@ type SelectedModelInput = {
   isLora?: boolean;
   ggufVariant?: string;
   loadingDescription?: string;
+  isDownloaded?: boolean;
+  expectedBytes?: number;
 };
 
 const LORA_SUFFIX_RE = /_(\d{9,})$/;
@@ -146,6 +149,7 @@ export function useChatModelRuntime() {
   const [loadingModel, setLoadingModel] = useState<{
     id: string;
     displayName: string;
+    isDownloaded?: boolean;
   } | null>(null);
   const [loadAbortController, setLoadAbortController] =
     useState<AbortController | null>(null);
@@ -189,6 +193,8 @@ export function useChatModelRuntime() {
         typeof selection === "string" ? undefined : selection.isLora;
       const extraLoadingDescription =
         typeof selection === "string" ? undefined : selection.loadingDescription;
+      const isDownloaded =
+        typeof selection === "string" ? false : selection.isDownloaded ?? false;
       const model = models.find((entry) => entry.id === modelId);
       const lora = loras.find((entry) => entry.id === modelId);
       const isLora =
@@ -210,13 +216,15 @@ export function useChatModelRuntime() {
       const loadingDescription = [
         currentCheckpoint ? "Unloading previous model first." : null,
         extraLoadingDescription ?? null,
-        "This may include downloading. Large models can take a while.",
+        isDownloaded
+          ? "Loading cached model into memory."
+          : "This may include downloading. Large models can take a while.",
       ]
         .filter(Boolean)
         .join(" ");
 
       setModelsError(null);
-      setLoadingModel({ id: modelId, displayName });
+      setLoadingModel({ id: modelId, displayName, isDownloaded });
       const abortCtrl = new AbortController();
       setLoadAbortController(abortCtrl);
       try {
@@ -280,21 +288,77 @@ export function useChatModelRuntime() {
           }
         }
 
-        const toastId = toast.loading("Loading model…", {
-          description: loadingDescription,
-          action: {
-            label: "Cancel",
-            onClick: () => {
-              abortCtrl.abort();
-              setLoadingModel(null);
-              setLoadAbortController(null);
-              unloadModel({ model_path: modelId }).catch(() => {});
-              clearCheckpoint();
-              toast.dismiss(toastId);
-              toast.info("Model loading cancelled");
+        const toastId = toast.loading(
+          isDownloaded ? "Loading model…" : "Downloading model…",
+          {
+            description: loadingDescription,
+            duration: Infinity,
+            action: {
+              label: "Cancel",
+              onClick: () => {
+                abortCtrl.abort();
+                setLoadingModel(null);
+                setLoadAbortController(null);
+                unloadModel({ model_path: modelId }).catch(() => {});
+                clearCheckpoint();
+                toast.dismiss(toastId);
+                toast.info("Model loading cancelled");
+              },
             },
           },
-        });
+        );
+
+        // Poll download progress for non-cached models
+        let progressInterval: ReturnType<typeof setInterval> | null = null;
+        if (!isDownloaded && ggufVariant) {
+          const expectedBytes =
+            typeof selection !== "string" ? selection.expectedBytes ?? 0 : 0;
+          if (expectedBytes > 0) {
+            progressInterval = setInterval(async () => {
+              if (abortCtrl.signal.aborted) {
+                if (progressInterval) clearInterval(progressInterval);
+                return;
+              }
+              try {
+                const prog = await getGgufDownloadProgress(modelId, expectedBytes);
+                if (prog.progress > 0 && prog.progress < 1) {
+                  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}%`,
+                    {
+                      id: toastId,
+                      description: `${dlGb.toFixed(1)} / ${totalGb.toFixed(1)} GB`,
+                      duration: Infinity,
+                      action: {
+                        label: "Cancel",
+                        onClick: () => {
+                          abortCtrl.abort();
+                          setLoadingModel(null);
+                          setLoadAbortController(null);
+                          unloadModel({ model_path: modelId }).catch(() => {});
+                          clearCheckpoint();
+                          toast.dismiss(toastId);
+                          toast.info("Model loading cancelled");
+                        },
+                      },
+                    },
+                  );
+                } else if (prog.progress >= 1) {
+                  toast.loading("Loading model…", {
+                    id: toastId,
+                    description: "Download complete. Starting inference server…",
+                    duration: Infinity,
+                  });
+                  if (progressInterval) clearInterval(progressInterval);
+                }
+              } catch {
+                // Ignore polling errors
+              }
+            }, 2000);
+          }
+        }
 
         try {
           await performLoad();
@@ -308,6 +372,7 @@ export function useChatModelRuntime() {
           }
           throw err;
         } finally {
+          if (progressInterval) clearInterval(progressInterval);
           setLoadingModel(null);
           setLoadAbortController(null);
         }