From 6a4bff0d453a48d581be7c134e863e185f085474 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Thu, 27 Aug 2026 09:53:09 -0400 Subject: [PATCH] fix(agent): a racing cancel can no longer wedge the session's turn slot Port of upstream grok-build 0.2.103. session/prompt runs an async intake preamble before enqueueing SessionCommand::Prompt on the actor mailbox; a concurrent cancel could enqueue Cancel first, land on an idle actor as a no-op, and the prompt then started a turn nothing would ever cancel. The per-session prompt-intake lock becomes the dispatch lock: prompt() holds it across intake as before, and cancel() now holds it around its Cancel send, so a cancel orders behind the prompt it targets. Regression test ported verbatim (unix-only harness, like its siblings). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01KSy9S31pZNWXtRZNVJ1sDf --- .../src/agent/mvp_agent/acp_agent.rs | 11 ++-- .../src/agent/mvp_agent/agent_ops.rs | 2 +- .../kigi-shell/src/agent/mvp_agent/mod.rs | 23 ++++---- .../src/agent/mvp_agent/session_lifecycle.rs | 13 ++--- .../kigi-shell/src/agent/mvp_agent/tests.rs | 53 +++++++++++++++++++ 5 files changed, 77 insertions(+), 25 deletions(-) diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs index 1c7dd16..b0c7cd0 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs @@ -1725,8 +1725,8 @@ impl acp::Agent for MvpAgent { return Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)); } } - let intake_lock = self.prompt_intake_lock(&arguments.session_id); - let intake_guard = intake_lock.lock().await; + let dispatch_lock = self.dispatch_lock(&arguments.session_id); + let dispatch_guard = dispatch_lock.lock().await; let meta_prompt_mode = arguments .meta .as_ref() @@ -1830,7 +1830,7 @@ impl acp::Agent for MvpAgent { acp::Error::internal_error() .data(format!("failed to dispatch prompt to session: {e}")) })?; - drop(intake_guard); + drop(dispatch_guard); self.push_roster_activity_delta( &arguments.session_id, crate::agent::roster::RosterActivity::Working, @@ -2025,6 +2025,11 @@ impl acp::Agent for MvpAgent { .and_then(|m| m.get("rewindIfPristine")) .and_then(|v| v.as_bool()) .unwrap_or(false); + // Ordered behind any in-flight prompt intake: a cancel that + // overtakes the prompt it targets lands on an idle actor, and the + // prompt then starts a turn nothing will ever cancel. + let dispatch_lock = self.dispatch_lock(&args.session_id); + let _dispatch_guard = dispatch_lock.lock().await; let _ = handle .cmd_tx .send(SessionCommand::Cancel { diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs index 51bb5bf..5893122 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs @@ -978,7 +978,7 @@ impl MvpAgent { sessions: RefCell::new(HashMap::new()), activity, loading_sessions: RefCell::new(HashMap::new()), - prompt_intake_locks: RefCell::new(HashMap::new()), + dispatch_locks: RefCell::new(HashMap::new()), session_threads: RefCell::new(HashMap::new()), resident_roster_titles: RefCell::new(HashMap::new()), initialize_request: OnceLock::new(), diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs index b9f23e8..6e31f61 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs @@ -462,19 +462,16 @@ pub struct MvpAgent { loading_sessions: RefCell< HashMap>, >, - /// Per-session prompt-intake serialization lock. LEADER-SAFE(per-session): - /// keyed by SessionId, mirrors `sessions` lifecycle. - /// - /// Each incoming `session/prompt` RPC is dispatched as its own task by the - /// ACP message loop, and [`Self::prompt`] runs an async preamble (prompt-mode - /// query, trace context, model lookup) BEFORE it enqueues - /// `SessionCommand::Prompt` onto the actor's FIFO mailbox. Without - /// serialization those preambles interleave across tasks, so the mailbox — - /// and therefore the authoritative prompt queue — receives prompts out of - /// submission order. `prompt()` holds this lock across the preamble and - /// releases it immediately after the enqueue (the turn itself runs unlocked), - /// which makes intake order match arrival order. - prompt_intake_locks: RefCell< + /// Per-session lock ordering dispatch onto the actor's mailbox: + /// [`Self::prompt`] holds it across its intake preamble (prompt-mode + /// query, trace context, model lookup) and releases it right after the + /// enqueue, so prompts land in submission order; `cancel` holds it + /// around its `Cancel` send, so a cancel cannot overtake the prompt it + /// targets and wedge the turn slot (see + /// `cancel_never_overtakes_in_flight_prompt_intake`). Cancels wait out + /// preambles held ahead of them — keep preambles lean. + /// LEADER-SAFE(per-session): mirrors `sessions` lifecycle. + dispatch_locks: RefCell< HashMap>>, >, /// LEADER-SAFE(per-session): keyed by SessionId. Mirrors `sessions` lifecycle. diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/session_lifecycle.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/session_lifecycle.rs index 0f70111..feb9174 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/session_lifecycle.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/session_lifecycle.rs @@ -38,19 +38,16 @@ impl MvpAgent { /// drops the `session_live_state` entry so that map stays bounded. pub(crate) fn remove_session(&self, id: &acp::SessionId) { self.sessions.borrow_mut().remove(id); - self.prompt_intake_locks.borrow_mut().remove(id); + self.dispatch_locks.borrow_mut().remove(id); self.session_threads.borrow_mut().remove(id); self.session_index_claims.borrow_mut().remove(id); self.require_gateway_sessions.borrow_mut().remove(id); self.session_live_state.borrow_mut().remove(id); } - /// Get-or-create the per-session prompt-intake lock (see - /// [`Self::prompt_intake_locks`]). - pub(super) fn prompt_intake_lock( - &self, - id: &acp::SessionId, - ) -> std::rc::Rc> { - self.prompt_intake_locks + /// Get-or-create the per-session dispatch lock (see + /// [`Self::dispatch_locks`]). Cheap clone of the shared `Rc`. + pub(super) fn dispatch_lock(&self, id: &acp::SessionId) -> std::rc::Rc> { + self.dispatch_locks .borrow_mut() .entry(id.clone()) .or_default() diff --git a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs index 9cc18f2..49b7548 100644 --- a/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs +++ b/crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs @@ -2152,6 +2152,59 @@ fn cancel_does_not_forward_to_bridge_in_local_mode() { ); }); } +/// Regression (post-cancel slot hang): a cancel racing the prompt it targets +/// must land on the actor mailbox AFTER that prompt, or it hits an idle +/// actor as a no-op and the prompt then starts a turn nothing cancels. +/// The dispatch lock (`dispatch_locks`) provides that ordering. +#[test] +fn cancel_never_overtakes_in_flight_prompt_intake() { + use crate::session::SessionCommand; + use acp::Agent as _; + run_local_for_bridge_test(|| async { + let agent = build_minimal_agent_for_tests(); + let sid = acp::SessionId::new("sess-cancel-intake-race"); + let (handle, _tx, mut cmd_rx) = make_live_session_handle(&sid, None); + agent.sessions.borrow_mut().insert(sid.clone(), handle); + let order: std::rc::Rc>> = + std::rc::Rc::new(std::cell::RefCell::new(Vec::new())); + let (intake_parked_tx, intake_parked_rx) = tokio::sync::oneshot::channel::<()>(); + let driver_order = order.clone(); + tokio::task::spawn_local(async move { + let mut intake_parked_tx = Some(intake_parked_tx); + while let Some(cmd) = cmd_rx.recv().await { + match cmd { + SessionCommand::GetCurrentPromptMode { .. } => { + if let Some(tx) = intake_parked_tx.take() { + let _ = tx.send(()); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + } + SessionCommand::Prompt { .. } => driver_order.borrow_mut().push("prompt"), + SessionCommand::Cancel { .. } => driver_order.borrow_mut().push("cancel"), + _ => {} + } + } + }); + let prompt_fut = agent.prompt(acp::PromptRequest::new( + sid.clone(), + vec![acp::ContentBlock::from("hi")], + )); + let cancel_fut = async { + intake_parked_rx + .await + .expect("prompt intake reaches the fake actor"); + let _ = agent + .cancel(acp::CancelNotification::new(sid.clone())) + .await; + }; + let _ = futures::join!(prompt_fut, cancel_fut); + assert_eq!( + order.borrow().as_slice(), + ["prompt", "cancel"], + "cancel must land on the actor mailbox after the prompt it targets" + ); + }); +} use crate::session::SessionCommand as TestSessionCommand; /// Build a session handle wired to a *live* command channel. Returns the /// handle (move into `sessions`) plus a probe `cmd_tx`/`cmd_rx` so a test