feat(agents): concurrent all-agent-driver upgrade

crates/dbx-core/src/agent_manager.rs

@@ -2,6 +2,7 @@ use std::ffi::OsStr;
 use std::fs::File;
 use std::io::Read;
 use std::path::{Path, PathBuf};
+use std::sync::{Arc, Mutex as StdMutex};
 
 use serde::{Deserialize, Serialize};
 use tokio::sync::Mutex;
@@ -435,6 +436,18 @@ pub struct AgentManager {
     pub(crate) connection_runtimes: Mutex<
         std::collections::HashMap<String, std::sync::Arc<tokio::sync::OnceCell<std::sync::Arc<AgentRuntimeClient>>>>,
     >,
+    /// Serializes `load_state` → modify → `save_state` to prevent lost updates
+    /// when multiple driver installs run concurrently.
+    pub(crate) state_lock: StdMutex<()>,
+    /// Per-JRE-key install locks so that concurrent driver installs sharing the
+    /// same JRE download it only once (DCL pattern: lock → re-check installed → download).
+    pub(crate) jre_install_locks: Mutex<std::collections::HashMap<String, Arc<Mutex<()>>>>,
+    /// Per-driver locks serialize install, import, and uninstall operations
+    /// targeting the same on-disk agent files.
+    pub(crate) driver_operation_locks: Mutex<std::collections::HashMap<String, Arc<Mutex<()>>>>,
+    /// Driver operations may run concurrently, but JRE replacement/removal
+    /// must exclude them until their dependent driver state is persisted.
+    pub(crate) installation_operation_lock: tokio::sync::RwLock<()>,
 }
 
 impl Default for AgentManager {
@@ -460,6 +473,10 @@ impl AgentManager {
             app_version: app_version.into(),
             daemons: Mutex::new(std::collections::HashMap::new()),
             connection_runtimes: Mutex::new(std::collections::HashMap::new()),
+            state_lock: StdMutex::new(()),
+            jre_install_locks: Mutex::new(std::collections::HashMap::new()),
+            driver_operation_locks: Mutex::new(std::collections::HashMap::new()),
+            installation_operation_lock: tokio::sync::RwLock::new(()),
         };
         mgr.migrate_legacy_jre();
         mgr.cleanup_pending_jre_dirs();
@@ -467,6 +484,19 @@ impl AgentManager {
         mgr
     }
 
+    /// Atomically load, modify, and persist the installation state.
+    ///
+    /// Keep the closure free of async work: this lock exists specifically to
+    /// prevent one installation operation from saving an older snapshot over
+    /// another operation's successful update.
+    pub fn mutate_state<T>(&self, mutate: impl FnOnce(&mut AgentState) -> T) -> Result<T, String> {
+        let _guard = self.state_lock.lock().map_err(|_| "Agent installation state lock was poisoned".to_string())?;
+        let mut state = self.load_state();
+        let result = mutate(&mut state);
+        self.save_state(&state)?;
+        Ok(result)
+    }
+
     fn migrate_legacy_jre(&self) {
         let legacy = self.base_dir.join("jre");
         let versioned = self.jre_dir(DEFAULT_JRE_KEY);
@@ -480,25 +510,26 @@ impl AgentManager {
     /// removals are pruned from the persisted state. Failures are kept for the
     /// next launch and never block startup. (Issue #1100, D6.)
     fn cleanup_pending_jre_dirs(&self) {
-        let mut state = self.load_state();
-        if state.pending_jre_cleanup.is_empty() {
+        if self.load_state().pending_jre_cleanup.is_empty() {
             return;
         }
-        let mut remaining = Vec::new();
-        for path in std::mem::take(&mut state.pending_jre_cleanup) {
-            if !path.exists() {
-                continue;
-            }
-            match std::fs::remove_dir_all(&path) {
-                Ok(()) => log::info!("Cleaned up pending JRE stash: {}", path.display()),
-                Err(err) => {
-                    log::warn!("Pending JRE cleanup failed for {}: {err}", path.display());
-                    remaining.push(path);
+
+        if let Err(err) = self.mutate_state(|state| {
+            let mut remaining = Vec::new();
+            for path in std::mem::take(&mut state.pending_jre_cleanup) {
+                if !path.exists() {
+                    continue;
+                }
+                match std::fs::remove_dir_all(&path) {
+                    Ok(()) => log::info!("Cleaned up pending JRE stash: {}", path.display()),
+                    Err(err) => {
+                        log::warn!("Pending JRE cleanup failed for {}: {err}", path.display());
+                        remaining.push(path);
+                    }
                 }
             }
-        }
-        state.pending_jre_cleanup = remaining;
-        if let Err(err) = self.save_state(&state) {
+            state.pending_jre_cleanup = remaining;
+        }) {
             log::warn!("Failed to persist post-cleanup AgentState: {err}");
         }
     }

crates/dbx-core/src/agent_service.rs

@@ -1,8 +1,11 @@
 use std::hash::{Hash, Hasher};
 use std::io::{Read, Seek, SeekFrom};
 use std::path::{Path, PathBuf};
+use std::sync::Arc;
 use std::time::Duration;
 
+use futures::{stream, StreamExt};
+
 use crate::agent_catalog;
 use crate::agent_manager::{
     AgentDriverInfo, AgentManager, AgentRegistry, InstalledDriver, JavaRuntimeMode, DEFAULT_JRE_KEY,
@@ -18,6 +21,10 @@ const JRE_REMOVE_ATTEMPTS: usize = if cfg!(windows) { 6 } else { 1 };
 /// Exponential-ish backoff between retries. Total wait ≈ 1.55s on Windows.
 const JRE_REMOVE_BACKOFF_MS: &[u64] = &[50, 100, 200, 400, 400, 400];
 
+/// Keep batch updates concurrent without allowing a large registry to exhaust
+/// the download server, local disk, or the application's file descriptors.
+const MAX_CONCURRENT_AGENT_UPDATES: usize = 4;
+
 /// Delete an old JRE directory, retrying on Windows to cover the daemon-exit
 /// and AV-scan release window. Returns the original `std::io::Error` when all
 /// retries fail so callers can decide whether to fall back to rename-stash.
@@ -93,7 +100,7 @@ fn stash_old_jre_dir(path: &Path) -> std::io::Result<PathBuf> {
 /// path (Some) if the rename fallback was used so the caller can persist it
 /// for startup cleanup, or None if the directory was deleted outright (or
 /// did not exist).
-fn replace_old_jre_dir(am: &AgentManager, path: &Path) -> Result<Option<PathBuf>, String> {
+fn replace_old_jre_dir(path: &Path) -> Result<Option<PathBuf>, String> {
     match remove_jre_dir_with_retry(path) {
         Ok(()) => Ok(None),
         Err(remove_err) => {
@@ -102,13 +109,8 @@ fn replace_old_jre_dir(am: &AgentManager, path: &Path) -> Result<Option<PathBuf>
                 match stash_old_jre_dir(path) {
                     Ok(stash) => {
                         log::warn!("remove_dir_all failed, stashed old JRE at {} ({remove_err})", stash.display());
-                        // Persist immediately so a crash before extraction
-                        // still leaves the stash recorded for cleanup.
-                        let mut state = am.load_state();
-                        state.pending_jre_cleanup.push(stash.clone());
-                        if let Err(save_err) = am.save_state(&state) {
-                            log::warn!("Failed to persist pending_jre_cleanup: {save_err}");
-                        }
+                        // The caller will persist this stash under
+                        // state_lock after extraction succeeds.
                         Ok(Some(stash))
                     }
                     Err(rename_err) => {
@@ -122,7 +124,6 @@ fn replace_old_jre_dir(am: &AgentManager, path: &Path) -> Result<Option<PathBuf>
             }
             #[cfg(not(windows))]
             {
-                let _ = am; // silence unused warning on POSIX
                 Err(format_jre_dir_remove_error(path, &remove_err))
             }
         }
@@ -138,6 +139,8 @@ static REGISTRY_CACHE: std::sync::LazyLock<
 
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
 pub struct AgentProgressEvent {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub operation_id: Option<String>,
     pub step: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub downloaded: Option<u64>,
@@ -165,7 +168,15 @@ pub struct UpgradeAllAgentDriversResult {
 
 impl AgentProgressEvent {
     pub fn step(step: impl Into<String>) -> Self {
-        Self { step: step.into(), downloaded: None, total: None, db_type: None, current: None, total_drivers: None }
+        Self {
+            operation_id: None,
+            step: step.into(),
+            downloaded: None,
+            total: None,
+            db_type: None,
+            current: None,
+            total_drivers: None,
+        }
     }
 
     pub fn transfer(step: impl Into<String>, downloaded: u64, total: u64) -> Self {
@@ -178,6 +189,11 @@ impl AgentProgressEvent {
         self.total_drivers = total_drivers;
         self
     }
+
+    pub fn with_operation_id(mut self, operation_id: &str) -> Self {
+        self.operation_id = Some(operation_id.to_string());
+        self
+    }
 }
 
 pub fn build_agent_list(am: &AgentManager, registry: Option<&AgentRegistry>) -> Vec<AgentDriverInfo> {
@@ -303,27 +319,32 @@ pub fn find_local_agent_jar(db_type: &str) -> Option<PathBuf> {
 }
 
 pub fn install_local_agent(am: &AgentManager, db_type: &str, source: PathBuf) -> Result<(), String> {
+    install_local_agent_file(am, db_type, &source)?;
+    am.mutate_state(|state| record_local_agent_install(state, db_type, DEFAULT_JRE_KEY))
+}
+
+fn install_local_agent_file(am: &AgentManager, db_type: &str, source: &Path) -> Result<(), String> {
     let jar_path = am.driver_jar_path(db_type);
     let parent = jar_path.parent().ok_or_else(|| format!("Invalid driver path: {}", jar_path.display()))?;
     std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
     let staging_path = parent.join(format!(".agent-jar-import-{}", uuid::Uuid::new_v4()));
-    std::fs::copy(&source, &staging_path).map_err(|e| format!("Failed to copy local agent jar: {e}"))?;
+    std::fs::copy(source, &staging_path).map_err(|e| format!("Failed to copy local agent jar: {e}"))?;
     if !is_valid_agent_jar(&staging_path) {
         std::fs::remove_file(&staging_path).ok();
         return Err(format!("Local agent jar is invalid or corrupt: {}", source.display()));
     }
-    replace_imported_agent_file(&staging_path, &jar_path)?;
+    replace_imported_agent_file(&staging_path, &jar_path)
+}
 
-    let mut local_state = am.load_state();
-    local_state.installed_drivers.insert(
+fn record_local_agent_install(state: &mut crate::agent_manager::AgentState, db_type: &str, jre_key: &str) {
+    state.installed_drivers.insert(
         db_type.to_string(),
         InstalledDriver {
             version: "0.1.0-local".to_string(),
             installed_at: chrono::Utc::now().to_rfc3339(),
-            jre: DEFAULT_JRE_KEY.to_string(),
+            jre: jre_key.to_string(),
         },
     );
-    am.save_state(&local_state)
 }
 
 fn is_valid_agent_jar(path: &Path) -> bool {
@@ -423,27 +444,45 @@ pub async fn upgrade_all_agent_drivers_from(
     progress: impl Fn(AgentProgressEvent),
 ) -> Result<UpgradeAllAgentDriversResult, String> {
     let registry = fetch_registry_from(source).await?;
-    let agents = build_agent_list(am, Some(&registry));
-    let updatable: Vec<&AgentDriverInfo> = agents.iter().filter(|agent| agent.update_available).collect();
-    let total = updatable.len() as u32;
-    let mut result = UpgradeAllAgentDriversResult::default();
+    upgrade_all_agent_drivers_with_registry(am, &registry, source, &progress).await
+}
 
-    for (index, agent) in updatable.iter().enumerate() {
-        match install_agent_driver_from_registry(
+async fn upgrade_all_agent_drivers_with_registry(
+    am: &AgentManager,
+    registry: &AgentRegistry,
+    source: DownloadSource,
+    progress: &impl Fn(AgentProgressEvent),
+) -> Result<UpgradeAllAgentDriversResult, String> {
+    let agents = build_agent_list(am, Some(registry));
+    let updatable: Vec<String> =
+        agents.iter().filter(|agent| agent.update_available).map(|agent| agent.db_type.clone()).collect();
+    let total = updatable.len() as u32;
+
+    // Run independent driver installs concurrently, with a fixed upper bound
+    // so a large registry cannot saturate download and file-system resources.
+    let installs = updatable.into_iter().enumerate().map(|(index, db_type)| async move {
+        let result = install_agent_driver_from_registry_locked(
             am,
-            &registry,
+            registry,
             source,
-            &agent.db_type,
-            &progress,
+            &db_type,
+            progress,
             Some((index + 1) as u32),
             Some(total),
         )
-        .await
-        {
+        .await;
+        (db_type, result)
+    });
+
+    let outcomes = stream::iter(installs).buffer_unordered(MAX_CONCURRENT_AGENT_UPDATES).collect::<Vec<_>>().await;
+
+    let mut result = UpgradeAllAgentDriversResult::default();
+    for (db_type, outcome) in outcomes {
+        match outcome {
             Ok(()) => result.upgraded += 1,
             Err(error) => {
-                log::warn!("Failed to update {} agent driver: {}", agent.db_type, error);
-                result.failed.push(AgentDriverUpdateIssue { db_type: agent.db_type.clone(), error });
+                log::warn!("Failed to update {} agent driver: {}", db_type, error);
+                result.failed.push(AgentDriverUpdateIssue { db_type, error });
             }
         }
     }
@@ -452,7 +491,20 @@ pub async fn upgrade_all_agent_drivers_from(
     Ok(result)
 }
 
+async fn driver_operation_lock(am: &AgentManager, db_type: &str) -> Arc<tokio::sync::Mutex<()>> {
+    let mut locks = am.driver_operation_locks.lock().await;
+    locks.entry(db_type.to_string()).or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))).clone()
+}
+
+async fn jre_operation_lock(am: &AgentManager, jre_key: &str) -> Arc<tokio::sync::Mutex<()>> {
+    let mut locks = am.jre_install_locks.lock().await;
+    locks.entry(jre_key.to_string()).or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))).clone()
+}
+
 pub async fn uninstall_agent_driver(am: &AgentManager, db_type: &str) -> Result<(), String> {
+    let _installation_guard = am.installation_operation_lock.read().await;
+    let driver_lock = driver_operation_lock(am, db_type).await;
+    let _driver_guard = driver_lock.lock().await;
     prune_driver_download_cache(am, db_type)?;
     let jar_path = am.driver_jar_path(db_type);
     if jar_path.exists() {
@@ -463,9 +515,7 @@ pub async fn uninstall_agent_driver(am: &AgentManager, db_type: &str) -> Result<
             std::fs::remove_dir_all(driver_dir).map_err(|err| err.to_string())?;
         }
     }
-    let mut local_state = am.load_state();
-    local_state.installed_drivers.remove(db_type);
-    am.save_state(&local_state)?;
+    am.mutate_state(|state| state.installed_drivers.remove(db_type))?;
     am.stop_daemon_by_key(db_type).await;
     Ok(())
 }
@@ -475,6 +525,11 @@ pub fn clear_agent_download_cache(am: &AgentManager) -> Result<(), String> {
 }
 
 pub async fn uninstall_agent_jre(am: &AgentManager, jre_key: &str) -> Result<(), String> {
+    // Keep the dependency check and removal atomic with respect to driver
+    // installs/uninstalls that may add or remove a dependency on this JRE.
+    let _installation_guard = am.installation_operation_lock.write().await;
+    let jre_lock = jre_operation_lock(am, jre_key).await;
+    let _jre_guard = jre_lock.lock().await;
     let local_state = am.load_state();
     let dependents: Vec<&str> = local_state
         .installed_drivers
@@ -492,9 +547,7 @@ pub async fn uninstall_agent_jre(am: &AgentManager, jre_key: &str) -> Result<(),
     if let Err(err) = remove_jre_dir_with_retry(&jre_dir) {
         return Err(format_jre_dir_remove_error(&jre_dir, &err));
     }
-    let mut local_state = am.load_state();
-    local_state.jre_versions.remove(jre_key);
-    am.save_state(&local_state)?;
+    am.mutate_state(|state| state.jre_versions.remove(jre_key))?;
     Ok(())
 }
 
@@ -512,6 +565,11 @@ pub async fn reinstall_agent_jre_from(
     source: DownloadSource,
     progress: impl Fn(AgentProgressEvent),
 ) -> Result<(), String> {
+    // Replacing a JRE must not race a driver operation that is using or about
+    // to persist a dependency on the same runtime.
+    let _installation_guard = am.installation_operation_lock.write().await;
+    let jre_lock = jre_operation_lock(am, jre_key).await;
+    let _jre_guard = jre_lock.lock().await;
     let registry = fetch_registry_from(source).await?;
     let jre_info = registry.resolve_jre(jre_key).ok_or_else(|| format!("No JRE definition for version: {jre_key}"))?;
     let platform = AgentManager::current_platform();
@@ -540,32 +598,33 @@ pub async fn reinstall_agent_jre_from(
     // handles on Windows (Issue #1100). Falls back to a rename-stash if the
     // directory still cannot be removed.
     am.stop_daemons().await;
-    replace_old_jre_dir(am, &jre_dir)?;
+    let stash = replace_old_jre_dir(&jre_dir)?;
+    persist_pending_jre_cleanup(am, stash.as_ref()).await?;
     extract_tar_gz(&jre_archive, &jre_dir)?;
     std::fs::remove_file(&jre_archive).ok();
-    let mut local_state = am.load_state();
-    local_state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone());
-    am.save_state(&local_state)?;
+    am.mutate_state(|state| state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone()))?;
     cleanup_jre_download_cache_after_success(am, jre_key);
     progress(AgentProgressEvent::step("done"));
     Ok(())
 }
 
-pub fn import_agents_from_zip(
+pub async fn import_agents_from_zip(
     am: &AgentManager,
     zip_path: &Path,
     progress: impl Fn(AgentProgressEvent),
 ) -> Result<OfflineImportResult, String> {
     import_offline_zip(am, zip_path, |p| {
         progress(AgentProgressEvent {
+            operation_id: None,
             step: p.step,
             downloaded: Some(p.current as u64),
             total: Some(p.total as u64),
-            db_type: Some(p.label),
+            db_type: p.db_type,
             current: Some(p.current),
             total_drivers: Some(p.total),
         });
     })
+    .await
 }
 
 async fn install_agent_driver_with_batch(
@@ -575,6 +634,35 @@ async fn install_agent_driver_with_batch(
     progress: &impl Fn(AgentProgressEvent),
     current: Option<u32>,
     total_drivers: Option<u32>,
+) -> Result<(), String> {
+    let _installation_guard = am.installation_operation_lock.read().await;
+    let driver_lock = driver_operation_lock(am, db_type).await;
+    let _driver_guard = driver_lock.lock().await;
+    install_agent_driver_with_batch_unlocked(am, db_type, source, progress, current, total_drivers).await
+}
+
+async fn install_agent_driver_from_registry_locked(
+    am: &AgentManager,
+    registry: &AgentRegistry,
+    source: DownloadSource,
+    db_type: &str,
+    progress: &impl Fn(AgentProgressEvent),
+    current: Option<u32>,
+    total_drivers: Option<u32>,
+) -> Result<(), String> {
+    let _installation_guard = am.installation_operation_lock.read().await;
+    let driver_lock = driver_operation_lock(am, db_type).await;
+    let _driver_guard = driver_lock.lock().await;
+    install_agent_driver_from_registry(am, registry, source, db_type, progress, current, total_drivers).await
+}
+
+async fn install_agent_driver_with_batch_unlocked(
+    am: &AgentManager,
+    db_type: &str,
+    source: DownloadSource,
+    progress: &impl Fn(AgentProgressEvent),
+    current: Option<u32>,
+    total_drivers: Option<u32>,
 ) -> Result<(), String> {
     match fetch_registry_from(source).await {
         Ok(registry) => {
@@ -623,13 +711,30 @@ async fn ensure_jre_from_registry(
     current: Option<u32>,
     total_drivers: Option<u32>,
 ) -> Result<(), String> {
+    // Fast path: already installed — return immediately without acquiring the
+    // per-JRE lock.  The lock is only needed when a download + extract may be
+    // required.
+    if !jre_needs_install(am, registry, jre_key) {
+        return Ok(());
+    }
+
+    // Acquire (or create) the per-JRE-key mutex so that concurrent driver
+    // installs sharing the same JRE download it exactly once.
+    let lock = jre_operation_lock(am, jre_key).await;
+    let _jre_guard = lock.lock().await;
+
+    // Double-check: the previous lock holder may have already installed.
+    if !jre_needs_install(am, registry, jre_key) {
+        return Ok(());
+    }
+
     let jre_info = registry.resolve_jre(jre_key).ok_or_else(|| format!("No JRE definition for version: {jre_key}"))?;
     let platform = AgentManager::current_platform();
     let platform_jre = jre_info
         .platforms
         .get(platform)
         .ok_or_else(|| format!("No JRE {jre_key} available for platform: {platform}"))?;
-    let jre_archive = am.base_dir().join("jre-download.tar.gz");
+    let jre_archive = am.base_dir().join(format!("jre-{}-download.tar.gz", jre_key));
     progress(AgentProgressEvent::transfer("jre", 0, platform_jre.size).with_batch(
         Some(db_type),
         current,
@@ -652,15 +757,67 @@ async fn ensure_jre_from_registry(
     .await?;
     progress(AgentProgressEvent::transfer("jre-extract", 0, 0).with_batch(Some(db_type), current, total_drivers));
     let jre_dir = am.jre_dir(jre_key);
-    // Stop daemons first (Windows ERROR_ACCESS_DENIED, Issue #1100).
-    am.stop_daemons().await;
-    replace_old_jre_dir(am, &jre_dir)?;
+    // Stop only daemons that use this JRE before replacing its directory
+    // (Windows ERROR_ACCESS_DENIED, Issue #1100).  In a concurrent
+    // upgrade-all this avoids killing unrelated daemons mid-install.
+    stop_daemons_using_jre(am, jre_key).await;
+    let stash = replace_old_jre_dir(&jre_dir)?;
+
+    // Persist the stash path *before* extraction so that a crash during
+    // extract_tar_gz (or a process kill) doesn't leave the renamed-stash
+    // directory as an orphan that never gets cleaned up.
+    persist_pending_jre_cleanup(am, stash.as_ref()).await?;
+
     extract_tar_gz(&jre_archive, &jre_dir)?;
     std::fs::remove_file(&jre_archive).ok();
     cleanup_jre_download_cache_after_success(am, jre_key);
+
+    // Persist the JRE version after extraction succeeds, while still holding
+    // the per-JRE lock.  This guarantees the DCL in another task's
+    // jre_needs_install() sees the installed version and skips download.
+    am.mutate_state(|state| state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone()))?;
+
     Ok(())
 }
 
+/// Stop daemons whose installed driver lists `jre_key` as its runtime.
+async fn stop_daemons_using_jre(am: &AgentManager, jre_key: &str) {
+    let state = am.load_state();
+    for (db_type, driver) in &state.installed_drivers {
+        if driver.jre == jre_key {
+            am.stop_daemon_by_key(db_type).await;
+        }
+    }
+}
+
+/// Record a rename-stashed JRE before extraction so startup can clean it up
+/// even if the process exits mid-install.
+async fn persist_pending_jre_cleanup(am: &AgentManager, stash: Option<&PathBuf>) -> Result<(), String> {
+    let Some(stash_path) = stash else {
+        return Ok(());
+    };
+
+    am.mutate_state(|state| {
+        if !state.pending_jre_cleanup.contains(stash_path) {
+            state.pending_jre_cleanup.push(stash_path.clone());
+        }
+    })
+}
+
+async fn persist_local_agent_install_state(
+    am: &AgentManager,
+    db_type: &str,
+    jre_key: &str,
+    jre_version: Option<&str>,
+) -> Result<(), String> {
+    am.mutate_state(|state| {
+        if let Some(version) = jre_version {
+            state.jre_versions.insert(jre_key.to_string(), version.to_string());
+        }
+        record_local_agent_install(state, db_type, jre_key);
+    })
+}
+
 async fn install_local_agent_with_registry_jre(
     am: &AgentManager,
     registry: &AgentRegistry,
@@ -675,12 +832,16 @@ async fn install_local_agent_with_registry_jre(
     if jre_needs_install(am, registry, jre_key) {
         ensure_jre_from_registry(am, registry, source, jre_key, db_type, progress, current, total_drivers).await?;
     }
-    install_local_agent(am, db_type, local_jar)?;
-    if let Some(jre_info) = registry.resolve_jre(jre_key) {
-        let mut local_state = am.load_state();
-        local_state.jre_versions.insert(jre_key.to_string(), jre_info.version.clone());
-        am.save_state(&local_state)?;
-    }
+    install_local_agent_file(am, db_type, &local_jar)?;
+    // This fallback can run for several drivers during upgrade-all. Keep its
+    // driver and JRE updates in one state_lock-protected transaction.
+    persist_local_agent_install_state(
+        am,
+        db_type,
+        jre_key,
+        registry.resolve_jre(jre_key).map(|jre| jre.version.as_str()),
+    )
+    .await?;
     am.stop_daemon_by_key(db_type).await;
     progress(AgentProgressEvent::step("done").with_batch(Some(db_type), current, total_drivers));
     Ok(())
@@ -765,21 +926,21 @@ async fn install_agent_driver_from_registry(
         std::fs::remove_file(am.driver_native_path(db_type)).ok();
     }
 
-    let mut local_state = am.load_state();
-    if requires_java_runtime {
-        if let Some(jre_info) = registry.resolve_jre(jre_key) {
-            local_state.jre_versions.insert(jre_key.clone(), jre_info.version.clone());
+    am.mutate_state(|state| {
+        if requires_java_runtime {
+            if let Some(jre_info) = registry.resolve_jre(jre_key) {
+                state.jre_versions.insert(jre_key.clone(), jre_info.version.clone());
+            }
         }
-    }
-    local_state.installed_drivers.insert(
-        db_type.to_string(),
-        InstalledDriver {
-            version: driver.version.clone(),
-            installed_at: chrono::Utc::now().to_rfc3339(),
-            jre: jre_key.clone(),
-        },
-    );
-    am.save_state(&local_state)?;
+        state.installed_drivers.insert(
+            db_type.to_string(),
+            InstalledDriver {
+                version: driver.version.clone(),
+                installed_at: chrono::Utc::now().to_rfc3339(),
+                jre: jre_key.clone(),
+            },
+        );
+    })?;
     am.stop_daemon_by_key(db_type).await;
     cleanup_driver_download_cache_after_success(am, db_type);
     progress(AgentProgressEvent::step("done").with_batch(Some(db_type), current, total_drivers));
@@ -1234,7 +1395,12 @@ pub struct OfflineImportProgress {
     pub step: String,
     pub current: u32,
     pub total: u32,
+    /// Display label for the current item (e.g. "MySQL", "JRE 21.0.12").
     pub label: String,
+    /// The real database-type key (e.g. "mysql"), used by the frontend for
+    /// per-driver progress routing. `None` for JRE-only steps.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub db_type: Option<String>,
 }
 
 #[derive(Debug, Clone)]
@@ -1265,11 +1431,16 @@ pub fn inspect_offline_zip(zip_path: &Path) -> Result<OfflineImportPlan, String>
     })
 }
 
-pub fn import_offline_zip(
+pub async fn import_offline_zip(
     am: &AgentManager,
     zip_path: &Path,
     progress: impl Fn(OfflineImportProgress),
 ) -> Result<OfflineImportResult, String> {
+    // Offline import can touch both JRE and driver directories — hold an
+    // exclusive installation-operation lock so that concurrent driver installs,
+    // JRE installs, Upgrade All, and uninstall operations are serialised.
+    let _installation_guard = am.installation_operation_lock.write().await;
+
     let file = std::fs::File::open(zip_path).map_err(|e| format!("Failed to open ZIP file: {e}"))?;
     let mut archive = zip::ZipArchive::new(file).map_err(|e| format!("Invalid ZIP file: {e}"))?;
 
@@ -1298,7 +1469,13 @@ pub fn import_offline_zip(
             continue;
         }
 
-        progress(OfflineImportProgress { step: "jre-extract".into(), current, total, label: format!("JRE {jre_key}") });
+        progress(OfflineImportProgress {
+            step: "jre-extract".into(),
+            current,
+            total,
+            label: format!("JRE {jre_key}"),
+            db_type: None,
+        });
 
         let mut entry = archive.by_name(entry_name).map_err(|e| format!("Failed to read {entry_name}: {e}"))?;
         let tmp_archive = am.base_dir().join(format!("jre-offline-{jre_key}.tar.gz"));
@@ -1352,6 +1529,7 @@ pub fn import_offline_zip(
             current,
             total,
             label: agent_catalog::label_for_key(db_type).unwrap_or(db_type).to_string(),
+            db_type: Some(db_type.clone()),
         });
 
         let driver_path = if *is_native { am.driver_native_path(db_type) } else { am.driver_jar_path(db_type) };
@@ -1396,7 +1574,23 @@ pub fn import_offline_zip(
         result.drivers_installed.push(db_type.clone());
     }
 
-    am.save_state(&local_state)?;
+    am.mutate_state(|state| {
+        for jre_key in &result.jre_installed {
+            if let Some(version) = local_state.jre_versions.get(jre_key) {
+                state.jre_versions.insert(jre_key.clone(), version.clone());
+            }
+        }
+        for path in &local_state.pending_jre_cleanup {
+            if !state.pending_jre_cleanup.contains(path) {
+                state.pending_jre_cleanup.push(path.clone());
+            }
+        }
+        for db_type in &result.drivers_installed {
+            if let Some(driver) = local_state.installed_drivers.get(db_type) {
+                state.installed_drivers.insert(db_type.clone(), driver.clone());
+            }
+        }
+    })?;
     Ok(result)
 }
 
@@ -1541,7 +1735,14 @@ fn extract_tar_gz(archive: &Path, dest: &Path) -> Result<(), String> {
     Ok(())
 }
 
-pub fn import_agent_driver(am: &AgentManager, db_type: &str, source_path: &Path) -> Result<(), String> {
+pub async fn import_agent_driver(am: &AgentManager, db_type: &str, source_path: &Path) -> Result<(), String> {
+    // Manual imports replace the same artifact paths as downloads. Reuse the
+    // install operation and per-driver locks so an import cannot race an
+    // install, Upgrade All, or uninstall for this driver.
+    let _installation_guard = am.installation_operation_lock.read().await;
+    let driver_lock = driver_operation_lock(am, db_type).await;
+    let _driver_guard = driver_lock.lock().await;
+
     if !source_path.is_file() {
         return Err(format!("File not found: {}", source_path.display()));
     }
@@ -1562,20 +1763,20 @@ pub fn import_agent_driver(am: &AgentManager, db_type: &str, source_path: &Path)
     replace_imported_agent_file(&staging_path, &native_path)?;
     std::fs::remove_file(am.driver_jar_path(db_type)).ok();
 
-    let mut local_state = am.load_state();
-    local_state.installed_drivers.insert(
-        db_type.to_string(),
-        InstalledDriver {
-            version: "0.1.0-local".to_string(),
-            installed_at: chrono::Utc::now().to_rfc3339(),
-            jre: DEFAULT_JRE_KEY.to_string(),
-        },
-    );
-    am.save_state(&local_state)
+    am.mutate_state(|state| {
+        state.installed_drivers.insert(
+            db_type.to_string(),
+            InstalledDriver {
+                version: "0.1.0-local".to_string(),
+                installed_at: chrono::Utc::now().to_rfc3339(),
+                jre: DEFAULT_JRE_KEY.to_string(),
+            },
+        );
+    })
 }
 
-pub fn import_agent_jar(am: &AgentManager, db_type: &str, jar_path: &Path) -> Result<(), String> {
-    import_agent_driver(am, db_type, jar_path)
+pub async fn import_agent_jar(am: &AgentManager, db_type: &str, jar_path: &Path) -> Result<(), String> {
+    import_agent_driver(am, db_type, jar_path).await
 }
 
 fn replace_imported_agent_file(staging_path: &Path, target_path: &Path) -> Result<(), String> {
@@ -1804,13 +2005,23 @@ mod agent_download_url_tests {
 #[cfg(test)]
 mod agent_registry_install_tests {
     use super::*;
-    use crate::agent_manager::{ArtifactInfo, DriverInfo, JavaRuntimeConfig};
+    use crate::agent_manager::{ArtifactInfo, DriverInfo, InstalledDriver, JavaRuntimeConfig, JreInfo};
 
     fn test_manager(name: &str) -> AgentManager {
         let dir = std::env::temp_dir().join(format!("dbx-agent-registry-install-{name}-{}", uuid::Uuid::new_v4()));
         AgentManager::new_with_base_dir(dir)
     }
 
+    fn write_test_agent_jar(path: &Path) {
+        use std::io::Write;
+
+        let file = std::fs::File::create(path).unwrap();
+        let mut archive = zip::ZipWriter::new(file);
+        archive.start_file("META-INF/MANIFEST.MF", zip::write::SimpleFileOptions::default()).unwrap();
+        archive.write_all(b"Manifest-Version: 1.0\nMain-Class: com.dbx.Agent\n").unwrap();
+        archive.finish().unwrap();
+    }
+
     fn registry_with_native_and_legacy_jar(
         db_type: &str,
         version: &str,
@@ -1871,6 +2082,57 @@ mod agent_registry_install_tests {
         cache_path
     }
 
+    fn registry_with_jre(jre_key: &str, version: &str, url: &str, size: u64) -> AgentRegistry {
+        AgentRegistry {
+            jre: None,
+            jres: [(
+                jre_key.to_string(),
+                JreInfo {
+                    version: version.to_string(),
+                    platforms: [(
+                        AgentManager::current_platform().to_string(),
+                        ArtifactInfo { url: url.to_string(), size },
+                    )]
+                    .into_iter()
+                    .collect(),
+                },
+            )]
+            .into_iter()
+            .collect(),
+            drivers: std::collections::HashMap::new(),
+        }
+    }
+
+    fn build_jre_archive(am: &AgentManager, jre_key: &str) -> Vec<u8> {
+        let archive_root = am.base_dir().join("jre-test-archive");
+        let payload = archive_root.join("payload");
+        let java_path = am.jre_java_path(jre_key);
+        let relative_java_path = java_path.strip_prefix(am.jre_dir(jre_key)).unwrap();
+        let java_path = payload.join(relative_java_path);
+        std::fs::create_dir_all(java_path.parent().unwrap()).unwrap();
+        std::fs::write(java_path, b"java").unwrap();
+        let archive = archive_root.join("runtime.tar.gz");
+        let status = crate::process::new_std_command("tar")
+            .args(["czf", &archive.to_string_lossy(), "-C", &archive_root.to_string_lossy(), "payload"])
+            .status()
+            .unwrap();
+        assert!(status.success());
+        std::fs::read(archive).unwrap()
+    }
+
+    fn write_cached_jre_download(am: &AgentManager, jre_key: &str, version: &str, url: &str, archive: &[u8]) {
+        let dest = am.base_dir().join(format!("jre-{jre_key}-download.tar.gz"));
+        let cache_path = cached_download_path(
+            am,
+            url,
+            archive.len() as u64,
+            Some(CacheIdentity::Jre { key: jre_key, version }),
+            &dest,
+        );
+        std::fs::create_dir_all(cache_path.parent().unwrap()).unwrap();
+        std::fs::write(cache_path, archive).unwrap();
+    }
+
     #[tokio::test]
     async fn registry_install_accepts_native_driver_with_legacy_jar_fallback() {
         let manager = test_manager("native-with-jar-fallback");
@@ -1908,6 +2170,263 @@ mod agent_registry_install_tests {
             .any(|event| event.step == "done" && event.db_type.as_deref() == Some(db_type)));
     }
 
+    #[tokio::test]
+    async fn batch_upgrade_preserves_successful_state_and_reports_independent_failure() {
+        let manager = test_manager("batch-upgrade");
+        let oracle_url = "https://example.com/dbx-agent-oracle";
+        let dameng_url = "https://example.com/dbx-agent-dameng";
+        let kingbase_url = "https://example.com/dbx-agent-kingbase.jar";
+        let oracle_bytes = b"oracle-native-agent";
+        let dameng_bytes = b"dameng-native-agent";
+        let corrupt_jar = b"not-a-jar";
+
+        let mut registry =
+            registry_with_native_and_legacy_jar("oracle", "2.0.0", oracle_url, oracle_bytes.len() as u64);
+        registry.drivers.extend(
+            registry_with_native_and_legacy_jar("dameng", "2.0.0", dameng_url, dameng_bytes.len() as u64).drivers,
+        );
+        registry.drivers.extend(registry_with_jar("kingbase", "2.0.0", kingbase_url, corrupt_jar.len() as u64).drivers);
+
+        let mut state = manager.load_state();
+        state.java_runtime = JavaRuntimeConfig { mode: JavaRuntimeMode::System, custom_java_path: None };
+        for (db_type, version) in [("oracle", "1.0.0"), ("dameng", "1.0.0"), ("kingbase", "1.0.0")] {
+            state.installed_drivers.insert(
+                db_type.to_string(),
+                InstalledDriver {
+                    version: version.to_string(),
+                    installed_at: "2026-01-01T00:00:00Z".to_string(),
+                    jre: DEFAULT_JRE_KEY.to_string(),
+                },
+            );
+        }
+        manager.save_state(&state).unwrap();
+
+        write_cached_driver_download(
+            &manager,
+            "oracle",
+            "2.0.0",
+            oracle_url,
+            &manager.driver_native_path("oracle"),
+            oracle_bytes,
+        );
+        write_cached_driver_download(
+            &manager,
+            "dameng",
+            "2.0.0",
+            dameng_url,
+            &manager.driver_native_path("dameng"),
+            dameng_bytes,
+        );
+        write_cached_driver_download(
+            &manager,
+            "kingbase",
+            "2.0.0",
+            kingbase_url,
+            &manager.driver_jar_path("kingbase"),
+            corrupt_jar,
+        );
+        let events = std::sync::Mutex::new(Vec::new());
+        let progress = |event| events.lock().unwrap().push(event);
+
+        let result = upgrade_all_agent_drivers_with_registry(&manager, &registry, DownloadSource::Official, &progress)
+            .await
+            .unwrap();
+
+        assert_eq!(result.upgraded, 2);
+        assert_eq!(result.failed.len(), 1);
+        assert_eq!(result.failed[0].db_type, "kingbase");
+        let state = manager.load_state();
+        assert_eq!(state.installed_drivers["oracle"].version, "2.0.0");
+        assert_eq!(state.installed_drivers["dameng"].version, "2.0.0");
+        assert_eq!(state.installed_drivers["kingbase"].version, "1.0.0");
+        assert_eq!(events.lock().unwrap().iter().filter(|event| event.step == "done").count(), 2);
+    }
+
+    #[tokio::test]
+    async fn shared_jre_is_not_downloaded_again_after_the_first_install_persists_its_version() {
+        let manager = test_manager("shared-jre-deduplication");
+        let jre_key = DEFAULT_JRE_KEY;
+        let version = "21.0.12";
+        let url = "https://example.com/dbx-jre.tar.gz";
+        let archive = build_jre_archive(&manager, jre_key);
+        let registry = registry_with_jre(jre_key, version, url, archive.len() as u64);
+        let events = std::sync::Mutex::new(Vec::new());
+        let progress = |event| events.lock().unwrap().push(event);
+
+        write_cached_jre_download(&manager, jre_key, version, url, &archive);
+        ensure_jre_from_registry(
+            &manager,
+            &registry,
+            DownloadSource::Official,
+            jre_key,
+            "oracle",
+            &progress,
+            Some(1),
+            Some(2),
+        )
+        .await
+        .unwrap();
+
+        // The successful install cleans its cache. Re-add it so the old
+        // implementation fails deterministically by consuming it again.
+        write_cached_jre_download(&manager, jre_key, version, url, &archive);
+        ensure_jre_from_registry(
+            &manager,
+            &registry,
+            DownloadSource::Official,
+            jre_key,
+            "dameng",
+            &progress,
+            Some(2),
+            Some(2),
+        )
+        .await
+        .unwrap();
+
+        assert_eq!(manager.load_state().jre_versions[jre_key], version);
+        assert!(events
+            .lock()
+            .unwrap()
+            .iter()
+            .any(|event| event.step == "jre" && event.db_type.as_deref() == Some("oracle")));
+        assert!(!events
+            .lock()
+            .unwrap()
+            .iter()
+            .any(|event| event.step == "jre" && event.db_type.as_deref() == Some("dameng")));
+    }
+
+    #[tokio::test]
+    async fn stash_is_recorded_before_jre_extraction() {
+        let manager = test_manager("stash-before-extract");
+        let stash = manager.base_dir().join("jre-21.old-test");
+
+        persist_pending_jre_cleanup(&manager, Some(&stash)).await.unwrap();
+
+        assert_eq!(manager.load_state().pending_jre_cleanup, vec![stash]);
+    }
+
+    #[test]
+    fn concurrent_local_agent_state_updates_preserve_each_driver() {
+        let manager = test_manager("concurrent-local-agent-state");
+        let start = std::sync::Arc::new(std::sync::Barrier::new(3));
+        std::thread::scope(|scope| {
+            for db_type in ["oracle", "dameng"] {
+                let start = start.clone();
+                let manager = &manager;
+                scope.spawn(move || {
+                    start.wait();
+                    manager
+                        .mutate_state(|state| {
+                            state.jre_versions.insert(DEFAULT_JRE_KEY.to_string(), "21.0.12".to_string());
+                            record_local_agent_install(state, db_type, DEFAULT_JRE_KEY);
+                        })
+                        .unwrap();
+                });
+            }
+            start.wait();
+        });
+
+        let state = manager.load_state();
+        assert!(state.installed_drivers.contains_key("oracle"));
+        assert!(state.installed_drivers.contains_key("dameng"));
+        assert_eq!(state.jre_versions[DEFAULT_JRE_KEY], "21.0.12");
+    }
+
+    #[tokio::test]
+    async fn batch_registry_install_waits_for_an_existing_driver_operation() {
+        let manager = test_manager("batch-driver-operation-lock");
+        let db_type = "oracle";
+        let version = "0.1.31";
+        let native_url = "https://example.com/dbx-agent-oracle";
+        let native_bytes = b"native-agent";
+        let registry = registry_with_native_and_legacy_jar(db_type, version, native_url, native_bytes.len() as u64);
+        write_cached_driver_download(
+            &manager,
+            db_type,
+            version,
+            native_url,
+            &manager.driver_native_path(db_type),
+            native_bytes,
+        );
+        let first_lock = driver_operation_lock(&manager, "oracle").await;
+        let first_guard = first_lock.lock().await;
+        let progress = |_| {};
+
+        let blocked = tokio::time::timeout(std::time::Duration::from_millis(50), async {
+            install_agent_driver_from_registry_locked(
+                &manager,
+                &registry,
+                DownloadSource::Official,
+                db_type,
+                &progress,
+                Some(1),
+                Some(1),
+            )
+            .await
+        })
+        .await;
+        assert!(blocked.is_err(), "batch install entered while another operation owned the driver files");
+
+        drop(first_guard);
+        tokio::time::timeout(
+            std::time::Duration::from_secs(1),
+            install_agent_driver_from_registry_locked(
+                &manager,
+                &registry,
+                DownloadSource::Official,
+                db_type,
+                &progress,
+                Some(1),
+                Some(1),
+            ),
+        )
+        .await
+        .expect("batch install did not resume after the driver lock was released")
+        .unwrap();
+        assert_eq!(std::fs::read(manager.driver_native_path(db_type)).unwrap(), native_bytes);
+    }
+
+    #[tokio::test]
+    async fn manual_import_waits_for_an_existing_driver_operation() {
+        let manager = test_manager("manual-import-driver-operation-lock");
+        let db_type = "h2";
+        let source = manager.base_dir().join("dbx-agent-h2.jar");
+        std::fs::create_dir_all(source.parent().unwrap()).unwrap();
+        write_test_agent_jar(&source);
+        let lock = driver_operation_lock(&manager, db_type).await;
+        let first_guard = lock.lock().await;
+
+        let blocked =
+            tokio::time::timeout(std::time::Duration::from_millis(50), import_agent_driver(&manager, db_type, &source))
+                .await;
+        assert!(blocked.is_err(), "manual import entered while another operation owned the driver files");
+
+        drop(first_guard);
+        tokio::time::timeout(std::time::Duration::from_secs(1), import_agent_driver(&manager, db_type, &source))
+            .await
+            .expect("manual import did not resume after the driver lock was released")
+            .unwrap();
+        assert_eq!(std::fs::read(manager.driver_jar_path(db_type)).unwrap(), std::fs::read(source).unwrap());
+    }
+
+    #[tokio::test]
+    async fn jre_exclusive_operation_waits_for_in_flight_driver_operation() {
+        let manager = test_manager("jre-exclusive-operation-lock");
+        let driver_guard = manager.installation_operation_lock.read().await;
+
+        let blocked =
+            tokio::time::timeout(std::time::Duration::from_millis(50), manager.installation_operation_lock.write())
+                .await;
+        assert!(blocked.is_err(), "JRE replacement entered before an in-flight driver operation completed");
+
+        drop(driver_guard);
+        let _jre_guard =
+            tokio::time::timeout(std::time::Duration::from_secs(1), manager.installation_operation_lock.write())
+                .await
+                .expect("JRE replacement did not resume after driver operations completed");
+    }
+
     #[tokio::test]
     async fn registry_install_rejects_corrupt_downloaded_jar() {
         let manager = test_manager("corrupt-jar");
@@ -1943,6 +2462,24 @@ mod agent_registry_install_tests {
         assert!(!jar_path.exists());
         assert!(!manager.load_state().installed_drivers.contains_key(db_type));
     }
+
+    #[tokio::test]
+    async fn offline_import_exclusive_lock_waits_for_in_flight_driver_operation() {
+        let manager = test_manager("offline-import-lock");
+        // Simulate an in-flight driver operation holding a read lock.
+        let driver_guard = manager.installation_operation_lock.read().await;
+        // import_offline_zip acquires the write lock — it must wait.
+        let blocked =
+            tokio::time::timeout(std::time::Duration::from_millis(50), manager.installation_operation_lock.write())
+                .await;
+        assert!(blocked.is_err(), "offline import entered before an in-flight driver operation completed");
+
+        drop(driver_guard);
+        let _offline_guard =
+            tokio::time::timeout(std::time::Duration::from_secs(1), manager.installation_operation_lock.write())
+                .await
+                .expect("offline import did not resume after driver operations completed");
+    }
 }
 
 #[cfg(test)]

crates/dbx-core/src/db/ssh_tunnel.rs

@@ -163,11 +163,11 @@ async fn connect_and_authenticate(
 async fn try_authenticate_with_agent(
     session: &mut Handle<SshClient>,
     ssh_user: &str,
-    ssh_agent_sock_path: &str,
+    _ssh_agent_sock_path: &str,
     connect_timeout: &Duration,
 ) -> Result<(), String> {
     #[cfg(unix)]
-    let mut agent = if ssh_agent_sock_path.is_empty() {
+    let mut agent = if _ssh_agent_sock_path.is_empty() {
         match AgentClient::connect_env().await {
             Ok(a) => a,
             Err(e) => {
@@ -175,12 +175,12 @@ async fn try_authenticate_with_agent(
             }
         }
     } else {
-        match AgentClient::connect_uds(ssh_agent_sock_path).await {
+        match AgentClient::connect_uds(_ssh_agent_sock_path).await {
             Ok(a) => a,
             Err(e) => {
                 return Err(format!(
                     "No SSH password or key provided, and ssh-agent at '{}' is unavailable: {e}",
-                    ssh_agent_sock_path
+                    _ssh_agent_sock_path
                 ));
             }
         }