diff --git a/desktop/src-tauri/src/managed_agents/runtime_commands.rs b/desktop/src-tauri/src/managed_agents/runtime_commands.rs index 9d1140e8b..74d96203f 100644 --- a/desktop/src-tauri/src/managed_agents/runtime_commands.rs +++ b/desktop/src-tauri/src/managed_agents/runtime_commands.rs @@ -820,9 +820,6 @@ pub(crate) fn drain_scope_runtimes( /// old workspace is as intact as possible. /// /// `captured_scope` is the workspace scope that was active when the drain began. -/// Before restoring any entry this function validates that the scope generation -/// has not changed — if a workspace switch raced the drain failure, compensation -/// is skipped entirely (the new scope will restore its own agents). /// /// `_rt_transition_held` is the caller's already-held /// `managed_agent_runtime_transition` guard. The caller must NOT drop it before @@ -830,66 +827,46 @@ pub(crate) fn drain_scope_runtimes( /// is held continuously from drain through all journal restarts, closing the /// drop-then-reacquire interleave window that a concurrent start could exploit. /// -/// This function additionally re-acquires `managed_agents_store_lock` so all -/// journal restarts run under both guards. The caller must have released -/// `managed_agents_store_lock` (only) before calling. -/// /// Returns a degradation message describing what could not be restarted. /// /// The journal-restore loop is implemented in [`compensate_drain_for`] with an -/// injected start function, allowing the generation-validation + iteration -/// contract to be unit-tested without an `AppHandle`. +/// injected start function, allowing the iteration contract to be unit-tested +/// without an `AppHandle`. // ──────────────────────────────────────────────────────────────────────────── -/// Testable core of [`compensate_drain`]: re-acquires the store lock, validates -/// the captured scope generation, then calls `start_fn` for each stopped entry. +/// Lock-free testable core of the journal-restore loop. /// -/// The transition lock must already be held by the caller. `compensate_drain` -/// passes a `start_fn` that invokes [`start_pair_under_held_locks`]; tests inject -/// a closure that records calls and returns synthetic results without spawning -/// processes or touching disk. +/// **Preconditions (enforced by the [`compensate_drain`] adapter before calling):** +/// - `managed_agent_runtime_transition` is held by the caller (passed by value +/// to `compensate_drain`). +/// - `managed_agents_store_lock` is acquired by the adapter BEFORE this call. +/// - `records` is loaded by the adapter AFTER acquiring the store lock. +/// +/// `compensate_drain` passes a `start_fn` that invokes +/// [`start_pair_under_held_locks`]; tests inject a closure that records calls +/// and returns synthetic results without spawning processes or touching disk. +/// +/// The mechanism serializing writers: the adapter holds `managed_agents_store_lock` +/// continuously across validate→load→restore→save, so any writer that takes only +/// the store lock is serialized here — not on the transition guard. /// /// Returns a degradation message when one or more restarts fail, `None` on full -/// success. A stale scope also returns a degradation message (compensation -/// skipped) — that is not a hard failure, so callers treat it the same way. +/// success. pub(crate) fn compensate_drain_for( - state: &AppState, stopped: &[DrainJournalEntry], - captured_scope: &crate::managed_agents::scope::WorkspaceAgentScope, + records: &mut [super::ManagedAgentRecord], mut start_fn: F, ) -> Option where - F: FnMut(&DrainJournalEntry) -> Result<(), String>, + F: FnMut(&DrainJournalEntry, &mut [super::ManagedAgentRecord]) -> Result<(), String>, { debug_assert!( !stopped.is_empty(), "compensate_drain_for called with empty stopped list" ); - // Re-acquire only the store lock — the transition lock is already held via - // the caller's guard so no concurrent start can enter the window. - let _store = match state.managed_agents_store_lock.lock() { - Ok(g) => g, - Err(e) => { - return Some(format!( - "compensation failed: could not acquire store lock: {e}" - )); - } - }; - - // Validate the captured scope before restoring. If the workspace switched - // between the drain failure and this compensation, the new scope's own - // restore pass will start the correct agents — we must not restart agents - // for a scope that is no longer active. - if let Err(stale_msg) = crate::managed_agents::scope::validate_scope_generation(captured_scope) - { - return Some(format!( - "compensation skipped: {stale_msg}; new scope will restore its own agents" - )); - } - let mut failed_restarts = Vec::new(); for entry in stopped { - if let Err(e) = start_fn(entry) { + if let Err(e) = start_fn(entry, records) { failed_restarts.push(format!("{}@{}: {e}", entry.key.pubkey, entry.key.relay_url)); } } @@ -904,6 +881,17 @@ where } } +/// Production adapter for [`compensate_drain_for`]. +/// +/// Lock acquisition order inside this function: +/// 1. `_rt_transition_held` — already held by caller (passed by value). +/// 2. Acquire `managed_agents_store_lock` — store-lock-only writers +/// (instance edits, status flushes) are serialized here, not on the +/// transition guard. +/// 3. Validate `captured_scope.generation` under the store lock. +/// 4. `load_managed_agents_at` — records loaded AFTER both locks held. +/// 5. Delegate to [`compensate_drain_for`] — store guard held through +/// every save performed by `start_pair_under_held_locks`. pub(crate) fn compensate_drain( app: &tauri::AppHandle, stopped: &[DrainJournalEntry], @@ -920,6 +908,28 @@ pub(crate) fn compensate_drain( let state = app.state::(); + // 2. Acquire the store lock — transition already held via _rt_transition_held. + let _store = match state.managed_agents_store_lock.lock() { + Ok(g) => g, + Err(e) => { + return Some(format!( + "compensation failed: could not acquire store lock: {e}" + )); + } + }; + + // 3. Validate the captured scope before restoring. If the workspace switched + // between the drain failure and this compensation, the new scope's own + // restore pass will start the correct agents — we must not restart agents + // for a scope that is no longer active. + if let Err(stale_msg) = crate::managed_agents::scope::validate_scope_generation(captured_scope) + { + return Some(format!( + "compensation skipped: {stale_msg}; new scope will restore its own agents" + )); + } + + // 4. Load records under the held store lock. let mut records = match load_managed_agents_at(&captured_scope.definitions_dir) { Ok(r) => r, Err(e) => { @@ -929,7 +939,8 @@ pub(crate) fn compensate_drain( } }; - compensate_drain_for(&state, stopped, captured_scope, |entry| { + // 5. Delegate to the lock-free core — store guard is held through every save. + compensate_drain_for(stopped, &mut records, |entry, recs| { start_pair_under_held_locks( app, &state, @@ -937,7 +948,7 @@ pub(crate) fn compensate_drain( entry.key.relay_url.clone(), entry.start_on_app_launch, None, - &mut records, + recs, ) .map(|_| ()) }) diff --git a/desktop/src-tauri/src/managed_agents/runtime_commands_tests.rs b/desktop/src-tauri/src/managed_agents/runtime_commands_tests.rs index 66d92c45d..919f30c2d 100644 --- a/desktop/src-tauri/src/managed_agents/runtime_commands_tests.rs +++ b/desktop/src-tauri/src/managed_agents/runtime_commands_tests.rs @@ -476,22 +476,17 @@ fn test_partial_drain_stop_failure_delivers_stopped_prefix_and_remaining_tail() // ── compensate_drain_for tests ────────────────────────────────────────────── // -// These tests call `compensate_drain_for` directly — the extracted production -// core — with an injected `start_fn` instead of the real `start_pair_under_held_locks`. -// The injected function records which entries were restarted and returns +// These tests call `compensate_drain_for` directly — the lock-free production +// core — with an injected `start_fn`. The function takes `&mut [ManagedAgentRecord]` +// (loaded by the adapter under the store lock) and calls start_fn for each +// stopped entry. Tests inject a closure that records calls and returns // synthetic success/failure without spawning processes or touching disk. // -// This exercises the two critical invariants of compensation: -// 1. Only the `stopped` prefix (not the full journal) is restarted. -// 2. A stale captured scope skips compensation entirely (generation guard). +// The stale-scope generation guard lives in the adapter (`compensate_drain`), +// not the core, and is tested at the adapter level below. // -// The `start_pair` gate invariant (concurrent ordinary starts are serialized by -// holding the transition lock) is structural: `compensate_drain_for` runs under -// the caller's held transition guard (see `compensate_drain`'s signature), and -// every `start_pair`/`start_managed_agent_process` path acquires that same lock -// before proceeding. The lock-order audit lives in `identity.rs` and -// `workspace.rs`; no concurrent-thread test is added here because the Rust -// type system enforces the ownership invariant at compile time. +// The serialization invariant (writers blocked by store lock, not transition +// guard) is documented in the adapter and verified in the adapter-level tests. fn make_captured_scope() -> super::super::scope::WorkspaceAgentScope { // Build a scope whose generation matches the current global counter. @@ -510,6 +505,9 @@ fn make_captured_scope() -> super::super::scope::WorkspaceAgentScope { /// Two-entry compensation: entry 1 stopped, entry 2 stop-failed. /// `compensate_drain_for` must call start_fn exactly for entry 1 and return /// None (full success), proving the stopped-prefix contract. +/// +/// The start_fn receives both the entry and the mutable records slice, +/// matching `start_pair_under_held_locks`'s contract. #[test] fn test_compensate_for_restarts_stopped_entries_in_order() { let pubkey1 = "aa".repeat(32); @@ -518,11 +516,9 @@ fn test_compensate_for_restarts_stopped_entries_in_order() { // entry2 was not stopped (stop failed), so it is NOT in the stopped slice. let stopped = vec![entry1.clone()]; - let state = crate::app_state::build_app_state(); - let captured_scope = make_captured_scope(); - + let mut records: Vec = Vec::new(); let mut restarted: Vec = Vec::new(); - let result = compensate_drain_for(&state, &stopped, &captured_scope, |entry| { + let result = compensate_drain_for(&stopped, &mut records, |entry, _recs| { restarted.push(entry.key.pubkey.clone()); Ok(()) }); @@ -540,42 +536,6 @@ fn test_compensate_for_restarts_stopped_entries_in_order() { ); } -/// Stale scope: compensation skips all restarts and returns a degradation message. -/// Proves the generation guard: a workspace switch between drain failure and -/// compensation must prevent stale-scope agents from restarting. -#[test] -fn test_compensate_for_skips_all_on_stale_scope() { - let pubkey1 = "aa".repeat(32); - let entry1 = make_drain_entry(&pubkey1, "wss://relay.example", true); - let stopped = vec![entry1.clone()]; - - let state = crate::app_state::build_app_state(); - let captured_scope = make_captured_scope(); - - // Advance the global generation — scope is now stale. - super::super::scope::next_scope_generation(); - - let mut restart_count = 0usize; - let result = compensate_drain_for(&state, &stopped, &captured_scope, |_entry| { - restart_count += 1; - Ok(()) - }); - - assert!( - result.is_some(), - "stale scope must return a degradation message" - ); - let msg = result.unwrap(); - assert!( - msg.contains("compensation skipped") || msg.contains("stale scope"), - "degradation message must describe stale scope: {msg}" - ); - assert_eq!( - restart_count, 0, - "start_fn must not be called when scope is stale" - ); -} - /// Partial restart failure: start_fn returns an error for entry1, success for /// entry2. The function must return a degradation message naming the failing /// entry and not abort early (entry2 is still attempted). @@ -587,11 +547,9 @@ fn test_compensate_for_reports_partial_restart_failure() { let entry2 = make_drain_entry(&pubkey2, "wss://relay.example", false); let stopped = vec![entry1.clone(), entry2.clone()]; - let state = crate::app_state::build_app_state(); - let captured_scope = make_captured_scope(); - + let mut records: Vec = Vec::new(); let pubkey1_clone = pubkey1.clone(); - let result = compensate_drain_for(&state, &stopped, &captured_scope, |entry| { + let result = compensate_drain_for(&stopped, &mut records, |entry, _recs| { if entry.key.pubkey == pubkey1_clone { Err(format!("injected failure for {}", entry.key.pubkey)) } else { @@ -610,6 +568,30 @@ fn test_compensate_for_reports_partial_restart_failure() { ); } +/// `compensate_drain_for` passes the records slice to start_fn so it can be +/// mutated (matching `start_pair_under_held_locks`'s &mut [ManagedAgentRecord]). +/// Prove start_fn receives and can mutate the slice. +#[test] +fn test_compensate_for_start_fn_receives_records_slice() { + let pubkey1 = "aa".repeat(32); + let entry1 = make_drain_entry(&pubkey1, "wss://relay.example", true); + let stopped = vec![entry1.clone()]; + + let mut records: Vec = Vec::new(); + let mut received_records_len: Option = None; + let result = compensate_drain_for(&stopped, &mut records, |_entry, recs| { + received_records_len = Some(recs.len()); + Ok(()) + }); + + assert!(result.is_none(), "must succeed: {result:?}"); + assert_eq!( + received_records_len, + Some(0), + "start_fn must receive the records slice (empty in this test)" + ); +} + // ── compensate_drain round-trip tests (via tauri::test::mock_app) ────────── // // These tests call `compensate_drain` directly — the real production function