github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-15 14:52:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(agents): share runtimes across connection sessions

Browse Source
This commit is contained in:
t8y2 committed 2026-07-13 15:13:48 +08:00
1 parent 1280450b98
commit 67541232a2
67 files changed
+2904 -354
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/agent_manager.rs
+38 -4
View File
@@ -6,7 +6,7 @@ use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec, AgentMethod};
use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec, AgentMethod, AgentRuntimeClient};
use crate::models::connection::DatabaseType;
pub const DEFAULT_JRE_KEY: &str = "21";
@@ -394,6 +394,9 @@ pub struct AgentManager {
base_dir: PathBuf,
app_version: String,
pub(crate) daemons: Mutex<std::collections::HashMap<String, AgentDriverClient>>,
pub(crate) connection_runtimes: Mutex<
std::collections::HashMap<String, std::sync::Arc<tokio::sync::OnceCell<std::sync::Arc<AgentRuntimeClient>>>>,
>,
}
impl Default for AgentManager {
@@ -414,8 +417,12 @@ impl AgentManager {
}
pub fn new_with_base_dir_and_app_version(base_dir: PathBuf, app_version: impl Into<String>) -> Self {
let mgr =
Self { base_dir, app_version: app_version.into(), daemons: Mutex::new(std::collections::HashMap::new()) };
let mgr = Self {
base_dir,
app_version: app_version.into(),
daemons: Mutex::new(std::collections::HashMap::new()),
connection_runtimes: Mutex::new(std::collections::HashMap::new()),
};
mgr.migrate_legacy_jre();
mgr.cleanup_pending_jre_dirs();
mgr.cleanup_orphan_jre_dirs();
@@ -760,7 +767,13 @@ impl AgentManager {
}
pub async fn active_daemon_keys(&self) -> Vec<String> {
self.daemons.lock().await.keys().cloned().collect()
let mut keys = self.daemons.lock().await.keys().cloned().collect::<std::collections::HashSet<_>>();
for runtime_key in self.connection_runtimes.lock().await.keys() {
if let Some((agent_key, _)) = runtime_key.split_once('|') {
keys.insert(agent_key.to_string());
}
}
keys.into_iter().collect()
}
pub fn db_type_to_agent_key(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option<&'static str> {
@@ -788,6 +801,27 @@ impl AgentManager {
crate::agent_runtime::spawn_connection_client(self, db_type, driver_profile, extra_java_args).await
}
pub async fn spawn_shared_connection_client(
&self,
db_type: &DatabaseType,
driver_profile: Option<&str>,
extra_java_args: &[String],
agent_session_id: String,
connect_params: serde_json::Value,
connect_timeout: std::time::Duration,
) -> Result<AgentDriverClient, String> {
crate::agent_runtime::spawn_shared_connection_client(
self,
db_type,
driver_profile,
extra_java_args,
agent_session_id,
connect_params,
connect_timeout,
)
.await
}
pub async fn call_daemon<T: serde::de::DeserializeOwned + Send + 'static>(
&self,
db_type: &DatabaseType,
crates/dbx-core/src/agent_runtime.rs
+74 -1
View File
@@ -3,7 +3,7 @@ use std::time::Duration;
use crate::agent_manager::{AgentManager, DEFAULT_JRE_KEY};
use crate::database_capabilities;
use crate::db::agent_driver::{AgentDriverClient, AgentMethod};
use crate::db::agent_driver::{AgentDriverClient, AgentMethod, AgentRuntimeClient};
use crate::models::connection::DatabaseType;
pub fn db_type_to_agent_key(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option<&'static str> {
@@ -16,10 +16,21 @@ pub fn is_agent_type(db_type: &DatabaseType) -> bool {
pub async fn stop_daemons(manager: &AgentManager) {
manager.daemons.lock().await.clear();
let runtimes = std::mem::take(&mut *manager.connection_runtimes.lock().await);
for runtime in runtimes.into_values().filter_map(|cell| cell.get().cloned()) {
runtime.kill();
}
}
pub async fn stop_daemon_by_key(manager: &AgentManager, agent_key: &str) {
manager.daemons.lock().await.remove(agent_key);
let mut runtimes = manager.connection_runtimes.lock().await;
let matching = runtimes.keys().filter(|key| key.starts_with(&format!("{agent_key}|"))).cloned().collect::<Vec<_>>();
for key in matching {
if let Some(runtime) = runtimes.remove(&key).and_then(|cell| cell.get().cloned()) {
runtime.kill();
}
}
}
pub async fn restart_daemon_by_key(manager: &AgentManager, agent_key: &str) -> Result<(), String> {
@@ -40,6 +51,54 @@ pub async fn spawn_connection_client(
spawn_first_available_client(manager, &keys, extra_java_args).await
}
pub async fn spawn_shared_connection_client(
manager: &AgentManager,
db_type: &DatabaseType,
driver_profile: Option<&str>,
extra_java_args: &[String],
agent_session_id: String,
connect_params: serde_json::Value,
connect_timeout: Duration,
) -> Result<AgentDriverClient, String> {
let keys = runtime_agent_key_candidates(db_type, driver_profile)
.ok_or_else(|| format!("{:?} is not an agent-driven database type", db_type))?;
let key = first_installed_agent_key(manager, &keys).unwrap_or(keys[0]);
let state = manager.load_state();
let jre_key = state.installed_drivers.get(key).map(|driver| driver.jre.as_str()).unwrap_or(DEFAULT_JRE_KEY);
let launch = manager.resolve_agent_launch_spec_with_extra_args(&state, key, jre_key, extra_java_args)?;
let runtime_key = shared_runtime_key(key, &launch);
let runtime_cell = {
let mut runtimes = manager.connection_runtimes.lock().await;
if runtimes.get(&runtime_key).and_then(|cell| cell.get()).is_some_and(|runtime| runtime.is_failed()) {
runtimes.remove(&runtime_key);
}
runtimes.entry(runtime_key).or_insert_with(|| std::sync::Arc::new(tokio::sync::OnceCell::new())).clone()
};
let runtime =
runtime_cell.get_or_try_init(|| AgentRuntimeClient::spawn(launch, manager.agent_app_version())).await?.clone();
let mut session_params = connect_params;
session_params
.as_object_mut()
.ok_or_else(|| "Agent connect parameters must be an object".to_string())?
.insert("agentSessionId".to_string(), serde_json::Value::String(agent_session_id.clone()));
runtime
.call::<serde_json::Value>(AgentMethod::OpenSession.as_str(), session_params, Some(connect_timeout), None)
.await?;
runtime.increment_session_count();
Ok(AgentDriverClient::shared_session(runtime, agent_session_id))
}
fn shared_runtime_key(agent_key: &str, launch: &crate::db::agent_driver::AgentLaunchSpec) -> String {
format!(
"{}|{}|{}|{}",
agent_key,
launch.program.display(),
launch.args.join("\u{1f}"),
launch.working_dir.as_ref().map(|path| path.display().to_string()).unwrap_or_default()
)
}
pub async fn call_daemon<T: DeserializeOwned + Send + 'static>(
manager: &AgentManager,
db_type: &DatabaseType,
@@ -167,6 +226,7 @@ async fn spawn_client_for_key(
#[cfg(test)]
mod tests {
use super::*;
use std::path::PathBuf;
#[test]
fn prestosql_does_not_use_agent_driver() {
@@ -177,4 +237,17 @@ mod tests {
fn trino_uses_only_trino_agent_driver() {
assert_eq!(runtime_agent_key_candidates(&DatabaseType::Trino, None).unwrap(), vec!["trino"]);
}
#[test]
fn shared_runtime_key_includes_launch_fingerprint() {
let base = crate::db::agent_driver::AgentLaunchSpec::new(PathBuf::from("oracle-agent"))
.with_args(["--mode", "stdio"])
.with_working_dir(PathBuf::from("/tmp/oracle"));
let different_args = crate::db::agent_driver::AgentLaunchSpec::new(PathBuf::from("oracle-agent"))
.with_args(["--mode", "debug"])
.with_working_dir(PathBuf::from("/tmp/oracle"));
assert_eq!(shared_runtime_key("oracle", &base), shared_runtime_key("oracle", &base));
assert_ne!(shared_runtime_key("oracle", &base), shared_runtime_key("oracle", &different_args));
}
}
crates/dbx-core/src/connection.rs
+175 -55
View File
@@ -1264,71 +1264,161 @@ impl AppState {
agent_connection_pool_database_type!() => {
let connect_params =
agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or(""));
// Kerberos JVM properties are connection-scoped; shared agent daemons must not inherit them.
let mut client = self
.agent_manager
.spawn_with_extra_java_args(
&db_config.db_type,
db_config.driver_profile.as_deref(),
&db_config.agent_java_options,
)
.await?;
let connect_result = client
.call_method_with_timeout::<serde_json::Value>(
AgentMethod::Connect,
connect_params,
Some(agent_connect_timeout(&db_config)),
)
.await;
if let Err(err) = connect_result {
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
if !alternate_configs.is_empty() {
log::warn!(
if db_config.db_type != DatabaseType::Etcd && db_config.db_type != DatabaseType::ZooKeeper {
let agent_session_id = uuid::Uuid::new_v4().simple().to_string();
let initial_result = self
.agent_manager
.spawn_shared_connection_client(
&db_config.db_type,
db_config.driver_profile.as_deref(),
&db_config.agent_java_options,
agent_session_id.clone(),
connect_params,
agent_connect_timeout(&db_config),
)
.await;
let client = match initial_result {
Ok(client) => client,
Err(err) => {
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
if alternate_configs.is_empty() {
if err.contains("does not support multi_session protocol v2") {
let mut client = self
.agent_manager
.spawn_with_extra_java_args(
&db_config.db_type,
db_config.driver_profile.as_deref(),
&db_config.agent_java_options,
)
.await?;
client
.call_method_with_timeout::<serde_json::Value>(
AgentMethod::Connect,
agent_connect_params(
&db_config,
&host,
port,
db_config.effective_database().unwrap_or(""),
),
Some(agent_connect_timeout(&db_config)),
)
.await?;
client
} else {
return Err(oracle_error_with_driver_hint(&db_config, &err));
}
} else {
let mut fallback_errors = Vec::new();
let mut connected = None;
for alternate_config in alternate_configs {
let label =
oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
.into_iter()
.next()
.unwrap_or_else(|| "alternate".to_string());
let alternate_params = agent_connect_params(
&alternate_config,
&host,
port,
alternate_config.effective_database().unwrap_or(""),
);
match self
.agent_manager
.spawn_shared_connection_client(
&alternate_config.db_type,
alternate_config.driver_profile.as_deref(),
&alternate_config.agent_java_options,
agent_session_id.clone(),
alternate_params,
agent_connect_timeout(&alternate_config),
)
.await
{
Ok(client) => {
connected = Some(client);
break;
}
Err(alternate_err) => fallback_errors.push(format!("{label}: {alternate_err}")),
}
}
connected.ok_or_else(|| {
format!(
"{err}\n\nFallback with alternate Oracle connection descriptors failed: {}",
fallback_errors.join("\n")
)
})?
}
}
};
PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client)))
} else {
// Kerberos JVM properties are connection-scoped; shared agent daemons must not inherit them.
let mut client = self
.agent_manager
.spawn_with_extra_java_args(
&db_config.db_type,
db_config.driver_profile.as_deref(),
&db_config.agent_java_options,
)
.await?;
let connect_result = client
.call_method_with_timeout::<serde_json::Value>(
AgentMethod::Connect,
connect_params,
Some(agent_connect_timeout(&db_config)),
)
.await;
if let Err(err) = connect_result {
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
if !alternate_configs.is_empty() {
log::warn!(
"Oracle connect failed with {:?} descriptor: {}. Retrying with Oracle JDBC URL variants: {:?}.",
db_config.oracle_connection_type,
err,
oracle_alternate_connect_config_labels(&alternate_configs)
);
let mut fallback_errors = Vec::new();
let mut connected = false;
for alternate_config in alternate_configs {
let label = oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
.into_iter()
.next()
.unwrap_or_else(|| "alternate".to_string());
match client
.call_method_with_timeout::<serde_json::Value>(
AgentMethod::Connect,
agent_connect_params(
&alternate_config,
&host,
port,
alternate_config.effective_database().unwrap_or(""),
),
Some(agent_connect_timeout(&alternate_config)),
)
.await
{
Ok(_) => {
connected = true;
break;
}
Err(alternate_err) => {
fallback_errors.push(format!("{label}: {alternate_err}"));
let mut fallback_errors = Vec::new();
let mut connected = false;
for alternate_config in alternate_configs {
let label =
oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
.into_iter()
.next()
.unwrap_or_else(|| "alternate".to_string());
match client
.call_method_with_timeout::<serde_json::Value>(
AgentMethod::Connect,
agent_connect_params(
&alternate_config,
&host,
port,
alternate_config.effective_database().unwrap_or(""),
),
Some(agent_connect_timeout(&alternate_config)),
)
.await
{
Ok(_) => {
connected = true;
break;
}
Err(alternate_err) => {
fallback_errors.push(format!("{label}: {alternate_err}"));
}
}
}
if !connected {
return Err(format!(
"{err}\n\nFallback with alternate Oracle JDBC URLs failed: {}",
fallback_errors.join("\n")
));
}
} else {
return Err(oracle_error_with_driver_hint(&db_config, &err));
}
if !connected {
return Err(format!(
"{err}\n\nFallback with alternate Oracle JDBC URLs failed: {}",
fallback_errors.join("\n")
));
}
} else {
return Err(oracle_error_with_driver_hint(&db_config, &err));
}
PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client)))
}
PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client)))
}
DatabaseType::PrestoSql => {
let jdbc_config = prestosql_jdbc_config_for_endpoint(&db_config, &host, port);
@@ -4003,6 +4093,17 @@ mod tests {
);
}
#[test]
fn oracle_uses_isolated_pool_keys_for_tab_sessions() {
let mut config = mysql_config(Some("ORCLPDB1"));
config.db_type = DatabaseType::Oracle;
assert_eq!(
super::session_scoped_pool_key_for(Some(&config), "oracle-conn".to_string(), Some("table-tab-1")),
"oracle-conn:session:table-tab-1"
);
}
#[test]
fn other_agent_single_connection_types_keep_database_scoped_pool_keys() {
assert_eq!(
@@ -4338,6 +4439,25 @@ mod tests {
let _ = std::fs::remove_dir_all(dir);
}
#[tokio::test]
async fn closing_oracle_table_tab_keeps_connection_scoped_pool() {
let (state, dir) = test_app_state().await;
let mut config = mysql_config(Some("ORCLPDB1"));
config.id = "oracle-conn".to_string();
config.db_type = DatabaseType::Oracle;
state.configs.write().await.insert(config.id.clone(), config);
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
state.connections.write().await.insert("oracle-conn".to_string(), PoolKind::Sqlite(pool));
state.pool_activity.write().await.insert("oracle-conn".to_string(), super::PoolActivity::now());
assert!(!state.close_client_session_pool("oracle-conn", Some("APP"), "table-tab-1").await.unwrap());
assert!(state.connections.read().await.contains_key("oracle-conn"));
assert!(state.pool_activity.read().await.contains_key("oracle-conn"));
let _ = std::fs::remove_dir_all(dir);
}
#[test]
fn agent_validate_connection_unknown_method_is_not_stale() {
assert!(super::is_agent_validate_connection_unsupported(
crates/dbx-core/src/db/agent_driver.rs
+510 -18
View File
@@ -1,7 +1,8 @@
use std::collections::VecDeque;
use std::collections::{HashMap, VecDeque};
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, Stdio};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
@@ -10,23 +11,261 @@ use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio_util::sync::CancellationToken;
pub const AGENT_PROTOCOL_VERSION: u32 = 1;
type PendingAgentResponse = tokio::sync::oneshot::Sender<Result<Value, String>>;
pub struct AgentRuntimeClient {
child: Arc<Mutex<Child>>,
stdin: Arc<Mutex<BufWriter<ChildStdin>>>,
pending: Arc<Mutex<HashMap<u64, PendingAgentResponse>>>,
stderr_tail: Arc<Mutex<StderrTail>>,
next_id: AtomicU64,
active_sessions: AtomicU64,
failed: Arc<AtomicBool>,
handshake: AgentHandshake,
}
impl AgentRuntimeClient {
pub async fn spawn(launch: AgentLaunchSpec, app_version: &str) -> Result<Arc<Self>, String> {
let mut command = crate::process::new_std_command(&launch.program);
command.args(&launch.args).stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped());
if let Some(working_dir) = &launch.working_dir {
command.current_dir(working_dir);
}
remove_agent_proxy_env(&mut command);
let mut child =
command.spawn().map_err(|e| format!("Failed to spawn agent process {}: {e}", launch_display(&launch)))?;
let child_stdin = child.stdin.take().ok_or("Failed to capture agent stdin")?;
let child_stdout = child.stdout.take().ok_or("Failed to capture agent stdout")?;
let child_stderr = child.stderr.take().ok_or("Failed to capture agent stderr")?;
let stderr_tail = Arc::new(Mutex::new(StderrTail::default()));
start_stderr_collector(child_stderr, stderr_tail.clone());
let mut stdout = BufReader::new(child_stdout);
let stdout = tokio::time::timeout(
Duration::from_secs(STARTUP_TIMEOUT_SECS),
tokio::task::spawn_blocking(move || loop {
let line = read_agent_line(&mut stdout, "startup line")?;
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
match serde_json::from_str::<Value>(trimmed) {
Ok(value) if value.get("ready") == Some(&Value::Bool(true)) => return Ok(stdout),
Ok(_) => return Err(format!("Agent did not send ready signal, got: {line}")),
Err(_) => log::warn!("[agent:stdout] ignoring non-JSON line during startup: {trimmed}"),
}
}),
)
.await
.map_err(|_| format!("Agent startup timed out ({STARTUP_TIMEOUT_SECS}s)"))?
.map_err(|e| format!("Agent startup task failed: {e}"))??;
let runtime = Arc::new(Self {
child: Arc::new(Mutex::new(child)),
stdin: Arc::new(Mutex::new(BufWriter::new(child_stdin))),
pending: Arc::new(Mutex::new(HashMap::new())),
stderr_tail,
next_id: AtomicU64::new(0),
active_sessions: AtomicU64::new(0),
failed: Arc::new(AtomicBool::new(false)),
handshake: AgentHandshake { protocol_version: 0, agent_protocol_version: 0, capabilities: Vec::new() },
});
runtime.start_response_reader(stdout);
let handshake = runtime
.call::<AgentHandshake>(
AgentMethod::Handshake.as_str(),
agent_handshake_params(app_version),
Some(Duration::from_secs(RPC_TIMEOUT_SECS)),
None,
)
.await?;
if handshake.protocol_version < 2 || !handshake.supports(AgentCapability::MultiSession) {
runtime.kill();
return Err("Agent runtime does not support multi_session protocol v2".to_string());
}
let runtime =
Arc::try_unwrap(runtime).map_err(|_| "Agent runtime initialization is still referenced".to_string())?;
Ok(Arc::new(Self { handshake, ..runtime }))
}
fn start_response_reader(self: &Arc<Self>, mut stdout: BufReader<ChildStdout>) {
let pending = self.pending.clone();
let failed = self.failed.clone();
std::thread::spawn(move || loop {
let line = match read_agent_line(&mut stdout, "response") {
Ok(line) => line,
Err(err) => {
failed.store(true, Ordering::Release);
fail_pending_requests(&pending, err);
return;
}
};
let response: Value = match serde_json::from_str(line.trim()) {
Ok(response) => response,
Err(err) => {
failed.store(true, Ordering::Release);
fail_pending_requests(&pending, format!("Invalid JSON response from agent: {err}"));
return;
}
};
let Some(id) = response.get("id").and_then(Value::as_u64) else {
continue;
};
if let Some(sender) = pending.lock().expect("agent pending response lock poisoned").remove(&id) {
let _ = sender.send(Ok(response));
}
});
}
pub async fn call<T: DeserializeOwned>(
&self,
method: &str,
params: Value,
timeout_duration: Option<Duration>,
cancel_token: Option<CancellationToken>,
) -> Result<T, String> {
if self.failed.load(Ordering::Acquire) {
return Err("Agent runtime is unavailable".to_string());
}
let id = self.next_id.fetch_add(1, Ordering::Relaxed) + 1;
let request = serde_json::json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
let request_line =
serde_json::to_string(&request).map_err(|e| format!("Failed to serialize JSON-RPC request: {e}"))?;
let (sender, receiver) = tokio::sync::oneshot::channel();
self.pending.lock().map_err(|_| "Agent pending response lock poisoned".to_string())?.insert(id, sender);
let write_result =
self.stdin.lock().map_err(|_| "Agent stdin lock poisoned".to_string()).and_then(|mut writer| {
writer
.write_all(request_line.as_bytes())
.and_then(|_| writer.write_all(b"\n"))
.and_then(|_| writer.flush())
.map_err(|e| format!("Failed to write agent request: {e}"))
});
if let Err(err) = write_result {
self.pending.lock().expect("agent pending response lock poisoned").remove(&id);
return Err(err);
}
let receive = async { receiver.await.map_err(|_| "Agent response channel closed".to_string())? };
let response = match (timeout_duration, cancel_token) {
(Some(duration), Some(token)) => tokio::select! {
_ = token.cancelled() => {
self.cancel_session_request(&params).await;
Err("Query canceled".to_string())
},
result = tokio::time::timeout(duration, receive) => result.map_err(|_| format!("Agent RPC call timed out ({}s)", duration.as_secs()))?,
},
(Some(duration), None) => tokio::time::timeout(duration, receive)
.await
.map_err(|_| format!("Agent RPC call timed out ({}s)", duration.as_secs()))?,
(None, Some(token)) => tokio::select! {
_ = token.cancelled() => {
self.cancel_session_request(&params).await;
Err("Query canceled".to_string())
},
result = receive => result,
},
(None, None) => receive.await,
};
if response.is_err() {
self.pending.lock().expect("agent pending response lock poisoned").remove(&id);
}
decode_agent_response(response?)
}
async fn cancel_session_request(&self, params: &Value) {
let Some(agent_session_id) = params.get("agentSessionId").and_then(Value::as_str) else {
return;
};
let _ = Box::pin(self.call::<Value>(
AgentMethod::CancelSession.as_str(),
serde_json::json!({ "agentSessionId": agent_session_id }),
Some(Duration::from_secs(5)),
None,
))
.await;
}
pub fn handshake(&self) -> &AgentHandshake {
&self.handshake
}
pub fn is_failed(&self) -> bool {
self.failed.load(Ordering::Acquire)
}
pub fn active_session_count(&self) -> u64 {
self.active_sessions.load(Ordering::Acquire)
}
pub fn increment_session_count(&self) {
self.active_sessions.fetch_add(1, Ordering::AcqRel);
}
pub fn decrement_session_count(runtime: &Arc<Self>) {
let previous = runtime
.active_sessions
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |value| Some(value.saturating_sub(1)))
.unwrap_or_default();
if previous <= 1 {
let runtime = runtime.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_secs(SHARED_RUNTIME_IDLE_GRACE_SECS)).await;
if runtime.active_session_count() == 0 {
runtime.kill();
}
});
}
}
pub fn kill(&self) {
self.failed.store(true, Ordering::Release);
if let Ok(mut child) = self.child.lock() {
let _ = child.kill();
}
fail_pending_requests(&self.pending, "Agent runtime terminated".to_string());
}
}
fn decode_agent_response<T: DeserializeOwned>(response: Value) -> Result<T, String> {
if let Some(err) = response.get("error") {
let message = err.get("message").and_then(Value::as_str).unwrap_or("Unknown agent error");
let code = err.get("code").and_then(Value::as_i64).unwrap_or(-1);
return Err(format!("Agent RPC error ({code}): {message}"));
}
let result =
response.get("result").ok_or_else(|| "Agent response missing both 'result' and 'error'".to_string())?;
serde_json::from_value(result.clone()).map_err(|e| format!("Failed to deserialize agent result: {e}"))
}
fn fail_pending_requests(pending: &Arc<Mutex<HashMap<u64, PendingAgentResponse>>>, error: String) {
let requests = std::mem::take(&mut *pending.lock().expect("agent pending response lock poisoned"));
for (_, sender) in requests {
let _ = sender.send(Err(error.clone()));
}
}
pub const AGENT_PROTOCOL_VERSION: u32 = 2;
const RPC_TIMEOUT_SECS: u64 = 30;
const STARTUP_TIMEOUT_SECS: u64 = 15;
const STDERR_TAIL_LINES: usize = 20;
const AGENT_EXIT_DIAGNOSTIC_WAIT_MS: u64 = 1_000;
const AGENT_EXIT_DIAGNOSTIC_POLL_MS: u64 = 10;
const SHARED_RUNTIME_IDLE_GRACE_SECS: u64 = 30;
const AGENT_JAVA_OPTS_ENV: &str = "DBX_AGENT_JAVA_OPTS";
const AGENT_JAVA_TOO_OLD_MESSAGE: &str =
"Agent requires Java 21, but DBX started it with an older Java runtime. Use DBX managed JRE 21 or select a Java 21 executable in Driver Manager.";
pub struct AgentDriverClient {
child: Child,
child: Option<Child>,
stdin: Option<BufWriter<ChildStdin>>,
stdout: Option<BufReader<ChildStdout>>,
stderr_tail: Arc<Mutex<StderrTail>>,
handshake: Option<AgentHandshake>,
next_id: u64,
shared_runtime: Option<Arc<AgentRuntimeClient>>,
agent_session_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -93,10 +332,11 @@ pub enum AgentCapability {
Transaction,
Ddl,
Kv,
MultiSession,
}
impl AgentCapability {
pub const ALL: [Self; 8] = [
pub const ALL: [Self; 9] = [
Self::Connect,
Self::TestConnection,
Self::Metadata,
@@ -105,6 +345,7 @@ impl AgentCapability {
Self::Transaction,
Self::Ddl,
Self::Kv,
Self::MultiSession,
];
pub fn as_str(self) -> &'static str {
@@ -117,6 +358,7 @@ impl AgentCapability {
Self::Transaction => "transaction",
Self::Ddl => "ddl",
Self::Kv => "kv",
Self::MultiSession => "multi_session",
}
}
}
@@ -125,6 +367,10 @@ impl AgentCapability {
pub enum AgentMethod {
Handshake,
Connect,
OpenSession,
CloseSession,
ValidateSession,
CancelSession,
TestConnection,
ValidateConnection,
ListDatabases,
@@ -154,9 +400,13 @@ pub enum AgentMethod {
}
impl AgentMethod {
pub const ALL: [Self; 28] = [
pub const ALL: [Self; 32] = [
Self::Handshake,
Self::Connect,
Self::OpenSession,
Self::CloseSession,
Self::ValidateSession,
Self::CancelSession,
Self::TestConnection,
Self::ValidateConnection,
Self::ListDatabases,
@@ -189,6 +439,10 @@ impl AgentMethod {
match self {
Self::Handshake => "handshake",
Self::Connect => "connect",
Self::OpenSession => "open_session",
Self::CloseSession => "close_session",
Self::ValidateSession => "validate_session",
Self::CancelSession => "cancel_session",
Self::TestConnection => "test_connection",
Self::ValidateConnection => "validate_connection",
Self::ListDatabases => "list_databases",
@@ -425,7 +679,29 @@ impl AgentDriverClient {
}
};
Ok(Self { child, stdin: Some(stdin), stdout: Some(ready_stdout), stderr_tail, handshake: None, next_id: 0 })
Ok(Self {
child: Some(child),
stdin: Some(stdin),
stdout: Some(ready_stdout),
stderr_tail,
handshake: None,
next_id: 0,
shared_runtime: None,
agent_session_id: None,
})
}
pub fn shared_session(runtime: Arc<AgentRuntimeClient>, agent_session_id: String) -> Self {
Self {
child: None,
stdin: None,
stdout: None,
stderr_tail: runtime.stderr_tail.clone(),
handshake: Some(runtime.handshake().clone()),
next_id: 0,
shared_runtime: Some(runtime),
agent_session_id: Some(agent_session_id),
}
}
/// Send a JSON-RPC 2.0 request and wait for the response.
@@ -458,6 +734,20 @@ impl AgentDriverClient {
timeout_duration: Option<Duration>,
cancel_token: Option<CancellationToken>,
) -> Result<T, String> {
if let Some(runtime) = &self.shared_runtime {
let mut params = params;
if method != AgentMethod::Handshake.as_str()
&& method != AgentMethod::TestConnection.as_str()
&& method != AgentMethod::Shutdown.as_str()
{
let session_id = self.agent_session_id.as_ref().ok_or("Shared Agent session id is missing")?;
params
.as_object_mut()
.ok_or_else(|| "Agent RPC parameters must be an object".to_string())?
.insert("agentSessionId".to_string(), Value::String(session_id.clone()));
}
return runtime.call(method, params, timeout_duration, cancel_token).await;
}
self.next_id += 1;
let id = self.next_id;
@@ -587,6 +877,31 @@ impl AgentDriverClient {
self.call_method(AgentMethod::Connect, params).await
}
pub async fn open_session(&mut self, agent_session_id: &str, mut params: Value) -> Result<Value, String> {
params
.as_object_mut()
.ok_or_else(|| "Agent session parameters must be an object".to_string())?
.insert("agentSessionId".to_string(), Value::String(agent_session_id.to_string()));
self.call_method(AgentMethod::OpenSession, params).await
}
pub async fn close_session(&mut self, agent_session_id: &str) -> Result<Value, String> {
self.call_method(AgentMethod::CloseSession, serde_json::json!({ "agentSessionId": agent_session_id })).await
}
pub async fn validate_session(
&mut self,
agent_session_id: &str,
timeout_duration: Option<Duration>,
) -> Result<Value, String> {
self.call_method_with_timeout(
AgentMethod::ValidateSession,
serde_json::json!({ "agentSessionId": agent_session_id }),
timeout_duration,
)
.await
}
pub async fn test_connection(&mut self, params: Value) -> Result<Value, String> {
self.call_method(AgentMethod::TestConnection, params).await
}
@@ -596,6 +911,18 @@ impl AgentDriverClient {
}
pub async fn disconnect(&mut self) -> Result<Value, String> {
if self.shared_runtime.is_some() {
let session_id = self.agent_session_id.as_ref().ok_or("Shared Agent session id is missing")?.clone();
let result =
self.call_method(AgentMethod::CloseSession, serde_json::json!({ "agentSessionId": session_id })).await;
if result.is_ok() {
if let Some(runtime) = &self.shared_runtime {
AgentRuntimeClient::decrement_session_count(runtime);
}
self.agent_session_id = None;
}
return result;
}
self.call_method(AgentMethod::Disconnect, serde_json::json!({})).await
}
@@ -1105,6 +1432,10 @@ impl AgentDriverClient {
/// Send a shutdown message to the agent and wait for the process to exit.
pub async fn shutdown(&mut self) {
if self.shared_runtime.is_some() {
let _ = self.disconnect().await;
return;
}
// Try to send a shutdown RPC; ignore errors if the agent is already gone
let shutdown_result: Result<Value, String> = self.call_method(AgentMethod::Shutdown, Value::Null).await;
if let Err(e) = &shutdown_result {
@@ -1115,7 +1446,8 @@ impl AgentDriverClient {
self.stdin.take();
// Wait for the child to exit
match self.child.wait() {
let Some(child) = self.child.as_mut() else { return };
match child.wait() {
Ok(status) => log::info!("Agent process exited with {status}"),
Err(e) => log::warn!("Failed to wait for agent process: {e}"),
}
@@ -1123,9 +1455,14 @@ impl AgentDriverClient {
/// Forcefully kill the agent process.
pub fn kill(&mut self) {
if let Some(runtime) = &self.shared_runtime {
runtime.kill();
return;
}
self.stdin.take();
self.stdout.take();
if let Err(e) = self.child.kill() {
let Some(child) = self.child.as_mut() else { return };
if let Err(e) = child.kill() {
log::warn!("Failed to kill agent process: {e}");
}
// Reap the child to avoid zombie processes.
@@ -1133,7 +1470,7 @@ impl AgentDriverClient {
// hanging in Drop during async cleanup. Poll up to 100ms for the
// process to exit after kill().
for _ in 0..10 {
match self.child.try_wait() {
match child.try_wait() {
Ok(Some(_status)) => return,
Ok(None) => std::thread::sleep(Duration::from_millis(10)),
Err(e) => {
@@ -1143,13 +1480,28 @@ impl AgentDriverClient {
}
}
// Final blocking wait as a last resort
if let Err(e) = self.child.wait() {
if let Err(e) = child.wait() {
log::warn!("Final wait failed for agent process: {e}");
}
}
pub fn pid(&self) -> u32 {
self.child.id()
if let Some(runtime) = &self.shared_runtime {
return runtime.child.lock().map(|child| child.id()).unwrap_or_default();
}
self.child.as_ref().map(Child::id).unwrap_or_default()
}
pub fn protocol_mode(&self) -> &'static str {
if self.shared_runtime.is_some() {
"multi_session"
} else {
"legacy"
}
}
pub fn active_session_count(&self) -> u64 {
self.shared_runtime.as_ref().map(|runtime| runtime.active_session_count()).unwrap_or(1)
}
pub fn stderr_tail_snapshot(&self) -> String {
@@ -1438,13 +1790,22 @@ impl AgentDriverClient {
fn format_agent_process_error(&mut self, base: &str) -> String {
// Runtime RPC errors are common SQL/driver paths. Do not wait for the
// child to exit unless startup diagnostics already expect the process to die.
format_agent_process_error(base, child_exit_status(&mut self.child), &stderr_tail_snapshot(&self.stderr_tail))
let exit_status = self.child.as_mut().and_then(child_exit_status);
format_agent_process_error(base, exit_status, &stderr_tail_snapshot(&self.stderr_tail))
}
}
impl Drop for AgentDriverClient {
fn drop(&mut self) {
self.kill();
// Shared-session clients do not own the runtime process. Session closure
// is explicit in pool cleanup; dropping one tab must not kill other tabs.
if let Some(runtime) = &self.shared_runtime {
if self.agent_session_id.take().is_some() {
AgentRuntimeClient::decrement_session_count(runtime);
}
} else {
self.kill();
}
}
}
@@ -1456,12 +1817,16 @@ mod tests {
agent_schema_table_params, agent_supports_capability, agent_transaction_params, format_agent_process_error,
format_agent_startup_error, is_unsupported_handshake_error, mongo_collection_params, mongo_database_params,
mongo_document_id_params, parse_agent_java_opts, read_agent_line, start_stderr_collector, AgentCapability,
AgentDriverClient, AgentHandshake, AgentKvMethod, AgentMethod, AgentTableReadCloseParams,
AgentTableReadPageParams, AgentTableReadStartParams, MongoAgentMethod, StderrTail, AGENT_PROTOCOL_VERSION,
AgentDriverClient, AgentHandshake, AgentKvMethod, AgentLaunchSpec, AgentMethod, AgentRuntimeClient,
AgentTableReadCloseParams, AgentTableReadPageParams, AgentTableReadStartParams, MongoAgentMethod, StderrTail,
AGENT_PROTOCOL_VERSION,
};
use std::io::Cursor;
use std::io::Write;
use std::process::{Command, Stdio};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio_util::sync::CancellationToken;
#[test]
fn agent_java_args_include_oracle_network_compatibility_flags() {
@@ -1638,12 +2003,14 @@ mod tests {
.spawn()
.expect("child should start");
let mut client = AgentDriverClient {
child,
child: Some(child),
stdin: None,
stdout: None,
stderr_tail: Arc::new(Mutex::new(StderrTail::default())),
handshake: None,
next_id: 0,
shared_runtime: None,
agent_session_id: None,
};
let started_at = std::time::Instant::now();
@@ -1654,6 +2021,129 @@ mod tests {
assert!(!message.contains("agent process exited"));
}
#[tokio::test]
async fn multiplexed_runtime_correlates_out_of_order_responses() {
let script_path = std::env::temp_dir().join(format!("dbx-agent-runtime-test-{}.py", uuid::Uuid::new_v4()));
let mut script = std::fs::File::create(&script_path).unwrap();
script
.write_all(
br#"import json, sys, threading, time
print(json.dumps({'ready': True}), flush=True)
def respond(req):
if req['method'] == 'handshake':
result = {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
else:
time.sleep(0.05 if req['params']['value'] == 1 else 0.0)
result = req['params']['value']
print(json.dumps({'jsonrpc': '2.0', 'id': req['id'], 'result': result}), flush=True)
for line in sys.stdin:
threading.Thread(target=respond, args=(json.loads(line),), daemon=True).start()
"#,
)
.unwrap();
drop(script);
let runtime = AgentRuntimeClient::spawn(
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
"test",
)
.await
.unwrap();
let first = runtime.call::<u64>("echo", serde_json::json!({"value": 1}), Some(Duration::from_secs(2)), None);
let second = runtime.call::<u64>("echo", serde_json::json!({"value": 2}), Some(Duration::from_secs(2)), None);
let (first, second) = tokio::join!(first, second);
assert_eq!(first.unwrap(), 1);
assert_eq!(second.unwrap(), 2);
runtime.kill();
let _ = std::fs::remove_file(script_path);
}
#[tokio::test]
async fn canceling_one_runtime_request_keeps_other_requests_alive() {
let script_path = std::env::temp_dir().join(format!("dbx-agent-cancel-test-{}.py", uuid::Uuid::new_v4()));
std::fs::write(
&script_path,
r#"import json, sys, threading, time
print(json.dumps({'ready': True}), flush=True)
def respond(req):
if req['method'] == 'handshake':
result = {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
else:
time.sleep(req['params'].get('delay', 0))
result = req['params'].get('value')
print(json.dumps({'jsonrpc': '2.0', 'id': req['id'], 'result': result}), flush=True)
for line in sys.stdin:
threading.Thread(target=respond, args=(json.loads(line),), daemon=True).start()
"#,
)
.unwrap();
let runtime = AgentRuntimeClient::spawn(
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
"test",
)
.await
.unwrap();
let token = CancellationToken::new();
let canceled = runtime.call::<u64>(
"echo",
serde_json::json!({"value": 1, "delay": 0.2}),
Some(Duration::from_secs(2)),
Some(token.clone()),
);
let healthy = runtime.call::<u64>(
"echo",
serde_json::json!({"value": 2, "delay": 0}),
Some(Duration::from_secs(2)),
None,
);
token.cancel();
let (canceled, healthy) = tokio::join!(canceled, healthy);
assert_eq!(canceled.unwrap_err(), "Query canceled");
assert_eq!(healthy.unwrap(), 2);
assert_eq!(
runtime
.call::<u64>("echo", serde_json::json!({"value": 3}), Some(Duration::from_secs(2)), None)
.await
.unwrap(),
3
);
runtime.kill();
let _ = std::fs::remove_file(script_path);
}
#[tokio::test]
async fn runtime_eof_fails_pending_requests_and_marks_runtime_failed() {
let script_path = std::env::temp_dir().join(format!("dbx-agent-eof-test-{}.py", uuid::Uuid::new_v4()));
std::fs::write(
&script_path,
r#"import json, sys
print(json.dumps({'ready': True}), flush=True)
for line in sys.stdin:
req = json.loads(line)
if req['method'] == 'handshake':
print(json.dumps({'jsonrpc':'2.0','id':req['id'],'result':{'protocolVersion':2,'agentProtocolVersion':2,'capabilities':['multi_session']}}), flush=True)
else:
sys.exit(0)
"#,
)
.unwrap();
let runtime = AgentRuntimeClient::spawn(
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
"test",
)
.await
.unwrap();
let error = runtime
.call::<serde_json::Value>("crash", serde_json::json!({}), Some(Duration::from_secs(2)), None)
.await
.unwrap_err();
assert!(error.contains("end of stream") || error.contains("response channel closed"));
assert!(runtime.is_failed());
let _ = std::fs::remove_file(script_path);
}
#[test]
fn stderr_tail_keeps_recent_lines_only() {
let mut stderr_tail = StderrTail::with_capacity(3);
@@ -1697,7 +2187,8 @@ mod tests {
assert_eq!(AgentCapability::Transaction.as_str(), "transaction");
assert_eq!(AgentCapability::Ddl.as_str(), "ddl");
assert_eq!(AgentCapability::Kv.as_str(), "kv");
assert_eq!(AgentCapability::ALL.len(), 8);
assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session");
assert_eq!(AgentCapability::ALL.len(), 9);
}
#[test]
@@ -1915,7 +2406,7 @@ mod tests {
#[test]
fn agent_protocol_matches_contract_file() {
let contract: serde_json::Value =
serde_json::from_str(include_str!("../../assets/agent-protocol-v1.json")).unwrap();
serde_json::from_str(include_str!("../../assets/agent-protocol-v2.json")).unwrap();
assert_eq!(contract["protocolVersion"], AGENT_PROTOCOL_VERSION);
assert_eq!(contract["handshakeMethod"], AgentMethod::Handshake.as_str());
@@ -1991,6 +2482,7 @@ mod tests {
AgentCapability::PagedQuery,
AgentCapability::Transaction,
AgentCapability::Ddl,
AgentCapability::MultiSession,
]
.iter()
.map(|capability| capability.as_str())
crates/dbx-core/src/driver_runtime.rs
+56 -12
View File
@@ -23,6 +23,10 @@ pub struct DriverRuntimeInfo {
pub can_stop: bool,
pub can_restart: bool,
pub control_unavailable_reason: Option<String>,
#[serde(default)]
pub protocol_mode: Option<String>,
#[serde(default)]
pub active_sessions: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
@@ -48,6 +52,8 @@ struct RuntimeSeed {
can_stop: bool,
can_restart: bool,
control_unavailable_reason: Option<String>,
protocol_mode: Option<String>,
active_sessions: Option<u64>,
}
#[derive(Debug, Clone, Default)]
@@ -65,16 +71,27 @@ pub async fn collect_driver_runtime_summary(state: &AppState) -> DriverRuntimeSu
}
pub async fn stop_driver_runtime(state: &AppState, runtime_id: &str) -> Result<(), String> {
let agent_key =
runtime_id.strip_prefix("agent:").ok_or_else(|| "Only agent daemon runtimes can be stopped".to_string())?;
let agent_key = runtime_agent_key(runtime_id).ok_or_else(|| "Only agent runtimes can be stopped".to_string())?;
state.agent_manager.stop_daemon_by_key(agent_key).await;
Ok(())
}
pub async fn restart_driver_runtime(state: &AppState, runtime_id: &str) -> Result<(), String> {
let agent_key =
runtime_id.strip_prefix("agent:").ok_or_else(|| "Only agent daemon runtimes can be restarted".to_string())?;
state.agent_manager.restart_daemon_by_key(agent_key).await
let agent_key = runtime_agent_key(runtime_id).ok_or_else(|| "Only agent runtimes can be restarted".to_string())?;
if runtime_id.starts_with("agent-runtime:") {
// Shared connection runtimes are recreated lazily with their exact launch
// fingerprint; eagerly spawning the generic daemon could use incompatible options.
state.agent_manager.stop_daemon_by_key(agent_key).await;
Ok(())
} else {
state.agent_manager.restart_daemon_by_key(agent_key).await
}
}
fn runtime_agent_key(runtime_id: &str) -> Option<&str> {
runtime_id.strip_prefix("agent:").or_else(|| {
runtime_id.strip_prefix("agent-runtime:").and_then(|value| value.rsplit_once(':').map(|(key, _)| key))
})
}
async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
@@ -97,6 +114,8 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
can_stop: false,
can_restart: state.agent_manager.is_driver_installed(key),
control_unavailable_reason: None,
protocol_mode: None,
active_sessions: None,
},
);
}
@@ -120,6 +139,8 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
can_stop: true,
can_restart: true,
control_unavailable_reason: None,
protocol_mode: Some("legacy".to_string()),
active_sessions: None,
},
);
}
@@ -140,21 +161,31 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
.unwrap_or("agent");
let client = client.lock().await;
let last_error = non_empty(client.stderr_tail_snapshot());
let protocol_mode = client.protocol_mode();
let runtime_id = if protocol_mode == "multi_session" {
format!("agent-runtime:{key}:{}", client.pid())
} else {
format!("agent-connection:{pool_key}")
};
seeds.insert(
format!("agent-connection:{pool_key}"),
runtime_id.clone(),
RuntimeSeed {
id: format!("agent-connection:{pool_key}"),
id: runtime_id,
driver_key: key.to_string(),
label: agent_catalog::label_for_key(key).unwrap_or(key).to_string(),
kind: "agent".to_string(),
source: "connection".to_string(),
source: if protocol_mode == "multi_session" { "shared-runtime" } else { "connection" }
.to_string(),
status: "running".to_string(),
pid: Some(client.pid()),
version: agent_state.installed_drivers.get(key).map(|driver| driver.version.clone()),
last_error,
can_stop: false,
can_restart: false,
control_unavailable_reason: Some("connection-owned".to_string()),
can_stop: protocol_mode == "multi_session",
can_restart: protocol_mode == "multi_session",
control_unavailable_reason: (protocol_mode != "multi_session")
.then(|| "connection-owned".to_string()),
protocol_mode: Some(protocol_mode.to_string()),
active_sessions: Some(client.active_session_count()),
},
);
}
@@ -174,6 +205,8 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
can_stop: false,
can_restart: false,
control_unavailable_reason: Some("connection-owned".to_string()),
protocol_mode: None,
active_sessions: None,
},
);
}
@@ -248,6 +281,8 @@ fn build_summary(seeds: Vec<RuntimeSeed>, stats: &HashMap<u32, ProcessStats>) ->
can_stop: seed.can_stop,
can_restart: seed.can_restart,
control_unavailable_reason: seed.control_unavailable_reason,
protocol_mode: seed.protocol_mode,
active_sessions: seed.active_sessions,
}
})
.collect::<Vec<_>>();
@@ -274,7 +309,7 @@ fn non_empty(value: String) -> Option<String> {
#[cfg(test)]
mod tests {
use super::{build_summary, ProcessStats, RuntimeSeed};
use super::{build_summary, runtime_agent_key, ProcessStats, RuntimeSeed};
use std::collections::HashMap;
fn seed(id: &str, status: &str, pid: Option<u32>) -> RuntimeSeed {
@@ -291,6 +326,8 @@ mod tests {
can_stop: true,
can_restart: true,
control_unavailable_reason: None,
protocol_mode: None,
active_sessions: None,
}
}
@@ -306,4 +343,11 @@ mod tests {
assert_eq!(summary.health, "healthy");
assert_eq!(summary.runtimes[0].memory_bytes, Some(1024));
}
#[test]
fn parses_legacy_and_shared_agent_runtime_ids() {
assert_eq!(runtime_agent_key("agent:oracle"), Some("oracle"));
assert_eq!(runtime_agent_key("agent-runtime:oracle:123"), Some("oracle"));
assert_eq!(runtime_agent_key("plugin-connection:test"), None);
}
}