diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index ac4d014cf6..3fc275858b 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -977,6 +977,8 @@ impl Drop for ReservedSteer { #[derive(Clone)] pub struct EngineHandle { run_completion: Option>>, + // Process-local cache retirement, not a model/user/wire operation. + idle_retirement: Option>, /// Send operations to the engine pub tx_op: mpsc::Sender, /// Receive events from the engine @@ -1274,6 +1276,7 @@ pub struct Engine { active_route_declared_limits: Option, active_route_capabilities: codewhale_config::route::RouteCapabilities, rx_op: mpsc::Receiver, + rx_idle_retirement: mpsc::Receiver, live_runtime_authority: Arc>, compaction_cancellation: Arc>, /// Clone of the op-channel sender, so the engine can self-dispatch ops @@ -1574,6 +1577,7 @@ const SHELL_WAKE_POLL_MS: u64 = 750; enum EngineRunInput { Operation(Box), + IdleRetirement(idle_shutdown::IdleRetirement), SubAgentCompletion(SubAgentCompletion), /// A background shell job finished while the engine sat idle with an /// active goal. Shell completion is pull-only (no channel), so without @@ -2238,6 +2242,10 @@ impl Engine { } let (tx_op, rx_op) = mpsc::channel(ENGINE_OP_CHANNEL_CAPACITY); + // Capacity 1 keeps at most one retirement proposal in flight: a second + // concurrent host blocks in `send` until this boundary dequeues the + // first, which is what serializes proposals. Do not raise it. + let (idle_retirement, rx_idle_retirement) = mpsc::channel(1); let (tx_event, rx_event) = mpsc::channel(256); let (tx_approval, rx_approval) = mpsc::channel(64); let (tx_user_input, rx_user_input) = mpsc::channel(32); @@ -2517,6 +2525,7 @@ impl Engine { active_route_declared_limits, active_route_capabilities: codewhale_config::route::RouteCapabilities::default(), rx_op, + rx_idle_retirement, live_runtime_authority: Arc::clone(&live_runtime_authority), compaction_cancellation: Arc::clone(&compaction_cancellation), tx_op: tx_op.clone(), @@ -2560,6 +2569,7 @@ impl Engine { }; let handle = EngineHandle { run_completion: None, + idle_retirement: Some(idle_retirement), tx_op, rx_event: Arc::new(RwLock::new(rx_event)), cancel_token: shared_cancel_token, @@ -3181,11 +3191,7 @@ impl Engine { // The synthetic token sits behind every operation that was // already queued when it was scheduled. Drain FIFO through // that token before accepting an idle child completion. - return self - .rx_op - .recv() - .await - .map(|op| EngineRunInput::Operation(Box::new(op))); + return self.next_operation_or_retirement().await; } if let Some(ready_at) = ready_at { @@ -3202,6 +3208,9 @@ impl Engine { op = self.rx_op.recv() => { return op.map(|op| EngineRunInput::Operation(Box::new(op))); } + Some(request) = self.rx_idle_retirement.recv() => { + return Some(EngineRunInput::IdleRetirement(request)); + } () = tokio::time::sleep(ready_at.saturating_duration_since(Instant::now())) => { if let Some(scheduled) = self.scheduled_goal_continuation.as_mut() && scheduled.ready_at == Some(ready_at) @@ -3219,11 +3228,7 @@ impl Engine { // A record that is ready but could not enter the full mailbox // waits for one queued control. The next loop pass retries the // same coalesced token, so there is no spin or duplicate turn. - return self - .rx_op - .recv() - .await - .map(|op| EngineRunInput::Operation(Box::new(op))); + return self.next_operation_or_retirement().await; } else { // A per-message restriction is removed from // `active_turn_tool_security` after the turn returns, while @@ -3239,6 +3244,9 @@ impl Engine { op = self.rx_op.recv() => { return op.map(|op| EngineRunInput::Operation(Box::new(op))); } + Some(request) = self.rx_idle_retirement.recv() => { + return Some(EngineRunInput::IdleRetirement(request)); + } completion = self.rx_subagent_completion.recv(), if automatic_followups_allowed => { return completion.map(EngineRunInput::SubAgentCompletion); } @@ -3434,6 +3442,29 @@ impl Engine { } match input { + EngineRunInput::IdleRetirement(request) => { + if self.accept_idle_shutdown(request).await { + // Committed: report both external inputs and autonomous + // work observed after the sealed wait. Pending event + // slots also include reserved permits, not just sends. + let (residual, event_slots, shell_work) = self.retirement_residuals(); + if residual > 0 || event_slots > 0 || shell_work { + tracing::warn!( + residual, + event_slots, + shell_work, + ops = self.rx_op.len(), + subagent_completions = self.rx_subagent_completion.len(), + steers = self.rx_steer.len(), + pending_steers = self.pending_steers.len(), + approvals = self.rx_approval.len(), + user_inputs = self.rx_user_input.len(), + "engine retired with residual work observed after the retirement commit" + ); + } + break; + } + } EngineRunInput::SubAgentCompletion(completion) => { self.handle_idle_subagent_completion(completion).await; } @@ -4207,6 +4238,22 @@ impl Engine { manager.flush_pending_persist(); } + self.shutdown_mcp_pool().await; + } + + async fn shutdown_mcp_pool(&mut self) { + // Boot owns the same pool in a background task. Drain it before the + // shutdown pass so a late handshake cannot insert a live transport + // after shutdown_all (or outlive the engine's op-channel closure). + // Scope note: this covers the current-generation boot only. A boot + // already invalidated by a workspace change holds only the dropped + // pool — nothing it inserts later is reachable from here, and its + // transports die with that task. + if self.mcp_boot_in_flight { + // The host may already have stopped consuming UI events. Join + // the worker without draining/emitting its progress notifications. + self.join_mcp_boot().await; + } // #420: graceful MCP shutdown — send SIGTERM and give stdio servers // a brief window to exit before drop fires SIGKILL via kill_on_drop. // Best-effort: pool may not exist (no MCP configured) and the lock @@ -8077,7 +8124,7 @@ impl Engine { } } - async fn wait_for_mcp_boot(&mut self) { + async fn join_mcp_boot(&mut self) { if let Some(rx) = self.mcp_boot_done.as_mut() { while !*rx.borrow() { if rx.changed().await.is_err() { @@ -8085,6 +8132,10 @@ impl Engine { } } } + } + + async fn wait_for_mcp_boot(&mut self) { + self.join_mcp_boot().await; self.drain_mcp_boot_updates().await; } @@ -9313,6 +9364,7 @@ pub(crate) fn mock_engine_handle() -> MockEngineHandle { steer_control_state.begin_turn(); let handle = EngineHandle { run_completion: None, + idle_retirement: None, tx_op, rx_event: Arc::new(RwLock::new(rx_event)), cancel_token: shared_cancel_token, @@ -9648,6 +9700,7 @@ impl SubAgentWiring { mod approval; mod context; mod handle; +mod idle_shutdown; mod run_completion; pub use run_completion::EngineRunCompletion; pub mod preview; diff --git a/crates/tui/src/core/engine/README.md b/crates/tui/src/core/engine/README.md index 0840ecf0d3..6245e8d486 100644 --- a/crates/tui/src/core/engine/README.md +++ b/crates/tui/src/core/engine/README.md @@ -8,7 +8,10 @@ Handles constructed directly with `Engine::new` return `NotObserved`. Each waiter receives the same terminal value. Cancelling one wait neither stops the engine nor consumes another waiter's result. Waiting does not initiate -shutdown; callers must use the existing shutdown protocol. The observation +shutdown; callers must use an existing shutdown protocol: send `Op::Shutdown` +for an immediate stop, or `EngineHandle::shutdown_when_idle` to retire an idle +engine at its own quiescent boundary without losing accepted work (see +`core/engine/idle_shutdown.rs`). The observation does not keep another engine or operation-mailbox sender alive. The one-shot exec consumer uses a two-second absolute deadline shared between diff --git a/crates/tui/src/core/engine/handle.rs b/crates/tui/src/core/engine/handle.rs index 23ebb78b5d..8ba18ae6e2 100644 --- a/crates/tui/src/core/engine/handle.rs +++ b/crates/tui/src/core/engine/handle.rs @@ -95,21 +95,98 @@ async fn ack_answer(rx_ack: &mut oneshot::Receiver, budget: Duration) -> A } impl EngineHandle { + /// Retire only at the actor's own idle boundary, preserving active turns, + /// queued inputs, child/shell results and goal continuations by refusing. + /// `false` also covers states whose safe retirement cannot be certified, + /// including an outstanding advisor watcher or an event still sitting in + /// the host's queue: the host must keep consuming events until the + /// handshake completes. It does not + /// promise that waiting alone will eventually make retirement available. + /// + /// The call blocks until the actor reaches its next input boundary, so + /// retry loops must back off instead of spinning. The embedding host must + /// exclude new external work until retirement is complete — including + /// work still held by an in-flight `reserve()` permit, which the actor + /// cannot observe. `host_is_idle` checks its own consumers synchronously + /// while the actor is sealed: dequeuing an event does not prove host + /// persistence or terminal delivery finished. + /// The synchronous predicate must return promptly: an already-latched + /// cancel token deliberately cannot interrupt this sealed wait. Residual + /// work observed after commit is diagnosed, not a reopened handshake. + /// + /// Refusing resumes the actor without shutdown; so does a cancel token + /// firing for the first time during an uncommitted wait (a token already + /// latched by an earlier cancelled turn does not refuse). `Err` means no + /// proposal is answerable — the request channel or a oneshot counterpart + /// was dropped — not that the engine is unhealthy; the cancel-beats-commit + /// race also surfaces as `Err("idle-retirement proposal expired")`, so + /// an `Err` is not proof that another proposer did not commit. Observe + /// `wait_run_completion` to distinguish a stopped incarnation, and retry + /// only while it is still running. `true` means the actor + /// consumed the commit — the send receipt alone is not proof, since a + /// send racing a simultaneous cancel can succeed on a proposal the actor + /// abandons; a lost commit ack therefore reports `false`. `true` means + /// committed, not completed; wait_run_completion is still required before + /// touching resources. Cancelling an uncommitted wait cannot cause a late + /// shutdown. + pub async fn shutdown_when_idle(&self, host_is_idle: impl FnOnce() -> bool) -> Result { + let (ready, receipt) = oneshot::channel(); + let (commit, committed) = oneshot::channel(); + let (commit_ack, ack) = oneshot::channel(); + self.idle_retirement + .as_ref() + .ok_or_else(|| anyhow!("engine idle-retirement is unavailable"))? + .send(super::idle_shutdown::IdleRetirement { + ready, + commit: committed, + commit_ack, + }) + .await?; + if !receipt + .await + .map_err(|_| anyhow!("engine idle-retirement request dropped"))? + { + return Ok(false); + } + // Keep both owners at the same boundary. There is no await between this + // host check and commit. A refused host or cancelled proposal drops the + // commit sender, so the actor resumes instead of leaving a closed handle. + if !host_is_idle() { + return Ok(false); + } + commit + .send(()) + .map_err(|_| anyhow!("engine idle-retirement proposal expired"))?; + // Only the actor's ack proves it consumed the commit; the send above + // can also succeed against a proposal the cancel arm just abandoned. + if ack.await.is_err() { + return Ok(false); + } + Ok(true) + } + /// Wait for this exact background run to return or terminate abnormally. /// Cancelling a wait does not stop the engine or consume another waiter's result. /// Directly constructed handles return `NotObserved`; completion does not prove /// persistence succeeded or detached children and blocking writers stopped. - pub async fn wait_run_completion(&self) -> EngineRunCompletion { - let Some(mut completion) = self.run_completion.clone() else { - return EngineRunCompletion::NotObserved; - }; - loop { - let outcome = *completion.borrow_and_update(); - if let Some(outcome) = outcome { - return outcome; - } - if completion.changed().await.is_err() { - return EngineRunCompletion::AbnormalTermination; + /// The wait owns only its receipt and may outlive the handle without keeping + /// the engine's operation sender or event queue alive. + pub fn wait_run_completion( + &self, + ) -> impl std::future::Future + Send + 'static + use<> { + let completion = self.run_completion.clone(); + async move { + let Some(mut completion) = completion else { + return EngineRunCompletion::NotObserved; + }; + loop { + let outcome = *completion.borrow_and_update(); + if let Some(outcome) = outcome { + return outcome; + } + if completion.changed().await.is_err() { + return EngineRunCompletion::AbnormalTermination; + } } } } diff --git a/crates/tui/src/core/engine/idle_shutdown.rs b/crates/tui/src/core/engine/idle_shutdown.rs new file mode 100644 index 0000000000..5f37f09c15 --- /dev/null +++ b/crates/tui/src/core/engine/idle_shutdown.rs @@ -0,0 +1,973 @@ +//! Actor-owned, non-destructive retirement for embedding hosts. The host owns +//! external admission; this boundary owns autonomous turns and their evidence. +//! +//! An in-flight MCP boot is deliberately not part of the ready check: the run +//! teardown joins the boot worker before the pool shuts down, and that join +//! drops the boot's settlement event and failure briefing for this incarnation +//! (an engine that is about to disappear has no future turn to brief). + +use super::*; + +// An incarnation-local embedding lifetime handshake, like spawn/completion; +// it grants no new user, model or remote operation. Keep the Op/wire protocol +// and its exhaustive parity checks unchanged. +pub(super) struct IdleRetirement { + pub(super) ready: tokio::sync::oneshot::Sender, + pub(super) commit: tokio::sync::oneshot::Receiver<()>, + /// The actor acknowledges only a *consumed* commit, so the host can tell + /// a committed retirement from a send that raced a simultaneous cancel. + pub(super) commit_ack: tokio::sync::oneshot::Sender<()>, +} + +impl Engine { + pub(super) async fn next_operation_or_retirement(&mut self) -> Option { + // Biased toward the queued operation, matching the neighboring timed + // wait: at a FIFO-drain site a queued control must not lose the + // arbitration race to a retirement request. Acceptance still refuses + // while any queue is non-empty, but the ordering intent should live + // here rather than depend on that downstream refusal. + tokio::select! { + biased; + op = self.rx_op.recv() => op.map(|op| EngineRunInput::Operation(Box::new(op))), + Some(request) = self.rx_idle_retirement.recv() => Some(EngineRunInput::IdleRetirement(request)), + } + } + + fn can_retire_without_losing_work(&mut self) -> bool { + if self.has_scheduled_goal_continuation() + // Only the actor creates watcher clones — the sanctioned clone + // site is the advisor dispatch inside the actor loop in engine.rs. + // A second clone site would permanently refuse retirement, so any + // new holder must be gated into this check. Each watcher holds its + // clone through event delivery. Reuse that ownership invariant: + // get_mut synchronizes with the final clone's drop; unlike the + // enabled setting, it also covers a watcher switched off mid-run. + || self.advisor_emission_guard.as_mut().is_some_and(|guard| Arc::get_mut(guard).is_none()) + { + return false; + } + // try_lock keeps this gate non-blocking like the sibling checks: the + // goal state is shared with the host, so a host holding it across + // this call must defer retirement instead of deadlocking the actor + // before the readiness receipt is even sent. Uncertain is refuse. + if self + .config + .goal_state + .try_lock() + .map(|state| state.snapshot().is_active()) + .unwrap_or(true) + { + return false; + } + let Ok(agents) = self.subagent_manager.try_read() else { + return false; + }; + if agents.has_live_work_for_session(&self.session.id) { + return false; + } + let Ok(mut shells) = self.shell_manager.try_lock() else { + return false; + }; + if shells.has_live_or_unreported_work_for_session(&self.session.id) { + return false; + } + // Delivery publishes before a child's task/Running status clears. + // Recheck AFTER the live-work checks, with both producer owners locked: + // a child that finished during preflight cannot leave a missed payload. + // The actor is the only remaining autonomous work producer. External + // submissions are excluded by the host contract alone — including + // submissions still held by an in-flight `reserve()` permit, which + // `is_empty()` cannot observe — so what these checks accept is exactly + // what that contract leaves behind. + self.rx_op.is_empty() + && self.rx_subagent_completion.is_empty() + && self.rx_steer.is_empty() + && self.pending_steers.is_empty() + && self.rx_approval.is_empty() + && self.rx_user_input.is_empty() + && self.tx_event.capacity() == self.tx_event.max_capacity() + } + + /// Snapshot work visible after the sealed commit wait. Autonomous shell + /// evidence and events can arrive there too, without an external send. + /// Event slots include reserved permits; shell lock contention is unknown + /// work, so report it conservatively. This does not observe arbitrary host + /// shared-state writes or arrivals after the snapshot. + pub(super) fn retirement_residuals(&mut self) -> (usize, usize, bool) { + let inputs = self.rx_op.len() + + self.rx_subagent_completion.len() + + self.rx_steer.len() + + self.rx_approval.len() + + self.rx_user_input.len() + + self.pending_steers.len(); + let event_slots = self.tx_event.max_capacity() - self.tx_event.capacity(); + let shell_work = self.shell_manager.try_lock().map_or(true, |mut shells| { + shells.has_live_or_unreported_work_for_session(&self.session.id) + }); + (inputs, event_slots, shell_work) + } + + pub(super) async fn accept_idle_shutdown(&mut self, request: IdleRetirement) -> bool { + let IdleRetirement { + ready, + commit, + commit_ack, + } = request; + // This handler is reached only between complete actor operations. + // An autonomous turn that won first finishes normally; no cancel token + // is fired to make it idle. No next_run_input runs during the proposal. + if ready.is_closed() { + return false; + } + // Snapshot before readiness is offered: the escape below must not + // depend on scheduling between `ready.send` and entering the wait. + let already_cancelled = self.cancel_token.is_cancelled(); + let idle = self.can_retire_without_losing_work(); + if ready.send(idle).is_err() || !idle { + return false; + } + // Wait for the commit at the same sealed boundary, but let a cancel + // fired *during* the proposal break the wait: on a fresh token, a host + // blocked forever inside its synchronous predicate must not pin the + // actor. A token already latched by an earlier cancelled turn does not + // escape — that stale state must not make retirement permanently + // unavailable, and here it deliberately trades actor liveness for it: + // with the escape disarmed, only the commit itself can break a + // blocked predicate. Commit still wins a simultaneous cancel: a host + // send that lands while the cancel arm resolves can return `Ok` + // without being seen, so `Ok(true)` requires the commit ack, never + // the send alone. + let wait_cancel = self.cancel_token.clone(); + let committed = tokio::select! { + biased; + result = commit => { + // Acknowledge only a consumed commit: the ack, not the send + // receipt, is what the host treats as "committed". + let ok = result.is_ok(); + if ok { + let _ = commit_ack.send(()); + } + ok + } + () = async { + if already_cancelled { + std::future::pending::<()>().await; + } else { + wait_cancel.cancelled().await; + } + } => false, + }; + committed + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn fixture(root: &std::path::Path) -> (Engine, EngineHandle) { + Engine::new( + EngineConfig { + workspace: root.to_owned(), + mcp_config_path: root.join("mcp.json"), + ..Default::default() + }, + &Config::default(), + ) + } + + #[tokio::test] + async fn forkguard_idle_shutdown_sealed_wait_reports_autonomous_work() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + assert_eq!(engine.retirement_residuals(), (0, 0, false)); + let shells = Arc::clone(&engine.shell_manager); + let completion = engine.tx_subagent_completion.clone(); + let events = engine.tx_event.clone(); + let session = engine.session.id.clone(); + let (ready, receipt) = tokio::sync::oneshot::channel(); + let (commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, ack) = tokio::sync::oneshot::channel(); + let pending = tokio::spawn(async move { + let accepted = engine + .accept_idle_shutdown(IdleRetirement { + ready, + commit: committed, + commit_ack, + }) + .await; + (accepted, engine) + }); + assert!(receipt.await.unwrap()); + // The actor is sealed between ready and commit, not draining wakes. + shells + .lock() + .unwrap() + .seed_session_shell_for_test("late-shell", &session, false); + completion + .send(SubAgentCompletion { + owner_session_id: session, + agent_id: "late-child".into(), + payload: "late completion".into(), + }) + .unwrap(); + events + .try_send(Event::SteerDropped { + steer_id: "late-event".into(), + }) + .unwrap(); + commit.send(()).unwrap(); + ack.await.unwrap(); + let (accepted, mut engine) = tokio::time::timeout(Duration::from_secs(10), pending) + .await + .unwrap() + .unwrap(); + assert!(accepted); + assert_eq!(engine.retirement_residuals(), (1, 1, true)); + assert_eq!( + engine.rx_subagent_completion.try_recv().unwrap().payload, + "late completion" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_does_not_overtake_accepted_operations() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, handle) = fixture(home.path()); + handle + .send(Op::SetAdvisorEnabled { enabled: false }) + .await + .unwrap(); + let (ready, receipt) = tokio::sync::oneshot::channel(); + let (_commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, _ack) = tokio::sync::oneshot::channel(); + // Timeout-wrapped so a regression that wedges the handshake fails the + // assertion instead of hanging the whole test process. The budgets in + // this module are generous wall-clock headroom for loaded machines, + // not timing assertions. + assert!( + !tokio::time::timeout( + Duration::from_secs(10), + engine.accept_idle_shutdown(IdleRetirement { + ready, + commit: committed, + commit_ack, + }) + ) + .await + .expect("the refusal path must not wait on the commit") + ); + assert!(!receipt.await.unwrap()); + assert!(matches!( + engine.rx_op.try_recv().unwrap(), + Op::SetAdvisorEnabled { enabled: false } + )); + assert!(engine.can_retire_without_losing_work()); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_closed_lifetime_channel_leaves_operations_usable() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, mut handle) = fixture(home.path()); + drop(handle.idle_retirement.take()); + assert!( + tokio::time::timeout( + Duration::from_millis(10), + engine.next_operation_or_retirement() + ) + .await + .is_err() + ); + handle + .send(Op::SetAdvisorEnabled { enabled: false }) + .await + .unwrap(); + assert!(matches!( + engine.next_operation_or_retirement().await, + Some(EngineRunInput::Operation(_)) + )); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_checks_host_completion_after_actor_ready() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, handle) = fixture(home.path()); + let host_idle = Arc::new(std::sync::atomic::AtomicBool::new(true)); + let retiring_handle = handle.clone(); + let observed = host_idle.clone(); + let retiring = tokio::spawn(async move { + retiring_handle + .shutdown_when_idle(|| observed.load(std::sync::atomic::Ordering::Acquire)) + .await + }); + let Some(EngineRunInput::IdleRetirement(request)) = + engine.next_operation_or_retirement().await + else { + panic!("the request must reach the actor before host completion changes"); + }; + // The actor's event queue is empty, but its host has received the final + // event and is still persisting/emitting it. A pre-request check was true. + host_idle.store(false, std::sync::atomic::Ordering::Release); + assert!(engine.can_retire_without_losing_work()); + assert!( + !tokio::time::timeout( + Duration::from_secs(10), + engine.accept_idle_shutdown(request) + ) + .await + .expect("the refusal path must not wait on the commit") + ); + assert!(!retiring.await.unwrap().unwrap()); + assert!(!engine.cancel_token.is_cancelled()); + handle + .send(Op::SetAdvisorEnabled { enabled: false }) + .await + .unwrap(); + assert!(matches!( + engine.next_operation_or_retirement().await, + Some(EngineRunInput::Operation(_)) + )); + + host_idle.store(true, std::sync::atomic::Ordering::Release); + let retiring_handle = handle.clone(); + let retiring = tokio::spawn(async move { + retiring_handle + .shutdown_when_idle(|| host_idle.load(std::sync::atomic::Ordering::Acquire)) + .await + }); + let Some(EngineRunInput::IdleRetirement(request)) = + engine.next_operation_or_retirement().await + else { + panic!("the idle retry must reach the actor"); + }; + assert!( + tokio::time::timeout( + Duration::from_secs(10), + engine.accept_idle_shutdown(request) + ) + .await + .expect("the committed path must resolve promptly") + ); + assert!(retiring.await.unwrap().unwrap()); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_cancelled_proposal_does_not_stop_the_actor() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + let (ready, receipt) = tokio::sync::oneshot::channel(); + let (commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, _ack) = tokio::sync::oneshot::channel(); + let pending = tokio::spawn(async move { + let accepted = tokio::time::timeout( + Duration::from_secs(10), + engine.accept_idle_shutdown(IdleRetirement { + ready, + commit: committed, + commit_ack, + }), + ) + .await + .expect("the cancel path must not wedge the actor"); + (accepted, engine) + }); + assert!(receipt.await.unwrap()); + drop(commit); + let (accepted, mut engine) = pending.await.unwrap(); + assert!(!accepted); + assert!(!engine.cancel_token.is_cancelled()); + assert!(engine.can_retire_without_losing_work()); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_preserves_queued_child_goal_and_live_worker() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + engine + .tx_subagent_completion + .send(SubAgentCompletion { + owner_session_id: engine.session.id.clone(), + agent_id: "finished-child".into(), + payload: "the exact completed child result".into(), + }) + .unwrap(); + assert!(!engine.can_retire_without_losing_work()); + assert_eq!( + engine.rx_subagent_completion.try_recv().unwrap().payload, + "the exact completed child result" + ); + let deadline = Instant::now() + Duration::from_secs(60); + engine.scheduled_goal_continuation = Some(ScheduledGoalContinuation { + id: 42, + dynamic_tools: Vec::new(), + enqueued: false, + ready_at: Some(deadline), + was_delayed: true, + }); + assert!(!engine.can_retire_without_losing_work()); + assert_eq!( + engine + .scheduled_goal_continuation + .as_ref() + .unwrap() + .ready_at, + Some(deadline) + ); + engine.scheduled_goal_continuation = None; + { + let mut agents = engine.subagent_manager.write().await; + let child = agents.insert_test_running_agent("child", home.path()); + agents.assign_test_session_owner(&child, &engine.session.id); + } + assert!( + !engine.can_retire_without_losing_work(), + "UI/heartbeat counts cannot certify retirement" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_advisor_clone_and_delivery_both_gate_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, handle) = fixture(home.path()); + let guard = Arc::new(tokio::sync::Mutex::new( + crate::tools::subagent::EmissionGuard::new(), + )); + let watcher = guard.clone(); + engine.advisor_emission_guard = Some(guard); + engine.config.advisor_config.enabled = false; + assert!( + !engine.can_retire_without_losing_work(), + "switching off is not a join receipt" + ); + engine + .tx_event + .try_send(Event::status("completed advisor output".to_string())) + .unwrap(); + drop(watcher); + assert!( + !engine.can_retire_without_losing_work(), + "undelivered output must survive the watcher" + ); + handle.rx_event.write().await.recv().await.unwrap(); + assert!( + engine.can_retire_without_losing_work(), + "a historical watcher must not permanently refuse retirement" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_every_pending_input_queue_blocks_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + // A queued steer must refuse retirement until it is drained. + let steer = SteerMessage { + id: "queued-steer".to_string(), + target: SteerTarget { + session_epoch: 0, + turn_generation: 1, + }, + content: "steer must not be lost".to_string(), + }; + engine.rx_steer = { + let (tx, rx) = mpsc::channel(1); + tx.send(steer.clone()).await.unwrap(); + rx + }; + assert!(!engine.can_retire_without_losing_work()); + engine.rx_steer = mpsc::channel(1).1; + engine.pending_steers.push(steer); + assert!( + !engine.can_retire_without_losing_work(), + "a committed pending steer is accepted work" + ); + engine.pending_steers.clear(); + // An undelivered approval decision and an undelivered user-input + // decision are the same class of accepted work. + engine.rx_approval = { + let (tx, rx) = mpsc::channel(1); + tx.send(ApprovalDecision::Approved { + id: "approval-1".to_string(), + ack: None, + }) + .await + .unwrap(); + rx + }; + assert!(!engine.can_retire_without_losing_work()); + engine.rx_approval = mpsc::channel(1).1; + engine.rx_user_input = { + let (tx, rx) = mpsc::channel(1); + tx.send(UserInputDecision::Cancelled { + id: "input-1".to_string(), + }) + .await + .unwrap(); + rx + }; + assert!(!engine.can_retire_without_losing_work()); + engine.rx_user_input = mpsc::channel(1).1; + assert!(engine.can_retire_without_losing_work()); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_live_and_unreported_shell_work_gates_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + let session = engine.session.id.clone(); + { + let mut shells = engine.shell_manager.lock().unwrap(); + shells.seed_session_shell_for_test("live-job", &session, true); + } + assert!( + !engine.can_retire_without_losing_work(), + "a running session shell is live work" + ); + { + let mut shells = engine.shell_manager.lock().unwrap(); + shells.finish_seeded_shell_for_test("live-job"); + } + assert!( + !engine.can_retire_without_losing_work(), + "finished output not yet delivered must survive retirement" + ); + { + let mut shells = engine.shell_manager.lock().unwrap(); + shells.mark_seeded_completion_reported_for_test("live-job"); + } + assert!( + engine.can_retire_without_losing_work(), + "delivered shell evidence no longer blocks retirement" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_timed_goal_wait_still_serves_the_retirement_arm() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, handle) = fixture(home.path()); + let deadline = Instant::now() + Duration::from_secs(60); + engine.scheduled_goal_continuation = Some(ScheduledGoalContinuation { + id: 7, + dynamic_tools: Vec::new(), + enqueued: false, + ready_at: Some(deadline), + was_delayed: true, + }); + let (ready, receipt) = tokio::sync::oneshot::channel(); + let (_commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, _ack) = tokio::sync::oneshot::channel(); + handle + .idle_retirement + .as_ref() + .unwrap() + .send(IdleRetirement { + ready, + commit: committed, + commit_ack, + }) + .await + .unwrap(); + // The biased timed wait for the goal continuation must serve the + // retirement arm instead of sleeping through the whole window. + let input = tokio::time::timeout(Duration::from_millis(2000), engine.next_run_input(false)) + .await + .expect("the armed retirement request must not wait out the goal timer"); + let Some(EngineRunInput::IdleRetirement(request)) = input else { + panic!("the retirement arm must fire while the goal wait is armed"); + }; + assert!( + !tokio::time::timeout( + Duration::from_secs(10), + engine.accept_idle_shutdown(request) + ) + .await + .expect("the refusal path must not wait on the commit"), + "a scheduled goal continuation refuses retirement" + ); + assert!(!receipt.await.unwrap()); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_actor_abort_fails_the_host_commit_send() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + let (ready, receipt) = tokio::sync::oneshot::channel(); + let (commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, _ack) = tokio::sync::oneshot::channel(); + let accept = tokio::spawn(async move { + engine + .accept_idle_shutdown(IdleRetirement { + ready, + commit: committed, + commit_ack, + }) + .await + }); + assert!( + receipt.await.unwrap(), + "a clean fixture must offer retirement" + ); + // The actor incarnation dies mid-handshake: the engine and the commit + // receiver it owns drop with the aborted task. Awaiting the handle + // makes the drop deterministic before the host-side send. + accept.abort(); + let _ = accept.await; + assert!( + commit.send(()).is_err(), + "the host must see the proposal expire instead of a silent no-op" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_fresh_cancel_escapes_a_wedged_commit_wait() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, handle) = fixture(home.path()); + let (ready, receipt) = tokio::sync::oneshot::channel(); + let (commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, ack) = tokio::sync::oneshot::channel(); + let pending = tokio::spawn(async move { + let accepted = tokio::time::timeout( + Duration::from_secs(10), + engine.accept_idle_shutdown(IdleRetirement { + ready, + commit: committed, + commit_ack, + }), + ) + .await + .expect("a fresh cancel must break the commit wait"); + (accepted, engine) + }); + assert!(receipt.await.unwrap()); + // The host cancels while its commit is still outstanding: the actor + // must resume instead of waiting forever on the commit. + handle.cancel(); + let (accepted, mut engine) = pending.await.unwrap(); + assert!(!accepted, "the cancel escape must refuse, not accept"); + assert!( + ack.await.is_err(), + "an escaped proposal must not be acknowledged as committed" + ); + assert!(engine.cancel_token.is_cancelled()); + drop(commit); + assert!( + engine.can_retire_without_losing_work(), + "the refused proposal leaves the actor usable" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_stale_cancelled_turn_does_not_refuse_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, handle) = fixture(home.path()); + // Latch a cancellation from a (synthetic) earlier turn: the idle slot + // keeps that token until the next turn swaps it. + handle.cancel(); + assert!(engine.cancel_token.is_cancelled()); + let (ready, receipt) = tokio::sync::oneshot::channel(); + let (commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, ack) = tokio::sync::oneshot::channel(); + let pending = tokio::spawn(async move { + let accepted = tokio::time::timeout( + Duration::from_secs(10), + engine.accept_idle_shutdown(IdleRetirement { + ready, + commit: committed, + commit_ack, + }), + ) + .await + .expect("a stale cancelled token must not wedge the proposal"); + (accepted, engine) + }); + assert!( + receipt.await.unwrap(), + "a latched cancellation must not refuse the proposal itself" + ); + commit.send(()).unwrap(); + assert!( + ack.await.is_ok(), + "a consumed commit must be acknowledged to the host" + ); + let (accepted, mut engine) = pending.await.unwrap(); + assert!( + accepted, + "commit must win over the already-fired token, or a committed host is silently refused" + ); + assert!( + engine.can_retire_without_losing_work(), + "the accepted proposal leaves nothing stuck" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_queued_operation_wins_the_retirement_arbitration() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, handle) = fixture(home.path()); + handle + .send(Op::SetAdvisorEnabled { enabled: false }) + .await + .unwrap(); + let (ready, _receipt) = tokio::sync::oneshot::channel(); + let (_commit, committed) = tokio::sync::oneshot::channel(); + let (commit_ack, _ack) = tokio::sync::oneshot::channel(); + handle + .idle_retirement + .as_ref() + .unwrap() + .send(IdleRetirement { + ready, + commit: committed, + commit_ack, + }) + .await + .unwrap(); + // Both inputs are queued before the drain runs, so the biased arm + // order — not luck — must decide: the FIFO drain site serves the + // operation first and still reaches the retirement request after. + let first = tokio::time::timeout( + Duration::from_millis(2000), + engine.next_operation_or_retirement(), + ) + .await + .expect("both inputs are queued; the drain must resolve"); + assert!( + matches!(first, Some(EngineRunInput::Operation(_))), + "a queued operation must not lose the arbitration race to a retirement request" + ); + let second = tokio::time::timeout( + Duration::from_millis(2000), + engine.next_operation_or_retirement(), + ) + .await + .expect("the retirement request must follow the queued operation"); + assert!( + matches!(second, Some(EngineRunInput::IdleRetirement(_))), + "the retirement request must still be served after the op" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_parked_child_queued_mail_gates_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + let child = { + let mut agents = engine.subagent_manager.write().await; + let (child, _continuation) = agents.insert_test_interrupted_continuable_agent( + "parked", + home.path(), + vec![crate::models::Message { + role: crate::models::Role::User, + content: vec![crate::models::ContentBlock::Text { + text: "recoverable work".into(), + cache_control: None, + }], + }], + ); + agents.assign_test_session_owner(&child, &engine.session.id); + child + }; + // A follow-up accepted for a parked child lands in the in-memory + // mailbox with a queue-only receipt. It is accepted work: the + // retirement gate must refuse until the owner session drains it. + { + let mut agents = engine.subagent_manager.write().await; + let receipt = agents + .followup_child(&child, "accepted while parked".to_string()) + .expect("queue-only follow-up receipt"); + assert!(!receipt.woke, "a parked child cannot be woken"); + } + assert!( + !engine.can_retire_without_losing_work(), + "queued parent mail for a parked child must survive retirement" + ); + { + let mut agents = engine.subagent_manager.write().await; + agents.clear_test_queued_mail(&child); + } + assert!( + engine.can_retire_without_losing_work(), + "draining the parked mail restores the idle certification" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_delivered_follow_up_counter_tracks_live_children_only() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + // Hold the receiver for the whole test: dropping it would close the + // child's input channel and turn the wake delivery into a failure. + let (_child_rx, child) = { + let mut agents = engine.subagent_manager.write().await; + let (child, rx) = agents.insert_test_running_agent_with_input("countered", home.path()); + agents.assign_test_session_owner(&child, &engine.session.id); + (rx, child) + }; + { + let mut agents = engine.subagent_manager.write().await; + let receipt = agents + .followup_child(&child, "delivered to a live child".to_string()) + .expect("wake receipt"); + assert!( + receipt.woke, + "a running child with a live input channel wakes" + ); + } + assert!( + !engine.can_retire_without_losing_work(), + "a delivered follow-up the child has not consumed is accepted work" + ); + { + let mut agents = engine.subagent_manager.write().await; + agents.commit_test_terminal_result( + &child, + crate::tools::subagent::SubAgentStatus::Cancelled, + ); + } + assert!( + engine.can_retire_without_losing_work(), + "a terminal child can never consume its counter: retirement must not refuse on bookkeeping no drain can clear" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_mail_on_a_dead_end_child_does_not_gate_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + let child = { + let mut agents = engine.subagent_manager.write().await; + let (child, _continuation) = agents.insert_test_interrupted_continuable_agent( + "dead-end", + home.path(), + vec![crate::models::Message { + role: crate::models::Role::User, + content: vec![crate::models::ContentBlock::Text { + text: "recoverable work".into(), + cache_control: None, + }], + }], + ); + agents.assign_test_session_owner(&child, &engine.session.id); + child + }; + { + let mut agents = engine.subagent_manager.write().await; + agents + .followup_child(&child, "queued for a resumable child".to_string()) + .expect("queue-only follow-up receipt"); + } + assert!( + !engine.can_retire_without_losing_work(), + "mail a resume could still deliver is accepted work" + ); + { + let mut agents = engine.subagent_manager.write().await; + agents.mark_test_agent_failed(&child); + } + assert!( + engine.can_retire_without_losing_work(), + "followup refuses a failed child, so its parked mail has no remaining delivery path and must not gate forever" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_child_owned_shell_evidence_does_not_block_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + let session = engine.session.id.clone(); + { + let mut shells = engine.shell_manager.lock().unwrap(); + shells.seed_agent_shell_for_test("child-job", &session, "agent_child", true); + } + assert!( + !engine.can_retire_without_losing_work(), + "a running agent-owned shell is live work regardless of owner" + ); + { + let mut shells = engine.shell_manager.lock().unwrap(); + shells.finish_seeded_shell_for_test("child-job"); + } + assert!( + engine.can_retire_without_losing_work(), + "the parent stream suppresses child-owned completions and a terminal owner never claims them, so finished child-owned evidence must not gate" + ); + { + let mut shells = engine.shell_manager.lock().unwrap(); + shells.seed_session_shell_for_test("parent-job", &session, false); + } + assert!( + !engine.can_retire_without_losing_work(), + "parent-owned finished-unreported evidence must still gate retirement" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_active_goal_state_gates_retirement() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + { + let mut state = engine.config.goal_state.lock().expect("goal state"); + state.sync_from_host_status( + Some("retirement-goal"), + None, + crate::tools::goal::GoalStatus::Active, + ); + } + assert!( + !engine.can_retire_without_losing_work(), + "an active goal state is a continuation the host has not observed" + ); + { + let mut state = engine.config.goal_state.lock().expect("goal state"); + state.sync_from_host_status( + Some("retirement-goal"), + None, + crate::tools::goal::GoalStatus::Complete, + ); + } + assert!( + engine.can_retire_without_losing_work(), + "a completed goal state no longer holds accepted work" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_empty_interrupt_checkpoint_does_not_gate_forever() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + { + let mut agents = engine.subagent_manager.write().await; + let (child, _rx) = agents.insert_test_running_agent_with_input("early", home.path()); + agents.assign_test_session_owner(&child, &engine.session.id); + let (_, interrupted) = agents + .interrupt_child(&child, None, "before first checkpoint".into()) + .unwrap(); + assert!(interrupted.checkpoint.as_ref().unwrap().messages.is_empty()); + let receipt = agents + .followup_child(&child, "no resume route".into()) + .unwrap(); + assert!(!receipt.woke); + assert!(receipt.continuation_handle.is_none()); + } + assert!( + engine.can_retire_without_losing_work(), + "an empty placeholder is not a usable resume route" + ); + } + + #[tokio::test] + async fn forkguard_idle_shutdown_goal_lock_contention_retries_after_release() { + let home = tempfile::tempdir().unwrap(); + let (mut engine, _handle) = fixture(home.path()); + let goal = Arc::clone(&engine.config.goal_state); + let held = goal.lock().unwrap(); + assert!( + !engine.can_retire_without_losing_work(), + "an unknown goal state must refuse retirement" + ); + drop(held); + assert!( + engine.can_retire_without_losing_work(), + "contention must not latch the gate closed" + ); + } +} diff --git a/crates/tui/src/core/engine/run_completion.rs b/crates/tui/src/core/engine/run_completion.rs index 6872959901..fbaba168a3 100644 --- a/crates/tui/src/core/engine/run_completion.rs +++ b/crates/tui/src/core/engine/run_completion.rs @@ -125,6 +125,26 @@ mod tests { (run, handle, release, dropped) } + #[tokio::test] + async fn forkguard_completion_wait_does_not_retain_engine_handle() { + let mut mock = mock_engine_handle(); + let (finished, receipt) = watch::channel(None); + mock.handle.run_completion = Some(receipt); + let waiting = mock.handle.wait_run_completion(); + drop(mock.handle); + assert!( + matches!( + mock.rx_op.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) + ), + "observing completion must not keep the engine alive through an op sender" + ); + finished + .send(Some(EngineRunCompletion::NormalReturn)) + .unwrap(); + assert_eq!(waiting.await, EngineRunCompletion::NormalReturn); + } + #[tokio::test] async fn run_completion_drops_future_before_notifying_all_waiters() { let (run, handle, release, dropped) = fixture(false); diff --git a/crates/tui/src/core/engine/tests.rs b/crates/tui/src/core/engine/tests.rs index d0eecaf5ea..d23e4e38ea 100644 --- a/crates/tui/src/core/engine/tests.rs +++ b/crates/tui/src/core/engine/tests.rs @@ -1161,6 +1161,105 @@ struct FirstRequestGatedGoalModelClient { release_request: std::sync::Arc, } +#[tokio::test] +async fn forkguard_idle_shutdown_preserves_an_autonomous_turn_before_host_observation() { + let home = tempdir().unwrap(); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let client = Arc::new(FirstRequestGatedGoalModelClient { + calls: std::sync::atomic::AtomicUsize::new(0), + request_entered: entered.clone(), + release_request: release.clone(), + }); + let config: Config = serde_json::from_value(json!({ + "provider": "local-fixture", + "providers": {"local-fixture": { + "kind": "openai-compatible", "base_url": "http://127.0.0.1:9/v1", + "model": "local-model", "api_key": "unused-local-fixture" + }} + })) + .unwrap(); + let mut engine_config = EngineConfig { + workspace: home.path().to_owned(), + model: "local-model".into(), + max_steps: 1, + snapshots_enabled: false, + terminal_chrome_enabled: false, + ..Default::default() + }; + engine_config + .features + .disable(crate::features::Feature::Mcp); + let (engine, mut handle) = Engine::new_with_model_client(engine_config, &config, client); + engine + .tx_subagent_completion + .send(SubAgentCompletion { + owner_session_id: engine.session.id.clone(), + agent_id: "completed-child".into(), + payload: "child result ready for the autonomous followup".into(), + }) + .unwrap(); + let (run, completion) = run_completion::observe(Box::pin(engine.run())); + handle.run_completion = Some(completion); + let running = tokio::spawn(run); + tokio::time::timeout(Duration::from_secs(5), entered.notified()) + .await + .unwrap(); + // Do not read TurnStarted: reproduce the host's unobserved-active window. + let active_token = handle.cancel_token.lock().unwrap().token.clone(); + let retirement_handle = handle.clone(); + let retiring = tokio::spawn(async move { retirement_handle.shutdown_when_idle(|| true).await }); + tokio::time::timeout(Duration::from_secs(5), async { + while handle.idle_retirement.as_ref().unwrap().capacity() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(!retiring.is_finished()); + assert!(!active_token.is_cancelled()); + release.notify_one(); + let snapshot = tokio::time::timeout(Duration::from_secs(5), handle.get_session_snapshot()) + .await + .unwrap() + .unwrap(); + assert!( + !tokio::time::timeout(Duration::from_secs(5), retiring) + .await + .unwrap() + .unwrap() + .unwrap() + ); + assert!(!active_token.is_cancelled()); + assert!( + serde_json::to_string(&snapshot.messages) + .unwrap() + .contains("still working") + ); + let mut completed = false; + let mut text = String::new(); + { + let mut events = handle.rx_event.write().await; + while let Ok(event) = events.try_recv() { + match event { + Event::MessageDelta { content, .. } => text.push_str(&content), + Event::TurnComplete { status, .. } => { + assert_eq!(status, TurnOutcomeStatus::Completed); + completed = true; + } + _ => {} + } + } + } + assert!(completed && text.contains("still working")); + assert!(handle.shutdown_when_idle(|| true).await.unwrap()); + assert_eq!( + handle.wait_run_completion().await, + EngineRunCompletion::NormalReturn + ); + running.await.unwrap(); +} + struct IndexedGatedGoalModelClient { calls: std::sync::atomic::AtomicUsize, gates: HashMap< @@ -29995,6 +30094,7 @@ fn engine_handle_try_send_does_not_block_when_op_channel_is_full() { let cancel_token = CancellationToken::new(); let handle = EngineHandle { run_completion: None, + idle_retirement: None, tx_op, rx_event: Arc::new(RwLock::new(mpsc::channel::(1).1)), cancel_token: Arc::new(StdMutex::new(super::TurnCancelSlot { @@ -30185,6 +30285,55 @@ async fn reload_mcp_op_recovers_from_invalid_initial_config_in_process() { task.await.expect("engine task"); } +#[tokio::test] +async fn forkguard_shutdown_waits_for_pending_mcp_boot_without_emitting_events() { + let tmp = tempdir().expect("tempdir"); + let (mut engine, _handle) = Engine::new( + EngineConfig { + workspace: tmp.path().to_path_buf(), + ..Default::default() + }, + &Config::default(), + ); + // A host reclaim may have stopped its event forwarder. Filling this queue + // pins the join (not the UI-emitting drain) being used here: the queued + // Finished update carries a connection failure, so `wait_for_mcp_boot`'s + // drain would inject the boot briefing and block forever on this full + // queue, while the join must leave the update unconsumed. + while engine.tx_event.try_send(Event::status("queued")).is_ok() {} + let (done_tx, done_rx) = tokio::sync::watch::channel(false); + let (progress_tx, progress_rx) = mpsc::unbounded_channel(); + engine.mcp_pool = Some(Arc::new(AsyncMutex::new( + McpPool::new(McpConfig::default()), + ))); + engine.mcp_event_generation = 1; + engine.mcp_boot_in_flight = true; + engine.mcp_boot_done = Some(done_rx); + engine.mcp_boot_rx = Some(progress_rx); + engine.mcp_boot_generation = Some(1); + progress_tx + .send(McpBootUpdate::Finished { + generation: 1, + authority_errors: Arc::new(HashMap::new()), + connection_errors: HashMap::from([( + "late-server".to_string(), + "connection refused".to_string(), + )]), + }) + .unwrap(); + let mut shutdown = Box::pin(engine.shutdown_mcp_pool()); + assert!( + tokio::time::timeout(Duration::from_millis(20), &mut shutdown) + .await + .is_err(), + "shutdown must not finish while a boot worker can still insert a connection" + ); + done_tx.send(true).unwrap(); + tokio::time::timeout(Duration::from_secs(1), shutdown) + .await + .expect("shutdown must not depend on an event consumer"); +} + #[tokio::test] async fn mcp_boot_updates_preserve_authority_errors_and_replace_ordinary_errors() { let tmp = tempdir().expect("tempdir"); diff --git a/crates/tui/src/tools/shell.rs b/crates/tui/src/tools/shell.rs index 2a788ed087..3d6598cab7 100644 --- a/crates/tui/src/tools/shell.rs +++ b/crates/tui/src/tools/shell.rs @@ -2071,6 +2071,73 @@ impl ShellManager { } } + /// Insert a session-owned record for retirement-check tests without + /// spawning: `running` selects the live vs finished shape. The record is + /// finished-but-unreported when not running, exactly the state the + /// retirement boundary must refuse on. + #[cfg(test)] + pub(crate) fn seed_session_shell_for_test( + &mut self, + id: impl Into, + owner_session_id: &str, + running: bool, + ) { + let now = Instant::now(); + let id = id.into(); + self.processes.insert( + id.clone(), + BackgroundShell { + id, + command: String::new(), + working_dir: self.default_workspace.clone(), + status: if running { + ShellStatus::Running + } else { + ShellStatus::Completed + }, + exit_code: (!running).then_some(0), + started_at: now, + finished_at: (!running).then_some(now), + last_output_at: now, + last_observed_output_len: 0, + sandbox_type: SandboxType::None, + linked_task_id: None, + owner_agent: None, + owner_session_id: owner_session_id.to_string(), + origin_tool_call_id: None, + origin_turn_id: None, + ownership: ShellOwnership::Managed, + stdout_buffer: new_shared_raw_output(), + stderr_buffer: Some(new_shared_raw_output()), + bounded_output: None, + heavy_permit: None, + stdout_cursor: 0, + stderr_cursor: 0, + completion_reported: false, + stdin: None, + child: None, + stop_unconfirmed: false, + #[cfg(windows)] + windows_job: None, + stdout_thread: None, + stderr_thread: None, + work_lifecycle: None, + lifecycle_seq: 0, + last_lifecycle_status: None, + last_lifecycle_bytes: 0, + }, + ); + } + + /// Flip a seeded record's completion to reported, modeling the engine's + /// claim of the finished job's evidence. + #[cfg(test)] + pub(crate) fn mark_seeded_completion_reported_for_test(&mut self, id: &str) { + if let Some(shell) = self.processes.get_mut(id) { + shell.completion_reported = true; + } + } + /// Model the state only an un-reaped stopped child can produce (see /// `stop_unconfirmed`): no portable test can spawn a child that survives /// SIGKILL, so tests set the flag — and the terminal status — the kill @@ -3542,6 +3609,55 @@ impl ShellManager { }) } + /// Retirement must protect agent-owned jobs while they run and finished + /// parent-owned output not yet delivered. Polling observes exit without + /// consuming any completion. Finished-unreported evidence gates only for + /// parent-owned jobs: the engine's boundary drain claims child-owned + /// completions but suppresses their delivery from the parent stream, and + /// once the owning child is terminal nothing will ever claim them — the + /// child's own liveness is gated on the sub-agent side. An unconfirmed + /// stop holds the frame like every other liveness predicate + /// (`holds_write_scope`) until the exit is observed; confirming here + /// keeps a seeded/reaped record from wedging retirement. + pub(crate) fn has_live_or_unreported_work_for_session( + &mut self, + active_session_id: &str, + ) -> bool { + !active_session_id.is_empty() + && self.processes.values_mut().any(|shell| { + if shell.owner_session_id != active_session_id { + return false; + } + shell.poll(); + shell.confirm_stop_when_reaped(); + shell.status == ShellStatus::Running + || shell.stop_unconfirmed + || (shell.owner_agent.is_none() && !shell.completion_reported) + }) + } + + /// Insert an agent-owned session record for retirement-check tests + /// without spawning: mirrors `seed_session_shell_for_test` with a + /// non-`None` `owner_agent`, the shape a nested child's background job + /// carries in production. + #[cfg(test)] + pub(crate) fn seed_agent_shell_for_test( + &mut self, + id: impl Into, + owner_session_id: &str, + owner_agent_id: &str, + running: bool, + ) { + let id = id.into(); + self.seed_session_shell_for_test(id.clone(), owner_session_id, running); + if let Some(shell) = self.processes.get_mut(&id) { + shell.owner_agent = Some(ShellJobOwner { + agent_id: owner_agent_id.to_string(), + agent_name: owner_agent_id.to_string(), + }); + } + } + /// Return agent owners whose tracked shell work is still running. The /// engine uses this to keep a worker's heartbeat alive while its only /// pending work is an explicitly tracked background shell task. diff --git a/crates/tui/src/tools/shell/tests.rs b/crates/tui/src/tools/shell/tests.rs index 7a26942d8d..3902e084ac 100644 --- a/crates/tui/src/tools/shell/tests.rs +++ b/crates/tui/src/tools/shell/tests.rs @@ -5255,6 +5255,48 @@ fn cleanup_keeps_an_unconfirmed_stop_until_the_exit_is_confirmed() { ); } +/// The retirement gate must agree with the coordination scope on unconfirmed +/// stops: a stop that was sent but whose exit nobody observed is still live +/// work even when the completion evidence was already claimed, and a +/// confirmed exit releases the frame. A real child, because the confirmation +/// inside the gate is a real reap. +#[cfg(unix)] +#[test] +fn retirement_gate_holds_an_unconfirmed_stop_until_the_exit_is_confirmed() { + let tmp = tempdir().expect("tempdir"); + let mut manager = ShellManager::new(tmp.path().to_path_buf()); + manager.seed_running_shell_for_test("job-1", "agent-1"); + let child = Command::new("sleep").arg("30").spawn().expect("spawn"); + let pid = child.id(); + { + let shell = manager.processes.get_mut("job-1").expect("seeded job"); + shell.owner_session_id = "session-1".to_string(); + shell.status = ShellStatus::Killed; + shell.completion_reported = true; + shell.stop_unconfirmed = true; + shell.child = Some(ShellChild::Process(child)); + } + assert!( + manager.has_live_or_unreported_work_for_session("session-1"), + "a stop whose exit was never observed is live work even with the completion claimed" + ); + + // The OS ends the child; the confirm inside the gate reaps it and the + // same observation certifies the session for retirement. + assert_eq!(unsafe { libc::kill(pid as libc::pid_t, libc::SIGKILL) }, 0); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + if !manager.has_live_or_unreported_work_for_session("session-1") { + break; + } + assert!( + Instant::now() < deadline, + "the retirement gate never released the confirmed stop" + ); + std::thread::sleep(Duration::from_millis(20)); + } +} + /// Capacity cleanup must not count — or evict — a job whose stop has not been /// confirmed: it still holds its write scope even though its status is /// finished. The production age bound keeps the fresh filler records, so the @@ -5367,3 +5409,36 @@ fn dropping_the_manager_stops_a_child_with_an_unconfirmed_stop() { std::thread::sleep(Duration::from_millis(20)); } } + +#[test] +fn forkguard_idle_shutdown_shell_predicate_gates_on_live_and_unreported_work() { + let workspace = tempdir().expect("tempdir"); + let mut manager = ShellManager::new(workspace.path().to_path_buf()); + assert!( + !manager.has_live_or_unreported_work_for_session("session-a"), + "an empty manager has no session work" + ); + assert!( + !manager.has_live_or_unreported_work_for_session(""), + "an empty session id matches nothing, like the sibling helpers" + ); + manager.seed_session_shell_for_test("live-job", "session-a", true); + assert!( + manager.has_live_or_unreported_work_for_session("session-a"), + "a running session shell is live work" + ); + assert!( + !manager.has_live_or_unreported_work_for_session("session-b"), + "another session's shell is not this engine's work" + ); + manager.finish_seeded_shell_for_test("live-job"); + assert!( + manager.has_live_or_unreported_work_for_session("session-a"), + "finished output not yet delivered must survive retirement" + ); + manager.mark_seeded_completion_reported_for_test("live-job"); + assert!( + !manager.has_live_or_unreported_work_for_session("session-a"), + "delivered shell evidence no longer blocks retirement" + ); +} diff --git a/crates/tui/src/tools/subagent/mod.rs b/crates/tui/src/tools/subagent/mod.rs index 45e96577cf..31a573e6ad 100644 --- a/crates/tui/src/tools/subagent/mod.rs +++ b/crates/tui/src/tools/subagent/mod.rs @@ -344,6 +344,9 @@ const MAX_AGENT_WORKER_EVENTS_PER_RECORD: usize = 128; /// this budget (always at least the last one, so continuability is /// preserved) and records how many older messages were omitted. Full tool /// outputs remain recoverable from the spillover files on disk. +/// The first checkpoint of freshly accepted parent input also retains that +/// input's complete suffix until the manager confirms its recovery handoff; +/// subsequent history checkpoints use the ordinary tail budget again. const SUBAGENT_CHECKPOINT_MESSAGE_BUDGET_BYTES: usize = 256 * 1024; /// Byte budget for the message tail embedded in a `subagent_full_transcript` /// handle (#3882). One handle is retained in memory per agent; the payload @@ -1669,29 +1672,48 @@ pub(crate) fn subagent_thinking_label(thinking: SubAgentThinking) -> &'static st } #[derive(Debug, Clone)] -struct SubAgentInput { +pub(crate) struct SubAgentInput { text: String, interrupt: bool, - /// Live "queued for the next round" counter shared with the manager. - /// Decremented when the child loop takes the input off its channel, so - /// the rail's `· N queued` count is the truthful number of follow-ups the - /// child has not yet folded into its next model round. - pending: Option>, + /// Shared recovery ticket. Taking it updates the rail's unread count; + /// checkpoint acceptance confirms that its content can be recovered. + pending: Option>, } +/// The manager retains accepted input until a checkpoint includes it. Taking +/// an item off the channel updates the UI count, but is not a recovery commit. +#[derive(Debug)] +struct PendingParentMail { + mail: QueuedParentMessage, + taken: std::sync::atomic::AtomicBool, + appended: std::sync::atomic::AtomicBool, + message_index: std::sync::atomic::AtomicUsize, +} + +impl PendingParentMail { + fn retained_index(&self, messages: &[Message], omitted: usize) -> Option { + if !self.appended.load(std::sync::atomic::Ordering::Acquire) { + return None; + } + let index = self + .message_index + .load(std::sync::atomic::Ordering::Acquire); + let message = messages.get(index.checked_sub(omitted)?)?; + (message.role == Role::User + && matches!(message.content.as_slice(), [ContentBlock::Text { text, .. }] if text == &self.mail.text)) + .then_some(index) + } +} + +type PendingParentMailQueue = Arc>>>; + impl SubAgentInput { /// Mark this input as consumed by the child loop. fn mark_taken(&self) { if let Some(pending) = self.pending.as_ref() { - // std renamed `fetch_update` to `try_update` and deprecated the - // old name in Rust 1.99. The workspace MSRV (1.88) predates the - // new name, so keep the old call until the MSRV passes 1.99. - #[allow(deprecated)] - let _ = pending.fetch_update( - std::sync::atomic::Ordering::AcqRel, - std::sync::atomic::Ordering::Acquire, - |value| Some(value.saturating_sub(1)), - ); + pending + .taken + .store(true, std::sync::atomic::Ordering::Release); } } } @@ -1702,6 +1724,11 @@ fn append_subagent_inputs_as_user_messages( ) { while let Some(input) = pending_inputs.pop_front() { if !input.text.trim().is_empty() { + if let Some(pending) = input.pending.as_ref() { + pending + .message_index + .store(messages.len(), std::sync::atomic::Ordering::Release); + } messages.push(Message { role: Role::User, content: vec![ContentBlock::Text { @@ -1710,6 +1737,11 @@ fn append_subagent_inputs_as_user_messages( }], }); } + if let Some(pending) = input.pending { + pending + .appended + .store(true, std::sync::atomic::Ordering::Release); + } } } @@ -3497,9 +3529,9 @@ pub struct SubAgentManager { /// Parent mail queued by `agents/message` without waking the child. /// `agents/followup` drains into `input_tx` when a live wake is possible. queued_mail: HashMap>, - /// Follow-ups handed to a running child's live input channel that the - /// child has not yet taken at a round boundary. Read by the rail rows. - pending_follow_ups: HashMap>, + /// Accepted follow-ups not yet included in a stored checkpoint. The rail + /// counts only untaken tickets; recovery retains their content longer. + pending_follow_ups: HashMap, /// Close a test receiver after an exact successful send, without scheduler races. #[cfg(test)] followup_delivery_hook: Option>, @@ -6330,7 +6362,13 @@ impl SubAgentManager { } fn followup_child_by_id(&mut self, agent_id: &str, text: String) -> Result { - let agent_id = agent_id.to_owned(); + // Terminal receipt ids stay immutable, but new follow-ups belong to + // their latest resumed incarnation. Otherwise accepted mail parks on + // a predecessor that will never drain again and prevents retirement. + let mut agent_id = agent_id.to_owned(); + while let Some(target) = self.resume_targets.get(&agent_id) { + agent_id = target.clone(); + } let status = self .agents .get(&agent_id) @@ -6371,15 +6409,24 @@ impl SubAgentManager { let mut undelivered = VecDeque::new(); let mut delivered = 0_usize; if let Some(tx) = input_tx { - let counter = + let inbox = Arc::clone(self.pending_follow_ups.entry(agent_id.clone()).or_default()); while let Some(mail) = pending.next() { - counter.fetch_add(1, std::sync::atomic::Ordering::AcqRel); + let ticket = Arc::new(PendingParentMail { + mail: mail.clone(), + taken: std::sync::atomic::AtomicBool::new(false), + appended: std::sync::atomic::AtomicBool::new(false), + message_index: std::sync::atomic::AtomicUsize::new(usize::MAX), + }); + inbox + .lock() + .unwrap_or_else(|e| e.into_inner()) + .push_back(Arc::clone(&ticket)); if tx .send(SubAgentInput { text: mail.text.clone(), interrupt: false, - pending: Some(Arc::clone(&counter)), + pending: Some(Arc::clone(&ticket)), }) .is_ok() { @@ -6389,7 +6436,10 @@ impl SubAgentManager { hook(delivered); } } else { - counter.fetch_sub(1, std::sync::atomic::Ordering::AcqRel); + inbox + .lock() + .unwrap_or_else(|e| e.into_inner()) + .retain(|entry| !Arc::ptr_eq(entry, &ticket)); undelivered.push_back(mail); undelivered.extend(pending); break; @@ -6757,7 +6807,12 @@ impl SubAgentManager { .and_then(|record| record.spec.child_route.clone()); ( agent.agent_type.clone(), - build_resume_prompt(&agent.prompt, checkpoint, followup_text), + build_resume_prompt( + &agent.prompt, + checkpoint, + self.queued_mail.get(&agent_id), + followup_text, + ), agent.assignment.clone(), agent.allowed_tools.clone(), agent.model.clone(), @@ -6833,6 +6888,28 @@ impl SubAgentManager { allowed_tools, options, )?; + // A second interrupt may win before the new task first polls. Its + // accepted initial input must already have a real recovery route. + let initial = self + .agents + .get(&resumed.agent_id) + .expect("spawn registered child") + .prompt + .clone(); + let checkpoint = build_subagent_checkpoint( + &resumed.agent_id, + "resume_admitted", + &[Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: initial, + cache_control: None, + }], + }], + 0, + true, + ); + self.update_checkpoint(&resumed.agent_id, checkpoint); let receipt = ParentMailReceipt { agent_id: resumed.agent_id.clone(), status: subagent_status_name(&resumed.status).to_string(), @@ -6845,6 +6922,10 @@ impl SubAgentManager { resumed.agent_id, resumed.model, agent_id ), }; + // The initial prompt already contains the accepted FIFO before the + // latest instruction. Remove it only after spawn succeeds; no later + // follow-up is needed to deliver it, and failed spawns retain it. + self.queued_mail.remove(&agent_id); self.resume_targets.insert(agent_id, resumed.agent_id); Ok(receipt) } @@ -6949,13 +7030,18 @@ impl SubAgentManager { ) -> HashMap { self.pending_follow_ups .iter() - .filter_map(|(agent_id, counter)| { + .filter_map(|(agent_id, inbox)| { if active_session_id.is_some_and(|session_id| { !self.agent_id_is_owned_by_session(agent_id, session_id) }) { return None; } - let count = counter.load(std::sync::atomic::Ordering::Acquire); + let count = inbox + .lock() + .unwrap_or_else(|e| e.into_inner()) + .iter() + .filter(|mail| !mail.taken.load(std::sync::atomic::Ordering::Acquire)) + .count(); (count > 0).then(|| (agent_id.clone(), count)) }) .collect() @@ -7112,10 +7198,39 @@ impl SubAgentManager { } } + /// Test helper: drop a child's parked mailbox so a scenario can verify the + /// retirement gate tracks the live queue instead of latching on history. + #[cfg(test)] + pub fn clear_test_queued_mail(&mut self, agent_id: &str) { + self.queued_mail.remove(agent_id); + } + + /// Test helper: commit a terminal result through the production path + /// (`claim` → `finish_terminal_result`) so the counter and mailbox + /// bookkeeping behave exactly as a real completion would. + #[cfg(test)] + pub fn commit_test_terminal_result(&mut self, agent_id: &str, status: SubAgentStatus) { + let mut result = self.get_result(agent_id).expect("test agent"); + result.status = status; + assert!( + self.finish_terminal_result(agent_id, result, true, false), + "test terminal commit must claim the running child" + ); + } + + /// Test helper: move an already-terminal child to a dead-end status, + /// modeling a terminal commit that no resume can undo. + #[cfg(test)] + pub fn mark_test_agent_failed(&mut self, agent_id: &str) { + let agent = self.agents.get_mut(agent_id).expect("test agent"); + agent.status = SubAgentStatus::Failed("test dead end".to_string()); + agent.input_tx = None; + } + /// Test helper exposing the receiving side so delivery and provenance can /// be verified rather than inferred from a still-present sender handle. #[cfg(test)] - fn insert_test_running_agent_with_input( + pub(crate) fn insert_test_running_agent_with_input( &mut self, name: &str, workspace: &Path, @@ -7232,6 +7347,57 @@ impl SubAgentManager { (agent_id, handle) } + /// Conservative retirement check, unlike UI/capacity counts: a timed-out + /// heartbeat does not establish that its task has stopped. Do not claim or + /// discard terminal delivery here; the actor checks its completion queue too. + /// Mail parked on non-running children (queue-only receipts) is accepted + /// work as well when a delivery path remains: it lives only in the + /// in-memory mailbox, so retirement must refuse until the owner session + /// drains it. Mail on a dead-end child (Completed/Failed/Cancelled/ + /// BudgetExhausted) has no such path — `followup` refuses those statuses — + /// so it is undeliverable bookkeeping, not accepted work, and must not + /// refuse retirement forever. + pub(crate) fn has_live_work_for_session(&self, active_session_id: &str) -> bool { + self.agents.values().any(|agent| { + self.agent_is_owned_by_session(agent, active_session_id) + && (agent.status == SubAgentStatus::Running + || agent + .task_handle + .as_ref() + .is_some_and(|task| !task.is_finished())) + }) || self.has_queued_follow_ups_for_session(active_session_id) + || self.queued_mail.iter().any(|(agent_id, queue)| { + !queue.is_empty() + && self.agent_id_is_owned_by_session(agent_id, active_session_id) + && self + .agents + .get(agent_id) + .is_some_and(|agent| self.can_still_deliver_mail(agent)) + }) + } + + /// Whether session-owned input still needs consumption or checkpoint + /// confirmation; channel receipt alone does not finish a recovery handoff. + pub(crate) fn has_queued_follow_ups_for_session(&self, active_session_id: &str) -> bool { + self.pending_follow_ups.iter().any(|(agent_id, inbox)| { + inbox + .try_lock() + .map(|mail| !mail.is_empty()) + .unwrap_or(true) + && self.agent_id_is_owned_by_session(agent_id, active_session_id) + }) + } + + /// Whether mail parked in the in-memory mailbox can still reach a child: + /// a running child drains it directly, and an interrupted child drains it + /// after a resume from a continuable checkpoint. A dead-end status has no + /// remaining delivery path. + fn can_still_deliver_mail(&self, agent: &SubAgent) -> bool { + agent.status == SubAgentStatus::Running + || matches!(&agent.status, SubAgentStatus::Interrupted(_) + if agent.checkpoint.as_ref().is_some_and(|cp| cp.continuable && !cp.messages.is_empty())) + } + /// Count running agents. pub fn running_count(&self) -> usize { self.admitted_count() @@ -8460,6 +8626,30 @@ impl SubAgentManager { agent.needs_input = result.needs_input.clone(); if result.status != SubAgentStatus::Running { agent.input_tx = None; + if let Some(inbox) = self.pending_follow_ups.remove(agent_id) { + let mut mail = inbox.lock().unwrap_or_else(|e| e.into_inner()); + if matches!(result.status, SubAgentStatus::Interrupted(_)) + && agent + .checkpoint + .as_ref() + .is_some_and(|cp| cp.continuable && !cp.messages.is_empty()) + { + // These inputs were accepted but not committed into the + // stored checkpoint. They precede any newer parked mail. + let mut recovered: VecDeque<_> = + mail.drain(..).map(|ticket| ticket.mail.clone()).collect(); + if !recovered.is_empty() { + recovered.extend(self.queued_mail.remove(agent_id).unwrap_or_default()); + self.queued_mail.insert(agent_id.to_owned(), recovered); + } + } else if !mail.is_empty() { + tracing::warn!( + target: "subagent", agent_id = %agent_id, + uncheckpointed = mail.len(), + "sub-agent went terminal without a resume path for accepted follow-ups" + ); + } + } } agent.completion_claimed = false; agent.task_handle = None; @@ -8476,9 +8666,27 @@ impl SubAgentManager { let Some(agent) = self.agents.get_mut(agent_id) else { return false; }; + if agent.status != SubAgentStatus::Running { + return false; + } agent.steps_taken = checkpoint.steps_taken; agent.checkpoint = Some(checkpoint); agent.last_activity_at = Instant::now(); + // This write lock serializes checkpoint acceptance with interrupt. + // A received input still in pending_inputs is deliberately retained. + if let Some(inbox) = self.pending_follow_ups.get(agent_id) { + let checkpoint = agent.checkpoint.as_ref().expect("checkpoint assigned"); + inbox + .lock() + .unwrap_or_else(|e| e.into_inner()) + .retain(|mail| { + !mail.appended.load(std::sync::atomic::Ordering::Acquire) + || (!mail.mail.text.trim().is_empty() + && mail + .retained_index(&checkpoint.messages, checkpoint.omitted_messages) + .is_none()) + }); + } // #freeze: hot per-step path — coalesce the full-fleet persist so 20 // agents stepping concurrently do not serialize the whole fleet (with // full transcripts) to disk under the write lock on every step. @@ -12123,9 +12331,25 @@ async fn checkpoint_subagent_progress( steps_taken: u32, continuable: bool, ) -> SubAgentCheckpoint { - let checkpoint = + let mut checkpoint = build_subagent_checkpoint(agent_id, reason, messages, steps_taken, continuable); let mut manager = runtime.manager.write().await; + // Projection preserves message positions. Protect only the newly accepted + // suffix; otherwise the normal history budget could omit an input just + // before the model sees it, yet falsely acknowledge its recovery ticket. + let preserve_from = manager.pending_follow_ups.get(agent_id).and_then(|inbox| { + inbox + .lock() + .unwrap_or_else(|e| e.into_inner()) + .iter() + .filter_map(|mail| mail.retained_index(messages, 0)) + .min() + }); + if let Some(start) = preserve_from.filter(|start| *start < checkpoint.omitted_messages) { + checkpoint.messages = + crate::image_attach::safe_tool_result_message_projection(&messages[start..]); + checkpoint.omitted_messages = start; + } manager.update_checkpoint(agent_id, checkpoint.clone()); checkpoint } @@ -12190,8 +12414,25 @@ fn checkpoint_messages_to_text(messages: &[Message]) -> String { fn build_resume_prompt( original_prompt: &str, checkpoint: &SubAgentCheckpoint, + parked_mail: Option<&VecDeque>, followup_text: &str, ) -> String { + // A checkpoint can already contain the complete prior resume prompt. + // Do not duplicate that nested input again on every incarnation. + let original_prompt = if checkpoint.messages.iter().any(|message| { + message.role == Role::User && message.content.iter().any(|block| { + matches!(block, ContentBlock::Text { text, .. } if text.contains(original_prompt)) + }) + }) { + "" + } else { + original_prompt + }; + let parked = parked_mail + .into_iter() + .flatten() + .map(|mail| format!("Parent follow-up: {}\n\n", mail.text)) + .collect::(); format!( "{original_prompt}\n\n\ [RESUMED SESSION — checkpoint {}]\n\ @@ -12199,7 +12440,7 @@ fn build_resume_prompt( continue the work from where it left off. The parent follow-up below is the latest \ instruction.\n\n\ --- prior conversation (tail) ---\n{}\n--- end prior conversation ---\n\n\ - Parent follow-up: {followup_text}", + {parked}Parent follow-up: {followup_text}", checkpoint.checkpoint_id, checkpoint.steps_taken, checkpoint_messages_to_text(&checkpoint.messages), diff --git a/crates/tui/src/tools/subagent/tests.rs b/crates/tui/src/tools/subagent/tests.rs index 6c89e4085c..b134e0e782 100644 --- a/crates/tui/src/tools/subagent/tests.rs +++ b/crates/tui/src/tools/subagent/tests.rs @@ -13873,6 +13873,496 @@ async fn interrupted_child_stops_gating_later_writers() { ); } +/// Mail accepted while a child was parked must follow the resume: the +/// queue-only receipt promised the owner session the mail survives a resume, +/// but the resume re-dispatches under a NEW id and `followup` on the old id +/// can only park further mail on the retained record — it can never deliver. +/// The resume delivers parked mail in its initial input, before the latest +/// instruction; no additional followup is required to wake the mailbox. +#[tokio::test] +async fn resumed_child_inherits_the_parked_mail_of_its_interrupted_predecessor() { + let tmp = tempdir().expect("tempdir"); + let manager = new_shared_subagent_manager(tmp.path().to_path_buf(), 4); + let mut runtime = stub_runtime(); + runtime.manager = Arc::clone(&manager); + runtime.context = ToolContext::new(tmp.path()); + + let mut guard = manager.write().await; + let first = guard + .spawn_background_with_assignment_options( + Arc::clone(&manager), + runtime.clone(), + FleetRole::Builder, + "edit the workspace".into(), + make_assignment(), + Some(vec![]), + SubAgentSpawnOptions { + name: Some("writer-mail".into()), + ..Default::default() + }, + ) + .expect("writer admitted"); + // Give the child a continuable checkpoint with content, as a real + // interrupted loop leaves behind, so the resume has something to seed. + guard + .agents + .get_mut(&first.agent_id) + .expect("agent") + .checkpoint = Some(build_subagent_checkpoint( + &first.agent_id, + "parent_interrupt", + &[Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "prior work".to_string(), + cache_control: None, + }], + }], + 1, + true, + )); + guard + .interrupt_child(&first.agent_id, None, "interrupted by parent".into()) + .expect("interrupt"); + guard + .followup_child(&first.agent_id, "accepted while parked".to_string()) + .expect("queue-only follow-up receipt"); + guard + .followup_child(&first.agent_id, "second parked mail".into()) + .unwrap(); + assert_eq!( + guard.queued_mail_depth(&first.agent_id), + Some(2), + "the parked mail sits on the interrupted predecessor" + ); + + let resumed = guard + .resume_from_checkpoint_for_session( + "workspace", + Arc::clone(&manager), + runtime, + &first.agent_id, + "continue", + ) + .expect("resume re-dispatches the checkpoint"); + assert_ne!( + resumed.agent_id, first.agent_id, + "resume runs under a new id" + ); + assert_eq!( + guard.queued_mail_depth(&resumed.agent_id), + None, + "the accepted FIFO is already included in the successor's initial input" + ); + assert_eq!( + guard.queued_mail_depth(&first.agent_id), + None, + "the predecessor's mailbox is emptied by the migration" + ); + let prompt = &guard.agents[&resumed.agent_id].prompt; + assert!( + prompt.find("accepted while parked").unwrap() < prompt.find("second parked mail").unwrap() + ); + assert!( + prompt.find("second parked mail").unwrap() + < prompt.rfind("Parent follow-up: continue").unwrap() + ); + assert_eq!(resumed.queue_depth, 0); + assert!(resumed.woke && resumed.continued_from_checkpoint); +} + +#[tokio::test] +async fn forkguard_parent_mail_interrupt_preserves_uncheckpointed_fifo() { + for stage in ["unread", "taken", "appended", "checkpointed"] { + let tmp = tempdir().unwrap(); + let manager = new_shared_subagent_manager(tmp.path().into(), 4); + let mut runtime = stub_runtime(); + runtime.manager = manager.clone(); + runtime.context = ToolContext::new(tmp.path()); + let mut guard = manager.write().await; + let (child, mut rx) = guard.insert_test_running_agent_with_input(stage, tmp.path()); + let mut messages = vec![Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "earlier state".into(), + cache_control: None, + }], + }]; + assert!(guard.update_checkpoint( + &child, + build_subagent_checkpoint(&child, "before_api_request", &messages, 1, true) + )); + for text in ["FIFO-M1", "FIFO-M2"] { + let receipt = guard.followup_child(&child, text.into()).unwrap(); + assert!(receipt.woke); + assert_eq!(receipt.queue_depth, 0); + } + assert_eq!(guard.queued_follow_up_counts()[&child], 2); + let mut taken = VecDeque::new(); + if stage != "unread" { + let input = rx.try_recv().unwrap(); + input.mark_taken(); + taken.push_back(input); + assert_eq!(guard.queued_follow_up_counts()[&child], 1); + } + if matches!(stage, "appended" | "checkpointed") { + append_subagent_inputs_as_user_messages(&mut messages, &mut taken); + } + if stage == "checkpointed" { + assert!(guard.update_checkpoint( + &child, + build_subagent_checkpoint(&child, "before_api_request", &messages, 2, true) + )); + } + guard + .queue_running_parent_message(&child, "FIFO-M3".into()) + .unwrap(); + guard + .interrupt_child(&child, None, "controlled interrupt".into()) + .unwrap(); + drop(rx); + drop(taken); + let expected = if stage == "checkpointed" { + vec!["FIFO-M2", "FIFO-M3"] + } else { + vec!["FIFO-M1", "FIFO-M2", "FIFO-M3"] + }; + assert_eq!( + guard.queued_mail[&child] + .iter() + .map(|mail| mail.text.as_str()) + .collect::>(), + expected, + "{stage}" + ); + assert!(guard.queued_follow_up_counts().is_empty()); + let checkpoint = guard.agents[&child].checkpoint.clone().unwrap(); + assert!(!guard.update_checkpoint( + &child, + build_subagent_checkpoint(&child, "late worker", &[], 99, true) + )); + assert_eq!( + guard.agents[&child] + .checkpoint + .as_ref() + .unwrap() + .checkpoint_id, + checkpoint.checkpoint_id + ); + let resumed = guard + .resume_from_checkpoint_for_session( + "workspace", + manager.clone(), + runtime, + &child, + "FIFO-LATEST", + ) + .unwrap(); + let prompt = &guard.agents[&resumed.agent_id].prompt; + assert_eq!( + prompt.matches("FIFO-M1").count(), + 1, + "no loss or replay: {stage}" + ); + assert!(prompt.find("FIFO-M1").unwrap() < prompt.find("FIFO-M2").unwrap()); + assert!(prompt.find("FIFO-M2").unwrap() < prompt.find("FIFO-M3").unwrap()); + assert!(prompt.find("FIFO-M3").unwrap() < prompt.find("FIFO-LATEST").unwrap()); + assert_eq!(resumed.queue_depth, 0); + assert_eq!(guard.queued_mail_depth(&child), None); + assert_eq!(guard.queued_mail_depth(&resumed.agent_id), None); + } +} + +#[tokio::test] +async fn forkguard_parent_mail_resume_reaches_real_request_without_another_followup() { + let tmp = tempdir().unwrap(); + let (sent, mut received) = mpsc::unbounded_channel(); + let app = Router::new().route("/{*path}", post(move |Json(body): Json| { + let sent = sent.clone(); + async move { + sent.send(body).unwrap(); + Json(json!({"id":"chatcmpl-mail","model":"deepseek-v4-flash","choices":[{"index":0,"message":{"role":"assistant","content":"done"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}})) + } + })); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let shutdown = CancellationToken::new(); + let stop_guard = shutdown.clone().drop_guard(); + let stopping = shutdown.clone(); + let server = tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(stopping.cancelled_owned()) + .await + .unwrap(); + }); + let config = crate::config::Config { + api_key: Some("test-key".into()), + base_url: Some(format!("http://{addr}/v1")), + ..Default::default() + }; + let manager = new_shared_subagent_manager(tmp.path().into(), 4); + let mut runtime = stub_runtime(); + runtime.client = DeepSeekClient::new(&config).unwrap(); + runtime.api_config = Some(Arc::new(config)); + runtime.manager = manager.clone(); + runtime.context = ToolContext::new(tmp.path()); + let task = { + let mut guard = manager.write().await; + let (child, _) = guard.insert_test_interrupted_continuable_agent( + "real-request", + tmp.path(), + vec![Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "earlier checkpoint".into(), + cache_control: None, + }], + }], + ); + for text in ["REAL-MAIL-ONE", "REAL-MAIL-TWO"] { + guard.followup_child(&child, text.into()).unwrap(); + } + let resumed = guard + .resume_from_checkpoint_for_session( + "workspace", + manager.clone(), + runtime, + &child, + "REAL-LATEST", + ) + .unwrap(); + assert_eq!(resumed.queue_depth, 0); + guard + .agents + .get_mut(&resumed.agent_id) + .unwrap() + .task_handle + .take() + .unwrap() + }; + let body = tokio::time::timeout(Duration::from_secs(10), received.recv()) + .await + .unwrap() + .unwrap(); + let messages = body["messages"].to_string(); + assert_eq!(messages.matches("REAL-MAIL-ONE").count(), 1); + assert_eq!(messages.matches("REAL-MAIL-TWO").count(), 1); + assert!(messages.find("REAL-MAIL-ONE").unwrap() < messages.find("REAL-MAIL-TWO").unwrap()); + assert!(messages.find("REAL-MAIL-TWO").unwrap() < messages.find("REAL-LATEST").unwrap()); + tokio::time::timeout(Duration::from_secs(10), task) + .await + .unwrap() + .unwrap(); + drop(stop_guard); + tokio::time::timeout(Duration::from_secs(5), server) + .await + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn forkguard_parent_mail_checkpoint_keeps_fresh_batch_before_history_trimming() { + let tmp = tempdir().unwrap(); + let manager = new_shared_subagent_manager(tmp.path().into(), 4); + let mut runtime = stub_runtime(); + runtime.manager = manager.clone(); + runtime.context = ToolContext::new(tmp.path()); + let mut messages = vec![Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "earlier work".into(), + cache_control: None, + }], + }]; + let child = { + let mut guard = manager.write().await; + let (child, mut rx) = guard.insert_test_running_agent_with_input("large-mail", tmp.path()); + guard.update_checkpoint( + &child, + build_subagent_checkpoint(&child, "prior", &messages, 1, true), + ); + for text in [ + "BATCH-FIRST".to_owned(), + format!( + "BATCH-SECOND{}", + "x".repeat(SUBAGENT_CHECKPOINT_MESSAGE_BUDGET_BYTES + 1) + ), + ] { + guard.followup_child(&child, text).unwrap(); + } + let mut inputs = VecDeque::new(); + while let Ok(input) = rx.try_recv() { + input.mark_taken(); + inputs.push_back(input); + } + append_subagent_inputs_as_user_messages(&mut messages, &mut inputs); + child + }; + let cp = + checkpoint_subagent_progress(&runtime, &child, "before_api_request", &messages, 2, true) + .await; + assert_eq!(cp.omitted_messages, 1); + assert_eq!( + cp.messages.len(), + 2, + "the first input must not disappear behind the oversized last input" + ); + assert_eq!(message_text(&cp.messages[0]), "BATCH-FIRST"); + assert!(message_text(&cp.messages[1]).starts_with("BATCH-SECOND")); + { + let mut guard = manager.write().await; + assert!( + guard.pending_follow_ups[&child].lock().unwrap().is_empty(), + "both inputs are actually recoverable in this checkpoint" + ); + guard + .interrupt_child(&child, None, "before API".into()) + .unwrap(); + let resumed = guard + .resume_from_checkpoint_for_session( + "workspace", + manager.clone(), + runtime.clone(), + &child, + "BATCH-LATEST", + ) + .unwrap(); + let prompt = &guard.agents[&resumed.agent_id].prompt; + assert_eq!(prompt.matches("BATCH-FIRST").count(), 1); + assert_eq!(prompt.matches("BATCH-SECOND").count(), 1); + } + // No pending batch: preserve #3882's ordinary bounded history behavior. + messages.push(Message { + role: Role::Assistant, + content: vec![ContentBlock::Text { + text: "processed".into(), + cache_control: None, + }], + }); + let other = manager + .write() + .await + .insert_test_running_agent("ordinary-history", tmp.path()); + let ordinary = + checkpoint_subagent_progress(&runtime, &other, "after_model_response", &messages, 3, true) + .await; + assert_eq!(ordinary.messages.len(), 1); + assert_eq!(message_text(&ordinary.messages[0]), "processed"); +} + +#[tokio::test] +async fn forkguard_parent_mail_resume_can_be_interrupted_before_its_first_poll() { + let tmp = tempdir().unwrap(); + let manager = new_shared_subagent_manager(tmp.path().into(), 4); + let mut runtime = stub_runtime(); + runtime.manager = manager.clone(); + runtime.context = ToolContext::new(tmp.path()); + let mut guard = manager.write().await; + let (child, _) = guard.insert_test_interrupted_continuable_agent( + "first-poll", + tmp.path(), + vec![Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "original context".into(), + cache_control: None, + }], + }], + ); + for text in ["EARLY-ONE", "EARLY-TWO"] { + guard.followup_child(&child, text.into()).unwrap(); + } + let first = guard + .resume_from_checkpoint_for_session( + "workspace", + manager.clone(), + runtime.clone(), + &child, + "EARLY-LATEST-1", + ) + .unwrap(); + // The manager write lock keeps the worker from constructing its first + // checkpoint. Interrupt through the real production entry in this gap. + assert_eq!( + guard.agents[&first.agent_id] + .checkpoint + .as_ref() + .unwrap() + .reason, + "resume_admitted" + ); + guard + .interrupt_child(&first.agent_id, None, "before first poll".into()) + .unwrap(); + let second = guard + .resume_from_checkpoint_for_session( + "workspace", + manager.clone(), + runtime, + &first.agent_id, + "EARLY-LATEST-2", + ) + .unwrap(); + let prompt = &guard.agents[&second.agent_id].prompt; + assert_eq!(prompt.matches("EARLY-ONE").count(), 1); + assert_eq!(prompt.matches("EARLY-TWO").count(), 1); + assert!(prompt.find("EARLY-ONE").unwrap() < prompt.find("EARLY-TWO").unwrap()); + assert!(prompt.find("EARLY-LATEST-1").unwrap() < prompt.find("EARLY-LATEST-2").unwrap()); + assert_eq!(second.queue_depth, 0); +} + +#[tokio::test] +async fn forkguard_parent_mail_failed_resume_keeps_fifo_for_the_next_attempt() { + let tmp = tempdir().unwrap(); + let manager = new_shared_subagent_manager(tmp.path().into(), 4); + let mut runtime = stub_runtime(); + runtime.manager = manager.clone(); + runtime.context = ToolContext::new(tmp.path()); + let mut guard = manager.write().await; + let (child, _) = guard.insert_test_interrupted_continuable_agent( + "retry-mail", + tmp.path(), + vec![Message { + role: Role::User, + content: vec![ContentBlock::Text { + text: "original context".into(), + cache_control: None, + }], + }], + ); + for text in ["RETRY-ONE", "RETRY-TWO"] { + guard.followup_child(&child, text.into()).unwrap(); + } + runtime.max_spawn_depth = 0; + assert!( + guard + .resume_from_checkpoint_for_session( + "workspace", + manager.clone(), + runtime.clone(), + &child, + "failed instruction" + ) + .is_err() + ); + assert_eq!(guard.queued_mail_depth(&child), Some(2)); + assert!(!guard.resume_targets.contains_key(&child)); + runtime.max_spawn_depth = DEFAULT_MAX_SPAWN_DEPTH; + let receipt = guard + .resume_from_checkpoint_for_session( + "workspace", + manager.clone(), + runtime, + &child, + "RETRY-LATEST", + ) + .unwrap(); + let prompt = &guard.agents[&receipt.agent_id].prompt; + assert!(prompt.find("RETRY-ONE").unwrap() < prompt.find("RETRY-TWO").unwrap()); + assert!(prompt.find("RETRY-TWO").unwrap() < prompt.find("RETRY-LATEST").unwrap()); + assert_eq!(guard.queued_mail_depth(&child), None); +} + /// The resume path is the same scenario from the other side: `followup` on /// an interrupted child re-dispatches its checkpoint under a NEW id and /// restores the old claim for that id. With the predecessor counted live, diff --git a/crates/tui/src/tools/subagent/tests/identity_selection.rs b/crates/tui/src/tools/subagent/tests/identity_selection.rs index 6cd70a6653..7b62bb338d 100644 --- a/crates/tui/src/tools/subagent/tests/identity_selection.rs +++ b/crates/tui/src/tools/subagent/tests/identity_selection.rs @@ -526,15 +526,34 @@ async fn forkguard_agent_identity_checkpoint_retry_keeps_old_receipt_and_resumed assert_eq!(delivery["delivered"], true); assert_eq!(delivery["message_preview"], "healthy next"); assert!(delivery["reason"].is_null()); + // A host without an attached resume runtime reaches followup_child + // directly. The retained predecessor must still route to the successor. { let mut guard = manager.write().await; + let receipt = guard + .followup_child_for_session("workspace", &a, "direct late follow-up".into()) + .unwrap(); + assert_eq!(receipt.agent_id, target); + assert!(receipt.woke); + assert_eq!(receipt.queue_depth, 0); + assert!(guard.queued_mail_depth(&a).is_none()); + } + assert_eq!(rx.try_recv().unwrap().text, "direct late follow-up"); + let task = { + let mut guard = manager.write().await; + let task = guard.agents.get_mut(&target).unwrap().task_handle.take(); let mut terminal = guard.get_result(&target).unwrap(); terminal.status = SubAgentStatus::Completed; assert!(guard.finish_terminal_result(&target, terminal, true, false)); - } + task + }; let rejected = json!({"action":"followup","agent_id":a,"message":"must not disappear"}); assert!(public.execute(rejected.clone(), &context).await.is_err()); assert!(hidden.execute(rejected, &context).await.is_err()); + if let Some(task) = task { + task.abort(); + let _ = task.await; + } let mut guard = manager.write().await; assert!( guard @@ -557,6 +576,26 @@ async fn forkguard_agent_identity_checkpoint_retry_keeps_old_receipt_and_resumed SubAgentStatus::Running ); assert_eq!(guard.queued_mail_depth(&c), None); + // The identity-collision controls remain running through all assertions + // above. Settle them too before asking whether the whole session is idle. + let mut tasks = Vec::new(); + for id in [&b, &c] { + if let Some(task) = guard.agents.get_mut(id).unwrap().task_handle.take() { + tasks.push(task); + } + let mut terminal = guard.get_result(id).unwrap(); + terminal.status = SubAgentStatus::Completed; + assert!(guard.finish_terminal_result(id, terminal, true, false)); + } + drop(guard); + for task in tasks { + task.abort(); + let _ = task.await; + } + assert!( + !manager.read().await.has_live_work_for_session("workspace"), + "delivered old-id mail must not strand retirement after all children settle" + ); } #[tokio::test]