fix: failed download should not be added to download state

core/src/types/model/index.ts

@@ -1,4 +1,3 @@
 export * from './modelEntity'
-export * from './modelInterface'
 export * from './modelImport'
 export * from './chatCompletion'

web/hooks/useCortex.ts

@@ -21,6 +21,7 @@ import { UpdateConfigMutationVariables } from './useEngineMutation'
 import { MessageCreateMutationVariables } from './useMessageCreateMutation'
 import { MessageDeleteMutationVariables } from './useMessageDeleteMutation'
 import { MessageUpdateMutationVariables } from './useMessageUpdateMutation'
+import { DownloadModelMutationVariables } from './useModelDownloadMutation'
 
 import { hostAtom } from '@/helpers/atoms/AppConfig.atom'
 
@@ -246,7 +247,8 @@ const useCortex = () => {
   )
 
   const downloadModel = useCallback(
-    async (modelId: string, fileName?: string, persistedModelId?: string) => {
+    async (variables: DownloadModelMutationVariables) => {
+      const { modelId, fileName, persistedModelId } = variables
       const response = await fetch(`${host}/models/${modelId}/pull`, {
         method: 'POST',
         headers: {

web/hooks/useHfModelFetchAndDownload.ts

@@ -3,15 +3,12 @@ import { useCallback } from 'react'
 import { HuggingFaceRepoData } from '@janhq/core'
 import { useQueryClient } from '@tanstack/react-query'
 
-import { useSetAtom } from 'jotai'
-
 import { toaster } from '@/containers/Toast'
 
 import { fetchHuggingFaceRepoData } from '@/utils/huggingface'
 
-import useCortex from './useCortex'
-import { addDownloadModelStateAtom } from './useDownloadState'
 import { fetchHuggingFaceRepoDataQueryKey } from './useHfRepoDataQuery'
+import useModelDownloadMutation from './useModelDownloadMutation'
 
 /**
  * Fetches the data for a Hugging Face model and downloads it.
@@ -21,8 +18,7 @@ import { fetchHuggingFaceRepoDataQueryKey } from './useHfRepoDataQuery'
  * @param modelHandle The model handle to fetch and download. E.g: "NousResearch/Hermes-2-Theta-Llama-3-8B-GGUF"
  */
 const useHfModelFetchAndDownload = () => {
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
   const queryClient = useQueryClient()
 
   const fetchData = useCallback(
@@ -90,15 +86,14 @@ const useHfModelFetchAndDownload = () => {
         .concat('_')
         .concat(recommendedModel.rfilename)
 
-      addDownloadState(persistModelId)
-      await downloadModel(
-        modelHandle,
-        recommendedModel.rfilename,
-        persistModelId
-      )
+      downloadModelMutation.mutate({
+        modelId: modelHandle,
+        fileName: recommendedModel.rfilename,
+        persistedModelId: persistModelId,
+      })
     },
 
-    [addDownloadState, downloadModel, fetchData]
+    [fetchData, downloadModelMutation]
   )
 
   return { fetchDataAndDownload }

web/hooks/useModelDownloadMutation.ts

@@ -0,0 +1,49 @@
+import { useMutation } from '@tanstack/react-query'
+
+import { useSetAtom } from 'jotai'
+
+import { toaster } from '@/containers/Toast'
+
+import useCortex from './useCortex'
+import { addDownloadModelStateAtom } from './useDownloadState'
+
+export type DownloadModelMutationVariables = {
+  modelId: string
+  fileName?: string
+  persistedModelId?: string
+}
+
+const useModelDownloadMutation = () => {
+  const { downloadModel } = useCortex()
+  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
+
+  return useMutation({
+    mutationFn: downloadModel,
+
+    onMutate: (variables) => {
+      console.debug('Downloading model', variables)
+    },
+
+    onSuccess: (data, variables) => {
+      console.debug('Download response success', data, variables)
+
+      const { persistedModelId, modelId } = variables
+      if (persistedModelId) {
+        addDownloadState(persistedModelId)
+      } else {
+        addDownloadState(modelId)
+      }
+    },
+
+    onError: (err, variables) => {
+      console.error('Failed to download model', err, variables)
+      toaster({
+        title: 'Failed to download model',
+        description: err.message,
+        type: 'error',
+      })
+    },
+  })
+}
+
+export default useModelDownloadMutation

web/screens/HubScreen2/components/BuiltInModelCard.tsx

@@ -8,11 +8,8 @@ import { toaster } from '@/containers/Toast'
 
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
-import useCortex from '@/hooks/useCortex'
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage } from '@/utils/converter'
@@ -73,9 +70,8 @@ type DownloadContainerProps = {
 const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
   const { abortDownload } = useAbortDownload()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const setDownloadLocalModelModalStage = useSetAtom(localModelModalStageAtom)
@@ -93,10 +89,9 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, modelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
-    addDownloadState(modelId)
-    await downloadModel(modelId)
-  }, [downloadModel, addDownloadState, modelId])
+  const onDownloadClick = useCallback(() => {
+    downloadModelMutation.mutate({ modelId })
+  }, [downloadModelMutation, modelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/HubScreen2/components/HfListModel.tsx

@@ -9,14 +9,10 @@ import { toaster } from '@/containers/Toast'
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
 
-import useCortex from '@/hooks/useCortex'
-
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
 
 import useHfRepoDataQuery from '@/hooks/useHfRepoDataQuery'
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage, toGibibytes } from '@/utils/converter'
@@ -97,9 +93,8 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   fileName,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
   const { abortDownload } = useAbortDownload()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const { data: assistants } = useAssistantQuery()
@@ -122,9 +117,12 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   )
 
   const onDownloadClick = useCallback(async () => {
-    addDownloadState(persistModelId)
-    await downloadModel(modelHandle, fileName, persistModelId)
-  }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId])
+    downloadModelMutation.mutate({
+      modelId: modelHandle,
+      fileName: fileName,
+      persistedModelId: persistModelId,
+    })
+  }, [downloadModelMutation, modelHandle, fileName, persistModelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/HubScreen2/components/ListModel.tsx

@@ -9,16 +9,12 @@ import { toaster } from '@/containers/Toast'
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
 
-import useCortex from '@/hooks/useCortex'
-
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
 
 import useEngineQuery from '@/hooks/useEngineQuery'
 import useHfEngineToBranchesQuery from '@/hooks/useHfEngineToBranchesQuery'
 
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage, toGibibytes } from '@/utils/converter'
@@ -150,9 +146,8 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   branch,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
   const { abortDownload } = useAbortDownload()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const setDownloadLocalModelModalStage = useSetAtom(localModelModalStageAtom)
@@ -172,10 +167,9 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, modelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
-    addDownloadState(modelId)
-    await downloadModel(modelId)
-  }, [downloadModel, addDownloadState, modelId])
+  const onDownloadClick = useCallback(() => {
+    downloadModelMutation.mutate({ modelId })
+  }, [downloadModelMutation, modelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/HubScreen2/components/SliderItem.tsx

@@ -10,11 +10,8 @@ import { toaster } from '@/containers/Toast'
 
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
-import useCortex from '@/hooks/useCortex'
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
 import { QuickStartModel } from '@/hooks/useModelHub'
 import useThreads from '@/hooks/useThreads'
 
@@ -77,8 +74,7 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   fileName,
 }) => {
-  const { downloadModel } = useCortex()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
+  const downloadModelMutation = useModelDownloadMutation()
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const { data: assistants } = useAssistantQuery()
@@ -102,10 +98,13 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, persistModelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
-    addDownloadState(persistModelId)
-    await downloadModel(modelHandle, fileName, persistModelId)
-  }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId])
+  const onDownloadClick = useCallback(() => {
+    downloadModelMutation.mutate({
+      modelId: modelHandle,
+      fileName: fileName,
+      persistedModelId: persistModelId,
+    })
+  }, [downloadModelMutation, modelHandle, fileName, persistModelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/Settings/HuggingFaceRepoDetailModal/ModelDownloadRow/index.tsx

@@ -10,12 +10,8 @@ import { toaster } from '@/containers/Toast'
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
 
-import useCortex from '@/hooks/useCortex'
-
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage, toGibibytes } from '@/utils/converter'
@@ -69,8 +65,7 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   fileName,
 }) => {
-  const { downloadModel } = useCortex()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
+  const downloadModelMutation = useModelDownloadMutation()
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const setHfImportingStage = useSetAtom(importHuggingFaceModelStageAtom)
   const { createThread } = useThreads()
@@ -93,10 +88,13 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, persistModelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
-    addDownloadState(persistModelId)
-    await downloadModel(modelHandle, fileName, persistModelId)
-  }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId])
+  const onDownloadClick = useCallback(() => {
+    downloadModelMutation.mutate({
+      modelId: modelHandle,
+      fileName: fileName,
+      persistedModelId: persistModelId,
+    })
+  }, [downloadModelMutation, modelHandle, fileName, persistModelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {