fix: #396 - allow user to cancel a model download (#530)

* fix: #396 - allow user to cancel model download

* chore: fix typo

electron/main.ts

@@ -20,6 +20,7 @@ const { autoUpdater } = require("electron-updater");
 const Store = require("electron-store");
 
 let requiredModules: Record<string, any> = {};
+const networkRequests: Record<string, any> = {};
 let mainWindow: BrowserWindow | undefined = undefined;
 
 app
@@ -48,18 +49,6 @@ app.on("quit", () => {
   app.quit();
 });
 
-ipcMain.handle("setNativeThemeLight", () => {
-  nativeTheme.themeSource = "light";
-});
-
-ipcMain.handle("setNativeThemeDark", () => {
-  nativeTheme.themeSource = "dark";
-});
-
-ipcMain.handle("setNativeThemeSystem", () => {
-  nativeTheme.themeSource = "system";
-});
-
 function createMainWindow() {
   mainWindow = new BrowserWindow({
     width: 1200,
@@ -138,6 +127,30 @@ function handleAppUpdates() {
  * Handles various IPC messages from the renderer process.
  */
 function handleIPCs() {
+  /**
+   * Handles the "setNativeThemeLight" IPC message by setting the native theme source to "light".
+   * This will change the appearance of the app to the light theme.
+   */
+  ipcMain.handle("setNativeThemeLight", () => {
+    nativeTheme.themeSource = "light";
+  });
+
+  /**
+   * Handles the "setNativeThemeDark" IPC message by setting the native theme source to "dark".
+   * This will change the appearance of the app to the dark theme.
+   */
+  ipcMain.handle("setNativeThemeDark", () => {
+    nativeTheme.themeSource = "dark";
+  });
+
+  /**
+   * Handles the "setNativeThemeSystem" IPC message by setting the native theme source to "system".
+   * This will change the appearance of the app to match the system's current theme.
+   */
+  ipcMain.handle("setNativeThemeSystem", () => {
+    nativeTheme.themeSource = "system";
+  });
+
   /**
    * Invokes a function from a plugin module in main node process.
    * @param _event - The IPC event object.
@@ -319,8 +332,9 @@ function handleIPCs() {
   ipcMain.handle("downloadFile", async (_event, url, fileName) => {
     const userDataPath = app.getPath("userData");
     const destination = resolve(userDataPath, fileName);
+    const rq = request(url);
 
-    progress(request(url), {})
+    progress(rq, {})
       .on("progress", function (state: any) {
         mainWindow?.webContents.send("FILE_DOWNLOAD_UPDATE", {
           ...state,
@@ -332,13 +346,54 @@ function handleIPCs() {
           fileName,
           err,
         });
+        networkRequests[fileName] = undefined;
       })
       .on("end", function () {
-        mainWindow?.webContents.send("FILE_DOWNLOAD_COMPLETE", {
+        if (networkRequests[fileName]) {
-          fileName,
+          mainWindow?.webContents.send("FILE_DOWNLOAD_COMPLETE", {
-        });
+            fileName,
+          });
+          networkRequests[fileName] = undefined;
+        } else {
+          mainWindow?.webContents.send("FILE_DOWNLOAD_ERROR", {
+            fileName,
+            err: "Download cancelled",
+          });
+        }
       })
       .pipe(createWriteStream(destination));
+
+    networkRequests[fileName] = rq;
+  });
+
+  /**
+   * Handles the "pauseDownload" IPC message by pausing the download associated with the provided fileName.
+   * @param _event - The IPC event object.
+   * @param fileName - The name of the file being downloaded.
+   */
+  ipcMain.handle("pauseDownload", async (_event, fileName) => {
+    networkRequests[fileName]?.pause();
+  });
+
+  /**
+   * Handles the "resumeDownload" IPC message by resuming the download associated with the provided fileName.
+   * @param _event - The IPC event object.
+   * @param fileName - The name of the file being downloaded.
+   */
+  ipcMain.handle("resumeDownload", async (_event, fileName) => {
+    networkRequests[fileName]?.resume();
+  });
+
+  /**
+   * Handles the "abortDownload" IPC message by aborting the download associated with the provided fileName.
+   * The network request associated with the fileName is then removed from the networkRequests object.
+   * @param _event - The IPC event object.
+   * @param fileName - The name of the file being downloaded.
+   */
+  ipcMain.handle("abortDownload", async (_event, fileName) => {
+    const rq = networkRequests[fileName];
+    networkRequests[fileName] = undefined;
+    rq?.abort();
   });
 
   /**

electron/preload.ts

@@ -38,6 +38,15 @@ contextBridge.exposeInMainWorld("electronAPI", {
   downloadFile: (url: string, path: string) =>
     ipcRenderer.invoke("downloadFile", url, path),
 
+  pauseDownload: (fileName: string) =>
+    ipcRenderer.invoke("pauseDownload", fileName),
+
+  resumeDownload: (fileName: string) =>
+    ipcRenderer.invoke("resumeDownload", fileName),
+
+  abortDownload: (fileName: string) =>
+    ipcRenderer.invoke("abortDownload", fileName),
+
   onFileDownloadUpdate: (callback: any) =>
     ipcRenderer.on("FILE_DOWNLOAD_UPDATE", callback),
 

web/app/_components/ConfirmationModal/index.tsx

@@ -0,0 +1,87 @@
+import React, { Fragment } from 'react'
+import { Dialog, Transition } from '@headlessui/react'
+import { QuestionMarkCircleIcon } from '@heroicons/react/24/outline'
+import { PrimitiveAtom, useAtom } from 'jotai'
+
+interface Props {
+  atom: PrimitiveAtom<boolean>
+  title: string
+  description: string
+  onConfirm: () => void
+}
+
+const ConfirmationModal: React.FC<Props> = ({ atom, title, description, onConfirm }) => {
+  const [show, setShow] = useAtom(atom)
+
+  return (
+    <Transition.Root show={show} as={Fragment}>
+      <Dialog as="div" className="relative z-10" onClose={setShow}>
+        <Transition.Child
+          as={Fragment}
+          enter="ease-out duration-300"
+          enterFrom="opacity-0"
+          enterTo="opacity-100"
+          leave="ease-in duration-200"
+          leaveFrom="opacity-100"
+          leaveTo="opacity-0"
+        >
+          <div className="fixed inset-0 bg-gray-500 bg-opacity-75 transition-opacity" />
+        </Transition.Child>
+
+        <div className="fixed inset-0 z-10 overflow-y-auto">
+          <div className="flex min-h-full items-end justify-center p-4 text-center sm:items-center sm:p-0">
+            <Transition.Child
+              as={Fragment}
+              enter="ease-out duration-300"
+              enterFrom="opacity-0 translate-y-4 sm:translate-y-0 sm:scale-95"
+              enterTo="opacity-100 translate-y-0 sm:scale-100"
+              leave="ease-in duration-200"
+              leaveFrom="opacity-100 translate-y-0 sm:scale-100"
+              leaveTo="opacity-0 translate-y-4 sm:translate-y-0 sm:scale-95"
+            >
+              <Dialog.Panel className="relative transform overflow-hidden rounded-lg bg-white px-4 pb-4 pt-5 text-left shadow-xl transition-all sm:my-8 sm:w-full sm:max-w-lg sm:p-6">
+                <div className="sm:flex sm:items-start">
+                  <div className="mx-auto flex h-12 w-12 flex-shrink-0 items-center justify-center rounded-full bg-red-100 sm:mx-0 sm:h-10 sm:w-10">
+                    <QuestionMarkCircleIcon
+                      className="h-6 w-6 text-green-600"
+                      aria-hidden="true"
+                    />
+                  </div>
+                  <div className="mt-3 text-center sm:ml-4 sm:mt-0 sm:text-left">
+                    <Dialog.Title
+                      as="h3"
+                      className="text-base font-semibold leading-6 text-gray-900"
+                    >
+                      {title}
+                    </Dialog.Title>
+                    <div className="mt-2">
+                      <p className="text-sm text-gray-500">{description}</p>
+                    </div>
+                  </div>
+                </div>
+                <div className="mt-5 sm:mt-4 sm:flex sm:flex-row-reverse">
+                  <button
+                    type="button"
+                    className="inline-flex w-full justify-center rounded-md bg-red-600 px-3 py-2 text-sm font-semibold text-white shadow-sm hover:bg-red-500 sm:ml-3 sm:w-auto"
+                    onClick={onConfirm}
+                  >
+                    OK
+                  </button>
+                  <button
+                    type="button"
+                    className="mt-3 inline-flex w-full justify-center rounded-md bg-white px-3 py-2 text-sm font-semibold text-gray-900 shadow-sm ring-1 ring-inset ring-gray-300 hover:bg-gray-50 sm:mt-0 sm:w-auto"
+                    onClick={() => setShow(false)}
+                  >
+                    Cancel
+                  </button>
+                </div>
+              </Dialog.Panel>
+            </Transition.Child>
+          </div>
+        </div>
+      </Dialog>
+    </Transition.Root>
+  )
+}
+
+export default React.memo(ConfirmationModal)

web/app/_components/ExploreModelItemHeader/index.tsx

@@ -11,6 +11,8 @@ import {
   MainViewState,
   setMainViewStateAtom,
 } from '@helpers/atoms/MainView.atom'
+import ConfirmationModal from '../ConfirmationModal'
+import { showingCancelDownloadModalAtom } from '@helpers/atoms/Modal.atom'
 
 type Props = {
   suitableModel: ModelVersion
@@ -31,6 +33,9 @@ const ExploreModelItemHeader: React.FC<Props> = ({
   )
   const downloadState = useAtomValue(downloadAtom)
   const setMainViewState = useSetAtom(setMainViewStateAtom)
+  const setShowingCancelDownloadModal = useSetAtom(
+    showingCancelDownloadModalAtom
+  )
 
   useEffect(() => {
     getPerformanceForModel(suitableModel)
@@ -70,17 +75,30 @@ const ExploreModelItemHeader: React.FC<Props> = ({
     // downloading
     downloadButton = (
       <Button
-        disabled
+        themes="outline"
-        themes="accent"
         onClick={() => {
-          setMainViewState(MainViewState.MyModel)
+          setShowingCancelDownloadModal(true)
         }}
       >
-        Downloading {formatDownloadPercentage(downloadState.percent)}
+        Cancel ({formatDownloadPercentage(downloadState.percent)})
       </Button>
     )
   }
 
+  let cancelDownloadModal =
+    downloadState != null ? (
+      <ConfirmationModal
+        atom={showingCancelDownloadModalAtom}
+        title="Cancel Download"
+        description={`Are you sure you want to cancel the download of ${downloadState?.fileName}?`}
+        onConfirm={() => {
+          window.coreAPI?.abortDownload(downloadState?.fileName)
+        }}
+      />
+    ) : (
+      <></>
+    )
+
   return (
     <div className="flex items-center justify-between rounded-t-md border-b border-border bg-background/50 px-4 py-2">
       <div className="flex items-center gap-2">
@@ -90,6 +108,7 @@ const ExploreModelItemHeader: React.FC<Props> = ({
         )}
       </div>
       {downloadButton}
+      {cancelDownloadModal}
     </div>
   )
 }

web/containers/BottomBar/index.tsx

@@ -57,14 +57,6 @@ const BottomBar = () => {
         {!stateModelStartStop.loading && (
           <SystemItem name="Active model:" value={activeModel?.name || '-'} />
         )}
-        {downloadStates.length > 0 && (
-          <SystemItem
-            name="Downloading:"
-            value={`${downloadStates[0].fileName} - ${formatDownloadPercentage(
-              downloadStates[0].percent
-            )}`}
-          />
-        )}
       </div>
       <div className="flex gap-x-2">
         <SystemItem name="CPU:" value={`${cpu}%`} />

web/helpers/EventListenerWrapper.tsx

@@ -35,6 +35,7 @@ export default function EventListenerWrapper({ children }: Props) {
       window.electronAPI.onFileDownloadError(
         (_event: string, callback: any) => {
           console.log('Download error', callback)
+          setDownloadStateSuccess(callback.fileName)
         }
       )
 

web/helpers/atoms/DownloadState.atom.ts

@@ -28,4 +28,4 @@ export const setDownloadStateSuccessAtom = atom(
     delete currentState[fileName]
     set(modelDownloadStateAtom, currentState)
   }
-)
+)

web/helpers/atoms/Modal.atom.ts

@@ -8,6 +8,7 @@ export const showingAdvancedPromptAtom = atom<boolean>(false)
 export const showingProductDetailAtom = atom<boolean>(false)
 export const showingMobilePaneAtom = atom<boolean>(false)
 export const showingBotListModalAtom = atom<boolean>(false)
+export const showingCancelDownloadModalAtom = atom<boolean>(false)
 
 export const switchingModelConfirmationModalPropsAtom = atom<
   SwitchingModelConfirmationModalProps | undefined

web/screens/MyModels/index.tsx

@@ -30,7 +30,7 @@ const MyModelsScreen = () => {
     return (
       <div className="flex h-full items-center justify-center px-4">
         <div className="text-center">
-          <LayoutGrid size={32} className="text-accent/50 mx-auto" />
+          <LayoutGrid size={32} className="mx-auto text-accent/50" />
           <div className="mt-4">
             {isDownloadingFirstModel ? (
               <div className="relative">
@@ -38,7 +38,7 @@ const MyModelsScreen = () => {
                   <h1 className="text-2xl font-bold leading-snug">
                     Donwloading your first model
                   </h1>
-                  <p className="text-muted-foreground mt-1">
+                  <p className="mt-1 text-muted-foreground">
                     {downloadStates[0].fileName} -{' '}
                     {formatDownloadPercentage(downloadStates[0].percent)}
                   </p>
@@ -47,7 +47,7 @@ const MyModelsScreen = () => {
             ) : (
               <Fragment>
                 <h1 className="text-2xl font-bold leading-snug">{`Ups, You don't have a model.`}</h1>
-                <p className="text-muted-foreground mt-1 text-base">{`let’s download your first model`}</p>
+                <p className="mt-1 text-base text-muted-foreground">{`let’s download your first model`}</p>
                 <Button
                   className="mt-4"
                   themes="accent"
@@ -65,7 +65,12 @@ const MyModelsScreen = () => {
   return (
     <div className="flex h-full w-full overflow-y-auto">
       <div className="w-full p-5">
-        <h1 data-testid="testid-mymodels-header" className="text-lg font-semibold">My Models</h1>
+        <h1
+          data-testid="testid-mymodels-header"
+          className="text-lg font-semibold"
+        >
+          My Models
+        </h1>
         <p className="mt-2 text-gray-600 dark:text-gray-400">
           You have <span>{downloadedModels.length}</span> models downloaded
         </p>