test(desktop): Phase 4 drain journal unit tests and execute_drain_journal extraction

Extract execute_drain_journal() as a pure inner function that takes the
runtime HashMap directly — no AppHandle needed — enabling deterministic
unit testing of the drain/compensate logic without a Tauri mock app.

Move the test block from runtime_commands.rs to the sibling
runtime_commands_tests.rs (following the storage_tests.rs pattern) to
keep the main file under the 1000-line size gate.

New tests:
- test_drain_empty_map_returns_success
- test_drain_exited_process_counts_as_stopped_and_clears_map
- test_drain_scope_id_propagates_from_runtime_starting
- test_drain_missing_key_treated_as_already_stopped
- test_drain_cleanup_fn_called_for_each_stopped_entry
- test_workspace_apply_result_drain_failed_returns_applied_false
- test_workspace_apply_result_degradation_accumulates

All 2125 existing tests continue to pass.

Co-authored-by: Will Pfleger <pfleger.will@gmail.com>
Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
This commit is contained in:
npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7
2026-08-03 00:25:53 -04:00
co-authored by Will Pfleger
parent 6eec536799
commit 8a35f5ed2b
2 changed files with 351 additions and 198 deletions
@@ -1,3 +1,4 @@
use std::collections::HashMap;
use std::sync::atomic::Ordering;
use tauri::{AppHandle, Emitter, Manager};
@@ -599,6 +600,69 @@ pub(crate) struct DrainJournalEntry {
pub start_on_app_launch: bool,
}
/// Execute a drain journal against the runtime map.
///
/// Pure inner function: takes the map directly so callers (including tests)
/// can drive it without an `AppHandle`. The `cleanup_fn` is called for each
/// successfully stopped entry to remove its receipt and clear the session
/// cache; the closure is a no-op in tests.
///
/// Returns `(stopped, remaining, first_stop_error)`:
/// - `stopped` — entries successfully killed (compensation restores these).
/// - `remaining` — entries NOT attempted due to an earlier stop failure.
/// - error — the first stop failure, if any; `None` on full success.
pub(crate) fn execute_drain_journal(
journal: &[DrainJournalEntry],
runtimes: &mut HashMap<ManagedAgentRuntimeKey, ManagedAgentPairRuntime>,
mut cleanup_fn: impl FnMut(&ManagedAgentRuntimeKey),
) -> (
Vec<DrainJournalEntry>,
Vec<DrainJournalEntry>,
Option<String>,
) {
let mut stopped: Vec<DrainJournalEntry> = Vec::new();
let mut first_error: Option<String> = None;
for (idx, entry) in journal.iter().enumerate() {
let key = &entry.key;
let stop_result = if let Some(mut runtime) = runtimes.remove(key) {
let kill_result = if super::process_is_running(runtime.child.id()) {
super::terminate_process(runtime.child.id())
} else {
Ok(())
}
.and_then(|()| runtime.child.wait().map_err(|e| e.to_string()));
match kill_result {
Ok(_) => {
cleanup_fn(key);
Ok(())
}
Err(e) => {
// Put it back so the map is consistent.
runtimes.insert(key.clone(), runtime);
Err(e)
}
}
} else {
// Nothing live at this key — treat as already stopped.
Ok(())
};
match stop_result {
Ok(()) => stopped.push(entry.clone()),
Err(e) => {
let msg = format!("failed to stop agent {}@{}: {e}", key.pubkey, key.relay_url);
first_error.get_or_insert(msg);
// Return the un-attempted tail (idx+1 onward) as remaining.
return (stopped, journal[idx + 1..].to_vec(), first_error);
}
}
}
(stopped, vec![], first_error)
}
/// Drain all live runtimes from the runtime map and return a drain journal
/// (keys + restart recipes) for use by compensation.
///
@@ -653,61 +717,21 @@ pub(crate) fn drain_scope_runtimes(
.collect()
};
let mut stopped: Vec<DrainJournalEntry> = Vec::new();
let mut first_error: Option<String> = None;
for (idx, entry) in journal.iter().enumerate() {
let key = &entry.key;
let stop_result = {
let mut runtimes = match state.managed_agent_processes.lock() {
Ok(r) => r,
Err(e) => {
first_error.get_or_insert_with(|| {
format!("runtime map lock poisoned during drain: {e}")
});
// Return remaining as the un-attempted tail.
return (stopped, journal[idx..].to_vec(), first_error);
}
};
if let Some(mut runtime) = runtimes.remove(key) {
let kill_result = if super::process_is_running(runtime.child.id()) {
super::terminate_process(runtime.child.id())
} else {
Ok(())
}
.and_then(|()| runtime.child.wait().map_err(|e| e.to_string()));
match kill_result {
Ok(_) => {
// Remove the receipt so sweep/restore see a clean slate.
super::remove_agent_runtime_receipt(app, key);
state.clear_agent_session_cache(key);
Ok(())
}
Err(e) => {
// Put it back so the map is consistent.
runtimes.insert(key.clone(), runtime);
Err(e)
}
}
} else {
// Nothing live at this key — treat as already stopped.
Ok(())
}
};
match stop_result {
Ok(()) => stopped.push(entry.clone()),
Err(e) => {
let msg = format!("failed to stop agent {}@{}: {e}", key.pubkey, key.relay_url);
first_error.get_or_insert(msg);
// Return the un-attempted tail (idx+1 onward) as remaining.
return (stopped, journal[idx + 1..].to_vec(), first_error);
}
let mut runtimes = match state.managed_agent_processes.lock() {
Ok(r) => r,
Err(e) => {
return (
vec![],
journal,
Some(format!("runtime map lock poisoned during drain: {e}")),
)
}
}
};
(stopped, vec![], first_error)
execute_drain_journal(&journal, &mut runtimes, |key| {
super::remove_agent_runtime_receipt(app, key);
state.clear_agent_session_cache(key);
})
}
/// Compensate a partial drain by restarting the entries that were successfully
@@ -743,148 +767,5 @@ pub(crate) fn compensate_drain(app: &AppHandle, stopped: &[DrainJournalEntry]) -
}
#[cfg(test)]
mod tests {
use super::*;
fn payload(
relay_url: &str,
lifecycle: ManagedAgentRuntimeLifecycle,
error: Option<&str>,
) -> super::super::ManagedAgentRuntimeLifecycleObserverPayload {
super::super::ManagedAgentRuntimeLifecycleObserverPayload {
pubkey: "aa".repeat(32),
relay_url: relay_url.into(),
start_nonce: "test-generation".into(),
lifecycle,
error: error.map(str::to_owned),
}
}
fn record_with_relay(relay_url: &str) -> super::super::ManagedAgentRecord {
serde_json::from_str(&format!(
r#"{{
"pubkey": "{}",
"name": "pin-test",
"relay_url": "{relay_url}",
"acp_command": "buzz-acp",
"agent_command": "goose",
"agent_args": [],
"mcp_command": "",
"turn_timeout_seconds": 320,
"system_prompt": "",
"created_at": "2026-01-01T00:00:00Z",
"updated_at": "2026-01-01T00:00:00Z"
}}"#,
"aa".repeat(32)
))
.unwrap()
}
#[test]
fn legacy_relay_pin_is_ignored_for_fan_out() {
// Zero-touch cutover (#2122): a record carrying a creation-era
// `relay_url` pin must fan out exactly like an unpinned one — the
// stored field is parsed but never consulted. See
// `effective_agent_relay_url`.
let unpinned = record_with_relay("");
let pinned = record_with_relay("wss://one.example");
for record in [&unpinned, &pinned] {
assert_eq!(
crate::relay::effective_agent_relay_url(&record.relay_url, "wss://two.example"),
"wss://two.example"
);
}
}
#[test]
fn unkeyable_relay_degrades_to_failed_row() {
// A requested URL that cannot form a pair key must still yield a
// Failed row keyed by the raw requested string, so one bad community
// never aborts the rest of the reconcile batch.
let record = record_with_relay("");
let status = unkeyable_failed_status(
&record,
"not a url".to_string(),
"relay access probe timed out".to_string(),
&[],
&super::super::GlobalAgentConfig::default(),
);
assert!(matches!(
status.lifecycle,
ManagedAgentRuntimeLifecycle::Failed
));
assert_eq!(status.relay_url, "not a url");
assert_eq!(status.requested_relay_url.as_deref(), Some("not a url"));
assert_eq!(status.pubkey, record.pubkey);
assert_eq!(
status.error.as_deref(),
Some("relay access probe timed out")
);
assert!(status.pid.is_none());
}
#[test]
fn runtime_key_rejects_non_hex_pubkeys() {
assert!(ManagedAgentRuntimeKey::new("../not-a-key", "wss://relay.example").is_err());
assert!(ManagedAgentRuntimeKey::new("gg".repeat(32), "wss://relay.example").is_err());
}
#[test]
fn runtime_key_canonicalizes_hex_pubkeys() {
let key = ManagedAgentRuntimeKey::new("AA".repeat(32), "wss://relay.example").unwrap();
assert_eq!(key.pubkey, "aa".repeat(32));
}
#[test]
fn observer_lifecycle_key_preserves_exact_canonical_pair() {
let first = payload(
"WSS://Relay.Example:443/",
ManagedAgentRuntimeLifecycle::Ready,
None,
);
let key = observer_lifecycle_key(&first.pubkey, &first).unwrap();
assert_eq!(key.pubkey, first.pubkey);
assert_eq!(key.relay_url, "wss://relay.example");
let other = payload(
"wss://other.example",
ManagedAgentRuntimeLifecycle::Ready,
None,
);
assert_ne!(key, observer_lifecycle_key(&other.pubkey, &other).unwrap());
}
#[test]
fn observer_lifecycle_rejects_cross_agent_and_desktop_states() {
let ready = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Ready,
None,
);
assert!(observer_lifecycle_key(&"bb".repeat(32), &ready).is_err());
let stopped = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Stopped,
None,
);
assert!(observer_lifecycle_key(&stopped.pubkey, &stopped).is_err());
}
#[test]
fn observer_lifecycle_enforces_failed_error_contract() {
let failed = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Failed,
None,
);
assert!(observer_lifecycle_key(&failed.pubkey, &failed).is_err());
let ready_with_error = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Ready,
Some("unexpected"),
);
assert!(observer_lifecycle_key(&ready_with_error.pubkey, &ready_with_error).is_err());
}
}
#[path = "runtime_commands_tests.rs"]
mod tests;
@@ -0,0 +1,272 @@
//! Unit tests for `managed_agents/runtime_commands.rs`.
//!
//! Kept in a sibling file so `runtime_commands.rs` stays under the
//! 1000-line size gate; `#[path]`-included from there.
use super::*;
fn payload(
relay_url: &str,
lifecycle: ManagedAgentRuntimeLifecycle,
error: Option<&str>,
) -> super::super::ManagedAgentRuntimeLifecycleObserverPayload {
super::super::ManagedAgentRuntimeLifecycleObserverPayload {
pubkey: "aa".repeat(32),
relay_url: relay_url.into(),
start_nonce: "test-generation".into(),
lifecycle,
error: error.map(str::to_owned),
}
}
fn record_with_relay(relay_url: &str) -> super::super::ManagedAgentRecord {
serde_json::from_str(&format!(
r#"{{
"pubkey": "{}",
"name": "pin-test",
"relay_url": "{relay_url}",
"acp_command": "buzz-acp",
"agent_command": "goose",
"agent_args": [],
"mcp_command": "",
"turn_timeout_seconds": 320,
"system_prompt": "",
"created_at": "2026-01-01T00:00:00Z",
"updated_at": "2026-01-01T00:00:00Z"
}}"#,
"aa".repeat(32)
))
.unwrap()
}
#[test]
fn legacy_relay_pin_is_ignored_for_fan_out() {
// Zero-touch cutover (#2122): a record carrying a creation-era
// `relay_url` pin must fan out exactly like an unpinned one — the
// stored field is parsed but never consulted. See
// `effective_agent_relay_url`.
let unpinned = record_with_relay("");
let pinned = record_with_relay("wss://one.example");
for record in [&unpinned, &pinned] {
assert_eq!(
crate::relay::effective_agent_relay_url(&record.relay_url, "wss://two.example"),
"wss://two.example"
);
}
}
#[test]
fn unkeyable_relay_degrades_to_failed_row() {
// A requested URL that cannot form a pair key must still yield a
// Failed row keyed by the raw requested string, so one bad community
// never aborts the rest of the reconcile batch.
let record = record_with_relay("");
let status = unkeyable_failed_status(
&record,
"not a url".to_string(),
"relay access probe timed out".to_string(),
&[],
&super::super::GlobalAgentConfig::default(),
);
assert!(matches!(
status.lifecycle,
ManagedAgentRuntimeLifecycle::Failed
));
assert_eq!(status.relay_url, "not a url");
assert_eq!(status.requested_relay_url.as_deref(), Some("not a url"));
assert_eq!(status.pubkey, record.pubkey);
assert_eq!(
status.error.as_deref(),
Some("relay access probe timed out")
);
assert!(status.pid.is_none());
}
#[test]
fn runtime_key_rejects_non_hex_pubkeys() {
assert!(ManagedAgentRuntimeKey::new("../not-a-key", "wss://relay.example").is_err());
assert!(ManagedAgentRuntimeKey::new("gg".repeat(32), "wss://relay.example").is_err());
}
#[test]
fn runtime_key_canonicalizes_hex_pubkeys() {
let key = ManagedAgentRuntimeKey::new("AA".repeat(32), "wss://relay.example").unwrap();
assert_eq!(key.pubkey, "aa".repeat(32));
}
#[test]
fn observer_lifecycle_key_preserves_exact_canonical_pair() {
let first = payload(
"WSS://Relay.Example:443/",
ManagedAgentRuntimeLifecycle::Ready,
None,
);
let key = observer_lifecycle_key(&first.pubkey, &first).unwrap();
assert_eq!(key.pubkey, first.pubkey);
assert_eq!(key.relay_url, "wss://relay.example");
let other = payload(
"wss://other.example",
ManagedAgentRuntimeLifecycle::Ready,
None,
);
assert_ne!(key, observer_lifecycle_key(&other.pubkey, &other).unwrap());
}
#[test]
fn observer_lifecycle_rejects_cross_agent_and_desktop_states() {
let ready = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Ready,
None,
);
assert!(observer_lifecycle_key(&"bb".repeat(32), &ready).is_err());
let stopped = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Stopped,
None,
);
assert!(observer_lifecycle_key(&stopped.pubkey, &stopped).is_err());
}
#[test]
fn observer_lifecycle_enforces_failed_error_contract() {
let failed = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Failed,
None,
);
assert!(observer_lifecycle_key(&failed.pubkey, &failed).is_err());
let ready_with_error = payload(
"wss://relay.example",
ManagedAgentRuntimeLifecycle::Ready,
Some("unexpected"),
);
assert!(observer_lifecycle_key(&ready_with_error.pubkey, &ready_with_error).is_err());
}
// ── drain journal / WorkspaceApplyResult tests ───────────────────────────
fn make_drain_entry(pubkey_hex: &str, relay: &str, auto: bool) -> DrainJournalEntry {
DrainJournalEntry {
key: ManagedAgentRuntimeKey::new(pubkey_hex, relay).unwrap(),
start_on_app_launch: auto,
}
}
fn make_exited_pair_runtime(scope_id: Option<String>) -> ManagedAgentPairRuntime {
use std::process::{Command, Stdio};
#[cfg(unix)]
let program = "/usr/bin/true";
#[cfg(windows)]
let program = "true";
let child = Command::new(program)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.expect("spawn /usr/bin/true");
let process = super::super::ManagedAgentProcess {
child,
log_path: std::path::PathBuf::new(),
spawn_config_hash: 0,
setup_mode: false,
adapter_availability: None,
start_nonce: "test-nonce".to_string(),
#[cfg(windows)]
job: None,
};
ManagedAgentPairRuntime::starting(process, scope_id)
}
#[test]
fn test_drain_empty_map_returns_success() {
let journal: Vec<DrainJournalEntry> = vec![];
let (stopped, remaining, err) = execute_drain_journal(&journal, &mut HashMap::new(), |_| {});
assert!(stopped.is_empty());
assert!(remaining.is_empty());
assert!(err.is_none());
}
#[test]
fn test_drain_exited_process_counts_as_stopped_and_clears_map() {
// `true` exits immediately with 0 — process_is_running returns false
// after a brief moment, so drain treats it as already-stopped and
// calls wait() to reap it.
let pubkey = "aa".repeat(32);
let key = ManagedAgentRuntimeKey::new(&pubkey, "wss://relay.example").unwrap();
let runtime = make_exited_pair_runtime(None);
// Give the process a moment to exit before drain tries to stop it.
std::thread::sleep(std::time::Duration::from_millis(50));
let entry = make_drain_entry(&pubkey, "wss://relay.example", true);
let mut map = HashMap::from([(key, runtime)]);
let (stopped, remaining, err) = execute_drain_journal(&[entry], &mut map, |_| {});
assert_eq!(stopped.len(), 1, "exited process must appear in stopped");
assert!(remaining.is_empty());
assert!(err.is_none());
assert!(map.is_empty(), "entry must be removed from the runtime map");
}
#[test]
fn test_drain_scope_id_propagates_from_runtime_starting() {
let scope_id = Some("test-scope-abc".to_string());
let runtime = make_exited_pair_runtime(scope_id.clone());
assert_eq!(
runtime.scope_id, scope_id,
"scope_id must be preserved through ManagedAgentPairRuntime::starting()"
);
}
#[test]
fn test_drain_missing_key_treated_as_already_stopped() {
// A key in the journal but absent from the map is treated as
// already stopped: it still appears in `stopped` so compensation
// would attempt a restart (safe-but-redundant, not silent loss).
let pubkey = "bb".repeat(32);
let entry = make_drain_entry(&pubkey, "wss://relay.example", false);
let mut map: HashMap<ManagedAgentRuntimeKey, ManagedAgentPairRuntime> = HashMap::new();
let (stopped, remaining, err) = execute_drain_journal(&[entry], &mut map, |_| {});
assert_eq!(stopped.len(), 1);
assert!(remaining.is_empty());
assert!(err.is_none());
}
#[test]
fn test_drain_cleanup_fn_called_for_each_stopped_entry() {
let pubkey = "cc".repeat(32);
let key = ManagedAgentRuntimeKey::new(&pubkey, "wss://relay.example").unwrap();
let runtime = make_exited_pair_runtime(None);
std::thread::sleep(std::time::Duration::from_millis(50));
let entry = make_drain_entry(&pubkey, "wss://relay.example", false);
let mut map = HashMap::from([(key.clone(), runtime)]);
let mut cleaned: Vec<ManagedAgentRuntimeKey> = Vec::new();
execute_drain_journal(&[entry], &mut map, |k| cleaned.push(k.clone()));
assert_eq!(
cleaned,
vec![key],
"cleanup_fn must be called once per stopped entry"
);
}
#[test]
fn test_workspace_apply_result_drain_failed_returns_applied_false() {
let r = super::super::scope::WorkspaceApplyResult::drain_failed("stop failed");
assert!(!r.applied);
assert_eq!(r.degraded, vec!["stop failed"]);
}
#[test]
fn test_workspace_apply_result_degradation_accumulates() {
let r = super::super::scope::WorkspaceApplyResult::success()
.with_degradation("nest failed")
.with_degradation("sync skipped");
assert!(
r.applied,
"degraded workspace must still report applied: true"
);
assert_eq!(r.degraded.len(), 2);
assert!(r.degraded[0].contains("nest"));
assert!(r.degraded[1].contains("sync"));
}