diff --git a/crates/dbx-core/src/agent_manager.rs b/crates/dbx-core/src/agent_manager.rs index d594c63b9..e5729d613 100644 --- a/crates/dbx-core/src/agent_manager.rs +++ b/crates/dbx-core/src/agent_manager.rs @@ -104,6 +104,14 @@ mod tests { assert!(err.contains("Custom Java runtime does not exist")); } + #[test] + fn stores_configured_app_version_for_agent_handshake() { + let dir = std::env::temp_dir().join(format!("dbx-agent-manager-version-{}", uuid::Uuid::new_v4())); + let manager = AgentManager::new_with_base_dir_and_app_version(dir, "0.5.13"); + + assert_eq!(manager.agent_app_version(), "0.5.13"); + } + #[test] fn resolves_system_java_runtime_from_path() { let manager = test_manager("system"); @@ -190,6 +198,7 @@ pub struct AgentDriverInfo { pub struct AgentManager { base_dir: PathBuf, + app_version: String, daemons: Mutex>, } @@ -201,7 +210,12 @@ impl AgentManager { } pub fn new_with_base_dir(base_dir: PathBuf) -> Self { - let mgr = Self { base_dir, daemons: Mutex::new(std::collections::HashMap::new()) }; + Self::new_with_base_dir_and_app_version(base_dir, env!("CARGO_PKG_VERSION")) + } + + pub fn new_with_base_dir_and_app_version(base_dir: PathBuf, app_version: impl Into) -> Self { + let mgr = + Self { base_dir, app_version: app_version.into(), daemons: Mutex::new(std::collections::HashMap::new()) }; mgr.migrate_legacy_jre(); mgr } @@ -218,6 +232,10 @@ impl AgentManager { &self.base_dir } + pub fn agent_app_version(&self) -> &str { + &self.app_version + } + pub fn jre_dir(&self, jre_key: &str) -> PathBuf { self.base_dir.join(format!("jre-{jre_key}")) } @@ -320,7 +338,9 @@ impl AgentManager { let java = self.resolve_java_runtime(&state, jre_key)?.to_string_lossy().to_string(); let jar = self.driver_jar_path(key).to_string_lossy().to_string(); - AgentDriverClient::spawn(&java, &jar).await + let mut client = AgentDriverClient::spawn(&java, &jar).await?; + client.try_optional_handshake(self.agent_app_version()).await; + Ok(client) } pub async fn call_daemon( @@ -345,7 +365,8 @@ impl AgentManager { } let java = self.resolve_java_runtime(&state, jre_key)?.to_string_lossy().to_string(); let jar = self.driver_jar_path(&key).to_string_lossy().to_string(); - let client = AgentDriverClient::spawn(&java, &jar).await?; + let mut client = AgentDriverClient::spawn(&java, &jar).await?; + client.try_optional_handshake(self.agent_app_version()).await; daemons.insert(key.clone(), client); } @@ -360,6 +381,7 @@ impl AgentManager { let java = self.resolve_java_runtime(&state, jre_key)?.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?; + new_client.try_optional_handshake(self.agent_app_version()).await; let result = new_client.call::(method, params).await?; daemons.insert(key, new_client); Ok(result) diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 814c3f6c2..706ff4924 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -84,6 +84,14 @@ impl AppState { } pub fn new_with_plugin_dir(storage: Storage, plugin_dir: PathBuf) -> Self { + Self::new_with_plugin_dir_and_app_version(storage, plugin_dir, env!("CARGO_PKG_VERSION")) + } + + pub fn new_with_plugin_dir_and_app_version( + storage: Storage, + plugin_dir: PathBuf, + app_version: impl Into, + ) -> Self { Self { connections: RwLock::new(HashMap::new()), configs: RwLock::new(HashMap::new()), @@ -92,7 +100,10 @@ impl AppState { proxy_tunnels: ProxyTunnelManager::new(), storage, plugins: PluginRegistry::new(plugin_dir), - agent_manager: crate::agent_manager::AgentManager::new(), + agent_manager: crate::agent_manager::AgentManager::new_with_base_dir_and_app_version( + default_agent_dir(), + app_version, + ), } } @@ -429,8 +440,16 @@ impl AppState { } fn default_plugin_dir() -> PathBuf { + default_dbx_dir().join("plugins") +} + +fn default_agent_dir() -> PathBuf { + default_dbx_dir().join("agents") +} + +fn default_dbx_dir() -> PathBuf { let home = std::env::var(if cfg!(windows) { "USERPROFILE" } else { "HOME" }).unwrap_or_else(|_| ".".to_string()); - PathBuf::from(home).join(".dbx").join("plugins") + PathBuf::from(home).join(".dbx") } pub fn connection_url_for_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> String { diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 54e162a61..336c18e7e 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -5,8 +5,10 @@ use std::sync::{Arc, Mutex}; use std::time::Duration; use serde::de::DeserializeOwned; +use serde::Deserialize; use serde_json::Value; +pub const AGENT_PROTOCOL_VERSION: u32 = 1; const RPC_TIMEOUT_SECS: u64 = 30; const STARTUP_TIMEOUT_SECS: u64 = 15; const STDERR_TAIL_LINES: usize = 20; @@ -21,6 +23,14 @@ pub struct AgentDriverClient { next_id: u64, } +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentHandshake { + pub protocol_version: u32, + pub agent_protocol_version: u32, + pub capabilities: Vec, +} + struct StderrTail { lines: VecDeque, capacity: usize, @@ -195,6 +205,28 @@ impl AgentDriverClient { result.map_err(|e| self.format_agent_process_error(&e)) } + pub async fn try_optional_handshake(&mut self, app_version: &str) -> Option { + match self.call::("handshake", agent_handshake_params(app_version)).await { + Ok(handshake) => { + log::info!( + "[agent] handshake complete: protocol={}, agent_protocol={}, capabilities={:?}", + handshake.protocol_version, + handshake.agent_protocol_version, + handshake.capabilities + ); + Some(handshake) + } + Err(err) if is_unsupported_handshake_error(&err) => { + log::info!("[agent] handshake unsupported by this driver; continuing with legacy protocol"); + None + } + Err(err) => { + log::warn!("[agent] handshake failed; continuing with legacy protocol: {err}"); + None + } + } + } + /// Send a shutdown message to the agent and wait for the process to exit. pub async fn shutdown(&mut self) { // Try to send a shutdown RPC; ignore errors if the agent is already gone @@ -225,6 +257,19 @@ impl AgentDriverClient { } } +pub fn agent_handshake_params(app_version: &str) -> Value { + serde_json::json!({ + "appVersion": app_version, + "supportedProtocolVersions": [AGENT_PROTOCOL_VERSION], + }) +} + +pub fn is_unsupported_handshake_error(error: &str) -> bool { + error.contains("Unknown method: handshake") + || error.contains("Method not found: handshake") + || error.contains("method not found: handshake") +} + fn agent_java_args(jar_path: &str) -> Vec { [ "-Dfile.encoding=UTF-8", @@ -331,7 +376,10 @@ impl Drop for AgentDriverClient { #[cfg(test)] mod tests { - use super::{agent_java_args, agent_proxy_env_vars, format_agent_process_error, read_agent_line, StderrTail}; + use super::{ + agent_handshake_params, agent_java_args, agent_proxy_env_vars, format_agent_process_error, + is_unsupported_handshake_error, read_agent_line, AgentHandshake, StderrTail, AGENT_PROTOCOL_VERSION, + }; use std::io::Cursor; #[test] @@ -402,4 +450,32 @@ mod tests { assert_eq!(stderr_tail.snapshot(), "line 2\nline 3\nline 4"); } + + #[test] + fn builds_agent_handshake_request_params() { + let params = agent_handshake_params("0.5.13"); + + assert_eq!(params["appVersion"], "0.5.13"); + assert_eq!(params["supportedProtocolVersions"], serde_json::json!([AGENT_PROTOCOL_VERSION])); + } + + #[test] + fn decodes_agent_handshake_response() { + let handshake: AgentHandshake = serde_json::from_value(serde_json::json!({ + "protocolVersion": 1, + "agentProtocolVersion": 1, + "capabilities": ["connect", "query", "metadata"] + })) + .unwrap(); + + assert_eq!(handshake.protocol_version, 1); + assert_eq!(handshake.agent_protocol_version, 1); + assert_eq!(handshake.capabilities, vec!["connect", "query", "metadata"]); + } + + #[test] + fn treats_unknown_handshake_method_as_compatible_fallback() { + assert!(is_unsupported_handshake_error("Agent RPC error (-1): Unknown method: handshake")); + assert!(!is_unsupported_handshake_error("Agent RPC error (-1): Connection failed")); + } } diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 585d99300..91072b3d4 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -44,7 +44,11 @@ async fn main() { let db_path = data_dir.join("dbx.db"); let storage = Storage::open(&db_path).await.expect("Failed to open storage"); storage.migrate_from_json(&data_dir).await.expect("Failed to migrate JSON data"); - Arc::new(AppState::new_with_plugin_dir(storage, data_dir.join("plugins"))) + Arc::new(AppState::new_with_plugin_dir_and_app_version( + storage, + data_dir.join("plugins"), + env!("CARGO_PKG_VERSION"), + )) }; // Password hash: env var takes priority, then database diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index f29fe73a0..a5d185bb0 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -119,7 +119,11 @@ pub fn run() { }); eprintln!("[STARTUP] storage ready in {:?}", t.elapsed()); - let state = Arc::new(AppState::new_with_plugin_dir(storage, data_dir.join("plugins"))); + let state = Arc::new(AppState::new_with_plugin_dir_and_app_version( + storage, + data_dir.join("plugins"), + env!("CARGO_PKG_VERSION"), + )); app.manage(state.clone()); app.manage(commands::external_sql::ExternalSqlOpenState::default());