add cancel download

Author: gau-nernst

extensions/model-extension/src/index.ts

@@ -160,6 +160,7 @@ export default class JanModelExtension extends ModelExtension {
    * @returns {Promise<void>} A promise that resolves when the download has been cancelled.
    */
   async cancelModelPull(model: string): Promise<void> {
+    return invoke<void>("cancel_download_task", {taskId: model}).catch(console.error)
     /**
      * Sending DELETE to /models/pull/{id} endpoint to cancel a model pull
      */

src-tauri/Cargo.toml

@@ -41,6 +41,7 @@ rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "mai
 uuid = { version = "1.7", features = ["v4"] }
 env = "1.0.1"
 futures-util = "0.3.31"
+tokio-util = "0.7.14"
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 tauri-plugin-updater = "2"

src-tauri/src/core/state.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashMap, sync::Arc};
 
+use crate::core::utils::download::DownloadManagerState;
 use rand::{distributions::Alphanumeric, Rng};
 use rmcp::{service::RunningService, RoleClient};
 use tokio::sync::Mutex;
@@ -8,6 +9,7 @@ use tokio::sync::Mutex;
 pub struct AppState {
     pub app_token: Option<String>,
     pub mcp_servers: Arc<Mutex<HashMap<String, RunningService<RoleClient, ()>>>>,
+    pub download_manager: Arc<Mutex<DownloadManagerState>>,
 }
 pub fn generate_app_token() -> String {
     rand::thread_rng()

src-tauri/src/core/utils/download.rs

@@ -1,8 +1,17 @@
+use crate::core::state::AppState;
 use futures_util::StreamExt;
+use std::collections::HashMap;
 use std::path::Path;
 use std::time::Duration;
+use tauri::State;
 use tokio::fs::File;
 use tokio::io::AsyncWriteExt;
+use tokio_util::sync::CancellationToken;
+
+#[derive(Default)]
+pub struct DownloadManagerState {
+    pub cancel_tokens: HashMap<String, CancellationToken>,
+}
 
 // this is to emulate the current way of downloading files by Cortex + Jan
 // we can change this later
@@ -27,10 +36,10 @@ pub struct DownloadEvent {
 pub async fn download<F>(
     url: &str,
     save_path: &Path,
+    cancel_token: Option<CancellationToken>,
     mut callback: Option<F>,
 ) -> Result<(), Box<dyn std::error::Error>>
 where
-    // F: FnMut(u64) + Send + 'static,
     F: FnMut(u64),
 {
     let client = reqwest::Client::builder()
@@ -58,7 +67,16 @@ where
 
     // write chunk to file, and call callback if needed (e.g. download progress)
     let mut stream = resp.bytes_stream();
+    let mut is_cancelled = false;
     while let Some(chunk) = stream.next().await {
+        if let Some(token) = cancel_token.as_ref() {
+            if token.is_cancelled() {
+                log::info!("Download cancelled: {}", url);
+                is_cancelled = true;
+                break;
+            }
+        }
+
         let chunk = chunk?;
         file.write_all(&chunk).await?;
 
@@ -67,6 +85,23 @@ where
             cb(chunk.len() as u64);
         }
     }
+
+    // cleanup
     file.flush().await?;
+    if is_cancelled {
+        // NOTE: we don't check error here
+        let _ = std::fs::remove_file(save_path);
+    }
     Ok(())
 }
+
+#[tauri::command]
+pub async fn cancel_download_task(state: State<'_, AppState>, task_id: &str) -> Result<(), String> {
+    let mut download_manager = state.download_manager.lock().await;
+    if let Some(token) = download_manager.cancel_tokens.remove(task_id) {
+        token.cancel();
+        Ok(())
+    } else {
+        Err(format!("No download task with id {}", task_id))
+    }
+}

src-tauri/src/core/utils/hf.rs

@@ -1,7 +1,9 @@
+use crate::core::state::AppState;
 use crate::core::utils::download::{download, DownloadEvent, DownloadEventType};
 use std::path::Path;
 use std::sync::{Arc, Mutex};
-use tauri::{AppHandle, Emitter};
+use tauri::{AppHandle, Emitter, State};
+use tokio_util::sync::CancellationToken;
 
 // task_id is only for compatibility with current Jan+Cortex we can remove it later
 // should we disallow custom save_dir? i.e. all HF models will be downloaded to
@@ -10,12 +12,22 @@ use tauri::{AppHandle, Emitter};
 #[tauri::command]
 pub async fn download_hf_repo(
     app: AppHandle,
+    state: State<'_, AppState>,
     task_id: &str,
     repo_id: &str,
     branch: &str,
     save_dir: &Path,
 ) -> Result<(), String> {
-    // TODO: check if it has been downloaded
+    // TODO: check if it has been/is being downloaded
+
+    // check if task_id already exists
+    {
+        let download_manager = state.download_manager.lock().await;
+        if download_manager.cancel_tokens.contains_key(task_id) {
+            return Err(format!("Task ID {} already exists", task_id));
+        }
+    }
+
     let files = list_files(repo_id, branch)
         .await
         .map_err(|e| format!("Failed to list files {}", e))?;
@@ -33,6 +45,15 @@ pub async fn download_hf_repo(
     app.emit("download", info_arc.lock().unwrap().clone())
         .unwrap();
 
+    // insert cancel tokens
+    let cancel_token = CancellationToken::new();
+    {
+        let mut download_manager = state.download_manager.lock().await;
+        download_manager
+            .cancel_tokens
+            .insert(task_id.to_string(), cancel_token.clone());
+    }
+
     let download_result = async {
         // NOTE: currently we are downloading sequentially. we can spawn tokio tasks
         // to download files in parallel.
@@ -55,13 +76,27 @@ pub async fn download_hf_repo(
                     app.emit("download", info.clone()).unwrap();
                 }
             };
-            download(&url, &full_path, Some(callback))
+            download(&url, &full_path, Some(cancel_token.clone()), Some(callback))
                 .await
                 .map_err(|e| format!("Failed to download file {}: {}", file.path, e))?;
         }
         Ok(())
     }
     .await;
+
+    // cleanup
+    {
+        let mut download_manager = state.download_manager.lock().await;
+        download_manager.cancel_tokens.remove(task_id);
+
+        if cancel_token.is_cancelled() && save_dir.exists() {
+            // NOTE: we don't check error here
+            let _ = std::fs::remove_dir_all(save_dir);
+        }
+    }
+
+    // report results
+    // TODO: what if it is cancelled? do we still emit success?
     match download_result {
         Ok(_) => {
             let mut info = info_arc.lock().unwrap();

src-tauri/src/lib.rs

@@ -3,6 +3,7 @@ use core::{
     cmd::get_jan_data_folder_path,
     setup::{self, setup_engine_binaries, setup_mcp, setup_sidecar},
     state::{generate_app_token, AppState},
+    utils::download::DownloadManagerState,
 };
 use std::{collections::HashMap, sync::Arc};
 
@@ -61,10 +62,12 @@ pub fn run() {
             core::threads::modify_thread_assistant,
             // Download
             core::utils::hf::download_hf_repo,
+            core::utils::download::cancel_download_task,
         ])
         .manage(AppState {
             app_token: Some(generate_app_token()),
             mcp_servers: Arc::new(Mutex::new(HashMap::new())),
+            download_manager: Arc::new(Mutex::new(DownloadManagerState::default())),
         })
         .setup(|app| {
             app.handle().plugin(