feat(agent): support cancelling installs and upgrades

src-tauri/src/commands/agents.rs

@@ -4,10 +4,12 @@ use tauri::{Emitter, State};
 
 use dbx_core::agent_manager::{AgentDriverInfo, DriverStoreUsage, JavaRuntimeConfig, JavaRuntimeMode, DEFAULT_JRE_KEY};
 use dbx_core::agent_service::{
-    build_agent_list, clear_agent_download_cache, fetch_registry_from, import_agent_driver,
-    import_agents_from_package as import_agents_from_package_core, inspect_offline_package, install_agent_driver_from,
-    invalidate_registry_cache, reinstall_agent_jre_from, uninstall_agent_driver, uninstall_agent_jre,
-    upgrade_all_agent_drivers_from, AgentProgressEvent, OfflineImportPlan, UpgradeAllAgentDriversResult,
+    batch_cancellation_key, build_agent_list, cancel_agent_batch_upgrade, cancel_agent_driver_install,
+    clear_agent_download_cache, fetch_registry_from, fetch_registry_from_claimed, import_agent_driver,
+    import_agents_from_package as import_agents_from_package_core, inspect_offline_package,
+    install_agent_driver_from_claimed, install_cancellation_key, invalidate_registry_cache, reinstall_agent_jre_from,
+    uninstall_agent_driver, uninstall_agent_jre, upgrade_all_agent_drivers_from_claimed, AgentProgressEvent,
+    OfflineImportPlan, UpgradeAllAgentDriversResult,
 };
 use dbx_core::connection::AppState;
 use dbx_core::driver_runtime::DriverRuntimeSummary;
@@ -71,13 +73,29 @@ pub async fn install_agent(
     source: Option<DownloadSource>,
     operation_id: Option<String>,
 ) -> Result<(), String> {
-    ensure_no_agent_update_blockers(state.inner().as_ref(), std::slice::from_ref(&db_type)).await?;
-    let app_handle = app.clone();
+    // Resolve the operation id first, then register the cancellation token
+    // under it BEFORE any awaitable setup (blocker check, lock wait, registry
+    // fetch) so a cancel fired while the UI shows the modal is observed by this
+    // exact install instead of being silently lost or crossing into a second
+    // same-driver install.
     let operation_id = operation_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
-    install_agent_driver_from(&state.agent_manager, &db_type, source.unwrap_or_default(), move |event| {
-        emit_agent_progress(&app_handle, &operation_id, event)
-    })
-    .await
+    let cancellation = state.agent_manager.begin_install_cancellation(&install_cancellation_key(&operation_id)).await;
+    let result = async {
+        ensure_no_agent_update_blockers(state.inner().as_ref(), std::slice::from_ref(&db_type)).await?;
+        let app_handle = app.clone();
+        let progress_operation_id = operation_id.clone();
+        install_agent_driver_from_claimed(
+            &state.agent_manager,
+            &db_type,
+            source.unwrap_or_default(),
+            move |event| emit_agent_progress(&app_handle, &progress_operation_id, event),
+            &cancellation,
+        )
+        .await
+    }
+    .await;
+    state.agent_manager.finish_install_cancellation(&install_cancellation_key(&operation_id), &cancellation).await;
+    result
 }
 
 #[tauri::command]
@@ -88,17 +106,48 @@ pub async fn upgrade_all_agents(
     operation_id: Option<String>,
 ) -> Result<UpgradeAllAgentDriversResult, String> {
     let source = source.unwrap_or_default();
-    let registry = fetch_registry_from(source).await?;
-    let agents = build_agent_list(&state.agent_manager, Some(&registry));
-    let updatable: Vec<String> =
-        agents.iter().filter(|agent| agent.update_available).map(|agent| agent.db_type.clone()).collect();
-    ensure_no_agent_update_blockers(state.inner().as_ref(), &updatable).await?;
-    let app_handle = app.clone();
+    // Resolve the batch operation id first, then register the batch token under
+    // it BEFORE the registry fetch + blocker check so a cancel fired while the
+    // batch is still setting up aborts it instead of being lost.
     let operation_id = operation_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
-    upgrade_all_agent_drivers_from(&state.agent_manager, source, move |event| {
-        emit_agent_progress(&app_handle, &operation_id, event)
-    })
-    .await
+    let cancellation = state.agent_manager.begin_install_cancellation(&batch_cancellation_key(&operation_id)).await;
+    let result = async {
+        let registry = fetch_registry_from_claimed(source, &[cancellation.as_ref()]).await?;
+        let agents = build_agent_list(&state.agent_manager, Some(&registry));
+        let updatable: Vec<String> =
+            agents.iter().filter(|agent| agent.update_available).map(|agent| agent.db_type.clone()).collect();
+        ensure_no_agent_update_blockers(state.inner().as_ref(), &updatable).await?;
+        let app_handle = app.clone();
+        let progress_operation_id = operation_id.clone();
+        upgrade_all_agent_drivers_from_claimed(
+            &state.agent_manager,
+            source,
+            move |event| emit_agent_progress(&app_handle, &progress_operation_id, event),
+            &cancellation,
+            &operation_id,
+        )
+        .await
+    }
+    .await;
+    state.agent_manager.finish_install_cancellation(&batch_cancellation_key(&operation_id), &cancellation).await;
+    result
+}
+
+#[tauri::command]
+pub async fn cancel_agent_install(
+    state: State<'_, Arc<AppState>>,
+    db_type: String,
+    operation_id: Option<String>,
+) -> Result<(), String> {
+    cancel_agent_driver_install(&state.agent_manager, &db_type, operation_id.as_deref()).await
+}
+
+#[tauri::command]
+pub async fn cancel_agent_upgrade_all(
+    state: State<'_, Arc<AppState>>,
+    operation_id: Option<String>,
+) -> Result<(), String> {
+    cancel_agent_batch_upgrade(&state.agent_manager, operation_id.as_deref()).await
 }
 
 #[tauri::command]