From f5f631e5d1b9b6caaa9a4cf305150845a6856e95 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 15 Mar 2026 06:24:32 +0000
Subject: [PATCH] studio: add cancel button for model loading/downloading

Adds a Cancel button next to the "Downloading model..." spinner so
users can abort long downloads. Clicking it aborts the in-flight load,
calls unloadModel to kill any running llama-server process, and clears
the loading state.
---
 .../frontend/src/features/chat/chat-page.tsx  |  9 +++++++-
 .../chat/hooks/use-chat-model-runtime.ts      | 21 +++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 4d9cffc784..9a086dd1d4 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -321,7 +321,7 @@ 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, loadingModel } =
+  const { refresh, selectModel, ejectModel, cancelLoading, loadingModel } =
     useChatModelRuntime();
   const refreshRef = useRef(refresh);
   const selectModelRef = useRef(selectModel);
@@ -613,6 +613,13 @@ export function ChatPage(): ReactElement {
                   <span className="text-xs">
                     Downloading model…
                   </span>
+                  <button
+                    type="button"
+                    onClick={cancelLoading}
+                    className="ml-1 rounded px-1.5 py-0.5 text-[10px] font-medium text-muted-foreground hover:bg-destructive/10 hover:text-destructive transition-colors"
+                  >
+                    Cancel
+                  </button>
                 </div>
               ) : null}
             </div>
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 107b449ccd..3f8d41d743 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
@@ -147,6 +147,8 @@ export function useChatModelRuntime() {
     id: string;
     displayName: string;
   } | null>(null);
+  const [loadAbortController, setLoadAbortController] =
+    useState<AbortController | null>(null);
 
   const refresh = useCallback(async () => {
     setModelsError(null);
@@ -215,8 +217,11 @@ export function useChatModelRuntime() {
 
       setModelsError(null);
       setLoadingModel({ id: modelId, displayName });
+      const abortCtrl = new AbortController();
+      setLoadAbortController(abortCtrl);
       try {
         async function performLoad(): Promise<void> {
+          if (abortCtrl.signal.aborted) throw new Error("Cancelled");
           let previousWasUnloaded = false;
           const currentCheckpoint =
             useChatRuntimeStore.getState().params.checkpoint;
@@ -277,6 +282,7 @@ export function useChatModelRuntime() {
 
         const loadPromise = performLoad().finally(() => {
           setLoadingModel(null);
+          setLoadAbortController(null);
         });
 
         await toast.promise(loadPromise, {
@@ -322,10 +328,25 @@ export function useChatModelRuntime() {
     }
   }, [clearCheckpoint, params.checkpoint, refresh, setModelsError]);
 
+  const cancelLoading = useCallback(async () => {
+    if (!loadingModel) return;
+    loadAbortController?.abort();
+    setLoadingModel(null);
+    setLoadAbortController(null);
+    try {
+      await unloadModel({ model_path: loadingModel.id });
+    } catch {
+      // Best-effort cleanup
+    }
+    clearCheckpoint();
+    toast.info("Model loading cancelled");
+  }, [loadingModel, loadAbortController, clearCheckpoint]);
+
   return {
     refresh,
     selectModel,
     ejectModel,
+    cancelLoading,
     loadingModel,
   };
 }