Merge pull request #3956 from janhq/fix/3911-inconsistent-model-hub-and-download-bar

fix: Inconsistent model hub and download bar

core/src/node/api/processors/download.ts

@@ -50,11 +50,6 @@ export class Downloader implements Processor {
     const initialDownloadState: DownloadState = {
       modelId,
       fileName,
-      time: {
-        elapsed: 0,
-        remaining: 0,
-      },
-      speed: 0,
       percent: 0,
       size: {
         total: 0,

core/src/types/file/index.ts

@@ -6,8 +6,8 @@ export type FileStat = {
 export type DownloadState = {
   modelId: string // TODO: change to download id
   fileName: string
-  time: DownloadTime
+  time?: DownloadTime
-  speed: number
+  speed?: number
 
   percent: number
   size: DownloadSize

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

@@ -2,15 +2,19 @@ import { Fragment } from 'react'
 
 import { Progress, Modal, Button } from '@janhq/joi'
 
-import { useAtomValue } from 'jotai'
+import { useAtomValue, useSetAtom } from 'jotai'
 
 import useDownloadModel from '@/hooks/useDownloadModel'
-import { modelDownloadStateAtom } from '@/hooks/useDownloadState'
+import {
+  modelDownloadStateAtom,
+  removeDownloadStateAtom,
+} from '@/hooks/useDownloadState'
 
 import { formatDownloadPercentage } from '@/utils/converter'
 
 export default function DownloadingState() {
   const downloadStates = useAtomValue(modelDownloadStateAtom)
+  const removeDownloadState = useSetAtom(removeDownloadStateAtom)
   const { abortModelDownload } = useDownloadModel()
 
   const totalCurrentProgress = Object.values(downloadStates)
@@ -73,6 +77,7 @@ export default function DownloadingState() {
                       theme="destructive"
                       onClick={() => {
                         if (item?.modelId) {
+                          removeDownloadState(item?.modelId)
                           abortModelDownload(item?.modelId)
                         }
                       }}

web/containers/ModalCancelDownload/index.tsx

@@ -8,12 +8,13 @@ import { useAtomValue, useSetAtom } from 'jotai'
 
 import useDownloadModel from '@/hooks/useDownloadModel'
 
-import { modelDownloadStateAtom } from '@/hooks/useDownloadState'
+import {
+  modelDownloadStateAtom,
+  removeDownloadStateAtom,
+} from '@/hooks/useDownloadState'
 
 import { formatDownloadPercentage } from '@/utils/converter'
 
-import { removeDownloadingModelAtom } from '@/helpers/atoms/Model.atom'
-
 type Props = {
   model: Model
   isFromList?: boolean
@@ -21,16 +22,16 @@ type Props = {
 
 const ModalCancelDownload = ({ model, isFromList }: Props) => {
   const { abortModelDownload } = useDownloadModel()
-  const removeModelDownload = useSetAtom(removeDownloadingModelAtom)
+  const removeDownloadState = useSetAtom(removeDownloadStateAtom)
   const allDownloadStates = useAtomValue(modelDownloadStateAtom)
   const downloadState = allDownloadStates[model.id]
 
   const cancelText = `Cancel ${formatDownloadPercentage(downloadState?.percent ?? 0)}`
 
   const onAbortDownloadClick = useCallback(() => {
-    removeModelDownload(model.id)
+    removeDownloadState(model.id)
     abortModelDownload(downloadState?.modelId ?? model.id)
-  }, [downloadState, abortModelDownload, removeModelDownload, model])
+  }, [downloadState, abortModelDownload, removeDownloadState, model])
 
   return (
     <Modal

web/hooks/useDownloadModel.ts

@@ -6,6 +6,8 @@ import { useSetAtom } from 'jotai'
 
 import { toaster } from '@/containers/Toast'
 
+import { setDownloadStateAtom } from './useDownloadState'
+
 import { extensionManager } from '@/extension/ExtensionManager'
 
 import {
@@ -16,10 +18,21 @@ import {
 export default function useDownloadModel() {
   const removeDownloadingModel = useSetAtom(removeDownloadingModelAtom)
   const addDownloadingModel = useSetAtom(addDownloadingModelAtom)
+  const setDownloadStates = useSetAtom(setDownloadStateAtom)
 
   const downloadModel = useCallback(
     async (model: string, id?: string, name?: string) => {
       addDownloadingModel(id ?? model)
+      setDownloadStates({
+        modelId: id ?? model,
+        downloadState: 'downloading',
+        fileName: id ?? model,
+        size: {
+          total: 0,
+          transferred: 0,
+        },
+        percent: 0,
+      })
       downloadLocalModel(model, id, name).catch((error) => {
         if (error.message) {
           toaster({
@@ -32,7 +45,7 @@ export default function useDownloadModel() {
         removeDownloadingModel(model)
       })
     },
-    [removeDownloadingModel, addDownloadingModel]
+    [removeDownloadingModel, addDownloadingModel, setDownloadStates]
   )
 
   const abortModelDownload = useCallback(async (model: string) => {

web/hooks/useDownloadState.ts

@@ -10,8 +10,18 @@ import {
 } from '@/helpers/atoms/Model.atom'
 
 // download states
+
 export const modelDownloadStateAtom = atom<Record<string, DownloadState>>({})
 
+/**
+ * Remove a download state for a particular model.
+ */
+export const removeDownloadStateAtom = atom(null, (get, set, id: string) => {
+  const currentState = { ...get(modelDownloadStateAtom) }
+  delete currentState[id]
+  set(modelDownloadStateAtom, currentState)
+  set(removeDownloadingModelAtom, id)
+})
 /**
  * Used to set the download state for a particular model.
  */

web/screens/Settings/MyModels/MyModelList/index.tsx

@@ -1,6 +1,6 @@
 import { memo, useState } from 'react'
 
-import { InferenceEngine, Model } from '@janhq/core'
+import { Model } from '@janhq/core'
 import { Badge, Button, Tooltip, useClickOutside } from '@janhq/joi'
 import { useAtom } from 'jotai'
 import {