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 { Downloading model… + ) : 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 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(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 { + 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, }; }