fix: download mutilple binaries (janhq#2043)

Signed-off-by: James <[email protected]>
Co-authored-by: James <[email protected]>

extensions/model-extension/src/index.ts

@@ -12,8 +12,9 @@ import {
   DownloadEvent,
   DownloadRoute,
   ModelEvent,
+  DownloadState,
 } from '@janhq/core'
-import { DownloadState } from '@janhq/core/.'
+
 import { extractFileName } from './helpers/path'
 
 /**

web/containers/Layout/BottomBar/DownloadingState/index.tsx

@@ -32,7 +32,8 @@ export default function DownloadingState() {
     .map((a) => a.size.total + a.size.total)
     .reduce((partialSum, a) => partialSum + a, 0)
 
-  const totalPercentage = ((totalCurrentProgress / totalSize) * 100).toFixed(2)
+  const totalPercentage =
+    totalSize !== 0 ? ((totalCurrentProgress / totalSize) * 100).toFixed(2) : 0
 
   return (
     <Fragment>

web/containers/ModalCancelDownload/index.tsx

@@ -1,4 +1,4 @@
-import { useMemo } from 'react'
+import { useCallback } from 'react'
 
 import { Model } from '@janhq/core'
 
@@ -14,7 +14,7 @@ import {
   Progress,
 } from '@janhq/uikit'
 
-import { atom, useAtomValue } from 'jotai'
+import { useAtomValue } from 'jotai'
 
 import useDownloadModel from '@/hooks/useDownloadModel'
 
@@ -30,14 +30,21 @@ type Props = {
 }
 
 const ModalCancelDownload: React.FC<Props> = ({ model, isFromList }) => {
+  const { abortModelDownload } = useDownloadModel()
   const downloadingModels = useAtomValue(getDownloadingModelAtom)
-  const downloadAtom = useMemo(
-    () => atom((get) => get(modelDownloadStateAtom)[model.id]),
-    [model.id]
-  )
-  const downloadState = useAtomValue(downloadAtom)
+  const allDownloadStates = useAtomValue(modelDownloadStateAtom)
+  const downloadState = allDownloadStates[model.id]
+
   const cancelText = `Cancel ${formatDownloadPercentage(downloadState.percent)}`
-  const { abortModelDownload } = useDownloadModel()
+
+  const onAbortDownloadClick = useCallback(() => {
+    if (downloadState?.modelId) {
+      const model = downloadingModels.find(
+        (model) => model.id === downloadState.modelId
+      )
+      if (model) abortModelDownload(model)
+    }
+  }, [downloadState, downloadingModels, abortModelDownload])
 
   return (
     <Modal>
@@ -77,17 +84,7 @@ const ModalCancelDownload: React.FC<Props> = ({ model, isFromList }) => {
               <Button themes="ghost">No</Button>
             </ModalClose>
             <ModalClose asChild>
-              <Button
-                themes="danger"
-                onClick={() => {
-                  if (downloadState?.modelId) {
-                    const model = downloadingModels.find(
-                      (model) => model.id === downloadState.modelId
-                    )
-                    if (model) abortModelDownload(model)
-                  }
-                }}
-              >
+              <Button themes="danger" onClick={onAbortDownloadClick}>
                 Yes
               </Button>
             </ModalClose>

web/containers/Providers/EventListener.tsx

@@ -3,7 +3,7 @@ import { PropsWithChildren, useCallback, useEffect } from 'react'
 
 import React from 'react'
 
-import { DownloadEvent, events } from '@janhq/core'
+import { DownloadEvent, events, DownloadState } from '@janhq/core'
 import { useSetAtom } from 'jotai'
 
 import { setDownloadStateAtom } from '@/hooks/useDownloadState'

web/hooks/useDownloadModel.ts

@@ -7,14 +7,13 @@ import {
   abortDownload,
   joinPath,
   ModelArtifact,
+  DownloadState,
 } from '@janhq/core'
 
 import { useSetAtom } from 'jotai'
 
 import { FeatureToggleContext } from '@/context/FeatureToggle'
 
-import { modelBinFileName } from '@/utils/model'
-
 import { setDownloadStateAtom } from './useDownloadState'
 
 import { extensionManager } from '@/extension/ExtensionManager'
@@ -29,7 +28,7 @@ export default function useDownloadModel() {
     async (model: Model) => {
       const childProgresses: DownloadState[] = model.sources.map(
         (source: ModelArtifact) => ({
-          filename: source.filename,
+          fileName: source.filename,
           modelId: model.id,
           time: {
             elapsed: 0,
@@ -47,7 +46,7 @@ export default function useDownloadModel() {
 
       // set an initial download state
       setDownloadState({
-        filename: '',
+        fileName: '',
         modelId: model.id,
         time: {
           elapsed: 0,
@@ -70,11 +69,12 @@ export default function useDownloadModel() {
     [ignoreSSL, proxy, addDownloadingModel, setDownloadState]
   )
 
-  const abortModelDownload = async (model: Model) => {
-    await abortDownload(
-      await joinPath(['models', model.id, modelBinFileName(model)])
-    )
-  }
+  const abortModelDownload = useCallback(async (model: Model) => {
+    for (const source of model.sources) {
+      const path = await joinPath(['models', model.id, source.filename])
+      await abortDownload(path)
+    }
+  }, [])
 
   return {
     downloadModel,

web/hooks/useDownloadState.ts

@@ -1,3 +1,4 @@
+import { DownloadState } from '@janhq/core'
 import { atom } from 'jotai'
 
 import { toaster } from '@/containers/Toast'
@@ -20,18 +21,35 @@ export const setDownloadStateAtom = atom(
     const currentState = { ...get(modelDownloadStateAtom) }
 
     if (state.downloadState === 'end') {
-      // download successfully
-      delete currentState[state.modelId]
-      set(removeDownloadingModelAtom, state.modelId)
-      const model = get(configuredModelsAtom).find(
-        (e) => e.id === state.modelId
+      const modelDownloadState = currentState[state.modelId]
+
+      const updatedChildren: DownloadState[] =
+        modelDownloadState.children!.filter(
+          (m) => m.fileName !== state.fileName
+        )
+      updatedChildren.push(state)
+      modelDownloadState.children = updatedChildren
+      currentState[state.modelId] = modelDownloadState
+
+      const isAllChildrenDownloadEnd = modelDownloadState.children?.every(
+        (m) => m.downloadState === 'end'
       )
-      if (model) set(downloadedModelsAtom, (prev) => [...prev, model])
-      toaster({
-        title: 'Download Completed',
-        description: `Download ${state.modelId} completed`,
-        type: 'success',
-      })
+
+      if (isAllChildrenDownloadEnd) {
+        // download successfully
+        delete currentState[state.modelId]
+        set(removeDownloadingModelAtom, state.modelId)
+
+        const model = get(configuredModelsAtom).find(
+          (e) => e.id === state.modelId
+        )
+        if (model) set(downloadedModelsAtom, (prev) => [...prev, model])
+        toaster({
+          title: 'Download Completed',
+          description: `Download ${state.modelId} completed`,
+          type: 'success',
+        })
+      }
     } else if (state.downloadState === 'error') {
       // download error
       delete currentState[state.modelId]
@@ -59,7 +77,62 @@ export const setDownloadStateAtom = atom(
       }
     } else {
       // download in progress
-      currentState[state.modelId] = state
+      if (state.size.total === 0) {
+        // this is initial state, just set the state
+        currentState[state.modelId] = state
+        set(modelDownloadStateAtom, currentState)
+        return
+      }
+
+      const modelDownloadState = currentState[state.modelId]
+      if (!modelDownloadState) {
+        console.debug('setDownloadStateAtom: modelDownloadState not found')
+        return
+      }
+
+      // delete the children if the filename is matched and replace the new state
+      const updatedChildren: DownloadState[] =
+        modelDownloadState.children!.filter(
+          (m) => m.fileName !== state.fileName
+        )
+
+      updatedChildren.push(state)
+
+      // re-calculate the overall progress if we have all the children download data
+      const isAnyChildDownloadNotReady = updatedChildren.some(
+        (m) => m.size.total === 0
+      )
+
+      modelDownloadState.children = updatedChildren
+
+      if (isAnyChildDownloadNotReady) {
+        // just update the children
+        currentState[state.modelId] = modelDownloadState
+        set(modelDownloadStateAtom, currentState)
+
+        return
+      }
+
+      const parentTotalSize = modelDownloadState.size.total
+      if (parentTotalSize === 0) {
+        // calculate the total size of the parent by sum all children total size
+        const totalSize = updatedChildren.reduce(
+          (acc, m) => acc + m.size.total,
+          0
+        )
+
+        modelDownloadState.size.total = totalSize
+      }
+
+      // calculate the total transferred size by sum all children transferred size
+      const transferredSize = updatedChildren.reduce(
+        (acc, m) => acc + m.size.transferred,
+        0
+      )
+      modelDownloadState.size.transferred = transferredSize
+      modelDownloadState.percent = transferredSize / parentTotalSize
+
+      currentState[state.modelId] = modelDownloadState
     }
 
     set(modelDownloadStateAtom, currentState)

web/screens/ExploreModels/ExploreModelItem/index.tsx

@@ -1,6 +1,4 @@
-/* eslint-disable react/display-name */
-
-import { forwardRef, useState } from 'react'
+import { useState } from 'react'
 
 import { Model } from '@janhq/core'
 import { Badge } from '@janhq/uikit'
@@ -11,7 +9,7 @@ type Props = {
   model: Model
 }
 
-const ExploreModelItem = forwardRef<HTMLDivElement, Props>(({ model }, ref) => {
+const ExploreModelItem: React.FC<Props> = ({ model }) => {
   const [open, setOpen] = useState('')
 
   const handleToggle = () => {
@@ -23,10 +21,7 @@ const ExploreModelItem = forwardRef<HTMLDivElement, Props>(({ model }, ref) => {
   }
 
   return (
-    <div
-      ref={ref}
-      className="mb-6 flex flex-col overflow-hidden rounded-xl border border-border bg-background/60"
-    >
+    <div className="mb-6 flex flex-col overflow-hidden rounded-xl border border-border bg-background/60">
       <ExploreModelItemHeader
         model={model}
         onClick={handleToggle}
@@ -82,17 +77,11 @@ const ExploreModelItem = forwardRef<HTMLDivElement, Props>(({ model }, ref) => {
               </span>
               <p className="mt-2 font-medium uppercase">{model.format}</p>
             </div>
-            {/* <div className="mt-4">
-              <span className="font-semibold text-muted-foreground">
-                Compatibility
-              </span>
-              <p className="mt-2 font-medium">-</p>
-            </div> */}
           </div>
         </div>
       )}
     </div>
   )
-})
+}
 
 export default ExploreModelItem