From a450405b0de88ade82a1f2b9eeea48ececa9b30f Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Wed, 13 May 2026 05:08:55 +0800
Subject: [PATCH] perf(agent): reuse daemon process for test_connection to
 avoid JVM cold start

Keep a per-db-type daemon agent process alive in AgentManager. test_connection
now reuses the warm daemon instead of spawning a new JVM each time. If the
daemon dies, it auto-respawns on next call.
---
 crates/dbx-core/src/agent_manager.rs | 48 +++++++++++++++++++++++++++-
 src-tauri/src/commands/connection.rs |  7 ++--
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/crates/dbx-core/src/agent_manager.rs b/crates/dbx-core/src/agent_manager.rs
index 09e650466..acf0859b5 100644
--- a/crates/dbx-core/src/agent_manager.rs
+++ b/crates/dbx-core/src/agent_manager.rs
@@ -1,6 +1,7 @@
 use std::path::{Path, PathBuf};
 
 use serde::{Deserialize, Serialize};
+use tokio::sync::Mutex;
 
 use crate::db::agent_driver::AgentDriverClient;
 use crate::models::connection::DatabaseType;
@@ -59,13 +60,17 @@ pub struct AgentDriverInfo {
 
 pub struct AgentManager {
     base_dir: PathBuf,
+    daemons: Mutex<std::collections::HashMap<String, AgentDriverClient>>,
 }
 
 impl AgentManager {
     pub fn new() -> Self {
         let home =
             std::env::var(if cfg!(windows) { "USERPROFILE" } else { "HOME" }).unwrap_or_else(|_| ".".to_string());
-        Self { base_dir: PathBuf::from(home).join(".dbx").join("agents") }
+        Self {
+            base_dir: PathBuf::from(home).join(".dbx").join("agents"),
+            daemons: Mutex::new(std::collections::HashMap::new()),
+        }
     }
 
     pub fn base_dir(&self) -> &PathBuf {
@@ -140,6 +145,47 @@ impl AgentManager {
         AgentDriverClient::spawn(&java, &jar).await
     }
 
+    pub async fn call_daemon<T: serde::de::DeserializeOwned + Send + 'static>(
+        &self,
+        db_type: &DatabaseType,
+        method: &str,
+        params: serde_json::Value,
+    ) -> Result<T, String> {
+        let key = Self::db_type_to_agent_key(db_type)
+            .ok_or_else(|| format!("{:?} is not an agent-driven database type", db_type))?
+            .to_string();
+
+        let mut daemons = self.daemons.lock().await;
+
+        if !daemons.contains_key(&key) {
+            if !self.is_jre_installed() {
+                return Err("JRE runtime is not installed. Please install it from the Driver Manager.".to_string());
+            }
+            if !self.is_driver_installed(&key) {
+                return Err(format!("{key} driver is not installed. Please install it from the Driver Manager."));
+            }
+            let java = self.jre_java_path().to_string_lossy().to_string();
+            let jar = self.driver_jar_path(&key).to_string_lossy().to_string();
+            let client = AgentDriverClient::spawn(&java, &jar).await?;
+            daemons.insert(key.clone(), client);
+        }
+
+        let client = daemons.get_mut(&key).unwrap();
+        match client.call::<T>(method, params.clone()).await {
+            Ok(result) => Ok(result),
+            Err(e) => {
+                log::warn!("[agent] daemon call failed, respawning: {e}");
+                daemons.remove(&key);
+                let java = self.jre_java_path().to_string_lossy().to_string();
+                let jar = self.driver_jar_path(&key).to_string_lossy().to_string();
+                let mut new_client = AgentDriverClient::spawn(&java, &jar).await?;
+                let result = new_client.call::<T>(method, params).await?;
+                daemons.insert(key, new_client);
+                Ok(result)
+            }
+        }
+    }
+
     pub async fn download_file(url: &str, dest: &Path) -> Result<(), String> {
         if let Some(parent) = dest.parent() {
             std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs
index 77d01f3e4..c38edd95c 100644
--- a/src-tauri/src/commands/connection.rs
+++ b/src-tauri/src/commands/connection.rs
@@ -113,9 +113,10 @@ pub async fn test_connection(state: State<'_, Arc<AppState>>, config: Connection
                 db::elasticsearch_driver::test_connection(&client).await.map(|_| "Connection successful".to_string())
             }
             DatabaseType::Dameng | DatabaseType::Kingbase | DatabaseType::Vastbase | DatabaseType::Goldendb => {
-                let mut client = state.agent_manager.spawn(&config.db_type).await?;
-                client
-                    .call::<serde_json::Value>(
+                state
+                    .agent_manager
+                    .call_daemon::<serde_json::Value>(
+                        &config.db_type,
                         "test_connection",
                         serde_json::json!({
                             "host": host,