From 2239d54a02a43d605a77b043fa34cb3c0affeae2 Mon Sep 17 00:00:00 2001 From: asto Date: Wed, 2 Sep 2026 18:56:01 +0800 Subject: [PATCH] fix(engine): bind shared cancel slot to the owning turn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Hosts cancel through a shared token slot that used to carry no turn identity: `cancel_with_mode` fired whatever token occupied the slot at call time. For host-driven turns the app-side generation gate made that safe, because the host reserves the next turn (advancing its epoch under the same lock the cancel checks) before the engine can start one. The three runtime self-start paths — idle sub-agent completion, background shell completion wake, and goal continuation — call handle_send_message inside the engine and swap the shared token before any host reserve, so a host cancel whose generation view was still on the finished turn passed the epoch check and fired the follow-up turn's token (Pinvou pinvou-agent#254). The slot is now a TurnCancelSlot { turn_id, token } swapped atomically at every turn start (handle_send_message and the user shell-command turn mint the turn id first, then install). Hosts get EngineHandle::cancel_turn(turn_id, reason, mode): the identity check and token clone happen under the same slot lock the install uses, so the decision is atomic against a concurrent turn start and the cloned token only ever fires the named turn. On identity mismatch nothing is cancelled and no steer disposition or cancel reason is published — the target turn is already gone. cancel_with_mode keeps its exact fire-current-token semantics for single-user frontends. Tests: a slot-contract unit test (unnamed slot skips, stale id skips without firing the follow-up, matching id fires) and a forkguard regression driving a real turn through an injected blocking client, asserting a foreign turn id leaves the in-flight request untouched while the observed TurnStarted id interrupts it. --- crates/tui/src/core/engine.rs | 92 +++++++++++++++++++--- crates/tui/src/core/engine/handle.rs | 57 ++++++++++++-- crates/tui/src/core/engine/tests.rs | 112 ++++++++++++++++++++++++++- 3 files changed, 241 insertions(+), 20 deletions(-) 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,