diff --git a/crates/tui/src/core/engine.rs b/crates/tui/src/core/engine.rs index 6fd398d6d7..e330ff1480 100644 --- a/crates/tui/src/core/engine.rs +++ b/crates/tui/src/core/engine.rs @@ -579,6 +579,32 @@ pub enum CancelMode { StopDropInbox, } +/// Turn-scoped cancellation slot shared between the engine and its hosts. +/// +/// The engine installs a fresh token under the starting turn's identity at +/// every turn start, atomically replacing the previous entry under one lock. +/// Hosts that cancel through [`EngineHandle::cancel_turn`] resolve the token +/// under that same lock only while the slot still names the target turn, so +/// a host whose turn view is stale — the engine swapped tokens for a +/// self-started follow-up turn (idle sub-agent completion, background shell +/// wake, goal continuation) before the host observed it — can never fire the +/// newer turn's token. +#[derive(Clone, Debug)] +pub struct TurnCancelSlot { + /// Identity of the turn owning [`Self::token`] — the engine-minted + /// `TurnContext::id`, the same value carried by `Event::TurnStarted`. + /// `None` before the first turn and after an engine-side reset that is + /// not bound to a turn (`Op::CancelRequest`): nothing cancellable is + /// bound to a named turn, so every turn-bound cancel must skip. + pub turn_id: Option, + pub token: CancellationToken, +} + +/// Lock-protected slot held jointly by the engine and every +/// [`EngineHandle`] clone. The engine swaps the whole entry at turn start; +/// hosts read-verify-clone under the same lock to cancel one exact turn. +pub type SharedCancelToken = Arc>; + /// Outcome of withdrawing a queued steer by id. Hosts that re-send the same /// input through another path (e.g. interrupt-and-send) need to know whether /// the engine copy can still be committed, otherwise the same message may be @@ -814,8 +840,11 @@ pub struct EngineHandle { pub tx_op: mpsc::Sender, /// Receive events from the engine pub rx_event: Arc>>, - /// Shared pointer to the cancellation token for the current request. - cancel_token: Arc>, + /// Shared pointer to the cancellation slot for the current request. The + /// engine swaps the whole slot (identity + token) at every turn start; + /// `cancel_with_mode` cancels whatever token currently occupies it, + /// `cancel_turn` cancels only the named turn's token. + cancel_token: SharedCancelToken, /// Latched reason for the most recent cancellation. Read by the /// approval / user-input handlers to enrich their error strings. /// Cleared by the engine when a fresh turn starts. @@ -934,7 +963,7 @@ pub struct Engine { /// delivery. delivered_subagent_completion_ids: HashSet, cancel_token: CancellationToken, - shared_cancel_token: Arc>, + shared_cancel_token: SharedCancelToken, /// Latched reason for the current cancellation, mirrored to /// `EngineHandle::cancel_reason`. Read by `approval.rs` when /// surfacing the "Request cancelled while awaiting …" error so the @@ -1280,15 +1309,37 @@ impl Engine { .await; } + /// Install a fresh cancellation token with no turn identity: used by the + /// engine-side `Op::CancelRequest` reset, where the cancelled turn is over + /// and the follow-up token belongs to no named turn yet. Turn starts must + /// use [`Self::install_turn_cancel_token`] so hosts can bind their cancels. fn reset_cancel_token(&mut self) { + self.install_cancel_slot(None); + } + + /// Install this turn's cancellation token bound to `turn_id` as one + /// atomic slot swap. Hosts cancel through the turn-bound slot, so the + /// token swap and the identity swap must be a single step: a cancel that + /// resolves the token under the slot lock either sees the previous turn's + /// identity and token (cancels them — correct) or this turn's (skips — + /// the newer turn is not the cancel's target). + fn install_turn_cancel_token(&mut self, turn_id: &str) { + self.install_cancel_slot(Some(turn_id.to_string())); + } + + fn install_cancel_slot(&mut self, turn_id: Option) { let token = CancellationToken::new(); self.cancel_token = token.clone(); + let slot = TurnCancelSlot { + turn_id, + token: token.clone(), + }; match self.shared_cancel_token.lock() { Ok(mut shared) => { - *shared = token; + *shared = slot; } Err(poisoned) => { - *poisoned.into_inner() = token; + *poisoned.into_inner() = slot; } } // Fresh turn → clear any latched cancellation reason from the @@ -1457,7 +1508,10 @@ impl Engine { let (tx_steer, rx_steer) = mpsc::channel(64); let (tx_subagent_completion, rx_subagent_completion) = mpsc::unbounded_channel(); let cancel_token = CancellationToken::new(); - let shared_cancel_token = Arc::new(StdMutex::new(cancel_token.clone())); + let shared_cancel_token: SharedCancelToken = Arc::new(StdMutex::new(TurnCancelSlot { + turn_id: None, + token: cancel_token.clone(), + })); let steer_control = Arc::new(StdMutex::new(SteerControlState::default())); let cancel_reason: Arc>> = Arc::new(StdMutex::new(None)); let shared_paused = Arc::new(StdMutex::new(false)); @@ -1713,7 +1767,6 @@ impl Engine { auto_approve: bool, approval_mode: crate::tui::approval::ApprovalMode, ) { - self.reset_cancel_token(); self.turn_counter = self.turn_counter.saturating_add(1); let turn_id = format!( @@ -1721,6 +1774,10 @@ impl Engine { USER_SHELL_TOOL_ID_PREFIX, seq = self.turn_counter ); + // Bind the fresh cancellation token to this turn's identity before + // anything else observes the turn (same turn-bound slot contract as + // `handle_send_message`). + self.install_turn_cancel_token(&turn_id); let tool_id = turn_id.clone(); let tool_name = "Bash".to_string(); let tool_input = json!({ "action": "run", "command": command, "source": "user" }); @@ -4573,8 +4630,18 @@ impl Engine { if let Some(status) = input_policy.status() { let _ = self.tx_event.send(Event::status(status)).await; } - // Reset cancel token for fresh turn (in case previous was cancelled) - self.reset_cancel_token(); + // Create turn context first so the fresh cancellation token can be + // installed under this turn's identity and the start event includes a + // stable turn id. The token swap and the identity swap are one atomic + // slot step: hosts cancel through the turn-bound slot, so a cancel + // racing this install either hits the previous turn's token (its + // target) or skips this turn entirely — it can never fire this new + // token while believing it targets the previous turn (Pinvou + // pinvou-agent#254: runtime self-started turns swapped the shared + // token before the host observed `TurnStarted`, so a stale + // generation-matched cancel killed the follow-up turn). + let mut turn = TurnContext::new(self.config.max_steps); + self.install_turn_cancel_token(&turn.id); // Track the complete effective mode policy so mid-turn metadata, `/edit`, // idle worker resumptions, and approval gates cannot read a stale policy @@ -4590,8 +4657,6 @@ impl Engine { // turns". Steers collected into `pending_steers` by an interrupted // turn were already parked back by the turn loop. - // Create turn context first so start event includes a stable turn id. - let mut turn = TurnContext::new(self.config.max_steps); // Publish the active steer destination before `TurnStarted`. Hosts // may steer or cancel as soon as they observe that event, and both // operations must resolve against this exact generation. @@ -6522,7 +6587,10 @@ pub(crate) fn mock_engine_handle() -> MockEngineHandle { let (tx_user_input, rx_user_input) = mpsc::channel(32); let (tx_steer, rx_steer) = mpsc::channel(64); let cancel_token = CancellationToken::new(); - let shared_cancel_token = Arc::new(StdMutex::new(cancel_token.clone())); + let shared_cancel_token: SharedCancelToken = Arc::new(StdMutex::new(TurnCancelSlot { + turn_id: None, + token: cancel_token.clone(), + })); let cancel_reason: Arc>> = Arc::new(StdMutex::new(None)); let shared_paused = Arc::new(StdMutex::new(false)); let live_runtime_authority = Arc::new(StdMutex::new(LiveRuntimeAuthorityState::new( diff --git a/crates/tui/src/core/engine/handle.rs b/crates/tui/src/core/engine/handle.rs index b186357274..c46873c01c 100644 --- a/crates/tui/src/core/engine/handle.rs +++ b/crates/tui/src/core/engine/handle.rs @@ -172,6 +172,54 @@ impl EngineHandle { /// A stop barrier is visible before the token fires, so concurrent or /// already-reserved sends cannot escape into a later turn. pub fn cancel_with_mode(&self, reason: CancelReason, mode: CancelMode) { + self.publish_cancel_disposition(reason, mode); + match self.cancel_token.lock() { + Ok(slot) => slot.token.cancel(), + Err(poisoned) => poisoned.into_inner().token.cancel(), + } + crate::retry_status::clear(); + } + + /// Cancel exactly the turn named by `turn_id`, and only while that turn + /// still holds the shared cancellation slot. + /// + /// Returns `true` when the slot named `turn_id` and its token fired; + /// `false` when the slot already moved on — the engine swapped tokens for + /// a newer turn (a host-driven send, or a runtime self-start such as an + /// idle sub-agent completion, a background shell wake, or a goal + /// continuation) before this cancel was observed. In the `false` case + /// nothing is cancelled and no steer disposition or cancel reason is + /// published: the target turn is already gone, and the newer turn belongs + /// to a different generation the caller never saw. + /// + /// The identity check and the token resolution happen under the same slot + /// lock the engine's turn-start install uses, so a concurrent turn start + /// can neither race the decision nor get its token cancelled: the token + /// is cloned out while locked, and firing it afterwards cancels only that + /// cloned token — even if the slot moves on in between, the fired object + /// is still the target turn's own token. + #[must_use = "a false return means the target turn already ended; do not treat the cancel as delivered"] + pub fn cancel_turn(&self, turn_id: &str, reason: CancelReason, mode: CancelMode) -> bool { + let token = { + let slot = self + .cancel_token + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if slot.turn_id.as_deref() != Some(turn_id) { + return false; + } + slot.token.clone() + }; + self.publish_cancel_disposition(reason, mode); + token.cancel(); + crate::retry_status::clear(); + true + } + + /// Publish the steer disposition and latch the cancel reason shared by + /// every cancel entry point. The token fire itself stays with the caller + /// so turn-bound cancels can resolve the exact token under the slot lock. + fn publish_cancel_disposition(&self, reason: CancelReason, mode: CancelMode) { self.steer_control .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -180,11 +228,6 @@ impl EngineHandle { Ok(mut slot) => *slot = Some(reason), Err(poisoned) => *poisoned.into_inner() = Some(reason), } - match self.cancel_token.lock() { - Ok(token) => token.cancel(), - Err(poisoned) => poisoned.into_inner().cancel(), - } - crate::retry_status::clear(); } /// Check if a request is currently cancelled @@ -192,8 +235,8 @@ impl EngineHandle { #[allow(dead_code)] pub fn is_cancelled(&self) -> bool { match self.cancel_token.lock() { - Ok(token) => token.is_cancelled(), - Err(poisoned) => poisoned.into_inner().is_cancelled(), + Ok(slot) => slot.token.is_cancelled(), + Err(poisoned) => poisoned.into_inner().token.is_cancelled(), } } diff --git a/crates/tui/src/core/engine/tests.rs b/crates/tui/src/core/engine/tests.rs index d634318bed..e13a6c01b5 100644 --- a/crates/tui/src/core/engine/tests.rs +++ b/crates/tui/src/core/engine/tests.rs @@ -6928,6 +6928,113 @@ fn engine_handle_cancel_tracks_latest_turn_token() { assert!(!stale_token.is_cancelled()); } +#[test] +fn engine_handle_cancel_turn_only_fires_the_named_turns_token() { + let (mut engine, handle) = Engine::new(EngineConfig::default(), &Config::default()); + // No turn installed yet: the slot carries no identity, so every + // turn-bound cancel must skip instead of firing the anonymous token. + assert!(!handle.cancel_turn("turn-1", CancelReason::User, CancelMode::StopDropInbox)); + assert!(!handle.is_cancelled()); + + engine.install_turn_cancel_token("turn-1"); + assert!(handle.cancel_turn("turn-1", CancelReason::User, CancelMode::StopDropInbox)); + assert!(engine.cancel_token.is_cancelled()); + + // A follow-up turn swaps the slot — exactly what a runtime self-started + // continuation (idle sub-agent completion, background shell wake, goal + // continuation) does before the host observes its `TurnStarted`. + engine.install_turn_cancel_token("turn-2"); + let followup_token = engine.cancel_token.clone(); + // A stale cancel still targeting turn-1 must skip: it may not fire the + // follow-up's token, and it must not latch a steer disposition for it. + assert!(!handle.cancel_turn("turn-1", CancelReason::User, CancelMode::StopDropInbox)); + assert!( + !followup_token.is_cancelled(), + "stale turn-bound cancel must not fire the follow-up turn's token" + ); + // The follow-up's own cancel still lands on its token. + assert!(handle.cancel_turn("turn-2", CancelReason::User, CancelMode::InterruptKeepInbox)); + assert!(followup_token.is_cancelled()); +} + +#[tokio::test] +async fn forkguard_cancel_turn_binding_spares_unnamed_turns_and_hits_the_observed_turn() { + let workspace = tempdir().expect("tempdir"); + let entered = std::sync::Arc::new(tokio::sync::Notify::new()); + let request_dropped = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let client: crate::core::model_client::SharedModelClient = + std::sync::Arc::new(BlockingModelClient { + entered: std::sync::Arc::clone(&entered), + request_dropped: std::sync::Arc::clone(&request_dropped), + }); + let (engine, handle) = Engine::new_with_model_client( + deterministic_engine_config(workspace.path()), + &Config::default(), + client, + ); + let task = tokio::spawn(engine.run()); + handle + .send(external_user_message_op( + "Cancel must bind to the observed turn id.", + AppMode::Agent, + &Config::default(), + )) + .await + .expect("send turn"); + let turn_id = { + let mut rx = handle.rx_event.write().await; + loop { + let event = tokio::time::timeout(model_turn_event_timeout(), rx.recv()) + .await + .expect("timed out waiting for TurnStarted") + .expect("engine event"); + if let Event::TurnStarted { turn_id, .. } = event { + break turn_id; + } + } + }; + tokio::time::timeout(model_turn_event_timeout(), entered.notified()) + .await + .expect("model request was never entered"); + + // A host cancel bound to a turn identity the engine never installed + // (stale target, or a follow-up turn the host has not observed yet — + // pinvou-agent#254) must not fire the running turn's token and must not + // disturb the in-flight provider request. + assert!(!handle.cancel_turn( + "turn-that-never-ran", + CancelReason::User, + CancelMode::StopDropInbox, + )); + assert!( + !handle.is_cancelled(), + "turn-bound cancel with a foreign turn id fired the live token" + ); + assert!(!request_dropped.load(std::sync::atomic::Ordering::SeqCst)); + + // Cancelling the turn id observed from `TurnStarted` lands on exactly + // that turn: same interrupted terminal, same dropped provider future as + // the mode-less cancel path. + assert!(handle.cancel_turn(&turn_id, CancelReason::User, CancelMode::StopDropInbox,)); + let mut rx = handle.rx_event.write().await; + while let Some(event) = tokio::time::timeout(model_turn_event_timeout(), rx.recv()) + .await + .expect("timed out waiting for cancellation") + { + if let Event::TurnComplete { status, error, .. } = event { + assert_eq!(status, TurnOutcomeStatus::Interrupted, "{error:?}"); + break; + } + } + drop(rx); + assert!( + request_dropped.load(std::sync::atomic::Ordering::SeqCst), + "turn-bound cancellation must drop the active provider future" + ); + handle.send(Op::Shutdown).await.expect("shutdown engine"); + task.await.expect("engine task"); +} + #[test] fn engine_initial_prompt_includes_configured_goal() { let config = EngineConfig { @@ -17148,7 +17255,10 @@ fn engine_handle_try_send_does_not_block_when_op_channel_is_full() { let handle = EngineHandle { tx_op, rx_event: Arc::new(RwLock::new(mpsc::channel::(1).1)), - cancel_token: Arc::new(StdMutex::new(cancel_token)), + cancel_token: Arc::new(StdMutex::new(super::TurnCancelSlot { + turn_id: None, + token: cancel_token, + })), cancel_reason: Arc::new(StdMutex::new(None)), tx_approval: mpsc::channel(1).0, tx_user_input: mpsc::channel(1).0,