diff --git a/web/containers/Providers/EventListener.tsx b/web/containers/Providers/EventListener.tsx index b82efb066..d91a877d6 100644 --- a/web/containers/Providers/EventListener.tsx +++ b/web/containers/Providers/EventListener.tsx @@ -22,8 +22,12 @@ export default function EventListenerWrapper({ children }: PropsWithChildren) { const modelsRef = useRef(models) const { setDownloadedModels, downloadedModels } = useGetDownloadedModels() - const { setDownloadState, setDownloadStateSuccess, setDownloadStateFailed } = - useDownloadState() + const { + setDownloadState, + setDownloadStateSuccess, + setDownloadStateFailed, + setDownloadStateCancelled, + } = useDownloadState() const downloadedModelRef = useRef(downloadedModels) useEffect(() => { @@ -52,13 +56,18 @@ export default function EventListenerWrapper({ children }: PropsWithChildren) { window.electronAPI.onFileDownloadError( async (_event: string, state: any) => { - if (state.err?.message !== 'aborted') - console.error('Download error', state) const modelName = await baseName(state.fileName) const model = modelsRef.current.find( (model) => modelBinFileName(model) === modelName ) - if (model) setDownloadStateFailed(model.id) + if (model) { + if (state.err?.message !== 'aborted') { + console.error('Download error', state) + setDownloadStateFailed(model.id, state.err.message) + } else { + setDownloadStateCancelled(model.id) + } + } } ) diff --git a/web/hooks/useDownloadState.ts b/web/hooks/useDownloadState.ts index 0d7967ac5..bf5177228 100644 --- a/web/hooks/useDownloadState.ts +++ b/web/hooks/useDownloadState.ts @@ -29,26 +29,49 @@ const setDownloadStateSuccessAtom = atom(null, (get, set, modelId: string) => { }) }) -const setDownloadStateFailedAtom = atom(null, (get, set, modelId: string) => { - const currentState = { ...get(modelDownloadStateAtom) } - const state = currentState[modelId] - if (!state) { - console.debug(`Cannot find download state for ${modelId}`) - toaster({ - title: 'Cancel Download', - description: `Model ${modelId} cancel download`, - }) - return +const setDownloadStateFailedAtom = atom( + null, + (get, set, modelId: string, error: string) => { + const currentState = { ...get(modelDownloadStateAtom) } + const state = currentState[modelId] + if (!state) { + console.debug(`Cannot find download state for ${modelId}`) + toaster({ + title: 'Download Failed', + description: `Model ${modelId} download failed: ${error}`, + type: 'error', + }) + return + } + delete currentState[modelId] + set(modelDownloadStateAtom, currentState) } - delete currentState[modelId] - set(modelDownloadStateAtom, currentState) -}) +) +const setDownloadStateCancelledAtom = atom( + null, + (get, set, modelId: string) => { + const currentState = { ...get(modelDownloadStateAtom) } + const state = currentState[modelId] + if (!state) { + console.debug(`Cannot find download state for ${modelId}`) + toaster({ + title: 'Cancel Download', + description: `Model ${modelId} cancel download`, + }) + + return + } + delete currentState[modelId] + set(modelDownloadStateAtom, currentState) + } +) export function useDownloadState() { const modelDownloadState = useAtomValue(modelDownloadStateAtom) const setDownloadState = useSetAtom(setDownloadStateAtom) const setDownloadStateSuccess = useSetAtom(setDownloadStateSuccessAtom) const setDownloadStateFailed = useSetAtom(setDownloadStateFailedAtom) + const setDownloadStateCancelled = useSetAtom(setDownloadStateCancelledAtom) const downloadStates: DownloadState[] = [] for (const [, value] of Object.entries(modelDownloadState)) { @@ -61,6 +84,7 @@ export function useDownloadState() { setDownloadState, setDownloadStateSuccess, setDownloadStateFailed, + setDownloadStateCancelled, downloadStates, } }