From 7a219c29e91ecdbb133c0b3a1b9b06daac4ab13c Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 28 Oct 2026 17:41:58 +1100 Subject: [PATCH 1/7] feat(acp-client): turn-control channel with steering and interrupt-and-reprompt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a second control channel to the session connection so a message can reach a running agent now instead of after the turn in flight ends: - `TurnControlHandle` (via `SessionConnection::turn_control_handle()`), with `steer()` and `interrupt_and_reprompt()`, cloneable for registration next to `AsyncTaskStopHandle`. - Steering: parse `_meta.steering.supported` at initialize onto `AcpSessionSetup`; inside the live turn's select, forward a steer as `_session/steering` with `idleBehavior: promptRequired` and map the answer (`injected` → `Ok(Injected)`, anything else or a transport error → `Err`) without stalling the select. Refused locally when the agent did not advertise it. - Interrupt: `session/cancel` the live turn, await its prompt response (5 s bound, `INTERRUPT_ACK_TIMEOUT`), and hand the message back as `TurnStep::Interrupted(next_turn)`. The session loop treats that as a turn boundary: same connection, lifetime stays `TurnLive`, no settle, and the interrupted `prompt()` stays pending for the final turn's result. The connection-wide permission token is deliberately left alone. A bridge that never acknowledges settles the turn cancelled without a re-prompt. - Hold and idle: a steer or interrupt read by the holding wait (or while the loop waits for a prompt) becomes a plain new turn with no extension frame; the caller gets `Ok(NewTurnStarted)` / `Ok(())` once the turn goes out. Requests are never rejected at hold boundaries; teardown answers anything still queued and flips `connection_alive` so later requests fail fast. - `MessageWriter::record_user_message(text, images)` default no-op: the connection persists messages it delivered itself, after finalizing the open assistant message and before the agent can reply. Tests (`driver/tests/turn_control.rs`) add an in-process fake ACP agent (initialize/new/prompt/cancel plus an `UntypedMessage` steering handler) and drive `run_prompt_turn` / `run_acp_session` end to end: capability parsing, injected steer with the frame shape and row ordering, promptRequired / method-not-found / unadvertised steers, steer and interrupt during a hold, interrupt-and-reprompt with the transcript order, the still-pending original reply and a turn-2 permission request that is still answered, hold-boundary pass-through, and teardown drains. Where this follows the code rather than the plan note: - The steered user row is persisted in the steer's ordered response callback and only on `injected`, not before the request is sent. The callback runs under the SDK's ordered-response barrier, so the assistant→user→assistant ordering holds with no interleaving window, and a `promptRequired`/`failed` steer leaves no row for the caller's fallback delivery to duplicate — no `persist_user_message` retry flag is needed. - `record_user_message` takes the `PromptImage`s, not image ids: acp -client has no ids, only the blocks it sends. - The plan's "fake-agent harness used by the hold tests" is a handler-level harness (direct `hold_for_background_quiescence` calls); the real in-process agent harness lives in the new test module. ACP request handlers run inline in the dispatch loop, so the fake agent stashes the prompt `Responder` and answers it later. - The interrupt reply resolves when the re-prompt starts, not at cancel acknowledgment; the two are adjacent on the connection task. - A control request arriving while the loop is idle before any prompt is also served as a new turn (same rule as the hold). Not covered by a test: the wedged-bridge timeout path (5 s real time). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- crates/acp-client/src/driver.rs | 657 +++++++++++- .../src/driver/tests/response_completion.rs | 10 +- .../acp-client/src/driver/tests/settlement.rs | 2 + .../src/driver/tests/turn_control.rs | 995 ++++++++++++++++++ 4 files changed, 1623 insertions(+), 41 deletions(-) create mode 100644 crates/acp-client/src/driver/tests/turn_control.rs diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index d98eba3fc..d3bfd8ec2 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -77,6 +77,18 @@ pub trait MessageWriter: Send + Sync { false } + /// Persist a user message the connection delivered itself: one that + /// reached the agent mid-session through a [`TurnControlHandle`] (steered + /// into the live turn, or sent as the prompt that follows an interrupt) + /// rather than through a caller that wrote its own row before calling + /// [`SessionConnection::prompt`]. + /// + /// Called on the connection task right after the open assistant message + /// is finalized and before anything the agent says in reply can be + /// appended, so the transcript reads assistant → user → assistant. The + /// default records nothing. + async fn record_user_message(&self, _text: &str, _images: &[PromptImage]) {} + /// Record a tool call with its ID, title, and optional raw input parameters. async fn record_tool_call( &self, @@ -1412,8 +1424,107 @@ type SessionTurnResult = Result; struct QueuedSessionTurn { prompt: String, images: Vec, - /// Resolved by the connection task with the turn's outcome. - reply: oneshot::Sender, + /// Resolved by the connection task with the turn's outcome. `None` for a + /// turn that did not come through [`SessionConnection::prompt`]: one the + /// turn-control channel started (served out of a hold, or re-prompted + /// after an interrupt). Its caller is answered through `control`, and its + /// outcome belongs to whichever `prompt()` is still pending on the + /// connection. + reply: Option>, + /// The turn-control request this turn serves, if any. Answered once the + /// turn is being sent; its presence also tells the session loop that the + /// connection must persist the user message itself (see + /// [`MessageWriter::record_user_message`]). + control: Option, +} + +/// How a message reached the agent through [`TurnControlHandle::steer`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SteerOutcome { + /// The agent injected the message into its in-flight turn. + Injected, + /// No turn was in flight on the connection (it was holding for background + /// work, or idle between turns), so the message was sent as a fresh + /// `session/prompt` on the same connection instead. + NewTurnStarted, +} + +/// A request queued onto a [`SessionConnection`]'s turn-control channel by a +/// [`TurnControlHandle`]: a message that must reach the agent now, not after +/// the turn in flight ends. +enum TurnControlRequest { + /// Inject into the in-flight turn via `_session/steering`. + Steer { + prompt: String, + images: Vec, + reply: oneshot::Sender>, + }, + /// `session/cancel` the in-flight turn, then send the message as the next + /// `session/prompt` on the same connection. + Interrupt { + prompt: String, + images: Vec, + reply: oneshot::Sender>, + }, +} + +impl TurnControlRequest { + /// Serve the request as a fresh turn on this connection — the answer + /// whenever nothing is in flight to steer or interrupt: a steer would come + /// back `promptRequired` (or, from codex-acp, start a turn this client + /// would not own), and an interrupt has nothing to cancel. + fn into_turn(self) -> QueuedSessionTurn { + let (prompt, images, control) = match self { + Self::Steer { + prompt, + images, + reply, + } => (prompt, images, TurnControlReply::Steer(reply)), + Self::Interrupt { + prompt, + images, + reply, + } => (prompt, images, TurnControlReply::Interrupt(reply)), + }; + QueuedSessionTurn { + prompt, + images, + reply: None, + control: Some(control), + } + } + + fn reject(self, reason: String) { + match self { + Self::Steer { reply, .. } => { + let _ = reply.send(Err(reason)); + } + Self::Interrupt { reply, .. } => { + let _ = reply.send(Err(reason)); + } + } + } +} + +/// The reply half of a [`TurnControlRequest`] whose message became a +/// [`QueuedSessionTurn`] of its own. +enum TurnControlReply { + Steer(oneshot::Sender>), + Interrupt(oneshot::Sender>), +} + +impl TurnControlReply { + /// Answer the caller: the message is going out as its own turn. + fn turn_started(self) { + match self { + Self::Steer(reply) => { + let _ = reply.send(Ok(SteerOutcome::NewTurnStarted)); + } + Self::Interrupt(reply) => { + let _ = reply.send(Ok(())); + } + } + } } /// A per-task stop queued onto a [`SessionConnection`]'s hold-control channel, @@ -1465,6 +1576,84 @@ impl AsyncTaskStopHandle { } } +/// Cloneable, thread-safe handle for delivering a message to a +/// [`SessionConnection`]'s agent *now* — into the turn in flight, or in place +/// of it — instead of queueing it for after the turn. Registered where a +/// client's own commands can reach it, like [`AsyncTaskStopHandle`]. +#[derive(Clone)] +pub struct TurnControlHandle { + turn_control_tx: mpsc::UnboundedSender, + /// Cleared when the connection task tears down, so a request issued after + /// the session settled fails fast instead of sitting queued. + connection_alive: Arc, +} + +impl TurnControlHandle { + /// Inject a message into the agent's in-flight turn with the + /// `_session/steering` extension; tool calls already running survive. + /// + /// `Ok(Injected)` means the agent took it into the live turn. With no turn + /// in flight (the connection holding for background work, or idle) the + /// message is sent as a fresh `session/prompt` on the same connection and + /// the answer is `Ok(NewTurnStarted)`. `Err` when the agent does not + /// advertise steering, answered anything but `injected` (`promptRequired` + /// when the turn ended under the request, `failed`, `startedNewTurn`), or + /// the connection is gone; callers then fall back to + /// [`TurnControlHandle::interrupt_and_reprompt`]. The user message is + /// persisted (see [`MessageWriter::record_user_message`]) only on `Ok`. + pub async fn steer( + &self, + prompt: &str, + images: &[PromptImage], + ) -> Result { + let (reply_tx, reply_rx) = oneshot::channel(); + self.send(TurnControlRequest::Steer { + prompt: prompt.to_string(), + images: images.to_vec(), + reply: reply_tx, + })?; + reply_rx + .await + .map_err(|_| "ACP connection closed before the steer resolved".to_string())? + } + + /// Cancel the agent's in-flight turn (`session/cancel`) and send the + /// message as the next `session/prompt` on the same connection — the + /// delivery that works with every ACP agent. The connection never settles + /// in between: the `prompt()` future for the interrupted turn stays + /// pending and resolves with the final turn's outcome, and nothing is + /// torn down or respawned. With no turn in flight the message is simply + /// sent as a fresh turn. + /// + /// Resolves once the new prompt is going out. `Err` if the agent did not + /// acknowledge the cancel within `INTERRUPT_ACK_TIMEOUT` (the turn is + /// then cancelled without a re-prompt) or the connection is gone. + pub async fn interrupt_and_reprompt( + &self, + prompt: &str, + images: &[PromptImage], + ) -> Result<(), String> { + let (reply_tx, reply_rx) = oneshot::channel(); + self.send(TurnControlRequest::Interrupt { + prompt: prompt.to_string(), + images: images.to_vec(), + reply: reply_tx, + })?; + reply_rx + .await + .map_err(|_| "ACP connection closed before the interrupt resolved".to_string())? + } + + fn send(&self, request: TurnControlRequest) -> Result<(), String> { + if !self.connection_alive.load(Ordering::Relaxed) { + return Err("ACP connection is no longer running".to_string()); + } + self.turn_control_tx + .send(request) + .map_err(|_| "ACP connection is no longer running".to_string()) + } +} + /// A session-scoped connection to an ACP agent. /// /// Owns the agent (bridge) child process and the JSON-RPC connection through @@ -1489,6 +1678,11 @@ pub struct SessionConnection { /// stop issued outside that window is rejected immediately (see /// [`AsyncTaskStopHandle`]). hold_active: Arc, + /// Messages for the turn in flight (see [`TurnControlHandle`]), served by + /// the live turn, the holding wait, and the between-turns wait alike. + turn_control_tx: mpsc::UnboundedSender, + /// Cleared by the connection task at teardown (see [`TurnControlHandle`]). + connection_alive: Arc, /// Handle to the connection task; awaited (once) to surface the session /// result when a turn's reply channel is dropped without an answer. connection_task: Option>, @@ -1530,6 +1724,17 @@ impl SessionConnection { } } + /// A cloneable, thread-safe handle for delivering a message to the turn + /// in flight ([`TurnControlHandle::steer`]) or in place of it + /// ([`TurnControlHandle::interrupt_and_reprompt`]), for callers that + /// route it from another thread (a UI command, an MCP tool). + pub fn turn_control_handle(&self) -> TurnControlHandle { + TurnControlHandle { + turn_control_tx: self.turn_control_tx.clone(), + connection_alive: Arc::clone(&self.connection_alive), + } + } + /// Take the receiver for the connection's [`SessionSettled`] event. /// /// The event fires exactly once, after the agent has been torn down. @@ -1560,7 +1765,8 @@ impl SessionConnection { let turn = QueuedSessionTurn { prompt: prompt.to_string(), images: images.to_vec(), - reply: reply_tx, + reply: Some(reply_tx), + control: None, }; if self.prompt_tx.send(turn).is_err() { // The connection task already exited (e.g. setup failed before @@ -1871,6 +2077,10 @@ impl AcpDriver { mpsc::unbounded_channel::(); let hold_active = Arc::new(AtomicBool::new(false)); let hold_active_for_task = Arc::clone(&hold_active); + let (turn_control_tx, mut turn_control_rx) = + mpsc::unbounded_channel::(); + let connection_alive = Arc::new(AtomicBool::new(true)); + let connection_alive_for_task = Arc::clone(&connection_alive); let (settled_tx, settled_rx) = oneshot::channel::(); // Owned copies for the connection task — its lifetime is the @@ -1978,6 +2188,7 @@ impl AcpDriver { &child_exited, &mut prompt_rx, &mut hold_control_rx, + &mut turn_control_rx, &hold_active_for_task, &pending_reply, ) @@ -1987,6 +2198,11 @@ impl AcpDriver { .await .map_err(|e| format!("ACP protocol failed: {e:?}")); + // Nothing serves the turn-control channel past this point: fail + // new requests fast at the handle and answer any already queued. + connection_alive_for_task.store(false, Ordering::Relaxed); + reject_queued_turn_control_requests(&mut turn_control_rx, "the ACP connection closed"); + writer.finalize().await; graceful_stop(&mut child, is_remote).await; @@ -2025,6 +2241,8 @@ impl AcpDriver { prompt_tx, hold_control_tx, hold_active, + turn_control_tx, + connection_alive, connection_task: Some(connection_task), exit_result: None, settled_rx: Some(settled_rx), @@ -2817,6 +3035,23 @@ fn task_tracking_mode_from_initialize(meta: Option<&Meta>) -> TaskTrackingMode { } } +/// Top-level `_meta` key under which an agent advertises the +/// `_session/steering` extension at initialize (a sibling of +/// `agentCapabilities`, not inside it): `_meta.steering.supported`. +const STEERING_META_KEY: &str = "steering"; + +/// Whether the agent advertised steering at initialize (claude-agent-acp +/// ≥ 0.98, codex-acp ≥ 2.2). Anything missing or malformed is an agent without +/// it, which [`TurnControlHandle::steer`] answers with an error so callers +/// interrupt instead — the only safe gate with managed bridge versions that +/// float. +fn steering_supported_from_initialize(meta: Option<&Meta>) -> bool { + meta.and_then(|meta| meta.get(STEERING_META_KEY)) + .and_then(|steering| steering.get("supported")) + .and_then(serde_json::Value::as_bool) + == Some(true) +} + /// Method name shared by every `session/update` notification, typed or not. const SESSION_UPDATE_METHOD: &str = "session/update"; @@ -2859,6 +3094,66 @@ fn async_task_stop_outcome( .map_err(|e| format!("Async task stop request failed: {e:?}")) } +/// Extension request that injects a message into the turn in flight: +/// `{sessionId, prompt: ContentBlock[], _meta.steering.idleBehavior}` → +/// `{outcome: "injected" | "startedNewTurn" | "promptRequired" | "failed"}`. +/// The prompt is the same content-block array as `session/prompt`. Priority is +/// the agent's call (`now`, or `later` while a permission request is open), +/// not a client field. Sent verbatim with its leading `_`, like +/// [`ASYNC_TASK_STOP_METHOD`]. +const STEERING_METHOD: &str = "_session/steering"; + +/// The `_session/steering` request frame for one message. +/// +/// `idleBehavior: promptRequired` asks an agent with no turn in flight to say +/// so instead of starting a turn this client would not own. claude-agent-acp +/// honours it; codex-acp ignores it and starts a turn — which is why the +/// connection only ever sends this frame from inside a live turn (see +/// [`serve_steer`]) and serves a steer that lands anywhere else as a turn of +/// its own. +fn steering_message( + agent_session_id: &str, + content_blocks: Vec, +) -> Result { + UntypedMessage::new( + STEERING_METHOD, + serde_json::json!({ + "sessionId": agent_session_id, + "prompt": content_blocks, + "_meta": { STEERING_META_KEY: { "idleBehavior": "promptRequired" } }, + }), + ) +} + +/// Fold a `_session/steering` response into the caller's answer. Only +/// `injected` is a delivery. `promptRequired` means the turn ended under the +/// request, `failed` is the agent's own refusal, `startedNewTurn` is a turn +/// this client is not awaiting (unreachable from a live turn in practice, but +/// mapped anyway); all read as "not injected" so the caller falls back, and +/// a transport or agent error (method-not-found from a bridge without the +/// extension) surfaces as such. +fn steer_outcome( + result: Result, +) -> Result { + let response = result.map_err(|e| format!("Steering request failed: {e:?}"))?; + match response.get("outcome").and_then(serde_json::Value::as_str) { + Some("injected") => Ok(SteerOutcome::Injected), + Some(outcome) => { + let reason = response + .get("reason") + .and_then(serde_json::Value::as_str) + .map(|reason| format!(" ({reason})")) + .unwrap_or_default(); + Err(format!( + "Steering not injected: agent answered '{outcome}'{reason}" + )) + } + None => Err(format!( + "Steering not injected: malformed response {response}" + )), + } +} + /// The three `sessionUpdate` discriminators of the typed asyncTasks /// extension. const ASYNC_TASK_SPAWNED_KIND: &str = "async_task_spawned"; @@ -4580,16 +4875,110 @@ fn serve_async_task_stop( } } +/// Serve one queued steer from inside a live turn: send `_session/steering` +/// for the message and schedule the agent's answer onto the request's reply +/// channel. Never awaits the answer inline, for the same reason as +/// [`serve_async_task_stop`] — the live turn's select must stay free to see +/// the prompt resolve, the cancel, or the next control request. +/// +/// The user message is persisted in the response callback, and only on +/// `injected`: the callback runs under the connection's ordered-response +/// barrier, so the assistant message open at that moment is finalized and the +/// user row written before any later frame from the agent can be dispatched +/// — the transcript reads assistant → user → assistant with no interleaving +/// window — while a `promptRequired` or failed steer leaves no row behind for +/// the caller's fallback delivery to duplicate. +fn serve_steer( + connection: &ConnectionTo, + setup: &AcpSessionSetup, + handler: &Arc, + our_session_id: &str, + prompt: String, + images: Vec, + reply: oneshot::Sender>, +) { + if !setup.steering_supported { + let _ = reply.send(Err( + "Steering not available: agent does not advertise steering".to_string(), + )); + return; + } + log::info!("ACP session {our_session_id}: steering the turn in flight"); + let supports_images = setup.agent_capabilities.prompt_capabilities.image; + let content_blocks = build_prompt_content_blocks(&prompt, &images, supports_images); + let message = match steering_message(&setup.agent_session_id, content_blocks) { + Ok(message) => message, + Err(e) => { + let _ = reply.send(Err(format!("Failed to build steering request: {e:?}"))); + return; + } + }; + let handler = Arc::clone(handler); + let session_id = our_session_id.to_string(); + let scheduled = + connection + .send_request(message) + .on_receiving_result(move |result| async move { + let outcome = steer_outcome(result); + match &outcome { + Ok(outcome) => { + log::info!("ACP session {session_id}: steer answered {outcome:?}"); + handler.writer.finalize().await; + handler.response.lock().await.close_message(); + handler.writer.record_user_message(&prompt, &images).await; + } + Err(e) => log::warn!("ACP session {session_id}: steer not injected: {e}"), + } + let _ = reply.send(outcome); + Ok(()) + }); + if let Err(e) = scheduled { + // The reply sender moved into the callback and drops with it, which + // already answers the caller; the send failure is only loggable here. + log::warn!("ACP session {our_session_id}: failed to send steering request: {e:?}"); + } +} + // ============================================================================= // Protocol helpers // ============================================================================= +/// How long an interrupt gives the agent to answer the in-flight +/// `session/prompt` after `session/cancel` before giving up on the re-prompt: +/// a bridge that cannot acknowledge a cancel should not be handed a second +/// prompt, so the turn then settles cancelled, as a plain cancel would. +/// Mirrors the in-turn cancel's bound on the same response. +const INTERRUPT_ACK_TIMEOUT: Duration = Duration::from_secs(5); + /// How long the holding wait gives the bridge to process a `session/cancel` /// before teardown. Out-of-turn there is no prompt response to await, so this /// bounded courtesy drain mirrors the in-turn cancel's bounded response wait /// (which the bridge answers by aborting the turn). const HOLDING_CANCEL_DRAIN: Duration = Duration::from_millis(500); +/// How one prompt turn ended, from the session loop's point of view. +enum TurnStep { + /// The turn ran to its own stop reason (or the user cancelled it). + Settled(AgentRunOutcome), + /// The turn was cancelled on behalf of a new message, which is the next + /// turn on this connection. A turn boundary, not a session end. + Interrupted(QueuedSessionTurn), +} + +// Hand-written for the same reason as [`HoldOutcome`]'s: the queued turn +// carries reply channels, not `Debug` types. +impl std::fmt::Debug for TurnStep { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TurnStep::Settled(outcome) => f.debug_tuple("Settled").field(outcome).finish(), + TurnStep::Interrupted(turn) => f + .debug_struct("Interrupted") + .field("prompt", &turn.prompt) + .finish_non_exhaustive(), + } + } +} + /// Terminal decisions of the holding wait's timer core. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum HoldSettle { @@ -4870,6 +5259,7 @@ async fn hold_for_background_quiescence( child_exited: &CancellationToken, prompt_rx: &mut mpsc::UnboundedReceiver, hold_control_rx: &mut mpsc::UnboundedReceiver, + turn_control_rx: &mut mpsc::UnboundedReceiver, hold_active: &AtomicBool, serve_stop: &dyn Fn(StopAsyncTaskRequest), our_session_id: &str, @@ -4924,6 +5314,7 @@ async fn hold_for_background_quiescence( // Closed channels disable their branch instead of busy-looping the select. let mut prompts_open = true; let mut control_open = true; + let mut turn_control_open = true; let mut activity_open = true; let outcome = loop { if activity_rx.borrow().finish_requested { @@ -4951,6 +5342,17 @@ async fn hold_for_background_quiescence( Some(request) => serve_stop(request), None => control_open = false, }, + // No turn is in flight to steer or interrupt (the bridge has no + // unsettled prompt by construction), so a message delivered now + // is simply the next turn — never a `_session/steering` frame, + // which Claude would answer `promptRequired` and codex-acp would + // turn into a prompt this client is not awaiting. Unlike per-task + // stops, these are not aimed at this hold and are never rejected + // at its boundaries. + request = turn_control_rx.recv(), if turn_control_open => match request { + Some(request) => break HoldOutcome::NewTurn(request.into_turn()), + None => turn_control_open = false, + }, changed = activity_rx.changed(), if activity_open => match changed { Ok(()) => { let activity = activity_rx.borrow_and_update().clone(); @@ -4995,6 +5397,19 @@ fn reject_queued_stop_requests( } } +/// Answer every currently-queued turn-control request with an error. Only +/// called at connection teardown: unlike per-task stops, a message is not +/// aimed at one phase of the connection, so one racing a turn → hold boundary +/// is left queued for whichever loop reads next. +fn reject_queued_turn_control_requests( + turn_control_rx: &mut mpsc::UnboundedReceiver, + reason: &str, +) { + while let Ok(request) = turn_control_rx.try_recv() { + request.reject(format!("Message not delivered: {reason}")); + } +} + /// Session-scoped connection loop: set up the ACP session once, then serve /// prompt turns received over `prompt_rx`. /// @@ -5020,6 +5435,7 @@ async fn run_acp_session( child_exited: &CancellationToken, prompt_rx: &mut mpsc::UnboundedReceiver, hold_control_rx: &mut mpsc::UnboundedReceiver, + turn_control_rx: &mut mpsc::UnboundedReceiver, hold_active: &AtomicBool, pending_reply: &RefCell>>, ) -> Result<(AgentRunOutcome, SessionSettleReason), String> { @@ -5093,22 +5509,32 @@ async fn run_acp_session( // live turn when a prompt arrives while holding. let mut lifetime: Option = None; let mut next_turn: Option = None; + let mut turn_control_open = true; loop { - let turn = match next_turn.take() { - // A prompt that arrived during the holding wait re-enters a live - // turn on the already-established session. + let mut turn = match next_turn.take() { + // A prompt that arrived during the holding wait (or the message + // an interrupt was made for) re-enters a live turn on the + // already-established session. Some(turn) => turn, - None => tokio::select! { - _ = cancel_token.cancelled() => { - handler.cancel_pending_permissions(); - return Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)); + None => loop { + tokio::select! { + _ = cancel_token.cancelled() => { + handler.cancel_pending_permissions(); + return Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)); + } + turn = prompt_rx.recv() => match turn { + Some(turn) => break turn, + // Every connection handle was dropped without sending a + // prompt — nothing to serve, tear down. + None => return Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)), + }, + // A message delivered while nothing is in flight has + // nothing to steer or interrupt: it is the next turn. + request = turn_control_rx.recv(), if turn_control_open => match request { + Some(request) => break request.into_turn(), + None => turn_control_open = false, + }, } - turn = prompt_rx.recv() => match turn { - Some(turn) => turn, - // Every connection handle was dropped without sending a - // prompt — nothing to serve, tear down. - None => return Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)), - }, }, }; transition_lifetime(our_session_id, &mut lifetime, SessionLifetime::TurnLive); @@ -5116,8 +5542,27 @@ async fn run_acp_session( // Stash the reply sender: unless the holding path below resolves it // early, the connection task resolves it after teardown, so prompt() // callers resume with the child already stopped, exactly like the - // old per-turn run(). - *pending_reply.borrow_mut() = Some(turn.reply); + // old per-turn run(). A turn the control channel started has none; + // whatever prompt() is still pending (the interrupted turn's) keeps + // waiting and gets the final turn's result. + if let Some(reply) = turn.reply.take() { + *pending_reply.borrow_mut() = Some(reply); + } + + if let Some(control) = turn.control.take() { + // The message came through the control channel, so no caller + // wrote its row. Close whatever assistant message is open (the + // interrupted turn's, or a continuation's) and persist it here, + // before the prompt goes out, so nothing the agent says in reply + // can land ahead of it. + handler.writer.finalize().await; + handler.response.lock().await.close_message(); + handler + .writer + .record_user_message(&turn.prompt, &turn.images) + .await; + control.turn_started(); + } let prompt = run_prompt_turn( connection, @@ -5125,10 +5570,11 @@ async fn run_acp_session( handler, our_session_id, cancel_token, + turn_control_rx, &turn.prompt, &turn.images, ); - let outcome = if background_hold.is_some() { + let step = if background_hold.is_some() { tokio::select! { biased; result = prompt => result, @@ -5143,6 +5589,17 @@ async fn run_acp_session( } else { prompt.await }; + let outcome = match step { + // A turn boundary, not a session end: the connection stays up, + // the lifetime stays `TurnLive`, nothing is torn down, and the + // stashed reply keeps waiting for the turn that follows. + Ok(TurnStep::Interrupted(turn)) => { + next_turn = Some(turn); + continue; + } + Ok(TurnStep::Settled(outcome)) => Ok(outcome), + Err(e) => Err(e), + }; // Feature off: preserve the old behavior byte-for-byte — finalize() // + graceful_stop() fire immediately after each prompt resolves, and @@ -5242,6 +5699,7 @@ async fn run_acp_session( child_exited, prompt_rx, hold_control_rx, + turn_control_rx, hold_active, &|request| serve_async_task_stop(connection, &setup, our_session_id, request), our_session_id, @@ -5319,16 +5777,24 @@ async fn run_acp_session( } /// Send one prompt turn over an established ACP session and map its stop -/// reason to an [`AgentRunOutcome`]. +/// reason to an [`AgentRunOutcome`] — or, when a message interrupts it, hand +/// that message back as the next turn. +/// +/// While the prompt is in flight the turn also serves the turn-control +/// channel: a steer is forwarded to the agent as `_session/steering` without +/// leaving the turn, an interrupt cancels the turn and ends it with +/// [`TurnStep::Interrupted`]. +#[allow(clippy::too_many_arguments)] async fn run_prompt_turn( connection: &ConnectionTo, setup: &AcpSessionSetup, handler: &Arc, our_session_id: &str, cancel_token: &CancellationToken, + turn_control_rx: &mut mpsc::UnboundedReceiver, prompt: &str, images: &[PromptImage], -) -> Result { +) -> Result { let supports_images = setup.agent_capabilities.prompt_capabilities.image; if !images.is_empty() && !supports_images { log::warn!( @@ -5349,24 +5815,85 @@ async fn run_prompt_turn( let prompt_task = connection.send_request(prompt_request).block_task(); tokio::pin!(prompt_task); - let prompt_response = tokio::select! { - result = &mut prompt_task => { - result.map_err(|error| describe_prompt_error(error, &setup.auth_methods))? - } - _ = cancel_token.cancelled() => { - handler.cancel_pending_permissions(); - if let Err(e) = send_session_cancel(connection, &setup.agent_session_id) { - log::warn!("{e}"); + // A closed control channel disables its branch instead of busy-looping. + let mut control_open = true; + let prompt_response = loop { + tokio::select! { + result = &mut prompt_task => { + break result.map_err(|error| describe_prompt_error(error, &setup.auth_methods))?; } + _ = cancel_token.cancelled() => { + handler.cancel_pending_permissions(); + if let Err(e) = send_session_cancel(connection, &setup.agent_session_id) { + log::warn!("{e}"); + } - match tokio::time::timeout(Duration::from_secs(5), &mut prompt_task).await { - Ok(result) => result.map_err(|e| format!("Prompt failed after cancellation: {e:?}"))?, - Err(_) => { - log::warn!( - "Timed out waiting for ACP prompt response after session/cancel for session {our_session_id}" + match tokio::time::timeout(Duration::from_secs(5), &mut prompt_task).await { + Ok(result) => { + break result.map_err(|e| format!("Prompt failed after cancellation: {e:?}"))?; + } + Err(_) => { + log::warn!( + "Timed out waiting for ACP prompt response after session/cancel for session {our_session_id}" + ); + return Ok(TurnStep::Settled(AgentRunOutcome::Cancelled)); + } + } + } + request = turn_control_rx.recv(), if control_open => match request { + Some(TurnControlRequest::Steer { prompt, images, reply }) => { + serve_steer(connection, setup, handler, our_session_id, prompt, images, reply); + } + Some(TurnControlRequest::Interrupt { prompt: next_prompt, images: next_images, reply }) => { + log::info!( + "ACP session {our_session_id}: interrupting the turn in flight for a new message" ); - return Ok(AgentRunOutcome::Cancelled); + // Same wire sequence as the user-cancel arm above, minus + // `cancel_pending_permissions()`: that token is created + // once per connection and, once cancelled, auto-answers + // every later permission request — which would silently + // break the next turn. The agent resolves its own pending + // permission requests on `session/cancel`. + if let Err(e) = send_session_cancel(connection, &setup.agent_session_id) { + let _ = reply.send(Err(e.clone())); + return Err(e); + } + match tokio::time::timeout(INTERRUPT_ACK_TIMEOUT, &mut prompt_task).await { + // Whatever the stop reason (`cancelled`, or `end_turn` + // if the turn finished under the cancel), the turn is + // over and the message is the next one. + Ok(Ok(response)) => { + if let Some(metadata) = prompt_response_metadata(&response) { + handler.writer.record_acp_event_metadata(metadata).await; + } + return Ok(TurnStep::Interrupted(QueuedSessionTurn { + prompt: next_prompt, + images: next_images, + reply: None, + control: Some(TurnControlReply::Interrupt(reply)), + })); + } + Ok(Err(e)) => { + let message = format!("Prompt failed after interrupt: {e:?}"); + let _ = reply.send(Err(message.clone())); + return Err(message); + } + // A wedged bridge should not get a second prompt: + // settle the turn cancelled, as a plain cancel would. + Err(_) => { + log::warn!( + "Timed out waiting for ACP prompt response after session/cancel \ + for session {our_session_id}; dropping the re-prompt" + ); + let _ = reply.send(Err(format!( + "Agent did not acknowledge session/cancel within {}s", + INTERRUPT_ACK_TIMEOUT.as_secs() + ))); + return Ok(TurnStep::Settled(AgentRunOutcome::Cancelled)); + } + } } + None => control_open = false, } } }; @@ -5383,9 +5910,9 @@ async fn run_prompt_turn( handler.finish_completed_response().await; } - Ok(AgentRunOutcome::from_stop_reason( + Ok(TurnStep::Settled(AgentRunOutcome::from_stop_reason( prompt_response.stop_reason, - )) + ))) } /// The text a failed `session/prompt` surfaces to the user. @@ -5919,6 +6446,9 @@ struct AcpSessionSetup { /// an `auth_required` raised at `session/prompt` can log the same /// inventory the setup path logs (see [`describe_prompt_error`]). auth_methods: Vec, + /// Whether the agent advertised `_session/steering` at initialize (see + /// [`steering_supported_from_initialize`]). + steering_supported: bool, } struct AcpSessionSetupContext<'a> { @@ -6067,6 +6597,8 @@ async fn setup_acp_session(context: AcpSessionSetupContext<'_>) -> Result { @@ -6116,6 +6648,7 @@ async fn setup_acp_session(context: AcpSessionSetupContext<'_>) -> Result { @@ -6156,6 +6689,7 @@ async fn setup_acp_session(context: AcpSessionSetupContext<'_>) -> Result(); let (_hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); let hold_active = AtomicBool::new(false); @@ -11032,6 +11586,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE &child_exited, &mut prompt_rx, &mut hold_control_rx, + &mut turn_control_rx, &hold_active, &|_| {}, "sess-1", @@ -11111,6 +11666,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE let handler = hold_test_handler(); let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); let hold_active = AtomicBool::new(false); @@ -11126,6 +11682,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE &child_exited, &mut prompt_rx, &mut hold_control_rx, + &mut turn_control_rx, &hold_active, &|_| {}, "sess-1", @@ -11189,6 +11746,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE let handler = hold_test_handler(); let (prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); let hold_active = AtomicBool::new(false); @@ -11200,7 +11758,8 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE .send(QueuedSessionTurn { prompt: String::from("follow-up"), images: vec![], - reply: reply_tx, + reply: Some(reply_tx), + control: None, }) .expect("queued turn should send"); @@ -11214,6 +11773,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE &child_exited, &mut prompt_rx, &mut hold_control_rx, + &mut turn_control_rx, &hold_active, &|_| {}, "sess-1", @@ -11238,6 +11798,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE handler.set_task_tracking_mode(TaskTrackingMode::Typed); let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); let hold_active = AtomicBool::new(false); @@ -11261,6 +11822,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE &child_exited, &mut prompt_rx, &mut hold_control_rx, + &mut turn_control_rx, &hold_active, &|_| {}, "sess-1", @@ -11309,6 +11871,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE let handler = hold_test_handler(); let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); // Shared with the hold: it flips true at entry, which is what lets @@ -11343,6 +11906,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE &child_exited, &mut prompt_rx, &mut hold_control_rx, + &mut turn_control_rx, &hold_active, &serve, "sess-1", @@ -11384,6 +11948,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE let handler = hold_test_handler(); let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); let hold_active = Arc::new(AtomicBool::new(false)); @@ -11422,6 +11987,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE &child_exited, &mut prompt_rx, &mut hold_control_rx, + &mut turn_control_rx, &hold_active, &|_| {}, "sess-1", @@ -11523,6 +12089,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE Append(String), Finalize, Resume, + UserMessage(String), Event(RecordedMetadata), ToolCall(String), ToolCallMetadata(RecordedMetadata), @@ -11546,7 +12113,10 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE .filter(|call| { matches!( call, - WriterCall::Append(_) | WriterCall::Finalize | WriterCall::Resume + WriterCall::Append(_) + | WriterCall::Finalize + | WriterCall::Resume + | WriterCall::UserMessage(_) ) }) .collect() @@ -11590,6 +12160,13 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE self.calls.lock().unwrap().push(WriterCall::Finalize); } + async fn record_user_message(&self, text: &str, _images: &[PromptImage]) { + self.calls + .lock() + .unwrap() + .push(WriterCall::UserMessage(text.to_string())); + } + async fn record_tool_call( &self, tool_call_id: &str, diff --git a/crates/acp-client/src/driver/tests/response_completion.rs b/crates/acp-client/src/driver/tests/response_completion.rs index 36f221832..70d9ecd80 100644 --- a/crates/acp-client/src/driver/tests/response_completion.rs +++ b/crates/acp-client/src/driver/tests/response_completion.rs @@ -173,6 +173,7 @@ async fn background_response_releases_wait_only_at_completion_and_clears_indicat handler.transition_to_background_holding().await; let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_control_tx, mut control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let child_exited = CancellationToken::new(); let active = AtomicBool::new(false); let (observer, reported) = recording_hold_observer(); @@ -184,6 +185,7 @@ async fn background_response_releases_wait_only_at_completion_and_clears_indicat &child_exited, &mut prompt_rx, &mut control_rx, + &mut turn_control_rx, &active, &|_| {}, "sess-1", @@ -251,6 +253,7 @@ async fn idle_without_a_new_artifact_keeps_waiting() { idle(&handler).await; let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_control_tx, mut control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let child_exited = CancellationToken::new(); let active = AtomicBool::new(false); let mut hold = std::pin::pin!(hold_for_background_quiescence( @@ -261,6 +264,7 @@ async fn idle_without_a_new_artifact_keeps_waiting() { &child_exited, &mut prompt_rx, &mut control_rx, + &mut turn_control_rx, &active, &|_| {}, "sess-1", @@ -373,12 +377,14 @@ async fn session_loop_finishes_successfully_for_initial_and_background_artifacts let store: Arc = Arc::new(RecordingStore::default()); let (prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_control_tx, mut control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let (reply, _reply_rx) = oneshot::channel(); prompt_tx .send(QueuedSessionTurn { prompt: "write".into(), images: vec![], - reply, + reply: Some(reply), + control: None, }) .unwrap(); let child_exited = CancellationToken::new(); @@ -430,6 +436,8 @@ async fn session_loop_finishes_successfully_for_initial_and_background_artifacts &child_exited, &mut prompt_rx, &mut control_rx, + + &mut turn_control_rx, &active, &pending_reply, ); diff --git a/crates/acp-client/src/driver/tests/settlement.rs b/crates/acp-client/src/driver/tests/settlement.rs index 8245b911f..e1c485838 100644 --- a/crates/acp-client/src/driver/tests/settlement.rs +++ b/crates/acp-client/src/driver/tests/settlement.rs @@ -177,6 +177,7 @@ async fn holding_wait_survives_the_old_cap_after_a_coalesced_settlement() { let handler = hold_test_handler(); let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_control_tx, mut control_rx) = mpsc::unbounded_channel(); + let (_turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); let hold_active = AtomicBool::new(false); @@ -191,6 +192,7 @@ async fn holding_wait_survives_the_old_cap_after_a_coalesced_settlement() { &child_exited, &mut prompt_rx, &mut control_rx, + &mut turn_control_rx, &hold_active, &|_| {}, "sess-1", diff --git a/crates/acp-client/src/driver/tests/turn_control.rs b/crates/acp-client/src/driver/tests/turn_control.rs new file mode 100644 index 000000000..ebd678295 --- /dev/null +++ b/crates/acp-client/src/driver/tests/turn_control.rs @@ -0,0 +1,995 @@ +//! Timely message delivery to a running connection: steering into the live +//! turn, interrupt-and-reprompt on the same connection, and the hold-time and +//! teardown behaviour of the turn-control channel. + +use super::*; +use crate::driver::{ + reject_queued_turn_control_requests, run_acp_session, steer_outcome, steering_message, + steering_supported_from_initialize, HandlerPhase, IncomingSessionUpdate, SessionTurnResult, + SteerOutcome, TurnControlHandle, TurnControlReply, TurnControlRequest, TurnStep, + STEERING_META_KEY, STEERING_METHOD, +}; +use agent_client_protocol::schema::v1::{ + CancelNotification, InitializeRequest, InitializeResponse, Meta, PromptResponse, +}; +use agent_client_protocol::{Responder, UntypedMessage}; +use std::cell::RefCell; +use std::future::Future; +use std::sync::atomic::Ordering; + +/// Everything a fake agent saw, plus the prompt responder it is holding back. +/// ACP request handlers run inline in the agent's dispatch loop, so a handler +/// that waited for the test would deadlock the connection: the prompt handler +/// stashes its responder here and the test (or a later `session/cancel`) +/// answers it. +#[derive(Default)] +struct FakeAgent { + prompts: Mutex>, + steers: Mutex>, + cancels: Mutex, + pending_prompt: Mutex>>, + prompt_arrived: tokio::sync::Notify, + /// What the agent answers to `_session/steering`; `None` leaves the + /// method unregistered (an agent without the extension). + steer_response: Option, + /// Whether to announce steering support at initialize. + advertise_steering: bool, + /// Text the agent streams at the start of each prompt, in order. + stream_on_prompt: Vec<&'static str>, + /// Ask permission for a tool call during the Nth (1-based) prompt, and + /// record the client's answer. + permission_on_prompt: Option, + permission_outcomes: Mutex>, + permission_answered: tokio::sync::Notify, +} + +impl FakeAgent { + fn prompts(&self) -> Vec { + self.prompts.lock().unwrap().clone() + } + + fn steers(&self) -> Vec { + self.steers.lock().unwrap().clone() + } + + fn cancels(&self) -> usize { + *self.cancels.lock().unwrap() + } + + async fn await_prompt(&self) { + tokio::time::timeout(SETTLE_TIMEOUT, self.prompt_arrived.notified()) + .await + .expect("agent should receive a prompt"); + } + + /// Answer the prompt the agent is holding back. + fn finish_prompt(&self, stop_reason: StopReason) { + let responder = self + .pending_prompt + .lock() + .unwrap() + .take() + .expect("a prompt should be pending"); + responder + .respond(PromptResponse::new(stop_reason)) + .expect("prompt response should send"); + } + + fn build( + self: &Arc, + ) -> impl agent_client_protocol::ConnectTo + use<> { + let init_agent = Arc::clone(self); + let prompt_agent = Arc::clone(self); + let cancel_agent = Arc::clone(self); + let steer_agent = Arc::clone(self); + let steer_response = self.steer_response.clone(); + let builder = agent_client_protocol::Agent + .builder() + .on_receive_request( + async move |request: InitializeRequest, responder, _cx| { + let mut response = InitializeResponse::new(request.protocol_version); + if init_agent.advertise_steering { + let mut meta = Meta::new(); + meta.insert( + STEERING_META_KEY.to_string(), + serde_json::json!({ "supported": true }), + ); + response = response.meta(meta); + } + responder.respond(response) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: NewSessionRequest, responder, _cx| { + responder.respond(NewSessionResponse::new("session-1")) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: PromptRequest, responder, cx| { + let text = request + .prompt + .iter() + .filter_map(|block| match block { + AcpContentBlock::Text(text) => Some(text.text.as_str()), + _ => None, + }) + .collect::>() + .join(""); + let ordinal = { + let mut prompts = prompt_agent.prompts.lock().unwrap(); + prompts.push(text); + prompts.len() + }; + if let Some(chunk) = prompt_agent.stream_on_prompt.get(ordinal - 1) { + cx.send_notification(agent_text(Some(&format!("msg-{ordinal}")), chunk))?; + } + if prompt_agent.permission_on_prompt == Some(ordinal) { + let tool_call_id = format!("tc-{ordinal}"); + cx.send_notification(tool_call_notification(&tool_call_id))?; + let agent = Arc::clone(&prompt_agent); + cx.send_request(permission_request(&tool_call_id)) + .on_receiving_result(move |result| async move { + let response = result?; + agent + .permission_outcomes + .lock() + .unwrap() + .push(response.outcome); + agent.permission_answered.notify_one(); + Ok(()) + })?; + } + *prompt_agent.pending_prompt.lock().unwrap() = Some(responder); + prompt_agent.prompt_arrived.notify_one(); + Ok(()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_notification( + async move |_notification: CancelNotification, _cx| { + *cancel_agent.cancels.lock().unwrap() += 1; + // The bridge aborts the turn and answers its prompt. + cancel_agent.finish_prompt(StopReason::Cancelled); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ); + // Registered last: `UntypedMessage` matches every method, so it must + // not shadow the typed handlers above. Left out entirely for an agent + // without the extension, which then answers method-not-found. + builder.on_receive_request( + async move |request: UntypedMessage, responder, _cx| { + assert_eq!(request.method(), STEERING_METHOD, "unexpected request"); + steer_agent + .steers + .lock() + .unwrap() + .push(request.params().clone()); + match &steer_response { + Some(response) => responder.respond(response.clone()), + None => responder + .respond_with_error(agent_client_protocol::Error::method_not_found()), + } + }, + agent_client_protocol::on_receive_request!(), + ) + } +} + +fn steering_setup(steering_supported: bool) -> AcpSessionSetup { + AcpSessionSetup { + agent_session_id: "session-1".to_string(), + agent_capabilities: AgentCapabilities::default(), + task_tracking_mode: TaskTrackingMode::Raw, + auth_methods: vec![], + steering_supported, + } +} + +fn turn_control_handle( + turn_control_tx: &mpsc::UnboundedSender, + alive: bool, +) -> TurnControlHandle { + TurnControlHandle { + turn_control_tx: turn_control_tx.clone(), + connection_alive: Arc::new(AtomicBool::new(alive)), + } +} + +/// The client side of `session/update`, as [`AcpDriver::connect`] wires it. +async fn dispatch_session_update( + handler: &Arc, + notification: IncomingSessionUpdate, +) -> agent_client_protocol::Result<()> { + match notification { + IncomingSessionUpdate::Standard(notification) => { + handler.session_notification(notification).await + } + IncomingSessionUpdate::AsyncTask(notification) => { + handler.async_task_update(notification).await + } + } +} + +/// Whether the handler is in its holding phase right now. `false` while the +/// lock is momentarily held elsewhere — the poll simply tries again. +fn holding(handler: &AcpNotificationHandler) -> bool { + handler + .phase + .try_lock() + .map(|phase| matches!(&*phase, HandlerPhase::BackgroundHolding { .. })) + .unwrap_or(false) +} + +/// Poll until `done` holds — for ordering that depends on the client's +/// dispatch loop catching up with a frame the agent already sent. +async fn wait_until(what: &str, mut done: impl FnMut() -> bool) { + tokio::time::timeout(SETTLE_TIMEOUT, async { + while !done() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {what}")); +} + +// ============================================================================= +// Capability +// ============================================================================= + +#[test] +fn steering_capability_is_a_top_level_meta_flag() { + let supported = |meta: serde_json::Value| { + let serde_json::Value::Object(meta) = meta else { + panic!("test meta must be an object"); + }; + steering_supported_from_initialize(Some(&meta)) + }; + + assert!(supported( + serde_json::json!({ "steering": { "supported": true } }) + )); + // Sibling namespaces (the asyncTasks mirror, the bridge's own) don't + // interfere either way. + assert!(supported(serde_json::json!({ + "claudeCode": { "version": "0.98.0" }, + "jetbrains": { "air": { "version": 1, "capabilities": ["asyncTasks"] } }, + "steering": { "supported": true }, + }))); + + assert!(!steering_supported_from_initialize(None)); + for (label, meta) in [ + ("empty meta", serde_json::json!({})), + ( + "no steering key", + serde_json::json!({ "claudeCode": { "version": "0.70.0" } }), + ), + ( + "explicitly unsupported", + serde_json::json!({ "steering": { "supported": false } }), + ), + ( + "supported is not a bool", + serde_json::json!({ "steering": { "supported": "true" } }), + ), + ( + "steering is not an object", + serde_json::json!({ "steering": true }), + ), + // Inside agentCapabilities' meta would be a different field; the + // flag is a top-level sibling. + ( + "nested under another key", + serde_json::json!({ "agentCapabilities": { "steering": { "supported": true } } }), + ), + ] { + assert!(!supported(meta), "{label} must read as unsupported"); + } +} + +#[test] +fn steering_frame_carries_the_prompt_blocks_and_prompt_required_idle_behavior() { + let message = steering_message( + "session-1", + build_prompt_content_blocks("look at the failing test first", &[], true), + ) + .expect("frame should build"); + assert_eq!(message.method(), "_session/steering"); + assert_eq!( + message.params(), + &serde_json::json!({ + "sessionId": "session-1", + "prompt": [{ "type": "text", "text": "look at the failing test first" }], + "_meta": { "steering": { "idleBehavior": "promptRequired" } }, + }) + ); +} + +#[test] +fn steer_outcome_only_counts_injected_as_delivered() { + assert_eq!( + steer_outcome(Ok(serde_json::json!({ "outcome": "injected" }))), + Ok(SteerOutcome::Injected) + ); + for (response, needle) in [ + ( + serde_json::json!({ "outcome": "promptRequired", "reason": "noRunningTurn" }), + "promptRequired' (noRunningTurn)", + ), + ( + serde_json::json!({ "outcome": "startedNewTurn" }), + "startedNewTurn", + ), + (serde_json::json!({ "outcome": "failed" }), "failed"), + (serde_json::json!({}), "malformed"), + ] { + let error = steer_outcome(Ok(response)).expect_err("not a delivery"); + assert!(error.contains(needle), "got: {error}"); + } + let error = steer_outcome(Err(agent_client_protocol::Error::method_not_found())) + .expect_err("an agent without the extension is not a delivery"); + assert!(error.contains("Steering request failed"), "got: {error}"); +} + +// ============================================================================= +// Steering inside a live turn +// ============================================================================= + +/// Run one `run_prompt_turn` against the fake agent while `drive` exercises +/// the turn-control handle; returns the turn's step and the writer. +async fn live_turn_with( + agent: &Arc, + steering_supported: bool, + drive: F, +) -> (Result, Arc) +where + F: FnOnce(TurnControlHandle, Arc, Arc) -> Fut, + Fut: Future, +{ + let writer = Arc::new(TranscriptWriter::default()); + let handler = Arc::new(transcript_handler(&writer)); + let setup = steering_setup(steering_supported); + let (turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); + let handle = turn_control_handle(&turn_control_tx, true); + let cancel = CancellationToken::new(); + let notification_handler = Arc::clone(&handler); + let step = agent_client_protocol::Client + .builder() + .on_receive_notification( + async move |notification: IncomingSessionUpdate, _cx| { + dispatch_session_update(¬ification_handler, notification).await + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent.build(), async |connection| { + let (step, ()) = tokio::join!( + run_prompt_turn( + &connection, + &setup, + &handler, + "local-session", + &cancel, + &mut turn_control_rx, + "first", + &[], + ), + drive(handle, Arc::clone(agent), Arc::clone(&writer)), + ); + Ok(step) + }) + .await + .expect("the connection itself stays healthy"); + (step, writer) +} + +#[tokio::test(flavor = "current_thread")] +async fn steer_in_a_live_turn_sends_the_extension_request_and_persists_the_user_row() { + let agent = Arc::new(FakeAgent { + steer_response: Some(serde_json::json!({ "outcome": "injected" })), + advertise_steering: true, + stream_on_prompt: vec!["working on it"], + ..FakeAgent::default() + }); + + let (step, writer) = live_turn_with(&agent, true, |handle, agent, writer| async move { + agent.await_prompt().await; + wait_until("the first chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("working on it".to_string())) + }) + .await; + + let outcome = handle + .steer( + "also run the linter", + &[PromptImage::new("aGk=", "image/png")], + ) + .await; + assert_eq!(outcome, Ok(SteerOutcome::Injected)); + // The user row is written in the ordered response callback: after + // the open assistant message is closed, before anything else from the + // agent can be dispatched. The turn itself is still in flight. + assert_eq!( + writer.text_flow(), + vec![ + WriterCall::Append("working on it".to_string()), + WriterCall::Finalize, + WriterCall::UserMessage("also run the linter".to_string()), + ] + ); + assert!(agent.pending_prompt.lock().unwrap().is_some()); + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + assert!(matches!( + step, + Ok(TurnStep::Settled(AgentRunOutcome::Completed)) + )); + assert_eq!(agent.cancels(), 0, "steering never cancels the turn"); + let steers = agent.steers(); + assert_eq!(steers.len(), 1); + assert_eq!(steers[0]["sessionId"], "session-1"); + assert_eq!(steers[0]["prompt"][0]["text"], "also run the linter"); + // No image capability on the fake agent: the attachment is noted, not sent. + assert!(steers[0]["prompt"][1]["text"] + .as_str() + .is_some_and(|text| text.contains("1 image block(s) omitted"))); + assert_eq!( + steers[0]["_meta"]["steering"]["idleBehavior"], + "promptRequired" + ); + // Nothing more than one row: the agent's echo of the steered message is + // not a session update this client records. + assert_eq!( + writer + .calls() + .iter() + .filter(|call| matches!(call, WriterCall::UserMessage(_))) + .count(), + 1 + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn steer_answered_prompt_required_is_an_error_without_a_user_row() { + let agent = Arc::new(FakeAgent { + steer_response: Some( + serde_json::json!({ "outcome": "promptRequired", "reason": "noRunningTurn" }), + ), + advertise_steering: true, + ..FakeAgent::default() + }); + + let (step, writer) = live_turn_with(&agent, true, |handle, agent, _writer| async move { + agent.await_prompt().await; + let error = handle + .steer("too late", &[]) + .await + .expect_err("promptRequired is not a delivery"); + assert!(error.contains("promptRequired"), "got: {error}"); + assert!(error.contains("noRunningTurn"), "got: {error}"); + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + assert!(matches!( + step, + Ok(TurnStep::Settled(AgentRunOutcome::Completed)) + )); + assert_eq!(agent.steers().len(), 1); + // Leaves nothing behind for the caller's fallback delivery to duplicate. + assert_eq!(writer.text_flow(), vec![]); +} + +#[tokio::test(flavor = "current_thread")] +async fn steer_to_an_agent_without_the_method_is_an_error() { + // Advertises steering but doesn't serve the method: the request itself + // fails (method-not-found), which must reach the caller as an error + // rather than hang the turn's select. + let agent = Arc::new(FakeAgent { + steer_response: None, + advertise_steering: true, + ..FakeAgent::default() + }); + + let (step, writer) = live_turn_with(&agent, true, |handle, agent, _writer| async move { + agent.await_prompt().await; + let error = handle + .steer("hello?", &[]) + .await + .expect_err("method-not-found is not a delivery"); + assert!(error.contains("Steering request failed"), "got: {error}"); + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + assert!(matches!( + step, + Ok(TurnStep::Settled(AgentRunOutcome::Completed)) + )); + assert_eq!(writer.text_flow(), vec![]); +} + +#[tokio::test(flavor = "current_thread")] +async fn steer_is_refused_locally_when_the_agent_did_not_advertise_it() { + let agent = Arc::new(FakeAgent { + steer_response: Some(serde_json::json!({ "outcome": "injected" })), + advertise_steering: false, + ..FakeAgent::default() + }); + + let (step, _writer) = live_turn_with(&agent, false, |handle, agent, _writer| async move { + agent.await_prompt().await; + let error = handle + .steer("hello?", &[]) + .await + .expect_err("no capability, no steer"); + assert!( + error.contains("does not advertise steering"), + "got: {error}" + ); + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + assert!(matches!( + step, + Ok(TurnStep::Settled(AgentRunOutcome::Completed)) + )); + assert!( + agent.steers().is_empty(), + "the gate must stop the frame before the wire" + ); +} + +// ============================================================================= +// Interrupt-and-reprompt on the same connection +// ============================================================================= + +/// Drive a whole `run_acp_session` (setup, prompt loop) against the fake +/// agent. `drive` runs alongside it once setup has completed; it receives the +/// handle, the agent, the writer and the handler. Returns the session's +/// result, the reply receiver of the first `prompt()`, and whatever reply +/// sender is still stashed when the loop returns — the one the connection +/// task would resolve after teardown. +#[allow(clippy::type_complexity)] +async fn session_with( + agent: &Arc, + background_hold: Option<&BackgroundHoldConfig>, + drive: F, +) -> ( + Result<(AgentRunOutcome, SessionSettleReason), String>, + oneshot::Receiver, + Option>, + Arc, + Arc, +) +where + F: FnOnce( + TurnControlHandle, + Arc, + Arc, + Arc, + ) -> Fut, + Fut: Future, +{ + let writer = Arc::new(TranscriptWriter::default()); + let handler = Arc::new(transcript_handler(&writer)); + let store: Arc = Arc::new(RecordingStore::default()); + let cancel = CancellationToken::new(); + let child_exited = CancellationToken::new(); + let (prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); + let (_hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); + let hold_active = AtomicBool::new(false); + let pending_reply = RefCell::new(None); + let handle = turn_control_handle(&turn_control_tx, true); + + // The first turn, queued the way `SessionConnection::prompt` queues it. + let (reply_tx, reply_rx) = oneshot::channel(); + prompt_tx + .send(QueuedSessionTurn { + prompt: "first".to_string(), + images: vec![], + reply: Some(reply_tx), + control: None, + }) + .expect("queue should accept"); + + let permission_handler = Arc::clone(&handler); + let notification_handler = Arc::clone(&handler); + let result = agent_client_protocol::Client + .builder() + .on_receive_request( + async move |args: RequestPermissionRequest, responder, _connection| { + let response = permission_handler.request_permission(args).await?; + responder.respond(response) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_notification( + async move |notification: IncomingSessionUpdate, _cx| { + dispatch_session_update(¬ification_handler, notification).await + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent.build(), async |connection| { + let (result, ()) = tokio::join!( + run_acp_session( + &connection, + Path::new("/tmp"), + &store, + "local-session", + None, + &[], + &handler, + &[], + "Fake", + &cancel, + None, + background_hold, + None, + &child_exited, + &mut prompt_rx, + &mut hold_control_rx, + &mut turn_control_rx, + &hold_active, + &pending_reply, + ), + drive( + handle, + Arc::clone(agent), + Arc::clone(&writer), + Arc::clone(&handler) + ), + ); + Ok(result) + }) + .await + .expect("the connection itself stays healthy"); + let stashed_reply = pending_reply.borrow_mut().take(); + (result, reply_rx, stashed_reply, writer, handler) +} + +#[tokio::test(flavor = "current_thread")] +async fn interrupt_cancels_the_turn_and_reprompts_on_the_same_connection() { + let agent = Arc::new(FakeAgent { + advertise_steering: false, + stream_on_prompt: vec!["half an answer", "the new answer"], + // Turn 2 asks permission: the connection-wide permission token must + // have survived the interrupt for the handler to answer it. + permission_on_prompt: Some(2), + ..FakeAgent::default() + }); + + let (result, mut reply_rx, stashed_reply, writer, handler) = + session_with(&agent, None, |handle, agent, writer, _handler| async move { + agent.await_prompt().await; + wait_until("the first chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("half an answer".to_string())) + }) + .await; + + // The agent's cancel handler answers the pending prompt + // `cancelled`; the interrupt then sends the new message. + handle + .interrupt_and_reprompt("actually, do this instead", &[]) + .await + .expect("the agent acknowledged the cancel"); + agent.await_prompt().await; + assert_eq!(agent.prompts(), ["first", "actually, do this instead"]); + assert_eq!(agent.cancels(), 1); + + tokio::time::timeout(SETTLE_TIMEOUT, agent.permission_answered.notified()) + .await + .expect("turn 2's permission request should be answered"); + wait_until("the second chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("the new answer".to_string())) + }) + .await; + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + // Without a background hold the session settles right after the *final* + // turn, as a completed turn — the interrupted one never produced an + // outcome of its own. + assert_eq!( + result, + Ok((AgentRunOutcome::Completed, SessionSettleReason::Immediate)) + ); + // The transcript: the interrupted assistant message is closed, the new + // user message recorded, and only then does the next answer stream. + assert_eq!( + writer.text_flow(), + vec![ + WriterCall::Append("half an answer".to_string()), + WriterCall::Finalize, + WriterCall::UserMessage("actually, do this instead".to_string()), + WriterCall::Append("the new answer".to_string()), + ] + ); + // The first prompt()'s reply was never resolved by the interrupt: it is + // still stashed for the connection task to answer at teardown with the + // final turn's result — from the caller's side the session simply kept + // running. + assert!(matches!( + reply_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + let stashed_reply = stashed_reply.expect("the first turn's reply is still pending"); + stashed_reply + .send(result.clone().map(|(outcome, _)| outcome)) + .expect("the prompt() caller is still waiting"); + assert_eq!(reply_rx.await, Ok(Ok(AgentRunOutcome::Completed))); + // The interrupt did not burn the connection's permission token: turn 2's + // request reached the writer (which auto-approves) instead of being + // auto-answered `cancelled`. + assert!(!handler.permission_cancel_token.is_cancelled()); + let outcomes = agent.permission_outcomes.lock().unwrap(); + assert_eq!(outcomes.len(), 1); + assert!( + matches!(outcomes[0], RequestPermissionOutcome::Selected(_)), + "got: {:?}", + outcomes[0] + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn interrupt_during_a_hold_is_served_as_a_new_turn_without_a_cancel() { + let agent = Arc::new(FakeAgent { + advertise_steering: true, + steer_response: Some(serde_json::json!({ "outcome": "injected" })), + ..FakeAgent::default() + }); + + let (result, _reply_rx, _stashed_reply, writer, _handler) = session_with( + &agent, + Some(&QUIESCENCE_PROBE), + |handle, agent, writer, handler| async move { + agent.await_prompt().await; + // A live background task keeps the hold open after turn 1. + feed_sdk_frame(&handler, task_started("task-1")).await; + agent.finish_prompt(StopReason::EndTurn); + wait_until("the hold", || holding(&handler)).await; + + // Nothing is in flight to steer: the message becomes a turn of + // its own, with no `_session/steering` frame on the wire. + let outcome = handle.steer("while you wait", &[]).await; + assert_eq!(outcome, Ok(SteerOutcome::NewTurnStarted)); + agent.await_prompt().await; + assert_eq!(agent.prompts(), ["first", "while you wait"]); + assert!(agent.steers().is_empty()); + assert_eq!( + writer.text_flow(), + vec![ + WriterCall::Finalize, + WriterCall::UserMessage("while you wait".to_string()) + ] + ); + agent.finish_prompt(StopReason::EndTurn); + + // And an interrupt with nothing to interrupt is the same. + wait_until("the second hold", || holding(&handler)).await; + handle + .interrupt_and_reprompt("and this", &[]) + .await + .expect("nothing to cancel, served as a turn"); + agent.await_prompt().await; + assert_eq!(agent.cancels(), 0); + assert_eq!(agent.prompts().len(), 3); + agent.finish_prompt(StopReason::EndTurn); + + // Drain the task set so the hold settles on its own. + wait_until("the third hold", || holding(&handler)).await; + feed_sdk_frame(&handler, empty_tasks_frame()).await; + }, + ) + .await; + + assert_eq!( + result, + Ok((AgentRunOutcome::Completed, SessionSettleReason::Quiescent)) + ); + assert_eq!( + writer + .calls() + .iter() + .filter(|call| matches!(call, WriterCall::UserMessage(_))) + .count(), + 2 + ); +} + +#[tokio::test] +async fn holding_wait_turns_a_steer_into_a_new_turn_without_rejecting_it_at_boundaries() { + let handler = hold_test_handler(); + let (_prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); + let (_hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); + let (turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); + let child_exited = CancellationToken::new(); + let hold_active = AtomicBool::new(false); + let handle = turn_control_handle(&turn_control_tx, true); + + feed_sdk_frame(&handler, task_started("task-1")).await; + + // Queued before the hold exists. Unlike a per-task stop, a message is not + // aimed at one hold: the entry boundary must not reject it. + let (early_reply_tx, early_reply_rx) = oneshot::channel(); + turn_control_tx + .send(TurnControlRequest::Steer { + prompt: "queued early".to_string(), + images: vec![], + reply: early_reply_tx, + }) + .expect("queue should accept"); + + let outcome = tokio::time::timeout( + SETTLE_TIMEOUT, + hold_for_background_quiescence( + &QUIESCENCE_PROBE, + TaskTrackingMode::Raw, + &handler, + &cancel, + &child_exited, + &mut prompt_rx, + &mut hold_control_rx, + &mut turn_control_rx, + &hold_active, + &|_| {}, + "sess-1", + None, + ), + ) + .await + .expect("a queued message ends the hold at once"); + let HoldOutcome::NewTurn(turn) = outcome else { + panic!("expected NewTurn, got {outcome:?}"); + }; + assert_eq!(turn.prompt, "queued early"); + assert!(turn.reply.is_none(), "no prompt() is waiting on this turn"); + // Answered by the session loop once the turn goes out, not by the hold. + assert!(matches!(turn.control, Some(TurnControlReply::Steer(_)))); + let mut early_reply_rx = early_reply_rx; + assert!(matches!( + early_reply_rx.try_recv(), + Err(oneshot::error::TryRecvError::Empty) + )); + turn.control.expect("control reply").turn_started(); + assert_eq!( + early_reply_rx.await.expect("answered"), + Ok(SteerOutcome::NewTurnStarted) + ); + + // A handle-issued interrupt during the hold reads the same way. + let served = async { + let outcome = hold_for_background_quiescence( + &QUIESCENCE_PROBE, + TaskTrackingMode::Raw, + &handler, + &cancel, + &child_exited, + &mut prompt_rx, + &mut hold_control_rx, + &mut turn_control_rx, + &hold_active, + &|_| {}, + "sess-1", + None, + ) + .await; + let HoldOutcome::NewTurn(turn) = outcome else { + panic!("expected NewTurn, got {outcome:?}"); + }; + turn.control.expect("control reply").turn_started(); + turn.prompt + }; + let (prompt, answer) = tokio::join!(served, handle.interrupt_and_reprompt("now", &[])); + assert_eq!(prompt, "now"); + assert_eq!(answer, Ok(())); +} + +// ============================================================================= +// Teardown and handle failure mapping +// ============================================================================= + +#[tokio::test] +async fn teardown_answers_every_queued_message_and_the_handle_fails_fast_after() { + let (turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); + let alive = Arc::new(AtomicBool::new(true)); + let handle = TurnControlHandle { + turn_control_tx: turn_control_tx.clone(), + connection_alive: Arc::clone(&alive), + }; + + let (steer_reply_tx, steer_reply_rx) = oneshot::channel(); + let (interrupt_reply_tx, interrupt_reply_rx) = oneshot::channel(); + turn_control_tx + .send(TurnControlRequest::Steer { + prompt: "a".to_string(), + images: vec![], + reply: steer_reply_tx, + }) + .expect("queue should accept"); + turn_control_tx + .send(TurnControlRequest::Interrupt { + prompt: "b".to_string(), + images: vec![], + reply: interrupt_reply_tx, + }) + .expect("queue should accept"); + + // The connection task's teardown sequence. + alive.store(false, Ordering::Relaxed); + reject_queued_turn_control_requests(&mut turn_control_rx, "the ACP connection closed"); + + let steer = steer_reply_rx + .await + .expect("every queued request is answered"); + assert!( + steer + .as_ref() + .is_err_and(|e| e.contains("connection closed")), + "got: {steer:?}" + ); + let interrupt = interrupt_reply_rx + .await + .expect("every queued request is answered"); + assert!( + interrupt + .as_ref() + .is_err_and(|e| e.contains("connection closed")), + "got: {interrupt:?}" + ); + assert!( + turn_control_rx.try_recv().is_err(), + "the queue must be drained" + ); + + // After settle the handle refuses without queueing. + let error = handle + .steer("c", &[]) + .await + .expect_err("connection is gone"); + assert!(error.contains("no longer running"), "got: {error}"); + let error = handle + .interrupt_and_reprompt("d", &[]) + .await + .expect_err("connection is gone"); + assert!(error.contains("no longer running"), "got: {error}"); + assert!(turn_control_rx.try_recv().is_err()); +} + +#[tokio::test] +async fn a_turn_control_handle_maps_channel_failures_to_connection_state() { + let (turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); + let handle = turn_control_handle(&turn_control_tx, true); + + // A reply dropped unanswered (teardown discarded the scheduled callback) + // reads as the connection closing mid-request. + let drop_unanswered = async { + let request = turn_control_rx.recv().await.expect("request should arrive"); + drop(request); + }; + let (steered, ()) = tokio::join!(handle.steer("a", &[]), drop_unanswered); + let error = steered.expect_err("a dropped reply is not an answer"); + assert!(error.contains("closed before"), "got: {error}"); + + // With the receiver gone the connection itself is gone, whatever the + // alive flag still says. + drop(turn_control_rx); + let error = handle + .interrupt_and_reprompt("b", &[]) + .await + .expect_err("no connection to serve the message"); + assert!(error.contains("no longer running"), "got: {error}"); +} From 71f7d193469eed8474cca15e8bff748792711ed7 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 28 Oct 2026 17:51:15 +1100 Subject: [PATCH 2/7] feat(staged): register the turn-control handle and persist connection-delivered user rows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second step of timely message delivery to running repo sessions: the runner-side plumbing around the acp-client turn-control channel from 7a219c29, so a later caller (the `continue_repo_session` MCP tool, a UI command) can reach a running session's connection from another thread and the message it delivers lands in the transcript. - `SessionRegistry`: `RunningSession` carries `turn_control: Mutex>` next to the async-task stop handle, with `set_turn_control_handle` / `turn_control_handle(session_id)` mirroring the stop handle's accessors (`None` before `connect`, for unknown sessions, and after deregister). Registered at the one `driver.connect` site in the turn loop, right after the stop handle — the Pikchr-correction re-connect goes through the same site, so it replaces the handle the same way. - `agent::writer::MessageWriter::record_user_message(text, images) -> Option`: finalizes the open assistant text, inserts a plain `user` row through `Store::add_session_message`, and returns its id. The `PromptImage`s the connection sent carry bytes, not image ids, so the row has no `image_ids`; the caller that owns stored images can attach them to the returned row. The `acp_client::MessageWriter` trait impl forwards to it. - acp-client: export `TurnControlHandle` and `SteerOutcome` from the crate root (the runner could not name the type otherwise), and give the handle `is_connection_alive()` (the fail-fast flag, for routing), `same_connection(&other)` (channel identity, so a re-connect's new handle is distinguishable from a clone), and a `#[doc(hidden)] detached()` placeholder whose requests fail fast — the test seam for plumbing that holds a handle without a live connection. Tests: registry handle lifecycle (absent before connect, set, replaced on re-connect, no-op for an unregistered id, gone after deregister); writer `finalize → record_user_message → append_text` yields assistant/user/assistant rows in order with the returned id on the user row and no event kind or images; the trait entry point closes unflushed assistant text itself. Verified: `cargo build`, `cargo test --lib` for src-tauri (writer and session_runner modules) and acp-client, clippy clean at CI settings (`just lint` omits `--tests`; the three `--all-targets` lints left are pre-existing in untouched code), rustfmt applied. Not in this commit: the MCP tool and the shared `continue_session_for_store` body (commit 3). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/agent/writer.rs | 107 +++++++++++++++++++- apps/staged/src-tauri/src/session_runner.rs | 76 +++++++++++++- crates/acp-client/src/driver.rs | 26 +++++ crates/acp-client/src/lib.rs | 3 +- 4 files changed, 209 insertions(+), 3 deletions(-) diff --git a/apps/staged/src-tauri/src/agent/writer.rs b/apps/staged/src-tauri/src/agent/writer.rs index 363dba370..0358a1309 100644 --- a/apps/staged/src-tauri/src/agent/writer.rs +++ b/apps/staged/src-tauri/src/agent/writer.rs @@ -20,7 +20,7 @@ use crate::store::{AcpMessageMetadata, MessageRole, Store}; use acp_client::{ autoapprove_permission_decision, strip_code_fences, AcpEventMetadata, AcpInitializeMetadata, - AcpPermissionDecision, AcpPermissionRequest, AcpToolCallMetadata, + AcpPermissionDecision, AcpPermissionRequest, AcpToolCallMetadata, PromptImage, }; use agent_client_protocol::schema::MaybeUndefined; @@ -164,6 +164,42 @@ impl MessageWriter { *self.text_interrupted.lock().await = false; } + /// Persist a user message the ACP connection delivered itself — steered + /// into the turn in flight, or sent as the prompt after an interrupt — + /// as an ordinary `user` row, and return its id. + /// + /// Finalizes first so the row lands after whatever the assistant had + /// streamed so far and the next [`append_text`](Self::append_text) opens + /// a fresh assistant row: the transcript reads assistant → user → + /// assistant. The connection finalizes before calling too; doing it here + /// as well costs nothing (finalize is idempotent) and keeps the ordering + /// a property of this method rather than of its caller. + /// + /// The `images` are the blocks the connection sent, not stored images — + /// acp-client has no image ids — so the row carries no `image_ids`. A + /// caller that owns stored images attaches them to the row it gets back. + pub async fn record_user_message(&self, text: &str, images: &[PromptImage]) -> Option { + self.finalize().await; + if !images.is_empty() { + log::debug!( + "Session {}: user message delivered mid-turn with {} image block(s); the \ + persisted row records the text only", + self.session_id, + images.len() + ); + } + match self + .store + .add_session_message(&self.session_id, MessageRole::User, text) + { + Ok(id) => Some(id), + Err(e) => { + log::error!("Failed to insert delivered user message: {e}"); + None + } + } + } + // ===================================================================== // Tool calls // ===================================================================== @@ -310,6 +346,10 @@ impl acp_client::MessageWriter for MessageWriter { self.finalize().await } + async fn record_user_message(&self, text: &str, images: &[PromptImage]) { + self.record_user_message(text, images).await; + } + async fn should_finish_after_response(&self) -> bool { // Inspect only the currently streamed response, never older transcript // rows. The ACP caller invokes this only at a completed response boundary. @@ -570,6 +610,71 @@ mod tests { } } + /// A message the connection delivers mid-session (steer or interrupt) is + /// persisted between the assistant text before it and the reply after it, + /// as three rows in transcript order, and the reply does not extend the + /// row that was open when the message arrived. + #[tokio::test] + async fn record_user_message_lands_between_assistant_rows() { + let (store, session_id, writer) = setup_writer(); + + writer.append_text("Working on the first ").await; + writer.append_text("instruction.").await; + writer.finalize().await; + let user_row = writer + .record_user_message("Actually, do the second thing instead.", &[]) + .await + .expect("user row inserted"); + writer.append_text("Switching to the second thing.").await; + writer.finalize().await; + + let messages = store.get_session_messages(&session_id).unwrap(); + let rows: Vec<(MessageRole, &str)> = messages + .iter() + .map(|m| (m.role.clone(), m.content.as_str())) + .collect(); + assert_eq!( + rows, + vec![ + (MessageRole::Assistant, "Working on the first instruction."), + (MessageRole::User, "Actually, do the second thing instead."), + (MessageRole::Assistant, "Switching to the second thing."), + ] + ); + assert_eq!(messages[1].id, user_row); + assert!(messages[1].image_ids.is_empty()); + assert!(messages[1].acp.acp_event_kind.is_none()); + } + + /// The trait entry point (what the connection task calls) closes any + /// still-open assistant text itself, so a caller that skipped the + /// finalize still gets the row order right; image blocks are accepted + /// and the text is what gets persisted. + #[tokio::test] + async fn record_user_message_closes_open_assistant_text_first() { + let (store, session_id, writer) = setup_writer(); + + writer.append_text("Unflushed partial text").await; + let image = acp_client::PromptImage::new("AAAA", "image/png"); + acp_client::MessageWriter::record_user_message(&writer, "Look at this.", &[image]).await; + writer.append_text("Reply.").await; + writer.finalize().await; + + let messages = store.get_session_messages(&session_id).unwrap(); + let rows: Vec<(MessageRole, &str)> = messages + .iter() + .map(|m| (m.role.clone(), m.content.as_str())) + .collect(); + assert_eq!( + rows, + vec![ + (MessageRole::Assistant, "Unflushed partial text"), + (MessageRole::User, "Look at this."), + (MessageRole::Assistant, "Reply."), + ] + ); + } + #[tokio::test] async fn ordinary_session_does_not_finish_for_note_shaped_output() { use acp_client::MessageWriter as _; diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 88afe95be..fe0070271 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -266,6 +266,11 @@ struct RunningSession { /// re-connects). `None` until then, or for sessions that never connect /// (e.g. externally registered ones). async_task_stop: std::sync::Mutex>, + /// Turn-control handle for the session's current ACP connection — how a + /// message reaches the agent *now* (steered into the turn in flight, or + /// sent in place of it) rather than after the turn. Set and replaced + /// alongside `async_task_stop`. + turn_control: std::sync::Mutex>, /// The last background hold the session's connection reported, so a client /// that mounts mid-hold can ask for the wait it missed. The hold is /// reported on change only; the driver reports the cleared default on the @@ -444,6 +449,7 @@ impl SessionRegistry { cancellation_completion_reason: std::sync::Mutex::new(None), finished: std::sync::atomic::AtomicBool::new(false), async_task_stop: std::sync::Mutex::new(None), + turn_control: std::sync::Mutex::new(None), background_hold: std::sync::Mutex::new(acp_client::BackgroundHoldStatus::default()), }); let pending = inner.pending_cancellations.remove(session_id); @@ -790,6 +796,27 @@ impl SessionRegistry { .and_then(|running| running.async_task_stop.lock().unwrap().clone()) } + /// Store the turn-control handle for `session_id`'s current ACP + /// connection, replacing the previous connection's handle (a Pikchr + /// correction re-connects on the same session). + fn set_turn_control_handle(&self, session_id: &str, handle: acp_client::TurnControlHandle) { + if let Some(running) = self.inner.lock().unwrap().running.get(session_id) { + *running.turn_control.lock().unwrap() = Some(handle); + } + } + + /// The turn-control handle for `session_id`'s current ACP connection — + /// `None` when the session isn't running or hasn't connected yet (the + /// startup window between `register` and `connect`). + pub fn turn_control_handle(&self, session_id: &str) -> Option { + self.inner + .lock() + .unwrap() + .running + .get(session_id) + .and_then(|running| running.turn_control.lock().unwrap().clone()) + } + /// Retain the hold `session_id`'s connection just reported, so a client /// that mounts after it was emitted can still learn about the wait. /// Recorded before the event is emitted, so a client reacting to the event @@ -1754,11 +1781,18 @@ pub fn start_session( // Expose the connection's per-task stop where the // `stop_session_async_task` command (another thread) // can reach it; requests are only served while the - // connection holds for background work. + // connection holds for background work. The + // turn-control handle goes up next to it so a message + // for this session can be steered into (or sent in + // place of) the turn in flight from another thread. registry.set_async_task_stop_handle( &config.session_id, connection.async_task_stop_handle(), ); + registry.set_turn_control_handle( + &config.session_id, + connection.turn_control_handle(), + ); let settled_rx = connection.take_settled_receiver(); let turn_result = connection.prompt(&prompt, images).await; match settled_rx { @@ -6083,6 +6117,46 @@ mod tests { ); } + /// The turn-control handle follows the connection, not the registration: + /// absent in the startup window before `connect`, replaced when a Pikchr + /// correction re-connects on the same session, and gone with the entry. + /// Reads for an unknown session answer `None` rather than erroring, since + /// the caller routes on that (a message for a not-yet-connected session + /// waits, one for a session that is not running takes the resume path). + #[test] + fn turn_control_handle_follows_the_connection() { + let registry = SessionRegistry::new(); + let (_, epoch) = registry.register("session-steered"); + + // Registered but not connected yet. + assert!(registry.turn_control_handle("session-steered").is_none()); + + let first = acp_client::TurnControlHandle::detached(); + registry.set_turn_control_handle("session-steered", first.clone()); + let read = registry + .turn_control_handle("session-steered") + .expect("handle set once the connection is up"); + assert!(read.same_connection(&first)); + + // A re-connect hands out the new connection's handle. + let second = acp_client::TurnControlHandle::detached(); + registry.set_turn_control_handle("session-steered", second.clone()); + let read = registry + .turn_control_handle("session-steered") + .expect("handle replaced on re-connect"); + assert!(read.same_connection(&second)); + assert!(!read.same_connection(&first)); + + // Setting for a session that isn't registered is a no-op... + registry + .set_turn_control_handle("session-unknown", acp_client::TurnControlHandle::detached()); + assert!(registry.turn_control_handle("session-unknown").is_none()); + + // ...and the handle leaves with the entry. + registry.deregister("session-steered", epoch); + assert!(registry.turn_control_handle("session-steered").is_none()); + } + /// Every way into a session but the pipeline handoff arrives at /// `start_session` with nothing registered, so an ordinary refusal has only /// the quit's own reason to record. diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index d3bfd8ec2..cb864afd8 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -1644,6 +1644,32 @@ impl TurnControlHandle { .map_err(|_| "ACP connection closed before the interrupt resolved".to_string())? } + /// Whether the connection behind this handle is still running. A request + /// made while this is `false` fails fast; `true` is advisory — the + /// connection can settle between the check and the request. + pub fn is_connection_alive(&self) -> bool { + self.connection_alive.load(Ordering::Relaxed) + } + + /// Whether `other` reaches the same connection as this handle (clones of + /// one connection's handle do; a re-connect on the same session hands out + /// a different one). + pub fn same_connection(&self, other: &Self) -> bool { + self.turn_control_tx.same_channel(&other.turn_control_tx) + } + + /// A handle to no connection: every request fails fast with the + /// connection-gone error. For registries that need a placeholder and for + /// tests of the plumbing around a handle. + #[doc(hidden)] + pub fn detached() -> Self { + let (turn_control_tx, _) = mpsc::unbounded_channel(); + Self { + turn_control_tx, + connection_alive: Arc::new(AtomicBool::new(false)), + } + } + fn send(&self, request: TurnControlRequest) -> Result<(), String> { if !self.connection_alive.load(Ordering::Relaxed) { return Err("ACP connection is no longer running".to_string()); diff --git a/crates/acp-client/src/lib.rs b/crates/acp-client/src/lib.rs index f2084d29a..f90578324 100644 --- a/crates/acp-client/src/lib.rs +++ b/crates/acp-client/src/lib.rs @@ -36,7 +36,8 @@ pub use driver::{ AgentRunOutcome, AsyncTaskStopHandle, BackgroundHoldConfig, BackgroundHoldObserver, BackgroundHoldStatus, BackgroundHoldTask, BasicMessageWriter, MessageWriter, OutOfTurnPermissionPolicy, PromptImage, ReplayBoundary, SessionConnection, SessionLifetime, - SessionSettleReason, SessionSettled, Store, BACKGROUND_CONTINUATION_ORIGIN, + SessionSettleReason, SessionSettled, SteerOutcome, Store, TurnControlHandle, + BACKGROUND_CONTINUATION_ORIGIN, }; pub use simple::{run_acp_prompt, run_acp_prompt_with_interpreter_env_snapshot}; pub use types::{ From 13de408161a219d473f3116d939eff39d0f64df5 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 28 Oct 2026 18:07:19 +1100 Subject: [PATCH 3/7] feat(staged): continue_repo_session tool and the shared continue-session body MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third step of timely message delivery to running repo sessions, on top of the acp-client turn-control channel (7a219c29) and the registry handle plumbing (71f7d193): a project session can now send a follow-up to a repo session, and a running one receives it within seconds without a stop. - `session_commands::continue_session_for_store(...) -> ContinueOutcome` routes on the row's status: `queued` → `RejectedQueued` (never claimed out of FIFO order); `running` → the registry's turn-control handle, `steer()` first, `interrupt_and_reprompt()` when the agent refuses (`Injected` / `StartedNewTurn` / `InterruptedAndReprompted`); `running` with no handle yet → poll the registry for up to 10 s (`CONTINUE_SESSION_STARTUP_WAIT`), else `Blocked`; terminal → the resume gate (`continue_resume_gate`: a resumable `agent_id`, the working dir on disk for local branches, `should_queue_branch_session_start` under `branch_session_launch_lock_for`) then `resume_session_for_store`, with the already-running refusal read back as `Blocked` rather than an error. Nothing here persists the user row: acp-client writes it for the live paths (only once the agent took the message) and `start_session` for the resume, so it lands in `session_messages` exactly once on every path. - `continue_repo_session { repo_session_id, message }` on `ProjectToolsHandler`, next to `cancel_repo_session`: the usual `repo_session_payload` plus `delivery` and, for `blocked` / `not_resumable`, `reason`. The activity payload gains `last_user_message_at` so a parent can confirm the message landed. The project-session system prompt's tool list (both the local and the remote-workspace variant) describes the new tool and when to prefer it over cancelling and starting over. - `session_runner`: `load_prompt_attachments` extracted from `start_session` so the live paths build the same prompt text and image blocks for a message with attachments; `set_turn_control_handle` and the `attachments` module widened to `pub(crate)`. - acp-client: `TurnControlHandle::scripted(steer, interrupt)`, a `#[doc(hidden)]` stand-in connection with canned answers and a request log, for tests of the routing around a handle. Tests: `route_continue_session` across queued / running with a handle (injected, new turn, steer refused → interrupt, interrupt refused → blocked) / running without a handle (times out blocked, delivers once the handle lands, re-routes when the row settles mid-wait) / terminal (no agent id, worktree gone, busy branch, provisioning branch, resumable → Resume); the `delivery` spellings; the description test for the new tool; and `project_mcp` round trips through the tool's body for the steer, interrupt and still-starting cases asserting the row never leaves `running`. Where this follows the code rather than the plan note: - The body is split into `route_continue_session` (everything but the resume) and the concrete `continue_session_for_store`, because `resume_session_for_store` needs a Wry `AppHandle` that tests cannot build; the MCP round trips drive the route plus `continue_repo_session_payload`, the same two halves the tool calls. - No `persist_user_message` retry flag: acp-client only records the row when the steer is `injected` (or the fallback turn starts), so a refused steer leaves nothing for the interrupt to duplicate. - An interrupt served from the hold is reported `interrupted_and_reprompted` — `interrupt_and_reprompt()` does not say whether it had a turn to cancel. - The branch gate is checked under the launch lock but the lock is not held into the resume's `running` write (the future must stay `Send`); that write is atomic on the status, so a second continue on the same branch comes back `blocked`, not doubled. - Pipeline (pr/push) sessions carry no `BranchSessionType`, so the resume gate checks only their worktree and agent id. Verified: `cargo build`, `cargo test --lib` for session_commands, project_mcp, agent::writer and session_runner (339 passed) and acp-client's turn_control tests, clippy clean at CI settings (`just lint` omits `--tests`; the three `--all-targets` lints left are pre-existing in untouched code), rustfmt applied. Not in this commit: the UI command and composer split (commit 4). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/project_mcp.rs | 272 +++++- apps/staged/src-tauri/src/session_commands.rs | 792 +++++++++++++++++- apps/staged/src-tauri/src/session_runner.rs | 85 +- crates/acp-client/src/driver.rs | 37 + 4 files changed, 1144 insertions(+), 42 deletions(-) diff --git a/apps/staged/src-tauri/src/project_mcp.rs b/apps/staged/src-tauri/src/project_mcp.rs index 0b1afd23e..c470ff986 100644 --- a/apps/staged/src-tauri/src/project_mcp.rs +++ b/apps/staged/src-tauri/src/project_mcp.rs @@ -1,5 +1,6 @@ //! MCP server for project sessions. -//! Exposes `start_repo_session` and `add_project_repo` tools to the agent. +//! Exposes the repo-session tools (`start_repo_session`, `wait_for_repo_session`, +//! `cancel_repo_session`, `continue_repo_session`) and `add_project_repo` to the agent. use std::sync::Arc; use tokio::task::JoinHandle; @@ -17,6 +18,7 @@ use tauri::AppHandle; use crate::actions::{ActionExecutor, ActionRegistry}; use crate::mcp_progress::send_progress_keepalives; +use crate::session_commands::ContinueOutcome; use crate::session_runner::SessionRegistry; use crate::store::{ AcpConfigSelection, Branch, CompletionReason, MessageRole, ProjectRepo, Session, @@ -125,6 +127,16 @@ struct CancelRepoSessionParams { pub repo_session_id: String, } +#[derive(serde::Deserialize, schemars::JsonSchema)] +struct ContinueRepoSessionParams { + /// Opaque handle returned by `start_repo_session`. + pub repo_session_id: String, + /// The follow-up for the repo session's agent: a correction, the answer to + /// a question it raised, or its next step. Written as an instruction to + /// that agent, not as a report to the user. + pub message: String, +} + fn default_wait_for_completion_seconds() -> u64 { 240 } @@ -134,6 +146,10 @@ const REPO_SESSION_ACTIVITY_PREVIEW_MAX_CHARS: usize = 240; #[derive(Debug, Default, serde::Serialize)] struct RepoSessionActivity { last_activity_at: Option, + /// When the session's agent last received a user message — the original + /// instructions or a `continue_repo_session` follow-up — so a parent that + /// continued the session can confirm the message landed in its transcript. + last_user_message_at: Option, last_message: Option, last_tool_call: Option, last_tool_result: Option, @@ -380,7 +396,9 @@ impl ProjectToolsHandler { activity.counts.tool_results += 1; activity.last_tool_result = Some(entry); } - MessageRole::User => {} + MessageRole::User => { + activity.last_user_message_at = Some(message.created_at); + } } } @@ -504,6 +522,22 @@ impl ProjectToolsHandler { Ok(payload) } + /// The `continue_repo_session` reply: the session payload with how the + /// message was delivered and, when it was not, why. + fn continue_repo_session_payload( + store: &Store, + repo_session_id: &str, + handle: &RepoSessionHandle, + outcome: &ContinueOutcome, + ) -> Result { + let mut payload = Self::repo_session_payload(store, repo_session_id, handle)?; + payload["delivery"] = serde_json::Value::String(outcome.delivery().to_string()); + if let Some(reason) = outcome.reason() { + payload["reason"] = serde_json::Value::String(reason.to_string()); + } + Ok(payload) + } + /// Returns true if the store still records this session as running. Used to /// decide whether to defer a cancellation during the startup window where /// the DB row is already `running` but the runner hasn't registered its @@ -869,6 +903,60 @@ impl ProjectToolsHandler { } } + #[tool( + description = "Send a follow-up message to a repo session started by `start_repo_session`: a correction, the answer to a question it raised, or its next step. A running session receives the message within seconds without being stopped, so its work so far is kept; a completed, cancelled, or failed session is resumed with the message as its next instruction. A session that has not started yet cannot take a message — wait for it to run, or cancel it and start a new one. Prefer this over cancelling and starting over whenever the session's work so far is still wanted. Returns the session's current state together with how the message was delivered, or why it was not." + )] + async fn continue_repo_session( + &self, + Parameters(p): Parameters, + ) -> String { + let handle = match Self::decode_repo_session_handle(&p.repo_session_id) { + Ok(handle) => handle, + Err(e) => return e, + }; + log::debug!( + "[project_mcp] continue_repo_session called: session={} message_chars={}", + handle.session_id, + p.message.chars().count() + ); + + // Repo sessions are not project sessions, so the resume never hands + // these to the run; `resume_session_for_store` just needs values. + let action_executor = self + .action_executor + .clone() + .unwrap_or_else(|| Arc::new(ActionExecutor::new())); + let action_registry = self + .action_registry + .clone() + .unwrap_or_else(|| Arc::new(ActionRegistry::new())); + let outcome = match crate::session_commands::continue_session_for_store( + Arc::clone(&self.store), + Arc::clone(&self.registry), + action_executor, + action_registry, + self.app_handle.clone(), + handle.session_id.clone(), + p.message, + Vec::new(), + ) + .await + { + Ok(outcome) => outcome, + Err(e) => return format!("Error continuing repo session: {e}"), + }; + + match Self::continue_repo_session_payload( + &self.store, + &p.repo_session_id, + &handle, + &outcome, + ) { + Ok(payload) => payload.to_string(), + Err(e) => e, + } + } + #[tool( description = "Add a GitHub repository to the current project. Use this when the task requires a repository that isn't yet in the project. Returns once the repository worktree is ready; any setup actions then run in the background." )] @@ -1182,11 +1270,14 @@ mod tests { REPO_SESSION_ACTIVITY_PREVIEW_MAX_CHARS, }; use crate::git::Span; + use crate::session_runner::SessionRegistry; use crate::store::{ AcpConfigSelection, AcpConfigValueSelection, Branch, MessageRole, Project, ProjectNote, - Session, SessionMessage, Store, + Session, SessionMessage, SessionStatus, Store, }; + use acp_client::{SteerOutcome, TurnControlHandle}; use std::path::Path; + use std::sync::Arc; #[test] fn comment_line_range_reports_one_indexed_inclusive_lines() { @@ -1407,6 +1498,7 @@ mod tests { let activity = ProjectToolsHandler::summarize_session_activity(&session, &messages); assert_eq!(activity.last_activity_at, Some(109)); + assert_eq!(activity.last_user_message_at, Some(106)); assert_eq!(activity.counts.assistant_messages, 1); assert_eq!(activity.counts.tool_calls, 1); assert_eq!(activity.counts.tool_results, 1); @@ -1518,6 +1610,180 @@ mod tests { assert!(!cancel_description.contains("taking a long time")); } + /// The continue tool's description has to carry the routing the parent + /// agent must know to choose it over cancel-and-restart: a running + /// session keeps its work, a stopped one is resumed, a queued one cannot + /// take a message — in words, not payload field names. + #[test] + fn continue_repo_session_description_explains_delivery_without_field_names() { + let router = ProjectToolsHandler::tool_router(); + let description = router + .get("continue_repo_session") + .and_then(|tool| tool.description.as_deref()) + .expect("continue tool description"); + + assert!(description.contains("Send a follow-up message")); + assert!(description.contains("without being stopped")); + assert!(description.contains("resumed with the message")); + assert!(description.contains("has not started yet cannot take a message")); + assert!(description.contains("Prefer this over cancelling")); + assert!(!description.contains("`delivery`")); + assert!(!description.contains("injected")); + assert!(!description.contains("interrupt")); + assert!(!description.contains("steer")); + } + + /// Fixture for the continue round trips: a `running` repo session with + /// its note artifact, registered as a live session so the registry can + /// carry a turn-control handle for it. + struct ContinueFixture { + store: Arc, + registry: Arc, + _registration: crate::session_runner::ExternalSessionRegistration, + session: Session, + repo_session_id: String, + handle: RepoSessionHandle, + } + + impl ContinueFixture { + fn running_note_session() -> Self { + let store = Arc::new(Store::in_memory().unwrap()); + let project = Project::new("owner/repo"); + store.create_project(&project).unwrap(); + let branch = Branch::new(&project.id, "feature", "main"); + store.create_branch(&branch).unwrap(); + let session = + Session::new_running("investigate", Path::new("/tmp")).with_branch(&branch.id); + store.create_session(&session).unwrap(); + let note = crate::store::Note::new(&branch.id, "", "").with_session(&session.id); + store.create_note(¬e).unwrap(); + + let registry = Arc::new(SessionRegistry::new()); + let registration = registry.register_external(&session.id); + let repo_session_id = ProjectToolsHandler::encode_repo_session_handle( + &session.id, + RepoArtifactKind::Note, + ¬e.id, + ) + .unwrap(); + let handle = ProjectToolsHandler::decode_repo_session_handle(&repo_session_id).unwrap(); + Self { + store, + registry, + _registration: registration, + session, + repo_session_id, + handle, + } + } + + /// What the `continue_repo_session` tool does between decoding the + /// handle and serialising the reply, with the resume step left out: + /// a resume needs the app handle the tool carries, which tests cannot + /// build, and every case here stays on the live connection anyway. + async fn continue_with(&self, message: &str) -> serde_json::Value { + let route = crate::session_commands::route_continue_session( + &self.store, + &self.registry, + &self.session.id, + message, + &[], + std::time::Duration::from_millis(200), + ) + .await + .unwrap(); + let outcome = match route { + crate::session_commands::ContinueRoute::Delivered(outcome) => outcome, + other => panic!("a running session is reached live, got {other:?}"), + }; + ProjectToolsHandler::continue_repo_session_payload( + &self.store, + &self.repo_session_id, + &self.handle, + &outcome, + ) + .unwrap() + } + + fn status(&self) -> SessionStatus { + self.store + .get_session(&self.session.id) + .unwrap() + .unwrap() + .status + } + } + + /// A steered message leaves the row `running`: a parent that calls + /// `wait_for_repo_session` next must not see a spurious terminal state + /// (and the tool reply itself still says `running`). + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_running_session_by_steering_never_leaves_running() { + let f = ContinueFixture::running_note_session(); + let (turn_control, served) = + TurnControlHandle::scripted(Ok(SteerOutcome::Injected), Ok(())); + f.registry + .set_turn_control_handle(&f.session.id, turn_control); + + let payload = f.continue_with("also check the tests").await; + + assert_eq!(payload["delivery"], "injected"); + assert_eq!(payload["state"], "running"); + assert!(payload.get("reason").is_none()); + assert_eq!(payload["session_id"], f.session.id); + assert_eq!(f.status(), SessionStatus::Running); + assert_eq!( + *served.lock().unwrap(), + vec!["steer:also check the tests".to_string()] + ); + } + + /// Same when the agent refused the steer and the turn was interrupted and + /// re-prompted on the same connection: the cancel is the turn's, not the + /// session's, so the row stays `running`. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_running_session_by_interrupting_never_leaves_running() { + let f = ContinueFixture::running_note_session(); + let (turn_control, served) = TurnControlHandle::scripted( + Err("agent does not advertise steering".to_string()), + Ok(()), + ); + f.registry + .set_turn_control_handle(&f.session.id, turn_control); + + let payload = f.continue_with("stop, different approach").await; + + assert_eq!(payload["delivery"], "interrupted_and_reprompted"); + assert_eq!(payload["state"], "running"); + assert!(payload.get("reason").is_none()); + assert_eq!(f.status(), SessionStatus::Running); + assert_eq!( + *served.lock().unwrap(), + vec![ + "steer:stop, different approach".to_string(), + "interrupt:stop, different approach".to_string(), + ] + ); + } + + /// A session in its startup window (row `running`, no connection yet) + /// that stays that way past the wait is reported blocked with a reason, + /// still `running`, so the parent retries rather than restarting it. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_starting_session_reports_blocked_with_a_reason() { + let f = ContinueFixture::running_note_session(); + + let payload = f.continue_with("hello?").await; + + assert_eq!(payload["delivery"], "blocked"); + assert_eq!(payload["state"], "running"); + assert!(payload["reason"] + .as_str() + .unwrap() + .contains("still starting")); + assert_eq!(f.status(), SessionStatus::Running); + } + /// A parent project agent reads a child's note through this payload, so /// attached media must arrive as paths it can open, not `staged-media://` /// references only the Staged renderer understands. diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 8cf4e5aec..bd93c13e4 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1534,7 +1534,7 @@ pub(crate) async fn resume_session_for_store( .transition_to_running_clearing_agent_id(&session_id, pending_start_agent_id.as_deref()) .map_err(|e| e.to_string())?; if !transitioned { - return Err("Session is already running".to_string()); + return Err(SESSION_ALREADY_RUNNING_ERROR.to_string()); } let config_branch_id = event_branch_id.clone(); @@ -1631,6 +1631,367 @@ pub(crate) async fn resume_session_for_store( Ok(()) } +/// The refusal `resume_session_for_store` returns when the row is already +/// `running` — `continue_session_for_store` reads it back as a blocked +/// delivery rather than an error. +pub(crate) const SESSION_ALREADY_RUNNING_ERROR: &str = "Session is already running"; + +/// How long `continue_session_for_store` waits for a `running` session that +/// has no turn-control handle yet — the startup window between the `running` +/// write and the runner's `connect` — before answering `Blocked`. +pub(crate) const CONTINUE_SESSION_STARTUP_WAIT: Duration = Duration::from_secs(10); + +/// How a message handed to `continue_session_for_store` reached (or failed to +/// reach) the session's agent. The wire spelling is `delivery()`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum ContinueOutcome { + /// Steered into the running turn (`_session/steering`); the agent's + /// in-flight tool calls survived. + Injected, + /// The running turn was cancelled and the message sent as the next + /// prompt on the same connection. + InterruptedAndReprompted, + /// Nothing was in flight to steer (the connection was holding for + /// background work or idle), so the message went out as a fresh turn on + /// the same connection. + StartedNewTurn, + /// The session was stopped; it was resumed with the message as its next + /// prompt (`resume_session_for_store`). + Resumed, + /// The session has not started yet; delivering would claim it out of + /// FIFO order, so the message was refused. + RejectedQueued, + /// Not delivered now, but a retry may succeed: the session is still + /// starting, its connection went away under the request, the branch is + /// busy, or another start won the race. + Blocked { reason: String }, + /// A stopped session that cannot be resumed: no agent conversation to + /// load, or its worktree is gone. The caller should start a new session. + NotResumable { reason: String }, +} + +impl ContinueOutcome { + pub(crate) fn delivery(&self) -> &'static str { + match self { + ContinueOutcome::Injected => "injected", + ContinueOutcome::InterruptedAndReprompted => "interrupted_and_reprompted", + ContinueOutcome::StartedNewTurn => "started_new_turn", + ContinueOutcome::Resumed => "resumed", + ContinueOutcome::RejectedQueued => "rejected_queued", + ContinueOutcome::Blocked { .. } => "blocked", + ContinueOutcome::NotResumable { .. } => "not_resumable", + } + } + + pub(crate) fn reason(&self) -> Option<&str> { + match self { + ContinueOutcome::Blocked { reason } | ContinueOutcome::NotResumable { reason } => { + Some(reason) + } + _ => None, + } + } + + fn blocked(reason: impl Into) -> Self { + ContinueOutcome::Blocked { + reason: reason.into(), + } + } + + fn not_resumable(reason: impl Into) -> Self { + ContinueOutcome::NotResumable { + reason: reason.into(), + } + } +} + +/// Where `route_continue_session` sent the message. +#[derive(Debug)] +pub(crate) enum ContinueRoute { + /// Delivered, refused, or blocked without touching the row. + Delivered(ContinueOutcome), + /// A stopped session that passed the resume gate: the caller resumes it + /// with the message as the next prompt. + Resume(Box), +} + +/// Deliver `message` to a session now, whatever state it is in. +/// +/// The Tauri/web `send_session_message_now` command and the project MCP +/// server's `continue_repo_session` tool share this body. A `running` +/// session is reached over its live ACP connection — steering first, +/// interrupt-and-reprompt when the agent refuses — so the row never leaves +/// `running` and the runner never sees a turn boundary; a stopped session is +/// resumed; a `queued` one is refused. The user row is written by whoever +/// delivers: acp-client for the live paths (only once the agent took the +/// message), `start_session` for the resume, so nothing here persists it. +/// +/// Resume failures other than the already-running race are returned as +/// `Err`, as `resume_session_for_store` reports them. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn continue_session_for_store( + store: Arc, + registry: Arc, + action_executor: Arc, + action_registry: Arc, + app_handle: tauri::AppHandle, + session_id: String, + message: String, + image_ids: Vec, +) -> Result { + let route = route_continue_session( + &store, + ®istry, + &session_id, + &message, + &image_ids, + CONTINUE_SESSION_STARTUP_WAIT, + ) + .await?; + let session = match route { + ContinueRoute::Delivered(outcome) => return Ok(outcome), + ContinueRoute::Resume(session) => *session, + }; + + // No `branch_id` override: the resume resolves the branch from the + // session's artifact, or from the row itself for pipeline sessions. + match resume_session_for_store( + store, + registry, + action_executor, + action_registry, + app_handle, + session.id, + message, + Some(image_ids), + None, + None, + None, + ) + .await + { + Ok(()) => Ok(ContinueOutcome::Resumed), + Err(e) if e == SESSION_ALREADY_RUNNING_ERROR => Ok(ContinueOutcome::blocked( + "the session started running again under this request; retry to reach it live", + )), + Err(e) => Err(e), + } +} + +/// Everything `continue_session_for_store` does short of the resume itself, +/// which needs the app handle. Loops because the status can move under the +/// request: a `running` session whose connection ends while this waits for +/// its handle is routed again on the status it ended with. +pub(crate) async fn route_continue_session( + store: &Store, + registry: &session_runner::SessionRegistry, + session_id: &str, + message: &str, + image_ids: &[String], + startup_wait: Duration, +) -> Result { + let deadline = Instant::now() + startup_wait; + // A handle whose connection went away under this request; the wait below + // skips it so only a re-connect's handle (or a status change) ends it. + let mut dead_handle: Option = None; + loop { + let session = store + .get_session(session_id) + .map_err(|e| e.to_string())? + .ok_or_else(|| format!("Session not found: {session_id}"))?; + + match session.status { + store::SessionStatus::Queued => { + return Ok(ContinueRoute::Delivered(ContinueOutcome::RejectedQueued)); + } + store::SessionStatus::Running => { + let handle = match wait_for_turn_control_handle( + store, + registry, + session_id, + deadline, + dead_handle.as_ref(), + ) + .await + { + HandleWait::Ready(handle) => handle, + HandleWait::LeftRunning => continue, + HandleWait::TimedOut => { + return Ok(ContinueRoute::Delivered(ContinueOutcome::blocked( + "the session is still starting; retry in a few seconds", + ))); + } + }; + match deliver_to_running_session(store, &handle, message, image_ids).await { + Ok(outcome) => return Ok(ContinueRoute::Delivered(outcome)), + Err(reason) if !handle.is_connection_alive() => { + // The connection settled under the request. The row + // goes terminal moments later (or a re-connect hands + // out a new handle); route on whichever comes first. + log::info!( + "continue_session: connection for {session_id} ended under the \ + request ({reason}); waiting for the row to settle" + ); + dead_handle = Some(handle); + continue; + } + Err(reason) => { + return Ok(ContinueRoute::Delivered(ContinueOutcome::blocked(reason))); + } + } + } + store::SessionStatus::Completed + | store::SessionStatus::Cancelled + | store::SessionStatus::Error => { + return Ok(match continue_resume_gate(store, &session)? { + Some(refused) => ContinueRoute::Delivered(refused), + None => ContinueRoute::Resume(Box::new(session)), + }); + } + } + } +} + +enum HandleWait { + Ready(acp_client::TurnControlHandle), + /// The row is no longer `running`; route on its new status. + LeftRunning, + TimedOut, +} + +/// Poll the registry for `session_id`'s live turn-control handle until +/// `deadline`. A handle equal to `skip` (the one whose connection just died) +/// does not count; neither does one that already reports its connection gone. +async fn wait_for_turn_control_handle( + store: &Store, + registry: &session_runner::SessionRegistry, + session_id: &str, + deadline: Instant, + skip: Option<&acp_client::TurnControlHandle>, +) -> HandleWait { + const POLL: Duration = Duration::from_millis(100); + loop { + if let Some(handle) = registry.turn_control_handle(session_id) { + let is_skipped = skip.is_some_and(|dead| dead.same_connection(&handle)); + if !is_skipped && handle.is_connection_alive() { + return HandleWait::Ready(handle); + } + } + match store.get_session(session_id) { + Ok(Some(session)) if session.status == store::SessionStatus::Running => {} + _ => return HandleWait::LeftRunning, + } + if Instant::now() >= deadline { + return HandleWait::TimedOut; + } + tokio::time::sleep(POLL).await; + } +} + +/// Reach a running session's agent over its connection: steer, and when the +/// agent refuses (no steering advertised, `promptRequired` because the turn +/// ended under the request, `failed`), interrupt the turn and re-prompt. +/// `Err` only when the interrupt fails too — the caller decides whether that +/// is a dead connection or a blocked delivery. +async fn deliver_to_running_session( + store: &Store, + handle: &acp_client::TurnControlHandle, + message: &str, + image_ids: &[String], +) -> Result { + let attachments = session_runner::load_prompt_attachments(store, image_ids); + let prompt = attachments.prompt(message); + let steer_refusal = match handle.steer(&prompt, &attachments.images).await { + Ok(acp_client::SteerOutcome::Injected) => return Ok(ContinueOutcome::Injected), + Ok(acp_client::SteerOutcome::NewTurnStarted) => return Ok(ContinueOutcome::StartedNewTurn), + Err(reason) => reason, + }; + log::info!("continue_session: steer refused ({steer_refusal}); interrupting the turn instead"); + handle + .interrupt_and_reprompt(&prompt, &attachments.images) + .await + .map(|()| ContinueOutcome::InterruptedAndReprompted) + .map_err(|reason| format!("could not reach the running session: {reason}")) +} + +/// Whether a stopped session may be resumed with a new message: it needs an +/// agent conversation to load and, for a local branch, its worktree on disk; +/// on a branch, the start must also be allowed by the branch schedule (no +/// provisioning in progress, nothing queued ahead of it, no conflicting +/// running session). `Some` is the refusal; `None` lets the resume proceed. +/// +/// The schedule check runs under the branch's launch lock like the branch +/// session starts do, but the lock is not held into the resume's `running` +/// write (that write is itself atomic on the status), so two continues on +/// one branch can both pass the gate; the second then finds the first's row +/// running and is reported blocked. +pub(crate) fn continue_resume_gate( + store: &Store, + session: &store::Session, +) -> Result, String> { + if resumable_agent_session_id(session.agent_id.as_deref()).is_none() { + return Ok(Some(ContinueOutcome::not_resumable( + "the session has no agent conversation to resume; start a new session instead", + ))); + } + + let (branch, session_type) = continue_branch_and_type(store, session)?; + let is_remote = branch + .as_ref() + .is_some_and(|branch| branch.branch_type == store::BranchType::Remote); + if !is_remote && !Path::new(&session.working_dir).is_dir() { + return Ok(Some(ContinueOutcome::not_resumable(format!( + "the session's working directory is gone ({}); start a new session instead", + session.working_dir + )))); + } + + if let (Some(branch), Some(session_type)) = (branch, session_type) { + let launch_lock = branch_session_launch_lock_for(&branch.id); + let _guard = launch_lock.lock().unwrap(); + if should_queue_branch_session_start(store, &branch.id, &session_type)? { + return Ok(Some(ContinueOutcome::blocked( + "the branch is provisioning or busy with another session; retry once it is idle", + ))); + } + } + + Ok(None) +} + +/// The branch a session runs on and the schedule kind it counts as, from its +/// artifact (commit, note, review). A session with no artifact (a pipeline +/// pr/push session) carries its branch on the row but no kind the branch +/// schedule gates on, so the resume gate does not queue it behind anything. +fn continue_branch_and_type( + store: &Store, + session: &store::Session, +) -> Result<(Option, Option), String> { + let lookup = |branch_id: &str| store.get_branch(branch_id).map_err(|e| e.to_string()); + if let Some(commit) = store + .get_commit_by_session(&session.id) + .map_err(|e| e.to_string())? + { + return Ok((lookup(&commit.branch_id)?, Some(BranchSessionType::Commit))); + } + if let Some(note) = store + .get_note_by_session(&session.id) + .map_err(|e| e.to_string())? + { + return Ok((lookup(¬e.branch_id)?, Some(BranchSessionType::Note))); + } + if let Some(review) = store + .get_review_by_session(&session.id) + .map_err(|e| e.to_string())? + { + return Ok((lookup(&review.branch_id)?, Some(BranchSessionType::Review))); + } + let Some(branch_id) = session.branch_id.as_deref() else { + return Ok((None, None)); + }; + Ok((lookup(branch_id)?, None)) +} + #[tauri::command] pub fn queue_session_message( store: tauri::State<'_, Mutex>>>, @@ -3159,7 +3520,13 @@ or `failed`), any available artifacts, and activity details. Prefer another \ `wait_for_repo_session` call when the returned activity shows recent progress.\n\ - cancel_repo_session: Use this to abort a queued or running repo session by `repo_session_id` \ when the user wants the session stopped. Cancellation is best used when the user wants to go \ -down a different path rather than when you are surprised at how long the session is taking." +down a different path rather than when you are surprised at how long the session is taking.\n\ +- continue_repo_session: Use this to send a follow-up message to a repo session by \ +`repo_session_id` — a correction, an answer to a question it raised, or a next step. A running \ +session receives it within seconds without being stopped and keeps its work; a completed, \ +cancelled, or failed session is resumed with it as the next instruction. A session that has not \ +started yet cannot take a message; wait for it to run or cancel it and start a new one. Prefer \ +this over cancelling and starting over when the session's work so far is still wanted." } else { "- start_repo_session: Use this to make changes or run tasks in one of the project's \ repositories. It enqueues work and returns a `repo_session_id` immediately. For notes, prefer \ @@ -3178,7 +3545,13 @@ or `failed`), any available artifacts, and activity details. Prefer another \ `wait_for_repo_session` call when the returned activity shows recent progress.\n\ - cancel_repo_session: Use this to abort a queued or running repo session by `repo_session_id` \ when the user wants the session stopped. Cancellation is best used when the user wants to go \ -down a different path rather than when you are surprised at how long the session is taking." +down a different path rather than when you are surprised at how long the session is taking.\n\ +- continue_repo_session: Use this to send a follow-up message to a repo session by \ +`repo_session_id` — a correction, an answer to a question it raised, or a next step. A running \ +session receives it within seconds without being stopped and keeps its work; a completed, \ +cancelled, or failed session is resumed with it as the next instruction. A session that has not \ +started yet cannot take a message; wait for it to run or cancel it and start a new one. Prefer \ +this over cancelling and starting over when the session's work so far is still wanted." }; let coordinator_reminder = if is_remote { @@ -7483,6 +7856,419 @@ mod tests { ); } + /// `continue_session_for_store` short of the resume itself (which needs + /// the app handle): the routing on status, the live-connection delivery, + /// and the resume gate, with a short startup wait so the blocked cases + /// run in milliseconds. + async fn continue_through_shared_body( + registry: &session_runner::SessionRegistry, + store: &Arc, + session_id: &str, + message: &str, + ) -> ContinueRoute { + route_continue_session( + store, + registry, + session_id, + message, + &[], + Duration::from_millis(300), + ) + .await + .unwrap() + } + + fn delivered(route: ContinueRoute) -> ContinueOutcome { + match route { + ContinueRoute::Delivered(outcome) => outcome, + ContinueRoute::Resume(session) => { + panic!("expected a delivery, got a resume of {}", session.id) + } + } + } + + /// A commit session ended `completed` with the agent conversation id a + /// resume needs, on `branch`. + fn completed_commit_session_with_agent_id( + store: &Arc, + branch_id: &str, + ) -> store::Session { + let session = create_branch_commit_session(store, branch_id, store::SessionStatus::Running); + store + .update_session_status( + &session.id, + store::SessionStatus::Completed, + None, + Some(&store::CompletionReason::TurnComplete), + ) + .unwrap(); + store + .set_agent_session_id(&session.id, "acp-session-1") + .unwrap(); + session + } + + /// A `queued` row is never claimed out of FIFO order by a message: the + /// drain owns its start. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_queued_session_is_rejected() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Queued); + + let outcome = delivered( + continue_through_shared_body(&test_registry(), &store, &session.id, "hi").await, + ); + + assert_eq!(outcome, ContinueOutcome::RejectedQueued); + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Queued + ); + } + + /// A running session with a connection is steered first; the agent took + /// the message into its turn, so nothing else is sent. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_running_session_steers_the_live_turn() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, served) = + acp_client::TurnControlHandle::scripted(Ok(acp_client::SteerOutcome::Injected), Ok(())); + registry.set_turn_control_handle(&session.id, handle); + + let outcome = + delivered(continue_through_shared_body(®istry, &store, &session.id, "also X").await); + + assert_eq!(outcome, ContinueOutcome::Injected); + assert_eq!(*served.lock().unwrap(), vec!["steer:also X".to_string()]); + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Running + ); + } + + /// A steer served while the connection was holding (or idle) went out as + /// a fresh turn on the same connection; the caller learns which. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_holding_session_reports_the_new_turn() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, _) = acp_client::TurnControlHandle::scripted( + Ok(acp_client::SteerOutcome::NewTurnStarted), + Ok(()), + ); + registry.set_turn_control_handle(&session.id, handle); + + let outcome = + delivered(continue_through_shared_body(®istry, &store, &session.id, "next").await); + + assert_eq!(outcome, ContinueOutcome::StartedNewTurn); + } + + /// When the agent refuses the steer — no steering advertised, or the turn + /// ended under the request — the fallback interrupts the turn and + /// re-prompts on the same connection, in that order. + #[tokio::test(flavor = "current_thread")] + async fn continuing_falls_back_to_interrupt_when_the_steer_is_refused() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, served) = acp_client::TurnControlHandle::scripted( + Err("steer not injected: promptRequired".to_string()), + Ok(()), + ); + registry.set_turn_control_handle(&session.id, handle); + + let outcome = + delivered(continue_through_shared_body(®istry, &store, &session.id, "redo").await); + + assert_eq!(outcome, ContinueOutcome::InterruptedAndReprompted); + assert_eq!( + *served.lock().unwrap(), + vec!["steer:redo".to_string(), "interrupt:redo".to_string()] + ); + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Running + ); + } + + /// Both deliveries refused on a connection that is still up (a bridge + /// that never acknowledged the cancel): blocked with the agent's reason, + /// the row untouched. + #[tokio::test(flavor = "current_thread")] + async fn continuing_is_blocked_when_the_interrupt_fails_on_a_live_connection() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, _) = acp_client::TurnControlHandle::scripted( + Err("agent does not advertise steering".to_string()), + Err("agent did not acknowledge the cancel".to_string()), + ); + registry.set_turn_control_handle(&session.id, handle); + + let outcome = + delivered(continue_through_shared_body(®istry, &store, &session.id, "x").await); + + match outcome { + ContinueOutcome::Blocked { reason } => { + assert!(reason.contains("did not acknowledge"), "{reason}") + } + other => panic!("expected blocked, got {other:?}"), + } + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Running + ); + } + + /// The startup window: the row is `running` but the runner has not + /// connected, so there is no handle. The body waits for one and, when it + /// never comes, answers blocked — never an error, never a status write. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_session_without_a_connection_waits_then_blocks() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + + let started = Instant::now(); + let outcome = + delivered(continue_through_shared_body(®istry, &store, &session.id, "x").await); + + assert!(started.elapsed() >= Duration::from_millis(300)); + match outcome { + ContinueOutcome::Blocked { reason } => { + assert!(reason.contains("still starting"), "{reason}") + } + other => panic!("expected blocked, got {other:?}"), + } + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Running + ); + } + + /// A handle registered while the body is waiting (the runner's `connect` + /// landing) is picked up and the message delivered through it. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_starting_session_delivers_once_it_connects() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, served) = + acp_client::TurnControlHandle::scripted(Ok(acp_client::SteerOutcome::Injected), Ok(())); + let connect = { + let registry = Arc::clone(®istry); + let session_id = session.id.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(50)).await; + registry.set_turn_control_handle(&session_id, handle); + }) + }; + + let outcome = + delivered(continue_through_shared_body(®istry, &store, &session.id, "late").await); + + connect.await.unwrap(); + assert_eq!(outcome, ContinueOutcome::Injected); + assert_eq!(*served.lock().unwrap(), vec!["steer:late".to_string()]); + } + + /// A session that stops while the body waits for its connection is + /// routed on the status it ended with, not left to time out. + #[tokio::test(flavor = "current_thread")] + async fn continuing_reroutes_when_the_session_stops_during_the_wait() { + let (store, branch) = setup_branch_store_with_workdir(); + let session = + create_branch_commit_session(&store, &branch.id, store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let settle = { + let store = Arc::clone(&store); + let session_id = session.id.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(50)).await; + store + .update_session_status( + &session_id, + store::SessionStatus::Completed, + None, + Some(&store::CompletionReason::TurnComplete), + ) + .unwrap(); + store.set_agent_session_id(&session_id, "acp-1").unwrap(); + }) + }; + + let route = continue_through_shared_body(®istry, &store, &session.id, "x").await; + + settle.await.unwrap(); + match route { + ContinueRoute::Resume(resumed) => assert_eq!(resumed.id, session.id), + other => panic!("expected a resume, got {other:?}"), + } + } + + /// A stopped session with nothing to load (the agent never returned a + /// conversation id, or the start was still pending) cannot take a + /// message; the caller is told to start a new session. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_stopped_session_without_an_agent_id_is_not_resumable() { + let (store, branch) = setup_branch_store_with_workdir(); + let session = + create_branch_commit_session(&store, &branch.id, store::SessionStatus::Running); + store + .update_session_status( + &session.id, + store::SessionStatus::Error, + Some("boom"), + Some(&store::CompletionReason::Crashed), + ) + .unwrap(); + + let outcome = delivered( + continue_through_shared_body(&test_registry(), &store, &session.id, "x").await, + ); + + match outcome { + ContinueOutcome::NotResumable { reason } => { + assert!(reason.contains("no agent conversation"), "{reason}") + } + other => panic!("expected not resumable, got {other:?}"), + } + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Error + ); + } + + /// Same when the worktree the session ran in is gone. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_stopped_session_whose_worktree_is_gone_is_not_resumable() { + let (store, branch) = setup_branch_store_with_workdir(); + let missing = std::env::temp_dir().join(format!("staged-gone-{}", uuid::Uuid::new_v4())); + let session = store::Session::new_running("commit", &missing).with_branch(&branch.id); + store.create_session(&session).unwrap(); + let commit = store::Commit::new_pending(&branch.id).with_session(&session.id); + store.create_commit(&commit).unwrap(); + store + .update_session_status( + &session.id, + store::SessionStatus::Completed, + None, + Some(&store::CompletionReason::TurnComplete), + ) + .unwrap(); + store.set_agent_session_id(&session.id, "acp-1").unwrap(); + + let outcome = delivered( + continue_through_shared_body(&test_registry(), &store, &session.id, "x").await, + ); + + match outcome { + ContinueOutcome::NotResumable { reason } => { + assert!(reason.contains("working directory is gone"), "{reason}") + } + other => panic!("expected not resumable, got {other:?}"), + } + } + + /// The resume takes the branch gate the branch session starts take: a + /// branch with another commit session running cannot take a second one, + /// so the message is blocked (a retry, not a new session). + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_stopped_session_on_a_busy_branch_is_blocked() { + let (store, branch) = setup_branch_store_with_workdir(); + let session = completed_commit_session_with_agent_id(&store, &branch.id); + let _other = + create_branch_commit_session(&store, &branch.id, store::SessionStatus::Running); + + let outcome = delivered( + continue_through_shared_body(&test_registry(), &store, &session.id, "x").await, + ); + + match outcome { + ContinueOutcome::Blocked { reason } => assert!(reason.contains("busy"), "{reason}"), + other => panic!("expected blocked, got {other:?}"), + } + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Completed + ); + } + + /// A branch still provisioning (no worktree) is the other half of the + /// gate: the resume would have nowhere to run. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_stopped_session_on_a_provisioning_branch_is_blocked() { + let (store, branch) = setup_branch_store(); + let session = completed_commit_session_with_agent_id(&store, &branch.id); + + let outcome = delivered( + continue_through_shared_body(&test_registry(), &store, &session.id, "x").await, + ); + + assert!( + matches!(outcome, ContinueOutcome::Blocked { .. }), + "{outcome:?}" + ); + } + + /// A stopped session with a conversation to load, its worktree on disk + /// and an idle branch goes to the resume, which `start_session` then + /// persists the message for. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_resumable_stopped_session_routes_to_the_resume() { + let (store, branch) = setup_branch_store_with_workdir(); + let session = completed_commit_session_with_agent_id(&store, &branch.id); + + let route = continue_through_shared_body(&test_registry(), &store, &session.id, "x").await; + + match route { + ContinueRoute::Resume(resumed) => assert_eq!(resumed.id, session.id), + other => panic!("expected a resume, got {other:?}"), + } + // The gate decides; it does not write. + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Completed + ); + } + + /// The wire spellings the MCP reply and the UI route on. + #[test] + fn continue_outcome_delivery_spellings_and_reasons() { + assert_eq!(ContinueOutcome::Injected.delivery(), "injected"); + assert_eq!( + ContinueOutcome::InterruptedAndReprompted.delivery(), + "interrupted_and_reprompted" + ); + assert_eq!( + ContinueOutcome::StartedNewTurn.delivery(), + "started_new_turn" + ); + assert_eq!(ContinueOutcome::Resumed.delivery(), "resumed"); + assert_eq!( + ContinueOutcome::RejectedQueued.delivery(), + "rejected_queued" + ); + let blocked = ContinueOutcome::blocked("busy"); + assert_eq!(blocked.delivery(), "blocked"); + assert_eq!(blocked.reason(), Some("busy")); + let not_resumable = ContinueOutcome::not_resumable("gone"); + assert_eq!(not_resumable.delivery(), "not_resumable"); + assert_eq!(not_resumable.reason(), Some("gone")); + assert_eq!(ContinueOutcome::Injected.reason(), None); + } + #[test] fn branch_start_decision_queues_local_branch_without_workdir() { let (store, branch) = setup_branch_store(); diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index fe0070271..8c076864d 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -36,7 +36,7 @@ //! - All write paths in the background thread use `let _ =` or //! `log::error!`, so FK failures are swallowed gracefully. -mod attachments; +pub(crate) mod attachments; use std::cell::Cell; use std::collections::HashMap; @@ -799,7 +799,11 @@ impl SessionRegistry { /// Store the turn-control handle for `session_id`'s current ACP /// connection, replacing the previous connection's handle (a Pikchr /// correction re-connects on the same session). - fn set_turn_control_handle(&self, session_id: &str, handle: acp_client::TurnControlHandle) { + pub(crate) fn set_turn_control_handle( + &self, + session_id: &str, + handle: acp_client::TurnControlHandle, + ) { if let Some(running) = self.inner.lock().unwrap().running.get(session_id) { *running.turn_control.lock().unwrap() = Some(handle); } @@ -1282,6 +1286,48 @@ pub(crate) fn finish_failed_before_run( transitioned } +/// Read the attachments for a prompt's content blocks: images as they are, +/// PDFs rendered page by page (see `attachments`). An id that cannot be read +/// is logged and skipped, so a bad attachment never blocks the message. +pub(crate) fn load_prompt_attachments( + store: &Store, + image_ids: &[String], +) -> attachments::PromptAttachments { + let mut attachments = attachments::PromptAttachments::default(); + for image_id in image_ids { + match store.get_image(image_id) { + Ok(Some(image)) => { + match crate::store::images::image_file_path( + &image.project_id, + &image.id, + &image.filename, + ) { + Ok(path) => { + if let Ok(bytes) = std::fs::read(&path) { + attachments.add(&image, &bytes); + } else { + log::warn!( + "Failed to read image file for image {image_id}: {}", + path.display() + ); + } + } + Err(e) => { + log::warn!("Failed to resolve file path for image {image_id}: {e}"); + } + } + } + Ok(None) => { + log::warn!("Image {image_id} not found in store, skipping"); + } + Err(e) => { + log::warn!("Failed to fetch image {image_id} from store: {e}"); + } + } + } + attachments +} + /// Start a session: persist the user message, spawn the agent, stream to DB. /// /// Returns immediately — the actual agent work happens on a background task. @@ -1670,40 +1716,7 @@ pub fn start_session( driver.with_mcp_servers(required_mcp_servers) }; - // Read the attachments for the prompt content blocks: images as - // they are, PDFs rendered page by page (see `attachments`). - let mut attachments = attachments::PromptAttachments::default(); - for image_id in &config.image_ids { - match store.get_image(image_id) { - Ok(Some(image)) => { - match crate::store::images::image_file_path( - &image.project_id, - &image.id, - &image.filename, - ) { - Ok(path) => { - if let Ok(bytes) = std::fs::read(&path) { - attachments.add(&image, &bytes); - } else { - log::warn!( - "Failed to read image file for image {image_id}: {}", - path.display() - ); - } - } - Err(e) => { - log::warn!("Failed to resolve file path for image {image_id}: {e}"); - } - } - } - Ok(None) => { - log::warn!("Image {image_id} not found in store, skipping"); - } - Err(e) => { - log::warn!("Failed to fetch image {image_id} from store: {e}"); - } - } - } + let attachments = load_prompt_attachments(&store, &config.image_ids); let mut prompt = attachments.prompt(&config.prompt); let mut agent_session_id = config.agent_session_id.clone(); diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index cb864afd8..dd43c6ec2 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -1670,6 +1670,43 @@ impl TurnControlHandle { } } + /// A handle to a stand-in connection that answers every steer with + /// `steer` and every interrupt with `interrupt`, and appends + /// `"steer:"` / `"interrupt:"` to the returned log as the + /// requests arrive. For tests of the routing around a handle (steer first, + /// interrupt on refusal) that must not stand up a real agent. Needs a + /// tokio runtime: the stand-in is a task. + #[doc(hidden)] + pub fn scripted( + steer: Result, + interrupt: Result<(), String>, + ) -> (Self, Arc>>) { + let (turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); + let log = Arc::new(std::sync::Mutex::new(Vec::new())); + let served = Arc::clone(&log); + tokio::spawn(async move { + while let Some(request) = turn_control_rx.recv().await { + match request { + TurnControlRequest::Steer { prompt, reply, .. } => { + served.lock().unwrap().push(format!("steer:{prompt}")); + let _ = reply.send(steer.clone()); + } + TurnControlRequest::Interrupt { prompt, reply, .. } => { + served.lock().unwrap().push(format!("interrupt:{prompt}")); + let _ = reply.send(interrupt.clone()); + } + } + } + }); + ( + Self { + turn_control_tx, + connection_alive: Arc::new(AtomicBool::new(true)), + }, + log, + ) + } + fn send(&self, request: TurnControlRequest) -> Result<(), String> { if !self.connection_alive.load(Ordering::Relaxed) { return Err("ACP connection is no longer running".to_string()); From 7125d9190c20efc370cc2e858fcd053c6975a89f Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Wed, 28 Oct 2026 22:20:43 +1100 Subject: [PATCH 4/7] fix(acp-client): verify turn control against real bridges and fix what broke MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Exercises steer and interrupt-and-reprompt (7a219c29, 71f7d193, 13de4081) against the managed claude-agent-acp 0.95.0 and codex-acp 2.2.2 that Staged runs, through `AcpDriver::connect` and through the app's writer/store/registry/`continue_repo_session` route. Three bugs fell out; each gets a fix and a fake-agent regression test. The real-bridge runs stay as `#[ignore]`d tests so they can be re-run: cargo test -p acp-client real_bridge -- --ignored --nocapture --test-threads=1 cargo test -p staged --lib continue_repo_session_live -- --ignored --nocapture --test-threads=1 `ACP_REAL_BRIDGE=codex` picks the other bridge, `ACP_REAL_BRIDGE_DIR` the install dir (default `~/.staged/packages/bin`), and `ACP_REAL_BRIDGE_LOG=debug` raises the log level. Every writer event and log line is printed with its offset from connect. Fixes - Idle wait: a control request queued next to a prompt (the steer-with- the-prompt race) was served before the prompt and had nothing to steer, so the user's follow-up became the turn and the original prompt ran second. The idle select is now biased, prompt first. Regression: `a_prompt_queued_before_a_control_request_goes_first` (8/20 failures without the bias). - Duplicate assistant row after a steer: claude-agent-acp aborts the generation a steer pre-empts and still emits that cycle's `result` with the already-streamed text as its final message. `serve_steer` called `close_message`, which drops the buffer, so `accept_result` re-injected the whole message as a chunk after the user row. New `ResponseState::detach_message` keeps the text and drops only the id. Regression: `result_of_a_cycle_a_steer_aborted_does_not_echo_the_closed_message`. - `INTERRUPT_ACK_TIMEOUT` 5 s -> 30 s: claude-agent-acp acknowledged a cancel with a `sleep 1` tool call in flight only after ~3.9 s; real tool calls run longer and an expired interrupt drops the message. - `AcpDriver::with_steering_disabled` forces the interrupt path against a bridge that would take the steer (no non-steering bridge is installed here; Goose is absent). `log_initialize_response` now logs `_meta` so steering advertisement is visible at debug. Scenarios 1. Steer during a live Claude turn: PASS (claude, codex). `Injected` while the turn was live, agent changed course at once (claude: first steered text 0 ms after the answer; codex: one more count step), one user row in assistant -> user -> assistant order, `prompt()` resolved once, one agent session id. 2. Steer while the bridge holds the turn for background work: PASS (claude) — bridge held the turn for a background subagent, steer `Injected`, reply landed, then SubagentHandback. codex-acp does not hold the turn for subagents; the client hold quiesced and tore down, and the steer failed fast with "connection is no longer running", which is the correct answer there (test accepts that branch). 3. Interrupt-and-reprompt with steering forced off: PASS (claude, codex). Steer refused locally with "does not advertise steering" and no row written, cancel acked (claude ~3.9 s -> timeout fix above, codex 20 ms), same bridge pid before and after, no terminal status mid-way, one user row, follow-up answered in the second turn. Neither bridge issued `session/request_permission` for its shell calls in these runs, so the "permissions in the next turn reach the handler" check was vacuous live; it stays covered by the fake-agent unit tests. 4. Steer racing turn end: FAIL -> fixed (duplicate echo, above); after the fix PASS (claude, codex) for steer sent with the prompt and steer sent at the first chunk. Claude answers `Injected` or `NewTurnStarted` depending on timing; both leave exactly one user row and no repeated assistant row. codex-acp ignores `promptRequired` (it starts a detached turn and answers `startedNewTurn`) per its source; the client maps that to an error so callers fall back to interrupt. Bridge limitation, not fixable here. 5. Through the project MCP tool: PASS for the live half — a real bridge session registered in `SessionRegistry`, `continue_with` returned `delivery: injected` (steering on) and `interrupted_and_reprompted` (steering forced off), row stayed `running` throughout, `activity.last_user_message_at` matched the one user row, follow-up answered. SKIPPED live: `resumed`, `queued` refusal and `wait_for_repo_session` — they go through `start_session`, which needs a Wry `AppHandle`; existing unit tests cover those routes. Also noted: Homebrew's claude-agent-acp 0.53.0 (no steering extension) shadows the managed bridge on the login-shell PATH, which is why the real-bridge tests pin `set_bundled_tools_dir` rather than resolving by name. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/project_mcp.rs | 168 ++++ crates/acp-client/src/driver.rs | 65 +- .../src/driver/response_completion.rs | 14 + .../src/driver/tests/real_bridge.rs | 826 ++++++++++++++++++ .../src/driver/tests/response_completion.rs | 35 + .../src/driver/tests/turn_control.rs | 52 ++ 6 files changed, 1153 insertions(+), 7 deletions(-) create mode 100644 crates/acp-client/src/driver/tests/real_bridge.rs diff --git a/apps/staged/src-tauri/src/project_mcp.rs b/apps/staged/src-tauri/src/project_mcp.rs index c470ff986..d45e5b832 100644 --- a/apps/staged/src-tauri/src/project_mcp.rs +++ b/apps/staged/src-tauri/src/project_mcp.rs @@ -1766,6 +1766,174 @@ mod tests { ); } + /// The tool's live half against the bridge Staged actually runs: the + /// session's real `MessageWriter` and `Store` behind an installed + /// `claude-agent-acp` (or `codex-acp`), its turn-control handle in the + /// registry the way the runner registers it, and + /// `continue_repo_session`'s routing and reply on top. Ignored by default + /// — it spawns the bridge as the local user and calls the model. + /// + /// ```sh + /// cargo test -p staged --lib continue_repo_session_live -- --ignored --nocapture --test-threads=1 + /// ``` + /// + /// `ACP_REAL_BRIDGE` picks the provider id (`claude` by default) and + /// `ACP_REAL_BRIDGE_DIR` the install dir (`~/.staged/packages/bin`). The + /// resume (`resumed`) and `queued` paths are not covered here: a resume + /// needs the app handle the tool carries, which tests cannot build. + async fn continue_repo_session_live(steering_disabled: bool, expected_delivery: &str) { + const LONG_TURN: &str = "Count from 1 to 40 slowly. For EACH number: first run the \ + shell command `sleep 1` with your shell tool as its own separate tool call (never a \ + loop, never a script), then write that number on its own line. Keep going until 40. \ + Do not explain what you are doing."; + const FOLLOW_UP: &str = "Stop counting right now. Run the shell command `pwd` with your \ + shell tool, then reply with exactly one line: FOLLOW-UP-DONE"; + + if let Some(dir) = std::env::var_os("ACP_REAL_BRIDGE_DIR") + .map(std::path::PathBuf::from) + .or_else(|| { + std::env::var_os("HOME") + .map(|home| std::path::PathBuf::from(home).join(".staged/packages/bin")) + }) + .filter(|dir| dir.is_dir()) + { + acp_client::set_bundled_tools_dir(dir); + } + let provider = std::env::var("ACP_REAL_BRIDGE").unwrap_or_else(|_| "claude".to_string()); + let driver = acp_client::AcpDriver::new(&provider) + .expect("the bridge must be installed (see the test docs)") + .with_env_snapshot(std::env::vars().collect()) + .with_background_hold(crate::session_commands::default_background_hold()) + .with_steering_disabled(steering_disabled); + let dir = std::env::temp_dir().join(format!( + "staged-continue-live-{}-{steering_disabled}", + std::process::id() + )); + std::fs::create_dir_all(&dir).unwrap(); + + let f = ContinueFixture::running_note_session(); + let writer: Arc = Arc::new( + crate::agent::writer::MessageWriter::new(f.session.id.clone(), Arc::clone(&f.store)), + ); + let acp_store: Arc = f.store.clone(); + let cancel = tokio_util::sync::CancellationToken::new(); + let messages = || f.store.get_session_messages(&f.session.id).unwrap(); + + tokio::task::LocalSet::new() + .run_until(async { + let mut connection = driver + .connect( + &f.session.id, + &dir, + &acp_store, + &writer, + &cancel, + None, + &[], + None, + ) + .await + .expect("bridge should spawn"); + // What the turn loop does right after `connect`. + f.registry + .set_turn_control_handle(&f.session.id, connection.turn_control_handle()); + let settled_rx = connection.take_settled_receiver().unwrap(); + let turn = + tokio::task::spawn_local( + async move { connection.prompt(LONG_TURN, &[]).await }, + ); + + // Into the turn: wait for its first transcript row, then a + // little longer so the agent is mid-work. + let started = std::time::Instant::now(); + while messages().is_empty() { + assert!( + started.elapsed() < std::time::Duration::from_secs(90), + "the agent never started working" + ); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + tokio::time::sleep(std::time::Duration::from_secs(2)).await; + let rows_before = messages().len(); + + let payload = f.continue_with(FOLLOW_UP).await; + eprintln!("continue_repo_session reply: {payload:#}"); + assert_eq!(payload["delivery"], expected_delivery); + assert_eq!(payload["state"], "running"); + assert!(payload.get("reason").is_none()); + assert_eq!(f.status(), SessionStatus::Running); + // The row the connection persisted is what the parent sees as + // the message having landed. + let user_rows: Vec = messages() + .into_iter() + .filter(|m| m.role == MessageRole::User) + .collect(); + assert_eq!(user_rows.len(), 1, "{user_rows:#?}"); + assert_eq!(user_rows[0].content, FOLLOW_UP); + assert_eq!( + payload["activity"]["last_user_message_at"], + serde_json::json!(user_rows[0].created_at) + ); + + let outcome = tokio::time::timeout(std::time::Duration::from_secs(240), turn) + .await + .expect("the continued turn should end") + .unwrap(); + assert_eq!(outcome, Ok(acp_client::AgentRunOutcome::Completed)); + // Still exactly one user row, after the interrupted work and + // before the reply to it; the row never left `running`. + let all = messages(); + let user_at = all + .iter() + .position(|m| m.role == MessageRole::User) + .unwrap(); + assert_eq!( + all.iter().filter(|m| m.role == MessageRole::User).count(), + 1 + ); + assert!(user_at >= rows_before.min(1), "{all:#?}"); + let reply: String = all[user_at + 1..] + .iter() + .filter(|m| m.role == MessageRole::Assistant) + .map(|m| m.content.as_str()) + .collect::>() + .join("\n"); + assert!( + reply.contains("FOLLOW-UP-DONE"), + "reply after the user row: {reply:?}" + ); + assert_eq!(f.status(), SessionStatus::Running); + for m in &all { + eprintln!( + " {:<10} | {}", + m.role.as_str(), + m.content.replace('\n', "⏎") + ); + } + + cancel.cancel(); + let settled = tokio::time::timeout(std::time::Duration::from_secs(30), settled_rx) + .await + .expect("teardown after cancel") + .unwrap(); + eprintln!("settled: {:?} / {:?}", settled.outcome, settled.reason); + }) + .await; + let _ = std::fs::remove_dir_all(&dir); + } + + #[tokio::test(flavor = "current_thread")] + #[ignore = "spawns the installed ACP bridge and calls the model"] + async fn continue_repo_session_live_steers_a_running_session() { + continue_repo_session_live(false, "injected").await; + } + + #[tokio::test(flavor = "current_thread")] + #[ignore = "spawns the installed ACP bridge and calls the model"] + async fn continue_repo_session_live_interrupts_when_steering_is_unavailable() { + continue_repo_session_live(true, "interrupted_and_reprompted").await; + } + /// A session in its startup window (row `running`, no connection yet) /// that stays that way past the wait is reported blocked with a reason, /// still `running`, so the parent retries rather than restarting it. diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index dd43c6ec2..ba6589ecf 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -1106,6 +1106,11 @@ pub struct AcpDriver { /// hold, so the client can present the wait. Only ever called with a /// `background_hold` configured. background_hold_observer: Option, + /// Treat the agent as not advertising steering even when it does, so + /// every [`TurnControlHandle::steer`] is refused locally and callers take + /// the interrupt-and-reprompt path. See + /// [`AcpDriver::with_steering_disabled`]. + steering_disabled: bool, } const REMOTE_ACP_MAX_PENDING_LINE_BYTES: usize = 256 * 1024; @@ -1130,6 +1135,7 @@ impl AcpDriver { remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, }) .ok_or_else(|| format!("Unknown or unavailable agent provider: {provider_id}")) } @@ -1149,6 +1155,7 @@ impl AcpDriver { remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, }) .ok_or_else(|| { "No ACP agent found. Install Goose, Claude Code, Codex, Pi, or Amp and ensure it's on your PATH." @@ -1169,6 +1176,7 @@ impl AcpDriver { remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, } } @@ -1193,6 +1201,7 @@ impl AcpDriver { remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, }) } @@ -1274,6 +1283,18 @@ impl AcpDriver { self.background_hold_observer = observer; self } + + /// Ignore the agent's steering advertisement: every + /// [`TurnControlHandle::steer`] on connections from this driver is refused + /// as if the agent lacked the extension, so a caller's fallback + /// ([`TurnControlHandle::interrupt_and_reprompt`]) is what delivers the + /// message. A diagnostics knob: it exercises the interrupt path against an + /// agent that would otherwise take the steer, which every supported agent + /// now does. + pub fn with_steering_disabled(mut self, disabled: bool) -> Self { + self.steering_disabled = disabled; + self + } } /// Shell-escape a value by wrapping it in single quotes with interior quotes @@ -2158,6 +2179,7 @@ impl AcpDriver { let auth_selection = auth_selection.cloned(); let mcp_servers = self.mcp_servers.clone(); let agent_label = self.agent_label.clone(); + let steering_disabled = self.steering_disabled; let is_remote = self.is_remote; let background_hold = self.background_hold.clone(); let background_hold_observer = self.background_hold_observer.clone(); @@ -2248,6 +2270,7 @@ impl AcpDriver { auth_selection.as_ref(), background_hold.as_ref(), background_hold_observer.as_ref(), + steering_disabled, &child_exited, &mut prompt_rx, &mut hold_control_rx, @@ -3103,8 +3126,8 @@ fn task_tracking_mode_from_initialize(meta: Option<&Meta>) -> TaskTrackingMode { /// `agentCapabilities`, not inside it): `_meta.steering.supported`. const STEERING_META_KEY: &str = "steering"; -/// Whether the agent advertised steering at initialize (claude-agent-acp -/// ≥ 0.98, codex-acp ≥ 2.2). Anything missing or malformed is an agent without +/// Whether the agent advertised steering at initialize (seen from +/// claude-agent-acp 0.95 and codex-acp 2.2.2; absent from the 0.53 line). Anything missing or malformed is an agent without /// it, which [`TurnControlHandle::steer`] answers with an error so callers /// interrupt instead — the only safe gate with managed bridge versions that /// float. @@ -4987,7 +5010,7 @@ fn serve_steer( Ok(outcome) => { log::info!("ACP session {session_id}: steer answered {outcome:?}"); handler.writer.finalize().await; - handler.response.lock().await.close_message(); + handler.response.lock().await.detach_message(); handler.writer.record_user_message(&prompt, &images).await; } Err(e) => log::warn!("ACP session {session_id}: steer not injected: {e}"), @@ -5010,8 +5033,14 @@ fn serve_steer( /// `session/prompt` after `session/cancel` before giving up on the re-prompt: /// a bridge that cannot acknowledge a cancel should not be handed a second /// prompt, so the turn then settles cancelled, as a plain cancel would. -/// Mirrors the in-turn cancel's bound on the same response. -const INTERRUPT_ACK_TIMEOUT: Duration = Duration::from_secs(5); +/// +/// Longer than the in-turn cancel's 5 s bound on the same response because +/// the cost of expiring is different: a plain cancel that times out loses +/// nothing the user wanted, an interrupt that times out drops the message it +/// was made for. claude-agent-acp 0.95 answered the cancel of a turn with a +/// `sleep 1` shell call in flight only after ~4 s — it lets the tool wind +/// down first — and real tool calls run longer than that. +const INTERRUPT_ACK_TIMEOUT: Duration = Duration::from_secs(30); /// How long the holding wait gives the bridge to process a `session/cancel` /// before teardown. Out-of-turn there is no prompt response to await, so this @@ -5495,6 +5524,7 @@ async fn run_acp_session( auth_selection: Option<&AcpAuthenticationSelection>, background_hold: Option<&BackgroundHoldConfig>, background_hold_observer: Option<&BackgroundHoldObserver>, + steering_disabled: bool, child_exited: &CancellationToken, prompt_rx: &mut mpsc::UnboundedReceiver, hold_control_rx: &mut mpsc::UnboundedReceiver, @@ -5517,7 +5547,7 @@ async fn run_acp_session( auth_selection, }), ); - let setup = tokio::select! { + let mut setup = tokio::select! { _ = cancel_token.cancelled() => { handler.cancel_pending_permissions(); return Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)); @@ -5532,6 +5562,12 @@ async fn run_acp_session( })?? } }; + if steering_disabled && setup.steering_supported { + log::info!( + "ACP session {our_session_id}: steering disabled by the driver; steers will be refused" + ); + setup.steering_supported = false; + } // If resuming, wait for replay to complete (content match OR idle timeout). // An absolute 10s timeout prevents a hang if the server sends zero replay @@ -5580,7 +5616,11 @@ async fn run_acp_session( // already-established session. Some(turn) => turn, None => loop { + // Biased: a prompt already queued when nothing is in flight + // goes first; a control request that landed next to it has + // nothing to steer yet and becomes the turn after. tokio::select! { + biased; _ = cancel_token.cancelled() => { handler.cancel_pending_permissions(); return Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)); @@ -6785,8 +6825,14 @@ fn log_initialize_response(agent_label: &str, init_response: &InitializeResponse .unwrap_or_else(|_| "".to_string()); let auth_methods = describe_auth_methods(&init_response.auth_methods); + let meta = init_response + .meta + .as_ref() + .map(|meta| serde_json::to_string(meta).unwrap_or_else(|_| "".to_string())) + .unwrap_or_else(|| "null".to_string()); + log::debug!( - "ACP initialized provider={agent_label} agent={agent_name} version={agent_version} protocol={} capabilities={} auth_methods=[{}]", + "ACP initialized provider={agent_label} agent={agent_name} version={agent_version} protocol={} capabilities={} auth_methods=[{}] _meta={meta}", init_response.protocol_version, capabilities, auth_methods @@ -7077,6 +7123,7 @@ impl MessageWriter for BasicMessageWriter { #[cfg(test)] mod tests { + mod real_bridge; mod response_completion; mod settlement; mod turn_control; @@ -8749,6 +8796,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, } .with_env_snapshot(snapshot.clone()); let remote = AcpDriver { @@ -8763,6 +8811,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, } .with_env_snapshot(snapshot); @@ -8788,6 +8837,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, } .with_interpreter_env_snapshot(snapshot.clone()); let remote = AcpDriver { @@ -8802,6 +8852,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE remote_working_dir: None, background_hold: None, background_hold_observer: None, + steering_disabled: false, } .with_interpreter_env_snapshot(snapshot); diff --git a/crates/acp-client/src/driver/response_completion.rs b/crates/acp-client/src/driver/response_completion.rs index 4cef53d18..01d437169 100644 --- a/crates/acp-client/src/driver/response_completion.rs +++ b/crates/acp-client/src/driver/response_completion.rs @@ -47,6 +47,20 @@ impl ResponseState { self.interrupted = false; } + /// The writer closed the streamed message mid-turn, with the turn going + /// on — a steered user row went in after it. Unlike [`Self::close_message`] + /// the streamed text is kept: the bridge aborts the generation a steer + /// pre-empts and still emits that cycle's `result`, whose final text is + /// the message just closed. With the buffer intact [`Self::accept_result`] + /// recognises it as already streamed instead of re-injecting the whole + /// message as a chunk after the user row. Only the id is dropped, so a + /// later chunk of the same message opens a row of its own rather than + /// resuming the closed one. + pub(super) fn detach_message(&mut self) { + self.message_id = None; + self.interrupted = false; + } + pub(super) fn suppress_fallback(&mut self, text: &str) -> bool { if !self.fallback.is_empty() && self.fallback.starts_with(text) { self.fallback.drain(..text.len()); diff --git a/crates/acp-client/src/driver/tests/real_bridge.rs b/crates/acp-client/src/driver/tests/real_bridge.rs new file mode 100644 index 000000000..948b9ccf8 --- /dev/null +++ b/crates/acp-client/src/driver/tests/real_bridge.rs @@ -0,0 +1,826 @@ +//! Turn-control against the bridges Staged actually runs: steering into a +//! live turn, interrupt-and-reprompt on the same child, and the steer / +//! turn-end race. Ignored by default — every test spawns an installed bridge +//! as the local user and calls the model, so it needs the bridge on `PATH`, a +//! signed-in CLI, network, and a minute or two. +//! +//! ```sh +//! cargo test -p acp-client real_bridge -- --ignored --nocapture --test-threads=1 +//! ``` +//! +//! The bridge is taken from `ACP_REAL_BRIDGE_DIR`, defaulting to Staged's +//! managed install (`~/.staged/packages/bin`) — pinned explicitly because +//! plain resolution goes through the login shell's `PATH`, where a stale +//! Homebrew `claude-agent-acp` can shadow it. `ACP_REAL_BRIDGE=` +//! picks the bridge (`claude`, the default, or `codex`). Every writer event and every `log` line is printed to stderr +//! with its offset from connect, so under `--nocapture` the output is the +//! record of what the bridge did; `ACP_REAL_BRIDGE_LOG=debug` raises the +//! `log` level. The checks are the invariants, not the agent's wording: one +//! user row, in transcript order, the turn resolving once, the same child +//! across an interrupt, no hang. + +use super::*; +use crate::driver::{SessionConnection, SessionSettled, SteerOutcome, TurnControlHandle}; +use std::cell::Cell; +use std::rc::Rc; +use std::sync::OnceLock; +use std::time::Instant; + +const TURN_START_TIMEOUT: Duration = Duration::from_secs(90); +const TURN_TIMEOUT: Duration = Duration::from_secs(240); +const REPLY_TIMEOUT: Duration = Duration::from_secs(90); +const SETTLE_WAIT: Duration = Duration::from_secs(30); + +/// A long, multi-step turn: one shell call per number keeps the agent busy +/// for a minute or more, with plenty of boundaries for a steer to land on. +const COUNT_PROMPT: &str = "Count from 1 to 40 slowly. For EACH number: first run the shell \ + command `sleep 1` with your shell tool as its own separate tool call (never a loop, never a \ + script, never several sleeps in one command), then write that number on its own line. Keep \ + going until you reach 40. Do not explain what you are doing."; +const STEER_MESSAGE: &str = "Stop counting right now and do not output any more numbers. \ + Instead reply with exactly one line: STEERED-DONE "; +const INTERRUPT_MESSAGE: &str = "Stop what you were doing. Run the shell command `pwd` with \ + your shell tool, then reply with exactly one line: INTERRUPTED-DONE "; +const RACE_PROMPT: &str = "Reply with exactly the single word: pong"; +const RACE_MESSAGE: &str = "Reply with exactly one line: RACE-DONE"; +const HELD_TURN_PROMPT: &str = "Use your Agent (subagent) tool with run_in_background set to \ + true to start one subagent whose task is exactly: run the shell command `sleep 30` and then \ + reply done. As soon as it is launched — without waiting for it — reply with exactly the single \ + word: LAUNCHED"; +const HELD_STEER_MESSAGE: &str = "New instruction: reply with exactly one line: HELD-DONE"; + +fn provider_id() -> String { + std::env::var("ACP_REAL_BRIDGE").unwrap_or_else(|_| "claude".to_string()) +} + +/// Pin bridge lookup to the managed install dir (process-wide, once). +fn pin_bridge_dir() { + let dir = std::env::var_os("ACP_REAL_BRIDGE_DIR") + .map(PathBuf::from) + .or_else(|| { + std::env::var_os("HOME").map(|home| PathBuf::from(home).join(".staged/packages/bin")) + }); + if let Some(dir) = dir.filter(|dir| dir.is_dir()) { + crate::types::set_bundled_tools_dir(dir); + } +} + +/// All assistant text so far, as the agent streamed it. +fn streamed_text(evs: &[Ev]) -> String { + evs.iter() + .filter_map(|ev| match ev { + Ev::Append(text) => Some(text.as_str()), + _ => None, + }) + .collect() +} + +// ----------------------------------------------------------------------------- +// Trace: the writer, as a timestamped event log +// ----------------------------------------------------------------------------- + +#[derive(Debug, Clone, PartialEq)] +enum Ev { + Append(String), + Finalize, + User(String), + ToolCall { id: String, title: String }, + ToolResult { id: String }, + Permission { title: String }, + Event(String), +} + +/// Transcript rows as a store would hold them: assistant text between +/// boundaries, user rows, tool calls. +#[derive(Debug, Clone, PartialEq)] +enum Row { + Assistant(String), + User(String), + Tool(String), +} + +impl Row { + fn is_user(&self) -> bool { + matches!(self, Row::User(_)) + } + + fn text(&self) -> &str { + match self { + Row::Assistant(text) | Row::User(text) | Row::Tool(text) => text, + } + } +} + +struct Trace { + label: String, + started: Instant, + events: std::sync::Mutex>, +} + +impl Trace { + fn new(label: &str) -> Self { + Self { + label: label.to_string(), + started: Instant::now(), + events: std::sync::Mutex::new(Vec::new()), + } + } + + fn elapsed(&self) -> Duration { + self.started.elapsed() + } + + fn push(&self, event: Ev) { + let at = self.elapsed(); + eprintln!("[{} {:>7.2}s] {:?}", self.label, at.as_secs_f64(), event); + self.events.lock().unwrap().push((at, event)); + } + + fn events(&self) -> Vec<(Duration, Ev)> { + self.events.lock().unwrap().clone() + } + + fn count(&self, pred: impl Fn(&Ev) -> bool) -> usize { + self.events().iter().filter(|(_, ev)| pred(ev)).count() + } + + /// Timestamp of the first event (after `from`) that `pred` accepts. + fn first_at(&self, from: Duration, pred: impl Fn(&Ev) -> bool) -> Option { + self.events() + .into_iter() + .find(|(at, ev)| *at >= from && pred(ev)) + .map(|(at, _)| at) + } + + async fn wait_for(&self, what: &str, timeout: Duration, pred: impl Fn(&[Ev]) -> bool) { + let deadline = Instant::now() + timeout; + loop { + let snapshot: Vec = self.events().into_iter().map(|(_, ev)| ev).collect(); + if pred(&snapshot) { + return; + } + assert!( + Instant::now() < deadline, + "[{}] timed out after {timeout:?} waiting for {what}; events: {:#?}", + self.label, + self.events() + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + + fn rows(&self) -> Vec { + let mut rows = Vec::new(); + let mut open = String::new(); + let flush = |open: &mut String, rows: &mut Vec| { + if !open.trim().is_empty() { + rows.push(Row::Assistant(std::mem::take(open))); + } else { + open.clear(); + } + }; + for (_, ev) in self.events() { + match ev { + Ev::Append(text) => open.push_str(&text), + Ev::Finalize => flush(&mut open, &mut rows), + Ev::User(text) => { + flush(&mut open, &mut rows); + rows.push(Row::User(text)); + } + Ev::ToolCall { title, .. } => { + flush(&mut open, &mut rows); + rows.push(Row::Tool(title)); + } + Ev::ToolResult { .. } | Ev::Permission { .. } | Ev::Event(_) => {} + } + } + flush(&mut open, &mut rows); + rows + } + + fn dump_rows(&self) { + eprintln!("[{}] transcript rows:", self.label); + for (index, row) in self.rows().iter().enumerate() { + let text: String = row.text().chars().take(160).collect(); + let text = text.replace('\n', "⏎"); + match row { + Row::Assistant(_) => eprintln!(" {index:>2} assistant | {text}"), + Row::User(_) => eprintln!(" {index:>2} USER | {text}"), + Row::Tool(_) => eprintln!(" {index:>2} tool | {text}"), + } + } + } + + /// Assistant text in the rows after the single user row. + fn assistant_text_after_user(&self) -> String { + let rows = self.rows(); + let user = rows + .iter() + .position(Row::is_user) + .unwrap_or_else(|| panic!("no user row in {rows:#?}")); + rows[user + 1..] + .iter() + .filter_map(|row| match row { + Row::Assistant(text) => Some(text.as_str()), + _ => None, + }) + .collect::>() + .join("\n") + } +} + +#[async_trait::async_trait] +impl MessageWriter for Trace { + async fn append_text(&self, text: &str) { + self.push(Ev::Append(text.to_string())); + } + + async fn finalize(&self) { + self.push(Ev::Finalize); + } + + async fn record_user_message(&self, text: &str, _images: &[PromptImage]) { + self.push(Ev::User(text.to_string())); + } + + async fn record_tool_call( + &self, + tool_call_id: &str, + title: &str, + _raw_input: Option<&serde_json::Value>, + ) { + self.push(Ev::ToolCall { + id: tool_call_id.to_string(), + title: title.to_string(), + }); + } + + async fn update_tool_call_title( + &self, + _tool_call_id: &str, + _title: Option<&str>, + _raw_input: Option<&serde_json::Value>, + ) { + } + + async fn record_tool_result(&self, tool_call_id: &str, _content: &str) { + self.push(Ev::ToolResult { + id: tool_call_id.to_string(), + }); + } + + async fn record_acp_event_metadata(&self, metadata: AcpEventMetadata) { + self.push(Ev::Event( + metadata.event_kind.unwrap_or_else(|| "?".to_string()), + )); + } + + async fn request_permission( + &self, + request: AcpPermissionRequest, + cancel_token: CancellationToken, + ) -> AcpPermissionDecision { + self.push(Ev::Permission { + title: request.tool_title.clone().unwrap_or_default(), + }); + if cancel_token.is_cancelled() { + AcpPermissionDecision::Cancelled + } else { + autoapprove_permission_decision(&request) + } + } +} + +// ----------------------------------------------------------------------------- +// Logger: the driver's own `log` lines, on the same clock +// ----------------------------------------------------------------------------- + +struct StderrLogger { + started: Instant, +} + +impl log::Log for StderrLogger { + fn enabled(&self, metadata: &log::Metadata) -> bool { + metadata.target().starts_with("acp_client") + } + + fn log(&self, record: &log::Record) { + if self.enabled(record.metadata()) { + eprintln!( + "[log {:>7.2}s] {} {}", + self.started.elapsed().as_secs_f64(), + record.level(), + record.args() + ); + } + } + + fn flush(&self) {} +} + +fn install_logger() { + static LOGGER: OnceLock = OnceLock::new(); + let logger = LOGGER.get_or_init(|| StderrLogger { + started: Instant::now(), + }); + if log::set_logger(logger).is_ok() { + let level = match std::env::var("ACP_REAL_BRIDGE_LOG").as_deref() { + Ok("debug") => log::LevelFilter::Debug, + Ok("trace") => log::LevelFilter::Trace, + _ => log::LevelFilter::Info, + }; + log::set_max_level(level); + } +} + +// ----------------------------------------------------------------------------- +// Session: one real connection +// ----------------------------------------------------------------------------- + +struct Session { + connection: SessionConnection, + handle: TurnControlHandle, + trace: Arc, + store: Arc, + cancel: CancellationToken, + settled_rx: Option>, + dir: PathBuf, +} + +impl Session { + /// Spawn the installed bridge in a fresh temp dir, with the process's own + /// environment as the snapshot (so the shim and its Node resolve the way + /// they do under Staged) and, like the app, a post-turn background hold. + async fn connect(label: &str, steering_disabled: bool) -> Self { + install_logger(); + pin_bridge_dir(); + let provider = provider_id(); + let driver = AcpDriver::new(&provider) + .unwrap_or_else(|e| panic!("{e}: put the bridge on PATH (see the module docs)")) + .with_env_snapshot(std::env::vars().collect()) + .with_background_hold(Some(BackgroundHoldConfig::default())) + .with_steering_disabled(steering_disabled); + let dir = std::env::temp_dir().join(format!( + "acp-real-bridge-{label}-{}-{}", + std::process::id(), + Instant::now().elapsed().as_nanos() + )); + std::fs::create_dir_all(&dir).expect("temp working dir"); + let trace = Arc::new(Trace::new(label)); + let writer: Arc = Arc::clone(&trace) as Arc; + let store = Arc::new(RecordingStore::default()); + let dyn_store: Arc = Arc::clone(&store) as Arc; + let cancel = CancellationToken::new(); + eprintln!( + "[{label}] connecting {provider} ({}) in {} (steering_disabled={steering_disabled})", + driver.binary_path.display(), + dir.display() + ); + let mut connection = driver + .connect( + &format!("real-{label}"), + &dir, + &dyn_store, + &writer, + &cancel, + None, + &[], + None, + ) + .await + .expect("bridge should spawn"); + let handle = connection.turn_control_handle(); + let settled_rx = connection.take_settled_receiver(); + Self { + connection, + handle, + trace, + store, + cancel, + settled_rx, + dir, + } + } + + fn basename(&self) -> String { + self.dir.file_name().unwrap().to_string_lossy().into_owned() + } + + /// The bridge processes this test spawned: direct children whose command + /// line names the bridge. + fn bridge_pids(&self) -> Vec { + let out = std::process::Command::new("pgrep") + .args(["-P", &std::process::id().to_string()]) + .output() + .expect("pgrep"); + let mut pids = Vec::new(); + for line in String::from_utf8_lossy(&out.stdout).lines() { + let Ok(pid) = line.trim().parse::() else { + continue; + }; + let ps = std::process::Command::new("ps") + .args(["-o", "command=", "-p", &pid.to_string()]) + .output() + .expect("ps"); + let command = String::from_utf8_lossy(&ps.stdout); + if command.contains("-acp") { + pids.push(pid); + } + } + pids + } + + /// Cancel the session (ending the hold) and wait for teardown. + async fn finish(mut self) { + self.trace.dump_rows(); + self.cancel.cancel(); + if let Some(rx) = self.settled_rx.take() { + match tokio::time::timeout(SETTLE_WAIT, rx).await { + Ok(Ok(settled)) => eprintln!( + "[{}] settled: {:?} / {:?}", + self.trace.label, settled.outcome, settled.reason + ), + Ok(Err(_)) => eprintln!("[{}] settled sender dropped", self.trace.label), + Err(_) => panic!("[{}] teardown hung after cancel", self.trace.label), + } + } + let _ = std::fs::remove_dir_all(&self.dir); + } +} + +fn activity(ev: &Ev) -> bool { + matches!( + ev, + Ev::Append(_) | Ev::ToolCall { .. } | Ev::Permission { .. } + ) +} + +fn user_rows(trace: &Trace) -> usize { + trace.count(|ev| matches!(ev, Ev::User(_))) +} + +// ----------------------------------------------------------------------------- +// Scenario 1: steer during a live turn +// ----------------------------------------------------------------------------- + +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_steer_mid_turn_redirects_the_live_turn() { + tokio::task::LocalSet::new() + .run_until(async { + let mut session = Session::connect("steer-live", false).await; + let trace = Arc::clone(&session.trace); + let handle = session.handle.clone(); + let prompt_done = Rc::new(Cell::new(false)); + let prompt_done_for_steer = Rc::clone(&prompt_done); + let steer = tokio::task::spawn_local(async move { + trace + .wait_for("the turn's first activity", TURN_START_TIMEOUT, |evs| { + evs.iter().any(activity) + }) + .await; + tokio::time::sleep(Duration::from_secs(2)).await; + let sent_at = trace.elapsed(); + let outcome = handle.steer(STEER_MESSAGE, &[]).await; + let answered_at = trace.elapsed(); + let turn_still_live = !prompt_done_for_steer.get(); + eprintln!( + "[steer-live] steer sent at {:.2}s, answered {outcome:?} at {:.2}s (turn still live: {turn_still_live})", + sent_at.as_secs_f64(), + answered_at.as_secs_f64() + ); + (turn_still_live, answered_at, outcome) + }); + + let turn = tokio::time::timeout(TURN_TIMEOUT, session.connection.prompt(COUNT_PROMPT, &[])) + .await + .expect("the steered turn should end"); + prompt_done.set(true); + let resolved_at = session.trace.elapsed(); + eprintln!("[steer-live] prompt() resolved {turn:?} at {:.2}s", resolved_at.as_secs_f64()); + let (turn_still_live, answered_at, outcome) = steer.await.expect("steer task"); + assert!(turn_still_live, "the steer must land inside the live turn, not after it"); + + assert_eq!(outcome, Ok(SteerOutcome::Injected)); + assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + let rows = session.trace.rows(); + let users: Vec = rows + .iter() + .enumerate() + .filter(|(_, row)| row.is_user()) + .map(|(index, _)| index) + .collect(); + assert_eq!(users.len(), 1, "exactly one user row: {rows:#?}"); + assert!( + users[0] > 0, + "the steer lands after the turn's own output: {rows:#?}" + ); + assert_eq!(rows[users[0]], Row::User(STEER_MESSAGE.to_string())); + let after = session.trace.assistant_text_after_user(); + assert!( + after.contains("STEERED-DONE"), + "the agent should answer the steer; text after the user row: {after:?}" + ); + assert!( + after.contains(&session.basename()), + "the agent should answer the steer's question; got {after:?}" + ); + // Course change: the first thing after the steer was taken should + // follow within model latency, not after the count ran out. + let user_at = session + .trace + .first_at(Duration::ZERO, |ev| matches!(ev, Ev::User(_))) + .unwrap(); + let next = session + .trace + .first_at(user_at, |ev| matches!(ev, Ev::Append(_))) + .expect("assistant text after the user row"); + eprintln!( + "[steer-live] first assistant text {:.2}s after the steer was answered", + (next - answered_at).as_secs_f64() + ); + assert!( + next - answered_at < Duration::from_secs(30), + "the agent should change course within a few seconds" + ); + assert_eq!(session.store.agent_session_ids.lock().unwrap().len(), 1); + session.finish().await; + }) + .await; +} + +// ----------------------------------------------------------------------------- +// Scenario 2: steer while the bridge holds the turn open for background work +// ----------------------------------------------------------------------------- + +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_steer_while_the_bridge_holds_the_turn_for_a_subagent() { + tokio::task::LocalSet::new() + .run_until(async { + let mut session = Session::connect("steer-held", false).await; + let trace = Arc::clone(&session.trace); + let handle = session.handle.clone(); + let prompt_done = Rc::new(Cell::new(false)); + let prompt_done_for_steer = Rc::clone(&prompt_done); + let steer = tokio::task::spawn_local(async move { + trace + .wait_for("LAUNCHED", TURN_START_TIMEOUT, |evs| { + streamed_text(evs).contains("LAUNCHED") + }) + .await; + tokio::time::sleep(Duration::from_secs(3)).await; + let held = !prompt_done_for_steer.get(); + eprintln!("[steer-held] bridge still holding the turn 3s after LAUNCHED: {held}"); + let outcome = handle.steer(HELD_STEER_MESSAGE, &[]).await; + eprintln!("[steer-held] steer answered {outcome:?} at {:.2}s", trace.elapsed().as_secs_f64()); + (held, outcome) + }); + + let turn = tokio::time::timeout(TURN_TIMEOUT, session.connection.prompt(HELD_TURN_PROMPT, &[])) + .await + .expect("the held turn should end"); + prompt_done.set(true); + eprintln!("[steer-held] prompt() resolved {turn:?} at {:.2}s", session.trace.elapsed().as_secs_f64()); + let (held, outcome) = steer.await.expect("steer task"); + if held { + eprintln!("SCENARIO 2: bridge held the turn; steer outcome {outcome:?}"); + assert_eq!(outcome, Ok(SteerOutcome::Injected)); + } else { + // Not every bridge holds a turn for background subagents + // (codex-acp does not): the steer then lands in the client's + // own post-turn hold — a fresh turn on the same connection — + // or, once that hold has quiesced and torn down, fails fast + // for the caller to route as a resume. + eprintln!( + "SCENARIO 2: bridge did NOT hold the turn open (prompt resolved before the steer): \ + {outcome:?}" + ); + match &outcome { + Ok(SteerOutcome::NewTurnStarted) => {} + Err(e) if e.contains("no longer running") => { + assert_eq!(user_rows(&session.trace), 0); + session.finish().await; + return; + } + other => panic!("unexpected steer outcome out of turn: {other:?}"), + } + } + session + .trace + .wait_for("the reply to the steer", REPLY_TIMEOUT, |evs| { + streamed_text(evs).contains("HELD-DONE") + }) + .await; + assert_eq!(user_rows(&session.trace), 1); + session.finish().await; + }) + .await; +} + +// ----------------------------------------------------------------------------- +// Scenario 3: interrupt-and-reprompt (steering disabled on the driver) +// ----------------------------------------------------------------------------- + +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_interrupt_reprompts_on_the_same_child() { + tokio::task::LocalSet::new() + .run_until(async { + let mut session = Session::connect("interrupt", true).await; + let pids_before = session.bridge_pids(); + assert_eq!(pids_before.len(), 1, "one bridge child: {pids_before:?}"); + let trace = Arc::clone(&session.trace); + let handle = session.handle.clone(); + let prompt_done = Rc::new(Cell::new(false)); + let prompt_done_for_interrupt = Rc::clone(&prompt_done); + let interrupt = tokio::task::spawn_local(async move { + trace + .wait_for("the turn's first activity", TURN_START_TIMEOUT, |evs| { + evs.iter().any(activity) + }) + .await; + tokio::time::sleep(Duration::from_secs(2)).await; + // The caller's routing: steer first, interrupt on refusal. + let steer = handle.steer(INTERRUPT_MESSAGE, &[]).await; + eprintln!("[interrupt] steer refused as expected: {steer:?}"); + assert!(steer + .as_ref() + .is_err_and(|e| e.contains("does not advertise steering"))); + assert_eq!(user_rows(&trace), 0, "a refused steer leaves no row"); + let sent_at = trace.elapsed(); + let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[]).await; + let answered_at = trace.elapsed(); + // No terminal state mid-way: the interrupted prompt() is + // still pending (and so, under the hold, is teardown). + let prompt_still_pending = !prompt_done_for_interrupt.get(); + eprintln!( + "[interrupt] interrupt sent at {:.2}s, answered {outcome:?} at {:.2}s (prompt still pending: {prompt_still_pending})", + sent_at.as_secs_f64(), + answered_at.as_secs_f64() + ); + (prompt_still_pending, outcome) + }); + + let turn = tokio::time::timeout(TURN_TIMEOUT, session.connection.prompt(COUNT_PROMPT, &[])) + .await + .expect("the re-prompted turn should end"); + prompt_done.set(true); + let resolved_at = session.trace.elapsed(); + eprintln!("[interrupt] prompt() resolved {turn:?} at {:.2}s", resolved_at.as_secs_f64()); + let (prompt_still_pending, outcome) = interrupt.await.expect("interrupt task"); + let pids_after = session.bridge_pids(); + + assert_eq!(outcome, Ok(())); + assert!( + prompt_still_pending, + "the interrupted prompt() resolves with the final turn, not at the cancel" + ); + assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + assert_eq!(pids_after, pids_before, "the re-prompt runs on the same child"); + assert_eq!(session.store.agent_session_ids.lock().unwrap().len(), 1); + let rows = session.trace.rows(); + let users: Vec = rows + .iter() + .enumerate() + .filter(|(_, row)| row.is_user()) + .map(|(index, _)| index) + .collect(); + assert_eq!(users.len(), 1, "exactly one user row: {rows:#?}"); + assert!(users[0] > 0, "the interrupted turn's output precedes the row: {rows:#?}"); + let after = session.trace.assistant_text_after_user(); + assert!(after.contains("INTERRUPTED-DONE"), "text after the user row: {after:?}"); + assert!(after.contains(&session.basename()), "text after the user row: {after:?}"); + // Permission requests in the second turn still reach the handler: + // if the bridge asked for the first turn's shell calls, it must + // have asked for the `pwd` too. + let user_at = session + .trace + .first_at(Duration::ZERO, |ev| matches!(ev, Ev::User(_))) + .unwrap(); + let permissions_before = session + .trace + .events() + .iter() + .filter(|(at, ev)| *at < user_at && matches!(ev, Ev::Permission { .. })) + .count(); + let permissions_after = session + .trace + .events() + .iter() + .filter(|(at, ev)| *at >= user_at && matches!(ev, Ev::Permission { .. })) + .count(); + eprintln!("[interrupt] permission requests before/after the interrupt: {permissions_before}/{permissions_after}"); + if permissions_before > 0 { + assert!(permissions_after > 0, "turn 2's permission request never reached the handler"); + } + session.finish().await; + }) + .await; +} + +// ----------------------------------------------------------------------------- +// Scenario 4: steer racing the end of the turn +// ----------------------------------------------------------------------------- + +/// Send a steer `when` the short turn is at (or past) its end; whatever the +/// bridge answers, the caller's routing must end with the message delivered +/// exactly once and nothing hung. +async fn race_steer(label: &str, steer_after_first_chunk: bool) { + let mut session = Session::connect(label, false).await; + let trace = Arc::clone(&session.trace); + let handle = session.handle.clone(); + let label_owned = label.to_string(); + let steer = tokio::task::spawn_local(async move { + if steer_after_first_chunk { + trace + .wait_for("the first chunk", TURN_START_TIMEOUT, |evs| { + evs.iter().any(|ev| matches!(ev, Ev::Append(_))) + }) + .await; + } + let sent_at = trace.elapsed(); + let outcome = tokio::time::timeout(REPLY_TIMEOUT, handle.steer(RACE_MESSAGE, &[])) + .await + .expect("a steer at the turn boundary must not hang"); + eprintln!( + "[{label_owned}] steer sent at {:.2}s answered {outcome:?} at {:.2}s", + sent_at.as_secs_f64(), + trace.elapsed().as_secs_f64() + ); + let delivered = match &outcome { + Ok(_) => true, + Err(e) => { + // The caller's fallback: the connection is holding, so this + // is served as a plain new turn. + let fallback = tokio::time::timeout( + REPLY_TIMEOUT, + handle.interrupt_and_reprompt(RACE_MESSAGE, &[]), + ) + .await + .expect("the fallback must not hang"); + eprintln!("[{label_owned}] steer failed ({e}); fallback answered {fallback:?}"); + fallback.is_ok() + } + }; + (outcome, delivered) + }); + + let turn = tokio::time::timeout(TURN_TIMEOUT, session.connection.prompt(RACE_PROMPT, &[])) + .await + .expect("the short turn should end"); + eprintln!( + "[{label}] prompt() resolved {turn:?} at {:.2}s", + session.trace.elapsed().as_secs_f64() + ); + let (outcome, delivered) = steer.await.expect("steer task"); + assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + assert!( + delivered, + "neither the steer nor the fallback delivered: {outcome:?}" + ); + session + .trace + .wait_for("the reply to the steer", REPLY_TIMEOUT, |evs| { + streamed_text(evs).contains("RACE-DONE") + }) + .await; + // Let any trailing frames land before counting rows. + tokio::time::sleep(Duration::from_secs(3)).await; + eprintln!( + "SCENARIO 4 ({label}): steer outcome {outcome:?}; user rows {}", + user_rows(&session.trace) + ); + let rows = session.trace.rows(); + assert_eq!(user_rows(&session.trace), 1, "{rows:#?}"); + // The aborted cycle's result must not re-echo the pre-steer message as a + // row after the user row (claude-agent-acp emits that result; the client + // has to recognise its text as already streamed). + let user = rows.iter().position(Row::is_user).unwrap(); + for before in rows[..user] + .iter() + .filter(|row| matches!(row, Row::Assistant(_))) + { + assert!( + !rows[user + 1..].contains(before), + "pre-steer assistant row echoed after the user row: {rows:#?}" + ); + } + session.finish().await; +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_steer_sent_with_the_prompt_lands_once() { + tokio::task::LocalSet::new() + .run_until(race_steer("race-at-send", false)) + .await; +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_steer_sent_at_the_first_chunk_lands_once() { + tokio::task::LocalSet::new() + .run_until(race_steer("race-at-chunk", true)) + .await; +} diff --git a/crates/acp-client/src/driver/tests/response_completion.rs b/crates/acp-client/src/driver/tests/response_completion.rs index 70d9ecd80..f0bc633c4 100644 --- a/crates/acp-client/src/driver/tests/response_completion.rs +++ b/crates/acp-client/src/driver/tests/response_completion.rs @@ -433,6 +433,7 @@ async fn session_loop_finishes_successfully_for_initial_and_background_artifacts None, if scenario == 7 { None } else { Some(&QUIESCENCE_PROBE) }, Some(&observer), + false, &child_exited, &mut prompt_rx, &mut control_rx, @@ -614,6 +615,40 @@ async fn result_after_a_tool_call_ended_the_message_lands_in_a_new_row() { assert!(finish_requested(&handler)); } +/// A steer pre-empts the generation in flight: the bridge aborts that cycle +/// and still emits its `result`, whose final text is the message the steer +/// callback just closed (seen against claude-agent-acp 0.95, which then +/// re-echoed the whole message as a new row after the user row). Detaching +/// keeps the streamed text, so the result certifies it instead of +/// re-injecting it, and the steered reply opens its own row. +#[tokio::test] +async fn result_of_a_cycle_a_steer_aborted_does_not_echo_the_closed_message() { + let writer = Arc::new(ArtifactWriter::default()); + let handler = artifact_handler(&writer); + start_task(&handler, TaskTrackingMode::Raw).await; + feed(&handler, agent_text(Some("msg-1"), "pong")).await; + // What `serve_steer`'s response callback does once the agent answers + // `injected`: close the row, detach it, record the user row. + handler.writer.finalize().await; + handler.response.lock().await.detach_message(); + feed_sdk_frame(&handler, successful_result("aborted", "pong")).await; + feed(&handler, agent_text(Some("msg-2"), "RACE-DONE")).await; + feed_sdk_frame(&handler, successful_result("steered", "RACE-DONE")).await; + + assert_eq!( + writer.flow(), + vec![ + RowFlow::Append("pong".into()), + RowFlow::Finalize, + // No second "pong": the aborted cycle's result matched the text + // already streamed. The steered reply starts a row of its own. + RowFlow::Finalize, + RowFlow::Append("RACE-DONE".into()), + ] + ); + assert!(writer.recorded_result()); +} + #[tokio::test] async fn result_extending_an_interrupted_message_resumes_its_row() { let writer = Arc::new(ArtifactWriter::default()); diff --git a/crates/acp-client/src/driver/tests/turn_control.rs b/crates/acp-client/src/driver/tests/turn_control.rs index ebd678295..31706e297 100644 --- a/crates/acp-client/src/driver/tests/turn_control.rs +++ b/crates/acp-client/src/driver/tests/turn_control.rs @@ -633,6 +633,7 @@ where None, background_hold, None, + false, &child_exited, &mut prompt_rx, &mut hold_control_rx, @@ -743,6 +744,57 @@ async fn interrupt_cancels_the_turn_and_reprompts_on_the_same_connection() { ); } +/// A control request queued next to an already-queued `prompt()` (both land +/// while setup is still running, say) does not jump it: the prompt is the +/// turn, and the request is served inside it — here an interrupt, which +/// cancels that turn and follows it. Seen against a real bridge, where the +/// unbiased select sometimes made the message the first turn and the +/// prompt the second. +#[tokio::test(flavor = "current_thread")] +async fn a_prompt_queued_before_a_control_request_goes_first() { + let agent = Arc::new(FakeAgent { + advertise_steering: false, + ..FakeAgent::default() + }); + + let (result, _reply_rx, stashed_reply, writer, _handler) = session_with( + &agent, + None, + |handle, agent, _writer, _handler| async move { + // No await before this: it is queued while setup is in flight, + // behind the prompt `session_with` queued before connecting. + let outcome = handle.interrupt_and_reprompt("second", &[]).await; + assert_eq!(outcome, Ok(())); + wait_until("both prompts", || agent.prompts().len() == 2).await; + assert_eq!(agent.prompts(), ["first", "second"]); + assert_eq!( + agent.cancels(), + 1, + "the interrupt cancelled the queued prompt's turn" + ); + agent.finish_prompt(StopReason::EndTurn); + }, + ) + .await; + + assert_eq!( + result, + Ok((AgentRunOutcome::Completed, SessionSettleReason::Immediate)) + ); + assert_eq!( + writer + .calls() + .iter() + .filter(|call| matches!(call, WriterCall::UserMessage(_))) + .count(), + 1 + ); + assert!( + stashed_reply.is_some(), + "the prompt() for the first turn is still pending" + ); +} + #[tokio::test(flavor = "current_thread")] async fn interrupt_during_a_hold_is_served_as_a_new_turn_without_a_cancel() { let agent = Arc::new(FakeAgent { From 986ae6476d77c61e1929339dbe5f1b6efa047bca Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Thu, 29 Oct 2026 11:52:48 +1100 Subject: [PATCH 5/7] fix(acp-client): verify turn control against Goose and resolve pending permissions on interrupt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runs the real-bridge turn-control verification (7125d919) against Goose 1.57.0 — the installed bridge that genuinely does not advertise steering, so the interrupt path is exercised as a caller reaches it rather than by forcing steering off on Claude. One bug fell out, with a fix and two fake-agent regression tests; the Goose runs stay `#[ignore]`d: GOOSE_DISABLE_KEYRING=1 ACP_REAL_BRIDGE=goose \ cargo test -p acp-client real_bridge -- --ignored --nocapture --test-threads=1 cd apps/staged/src-tauri && GOOSE_DISABLE_KEYRING=1 ACP_REAL_BRIDGE=goose \ cargo test --lib continue_repo_session_live -- --ignored --nocapture --test-threads=1 Goose is not a managed bridge: `ACP_REAL_BRIDGE=goose` resolves `goose` from the login shell's PATH (the bundled-dir pin falls through) and runs `goose acp --with-builtin developer --with-builtin extensionmanager`, the provider-table command. It used the machine's configured provider (`databricks_v2`, `databricks-gpt-5-5`). `GOOSE_DISABLE_KEYRING=1` is required headless: without it the CLI blocks in `SecKeychainFindGenericPassword` on a Keychain prompt nobody can click and never answers `session/new` (sampled; `goose run` hung 120 s the same way). Fix - Interrupt while a permission request is pending: the interrupt arm left the request alone (the connection-wide token, once cancelled, would auto-answer every later request). With a request nobody had answered, the cancel went unacknowledged for the full 30 s, the re-prompt was dropped and the turn settled `Cancelled` (scenario c, below). The handler now holds a per-turn child token (`begin_turn_permissions` at every prompt, `cancel_turn_permissions` in the interrupt arm): the interrupted turn's requests resolve `Cancelled` — ACP's rule for a client that cancels — and the next turn's requests are presented as before. Regression: `interrupt_resolves_the_turns_pending_permission_request_before_reprompting` (fails without the cancel: the interrupt waits on the dialog). - The permission decision is awaited off the SDK's dispatch task (`connection.spawn` in the handler), per the SDK's "handlers must not block": inline, a pending dialog stalled every other frame — the prompt response that acknowledges a cancel, a steer's answer, session/update. Regression: `a_pending_permission_decision_does_not_stall_the_connection` (fails with the handler inline: the chunk behind the request never arrives). Harness - `bridge_pids` recognises `goose acp`; the same-child check reads the pid while the re-prompt is live, because Goose emits no raw SDK frames, so the client skips the post-turn hold and the child is gone the moment `prompt()` resolves. - The logger captures every debug line for assertions; scenario 3 asserts the initialize-time `steering advertised: false` and forces steering off only on bridges that advertise it. - `ConnectOptions`: env overrides (`GOOSE_MODE=approve` makes Goose ask permission for every tool call) and `Resume` (agent session id, legacy transcript rows, same dir) over a `BridgeStore` that serves the replay matcher. The writer can hold a permission request open until its token cancels. New scenarios: interrupt with a permission pending, interrupt racing turn end (two timings), resume via `session/load`. Steer scenarios report N/A on a non-steering bridge; the steer/interrupt race scenarios accept the hold-less bridge's fast "no longer running" failure as the resume route. - The `staged` live MCP test expects `interrupted_and_reprompted` on Goose without forcing. Scenarios (Goose; timings from the post-fix run) a. Initialize: PASS. `agentCapabilities.loadSession: true`, no top-level `_meta.steering`; `steer()` refused locally ("does not advertise steering"), no user row, caller fell through to interrupt. b. Interrupt-and-reprompt mid-turn: PASS. Cancel acked in 9 ms (10 ms in the first run), same child pid before and during the re-prompt, no terminal status mid-way, one user row in assistant → user → assistant order, `INTERRUPTED-DONE ` in turn 2, `prompt()` resolved once. c. Interrupt with a `session/request_permission` pending (`GOOSE_MODE= approve`; Goose asks for each shell call): FAIL → fixed. Before: no ack in 30 s, re-prompt dropped, `prompt()` → Cancelled. After: held request resolved `Cancelled` by the interrupt, cancel acked in 27 ms, user row written, turn 2's `pwd` permission request reached the handler and was answered, `INTERRUPTED-DONE`. d. Interrupt racing turn end: PASS. Sent with the prompt: served in the live turn, acked after 2.5–3.3 s (Goose finishes its in-flight model call first), one user row, no echoed rows. Sent at the first chunk: the one-word turn was already over and the hold-less connection torn down; the interrupt failed fast ("no longer running") with no row — the caller's resume route. The steer-race variants behave the same. e. Resume: PASS. Goose advertises `loadSession`; after settle a new child loaded the session in the same dir, replayed history was not recorded as rows, and the agent returned the code word from turn 1. So `continue_repo_session` resumes terminal Goose sessions rather than returning `not_resumable`; the capability-less case still surfaces as the existing connect error, unchanged here. f. Through `continue_repo_session` with a Goose session in `SessionRegistry`: PASS, forced and unforced — `delivery: interrupted_and_reprompted`, row `running` throughout, `activity.last_user_message_at` = the single user row, `FOLLOW-UP-DONE`. Also noted: Goose's first `session/new` on this machine failed inside its own sessions DB (sqlx worker panic right after its v16→v17 migration) and worked on every later run; Goose-side, not reproducible after migration. Verified: `cargo test -p acp-client --lib` (199 passed; regression tests confirmed failing with their fix reverted), `cargo clippy --workspace` and src-tauri clippy at `-D warnings`, rustfmt clean. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/project_mcp.rs | 17 +- crates/acp-client/src/driver.rs | 121 ++- .../src/driver/tests/real_bridge.rs | 716 ++++++++++++++++-- .../src/driver/tests/turn_control.rs | 183 ++++- 4 files changed, 926 insertions(+), 111 deletions(-) diff --git a/apps/staged/src-tauri/src/project_mcp.rs b/apps/staged/src-tauri/src/project_mcp.rs index d45e5b832..15b2abd1c 100644 --- a/apps/staged/src-tauri/src/project_mcp.rs +++ b/apps/staged/src-tauri/src/project_mcp.rs @@ -1777,11 +1777,13 @@ mod tests { /// cargo test -p staged --lib continue_repo_session_live -- --ignored --nocapture --test-threads=1 /// ``` /// - /// `ACP_REAL_BRIDGE` picks the provider id (`claude` by default) and + /// `ACP_REAL_BRIDGE` picks the provider id (`claude` by default; `goose` + /// resolves from the login shell's `PATH` and, headless, needs + /// `GOOSE_DISABLE_KEYRING=1` or it blocks on a Keychain prompt) and /// `ACP_REAL_BRIDGE_DIR` the install dir (`~/.staged/packages/bin`). The /// resume (`resumed`) and `queued` paths are not covered here: a resume /// needs the app handle the tool carries, which tests cannot build. - async fn continue_repo_session_live(steering_disabled: bool, expected_delivery: &str) { + async fn continue_repo_session_live(steering_disabled: bool) { const LONG_TURN: &str = "Count from 1 to 40 slowly. For EACH number: first run the \ shell command `sleep 1` with your shell tool as its own separate tool call (never a \ loop, never a script), then write that number on its own line. Keep going until 40. \ @@ -1800,6 +1802,13 @@ mod tests { acp_client::set_bundled_tools_dir(dir); } let provider = std::env::var("ACP_REAL_BRIDGE").unwrap_or_else(|_| "claude".to_string()); + // Goose does not advertise steering, so the route falls through to + // the interrupt whether or not it is forced off. + let expected_delivery = if steering_disabled || provider == "goose" { + "interrupted_and_reprompted" + } else { + "injected" + }; let driver = acp_client::AcpDriver::new(&provider) .expect("the bridge must be installed (see the test docs)") .with_env_snapshot(std::env::vars().collect()) @@ -1925,13 +1934,13 @@ mod tests { #[tokio::test(flavor = "current_thread")] #[ignore = "spawns the installed ACP bridge and calls the model"] async fn continue_repo_session_live_steers_a_running_session() { - continue_repo_session_live(false, "injected").await; + continue_repo_session_live(false).await; } #[tokio::test(flavor = "current_thread")] #[ignore = "spawns the installed ACP bridge and calls the model"] async fn continue_repo_session_live_interrupts_when_steering_is_unavailable() { - continue_repo_session_live(true, "interrupted_and_reprompted").await; + continue_repo_session_live(true).await; } /// A session in its startup window (row `running`, no connection yet) diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index ba6589ecf..ef2245ace 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -2198,9 +2198,17 @@ impl AcpDriver { .builder() .name("staged-acp-client") .on_receive_request( - async move |args: RequestPermissionRequest, responder, _connection| { - let response = permission_handler.request_permission(args).await?; - responder.respond(response) + async move |args: RequestPermissionRequest, responder, connection| { + // The decision may wait on a person. Handlers run on + // the connection's single dispatch task, so awaiting + // it here would stall every other frame — the prompt + // response that acknowledges a cancel, a steer's + // answer — until someone clicks. + let handler = Arc::clone(&permission_handler); + connection.spawn(async move { + let response = handler.request_permission(args).await?; + responder.respond(response) + }) }, agent_client_protocol::on_receive_request!(), ) @@ -3691,7 +3699,15 @@ struct AcpNotificationHandler { phase: Mutex, /// Signalled when replay matching determines all DB messages have been replayed. replay_done: tokio::sync::Notify, + /// Connection-wide: cancelled once, at a user cancel or teardown, after + /// which every permission request is answered `Cancelled`. permission_cancel_token: CancellationToken, + /// The turn in flight's child of `permission_cancel_token` — what pending + /// requests actually wait on, so an interrupt can resolve *this* turn's + /// requests `Cancelled` (ACP's rule for a client that cancels) without + /// burning the connection's token for the turns that follow. Rotated by + /// [`Self::begin_turn_permissions`] at every prompt. + turn_permission_token: std::sync::Mutex, /// This connection's task-tracking mode, fixed by [`setup_acp_session`] /// right after the initialize response — race-free, because `session/new` /// (and hence any update) only happens after initialize resolves. Unset @@ -3745,12 +3761,14 @@ impl AcpNotificationHandler { } }; + let turn_permission_token = std::sync::Mutex::new(permission_cancel_token.child_token()); Self { writer, response: Mutex::new(response_completion::ResponseState::default()), phase: Mutex::new(phase), replay_done: tokio::sync::Notify::new(), permission_cancel_token, + turn_permission_token, task_tracking_mode: OnceLock::new(), background_tasks: Mutex::new(BackgroundTaskSet::default()), typed_tasks: Mutex::new(TypedAsyncTaskSet::default()), @@ -3920,6 +3938,24 @@ impl AcpNotificationHandler { self.permission_cancel_token.cancel(); } + /// A turn is about to be prompted: permission requests from here on wait + /// on a fresh token of its own. + fn begin_turn_permissions(&self) { + *self.turn_permission_token.lock().unwrap() = self.permission_cancel_token.child_token(); + } + + /// Resolve the turn in flight's pending permission requests `Cancelled`, + /// leaving the connection's token — and so the next turn's requests — + /// untouched. Until [`Self::begin_turn_permissions`] rotates the token, a + /// late request from the cancelled turn is answered `Cancelled` too. + fn cancel_turn_permissions(&self) { + self.turn_permission_token.lock().unwrap().cancel(); + } + + fn turn_permission_token(&self) -> CancellationToken { + self.turn_permission_token.lock().unwrap().clone() + } + async fn is_replay_complete(&self) -> bool { let phase = self.phase.lock().await; match &*phase { @@ -3957,7 +3993,8 @@ impl AcpNotificationHandler { &self, args: RequestPermissionRequest, ) -> agent_client_protocol::Result { - if self.permission_cancel_token.is_cancelled() { + let turn_token = self.turn_permission_token(); + if turn_token.is_cancelled() { return Ok(RequestPermissionResponse::new( RequestPermissionOutcome::Cancelled, )); @@ -3985,7 +4022,7 @@ impl AcpNotificationHandler { )); } - let response = self.resolve_permission(request.clone()).await; + let response = self.resolve_permission(request.clone(), turn_token).await; self.pending_permissions .lock() .await @@ -3993,7 +4030,11 @@ impl AcpNotificationHandler { Ok(response) } - async fn resolve_permission(&self, request: AcpPermissionRequest) -> RequestPermissionResponse { + async fn resolve_permission( + &self, + request: AcpPermissionRequest, + turn_token: CancellationToken, + ) -> RequestPermissionResponse { if let Some(origin) = request.origin.clone() { // Out-of-turn: the tool call must be one this connection actually // saw announced, or there is nothing to correlate the request to. @@ -4036,10 +4077,7 @@ impl AcpNotificationHandler { ); } - let decision = self - .writer - .request_permission(request, self.permission_cancel_token.clone()) - .await; + let decision = self.writer.request_permission(request, turn_token).await; permission_response_for_decision(decision) } @@ -5908,6 +5946,7 @@ async fn run_prompt_turn( let content_blocks = build_prompt_content_blocks(prompt, images, supports_images); let prompt_request = PromptRequest::new(setup.agent_session_id.clone(), content_blocks); + handler.begin_turn_permissions(); handler.transition_to_live().await; handler .response @@ -5951,12 +5990,16 @@ async fn run_prompt_turn( log::info!( "ACP session {our_session_id}: interrupting the turn in flight for a new message" ); - // Same wire sequence as the user-cancel arm above, minus - // `cancel_pending_permissions()`: that token is created - // once per connection and, once cancelled, auto-answers - // every later permission request — which would silently - // break the next turn. The agent resolves its own pending - // permission requests on `session/cancel`. + // Same wire sequence as the user-cancel arm above, but + // only *this turn's* pending permission requests are + // resolved `Cancelled` — the connection-wide token would + // auto-answer every later request and silently break the + // next turn. Resolving them is not optional: with a + // request nobody had answered still pending, the cancel + // went unacknowledged for the full timeout and the + // message was dropped (seen against Goose) — and once + // the turn is gone the dialog is stale anyway. + handler.cancel_turn_permissions(); if let Err(e) = send_session_cancel(connection, &setup.agent_session_id) { let _ = reply.send(Err(e.clone())); return Err(e); @@ -12207,11 +12250,19 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE Event(RecordedMetadata), ToolCall(String), ToolCallMetadata(RecordedMetadata), + /// A permission request reached the writer (by tool-call id). + Permission(String), + /// A held permission request was resolved by its token cancelling. + PermissionCancelled(String), } #[derive(Default)] struct TranscriptWriter { calls: Mutex>, + /// Hold the next permission request open — a dialog nobody has + /// answered — until its cancel token fires or the test releases it. + hold_next_permission: AtomicBool, + release_permission: tokio::sync::Notify, } impl TranscriptWriter { @@ -12241,7 +12292,13 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE self.calls() .into_iter() .filter(|call| { - !matches!(call, WriterCall::Event(_) | WriterCall::ToolCallMetadata(_)) + !matches!( + call, + WriterCall::Event(_) + | WriterCall::ToolCallMetadata(_) + | WriterCall::Permission(_) + | WriterCall::PermissionCancelled(_) + ) }) .collect() } @@ -12307,6 +12364,36 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE async fn record_tool_result(&self, _tool_call_id: &str, _content: &str) {} + async fn request_permission( + &self, + request: AcpPermissionRequest, + cancel_token: CancellationToken, + ) -> AcpPermissionDecision { + self.calls + .lock() + .unwrap() + .push(WriterCall::Permission(request.tool_call_id.clone())); + if cancel_token.is_cancelled() { + return AcpPermissionDecision::Cancelled; + } + if self + .hold_next_permission + .swap(false, std::sync::atomic::Ordering::SeqCst) + { + tokio::select! { + _ = cancel_token.cancelled() => { + self.calls + .lock() + .unwrap() + .push(WriterCall::PermissionCancelled(request.tool_call_id.clone())); + return AcpPermissionDecision::Cancelled; + } + _ = self.release_permission.notified() => {} + } + } + autoapprove_permission_decision(&request) + } + async fn record_acp_event_metadata(&self, metadata: AcpEventMetadata) { self.calls .lock() diff --git a/crates/acp-client/src/driver/tests/real_bridge.rs b/crates/acp-client/src/driver/tests/real_bridge.rs index 948b9ccf8..fc5f5e4f9 100644 --- a/crates/acp-client/src/driver/tests/real_bridge.rs +++ b/crates/acp-client/src/driver/tests/real_bridge.rs @@ -12,17 +12,23 @@ //! managed install (`~/.staged/packages/bin`) — pinned explicitly because //! plain resolution goes through the login shell's `PATH`, where a stale //! Homebrew `claude-agent-acp` can shadow it. `ACP_REAL_BRIDGE=` -//! picks the bridge (`claude`, the default, or `codex`). Every writer event and every `log` line is printed to stderr -//! with its offset from connect, so under `--nocapture` the output is the -//! record of what the bridge did; `ACP_REAL_BRIDGE_LOG=debug` raises the -//! `log` level. The checks are the invariants, not the agent's wording: one -//! user row, in transcript order, the turn resolving once, the same child -//! across an interrupt, no hang. +//! picks the bridge (`claude`, the default, `codex`, or `goose`). Goose is not +//! a managed bridge: it resolves from the login shell's `PATH` and runs as +//! `goose acp`, and under a headless runner it needs `GOOSE_DISABLE_KEYRING=1` +//! — otherwise the CLI blocks in `SecKeychainFindGenericPassword` on a Keychain +//! prompt nobody can click, and never answers `session/new`. Every writer +//! event and every `log` line is printed to stderr with its offset from +//! connect, so under `--nocapture` the output is the record of what the bridge +//! did; `ACP_REAL_BRIDGE_LOG=debug` raises the printed `log` level (debug +//! lines are always captured for assertions). The checks are the invariants, +//! not the agent's wording: one user row, in transcript order, the turn +//! resolving once, the same child across an interrupt, no hang. use super::*; use crate::driver::{SessionConnection, SessionSettled, SteerOutcome, TurnControlHandle}; use std::cell::Cell; use std::rc::Rc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::OnceLock; use std::time::Instant; @@ -48,11 +54,25 @@ const HELD_TURN_PROMPT: &str = "Use your Agent (subagent) tool with run_in_backg reply done. As soon as it is launched — without waiting for it — reply with exactly the single \ word: LAUNCHED"; const HELD_STEER_MESSAGE: &str = "New instruction: reply with exactly one line: HELD-DONE"; +const REMEMBER_PROMPT: &str = "Remember this code word: PELICAN-7342. Do not use any tools. Reply \ + with exactly the single word: REMEMBERED"; +const RECALL_PROMPT: &str = "What was the code word I asked you to remember earlier in this \ + conversation? Do not use any tools. Reply with exactly one line: CODEWORD "; +/// How long a held permission request stays open before the writer gives up +/// waiting for the connection to resolve it. +const PERMISSION_HOLD: Duration = Duration::from_secs(45); fn provider_id() -> String { std::env::var("ACP_REAL_BRIDGE").unwrap_or_else(|_| "claude".to_string()) } +/// Goose is the installed bridge that genuinely does not advertise steering; +/// against the others the interrupt path has to be forced with +/// `with_steering_disabled`. +fn bridge_advertises_steering() -> bool { + provider_id() != "goose" +} + /// Pin bridge lookup to the managed install dir (process-wide, once). fn pin_bridge_dir() { let dir = std::env::var_os("ACP_REAL_BRIDGE_DIR") @@ -84,9 +104,20 @@ enum Ev { Append(String), Finalize, User(String), - ToolCall { id: String, title: String }, - ToolResult { id: String }, - Permission { title: String }, + ToolCall { + id: String, + title: String, + }, + ToolResult { + id: String, + }, + Permission { + title: String, + }, + /// A held permission request resolved by the connection cancelling it. + PermissionCancelled { + title: String, + }, Event(String), } @@ -115,6 +146,11 @@ struct Trace { label: String, started: Instant, events: std::sync::Mutex>, + /// Hold the next permission request open — the way a dialog nobody has + /// answered yet does — until the connection cancels it, the test releases + /// it, or [`PERMISSION_HOLD`] passes. + hold_next_permission: AtomicBool, + release_permission: tokio::sync::Notify, } impl Trace { @@ -123,6 +159,8 @@ impl Trace { label: label.to_string(), started: Instant::now(), events: std::sync::Mutex::new(Vec::new()), + hold_next_permission: AtomicBool::new(false), + release_permission: tokio::sync::Notify::new(), } } @@ -191,7 +229,10 @@ impl Trace { flush(&mut open, &mut rows); rows.push(Row::Tool(title)); } - Ev::ToolResult { .. } | Ev::Permission { .. } | Ev::Event(_) => {} + Ev::ToolResult { .. } + | Ev::Permission { .. } + | Ev::PermissionCancelled { .. } + | Ev::Event(_) => {} } } flush(&mut open, &mut rows); @@ -280,14 +321,28 @@ impl MessageWriter for Trace { request: AcpPermissionRequest, cancel_token: CancellationToken, ) -> AcpPermissionDecision { + let title = request.tool_title.clone().unwrap_or_default(); self.push(Ev::Permission { - title: request.tool_title.clone().unwrap_or_default(), + title: title.clone(), }); if cancel_token.is_cancelled() { - AcpPermissionDecision::Cancelled - } else { - autoapprove_permission_decision(&request) + return AcpPermissionDecision::Cancelled; + } + if self.hold_next_permission.swap(false, Ordering::SeqCst) { + tokio::select! { + _ = cancel_token.cancelled() => { + self.push(Ev::PermissionCancelled { title }); + return AcpPermissionDecision::Cancelled; + } + _ = self.release_permission.notified() => { + self.push(Ev::Event("permission-hold-released".to_string())); + } + _ = tokio::time::sleep(PERMISSION_HOLD) => { + self.push(Ev::Event("permission-hold-expired".to_string())); + } + } } + autoapprove_permission_decision(&request) } } @@ -297,6 +352,9 @@ impl MessageWriter for Trace { struct StderrLogger { started: Instant, + /// Lines at or above this level are printed; every line is captured. + print_level: log::LevelFilter, + lines: std::sync::Mutex>, } impl log::Log for StderrLogger { @@ -306,42 +364,98 @@ impl log::Log for StderrLogger { fn log(&self, record: &log::Record) { if self.enabled(record.metadata()) { - eprintln!( - "[log {:>7.2}s] {} {}", - self.started.elapsed().as_secs_f64(), - record.level(), - record.args() - ); + let line = format!("{} {}", record.level(), record.args()); + if record.level() <= self.print_level { + eprintln!( + "[log {:>7.2}s] {line}", + self.started.elapsed().as_secs_f64() + ); + } + self.lines.lock().unwrap().push(line); } } fn flush(&self) {} } +static LOGGER: OnceLock = OnceLock::new(); + fn install_logger() { - static LOGGER: OnceLock = OnceLock::new(); let logger = LOGGER.get_or_init(|| StderrLogger { started: Instant::now(), - }); - if log::set_logger(logger).is_ok() { - let level = match std::env::var("ACP_REAL_BRIDGE_LOG").as_deref() { + print_level: match std::env::var("ACP_REAL_BRIDGE_LOG").as_deref() { Ok("debug") => log::LevelFilter::Debug, Ok("trace") => log::LevelFilter::Trace, _ => log::LevelFilter::Info, - }; - log::set_max_level(level); + }, + lines: std::sync::Mutex::new(Vec::new()), + }); + if log::set_logger(logger).is_ok() { + log::set_max_level(log::LevelFilter::Debug); } } +/// Whether the driver logged a line containing `needle` (any level). +fn logged(needle: &str) -> bool { + LOGGER.get().is_some_and(|logger| { + logger + .lines + .lock() + .unwrap() + .iter() + .any(|line| line.contains(needle)) + }) +} + // ----------------------------------------------------------------------------- // Session: one real connection // ----------------------------------------------------------------------------- +/// The store as the app would hold it: the agent's session id, and the +/// transcript rows the replay matcher is handed on a resume. +#[derive(Default)] +struct BridgeStore { + agent_session_ids: std::sync::Mutex>, + messages: std::sync::Mutex>, +} + +#[async_trait::async_trait] +impl Store for BridgeStore { + fn set_agent_session_id(&self, session_id: &str, agent_session_id: &str) -> Result<(), String> { + self.agent_session_ids + .lock() + .unwrap() + .push((session_id.to_string(), agent_session_id.to_string())); + Ok(()) + } + + fn get_session_messages(&self, _session_id: &str) -> Result, String> { + Ok(self.messages.lock().unwrap().clone()) + } +} + +/// An earlier connection's session to load instead of creating a new one. +struct Resume { + agent_session_id: String, + /// That connection's transcript as `(role, content)` rows. + messages: Vec<(String, String)>, + /// Its working dir — a load happens where the session was created. + dir: PathBuf, +} + +#[derive(Default)] +struct ConnectOptions { + steering_disabled: bool, + /// Environment set over the process's own for the bridge. + env: Vec<(&'static str, &'static str)>, + resume: Option, +} + struct Session { connection: SessionConnection, handle: TurnControlHandle, trace: Arc, - store: Arc, + store: Arc, cancel: CancellationToken, settled_rx: Option>, dir: PathBuf, @@ -352,29 +466,54 @@ impl Session { /// environment as the snapshot (so the shim and its Node resolve the way /// they do under Staged) and, like the app, a post-turn background hold. async fn connect(label: &str, steering_disabled: bool) -> Self { + Self::connect_with( + label, + ConnectOptions { + steering_disabled, + ..ConnectOptions::default() + }, + ) + .await + } + + async fn connect_with(label: &str, options: ConnectOptions) -> Self { install_logger(); pin_bridge_dir(); let provider = provider_id(); + let mut env: Vec<(String, String)> = std::env::vars().collect(); + for (key, value) in &options.env { + env.retain(|(name, _)| name != key); + env.push((key.to_string(), value.to_string())); + } let driver = AcpDriver::new(&provider) .unwrap_or_else(|e| panic!("{e}: put the bridge on PATH (see the module docs)")) - .with_env_snapshot(std::env::vars().collect()) + .with_env_snapshot(env) .with_background_hold(Some(BackgroundHoldConfig::default())) - .with_steering_disabled(steering_disabled); - let dir = std::env::temp_dir().join(format!( - "acp-real-bridge-{label}-{}-{}", - std::process::id(), - Instant::now().elapsed().as_nanos() - )); + .with_steering_disabled(options.steering_disabled); + let dir = match &options.resume { + Some(resume) => resume.dir.clone(), + None => std::env::temp_dir().join(format!( + "acp-real-bridge-{label}-{}-{}", + std::process::id(), + Instant::now().elapsed().as_nanos() + )), + }; std::fs::create_dir_all(&dir).expect("temp working dir"); let trace = Arc::new(Trace::new(label)); let writer: Arc = Arc::clone(&trace) as Arc; - let store = Arc::new(RecordingStore::default()); + let store = Arc::new(BridgeStore::default()); + if let Some(resume) = &options.resume { + *store.messages.lock().unwrap() = resume.messages.clone(); + } let dyn_store: Arc = Arc::clone(&store) as Arc; let cancel = CancellationToken::new(); eprintln!( - "[{label}] connecting {provider} ({}) in {} (steering_disabled={steering_disabled})", + "[{label}] connecting {provider} ({}) in {} (steering_disabled={}, env={:?}, resume={:?})", driver.binary_path.display(), - dir.display() + dir.display(), + options.steering_disabled, + options.env, + options.resume.as_ref().map(|resume| &resume.agent_session_id) ); let mut connection = driver .connect( @@ -383,7 +522,10 @@ impl Session { &dyn_store, &writer, &cancel, - None, + options + .resume + .as_ref() + .map(|resume| resume.agent_session_id.as_str()), &[], None, ) @@ -406,32 +548,43 @@ impl Session { self.dir.file_name().unwrap().to_string_lossy().into_owned() } - /// The bridge processes this test spawned: direct children whose command - /// line names the bridge. - fn bridge_pids(&self) -> Vec { - let out = std::process::Command::new("pgrep") - .args(["-P", &std::process::id().to_string()]) - .output() - .expect("pgrep"); - let mut pids = Vec::new(); - for line in String::from_utf8_lossy(&out.stdout).lines() { - let Ok(pid) = line.trim().parse::() else { - continue; - }; - let ps = std::process::Command::new("ps") - .args(["-o", "command=", "-p", &pid.to_string()]) - .output() - .expect("ps"); - let command = String::from_utf8_lossy(&ps.stdout); - if command.contains("-acp") { - pids.push(pid); + /// The agent's id for this session, as the store recorded it at setup. + fn agent_session_id(&self) -> String { + self.store + .agent_session_ids + .lock() + .unwrap() + .first() + .map(|(_, id)| id.clone()) + .expect("the store should have the agent's session id") + } + + /// This session's transcript as the app's store would replay it: the + /// prompt as the first user row, then the recorded assistant and user + /// rows. + fn legacy_messages(&self, prompt: &str) -> Vec<(String, String)> { + let mut messages = vec![("user".to_string(), prompt.to_string())]; + for row in self.trace.rows() { + match row { + Row::Assistant(text) => messages.push(("assistant".to_string(), text)), + Row::User(text) => messages.push(("user".to_string(), text)), + Row::Tool(_) => {} } } - pids + messages } /// Cancel the session (ending the hold) and wait for teardown. - async fn finish(mut self) { + async fn finish(self) { + self.finish_inner(true).await; + } + + /// [`Self::finish`], keeping the working dir for a later resume. + async fn finish_keeping_dir(self) { + self.finish_inner(false).await; + } + + async fn finish_inner(mut self, remove_dir: bool) { self.trace.dump_rows(); self.cancel.cancel(); if let Some(rx) = self.settled_rx.take() { @@ -444,10 +597,61 @@ impl Session { Err(_) => panic!("[{}] teardown hung after cancel", self.trace.label), } } - let _ = std::fs::remove_dir_all(&self.dir); + if remove_dir { + let _ = std::fs::remove_dir_all(&self.dir); + } } } +/// The bridge processes this test spawned: direct children whose command +/// line names the bridge. +fn bridge_pids() -> Vec { + let out = std::process::Command::new("pgrep") + .args(["-P", &std::process::id().to_string()]) + .output() + .expect("pgrep"); + let mut pids = Vec::new(); + for line in String::from_utf8_lossy(&out.stdout).lines() { + let Ok(pid) = line.trim().parse::() else { + continue; + }; + let ps = std::process::Command::new("ps") + .args(["-o", "command=", "-p", &pid.to_string()]) + .output() + .expect("ps"); + let command = String::from_utf8_lossy(&ps.stdout); + // `claude-agent-acp` / `codex-acp` by name, `goose acp` by subcommand. + if command.contains("-acp") || command.contains(" acp") { + pids.push(pid); + } + } + pids +} + +/// Wait for every bridge child of this process to be gone. +async fn wait_for_bridge_exit(timeout: Duration) { + let deadline = Instant::now() + timeout; + while !bridge_pids().is_empty() { + assert!( + Instant::now() < deadline, + "bridge children still running after {timeout:?}: {:?}", + bridge_pids() + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } +} + +/// Assistant rows before the user row that reappear after it. +fn echoed_rows(rows: &[Row]) -> Vec<&Row> { + let Some(user) = rows.iter().position(Row::is_user) else { + return Vec::new(); + }; + rows[..user] + .iter() + .filter(|row| matches!(row, Row::Assistant(_)) && rows[user + 1..].contains(row)) + .collect() +} + fn activity(ev: &Ev) -> bool { matches!( ev, @@ -471,6 +675,7 @@ async fn real_bridge_steer_mid_turn_redirects_the_live_turn() { let mut session = Session::connect("steer-live", false).await; let trace = Arc::clone(&session.trace); let handle = session.handle.clone(); + let cancel = session.cancel.clone(); let prompt_done = Rc::new(Cell::new(false)); let prompt_done_for_steer = Rc::clone(&prompt_done); let steer = tokio::task::spawn_local(async move { @@ -489,6 +694,10 @@ async fn real_bridge_steer_mid_turn_redirects_the_live_turn() { sent_at.as_secs_f64(), answered_at.as_secs_f64() ); + if outcome.is_err() { + // Nothing more to learn from the long turn. + cancel.cancel(); + } (turn_still_live, answered_at, outcome) }); @@ -499,6 +708,19 @@ async fn real_bridge_steer_mid_turn_redirects_the_live_turn() { let resolved_at = session.trace.elapsed(); eprintln!("[steer-live] prompt() resolved {turn:?} at {:.2}s", resolved_at.as_secs_f64()); let (turn_still_live, answered_at, outcome) = steer.await.expect("steer task"); + if !bridge_advertises_steering() { + eprintln!( + "SCENARIO 1: N/A — {} does not advertise steering; the steer was refused \ + locally ({outcome:?}) and nothing was written", + provider_id() + ); + assert!(outcome + .as_ref() + .is_err_and(|e| e.contains("does not advertise steering"))); + assert_eq!(user_rows(&session.trace), 0); + session.finish().await; + return; + } assert!(turn_still_live, "the steer must land inside the live turn, not after it"); assert_eq!(outcome, Ok(SteerOutcome::Injected)); @@ -556,6 +778,13 @@ async fn real_bridge_steer_mid_turn_redirects_the_live_turn() { #[tokio::test(flavor = "current_thread")] #[ignore = "spawns the installed ACP bridge and calls the model"] async fn real_bridge_steer_while_the_bridge_holds_the_turn_for_a_subagent() { + if !bridge_advertises_steering() { + eprintln!( + "SCENARIO 2: N/A — {} does not advertise steering (and has no background subagent tool)", + provider_id() + ); + return; + } tokio::task::LocalSet::new() .run_until(async { let mut session = Session::connect("steer-held", false).await; @@ -619,7 +848,8 @@ async fn real_bridge_steer_while_the_bridge_holds_the_turn_for_a_subagent() { } // ----------------------------------------------------------------------------- -// Scenario 3: interrupt-and-reprompt (steering disabled on the driver) +// Scenario 3: interrupt-and-reprompt (steering unavailable: not advertised by +// the bridge, or disabled on the driver for one that does) // ----------------------------------------------------------------------------- #[tokio::test(flavor = "current_thread")] @@ -627,8 +857,9 @@ async fn real_bridge_steer_while_the_bridge_holds_the_turn_for_a_subagent() { async fn real_bridge_interrupt_reprompts_on_the_same_child() { tokio::task::LocalSet::new() .run_until(async { - let mut session = Session::connect("interrupt", true).await; - let pids_before = session.bridge_pids(); + let advertises = bridge_advertises_steering(); + let mut session = Session::connect("interrupt", advertises).await; + let pids_before = bridge_pids(); assert_eq!(pids_before.len(), 1, "one bridge child: {pids_before:?}"); let trace = Arc::clone(&session.trace); let handle = session.handle.clone(); @@ -651,15 +882,21 @@ async fn real_bridge_interrupt_reprompts_on_the_same_child() { let sent_at = trace.elapsed(); let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[]).await; let answered_at = trace.elapsed(); + // The re-prompt is live now: the child it runs on. (Read + // here, not after the turn — a bridge with no observable + // background work gets no post-turn hold, so the child is + // gone the moment prompt() resolves.) + let pids_during_reprompt = bridge_pids(); // No terminal state mid-way: the interrupted prompt() is // still pending (and so, under the hold, is teardown). let prompt_still_pending = !prompt_done_for_interrupt.get(); eprintln!( - "[interrupt] interrupt sent at {:.2}s, answered {outcome:?} at {:.2}s (prompt still pending: {prompt_still_pending})", + "[interrupt] interrupt sent at {:.2}s, answered {outcome:?} at {:.2}s — cancel acked after {:.3}s (prompt still pending: {prompt_still_pending}, child: {pids_during_reprompt:?})", sent_at.as_secs_f64(), - answered_at.as_secs_f64() + answered_at.as_secs_f64(), + (answered_at - sent_at).as_secs_f64() ); - (prompt_still_pending, outcome) + (prompt_still_pending, outcome, pids_during_reprompt) }); let turn = tokio::time::timeout(TURN_TIMEOUT, session.connection.prompt(COUNT_PROMPT, &[])) @@ -668,16 +905,24 @@ async fn real_bridge_interrupt_reprompts_on_the_same_child() { prompt_done.set(true); let resolved_at = session.trace.elapsed(); eprintln!("[interrupt] prompt() resolved {turn:?} at {:.2}s", resolved_at.as_secs_f64()); - let (prompt_still_pending, outcome) = interrupt.await.expect("interrupt task"); - let pids_after = session.bridge_pids(); + let (prompt_still_pending, outcome, pids_during_reprompt) = + interrupt.await.expect("interrupt task"); + // What initialize said, before any forcing (logged by the + // connection task during setup, so read after the turn). + let advertised = format!("steering advertised: {advertises}"); + assert!(logged(&advertised), "driver did not log {advertised:?}"); + eprintln!("[interrupt] {advertised} (forced off on the driver: {advertises})"); assert_eq!(outcome, Ok(())); assert!( prompt_still_pending, "the interrupted prompt() resolves with the final turn, not at the cancel" ); assert_eq!(turn, Ok(AgentRunOutcome::Completed)); - assert_eq!(pids_after, pids_before, "the re-prompt runs on the same child"); + assert_eq!( + pids_during_reprompt, pids_before, + "the re-prompt runs on the same child" + ); assert_eq!(session.store.agent_session_ids.lock().unwrap().len(), 1); let rows = session.trace.rows(); let users: Vec = rows @@ -749,10 +994,12 @@ async fn race_steer(label: &str, steer_after_first_chunk: bool) { trace.elapsed().as_secs_f64() ); let delivered = match &outcome { - Ok(_) => true, + Ok(_) => Some(true), Err(e) => { // The caller's fallback: the connection is holding, so this - // is served as a plain new turn. + // is served as a plain new turn — or, on a hold-less bridge + // whose connection tore down at turn end, fails fast for + // the caller to route as a resume. let fallback = tokio::time::timeout( REPLY_TIMEOUT, handle.interrupt_and_reprompt(RACE_MESSAGE, &[]), @@ -760,7 +1007,11 @@ async fn race_steer(label: &str, steer_after_first_chunk: bool) { .await .expect("the fallback must not hang"); eprintln!("[{label_owned}] steer failed ({e}); fallback answered {fallback:?}"); - fallback.is_ok() + match &fallback { + Ok(()) => Some(true), + Err(e) if e.contains("no longer running") => None, + Err(_) => Some(false), + } } }; (outcome, delivered) @@ -775,6 +1026,15 @@ async fn race_steer(label: &str, steer_after_first_chunk: bool) { ); let (outcome, delivered) = steer.await.expect("steer task"); assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + let Some(delivered) = delivered else { + eprintln!( + "SCENARIO 4 ({label}): the turn ended and the connection tore down before the \ + message arrived; the caller routes that as a resume" + ); + assert_eq!(user_rows(&session.trace), 0); + session.finish().await; + return; + }; assert!( delivered, "neither the steer nor the fallback delivered: {outcome:?}" @@ -796,16 +1056,10 @@ async fn race_steer(label: &str, steer_after_first_chunk: bool) { // The aborted cycle's result must not re-echo the pre-steer message as a // row after the user row (claude-agent-acp emits that result; the client // has to recognise its text as already streamed). - let user = rows.iter().position(Row::is_user).unwrap(); - for before in rows[..user] - .iter() - .filter(|row| matches!(row, Row::Assistant(_))) - { - assert!( - !rows[user + 1..].contains(before), - "pre-steer assistant row echoed after the user row: {rows:#?}" - ); - } + assert!( + echoed_rows(&rows).is_empty(), + "pre-steer assistant row echoed after the user row: {rows:#?}" + ); session.finish().await; } @@ -824,3 +1078,301 @@ async fn real_bridge_steer_sent_at_the_first_chunk_lands_once() { .run_until(race_steer("race-at-chunk", true)) .await; } + +// ----------------------------------------------------------------------------- +// Scenario 5: interrupt while the bridge waits on a permission decision +// ----------------------------------------------------------------------------- + +/// The first permission request is held open (a dialog nobody has answered) +/// when the interrupt arrives. The cancel must still be acknowledged, the +/// held request resolved rather than left dangling, and the next turn's own +/// permission request must reach the handler. Skipped when the bridge asks +/// for nothing: `GOOSE_MODE=approve` makes Goose ask for every tool call; +/// the managed bridges ignore the variable and approve their shell calls +/// themselves. +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_interrupt_while_a_permission_request_is_pending() { + tokio::task::LocalSet::new() + .run_until(async { + let mut session = Session::connect_with( + "interrupt-perm", + ConnectOptions { + steering_disabled: bridge_advertises_steering(), + env: vec![("GOOSE_MODE", "approve")], + ..ConnectOptions::default() + }, + ) + .await; + session + .trace + .hold_next_permission + .store(true, Ordering::SeqCst); + let pids_before = bridge_pids(); + let trace = Arc::clone(&session.trace); + let handle = session.handle.clone(); + let interrupt = tokio::task::spawn_local(async move { + let deadline = Instant::now() + TURN_START_TIMEOUT; + loop { + if trace.count(|ev| matches!(ev, Ev::Permission { .. })) > 0 { + break; + } + // A tool call that completed without asking: this bridge + // does not request permission for shell calls. + if trace.count(|ev| matches!(ev, Ev::ToolResult { .. })) > 0 + || Instant::now() > deadline + { + return None; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + tokio::time::sleep(Duration::from_secs(1)).await; + let sent_at = trace.elapsed(); + let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[]).await; + let answered_at = trace.elapsed(); + let pids_during_reprompt = bridge_pids(); + eprintln!( + "[interrupt-perm] interrupt sent at {:.2}s with a permission request pending, answered {outcome:?} at {:.2}s — cancel acked after {:.3}s (child: {pids_during_reprompt:?})", + sent_at.as_secs_f64(), + answered_at.as_secs_f64(), + (answered_at - sent_at).as_secs_f64() + ); + Some((sent_at, answered_at, outcome, pids_during_reprompt)) + }); + + let turn = tokio::time::timeout(TURN_TIMEOUT, session.connection.prompt(COUNT_PROMPT, &[])) + .await + .expect("the turn should end"); + eprintln!( + "[interrupt-perm] prompt() resolved {turn:?} at {:.2}s", + session.trace.elapsed().as_secs_f64() + ); + let Some((sent_at, answered_at, outcome, pids_during_reprompt)) = + interrupt.await.expect("interrupt task") + else { + eprintln!( + "SCENARIO 5: SKIPPED — the bridge issued no session/request_permission for its shell calls" + ); + session.finish().await; + return; + }; + + let events = session.trace.events(); + let user_at = events + .iter() + .find(|(_, ev)| matches!(ev, Ev::User(_))) + .map(|(at, _)| *at); + let cancelled_at = events + .iter() + .find(|(_, ev)| matches!(ev, Ev::PermissionCancelled { .. })) + .map(|(at, _)| *at); + let hold_ended_otherwise = events.iter().any(|(_, ev)| { + matches!(ev, Ev::Event(kind) if kind.starts_with("permission-hold-")) + }); + eprintln!( + "SCENARIO 5: cancel acked after {:.3}s; held permission cancelled at {:?} (hold ended another way: {hold_ended_otherwise}); user row at {:?}", + (answered_at - sent_at).as_secs_f64(), + cancelled_at.map(|at| at.as_secs_f64()), + user_at.map(|at| at.as_secs_f64()) + ); + assert_eq!( + outcome, + Ok(()), + "the interrupt must be acknowledged while a permission request is pending" + ); + assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + assert_eq!(pids_during_reprompt, pids_before, "same child"); + let rows = session.trace.rows(); + assert_eq!(user_rows(&session.trace), 1, "exactly one user row: {rows:#?}"); + let user_at = user_at.unwrap(); + // The held request was resolved by the interrupt, before the + // next turn, not left for someone to answer later. + assert!( + cancelled_at.is_some_and(|at| at <= user_at), + "the pending permission request was not resolved by the interrupt: {events:#?}" + ); + let after = session.trace.assistant_text_after_user(); + assert!(after.contains("INTERRUPTED-DONE"), "text after the user row: {after:?}"); + // Turn 2's own request (the `pwd`) reached the handler and was + // answered: the turn could not have completed otherwise. + let permissions_after = events + .iter() + .filter(|(at, ev)| *at > user_at && matches!(ev, Ev::Permission { .. })) + .count(); + eprintln!("[interrupt-perm] permission requests in the second turn: {permissions_after}"); + assert!(permissions_after > 0, "turn 2's permission request never reached the handler"); + session.finish().await; + }) + .await; +} + +// ----------------------------------------------------------------------------- +// Scenario 6: interrupt racing the end of the turn +// ----------------------------------------------------------------------------- + +/// Send an interrupt `when` the short turn is at (or past) its end. Whatever +/// the timing, nothing hangs and the message lands at most once: as an +/// interrupt of the live turn, as a new turn on the held connection, or — +/// once a hold-less connection has torn down — as a fast failure the caller +/// routes to a resume. +async fn race_interrupt(label: &str, after_first_chunk: bool) { + let mut session = Session::connect(label, bridge_advertises_steering()).await; + let trace = Arc::clone(&session.trace); + let handle = session.handle.clone(); + let label_owned = label.to_string(); + let interrupt = tokio::task::spawn_local(async move { + if after_first_chunk { + trace + .wait_for("the first chunk", TURN_START_TIMEOUT, |evs| { + evs.iter().any(|ev| matches!(ev, Ev::Append(_))) + }) + .await; + } + let sent_at = trace.elapsed(); + let outcome = tokio::time::timeout( + REPLY_TIMEOUT, + handle.interrupt_and_reprompt(RACE_MESSAGE, &[]), + ) + .await + .expect("an interrupt at the turn boundary must not hang"); + eprintln!( + "[{label_owned}] interrupt sent at {:.2}s answered {outcome:?} at {:.2}s", + sent_at.as_secs_f64(), + trace.elapsed().as_secs_f64() + ); + outcome + }); + + let turn = tokio::time::timeout(TURN_TIMEOUT, session.connection.prompt(RACE_PROMPT, &[])) + .await + .expect("the short turn should end"); + eprintln!( + "[{label}] prompt() resolved {turn:?} at {:.2}s", + session.trace.elapsed().as_secs_f64() + ); + let outcome = interrupt.await.expect("interrupt task"); + assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + match &outcome { + Ok(()) => { + session + .trace + .wait_for("the reply to the interrupt", REPLY_TIMEOUT, |evs| { + streamed_text(evs).contains("RACE-DONE") + }) + .await; + // Let any trailing frames land before counting rows. + tokio::time::sleep(Duration::from_secs(3)).await; + let rows = session.trace.rows(); + eprintln!( + "SCENARIO 6 ({label}): delivered; user rows {}", + user_rows(&session.trace) + ); + assert_eq!(user_rows(&session.trace), 1, "{rows:#?}"); + assert!( + echoed_rows(&rows).is_empty(), + "pre-interrupt assistant row echoed after the user row: {rows:#?}" + ); + } + Err(e) if e.contains("no longer running") => { + eprintln!( + "SCENARIO 6 ({label}): the turn ended and the connection tore down before the \ + interrupt arrived ({e}); the caller routes that as a resume" + ); + assert_eq!(user_rows(&session.trace), 0); + } + other => panic!("unexpected interrupt outcome at the turn boundary: {other:?}"), + } + session.finish().await; +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_interrupt_sent_with_the_prompt_lands_once() { + tokio::task::LocalSet::new() + .run_until(race_interrupt("irace-at-send", false)) + .await; +} + +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_interrupt_sent_at_the_first_chunk_lands_once() { + tokio::task::LocalSet::new() + .run_until(race_interrupt("irace-at-chunk", true)) + .await; +} + +// ----------------------------------------------------------------------------- +// Scenario 7: resume a settled session with session/load +// ----------------------------------------------------------------------------- + +/// One connection teaches the agent a code word and settles; a second, +/// on a new child, loads the agent's session in the same dir and asks for +/// it back. The replayed history must be matched or dropped, never recorded +/// as fresh rows. +#[tokio::test(flavor = "current_thread")] +#[ignore = "spawns the installed ACP bridge and calls the model"] +async fn real_bridge_resume_loads_the_session_on_a_new_child() { + tokio::task::LocalSet::new() + .run_until(async { + let mut first = Session::connect("resume-1", bridge_advertises_steering()).await; + let pids_first = bridge_pids(); + let turn = tokio::time::timeout(TURN_TIMEOUT, first.connection.prompt(REMEMBER_PROMPT, &[])) + .await + .expect("the first turn should end"); + assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + let first_text = streamed_text( + &first.trace.events().into_iter().map(|(_, ev)| ev).collect::>(), + ); + assert!(first_text.contains("REMEMBERED"), "first turn: {first_text:?}"); + let agent_session_id = first.agent_session_id(); + let messages = first.legacy_messages(REMEMBER_PROMPT); + let dir = first.dir.clone(); + first.finish_keeping_dir().await; + wait_for_bridge_exit(SETTLE_WAIT).await; + + let mut second = Session::connect_with( + "resume-2", + ConnectOptions { + steering_disabled: bridge_advertises_steering(), + resume: Some(Resume { + agent_session_id: agent_session_id.clone(), + messages, + dir: dir.clone(), + }), + ..ConnectOptions::default() + }, + ) + .await; + let pids_second = bridge_pids(); + assert_eq!(pids_second.len(), 1, "one new child: {pids_second:?}"); + assert_ne!(pids_second, pids_first, "the resume runs on a new child"); + let turn = tokio::time::timeout(TURN_TIMEOUT, second.connection.prompt(RECALL_PROMPT, &[])) + .await + .expect("the resumed turn should end"); + eprintln!( + "[resume-2] prompt() resolved {turn:?} at {:.2}s", + second.trace.elapsed().as_secs_f64() + ); + assert_eq!(turn, Ok(AgentRunOutcome::Completed)); + let rows = second.trace.rows(); + let text = streamed_text( + &second.trace.events().into_iter().map(|(_, ev)| ev).collect::>(), + ); + eprintln!( + "SCENARIO 7: loaded {agent_session_id} on {pids_second:?}; store ids {:?}; reply {text:?}", + second.store.agent_session_ids.lock().unwrap() + ); + assert_eq!(user_rows(&second.trace), 0, "{rows:#?}"); + assert!( + !rows.iter().any(|row| row.text().contains("REMEMBERED")), + "replayed history recorded as new rows: {rows:#?}" + ); + assert!( + text.contains("PELICAN-7342"), + "the loaded session should remember the code word: {text:?}" + ); + second.finish().await; + let _ = std::fs::remove_dir_all(&dir); + }) + .await; +} diff --git a/crates/acp-client/src/driver/tests/turn_control.rs b/crates/acp-client/src/driver/tests/turn_control.rs index 31706e297..c4c1e5c93 100644 --- a/crates/acp-client/src/driver/tests/turn_control.rs +++ b/crates/acp-client/src/driver/tests/turn_control.rs @@ -36,9 +36,12 @@ struct FakeAgent { advertise_steering: bool, /// Text the agent streams at the start of each prompt, in order. stream_on_prompt: Vec<&'static str>, - /// Ask permission for a tool call during the Nth (1-based) prompt, and - /// record the client's answer. - permission_on_prompt: Option, + /// Ask permission for a tool call during these (1-based) prompts, and + /// record the client's answers in order. + permission_on_prompts: Vec, + /// Text streamed right behind a permission request, while the client + /// has not answered it. + stream_after_permission: Option<&'static str>, permission_outcomes: Mutex>, permission_answered: tokio::sync::Notify, } @@ -125,7 +128,7 @@ impl FakeAgent { if let Some(chunk) = prompt_agent.stream_on_prompt.get(ordinal - 1) { cx.send_notification(agent_text(Some(&format!("msg-{ordinal}")), chunk))?; } - if prompt_agent.permission_on_prompt == Some(ordinal) { + if prompt_agent.permission_on_prompts.contains(&ordinal) { let tool_call_id = format!("tc-{ordinal}"); cx.send_notification(tool_call_notification(&tool_call_id))?; let agent = Arc::clone(&prompt_agent); @@ -140,6 +143,12 @@ impl FakeAgent { agent.permission_answered.notify_one(); Ok(()) })?; + if let Some(chunk) = prompt_agent.stream_after_permission { + cx.send_notification(agent_text( + Some(&format!("msg-{ordinal}")), + chunk, + ))?; + } } *prompt_agent.pending_prompt.lock().unwrap() = Some(responder); prompt_agent.prompt_arrived.notify_one(); @@ -605,9 +614,13 @@ where let result = agent_client_protocol::Client .builder() .on_receive_request( - async move |args: RequestPermissionRequest, responder, _connection| { - let response = permission_handler.request_permission(args).await?; - responder.respond(response) + // Off the dispatch task, as `connect` registers it. + async move |args: RequestPermissionRequest, responder, connection| { + let handler = Arc::clone(&permission_handler); + connection.spawn(async move { + let response = handler.request_permission(args).await?; + responder.respond(response) + }) }, agent_client_protocol::on_receive_request!(), ) @@ -663,7 +676,7 @@ async fn interrupt_cancels_the_turn_and_reprompts_on_the_same_connection() { stream_on_prompt: vec!["half an answer", "the new answer"], // Turn 2 asks permission: the connection-wide permission token must // have survived the interrupt for the handler to answer it. - permission_on_prompt: Some(2), + permission_on_prompts: vec![2], ..FakeAgent::default() }); @@ -1045,3 +1058,157 @@ async fn a_turn_control_handle_maps_channel_failures_to_connection_state() { .expect_err("no connection to serve the message"); assert!(error.contains("no longer running"), "got: {error}"); } + +/// The interrupt arrives while the turn's permission request is still +/// unanswered — a dialog nobody has clicked. The interrupt resolves it +/// `Cancelled` (ACP's rule for a client that cancels the turn), so the cancel +/// is acknowledged at once instead of waiting on the dialog, and the next +/// turn's own request is still presented and answered: the token cancelled +/// was the turn's, not the connection's. Seen against Goose, where the +/// pending request left the cancel unacknowledged for the full 30 s and the +/// message was dropped. +#[tokio::test(flavor = "current_thread")] +async fn interrupt_resolves_the_turns_pending_permission_request_before_reprompting() { + let agent = Arc::new(FakeAgent { + advertise_steering: false, + stream_on_prompt: vec!["half an answer", "the new answer"], + permission_on_prompts: vec![1, 2], + ..FakeAgent::default() + }); + + let (result, _reply_rx, stashed_reply, writer, handler) = + session_with(&agent, None, |handle, agent, writer, _handler| async move { + writer.hold_next_permission.store(true, Ordering::SeqCst); + agent.await_prompt().await; + wait_until("turn 1's permission request to reach the writer", || { + writer + .calls() + .contains(&WriterCall::Permission("tc-1".to_string())) + }) + .await; + assert!(agent.permission_outcomes.lock().unwrap().is_empty()); + + let started = tokio::time::Instant::now(); + handle + .interrupt_and_reprompt("actually, do this instead", &[]) + .await + .expect("the cancel is acknowledged with the permission request pending"); + assert!( + started.elapsed() < Duration::from_secs(5), + "the interrupt waited on the dialog" + ); + agent.await_prompt().await; + assert_eq!(agent.prompts(), ["first", "actually, do this instead"]); + assert_eq!(agent.cancels(), 1); + + wait_until("both permission requests to be answered", || { + agent.permission_outcomes.lock().unwrap().len() == 2 + }) + .await; + wait_until("the second chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("the new answer".to_string())) + }) + .await; + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + assert_eq!( + result, + Ok((AgentRunOutcome::Completed, SessionSettleReason::Immediate)) + ); + let outcomes = agent.permission_outcomes.lock().unwrap().clone(); + assert!( + matches!(outcomes[0], RequestPermissionOutcome::Cancelled), + "turn 1's request: {:?}", + outcomes[0] + ); + assert!( + matches!(outcomes[1], RequestPermissionOutcome::Selected(_)), + "turn 2's request: {:?}", + outcomes[1] + ); + let calls = writer.calls(); + assert!( + calls.contains(&WriterCall::PermissionCancelled("tc-1".to_string())), + "{calls:#?}" + ); + assert!( + calls.contains(&WriterCall::Permission("tc-2".to_string())), + "{calls:#?}" + ); + assert_eq!( + writer.text_flow(), + vec![ + WriterCall::Append("half an answer".to_string()), + WriterCall::Finalize, + WriterCall::UserMessage("actually, do this instead".to_string()), + WriterCall::Append("the new answer".to_string()), + ] + ); + assert!(!handler.permission_cancel_token.is_cancelled()); + if let Some(reply) = stashed_reply { + let _ = reply.send(result.map(|(outcome, _)| outcome)); + } +} + +/// A permission decision nobody has made yet does not stall the connection: +/// the frames behind the request keep being dispatched while the dialog is +/// up. (The SDK runs handlers on one task; the decision is awaited off it.) +#[tokio::test(flavor = "current_thread")] +async fn a_pending_permission_decision_does_not_stall_the_connection() { + let agent = Arc::new(FakeAgent { + advertise_steering: false, + stream_on_prompt: vec!["half an answer"], + permission_on_prompts: vec![1], + stream_after_permission: Some("still streaming"), + ..FakeAgent::default() + }); + + let (result, _reply_rx, stashed_reply, _writer, _handler) = session_with( + &agent, + None, + |_handle, agent, writer, _handler| async move { + writer.hold_next_permission.store(true, Ordering::SeqCst); + agent.await_prompt().await; + wait_until( + "the chunk sent behind the pending permission request", + || { + writer + .text_flow() + .contains(&WriterCall::Append("still streaming".to_string())) + }, + ) + .await; + assert!(writer + .calls() + .contains(&WriterCall::Permission("tc-1".to_string()))); + assert!( + agent.permission_outcomes.lock().unwrap().is_empty(), + "the decision is still pending" + ); + writer.release_permission.notify_one(); + wait_until("the released request to be answered", || { + agent.permission_outcomes.lock().unwrap().len() == 1 + }) + .await; + agent.finish_prompt(StopReason::EndTurn); + }, + ) + .await; + + assert_eq!( + result, + Ok((AgentRunOutcome::Completed, SessionSettleReason::Immediate)) + ); + let outcomes = agent.permission_outcomes.lock().unwrap().clone(); + assert!( + matches!(outcomes[0], RequestPermissionOutcome::Selected(_)), + "{outcomes:?}" + ); + if let Some(reply) = stashed_reply { + let _ = reply.send(result.map(|(outcome, _)| outcome)); + } +} From 02187bd8330185127d41779e77cc3eed5c08059e Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Thu, 29 Oct 2026 15:44:20 +1100 Subject: [PATCH 6/7] fix(staged,acp-client): address review of message-to-running-session MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses every issue, warning and suggestion of review 480243de-4c34-4340-bff4-d73c5244dc3b on the continue-session work, each behavioural fix with a regression test. Issue — resume-gate race in continue_session_for_store: the schedule check and the `running` CAS (transition_to_running_clearing_agent_id) now run together under branch_session_launch_lock_for(branch_id) with no await between them (continue_resume_gate); the rest of the resume runs outside the lock on the claimed session, and resume_session_for_store keeps its own CAS while the continue path no longer repeats it. Test: two stopped exclusive sessions on one branch continued concurrently → exactly one Resumed, the other Blocked. Warning — stop during the interrupt ack wait: the INTERRUPT_ACK_TIMEOUT wait selects on cancel_token, failing the interrupt request, cancelling pending permissions and settling Cancelled instead of re-prompting; a queued turn is also refused when the token is already cancelled at the top of the run loop. Tests: stop_during_the_interrupt_ack_wait_… and stop_landing_with_the_cancel_ack_…. Warning — response_completion relied on the aborted cycle's result arriving before the steered reply's first chunk: detach_message now moves the closed text into a `detached` list that a late result is matched against, so the reply row is never split and the final result still completes normally. Test: aborted_cycle_result_arriving_after_the_steered_reply_began_…. Warning — Running with no registry entry owned by another instance: route_continue_session compares session.owner_pid with the current pid first and returns Blocked immediately; the startup wait, the remote branch startup wait and the post-interrupt settle wait are separate budgets (ContinueWaits: 10 s / 30 s / 10 s). Tests: owner-pid block without waiting, remote budget, settle budget after a dead connection. Warning — image ids dropped on the live path: attachment ids travel through TurnControlRequest/QueuedSessionTurn to MessageWriter::record_user_message, and the Staged writer persists them with the user row and re-associates the images. Tests in acp-client (attachment_ids_reach_the_writer_…), session_commands (image_ids_travel_with_the_message_on_the_live_path) and agent::writer. Suggestion — load_prompt_attachments on the continue body runs under tokio::task::spawn_blocking. Suggestion — steer() is bounded by STEER_ACK_TIMEOUT (30 s, matching INTERRUPT_ACK_TIMEOUT); the reply sender sits in a slot the callback and the timer race for, so a late `injected` answer persists nothing. Test: steer_unanswered_within_the_timeout_fails_and_a_late_answer_persists_nothing. Suggestion — a failed send_session_cancel on the interrupt path fails only that interrupt request; the running turn continues. Suggestion — the live project_mcp test asserts the user row index against the row count captured right before continue_with instead of rows_before.min(1). Also: the continue_repo_session tool description explains that the activity's last user-message time reflects any user message the agent received, so the caller should compare it with its own call time. Verified: cargo build, acp-client tests, staged lib tests for the touched modules, clippy (-D warnings) and fmt; Claude and Goose real-bridge interrupt/steer scenarios and the live continue_repo_session MCP tests re-run. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/agent/writer.rs | 96 ++- apps/staged/src-tauri/src/project_mcp.rs | 21 +- apps/staged/src-tauri/src/session_commands.rs | 712 +++++++++++++++--- crates/acp-client/src/driver.rs | 290 ++++++- .../src/driver/response_completion.rs | 66 +- .../src/driver/tests/real_bridge.rs | 23 +- .../src/driver/tests/response_completion.rs | 63 +- .../src/driver/tests/turn_control.rs | 363 ++++++++- 8 files changed, 1448 insertions(+), 186 deletions(-) diff --git a/apps/staged/src-tauri/src/agent/writer.rs b/apps/staged/src-tauri/src/agent/writer.rs index 0358a1309..1cf3d5c6d 100644 --- a/apps/staged/src-tauri/src/agent/writer.rs +++ b/apps/staged/src-tauri/src/agent/writer.rs @@ -178,26 +178,48 @@ impl MessageWriter { /// The `images` are the blocks the connection sent, not stored images — /// acp-client has no image ids — so the row carries no `image_ids`. A /// caller that owns stored images attaches them to the row it gets back. - pub async fn record_user_message(&self, text: &str, images: &[PromptImage]) -> Option { + pub async fn record_user_message( + &self, + text: &str, + images: &[PromptImage], + image_ids: &[String], + ) -> Option { self.finalize().await; - if !images.is_empty() { + if images.len() != image_ids.len() { log::debug!( - "Session {}: user message delivered mid-turn with {} image block(s); the \ - persisted row records the text only", + "Session {}: user message delivered mid-turn with {} image block(s) and {} \ + stored attachment id(s)", self.session_id, - images.len() + images.len(), + image_ids.len() ); } - match self - .store - .add_session_message(&self.session_id, MessageRole::User, text) - { - Ok(id) => Some(id), + let id = match self.store.add_session_message_with_images( + &self.session_id, + MessageRole::User, + text, + image_ids, + ) { + Ok(id) => id, Err(e) => { log::error!("Failed to insert delivered user message: {e}"); - None + return None; + } + }; + // Same association `start_session` makes for the prompt's images, so + // the attachments show under this session in the branch timeline. + if !image_ids.is_empty() { + if let Err(e) = self + .store + .set_images_session_id(image_ids, &self.session_id) + { + log::warn!( + "Failed to associate images {image_ids:?} with session {}: {e}", + self.session_id + ); } } + Some(id) } // ===================================================================== @@ -346,8 +368,13 @@ impl acp_client::MessageWriter for MessageWriter { self.finalize().await } - async fn record_user_message(&self, text: &str, images: &[PromptImage]) { - self.record_user_message(text, images).await; + async fn record_user_message( + &self, + text: &str, + images: &[PromptImage], + attachment_ids: &[String], + ) { + self.record_user_message(text, images, attachment_ids).await; } async fn should_finish_after_response(&self) -> bool { @@ -622,7 +649,7 @@ mod tests { writer.append_text("instruction.").await; writer.finalize().await; let user_row = writer - .record_user_message("Actually, do the second thing instead.", &[]) + .record_user_message("Actually, do the second thing instead.", &[], &[]) .await .expect("user row inserted"); writer.append_text("Switching to the second thing.").await; @@ -656,7 +683,8 @@ mod tests { writer.append_text("Unflushed partial text").await; let image = acp_client::PromptImage::new("AAAA", "image/png"); - acp_client::MessageWriter::record_user_message(&writer, "Look at this.", &[image]).await; + acp_client::MessageWriter::record_user_message(&writer, "Look at this.", &[image], &[]) + .await; writer.append_text("Reply.").await; writer.finalize().await; @@ -675,6 +703,44 @@ mod tests { ); } + /// The attachment ids the connection hands back land on the user row as + /// its `image_ids`, and the stored images are associated with the session + /// — what `start_session` does for a prompt's images, so a message with + /// attachments reads the same whether it was delivered live or on a + /// resume. + #[tokio::test] + async fn record_user_message_persists_attachment_ids_and_associates_the_images() { + use crate::store::{Image, Project}; + let (store, session_id, writer) = setup_writer(); + let project = Project::new("owner/repo"); + store.create_project(&project).unwrap(); + let image = Image::new(None, &project.id, "shot.png", "image/png", 3, true); + store.create_image(&image).unwrap(); + assert_ne!( + store.get_image(&image.id).unwrap().unwrap().session_id, + Some(session_id.clone()) + ); + + let user_row = writer + .record_user_message( + "See attached.", + &[acp_client::PromptImage::new("AAAA", "image/png")], + std::slice::from_ref(&image.id), + ) + .await + .expect("user row inserted"); + + let messages = store.get_session_messages(&session_id).unwrap(); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].id, user_row); + assert_eq!(messages[0].role, MessageRole::User); + assert_eq!(messages[0].image_ids, vec![image.id.clone()]); + assert_eq!( + store.get_image(&image.id).unwrap().unwrap().session_id, + Some(session_id) + ); + } + #[tokio::test] async fn ordinary_session_does_not_finish_for_note_shaped_output() { use acp_client::MessageWriter as _; diff --git a/apps/staged/src-tauri/src/project_mcp.rs b/apps/staged/src-tauri/src/project_mcp.rs index 15b2abd1c..bd91624d1 100644 --- a/apps/staged/src-tauri/src/project_mcp.rs +++ b/apps/staged/src-tauri/src/project_mcp.rs @@ -147,8 +147,10 @@ const REPO_SESSION_ACTIVITY_PREVIEW_MAX_CHARS: usize = 240; struct RepoSessionActivity { last_activity_at: Option, /// When the session's agent last received a user message — the original - /// instructions or a `continue_repo_session` follow-up — so a parent that - /// continued the session can confirm the message landed in its transcript. + /// instructions or a `continue_repo_session` follow-up, whichever is + /// latest — so a parent that continued the session can confirm the + /// message landed in its transcript by comparing this with the time of + /// its own call (the tool description says so). last_user_message_at: Option, last_message: Option, last_tool_call: Option, @@ -904,7 +906,7 @@ impl ProjectToolsHandler { } #[tool( - description = "Send a follow-up message to a repo session started by `start_repo_session`: a correction, the answer to a question it raised, or its next step. A running session receives the message within seconds without being stopped, so its work so far is kept; a completed, cancelled, or failed session is resumed with the message as its next instruction. A session that has not started yet cannot take a message — wait for it to run, or cancel it and start a new one. Prefer this over cancelling and starting over whenever the session's work so far is still wanted. Returns the session's current state together with how the message was delivered, or why it was not." + description = "Send a follow-up message to a repo session started by `start_repo_session`: a correction, the answer to a question it raised, or its next step. A running session receives the message within seconds without being stopped, so its work so far is kept; a completed, cancelled, or failed session is resumed with the message as its next instruction. A session that has not started yet cannot take a message — wait for it to run, or cancel it and start a new one. Prefer this over cancelling and starting over whenever the session's work so far is still wanted. Returns the session's current state together with how the message was delivered, or why it was not. The activity's last user-message time is when the session's agent most recently received any user message — the original instructions included, not only this one — so to confirm this message landed, compare it with the time you made this call." )] async fn continue_repo_session( &self, @@ -1627,7 +1629,9 @@ mod tests { assert!(description.contains("resumed with the message")); assert!(description.contains("has not started yet cannot take a message")); assert!(description.contains("Prefer this over cancelling")); + assert!(description.contains("compare it with the time you made this call")); assert!(!description.contains("`delivery`")); + assert!(!description.contains("last_user_message_at")); assert!(!description.contains("injected")); assert!(!description.contains("interrupt")); assert!(!description.contains("steer")); @@ -1688,7 +1692,9 @@ mod tests { &self.session.id, message, &[], - std::time::Duration::from_millis(200), + crate::session_commands::ContinueWaits::uniform(std::time::Duration::from_millis( + 200, + )), ) .await .unwrap(); @@ -1863,6 +1869,8 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(100)).await; } tokio::time::sleep(std::time::Duration::from_secs(2)).await; + // Captured right before the continue: every row streamed so + // far precedes the user row the connection persists. let rows_before = messages().len(); let payload = f.continue_with(FOLLOW_UP).await; @@ -1900,7 +1908,10 @@ mod tests { all.iter().filter(|m| m.role == MessageRole::User).count(), 1 ); - assert!(user_at >= rows_before.min(1), "{all:#?}"); + assert!( + user_at >= rows_before, + "user row at {user_at}, {rows_before} rows preceded the continue: {all:#?}" + ); let reply: String = all[user_at + 1..] .iter() .filter(|m| m.role == MessageRole::Assistant) diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index bd93c13e4..1575fb91a 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -1410,6 +1410,67 @@ pub(crate) async fn resume_session_for_store( .map_err(|e| e.to_string())? .ok_or_else(|| format!("Session not found: {session_id}"))?; + let plan = prepare_resume(&store, &session, branch_id, acp_config_selection).await?; + + // Live from before the `running` write until the runner has registered + // (see `SessionRegistry::hold`); released when this function returns. + let _liveness = registry.hold(&session_id); + let transitioned = store + .transition_to_running_clearing_agent_id( + &session_id, + plan.pending_start_agent_id.as_deref(), + ) + .map_err(|e| e.to_string())?; + if !transitioned { + return Err(SESSION_ALREADY_RUNNING_ERROR.to_string()); + } + + launch_resume( + plan, + store, + registry, + action_executor, + action_registry, + app_handle, + &session, + prompt, + image_ids, + queued_message_id, + ) + .await +} + +/// Everything a resume reads before the row is claimed: the session's +/// artifact and branch, the provider and config selection, the workspace +/// path and the pre-run HEAD. None of it needs the row to be `running`, so +/// `resume_session_for_store` reads it all first and the claim is a single +/// statement; the continue path, which claims under the branch launch lock +/// before it can read any of this, runs the same preparation afterwards. +struct ResumePlan { + provider: Option, + pending_start_agent_id: Option, + agent_session_id: Option, + working_dir: PathBuf, + effective_acp_config_selection: Option, + parent_project_note_id: Option, + mcp_project_id: Option, + linked_commit: bool, + linked_note: bool, + session_type: Option, + event_branch_id: Option, + event_project_id: Option, + pre_head_sha: Option, + workspace_name: Option, + remote_working_dir: Option, +} + +async fn prepare_resume( + store: &Arc, + session: &store::Session, + branch_id: Option, + acp_config_selection: Option, +) -> Result { + let session_id = session.id.clone(); // Use the provider that originally created this session so the // agent's conversation history can be restored correctly. let provider = session.provider.clone(); @@ -1418,8 +1479,7 @@ pub(crate) async fn resume_session_for_store( let agent_session_id = resumable_agent_session_id(session.agent_id.as_deref()); let working_dir = PathBuf::from(&session.working_dir); let effective_acp_config_selection = - resolve_resume_acp_config_selection(&store, &session, acp_config_selection); - let acp_config_selection_to_persist = effective_acp_config_selection.clone(); + resolve_resume_acp_config_selection(store, session, acp_config_selection); // Check if this session is linked to a project note — if so, we need // to start the MCP server so the agent has access to project tools. @@ -1485,7 +1545,7 @@ pub(crate) async fn resume_session_for_store( if let Some(ref branch) = linked_branch { let ws_name = branch.workspace_name.clone(); let head = if linked_commit.is_some() { - crate::branches::branch_head_sha(&store, branch, &working_dir) + crate::branches::branch_head_sha(store, branch, &working_dir) .await .ok() } else { @@ -1502,7 +1562,7 @@ pub(crate) async fn resume_session_for_store( let remote_working_dir = if let Some(ref branch) = linked_branch { if branch.workspace_name.is_some() { let ws_name = branch.workspace_name.as_deref().unwrap().to_string(); - let store_for_resolve = Arc::clone(&store); + let store_for_resolve = Arc::clone(store); let branch_for_resolve = branch.clone(); match tauri::async_runtime::spawn_blocking(move || { crate::branches::resolve_branch_workspace_subpath( @@ -1527,22 +1587,69 @@ pub(crate) async fn resume_session_for_store( None }; - // Live from before the `running` write until the runner has registered - // (see `SessionRegistry::hold`); released when this function returns. - let _liveness = registry.hold(&session_id); - let transitioned = store - .transition_to_running_clearing_agent_id(&session_id, pending_start_agent_id.as_deref()) - .map_err(|e| e.to_string())?; - if !transitioned { - return Err(SESSION_ALREADY_RUNNING_ERROR.to_string()); - } + Ok(ResumePlan { + provider, + pending_start_agent_id, + agent_session_id, + working_dir, + effective_acp_config_selection, + parent_project_note_id: project_note.as_ref().map(|note| note.id.clone()), + mcp_project_id, + linked_commit: linked_commit.is_some(), + linked_note: linked_note.is_some(), + session_type, + event_branch_id, + event_project_id, + pre_head_sha, + workspace_name, + remote_working_dir, + }) +} + +/// The second half of a resume, after the row has been claimed `running` +/// (and with the caller's liveness hold still in place): announce the +/// status, persist the config selection and hand the row to the runner. A +/// failure from here on is recorded on the row, which is already `running` +/// and announced. +#[allow(clippy::too_many_arguments)] +async fn launch_resume( + plan: ResumePlan, + store: Arc, + registry: Arc, + action_executor: Arc, + action_registry: Arc, + app_handle: tauri::AppHandle, + session: &store::Session, + prompt: String, + image_ids: Option>, + queued_message_id: Option, +) -> Result<(), String> { + let session_id = session.id.clone(); + let ResumePlan { + provider, + pending_start_agent_id: _, + agent_session_id, + working_dir, + effective_acp_config_selection, + parent_project_note_id, + mcp_project_id, + linked_commit, + linked_note, + session_type, + event_branch_id, + event_project_id, + pre_head_sha, + workspace_name, + remote_working_dir, + } = plan; + let acp_config_selection_to_persist = effective_acp_config_selection.clone(); let config_branch_id = event_branch_id.clone(); let config_project_id = event_project_id.clone().or(mcp_project_id.clone()); // Note follow-ups (project notes or local branch notes) get the pikchr // tool; remote branch notes can't reach localhost. let expose_pikchr_tools = local_note_pikchr_tools_available( - project_note.is_some() || linked_note.is_some(), + parent_project_note_id.is_some() || linked_note, workspace_name.as_deref(), ); @@ -1584,7 +1691,7 @@ pub(crate) async fn resume_session_for_store( pre_head_sha, provider, workspace_name, - extra_env: if linked_commit.is_some() { + extra_env: if linked_commit { session_runner::git_identity_env_from_global_config() } else { vec![] @@ -1609,7 +1716,7 @@ pub(crate) async fn resume_session_for_store( expose_pikchr_tools, // When resuming a project session, keep its parent project note in // scope so `child_note` repo sessions still attach to it. - parent_project_note_id: project_note.as_ref().map(|note| note.id.clone()), + parent_project_note_id, background_hold: default_background_hold(), }, Arc::clone(&store), @@ -1636,10 +1743,43 @@ pub(crate) async fn resume_session_for_store( /// delivery rather than an error. pub(crate) const SESSION_ALREADY_RUNNING_ERROR: &str = "Session is already running"; -/// How long `continue_session_for_store` waits for a `running` session that -/// has no turn-control handle yet — the startup window between the `running` -/// write and the runner's `connect` — before answering `Blocked`. -pub(crate) const CONTINUE_SESSION_STARTUP_WAIT: Duration = Duration::from_secs(10); +/// How long each wait inside `route_continue_session` may take. Each has its +/// own budget: a startup wait that ran long must not eat into the wait for a +/// row to settle after its connection ended under the request. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct ContinueWaits { + /// A `running` session with no turn-control handle yet, on a local + /// branch (or no branch): the window between the `running` write and the + /// runner's `connect`. + pub(crate) startup: Duration, + /// The same window on a `Remote` branch, where the connect goes over SSH + /// to the workspace and takes correspondingly longer. + pub(crate) remote_startup: Duration, + /// After a delivery failed because the connection went away under it: + /// how long to wait for the row to leave `running` (or a re-connect to + /// hand out a new handle) before answering blocked. + pub(crate) settle: Duration, +} + +impl ContinueWaits { + /// The production budgets. + pub(crate) const DEFAULT: ContinueWaits = ContinueWaits { + startup: Duration::from_secs(10), + remote_startup: Duration::from_secs(30), + settle: Duration::from_secs(10), + }; + + /// Every budget the same — for tests that need the blocked cases to run + /// in milliseconds. + #[cfg(test)] + pub(crate) fn uniform(wait: Duration) -> Self { + Self { + startup: wait, + remote_startup: wait, + settle: wait, + } + } +} /// How a message handed to `continue_session_for_store` reached (or failed to /// reach) the session's agent. The wire spelling is `delivery()`. @@ -1663,7 +1803,8 @@ pub(crate) enum ContinueOutcome { RejectedQueued, /// Not delivered now, but a retry may succeed: the session is still /// starting, its connection went away under the request, the branch is - /// busy, or another start won the race. + /// busy, another start won the race, or the session runs in another + /// Staged instance that shares this database. Blocked { reason: String }, /// A stopped session that cannot be resumed: no agent conversation to /// load, or its worktree is gone. The caller should start a new session. @@ -1705,14 +1846,36 @@ impl ContinueOutcome { } } +/// A stopped session `continue_resume_gate` claimed for a resume. Its row is +/// already `running` under this process (the claim is the gate's schedule +/// check and the status write in one step under the branch's launch lock, +/// so two continues on one branch cannot both pass); the caller must now +/// finish the resume with `resume_claimed_session_for_store`, or the row +/// stays `running` until the stranded-session sweep ends it. +pub(crate) struct ClaimedResume { + pub(crate) session: store::Session, + /// Live from the claim until the runner has registered the session (see + /// `SessionRegistry::hold`); released when the claim is dropped. + _liveness: session_runner::LivenessHold, +} + +impl std::fmt::Debug for ClaimedResume { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ClaimedResume") + .field("session", &self.session.id) + .finish_non_exhaustive() + } +} + /// Where `route_continue_session` sent the message. #[derive(Debug)] pub(crate) enum ContinueRoute { /// Delivered, refused, or blocked without touching the row. Delivered(ContinueOutcome), - /// A stopped session that passed the resume gate: the caller resumes it - /// with the message as the next prompt. - Resume(Box), + /// A stopped session that passed the resume gate and was claimed + /// `running` for it: the caller resumes it with the message as the next + /// prompt. + Resume(Box), } /// Deliver `message` to a session now, whatever state it is in. @@ -1724,10 +1887,11 @@ pub(crate) enum ContinueRoute { /// `running` and the runner never sees a turn boundary; a stopped session is /// resumed; a `queued` one is refused. The user row is written by whoever /// delivers: acp-client for the live paths (only once the agent took the -/// message), `start_session` for the resume, so nothing here persists it. +/// message, with `image_ids` on the row), `start_session` for the resume, so +/// nothing here persists it. /// -/// Resume failures other than the already-running race are returned as -/// `Err`, as `resume_session_for_store` reports them. +/// Resume failures are returned as `Err`, as `resume_session_for_store` +/// reports them; the row is then recorded failed by the resume itself. #[allow(clippy::too_many_arguments)] pub(crate) async fn continue_session_for_store( store: Arc, @@ -1745,52 +1909,91 @@ pub(crate) async fn continue_session_for_store( &session_id, &message, &image_ids, - CONTINUE_SESSION_STARTUP_WAIT, + ContinueWaits::DEFAULT, ) .await?; - let session = match route { + let claim = match route { ContinueRoute::Delivered(outcome) => return Ok(outcome), - ContinueRoute::Resume(session) => *session, + ContinueRoute::Resume(claim) => *claim, }; - // No `branch_id` override: the resume resolves the branch from the - // session's artifact, or from the row itself for pipeline sessions. - match resume_session_for_store( + resume_claimed_session_for_store( store, registry, action_executor, action_registry, app_handle, - session.id, + claim, + message, + image_ids, + ) + .await + .map(|()| ContinueOutcome::Resumed) +} + +/// Finish the resume of a session `continue_resume_gate` claimed: the same +/// preparation and launch as `resume_session_for_store`, without the claim +/// it already made. No `branch_id` override and no config selection: the +/// resume resolves the branch from the session's artifact (or from the row +/// itself for pipeline sessions) and keeps the session's selection. A +/// preparation failure is recorded on the row like a launch failure — the +/// row is `running` and ours from the claim on. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn resume_claimed_session_for_store( + store: Arc, + registry: Arc, + action_executor: Arc, + action_registry: Arc, + app_handle: tauri::AppHandle, + claim: ClaimedResume, + message: String, + image_ids: Vec, +) -> Result<(), String> { + let ClaimedResume { session, _liveness } = claim; + let plan = match prepare_resume(&store, &session, None, None).await { + Ok(plan) => plan, + Err(e) => { + let branch_id = store.get_branch_id_for_session(&session.id).ok().flatten(); + let project_id = store.get_project_id_for_session(&session.id).ok().flatten(); + return fail_session_start_on_error( + &store, + ®istry, + &app_handle, + &session.id, + branch_id, + project_id, + Err(e), + ); + } + }; + launch_resume( + plan, + store, + registry, + action_executor, + action_registry, + app_handle, + &session, message, Some(image_ids), None, - None, - None, ) .await - { - Ok(()) => Ok(ContinueOutcome::Resumed), - Err(e) if e == SESSION_ALREADY_RUNNING_ERROR => Ok(ContinueOutcome::blocked( - "the session started running again under this request; retry to reach it live", - )), - Err(e) => Err(e), - } } -/// Everything `continue_session_for_store` does short of the resume itself, -/// which needs the app handle. Loops because the status can move under the -/// request: a `running` session whose connection ends while this waits for -/// its handle is routed again on the status it ended with. +/// Everything `continue_session_for_store` does short of the resume's +/// preparation and launch, which need the app handle. Loops because the +/// status can move under the request: a `running` session whose connection +/// ends while this waits for its handle is routed again on the status it +/// ended with. Each wait takes its own budget from `waits`. pub(crate) async fn route_continue_session( - store: &Store, - registry: &session_runner::SessionRegistry, + store: &Arc, + registry: &Arc, session_id: &str, message: &str, image_ids: &[String], - startup_wait: Duration, + waits: ContinueWaits, ) -> Result { - let deadline = Instant::now() + startup_wait; // A handle whose connection went away under this request; the wait below // skips it so only a re-connect's handle (or a status change) ends it. let mut dead_handle: Option = None; @@ -1805,17 +2008,45 @@ pub(crate) async fn route_continue_session( return Ok(ContinueRoute::Delivered(ContinueOutcome::RejectedQueued)); } store::SessionStatus::Running => { + // A `running` row this process does not own is not starting + // up: another Staged instance sharing this database runs it + // (the dead-session sweep leaves such rows alone while that + // process lives), and its connection is in that process. + // Waiting for a handle here would never end in one. + if let Some(owner_pid) = session.owner_pid.filter(|pid| *pid != std::process::id()) + { + return Ok(ContinueRoute::Delivered(ContinueOutcome::blocked(format!( + "the session is owned by another Staged instance (pid {owner_pid}); it \ + can only be reached from that instance" + )))); + } + let budget = if dead_handle.is_some() { + waits.settle + } else if continue_branch_and_type(store, &session)? + .0 + .is_some_and(|branch| branch.branch_type == store::BranchType::Remote) + { + waits.remote_startup + } else { + waits.startup + }; let handle = match wait_for_turn_control_handle( store, registry, session_id, - deadline, + Instant::now() + budget, dead_handle.as_ref(), ) .await { HandleWait::Ready(handle) => handle, HandleWait::LeftRunning => continue, + HandleWait::TimedOut if dead_handle.is_some() => { + return Ok(ContinueRoute::Delivered(ContinueOutcome::blocked( + "the session's connection ended under the request and the session \ + has not settled yet; retry in a few seconds", + ))); + } HandleWait::TimedOut => { return Ok(ContinueRoute::Delivered(ContinueOutcome::blocked( "the session is still starting; retry in a few seconds", @@ -1843,9 +2074,9 @@ pub(crate) async fn route_continue_session( store::SessionStatus::Completed | store::SessionStatus::Cancelled | store::SessionStatus::Error => { - return Ok(match continue_resume_gate(store, &session)? { - Some(refused) => ContinueRoute::Delivered(refused), - None => ContinueRoute::Resume(Box::new(session)), + return Ok(match continue_resume_gate(store, registry, session)? { + Ok(claim) => ContinueRoute::Resume(Box::new(claim)), + Err(refused) => ContinueRoute::Delivered(refused), }); } } @@ -1890,73 +2121,113 @@ async fn wait_for_turn_control_handle( /// Reach a running session's agent over its connection: steer, and when the /// agent refuses (no steering advertised, `promptRequired` because the turn -/// ended under the request, `failed`), interrupt the turn and re-prompt. -/// `Err` only when the interrupt fails too — the caller decides whether that -/// is a dead connection or a blocked delivery. +/// ended under the request, `failed`, no answer in time), interrupt the turn +/// and re-prompt. `Err` only when the interrupt fails too — the caller +/// decides whether that is a dead connection or a blocked delivery. The +/// `image_ids` ride along as the row's attachment ids, like `start_session` +/// records them, whether or not every file could be read. async fn deliver_to_running_session( - store: &Store, + store: &Arc, handle: &acp_client::TurnControlHandle, message: &str, image_ids: &[String], ) -> Result { - let attachments = session_runner::load_prompt_attachments(store, image_ids); + // File reads and, for PDFs, page rendering: off the async runtime that + // serves the MCP server and the Tauri commands. + let attachments = { + let store = Arc::clone(store); + let image_ids = image_ids.to_vec(); + tokio::task::spawn_blocking(move || { + session_runner::load_prompt_attachments(&store, &image_ids) + }) + .await + .map_err(|e| format!("could not load the message's attachments: {e}"))? + }; let prompt = attachments.prompt(message); - let steer_refusal = match handle.steer(&prompt, &attachments.images).await { + let steer_refusal = match handle.steer(&prompt, &attachments.images, image_ids).await { Ok(acp_client::SteerOutcome::Injected) => return Ok(ContinueOutcome::Injected), Ok(acp_client::SteerOutcome::NewTurnStarted) => return Ok(ContinueOutcome::StartedNewTurn), Err(reason) => reason, }; log::info!("continue_session: steer refused ({steer_refusal}); interrupting the turn instead"); handle - .interrupt_and_reprompt(&prompt, &attachments.images) + .interrupt_and_reprompt(&prompt, &attachments.images, image_ids) .await .map(|()| ContinueOutcome::InterruptedAndReprompted) .map_err(|reason| format!("could not reach the running session: {reason}")) } -/// Whether a stopped session may be resumed with a new message: it needs an -/// agent conversation to load and, for a local branch, its worktree on disk; -/// on a branch, the start must also be allowed by the branch schedule (no -/// provisioning in progress, nothing queued ahead of it, no conflicting -/// running session). `Some` is the refusal; `None` lets the resume proceed. +/// Whether a stopped session may be resumed with a new message — and if so, +/// claim it. It needs an agent conversation to load and, for a local branch, +/// its worktree on disk; on a branch, the start must also be allowed by the +/// branch schedule (no provisioning in progress, nothing queued ahead of it, +/// no conflicting running session). `Err` is the refusal; `Ok` is the claim: +/// the row is `running` under this process. /// -/// The schedule check runs under the branch's launch lock like the branch -/// session starts do, but the lock is not held into the resume's `running` -/// write (that write is itself atomic on the status), so two continues on -/// one branch can both pass the gate; the second then finds the first's row -/// running and is reported blocked. +/// The schedule check and the `running` write happen together under the +/// branch's launch lock with no await between them, as every branch session +/// start claims (see `QueuedSessionClaim::claim`). Checked and written apart, +/// two different stopped exclusive sessions on one branch continued at once +/// both pass the check and both resume; held together, the second finds the +/// first's row running and is blocked. A session off any branch is claimed +/// without the lock — nothing schedules against it — and the status CAS +/// alone refuses a concurrent claim of the same row. pub(crate) fn continue_resume_gate( - store: &Store, - session: &store::Session, -) -> Result, String> { + store: &Arc, + registry: &Arc, + session: store::Session, +) -> Result, String> { if resumable_agent_session_id(session.agent_id.as_deref()).is_none() { - return Ok(Some(ContinueOutcome::not_resumable( + return Ok(Err(ContinueOutcome::not_resumable( "the session has no agent conversation to resume; start a new session instead", ))); } - let (branch, session_type) = continue_branch_and_type(store, session)?; + let (branch, session_type) = continue_branch_and_type(store, &session)?; let is_remote = branch .as_ref() .is_some_and(|branch| branch.branch_type == store::BranchType::Remote); if !is_remote && !Path::new(&session.working_dir).is_dir() { - return Ok(Some(ContinueOutcome::not_resumable(format!( + return Ok(Err(ContinueOutcome::not_resumable(format!( "the session's working directory is gone ({}); start a new session instead", session.working_dir )))); } - if let (Some(branch), Some(session_type)) = (branch, session_type) { - let launch_lock = branch_session_launch_lock_for(&branch.id); - let _guard = launch_lock.lock().unwrap(); - if should_queue_branch_session_start(store, &branch.id, &session_type)? { - return Ok(Some(ContinueOutcome::blocked( - "the branch is provisioning or busy with another session; retry once it is idle", - ))); + let pending_start_agent_id = + pending_branch_start_agent_id(session.agent_id.as_deref()).map(ToString::to_string); + // The status CAS, with the liveness hold taken first (see + // `SessionRegistry::hold`): `Some` holds the claimed row live, `None` + // (the row is already `running`) releases the hold again. + let claim = || -> Result, String> { + let liveness = registry.hold(&session.id); + let claimed = store + .transition_to_running_clearing_agent_id(&session.id, pending_start_agent_id.as_deref()) + .map_err(|e| e.to_string())?; + Ok(claimed.then_some(liveness)) + }; + let liveness = match (branch, session_type) { + (Some(branch), Some(session_type)) => { + let launch_lock = branch_session_launch_lock_for(&branch.id); + let _guard = launch_lock.lock().unwrap(); + if should_queue_branch_session_start(store, &branch.id, &session_type)? { + return Ok(Err(ContinueOutcome::blocked( + "the branch is provisioning or busy with another session; retry once it is idle", + ))); + } + claim()? } - } - - Ok(None) + _ => claim()?, + }; + let Some(liveness) = liveness else { + return Ok(Err(ContinueOutcome::blocked( + "the session started running again under this request; retry to reach it live", + ))); + }; + Ok(Ok(ClaimedResume { + session, + _liveness: liveness, + })) } /// The branch a session runs on and the schedule kind it counts as, from its @@ -7861,18 +8132,28 @@ mod tests { /// and the resume gate, with a short startup wait so the blocked cases /// run in milliseconds. async fn continue_through_shared_body( - registry: &session_runner::SessionRegistry, + registry: &Arc, store: &Arc, session_id: &str, message: &str, + ) -> ContinueRoute { + continue_with_images(registry, store, session_id, message, &[]).await + } + + async fn continue_with_images( + registry: &Arc, + store: &Arc, + session_id: &str, + message: &str, + image_ids: &[String], ) -> ContinueRoute { route_continue_session( store, registry, session_id, message, - &[], - Duration::from_millis(300), + image_ids, + ContinueWaits::uniform(Duration::from_millis(300)), ) .await .unwrap() @@ -7881,8 +8162,8 @@ mod tests { fn delivered(route: ContinueRoute) -> ContinueOutcome { match route { ContinueRoute::Delivered(outcome) => outcome, - ContinueRoute::Resume(session) => { - panic!("expected a delivery, got a resume of {}", session.id) + ContinueRoute::Resume(claim) => { + panic!("expected a delivery, got a resume of {}", claim.session.id) } } } @@ -8112,7 +8393,7 @@ mod tests { settle.await.unwrap(); match route { - ContinueRoute::Resume(resumed) => assert_eq!(resumed.id, session.id), + ContinueRoute::Resume(resumed) => assert_eq!(resumed.session.id, session.id), other => panic!("expected a resume, got {other:?}"), } } @@ -8224,22 +8505,253 @@ mod tests { /// A stopped session with a conversation to load, its worktree on disk /// and an idle branch goes to the resume, which `start_session` then - /// persists the message for. + /// persists the message for. The gate claims the row: it is `running` + /// under this process from here, held live until the runner registers. #[tokio::test(flavor = "current_thread")] async fn continuing_a_resumable_stopped_session_routes_to_the_resume() { let (store, branch) = setup_branch_store_with_workdir(); let session = completed_commit_session_with_agent_id(&store, &branch.id); + let registry = test_registry(); - let route = continue_through_shared_body(&test_registry(), &store, &session.id, "x").await; + let route = continue_through_shared_body(®istry, &store, &session.id, "x").await; + + let claim = match route { + ContinueRoute::Resume(claim) => claim, + other => panic!("expected a resume, got {other:?}"), + }; + assert_eq!(claim.session.id, session.id); + let row = store.get_session(&session.id).unwrap().unwrap(); + assert_eq!(row.status, store::SessionStatus::Running); + assert_eq!(row.owner_pid, Some(std::process::id())); + assert_eq!( + row.agent_id.as_deref(), + Some("acp-session-1"), + "the conversation id the resume loads is kept" + ); + assert!(registry.is_live(&session.id), "held live by the claim"); + drop(claim); + assert!(!registry.is_live(&session.id)); + } + + /// The race the claim exists for: two *different* stopped exclusive + /// sessions on one branch continued at the same time. The schedule check + /// and the `running` write are one step under the branch launch lock, so + /// exactly one is claimed and the other sees the first's row running and + /// is blocked — never two running commit sessions on one branch. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn concurrent_continues_of_two_stopped_sessions_on_one_branch_claim_only_one() { + let (store, branch) = setup_branch_store_with_workdir(); + let first = completed_commit_session_with_agent_id(&store, &branch.id); + let second = completed_commit_session_with_agent_id(&store, &branch.id); + let registry = test_registry(); + let spawn_continue = |session_id: String| { + let store = Arc::clone(&store); + let registry = Arc::clone(®istry); + tokio::spawn(async move { + route_continue_session( + &store, + ®istry, + &session_id, + "x", + &[], + ContinueWaits::uniform(Duration::from_millis(300)), + ) + .await + .unwrap() + }) + }; + let (first_route, second_route) = tokio::join!( + spawn_continue(first.id.clone()), + spawn_continue(second.id.clone()) + ); + + let mut claimed = Vec::new(); + let mut blocked = Vec::new(); + for route in [first_route, second_route] { + match route.unwrap() { + ContinueRoute::Resume(claim) => claimed.push(claim), + ContinueRoute::Delivered(ContinueOutcome::Blocked { reason }) => { + blocked.push(reason) + } + other => panic!("unexpected route {other:?}"), + } + } + assert_eq!(claimed.len(), 1, "exactly one resume: {blocked:?}"); + assert_eq!(blocked.len(), 1); + assert!(blocked[0].contains("busy"), "{}", blocked[0]); + let running: Vec = store + .get_running_sessions() + .unwrap() + .into_iter() + .map(|session| session.id) + .collect(); + assert_eq!(running, vec![claimed[0].session.id.clone()]); + } + + /// A `running` row owned by another Staged instance (two dev builds on + /// one `~/.staged/data.db`) is not in its startup window: its connection + /// lives in that process and no handle will ever appear here. Answered + /// blocked at once, with the owner named, instead of after the full + /// startup wait on every retry. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_session_owned_by_another_instance_is_blocked_without_waiting() { + let store = Arc::new(Store::in_memory().unwrap()); + let mut session = store::Session::new_running("prompt", Path::new("/tmp")); + let other_pid = std::process::id() + 1; + session.owner_pid = Some(other_pid); + store.create_session(&session).unwrap(); + let registry = test_registry(); + + let started = Instant::now(); + let route = route_continue_session( + &store, + ®istry, + &session.id, + "x", + &[], + ContinueWaits::uniform(Duration::from_secs(5)), + ) + .await + .unwrap(); + + assert!(started.elapsed() < Duration::from_secs(1)); + match delivered(route) { + ContinueOutcome::Blocked { reason } => { + assert!(reason.contains("another Staged instance"), "{reason}"); + assert!(reason.contains(&other_pid.to_string()), "{reason}"); + } + other => panic!("expected blocked, got {other:?}"), + } + } + + /// The startup window on a `Remote` branch gets the longer budget: the + /// runner connects over SSH to the workspace before it can register a + /// handle. + #[tokio::test(flavor = "current_thread")] + async fn continuing_a_starting_remote_session_waits_the_remote_budget() { + let (store, branch) = setup_remote_branch_store(store::WorkspaceStatus::Running); + let session = + create_branch_commit_session(&store, &branch.id, store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + + let started = Instant::now(); + let route = route_continue_session( + &store, + ®istry, + &session.id, + "x", + &[], + ContinueWaits { + startup: Duration::from_millis(100), + remote_startup: Duration::from_millis(500), + settle: Duration::from_millis(100), + }, + ) + .await + .unwrap(); + + assert!( + started.elapsed() >= Duration::from_millis(500), + "waited only {:?}", + started.elapsed() + ); + match delivered(route) { + ContinueOutcome::Blocked { reason } => { + assert!(reason.contains("still starting"), "{reason}") + } + other => panic!("expected blocked, got {other:?}"), + } + } + + /// The wait for a row to settle after its connection ended under the + /// request has its own budget: here the handle lands late in the + /// startup window, the delivery finds the connection gone, and the row + /// settles after the startup budget would have run out — on the settle + /// budget the message is routed to the resume rather than blocked. + #[tokio::test(flavor = "current_thread")] + async fn settle_wait_after_a_dead_connection_has_its_own_budget() { + let (store, branch) = setup_branch_store_with_workdir(); + let session = + create_branch_commit_session(&store, &branch.id, store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, served) = acp_client::TurnControlHandle::scripted_then_gone( + Err("connection closed".to_string()), + Err("ACP connection is no longer running".to_string()), + ); + let script = { + let registry = Arc::clone(®istry); + let store = Arc::clone(&store); + let session_id = session.id.clone(); + tokio::spawn(async move { + // Most of a 300 ms startup budget goes by before the connect. + tokio::time::sleep(Duration::from_millis(220)).await; + registry.set_turn_control_handle(&session_id, handle); + // The row settles 250 ms later: past the startup budget, + // inside the settle budget. + tokio::time::sleep(Duration::from_millis(250)).await; + store + .update_session_status( + &session_id, + store::SessionStatus::Completed, + None, + Some(&store::CompletionReason::TurnComplete), + ) + .unwrap(); + store.set_agent_session_id(&session_id, "acp-1").unwrap(); + }) + }; + + let route = continue_through_shared_body(®istry, &store, &session.id, "x").await; + + script.await.unwrap(); match route { - ContinueRoute::Resume(resumed) => assert_eq!(resumed.id, session.id), + ContinueRoute::Resume(claim) => assert_eq!(claim.session.id, session.id), other => panic!("expected a resume, got {other:?}"), } - // The gate decides; it does not write. + // The steer's answer came with the connection already gone, so the + // interrupt failed fast at the handle and never reached the agent. + assert_eq!(*served.lock().unwrap(), vec!["steer:x".to_string()]); + } + + /// The message's image ids go to the connection with the prompt and come + /// back on the persisted user row (acp-client hands them to the writer), + /// on the steer and on the interrupt alike; the resume path records them + /// through `start_session` as before. + #[tokio::test(flavor = "current_thread")] + async fn image_ids_travel_with_the_message_on_the_live_path() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, served) = acp_client::TurnControlHandle::scripted( + Err("agent does not advertise steering".to_string()), + Ok(()), + ); + registry.set_turn_control_handle(&session.id, handle); + + let outcome = delivered( + continue_with_images( + ®istry, + &store, + &session.id, + "see attached", + &ids(&["img-1", "img-2"]), + ) + .await, + ); + + assert_eq!(outcome, ContinueOutcome::InterruptedAndReprompted); assert_eq!( - store.get_session(&session.id).unwrap().unwrap().status, - store::SessionStatus::Completed + *served.lock().unwrap(), + vec![ + "steer:see attached".to_string(), + "attachments:img-1,img-2".to_string(), + "interrupt:see attached".to_string(), + "attachments:img-1,img-2".to_string(), + ] ); } diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index ef2245ace..8c6af0b81 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -85,9 +85,18 @@ pub trait MessageWriter: Send + Sync { /// /// Called on the connection task right after the open assistant message /// is finalized and before anything the agent says in reply can be - /// appended, so the transcript reads assistant → user → assistant. The - /// default records nothing. - async fn record_user_message(&self, _text: &str, _images: &[PromptImage]) {} + /// appended, so the transcript reads assistant → user → assistant. + /// `attachment_ids` are the caller's own ids for the attachments behind + /// `images` (opaque here: this crate sends bytes, not ids), handed back + /// so the row can reference the stored attachments. The default records + /// nothing. + async fn record_user_message( + &self, + _text: &str, + _images: &[PromptImage], + _attachment_ids: &[String], + ) { + } /// Record a tool call with its ID, title, and optional raw input parameters. async fn record_tool_call( @@ -1445,6 +1454,9 @@ type SessionTurnResult = Result; struct QueuedSessionTurn { prompt: String, images: Vec, + /// The caller's ids for `images`, recorded with the user row when the + /// connection persists it (see [`MessageWriter::record_user_message`]). + attachment_ids: Vec, /// Resolved by the connection task with the turn's outcome. `None` for a /// turn that did not come through [`SessionConnection::prompt`]: one the /// turn-control channel started (served out of a hold, or re-prompted @@ -1478,6 +1490,7 @@ enum TurnControlRequest { Steer { prompt: String, images: Vec, + attachment_ids: Vec, reply: oneshot::Sender>, }, /// `session/cancel` the in-flight turn, then send the message as the next @@ -1485,6 +1498,7 @@ enum TurnControlRequest { Interrupt { prompt: String, images: Vec, + attachment_ids: Vec, reply: oneshot::Sender>, }, } @@ -1495,21 +1509,34 @@ impl TurnControlRequest { /// back `promptRequired` (or, from codex-acp, start a turn this client /// would not own), and an interrupt has nothing to cancel. fn into_turn(self) -> QueuedSessionTurn { - let (prompt, images, control) = match self { + let (prompt, images, attachment_ids, control) = match self { Self::Steer { prompt, images, + attachment_ids, reply, - } => (prompt, images, TurnControlReply::Steer(reply)), + } => ( + prompt, + images, + attachment_ids, + TurnControlReply::Steer(reply), + ), Self::Interrupt { prompt, images, + attachment_ids, reply, - } => (prompt, images, TurnControlReply::Interrupt(reply)), + } => ( + prompt, + images, + attachment_ids, + TurnControlReply::Interrupt(reply), + ), }; QueuedSessionTurn { prompt, images, + attachment_ids, reply: None, control: Some(control), } @@ -1546,6 +1573,18 @@ impl TurnControlReply { } } } + + /// Answer the caller: the message is not going out after all. + fn reject(self, reason: String) { + match self { + Self::Steer(reply) => { + let _ = reply.send(Err(reason)); + } + Self::Interrupt(reply) => { + let _ = reply.send(Err(reason)); + } + } + } } /// A per-task stop queued onto a [`SessionConnection`]'s hold-control channel, @@ -1618,19 +1657,23 @@ impl TurnControlHandle { /// message is sent as a fresh `session/prompt` on the same connection and /// the answer is `Ok(NewTurnStarted)`. `Err` when the agent does not /// advertise steering, answered anything but `injected` (`promptRequired` - /// when the turn ended under the request, `failed`, `startedNewTurn`), or - /// the connection is gone; callers then fall back to - /// [`TurnControlHandle::interrupt_and_reprompt`]. The user message is - /// persisted (see [`MessageWriter::record_user_message`]) only on `Ok`. + /// when the turn ended under the request, `failed`, `startedNewTurn`), + /// did not answer within `STEER_ACK_TIMEOUT`, or the connection is gone; + /// callers then fall back to [`TurnControlHandle::interrupt_and_reprompt`]. + /// The user message is persisted (see + /// [`MessageWriter::record_user_message`], with `attachment_ids`) only on + /// `Ok` — an answer that arrives after the timeout persists nothing. pub async fn steer( &self, prompt: &str, images: &[PromptImage], + attachment_ids: &[String], ) -> Result { let (reply_tx, reply_rx) = oneshot::channel(); self.send(TurnControlRequest::Steer { prompt: prompt.to_string(), images: images.to_vec(), + attachment_ids: attachment_ids.to_vec(), reply: reply_tx, })?; reply_rx @@ -1648,16 +1691,21 @@ impl TurnControlHandle { /// /// Resolves once the new prompt is going out. `Err` if the agent did not /// acknowledge the cancel within `INTERRUPT_ACK_TIMEOUT` (the turn is - /// then cancelled without a re-prompt) or the connection is gone. + /// then cancelled without a re-prompt), the session was stopped while + /// the interrupt was pending (the stop wins; nothing is re-prompted), the + /// cancel could not be sent (the turn keeps running), or the connection + /// is gone. `attachment_ids` go onto the persisted user row. pub async fn interrupt_and_reprompt( &self, prompt: &str, images: &[PromptImage], + attachment_ids: &[String], ) -> Result<(), String> { let (reply_tx, reply_rx) = oneshot::channel(); self.send(TurnControlRequest::Interrupt { prompt: prompt.to_string(), images: images.to_vec(), + attachment_ids: attachment_ids.to_vec(), reply: reply_tx, })?; reply_rx @@ -1694,26 +1742,70 @@ impl TurnControlHandle { /// A handle to a stand-in connection that answers every steer with /// `steer` and every interrupt with `interrupt`, and appends /// `"steer:"` / `"interrupt:"` to the returned log as the - /// requests arrive. For tests of the routing around a handle (steer first, - /// interrupt on refusal) that must not stand up a real agent. Needs a - /// tokio runtime: the stand-in is a task. + /// requests arrive (followed by `"attachments:"` when the request + /// carried attachment ids). For tests of the routing around a handle + /// (steer first, interrupt on refusal) that must not stand up a real + /// agent. Needs a tokio runtime: the stand-in is a task. #[doc(hidden)] pub fn scripted( steer: Result, interrupt: Result<(), String>, + ) -> (Self, Arc>>) { + Self::scripted_connection(steer, interrupt, false) + } + + /// [`Self::scripted`] whose connection reports itself gone right after it + /// answers its first request — the connection that settles under the + /// request. The stand-in keeps answering (every later request is logged + /// and answered the same way), but [`Self::is_connection_alive`] reads + /// `false` from the first answer on. + #[doc(hidden)] + pub fn scripted_then_gone( + steer: Result, + interrupt: Result<(), String>, + ) -> (Self, Arc>>) { + Self::scripted_connection(steer, interrupt, true) + } + + fn scripted_connection( + steer: Result, + interrupt: Result<(), String>, + gone_after_first_answer: bool, ) -> (Self, Arc>>) { let (turn_control_tx, mut turn_control_rx) = mpsc::unbounded_channel(); let log = Arc::new(std::sync::Mutex::new(Vec::new())); let served = Arc::clone(&log); + let connection_alive = Arc::new(AtomicBool::new(true)); + let alive = Arc::clone(&connection_alive); tokio::spawn(async move { while let Some(request) = turn_control_rx.recv().await { + let (entry, attachment_ids): (String, Vec) = match &request { + TurnControlRequest::Steer { + prompt, + attachment_ids, + .. + } => (format!("steer:{prompt}"), attachment_ids.clone()), + TurnControlRequest::Interrupt { + prompt, + attachment_ids, + .. + } => (format!("interrupt:{prompt}"), attachment_ids.clone()), + }; + { + let mut served = served.lock().unwrap(); + served.push(entry); + if !attachment_ids.is_empty() { + served.push(format!("attachments:{}", attachment_ids.join(","))); + } + } + if gone_after_first_answer { + alive.store(false, Ordering::Relaxed); + } match request { - TurnControlRequest::Steer { prompt, reply, .. } => { - served.lock().unwrap().push(format!("steer:{prompt}")); + TurnControlRequest::Steer { reply, .. } => { let _ = reply.send(steer.clone()); } - TurnControlRequest::Interrupt { prompt, reply, .. } => { - served.lock().unwrap().push(format!("interrupt:{prompt}")); + TurnControlRequest::Interrupt { reply, .. } => { let _ = reply.send(interrupt.clone()); } } @@ -1722,7 +1814,7 @@ impl TurnControlHandle { ( Self { turn_control_tx, - connection_alive: Arc::new(AtomicBool::new(true)), + connection_alive, }, log, ) @@ -1849,6 +1941,7 @@ impl SessionConnection { let turn = QueuedSessionTurn { prompt: prompt.to_string(), images: images.to_vec(), + attachment_ids: vec![], reply: Some(reply_tx), control: None, }; @@ -5012,6 +5105,14 @@ fn serve_async_task_stop( /// — the transcript reads assistant → user → assistant with no interleaving /// window — while a `promptRequired` or failed steer leaves no row behind for /// the caller's fallback delivery to duplicate. +/// +/// The answer is bounded by [`STEER_ACK_TIMEOUT`]: the caller's reply sender +/// sits in a slot that the response callback and a timer race for, and only +/// the one that takes it answers. A `_session/steering` answer that lands +/// after the timer has answered `Err` is logged and otherwise ignored — in +/// particular it persists no user row, because the caller has by then fallen +/// back to an interrupt that persists the message itself. +#[allow(clippy::too_many_arguments)] fn serve_steer( connection: &ConnectionTo, setup: &AcpSessionSetup, @@ -5019,6 +5120,7 @@ fn serve_steer( our_session_id: &str, prompt: String, images: Vec, + attachment_ids: Vec, reply: oneshot::Sender>, ) { if !setup.steering_supported { @@ -5039,17 +5141,31 @@ fn serve_steer( }; let handler = Arc::clone(handler); let session_id = our_session_id.to_string(); + let reply_slot = Arc::new(std::sync::Mutex::new(Some(reply))); + let timed_out_slot = Arc::clone(&reply_slot); + let timed_out_session_id = session_id.clone(); let scheduled = connection .send_request(message) .on_receiving_result(move |result| async move { let outcome = steer_outcome(result); + let Some(reply) = reply_slot.lock().unwrap().take() else { + log::warn!( + "ACP session {session_id}: steer answered {outcome:?} after the \ + {}s timeout already failed it; ignoring the late answer", + STEER_ACK_TIMEOUT.as_secs() + ); + return Ok(()); + }; match &outcome { Ok(outcome) => { log::info!("ACP session {session_id}: steer answered {outcome:?}"); handler.writer.finalize().await; handler.response.lock().await.detach_message(); - handler.writer.record_user_message(&prompt, &images).await; + handler + .writer + .record_user_message(&prompt, &images, &attachment_ids) + .await; } Err(e) => log::warn!("ACP session {session_id}: steer not injected: {e}"), } @@ -5060,6 +5176,27 @@ fn serve_steer( // The reply sender moved into the callback and drops with it, which // already answers the caller; the send failure is only loggable here. log::warn!("ACP session {our_session_id}: failed to send steering request: {e:?}"); + return; + } + // The timer runs on the connection's task set, like the permission + // handler's decision, so it is dropped with the connection: a steer the + // teardown drain answers is not answered a second time here. + let timer = connection.spawn(async move { + tokio::time::sleep(STEER_ACK_TIMEOUT).await; + if let Some(reply) = timed_out_slot.lock().unwrap().take() { + log::warn!( + "ACP session {timed_out_session_id}: agent did not answer the steer within {}s", + STEER_ACK_TIMEOUT.as_secs() + ); + let _ = reply.send(Err(format!( + "Agent did not answer the steer within {}s", + STEER_ACK_TIMEOUT.as_secs() + ))); + } + Ok(()) + }); + if let Err(e) = timer { + log::warn!("ACP session {our_session_id}: failed to arm the steer timeout: {e:?}"); } } @@ -5080,6 +5217,13 @@ fn serve_steer( /// down first — and real tool calls run longer than that. const INTERRUPT_ACK_TIMEOUT: Duration = Duration::from_secs(30); +/// How long a steer waits for the agent's `_session/steering` answer before +/// the request fails and the caller falls back to an interrupt. The same +/// bound as [`INTERRUPT_ACK_TIMEOUT`], for the same reason: a bridge that +/// takes the frame and never answers while the turn runs on must not wedge +/// the caller's tool call (which has no keepalive of its own). +const STEER_ACK_TIMEOUT: Duration = Duration::from_secs(30); + /// How long the holding wait gives the bridge to process a `session/cancel` /// before teardown. Out-of-turn there is no prompt response to await, so this /// bounded courtesy drain mirrors the in-turn cancel's bounded response wait @@ -5649,6 +5793,19 @@ async fn run_acp_session( let mut turn_control_open = true; loop { let mut turn = match next_turn.take() { + // A Stop that landed while the interrupt (or the hold) was + // handing this turn over wins: the message is not sent and its + // caller is told so, instead of a prompt going out that the + // next select cancels at once — with a user row persisted and + // the caller told it was delivered. + Some(mut turn) if cancel_token.is_cancelled() => { + if let Some(control) = turn.control.take() { + control.reject("Session stopped before the message could be sent".to_string()); + } + handler.cancel_pending_permissions(); + transition_lifetime(our_session_id, &mut lifetime, SessionLifetime::Cancelled); + return Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)); + } // A prompt that arrived during the holding wait (or the message // an interrupt was made for) re-enters a live turn on the // already-established session. @@ -5700,7 +5857,7 @@ async fn run_acp_session( handler.response.lock().await.close_message(); handler .writer - .record_user_message(&turn.prompt, &turn.images) + .record_user_message(&turn.prompt, &turn.images, &turn.attachment_ids) .await; control.turn_started(); } @@ -5983,10 +6140,24 @@ async fn run_prompt_turn( } } request = turn_control_rx.recv(), if control_open => match request { - Some(TurnControlRequest::Steer { prompt, images, reply }) => { - serve_steer(connection, setup, handler, our_session_id, prompt, images, reply); + Some(TurnControlRequest::Steer { prompt, images, attachment_ids, reply }) => { + serve_steer( + connection, + setup, + handler, + our_session_id, + prompt, + images, + attachment_ids, + reply, + ); } - Some(TurnControlRequest::Interrupt { prompt: next_prompt, images: next_images, reply }) => { + Some(TurnControlRequest::Interrupt { + prompt: next_prompt, + images: next_images, + attachment_ids: next_attachment_ids, + reply, + }) => { log::info!( "ACP session {our_session_id}: interrupting the turn in flight for a new message" ); @@ -6001,32 +6172,49 @@ async fn run_prompt_turn( // the turn is gone the dialog is stale anyway. handler.cancel_turn_permissions(); if let Err(e) = send_session_cancel(connection, &setup.agent_session_id) { - let _ = reply.send(Err(e.clone())); - return Err(e); + // Like the user-cancel arm: a cancel that could not + // be sent is logged and the turn goes on. Only the + // interrupt fails — its message was not delivered, + // but the turn's work is not lost over it. + log::warn!("{e}"); + let _ = reply.send(Err(e)); + continue; } - match tokio::time::timeout(INTERRUPT_ACK_TIMEOUT, &mut prompt_task).await { + // Biased towards the acknowledgment: a Stop that is + // already in by the time the ack arrives is caught by + // the session loop before the re-prompt goes out. A Stop + // that lands while the ack is still outstanding ends the + // wait here, so the re-prompt is never sent and the + // caller is not told it was. + let acked = tokio::select! { + biased; + result = tokio::time::timeout(INTERRUPT_ACK_TIMEOUT, &mut prompt_task) => Some(result), + _ = cancel_token.cancelled() => None, + }; + match acked { // Whatever the stop reason (`cancelled`, or `end_turn` // if the turn finished under the cancel), the turn is // over and the message is the next one. - Ok(Ok(response)) => { + Some(Ok(Ok(response))) => { if let Some(metadata) = prompt_response_metadata(&response) { handler.writer.record_acp_event_metadata(metadata).await; } return Ok(TurnStep::Interrupted(QueuedSessionTurn { prompt: next_prompt, images: next_images, + attachment_ids: next_attachment_ids, reply: None, control: Some(TurnControlReply::Interrupt(reply)), })); } - Ok(Err(e)) => { + Some(Ok(Err(e))) => { let message = format!("Prompt failed after interrupt: {e:?}"); let _ = reply.send(Err(message.clone())); return Err(message); } // A wedged bridge should not get a second prompt: // settle the turn cancelled, as a plain cancel would. - Err(_) => { + Some(Err(_)) => { log::warn!( "Timed out waiting for ACP prompt response after session/cancel \ for session {our_session_id}; dropping the re-prompt" @@ -6037,6 +6225,30 @@ async fn run_prompt_turn( ))); return Ok(TurnStep::Settled(AgentRunOutcome::Cancelled)); } + // The session was stopped under the interrupt. The + // cancel is already on the wire; finish as the + // user-cancel arm does, with the message undelivered. + None => { + log::info!( + "ACP session {our_session_id}: stopped while an interrupt awaited \ + the cancel acknowledgment; dropping the re-prompt" + ); + let _ = reply.send(Err( + "Session stopped before the message could be sent".to_string(), + )); + handler.cancel_pending_permissions(); + match tokio::time::timeout(Duration::from_secs(5), &mut prompt_task).await { + Ok(result) => { + break result.map_err(|e| format!("Prompt failed after cancellation: {e:?}"))?; + } + Err(_) => { + log::warn!( + "Timed out waiting for ACP prompt response after session/cancel for session {our_session_id}" + ); + return Ok(TurnStep::Settled(AgentRunOutcome::Cancelled)); + } + } + } } } None => control_open = false, @@ -11648,6 +11860,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE .send(QueuedSessionTurn { prompt: String::from("follow-up"), images: vec![], + attachment_ids: vec![], reply: Some(reply_tx), control: None, }) @@ -11915,6 +12128,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE .send(QueuedSessionTurn { prompt: String::from("follow-up"), images: vec![], + attachment_ids: vec![], reply: Some(reply_tx), control: None, }) @@ -12247,6 +12461,9 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE Finalize, Resume, UserMessage(String), + /// The attachment ids handed back with a `UserMessage` (only when + /// there were any). + UserAttachments(Vec), Event(RecordedMetadata), ToolCall(String), ToolCallMetadata(RecordedMetadata), @@ -12331,11 +12548,22 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE self.calls.lock().unwrap().push(WriterCall::Finalize); } - async fn record_user_message(&self, text: &str, _images: &[PromptImage]) { + async fn record_user_message( + &self, + text: &str, + _images: &[PromptImage], + attachment_ids: &[String], + ) { self.calls .lock() .unwrap() .push(WriterCall::UserMessage(text.to_string())); + if !attachment_ids.is_empty() { + self.calls + .lock() + .unwrap() + .push(WriterCall::UserAttachments(attachment_ids.to_vec())); + } } async fn record_tool_call( diff --git a/crates/acp-client/src/driver/response_completion.rs b/crates/acp-client/src/driver/response_completion.rs index 01d437169..2fd3d5b4d 100644 --- a/crates/acp-client/src/driver/response_completion.rs +++ b/crates/acp-client/src/driver/response_completion.rs @@ -15,6 +15,10 @@ pub(super) struct ResponseState { pub(super) completed: bool, fallback: String, seen_results: HashSet, + /// Text of messages a steer closed mid-turn (see [`Self::detach_message`]), + /// each awaiting the `result` of the cycle the steer aborted. Several + /// steers can land before any such result arrives, so this is a list. + detached: Vec, } fn successful_result(message: &serde_json::Value) -> bool { @@ -38,6 +42,7 @@ impl ResponseState { self.close_message(); self.completed = false; self.fallback.clear(); + self.detached.clear(); } /// Forget the streamed message: the writer has closed (or will close) it. @@ -49,18 +54,46 @@ impl ResponseState { /// The writer closed the streamed message mid-turn, with the turn going /// on — a steered user row went in after it. Unlike [`Self::close_message`] - /// the streamed text is kept: the bridge aborts the generation a steer - /// pre-empts and still emits that cycle's `result`, whose final text is - /// the message just closed. With the buffer intact [`Self::accept_result`] - /// recognises it as already streamed instead of re-injecting the whole - /// message as a chunk after the user row. Only the id is dropped, so a - /// later chunk of the same message opens a row of its own rather than - /// resuming the closed one. + /// the streamed text is remembered, in `detached`: the bridge aborts the + /// generation a steer pre-empts and still emits that cycle's `result`, + /// whose final text is the message just closed, and + /// [`Self::accept_result`] matches that result against the detached text + /// instead of re-injecting the whole message as a chunk after the user + /// row. The text is moved out of the live buffer rather than kept in it + /// so the steered reply's own chunks — which may arrive *before* the + /// aborted cycle's result — start a clean buffer: with the two + /// concatenated the aborted result matched nothing, retired the buffer, + /// and the reply was split into a second row at its next chunk. pub(super) fn detach_message(&mut self) { + let text = std::mem::take(&mut self.text); + if !text.is_empty() { + self.detached.push(text); + } self.message_id = None; self.interrupted = false; } + /// Whether `final_text` is the final text of a cycle a steer aborted + /// (see [`Self::detach_message`]); consumes the match. The bridge's + /// result carries that cycle's last assistant message, which is the + /// detached text or its tail. + fn take_detached_match(&mut self, final_text: &str) -> bool { + if final_text.is_empty() { + return false; + } + match self + .detached + .iter() + .position(|text| text.ends_with(final_text)) + { + Some(index) => { + self.detached.remove(index); + true + } + None => false, + } + } + pub(super) fn suppress_fallback(&mut self, text: &str) -> bool { if !self.fallback.is_empty() && self.fallback.starts_with(text) { self.fallback.drain(..text.len()); @@ -97,6 +130,25 @@ impl ResponseState { } let successful = successful_result(message); let final_text = message["result"].as_str().unwrap_or_default(); + // The result of a cycle a steer aborted: its text was streamed and + // closed before the steered user row, so there is nothing to inject + // and nothing to certify — the turn goes on with the steered reply. + // Recorded, so the transcript keeps the cycle's usage, but not a + // response boundary: `completed` stays as it is and the live buffer + // (possibly already holding the reply's first chunk) is untouched. + if self.take_detached_match(final_text) { + handler + .writer + .record_acp_event_metadata(AcpEventMetadata { + event_kind: Some("response_result".into()), + usage: message.get("usage").cloned(), + content: Some(message.clone()), + origin: handler.out_of_turn_origin().await, + ..Default::default() + }) + .await; + return Ok(()); + } // A result is the final assistant text, not necessarily all text in the // turn. Require agreement with streamed output before accepting it. let (missing, extends_buffer) = if final_text.is_empty() || self.text.ends_with(final_text) diff --git a/crates/acp-client/src/driver/tests/real_bridge.rs b/crates/acp-client/src/driver/tests/real_bridge.rs index fc5f5e4f9..c58fe52b6 100644 --- a/crates/acp-client/src/driver/tests/real_bridge.rs +++ b/crates/acp-client/src/driver/tests/real_bridge.rs @@ -280,7 +280,12 @@ impl MessageWriter for Trace { self.push(Ev::Finalize); } - async fn record_user_message(&self, text: &str, _images: &[PromptImage]) { + async fn record_user_message( + &self, + text: &str, + _images: &[PromptImage], + _attachment_ids: &[String], + ) { self.push(Ev::User(text.to_string())); } @@ -686,7 +691,7 @@ async fn real_bridge_steer_mid_turn_redirects_the_live_turn() { .await; tokio::time::sleep(Duration::from_secs(2)).await; let sent_at = trace.elapsed(); - let outcome = handle.steer(STEER_MESSAGE, &[]).await; + let outcome = handle.steer(STEER_MESSAGE, &[], &[]).await; let answered_at = trace.elapsed(); let turn_still_live = !prompt_done_for_steer.get(); eprintln!( @@ -801,7 +806,7 @@ async fn real_bridge_steer_while_the_bridge_holds_the_turn_for_a_subagent() { tokio::time::sleep(Duration::from_secs(3)).await; let held = !prompt_done_for_steer.get(); eprintln!("[steer-held] bridge still holding the turn 3s after LAUNCHED: {held}"); - let outcome = handle.steer(HELD_STEER_MESSAGE, &[]).await; + let outcome = handle.steer(HELD_STEER_MESSAGE, &[], &[]).await; eprintln!("[steer-held] steer answered {outcome:?} at {:.2}s", trace.elapsed().as_secs_f64()); (held, outcome) }); @@ -873,14 +878,14 @@ async fn real_bridge_interrupt_reprompts_on_the_same_child() { .await; tokio::time::sleep(Duration::from_secs(2)).await; // The caller's routing: steer first, interrupt on refusal. - let steer = handle.steer(INTERRUPT_MESSAGE, &[]).await; + let steer = handle.steer(INTERRUPT_MESSAGE, &[], &[]).await; eprintln!("[interrupt] steer refused as expected: {steer:?}"); assert!(steer .as_ref() .is_err_and(|e| e.contains("does not advertise steering"))); assert_eq!(user_rows(&trace), 0, "a refused steer leaves no row"); let sent_at = trace.elapsed(); - let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[]).await; + let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[], &[]).await; let answered_at = trace.elapsed(); // The re-prompt is live now: the child it runs on. (Read // here, not after the turn — a bridge with no observable @@ -985,7 +990,7 @@ async fn race_steer(label: &str, steer_after_first_chunk: bool) { .await; } let sent_at = trace.elapsed(); - let outcome = tokio::time::timeout(REPLY_TIMEOUT, handle.steer(RACE_MESSAGE, &[])) + let outcome = tokio::time::timeout(REPLY_TIMEOUT, handle.steer(RACE_MESSAGE, &[], &[])) .await .expect("a steer at the turn boundary must not hang"); eprintln!( @@ -1002,7 +1007,7 @@ async fn race_steer(label: &str, steer_after_first_chunk: bool) { // the caller to route as a resume. let fallback = tokio::time::timeout( REPLY_TIMEOUT, - handle.interrupt_and_reprompt(RACE_MESSAGE, &[]), + handle.interrupt_and_reprompt(RACE_MESSAGE, &[], &[]), ) .await .expect("the fallback must not hang"); @@ -1128,7 +1133,7 @@ async fn real_bridge_interrupt_while_a_permission_request_is_pending() { } tokio::time::sleep(Duration::from_secs(1)).await; let sent_at = trace.elapsed(); - let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[]).await; + let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[], &[]).await; let answered_at = trace.elapsed(); let pids_during_reprompt = bridge_pids(); eprintln!( @@ -1231,7 +1236,7 @@ async fn race_interrupt(label: &str, after_first_chunk: bool) { let sent_at = trace.elapsed(); let outcome = tokio::time::timeout( REPLY_TIMEOUT, - handle.interrupt_and_reprompt(RACE_MESSAGE, &[]), + handle.interrupt_and_reprompt(RACE_MESSAGE, &[], &[]), ) .await .expect("an interrupt at the turn boundary must not hang"); diff --git a/crates/acp-client/src/driver/tests/response_completion.rs b/crates/acp-client/src/driver/tests/response_completion.rs index f0bc633c4..12a0fdc9e 100644 --- a/crates/acp-client/src/driver/tests/response_completion.rs +++ b/crates/acp-client/src/driver/tests/response_completion.rs @@ -31,11 +31,23 @@ impl ArtifactWriter { } fn recorded_result(&self) -> bool { + self.recorded_results() > 0 + } + + fn recorded_results(&self) -> usize { self.metadata .lock() .unwrap() .iter() - .any(|event| event.event_kind.as_deref() == Some("response_result")) + .filter(|event| event.event_kind.as_deref() == Some("response_result")) + .count() + } + + /// How many results reached the completion boundary (the writer's + /// `should_finish_after_response`). + fn completion_checks(&self) -> usize { + self.acceptance_checks + .load(std::sync::atomic::Ordering::Relaxed) } } @@ -383,6 +395,7 @@ async fn session_loop_finishes_successfully_for_initial_and_background_artifacts .send(QueuedSessionTurn { prompt: "write".into(), images: vec![], + attachment_ids: vec![], reply: Some(reply), control: None, }) @@ -641,12 +654,54 @@ async fn result_of_a_cycle_a_steer_aborted_does_not_echo_the_closed_message() { RowFlow::Append("pong".into()), RowFlow::Finalize, // No second "pong": the aborted cycle's result matched the text - // already streamed. The steered reply starts a row of its own. - RowFlow::Finalize, + // the steer detached. The steered reply starts a row of its own + // (the steer's finalize already closed the previous one). RowFlow::Append("RACE-DONE".into()), ] ); - assert!(writer.recorded_result()); + // Both results are recorded (the aborted cycle's usage is real), but only + // the steered reply's reaches the completion boundary: the aborted cycle + // is not a response of its own. + assert_eq!(writer.recorded_results(), 2); + assert_eq!(writer.completion_checks(), 1); +} + +/// The other order: the steered reply's first chunk lands *before* the +/// aborted cycle's `result`. The detached text is kept apart from the live +/// buffer, so the aborted result still matches it rather than retiring a +/// buffer that now holds the reply's first chunk — which would split the +/// reply into a second row at its next chunk and leave the turn's final +/// result unrecorded. +#[tokio::test] +async fn aborted_cycle_result_arriving_after_the_steered_reply_began_does_not_split_it() { + let writer = Arc::new(ArtifactWriter::default()); + let handler = artifact_handler(&writer); + start_task(&handler, TaskTrackingMode::Raw).await; + feed(&handler, agent_text(Some("msg-1"), "pong")).await; + handler.writer.finalize().await; + handler.response.lock().await.detach_message(); + feed(&handler, agent_text(Some("msg-2"), "RACE-")).await; + feed_sdk_frame(&handler, successful_result("aborted", "pong")).await; + feed(&handler, agent_text(Some("msg-2"), "DONE")).await; + feed_sdk_frame(&handler, successful_result("steered", "RACE-DONE")).await; + + assert_eq!( + writer.flow(), + vec![ + RowFlow::Append("pong".into()), + RowFlow::Finalize, + // One row for the reply, no boundary between its two chunks and + // no echoed "pong". + RowFlow::Append("RACE-".into()), + RowFlow::Append("DONE".into()), + ] + ); + assert_eq!(handler.response.lock().await.text, "RACE-DONE"); + // The final result is still honoured: recorded, and the one that + // completes the response. + assert_eq!(writer.recorded_results(), 2); + assert_eq!(writer.completion_checks(), 1); + assert!(handler.response.lock().await.completed); } #[tokio::test] diff --git a/crates/acp-client/src/driver/tests/turn_control.rs b/crates/acp-client/src/driver/tests/turn_control.rs index c4c1e5c93..b91dddefb 100644 --- a/crates/acp-client/src/driver/tests/turn_control.rs +++ b/crates/acp-client/src/driver/tests/turn_control.rs @@ -7,7 +7,7 @@ use crate::driver::{ reject_queued_turn_control_requests, run_acp_session, steer_outcome, steering_message, steering_supported_from_initialize, HandlerPhase, IncomingSessionUpdate, SessionTurnResult, SteerOutcome, TurnControlHandle, TurnControlReply, TurnControlRequest, TurnStep, - STEERING_META_KEY, STEERING_METHOD, + STEERING_META_KEY, STEERING_METHOD, STEER_ACK_TIMEOUT, }; use agent_client_protocol::schema::v1::{ CancelNotification, InitializeRequest, InitializeResponse, Meta, PromptResponse, @@ -44,6 +44,17 @@ struct FakeAgent { stream_after_permission: Option<&'static str>, permission_outcomes: Mutex>, permission_answered: tokio::sync::Notify, + /// Record `session/cancel` without answering the prompt it cancels: the + /// bridge that takes its time acknowledging (the test answers it). + hold_cancel_ack: bool, + /// Cancel this token right after acknowledging a `session/cancel`: the + /// user's Stop landing together with the ack. + stop_on_cancel: Option, + /// Hold the `_session/steering` responder instead of answering it; the + /// test answers it with [`FakeAgent::finish_steer`]. + hold_steer: bool, + pending_steer: Mutex>, + steer_arrived: tokio::sync::Notify, } impl FakeAgent { @@ -65,6 +76,25 @@ impl FakeAgent { .expect("agent should receive a prompt"); } + async fn await_steer(&self) { + tokio::time::timeout(SETTLE_TIMEOUT, self.steer_arrived.notified()) + .await + .expect("agent should receive a steer"); + } + + /// Answer the `_session/steering` request the agent is holding back. + fn finish_steer(&self, response: serde_json::Value) { + let responder = self + .pending_steer + .lock() + .unwrap() + .take() + .expect("a steer should be pending"); + responder + .respond(response) + .expect("steer response should send"); + } + /// Answer the prompt the agent is holding back. fn finish_prompt(&self, stop_reason: StopReason) { let responder = self @@ -159,8 +189,14 @@ impl FakeAgent { .on_receive_notification( async move |_notification: CancelNotification, _cx| { *cancel_agent.cancels.lock().unwrap() += 1; + if cancel_agent.hold_cancel_ack { + return Ok(()); + } // The bridge aborts the turn and answers its prompt. cancel_agent.finish_prompt(StopReason::Cancelled); + if let Some(stop) = &cancel_agent.stop_on_cancel { + stop.cancel(); + } Ok(()) }, agent_client_protocol::on_receive_notification!(), @@ -176,6 +212,11 @@ impl FakeAgent { .lock() .unwrap() .push(request.params().clone()); + if steer_agent.hold_steer { + *steer_agent.pending_steer.lock().unwrap() = Some(responder); + steer_agent.steer_arrived.notify_one(); + return Ok(()); + } match &steer_response { Some(response) => responder.respond(response.clone()), None => responder @@ -415,6 +456,7 @@ async fn steer_in_a_live_turn_sends_the_extension_request_and_persists_the_user_ .steer( "also run the linter", &[PromptImage::new("aGk=", "image/png")], + &[], ) .await; assert_eq!(outcome, Ok(SteerOutcome::Injected)); @@ -476,7 +518,7 @@ async fn steer_answered_prompt_required_is_an_error_without_a_user_row() { let (step, writer) = live_turn_with(&agent, true, |handle, agent, _writer| async move { agent.await_prompt().await; let error = handle - .steer("too late", &[]) + .steer("too late", &[], &[]) .await .expect_err("promptRequired is not a delivery"); assert!(error.contains("promptRequired"), "got: {error}"); @@ -508,7 +550,7 @@ async fn steer_to_an_agent_without_the_method_is_an_error() { let (step, writer) = live_turn_with(&agent, true, |handle, agent, _writer| async move { agent.await_prompt().await; let error = handle - .steer("hello?", &[]) + .steer("hello?", &[], &[]) .await .expect_err("method-not-found is not a delivery"); assert!(error.contains("Steering request failed"), "got: {error}"); @@ -534,7 +576,7 @@ async fn steer_is_refused_locally_when_the_agent_did_not_advertise_it() { let (step, _writer) = live_turn_with(&agent, false, |handle, agent, _writer| async move { agent.await_prompt().await; let error = handle - .steer("hello?", &[]) + .steer("hello?", &[], &[]) .await .expect_err("no capability, no steer"); assert!( @@ -577,6 +619,33 @@ async fn session_with( Arc, Arc, ) +where + F: FnOnce( + TurnControlHandle, + Arc, + Arc, + Arc, + ) -> Fut, + Fut: Future, +{ + session_with_cancel(agent, background_hold, CancellationToken::new(), drive).await +} + +/// [`session_with`] on a session cancel token the test holds — the user's +/// Stop. +#[allow(clippy::type_complexity)] +async fn session_with_cancel( + agent: &Arc, + background_hold: Option<&BackgroundHoldConfig>, + cancel: CancellationToken, + drive: F, +) -> ( + Result<(AgentRunOutcome, SessionSettleReason), String>, + oneshot::Receiver, + Option>, + Arc, + Arc, +) where F: FnOnce( TurnControlHandle, @@ -589,7 +658,6 @@ where let writer = Arc::new(TranscriptWriter::default()); let handler = Arc::new(transcript_handler(&writer)); let store: Arc = Arc::new(RecordingStore::default()); - let cancel = CancellationToken::new(); let child_exited = CancellationToken::new(); let (prompt_tx, mut prompt_rx) = mpsc::unbounded_channel(); let (_hold_control_tx, mut hold_control_rx) = mpsc::unbounded_channel(); @@ -604,6 +672,7 @@ where .send(QueuedSessionTurn { prompt: "first".to_string(), images: vec![], + attachment_ids: vec![], reply: Some(reply_tx), control: None, }) @@ -693,7 +762,7 @@ async fn interrupt_cancels_the_turn_and_reprompts_on_the_same_connection() { // The agent's cancel handler answers the pending prompt // `cancelled`; the interrupt then sends the new message. handle - .interrupt_and_reprompt("actually, do this instead", &[]) + .interrupt_and_reprompt("actually, do this instead", &[], &[]) .await .expect("the agent acknowledged the cancel"); agent.await_prompt().await; @@ -757,6 +826,267 @@ async fn interrupt_cancels_the_turn_and_reprompts_on_the_same_connection() { ); } +/// The user's Stop lands while the interrupt is waiting for the bridge to +/// acknowledge the cancel. The stop wins: the interrupt fails, the message is +/// not re-prompted and no row is written for it, and the turn settles +/// cancelled as a plain Stop would — instead of a prompt going out that the +/// next select cancels at once, with the caller told it was delivered. +#[tokio::test(flavor = "current_thread")] +async fn stop_during_the_interrupt_ack_wait_fails_the_interrupt_without_a_reprompt() { + let agent = Arc::new(FakeAgent { + advertise_steering: false, + stream_on_prompt: vec!["half an answer"], + hold_cancel_ack: true, + ..FakeAgent::default() + }); + let cancel = CancellationToken::new(); + let stop = cancel.clone(); + + let (result, _reply_rx, stashed_reply, writer, _handler) = session_with_cancel( + &agent, + None, + cancel, + |handle, agent, writer, _handler| async move { + agent.await_prompt().await; + wait_until("the first chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("half an answer".to_string())) + }) + .await; + + let interrupt = tokio::spawn(async move { + handle + .interrupt_and_reprompt("actually, do this instead", &[], &[]) + .await + }); + // The cancel is on the wire and unacknowledged when the Stop lands. + wait_until("the cancel", || agent.cancels() == 1).await; + stop.cancel(); + let outcome = interrupt.await.unwrap(); + let error = outcome.expect_err("the stop pre-empts the interrupt"); + assert!(error.contains("stopped"), "got: {error}"); + + // The bridge gets around to acknowledging; nothing follows it. + agent.finish_prompt(StopReason::Cancelled); + }, + ) + .await; + + assert_eq!( + result, + Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)) + ); + assert_eq!( + agent.prompts(), + ["first"], + "the message was not re-prompted" + ); + assert_eq!( + writer.text_flow(), + vec![WriterCall::Append("half an answer".to_string())], + "no user row for an undelivered message" + ); + assert!(stashed_reply.is_some()); +} + +/// The Stop lands together with the acknowledgment: whichever of the two +/// checks sees it first (the ack wait, or the session loop before it serves +/// the re-prompt), the answer is the same — the message is refused and the +/// session settles cancelled. The fake agent cancels the session token the +/// moment it acknowledges, so the ack and the stop are ready together. +#[tokio::test(flavor = "current_thread")] +async fn stop_landing_with_the_cancel_ack_refuses_the_reprompt() { + let cancel = CancellationToken::new(); + let agent = Arc::new(FakeAgent { + advertise_steering: false, + stream_on_prompt: vec!["half an answer"], + stop_on_cancel: Some(cancel.clone()), + ..FakeAgent::default() + }); + + let (result, _reply_rx, _stashed_reply, writer, _handler) = session_with_cancel( + &agent, + None, + cancel, + |handle, agent, writer, _handler| async move { + agent.await_prompt().await; + wait_until("the first chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("half an answer".to_string())) + }) + .await; + + let error = handle + .interrupt_and_reprompt("actually, do this instead", &[], &[]) + .await + .expect_err("the stop pre-empts the interrupt"); + assert!(error.contains("stopped"), "got: {error}"); + assert_eq!(agent.cancels(), 1); + }, + ) + .await; + + assert_eq!( + result, + Ok((AgentRunOutcome::Cancelled, SessionSettleReason::Cancelled)) + ); + assert_eq!( + agent.prompts(), + ["first"], + "the message was not re-prompted" + ); + assert!( + !writer + .calls() + .iter() + .any(|call| matches!(call, WriterCall::UserMessage(_))), + "no user row for an undelivered message: {:?}", + writer.calls() + ); +} + +/// The caller's attachment ids ride along with the message and come back to +/// the writer with the user row it persists — on the steer, and on the +/// interrupt's re-prompt. This crate sends image bytes, so without the ids +/// the row could never reference the stored attachments. +#[tokio::test(flavor = "current_thread")] +async fn attachment_ids_reach_the_writer_with_the_user_row_on_both_live_paths() { + let agent = Arc::new(FakeAgent { + advertise_steering: true, + steer_response: Some(serde_json::json!({ "outcome": "injected" })), + stream_on_prompt: vec!["working", "redoing"], + ..FakeAgent::default() + }); + let image = PromptImage::new("aGk=", "image/png"); + + let (result, _reply_rx, _stashed_reply, writer, _handler) = + session_with(&agent, None, |handle, agent, writer, _handler| async move { + agent.await_prompt().await; + wait_until("the first chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("working".to_string())) + }) + .await; + + let outcome = handle + .steer( + "look at this", + std::slice::from_ref(&image), + &["img-1".to_string()], + ) + .await; + assert_eq!(outcome, Ok(SteerOutcome::Injected)); + + handle + .interrupt_and_reprompt( + "and this instead", + std::slice::from_ref(&image), + &["img-2".to_string(), "img-3".to_string()], + ) + .await + .expect("the cancel is acknowledged"); + wait_until("the re-prompt", || agent.prompts().len() == 2).await; + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + assert_eq!( + result, + Ok((AgentRunOutcome::Completed, SessionSettleReason::Immediate)) + ); + let user_rows: Vec = writer + .calls() + .into_iter() + .filter(|call| { + matches!( + call, + WriterCall::UserMessage(_) | WriterCall::UserAttachments(_) + ) + }) + .collect(); + assert_eq!( + user_rows, + vec![ + WriterCall::UserMessage("look at this".to_string()), + WriterCall::UserAttachments(vec!["img-1".to_string()]), + WriterCall::UserMessage("and this instead".to_string()), + WriterCall::UserAttachments(vec!["img-2".to_string(), "img-3".to_string()]), + ] + ); +} + +/// A bridge that takes the `_session/steering` frame and never answers while +/// the turn runs on must not wedge the caller: the steer fails after +/// `STEER_ACK_TIMEOUT` so the caller can fall back to an interrupt. An answer +/// that arrives after that — `injected`, even — persists no user row: the +/// caller's fallback has by then persisted the message itself, and a second +/// row would duplicate it. Runs on tokio's paused clock: the fake agent is +/// in-process, so the 30 s elapse the moment nothing else can run. +#[tokio::test(flavor = "current_thread", start_paused = true)] +async fn steer_unanswered_within_the_timeout_fails_and_a_late_answer_persists_nothing() { + let agent = Arc::new(FakeAgent { + advertise_steering: true, + steer_response: Some(serde_json::json!({ "outcome": "injected" })), + hold_steer: true, + stream_on_prompt: vec!["working on it"], + ..FakeAgent::default() + }); + + let (step, _writer) = live_turn_with(&agent, true, |handle, agent, writer| async move { + agent.await_prompt().await; + wait_until("the first chunk", || { + writer + .text_flow() + .contains(&WriterCall::Append("working on it".to_string())) + }) + .await; + + let started = tokio::time::Instant::now(); + let error = handle + .steer("also run the linter", &[], &[]) + .await + .expect_err("an unanswered steer times out"); + assert!(error.contains("did not answer the steer"), "got: {error}"); + assert!( + started.elapsed() >= STEER_ACK_TIMEOUT, + "failed before the timeout: {:?}", + started.elapsed() + ); + assert!( + agent.pending_prompt.lock().unwrap().is_some(), + "the turn is still live" + ); + assert_eq!( + writer.text_flow(), + vec![WriterCall::Append("working on it".to_string())], + "a failed steer leaves no user row" + ); + + // The bridge finally answers `injected`. Too late: the caller has + // moved on, so nothing is persisted and nothing is replied to. + agent.await_steer().await; + agent.finish_steer(serde_json::json!({ "outcome": "injected" })); + // Let the late answer's callback run. + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!( + writer.text_flow(), + vec![WriterCall::Append("working on it".to_string())], + "a late answer persists nothing" + ); + agent.finish_prompt(StopReason::EndTurn); + }) + .await; + + assert!(matches!( + step, + Ok(TurnStep::Settled(AgentRunOutcome::Completed)) + )); + assert_eq!(agent.steers().len(), 1); +} + /// A control request queued next to an already-queued `prompt()` (both land /// while setup is still running, say) does not jump it: the prompt is the /// turn, and the request is served inside it — here an interrupt, which @@ -776,7 +1106,7 @@ async fn a_prompt_queued_before_a_control_request_goes_first() { |handle, agent, _writer, _handler| async move { // No await before this: it is queued while setup is in flight, // behind the prompt `session_with` queued before connecting. - let outcome = handle.interrupt_and_reprompt("second", &[]).await; + let outcome = handle.interrupt_and_reprompt("second", &[], &[]).await; assert_eq!(outcome, Ok(())); wait_until("both prompts", || agent.prompts().len() == 2).await; assert_eq!(agent.prompts(), ["first", "second"]); @@ -828,7 +1158,7 @@ async fn interrupt_during_a_hold_is_served_as_a_new_turn_without_a_cancel() { // Nothing is in flight to steer: the message becomes a turn of // its own, with no `_session/steering` frame on the wire. - let outcome = handle.steer("while you wait", &[]).await; + let outcome = handle.steer("while you wait", &[], &[]).await; assert_eq!(outcome, Ok(SteerOutcome::NewTurnStarted)); agent.await_prompt().await; assert_eq!(agent.prompts(), ["first", "while you wait"]); @@ -845,7 +1175,7 @@ async fn interrupt_during_a_hold_is_served_as_a_new_turn_without_a_cancel() { // And an interrupt with nothing to interrupt is the same. wait_until("the second hold", || holding(&handler)).await; handle - .interrupt_and_reprompt("and this", &[]) + .interrupt_and_reprompt("and this", &[], &[]) .await .expect("nothing to cancel, served as a turn"); agent.await_prompt().await; @@ -894,6 +1224,7 @@ async fn holding_wait_turns_a_steer_into_a_new_turn_without_rejecting_it_at_boun .send(TurnControlRequest::Steer { prompt: "queued early".to_string(), images: vec![], + attachment_ids: vec![], reply: early_reply_tx, }) .expect("queue should accept"); @@ -958,7 +1289,7 @@ async fn holding_wait_turns_a_steer_into_a_new_turn_without_rejecting_it_at_boun turn.control.expect("control reply").turn_started(); turn.prompt }; - let (prompt, answer) = tokio::join!(served, handle.interrupt_and_reprompt("now", &[])); + let (prompt, answer) = tokio::join!(served, handle.interrupt_and_reprompt("now", &[], &[])); assert_eq!(prompt, "now"); assert_eq!(answer, Ok(())); } @@ -982,6 +1313,7 @@ async fn teardown_answers_every_queued_message_and_the_handle_fails_fast_after() .send(TurnControlRequest::Steer { prompt: "a".to_string(), images: vec![], + attachment_ids: vec![], reply: steer_reply_tx, }) .expect("queue should accept"); @@ -989,6 +1321,7 @@ async fn teardown_answers_every_queued_message_and_the_handle_fails_fast_after() .send(TurnControlRequest::Interrupt { prompt: "b".to_string(), images: vec![], + attachment_ids: vec![], reply: interrupt_reply_tx, }) .expect("queue should accept"); @@ -1022,12 +1355,12 @@ async fn teardown_answers_every_queued_message_and_the_handle_fails_fast_after() // After settle the handle refuses without queueing. let error = handle - .steer("c", &[]) + .steer("c", &[], &[]) .await .expect_err("connection is gone"); assert!(error.contains("no longer running"), "got: {error}"); let error = handle - .interrupt_and_reprompt("d", &[]) + .interrupt_and_reprompt("d", &[], &[]) .await .expect_err("connection is gone"); assert!(error.contains("no longer running"), "got: {error}"); @@ -1045,7 +1378,7 @@ async fn a_turn_control_handle_maps_channel_failures_to_connection_state() { let request = turn_control_rx.recv().await.expect("request should arrive"); drop(request); }; - let (steered, ()) = tokio::join!(handle.steer("a", &[]), drop_unanswered); + let (steered, ()) = tokio::join!(handle.steer("a", &[], &[]), drop_unanswered); let error = steered.expect_err("a dropped reply is not an answer"); assert!(error.contains("closed before"), "got: {error}"); @@ -1053,7 +1386,7 @@ async fn a_turn_control_handle_maps_channel_failures_to_connection_state() { // alive flag still says. drop(turn_control_rx); let error = handle - .interrupt_and_reprompt("b", &[]) + .interrupt_and_reprompt("b", &[], &[]) .await .expect_err("no connection to serve the message"); assert!(error.contains("no longer running"), "got: {error}"); @@ -1090,7 +1423,7 @@ async fn interrupt_resolves_the_turns_pending_permission_request_before_reprompt let started = tokio::time::Instant::now(); handle - .interrupt_and_reprompt("actually, do this instead", &[]) + .interrupt_and_reprompt("actually, do this instead", &[], &[]) .await .expect("the cancel is acknowledged with the permission request pending"); assert!( From 82a3409f51a3ca5deeca958e61799d8b3e7361c0 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Thu, 29 Oct 2026 16:13:09 +1100 Subject: [PATCH 7/7] fix(staged,acp-client): do not interrupt after an unacknowledged steer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to review 73a96644-306e-47dd-89fb-a95672507b02 (round 2 of message-to-running-session), two items. Warning — steer timeout → interrupt fallback could double-deliver. When `steer()` ran out `STEER_ACK_TIMEOUT` the bridge may still have injected the message, yet `continue_session_for_store` fell back to `interrupt_and_reprompt` with the same text: the agent read the instruction twice against a transcript with one user row, and a slow bridge cost the caller 30 s + 30 s. `TurnControlHandle::steer` now returns `Result` with two variants: `Refused(String)` for every failure that means "not delivered" (no steering advertised, a non-`injected` answer, request not sendable, connection gone, teardown drain) and `Unacknowledged { timeout }` for the ack timer alone. `deliver_to_running_session` falls back to the interrupt only on `Refused`; `Unacknowledged` returns `Blocked` with "steer not acknowledged within 30s; the message may still have been delivered — check the transcript before retrying". `STEER_ACK_TIMEOUT` is unchanged. `SteerError` is exported from the crate root, implements `Display`/`Error`, and the `scripted` stand-ins take it. Test: `continuing_does_not_interrupt_after_an_unacknowledged_steer` (blocked with the reason, the stand-in saw only the steer, row still `running`). The acp-client steer tests assert the variant (the timeout test pins `Unacknowledged { STEER_ACK_TIMEOUT }`; refusals match `Refused`), and the real-bridge scenarios compare via `Display`. Information — the `record_user_message` doc in agent/writer.rs still said the row carries no `image_ids` and the caller attaches them; it now describes the current behaviour: ids threaded through `TurnControlRequest`, persisted on the row, images re-associated via `set_images_session_id`. Verified: cargo build; `cargo test -p acp-client --lib` (204 passed); src-tauri `--lib` continue/route/project_mcp/writer tests (49 passed); clippy `-D warnings` on both crates; rustfmt clean. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/agent/writer.rs | 11 +- apps/staged/src-tauri/src/project_mcp.rs | 6 +- apps/staged/src-tauri/src/session_commands.rs | 87 ++++++++++-- crates/acp-client/src/driver.rs | 130 ++++++++++++------ .../src/driver/tests/real_bridge.rs | 6 +- .../src/driver/tests/turn_control.rs | 54 ++++++-- crates/acp-client/src/lib.rs | 2 +- 7 files changed, 222 insertions(+), 74 deletions(-) diff --git a/apps/staged/src-tauri/src/agent/writer.rs b/apps/staged/src-tauri/src/agent/writer.rs index 1cf3d5c6d..c86919e24 100644 --- a/apps/staged/src-tauri/src/agent/writer.rs +++ b/apps/staged/src-tauri/src/agent/writer.rs @@ -175,9 +175,14 @@ impl MessageWriter { /// as well costs nothing (finalize is idempotent) and keeps the ordering /// a property of this method rather than of its caller. /// - /// The `images` are the blocks the connection sent, not stored images — - /// acp-client has no image ids — so the row carries no `image_ids`. A - /// caller that owns stored images attaches them to the row it gets back. + /// The `images` are the blocks the connection sent; `image_ids` are the + /// stored attachments they were built from, threaded through the + /// connection's `TurnControlRequest` by the caller that loaded them. They + /// are persisted on the row as its `image_ids` and the images are + /// re-associated with this session (`set_images_session_id`), as + /// `start_session` does for a first prompt. The two lists normally pair + /// up one to one; a shortfall (an attachment whose file could not be + /// read still keeps its id) is logged and otherwise harmless. pub async fn record_user_message( &self, text: &str, diff --git a/apps/staged/src-tauri/src/project_mcp.rs b/apps/staged/src-tauri/src/project_mcp.rs index bd91624d1..ad806db4d 100644 --- a/apps/staged/src-tauri/src/project_mcp.rs +++ b/apps/staged/src-tauri/src/project_mcp.rs @@ -1277,7 +1277,7 @@ mod tests { AcpConfigSelection, AcpConfigValueSelection, Branch, MessageRole, Project, ProjectNote, Session, SessionMessage, SessionStatus, Store, }; - use acp_client::{SteerOutcome, TurnControlHandle}; + use acp_client::{SteerError, SteerOutcome, TurnControlHandle}; use std::path::Path; use std::sync::Arc; @@ -1751,7 +1751,9 @@ mod tests { async fn continuing_a_running_session_by_interrupting_never_leaves_running() { let f = ContinueFixture::running_note_session(); let (turn_control, served) = TurnControlHandle::scripted( - Err("agent does not advertise steering".to_string()), + Err(SteerError::Refused( + "agent does not advertise steering".to_string(), + )), Ok(()), ); f.registry diff --git a/apps/staged/src-tauri/src/session_commands.rs b/apps/staged/src-tauri/src/session_commands.rs index 1575fb91a..9ed67375f 100644 --- a/apps/staged/src-tauri/src/session_commands.rs +++ b/apps/staged/src-tauri/src/session_commands.rs @@ -2121,11 +2121,18 @@ async fn wait_for_turn_control_handle( /// Reach a running session's agent over its connection: steer, and when the /// agent refuses (no steering advertised, `promptRequired` because the turn -/// ended under the request, `failed`, no answer in time), interrupt the turn -/// and re-prompt. `Err` only when the interrupt fails too — the caller -/// decides whether that is a dead connection or a blocked delivery. The -/// `image_ids` ride along as the row's attachment ids, like `start_session` -/// records them, whether or not every file could be read. +/// ended under the request, `failed`), interrupt the turn and re-prompt. A +/// refusal is the only failure that falls back: it means the message did not +/// reach the agent, so sending the same text again cannot double-deliver it. +/// A steer the agent never answered (`SteerError::Unacknowledged`) is +/// different — the bridge may have injected the message and simply not said +/// so — and is reported `Blocked` with the ambiguity spelled out, rather +/// than re-sent as an interrupt that would have the agent read the same +/// instruction twice (and the caller wait out a second ack timeout). `Err` +/// only when the interrupt fails too — the caller decides whether that is a +/// dead connection or a blocked delivery. The `image_ids` ride along as the +/// row's attachment ids, like `start_session` records them, whether or not +/// every file could be read. async fn deliver_to_running_session( store: &Arc, handle: &acp_client::TurnControlHandle, @@ -2147,7 +2154,19 @@ async fn deliver_to_running_session( let steer_refusal = match handle.steer(&prompt, &attachments.images, image_ids).await { Ok(acp_client::SteerOutcome::Injected) => return Ok(ContinueOutcome::Injected), Ok(acp_client::SteerOutcome::NewTurnStarted) => return Ok(ContinueOutcome::StartedNewTurn), - Err(reason) => reason, + Err(acp_client::SteerError::Unacknowledged { timeout }) => { + log::warn!( + "continue_session: steer not acknowledged within {}s; not interrupting, the \ + message may already have been delivered", + timeout.as_secs() + ); + return Ok(ContinueOutcome::blocked(format!( + "steer not acknowledged within {}s; the message may still have been delivered \ + — check the transcript before retrying", + timeout.as_secs() + ))); + } + Err(acp_client::SteerError::Refused(reason)) => reason, }; log::info!("continue_session: steer refused ({steer_refusal}); interrupting the turn instead"); handle @@ -8260,7 +8279,9 @@ mod tests { let registry = test_registry(); let _registration = registry.register_external(&session.id); let (handle, served) = acp_client::TurnControlHandle::scripted( - Err("steer not injected: promptRequired".to_string()), + Err(acp_client::SteerError::Refused( + "steer not injected: promptRequired".to_string(), + )), Ok(()), ); registry.set_turn_control_handle(&session.id, handle); @@ -8279,6 +8300,46 @@ mod tests { ); } + /// A steer the agent never answered is not a refusal: the bridge may have + /// injected the message and just not said so, so the body must not + /// re-send the same text as an interrupt (the agent would read it twice, + /// and the caller would wait out a second ack timeout). It answers + /// blocked, says the message may have landed, and sends nothing else. + #[tokio::test(flavor = "current_thread")] + async fn continuing_does_not_interrupt_after_an_unacknowledged_steer() { + let store = Arc::new(Store::in_memory().unwrap()); + let session = create_session_with_status(&store, "prompt", store::SessionStatus::Running); + let registry = test_registry(); + let _registration = registry.register_external(&session.id); + let (handle, served) = acp_client::TurnControlHandle::scripted( + Err(acp_client::SteerError::Unacknowledged { + timeout: Duration::from_secs(30), + }), + Ok(()), + ); + registry.set_turn_control_handle(&session.id, handle); + + let outcome = + delivered(continue_through_shared_body(®istry, &store, &session.id, "redo").await); + + match outcome { + ContinueOutcome::Blocked { reason } => { + assert!(reason.contains("not acknowledged within 30s"), "{reason}"); + assert!(reason.contains("may still have been delivered"), "{reason}"); + } + other => panic!("expected blocked, got {other:?}"), + } + assert_eq!( + *served.lock().unwrap(), + vec!["steer:redo".to_string()], + "an unacknowledged steer must not be followed by an interrupt" + ); + assert_eq!( + store.get_session(&session.id).unwrap().unwrap().status, + store::SessionStatus::Running + ); + } + /// Both deliveries refused on a connection that is still up (a bridge /// that never acknowledged the cancel): blocked with the agent's reason, /// the row untouched. @@ -8289,7 +8350,9 @@ mod tests { let registry = test_registry(); let _registration = registry.register_external(&session.id); let (handle, _) = acp_client::TurnControlHandle::scripted( - Err("agent does not advertise steering".to_string()), + Err(acp_client::SteerError::Refused( + "agent does not advertise steering".to_string(), + )), Err("agent did not acknowledge the cancel".to_string()), ); registry.set_turn_control_handle(&session.id, handle); @@ -8678,7 +8741,9 @@ mod tests { let registry = test_registry(); let _registration = registry.register_external(&session.id); let (handle, served) = acp_client::TurnControlHandle::scripted_then_gone( - Err("connection closed".to_string()), + Err(acp_client::SteerError::Refused( + "connection closed".to_string(), + )), Err("ACP connection is no longer running".to_string()), ); let script = { @@ -8727,7 +8792,9 @@ mod tests { let registry = test_registry(); let _registration = registry.register_external(&session.id); let (handle, served) = acp_client::TurnControlHandle::scripted( - Err("agent does not advertise steering".to_string()), + Err(acp_client::SteerError::Refused( + "agent does not advertise steering".to_string(), + )), Ok(()), ); registry.set_turn_control_handle(&session.id, handle); diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index 8c6af0b81..6149be48c 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -1482,6 +1482,40 @@ pub enum SteerOutcome { NewTurnStarted, } +/// Why [`TurnControlHandle::steer`] did not deliver the message. The split +/// matters to the caller's fallback: a `Refused` steer is known not to have +/// reached the agent, so sending the same text again by another route cannot +/// double-deliver it; an `Unacknowledged` one may have — the bridge took the +/// frame and never answered — so the caller must not re-send blindly. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SteerError { + /// The agent did not answer the `_session/steering` request within + /// `timeout` ([`STEER_ACK_TIMEOUT`]). Whether the message was injected is + /// unknown; no user row was persisted for it, and an answer landing after + /// this persists nothing either. + Unacknowledged { timeout: Duration }, + /// The steer did not reach the agent's turn and the message was not + /// delivered: the agent does not advertise steering, it answered anything + /// but `injected` (`promptRequired`, `failed`, `startedNewTurn`), the + /// request could not be built or sent, or the connection is gone. + Refused(String), +} + +impl std::fmt::Display for SteerError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Unacknowledged { timeout } => write!( + f, + "Agent did not answer the steer within {}s", + timeout.as_secs() + ), + Self::Refused(reason) => f.write_str(reason), + } + } +} + +impl std::error::Error for SteerError {} + /// A request queued onto a [`SessionConnection`]'s turn-control channel by a /// [`TurnControlHandle`]: a message that must reach the agent now, not after /// the turn in flight ends. @@ -1491,7 +1525,7 @@ enum TurnControlRequest { prompt: String, images: Vec, attachment_ids: Vec, - reply: oneshot::Sender>, + reply: oneshot::Sender>, }, /// `session/cancel` the in-flight turn, then send the message as the next /// `session/prompt` on the same connection. @@ -1545,7 +1579,7 @@ impl TurnControlRequest { fn reject(self, reason: String) { match self { Self::Steer { reply, .. } => { - let _ = reply.send(Err(reason)); + let _ = reply.send(Err(SteerError::Refused(reason))); } Self::Interrupt { reply, .. } => { let _ = reply.send(Err(reason)); @@ -1557,7 +1591,7 @@ impl TurnControlRequest { /// The reply half of a [`TurnControlRequest`] whose message became a /// [`QueuedSessionTurn`] of its own. enum TurnControlReply { - Steer(oneshot::Sender>), + Steer(oneshot::Sender>), Interrupt(oneshot::Sender>), } @@ -1578,7 +1612,7 @@ impl TurnControlReply { fn reject(self, reason: String) { match self { Self::Steer(reply) => { - let _ = reply.send(Err(reason)); + let _ = reply.send(Err(SteerError::Refused(reason))); } Self::Interrupt(reply) => { let _ = reply.send(Err(reason)); @@ -1655,30 +1689,35 @@ impl TurnControlHandle { /// `Ok(Injected)` means the agent took it into the live turn. With no turn /// in flight (the connection holding for background work, or idle) the /// message is sent as a fresh `session/prompt` on the same connection and - /// the answer is `Ok(NewTurnStarted)`. `Err` when the agent does not - /// advertise steering, answered anything but `injected` (`promptRequired` - /// when the turn ended under the request, `failed`, `startedNewTurn`), - /// did not answer within `STEER_ACK_TIMEOUT`, or the connection is gone; - /// callers then fall back to [`TurnControlHandle::interrupt_and_reprompt`]. - /// The user message is persisted (see - /// [`MessageWriter::record_user_message`], with `attachment_ids`) only on - /// `Ok` — an answer that arrives after the timeout persists nothing. + /// the answer is `Ok(NewTurnStarted)`. [`SteerError::Refused`] when the + /// agent does not advertise steering, answered anything but `injected` + /// (`promptRequired` when the turn ended under the request, `failed`, + /// `startedNewTurn`), or the connection is gone — the message did not + /// reach the agent, and callers fall back to + /// [`TurnControlHandle::interrupt_and_reprompt`] with the same text. + /// [`SteerError::Unacknowledged`] when the agent did not answer within + /// `STEER_ACK_TIMEOUT`: the bridge may still have injected the message, + /// so callers must not re-send it by another route. The user message is + /// persisted (see [`MessageWriter::record_user_message`], with + /// `attachment_ids`) only on `Ok` — an answer that arrives after the + /// timeout persists nothing. pub async fn steer( &self, prompt: &str, images: &[PromptImage], attachment_ids: &[String], - ) -> Result { + ) -> Result { let (reply_tx, reply_rx) = oneshot::channel(); self.send(TurnControlRequest::Steer { prompt: prompt.to_string(), images: images.to_vec(), attachment_ids: attachment_ids.to_vec(), reply: reply_tx, - })?; - reply_rx - .await - .map_err(|_| "ACP connection closed before the steer resolved".to_string())? + }) + .map_err(SteerError::Refused)?; + reply_rx.await.map_err(|_| { + SteerError::Refused("ACP connection closed before the steer resolved".to_string()) + })? } /// Cancel the agent's in-flight turn (`session/cancel`) and send the @@ -1748,7 +1787,7 @@ impl TurnControlHandle { /// agent. Needs a tokio runtime: the stand-in is a task. #[doc(hidden)] pub fn scripted( - steer: Result, + steer: Result, interrupt: Result<(), String>, ) -> (Self, Arc>>) { Self::scripted_connection(steer, interrupt, false) @@ -1761,14 +1800,14 @@ impl TurnControlHandle { /// `false` from the first answer on. #[doc(hidden)] pub fn scripted_then_gone( - steer: Result, + steer: Result, interrupt: Result<(), String>, ) -> (Self, Arc>>) { Self::scripted_connection(steer, interrupt, true) } fn scripted_connection( - steer: Result, + steer: Result, interrupt: Result<(), String>, gone_after_first_answer: bool, ) -> (Self, Arc>>) { @@ -3321,8 +3360,9 @@ fn steering_message( /// extension) surfaces as such. fn steer_outcome( result: Result, -) -> Result { - let response = result.map_err(|e| format!("Steering request failed: {e:?}"))?; +) -> Result { + let response = + result.map_err(|e| SteerError::Refused(format!("Steering request failed: {e:?}")))?; match response.get("outcome").and_then(serde_json::Value::as_str) { Some("injected") => Ok(SteerOutcome::Injected), Some(outcome) => { @@ -3331,13 +3371,13 @@ fn steer_outcome( .and_then(serde_json::Value::as_str) .map(|reason| format!(" ({reason})")) .unwrap_or_default(); - Err(format!( + Err(SteerError::Refused(format!( "Steering not injected: agent answered '{outcome}'{reason}" - )) + ))) } - None => Err(format!( + None => Err(SteerError::Refused(format!( "Steering not injected: malformed response {response}" - )), + ))), } } @@ -5108,10 +5148,13 @@ fn serve_async_task_stop( /// /// The answer is bounded by [`STEER_ACK_TIMEOUT`]: the caller's reply sender /// sits in a slot that the response callback and a timer race for, and only -/// the one that takes it answers. A `_session/steering` answer that lands -/// after the timer has answered `Err` is logged and otherwise ignored — in -/// particular it persists no user row, because the caller has by then fallen -/// back to an interrupt that persists the message itself. +/// the one that takes it answers. The timer answers +/// [`SteerError::Unacknowledged`] — the one steer failure that does not mean +/// "not delivered", so the caller knows not to re-send the text by another +/// route. A `_session/steering` answer that lands after the timer has +/// answered is logged and otherwise ignored — in particular it persists no +/// user row, because the caller has by then reported the message as possibly +/// delivered and left the transcript to the user to check. #[allow(clippy::too_many_arguments)] fn serve_steer( connection: &ConnectionTo, @@ -5121,12 +5164,12 @@ fn serve_steer( prompt: String, images: Vec, attachment_ids: Vec, - reply: oneshot::Sender>, + reply: oneshot::Sender>, ) { if !setup.steering_supported { - let _ = reply.send(Err( + let _ = reply.send(Err(SteerError::Refused( "Steering not available: agent does not advertise steering".to_string(), - )); + ))); return; } log::info!("ACP session {our_session_id}: steering the turn in flight"); @@ -5135,7 +5178,9 @@ fn serve_steer( let message = match steering_message(&setup.agent_session_id, content_blocks) { Ok(message) => message, Err(e) => { - let _ = reply.send(Err(format!("Failed to build steering request: {e:?}"))); + let _ = reply.send(Err(SteerError::Refused(format!( + "Failed to build steering request: {e:?}" + )))); return; } }; @@ -5188,10 +5233,9 @@ fn serve_steer( "ACP session {timed_out_session_id}: agent did not answer the steer within {}s", STEER_ACK_TIMEOUT.as_secs() ); - let _ = reply.send(Err(format!( - "Agent did not answer the steer within {}s", - STEER_ACK_TIMEOUT.as_secs() - ))); + let _ = reply.send(Err(SteerError::Unacknowledged { + timeout: STEER_ACK_TIMEOUT, + })); } Ok(()) }); @@ -5218,10 +5262,12 @@ fn serve_steer( const INTERRUPT_ACK_TIMEOUT: Duration = Duration::from_secs(30); /// How long a steer waits for the agent's `_session/steering` answer before -/// the request fails and the caller falls back to an interrupt. The same -/// bound as [`INTERRUPT_ACK_TIMEOUT`], for the same reason: a bridge that -/// takes the frame and never answers while the turn runs on must not wedge -/// the caller's tool call (which has no keepalive of its own). +/// the request fails as [`SteerError::Unacknowledged`]. The same bound as +/// [`INTERRUPT_ACK_TIMEOUT`], for the same reason: a bridge that takes the +/// frame and never answers while the turn runs on must not wedge the caller's +/// tool call (which has no keepalive of its own). Expiring is not a licence +/// to re-send — the bridge may have injected the message anyway — which is +/// why the caller gets a distinct error rather than a refusal. const STEER_ACK_TIMEOUT: Duration = Duration::from_secs(30); /// How long the holding wait gives the bridge to process a `session/cancel` diff --git a/crates/acp-client/src/driver/tests/real_bridge.rs b/crates/acp-client/src/driver/tests/real_bridge.rs index c58fe52b6..1dc05d4bc 100644 --- a/crates/acp-client/src/driver/tests/real_bridge.rs +++ b/crates/acp-client/src/driver/tests/real_bridge.rs @@ -721,7 +721,7 @@ async fn real_bridge_steer_mid_turn_redirects_the_live_turn() { ); assert!(outcome .as_ref() - .is_err_and(|e| e.contains("does not advertise steering"))); + .is_err_and(|e| e.to_string().contains("does not advertise steering"))); assert_eq!(user_rows(&session.trace), 0); session.finish().await; return; @@ -832,7 +832,7 @@ async fn real_bridge_steer_while_the_bridge_holds_the_turn_for_a_subagent() { ); match &outcome { Ok(SteerOutcome::NewTurnStarted) => {} - Err(e) if e.contains("no longer running") => { + Err(e) if e.to_string().contains("no longer running") => { assert_eq!(user_rows(&session.trace), 0); session.finish().await; return; @@ -882,7 +882,7 @@ async fn real_bridge_interrupt_reprompts_on_the_same_child() { eprintln!("[interrupt] steer refused as expected: {steer:?}"); assert!(steer .as_ref() - .is_err_and(|e| e.contains("does not advertise steering"))); + .is_err_and(|e| e.to_string().contains("does not advertise steering"))); assert_eq!(user_rows(&trace), 0, "a refused steer leaves no row"); let sent_at = trace.elapsed(); let outcome = handle.interrupt_and_reprompt(INTERRUPT_MESSAGE, &[], &[]).await; diff --git a/crates/acp-client/src/driver/tests/turn_control.rs b/crates/acp-client/src/driver/tests/turn_control.rs index b91dddefb..dd795133a 100644 --- a/crates/acp-client/src/driver/tests/turn_control.rs +++ b/crates/acp-client/src/driver/tests/turn_control.rs @@ -6,7 +6,7 @@ use super::*; use crate::driver::{ reject_queued_turn_control_requests, run_acp_session, steer_outcome, steering_message, steering_supported_from_initialize, HandlerPhase, IncomingSessionUpdate, SessionTurnResult, - SteerOutcome, TurnControlHandle, TurnControlReply, TurnControlRequest, TurnStep, + SteerError, SteerOutcome, TurnControlHandle, TurnControlReply, TurnControlRequest, TurnStep, STEERING_META_KEY, STEERING_METHOD, STEER_ACK_TIMEOUT, }; use agent_client_protocol::schema::v1::{ @@ -376,11 +376,17 @@ fn steer_outcome_only_counts_injected_as_delivered() { (serde_json::json!({}), "malformed"), ] { let error = steer_outcome(Ok(response)).expect_err("not a delivery"); - assert!(error.contains(needle), "got: {error}"); + assert!( + matches!(&error, SteerError::Refused(reason) if reason.contains(needle)), + "got: {error}" + ); } let error = steer_outcome(Err(agent_client_protocol::Error::method_not_found())) .expect_err("an agent without the extension is not a delivery"); - assert!(error.contains("Steering request failed"), "got: {error}"); + assert!( + matches!(&error, SteerError::Refused(reason) if reason.contains("Steering request failed")), + "got: {error}" + ); } // ============================================================================= @@ -521,8 +527,11 @@ async fn steer_answered_prompt_required_is_an_error_without_a_user_row() { .steer("too late", &[], &[]) .await .expect_err("promptRequired is not a delivery"); - assert!(error.contains("promptRequired"), "got: {error}"); - assert!(error.contains("noRunningTurn"), "got: {error}"); + let SteerError::Refused(reason) = &error else { + panic!("a bridge answer is a refusal, not a timeout: {error:?}"); + }; + assert!(reason.contains("promptRequired"), "got: {error}"); + assert!(reason.contains("noRunningTurn"), "got: {error}"); agent.finish_prompt(StopReason::EndTurn); }) .await; @@ -553,7 +562,10 @@ async fn steer_to_an_agent_without_the_method_is_an_error() { .steer("hello?", &[], &[]) .await .expect_err("method-not-found is not a delivery"); - assert!(error.contains("Steering request failed"), "got: {error}"); + assert!( + matches!(&error, SteerError::Refused(reason) if reason.contains("Steering request failed")), + "got: {error}" + ); agent.finish_prompt(StopReason::EndTurn); }) .await; @@ -580,7 +592,7 @@ async fn steer_is_refused_locally_when_the_agent_did_not_advertise_it() { .await .expect_err("no capability, no steer"); assert!( - error.contains("does not advertise steering"), + matches!(&error, SteerError::Refused(reason) if reason.contains("does not advertise steering")), "got: {error}" ); agent.finish_prompt(StopReason::EndTurn); @@ -1049,7 +1061,17 @@ async fn steer_unanswered_within_the_timeout_fails_and_a_late_answer_persists_no .steer("also run the linter", &[], &[]) .await .expect_err("an unanswered steer times out"); - assert!(error.contains("did not answer the steer"), "got: {error}"); + // Distinct from a refusal: the caller must not re-send the text. + assert_eq!( + error, + SteerError::Unacknowledged { + timeout: STEER_ACK_TIMEOUT + } + ); + assert!( + error.to_string().contains("did not answer the steer"), + "got: {error}" + ); assert!( started.elapsed() >= STEER_ACK_TIMEOUT, "failed before the timeout: {:?}", @@ -1334,9 +1356,9 @@ async fn teardown_answers_every_queued_message_and_the_handle_fails_fast_after() .await .expect("every queued request is answered"); assert!( - steer - .as_ref() - .is_err_and(|e| e.contains("connection closed")), + steer.as_ref().is_err_and( + |e| matches!(e, SteerError::Refused(reason) if reason.contains("connection closed")) + ), "got: {steer:?}" ); let interrupt = interrupt_reply_rx @@ -1358,7 +1380,10 @@ async fn teardown_answers_every_queued_message_and_the_handle_fails_fast_after() .steer("c", &[], &[]) .await .expect_err("connection is gone"); - assert!(error.contains("no longer running"), "got: {error}"); + assert!( + matches!(&error, SteerError::Refused(reason) if reason.contains("no longer running")), + "got: {error}" + ); let error = handle .interrupt_and_reprompt("d", &[], &[]) .await @@ -1380,7 +1405,10 @@ async fn a_turn_control_handle_maps_channel_failures_to_connection_state() { }; let (steered, ()) = tokio::join!(handle.steer("a", &[], &[]), drop_unanswered); let error = steered.expect_err("a dropped reply is not an answer"); - assert!(error.contains("closed before"), "got: {error}"); + assert!( + matches!(&error, SteerError::Refused(reason) if reason.contains("closed before")), + "got: {error}" + ); // With the receiver gone the connection itself is gone, whatever the // alive flag still says. diff --git a/crates/acp-client/src/lib.rs b/crates/acp-client/src/lib.rs index f90578324..3f44c6d61 100644 --- a/crates/acp-client/src/lib.rs +++ b/crates/acp-client/src/lib.rs @@ -36,7 +36,7 @@ pub use driver::{ AgentRunOutcome, AsyncTaskStopHandle, BackgroundHoldConfig, BackgroundHoldObserver, BackgroundHoldStatus, BackgroundHoldTask, BasicMessageWriter, MessageWriter, OutOfTurnPermissionPolicy, PromptImage, ReplayBoundary, SessionConnection, SessionLifetime, - SessionSettleReason, SessionSettled, SteerOutcome, Store, TurnControlHandle, + SessionSettleReason, SessionSettled, SteerError, SteerOutcome, Store, TurnControlHandle, BACKGROUND_CONTINUATION_ORIGIN, }; pub use simple::{run_acp_prompt, run_acp_prompt_with_interpreter_env_snapshot};