From cf285dce271fa7684b16f2c91c760654123fbbb1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20M=C3=BCller?= Date: Thu, 1 Oct 2026 20:42:19 -0700 Subject: [PATCH] fix(acp): fail pending link requests on close and drop the turn cap (#3249) Co-authored-by: Kevin Thomas --- crates/pa-cli/tests/acp_mode_e2e.rs | 99 ++++++++++--- crates/pa-daemon/src/acp/daemon.rs | 182 +++++++++++++----------- crates/pa-daemon/src/acp/wire_config.rs | 60 ++++---- 3 files changed, 204 insertions(+), 137 deletions(-) diff --git a/crates/pa-cli/tests/acp_mode_e2e.rs b/crates/pa-cli/tests/acp_mode_e2e.rs index bb60b2c0f..9cb427cea 100644 --- a/crates/pa-cli/tests/acp_mode_e2e.rs +++ b/crates/pa-cli/tests/acp_mode_e2e.rs @@ -186,6 +186,28 @@ impl AcpChild { } } } + + /// Read frames until one `sessionUpdate` of `kind` arrives — the + /// observed-event readiness signal, never a timer. The frames + /// before it are dropped. + fn wait_update(&mut self, kind: &str, timeout: Duration) { + let deadline = Instant::now() + timeout; + loop { + let timeout_left = deadline.saturating_duration_since(Instant::now()); + assert!( + !timeout_left.is_zero(), + "the stream never published a {kind} update" + ); + let line = self + .lines + .recv_timeout(timeout_left) + .expect("the ACP stream stayed open"); + let frame: Value = serde_json::from_str(&line).expect("valid JSON line"); + if frame["params"]["update"]["sessionUpdate"] == kind { + return; + } + } + } } impl Drop for AcpChild { @@ -852,22 +874,7 @@ fn acp_daemon_attached_prompt_after_cancel_runs() { &json!({ "sessionId": session_id, "prompt": [{ "type": "text", "text": "a slow question" }] }), ); // Readiness is the turn's own first streamed chunk, not a timer. - let deadline = Instant::now() + TIMEOUT; - loop { - let timeout_left = deadline.saturating_duration_since(Instant::now()); - assert!( - !timeout_left.is_zero(), - "the paced turn never streamed a chunk" - ); - let line = client - .lines - .recv_timeout(timeout_left) - .expect("the ACP stream stayed open"); - let frame: Value = serde_json::from_str(&line).expect("valid JSON line"); - if frame["params"]["update"]["sessionUpdate"] == "agent_message_chunk" { - break; - } - } + client.wait_update("agent_message_chunk", TIMEOUT); client.notify("session/cancel", &json!({ "sessionId": session_id })); let (first_response, _) = client.wait_response(first, TIMEOUT); assert_eq!( @@ -893,6 +900,66 @@ fn acp_daemon_attached_prompt_after_cancel_runs() { ); } +#[cfg(unix)] +#[test] +fn acp_daemon_attached_supervisor_loss_fails_the_prompt() { + // The supervisor dies while a prompt is in flight: the pending + // request must fail fast with an error response, never hang, because + // the link close is the turn's liveness bound (turn-long requests + // carry no timer). + let (mut client, socket) = AcpChild::spawn_daemon_attached( + &["--mode", "acp", "--no-session"], + &json!({ "engine": "faux", "tokensPerSecond": 2, "responses": [ + { "text": "a paced answer that streams slowly enough to outlive the supervisor" }, + ] }), + ); + let init = client.request("initialize", &initialize_params()); + let _ = client.wait_response(init, TIMEOUT); + let new = client.request("session/new", &json!({ "mcpServers": [] })); + let (new_response, _) = client.wait_response(new, TIMEOUT); + let session_id = new_response["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + let prompt = client.request( + "session/prompt", + &json!({ "sessionId": session_id, "prompt": [{ "type": "text", "text": "a slow question" }] }), + ); + // Readiness is the turn's own first streamed chunk: the prompt is + // provably in flight before the supervisor goes away. + client.wait_update("agent_message_chunk", TIMEOUT); + // Crash the supervisor (no graceful drain, so the in-flight response + // can never arrive). Its pid comes from the hello frame every new + // connection receives. + let probe = pa_types::platform::transport::connect_blocking(&socket).expect("daemon socket"); + let reader = probe.try_clone_box().expect("daemon socket clone"); + let _ = reader.set_read_timeout(Duration::from_mins(2)); + let hello = BufReader::new(reader) + .lines() + .map(|line| { + serde_json::from_str::(&line.expect("daemon line")).expect("daemon JSON") + }) + .find(|frame| frame["type"] == json!("daemon_hello")) + .expect("the daemon closed without a hello"); + let pid = hello["supervisorPid"].as_u64().expect("supervisor pid"); + drop(probe); + let killed = std::process::Command::new("kill") + .args(["-9", &pid.to_string()]) + .status() + .expect("kill the supervisor"); + assert!(killed.success(), "the supervisor crash did not run"); + let (prompt_response, _) = client.wait_response(prompt, TIMEOUT); + assert_eq!( + prompt_response["error"], + json!({ + "code": -32603, + "message": "Internal error", + "data": { "details": "the daemon connection closed mid-request" } + }), + "the in-flight prompt fails fast on supervisor loss: {prompt_response}" + ); +} + #[test] fn acp_daemon_attached_image_only_prompt_reaches_the_session() { let (mut client, socket) = AcpChild::spawn_daemon_attached( diff --git a/crates/pa-daemon/src/acp/daemon.rs b/crates/pa-daemon/src/acp/daemon.rs index 17f37ed35..6b5712801 100644 --- a/crates/pa-daemon/src/acp/daemon.rs +++ b/crates/pa-daemon/src/acp/daemon.rs @@ -40,10 +40,8 @@ use super::wire_config::{ }; use super::wire_events::{self, WireMappingState}; -/// Response timeout for turn-long commands (the turn itself bounds them). -pub(crate) const TURN_TIMEOUT_MS: u64 = 600_000; /// Response timeout for session-scoped commands. -pub(crate) const REQUEST_TIMEOUT_MS: u64 = 30_000; +const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); /// One inbound supervisor frame, classified by the reader. enum LinkFrame { @@ -56,6 +54,10 @@ enum LinkFrame { pub(crate) struct DaemonLink { writer: mpsc::UnboundedSender, pending: Arc>>>, + /// Set once the frame channel ends (the socket closed and every + /// frame the reader read is delivered): every new request fails + /// fast instead of waiting for a response no one will send. + closed: Arc, frames: Mutex>, protocol_version: u64, next_request_id: std::sync::atomic::AtomicU64, @@ -76,6 +78,8 @@ impl DaemonLink { let (hello_tx, hello_rx) = oneshot::channel::(); let pending: Arc>>> = Arc::new(std::sync::Mutex::new(HashMap::new())); + let closed: Arc = + Arc::new(std::sync::atomic::AtomicBool::new(false)); tokio::spawn(async move { let mut writer = writer_half; @@ -146,17 +150,33 @@ impl DaemonLink { Ok(DaemonLink { writer: line_tx, pending, + closed, frames: Mutex::new(frame_rx), protocol_version: version, next_request_id: std::sync::atomic::AtomicU64::new(0), }) } - /// Send one command envelope and wait for the matching response. - pub(crate) async fn request( + /// Send one command envelope and wait for the matching response + /// under the session-scoped response timeout. + pub(crate) async fn request(&self, command: DaemonCommand) -> anyhow::Result { + self.exchange(command, Some(REQUEST_TIMEOUT)).await + } + + /// Send one command envelope with no fixed cap: the response or the + /// link close ends the wait (turn-long commands — declared + /// divergence, TS caps them at 24 h; the link-close signal is the + /// liveness bound here, not a timer). + async fn request_until_close(&self, command: DaemonCommand) -> anyhow::Result { + self.exchange(command, None).await + } + + /// One exchange: register the response slot, send the envelope, and + /// wait for the answer. + async fn exchange( &self, command: DaemonCommand, - timeout_ms: u64, + timeout: Option, ) -> anyhow::Result { use std::sync::atomic::Ordering; let id = format!( @@ -176,6 +196,14 @@ impl DaemonLink { let line = serde_json::to_string(&envelope)?; let (tx, rx) = oneshot::channel::(); self.pending.lock().unwrap().insert(id.clone(), tx); + // A request that races the close loses either way: the flag + // fails it here, or the pending clear already dropped its + // sender. The flag is stored before the clear, so no request + // parks unnoticed. + if self.closed.load(Ordering::SeqCst) { + self.pending.lock().unwrap().remove(&id); + anyhow::bail!("the daemon connection is closed"); + } if self.writer.send(line).is_err() { // A closed writer leaves the pending slot behind otherwise; a // link that never answers again would grow one entry per @@ -183,13 +211,14 @@ impl DaemonLink { self.pending.lock().unwrap().remove(&id); anyhow::bail!("the daemon connection is closed"); } - tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), rx) - .await - .map_err(|_| { + match timeout { + Some(timeout) => tokio::time::timeout(timeout, rx).await.map_err(|_| { self.pending.lock().unwrap().remove(&id); anyhow::anyhow!("timed out waiting for the daemon response") - })? - .map_err(|_| anyhow::anyhow!("the daemon connection closed mid-request")) + })?, + None => rx.await, + } + .map_err(|_| anyhow::anyhow!("the daemon connection closed mid-request")) } } @@ -370,6 +399,13 @@ pub async fn run_daemon_attached_acp_mode(options: DaemonAcpOptions) -> anyhow:: } } } + // The reader hit the socket close and every frame it read + // has been delivered: fail the still-pending requests + // (dropping a sender answers with the mid-request error) + // and fail every later request fast. The flag goes up + // before the clear; `request` covers the race. + link.closed.store(true, std::sync::atomic::Ordering::SeqCst); + link.pending.lock().unwrap().clear(); }); } @@ -558,7 +594,7 @@ async fn handle_session_new( launch_env: None, rest: Map::default(), }; - let create_response = match link.request(create, REQUEST_TIMEOUT_MS).await { + let create_response = match link.request(create).await { Ok(response) => response, Err(error) => { *state.lock().await = DaemonAcpState::default(); @@ -593,20 +629,17 @@ async fn handle_session_new( launch_env: None, rest: Map::default(), }; - if let Ok(response) = link.request(attach, REQUEST_TIMEOUT_MS).await { + if let Ok(response) = link.request(attach).await { if !response.success { let failure = response .error .unwrap_or_else(|| "unknown error".to_string()); let _ = link - .request( - DaemonCommand::Kill { - id: None, - active_session_id: daemon_active_session_id.clone(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::Kill { + id: None, + active_session_id: daemon_active_session_id.clone(), + rest: Map::default(), + }) .await; *state.lock().await = DaemonAcpState::default(); let _ = tx.send(super::internal_error(&id, &failure)); @@ -649,14 +682,11 @@ async fn handle_session_new( if let Err(error) = replace_session_servers(link, &hosted, &resolved).await { let failure = error.to_string(); let _ = link - .request( - DaemonCommand::Kill { - id: None, - active_session_id: hosted.daemon_active_session_id.clone(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::Kill { + id: None, + active_session_id: hosted.daemon_active_session_id.clone(), + rest: Map::default(), + }) .await; *state.lock().await = DaemonAcpState::default(); let _ = tx.send(super::internal_error(&id, &failure)); @@ -710,16 +740,13 @@ async fn replace_session_servers( resolved: &[pa_core::mcp::AcpMcpServerConfig], ) -> anyhow::Result<()> { let response = link - .request( - DaemonCommand::ReplaceAcpMcpServers { - id: None, - active_session_id: hosted.daemon_active_session_id.clone(), - owner_id: hosted.mcp_owner_id.clone(), - servers: serde_json::to_value(resolved)?, - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::ReplaceAcpMcpServers { + id: None, + active_session_id: hosted.daemon_active_session_id.clone(), + owner_id: hosted.mcp_owner_id.clone(), + servers: serde_json::to_value(resolved)?, + rest: Map::default(), + }) .await?; if !response.success { anyhow::bail!(response @@ -823,7 +850,7 @@ async fn handle_session_prompt( }, rest: Map::default(), }; - let response = match link.request(prompt, TURN_TIMEOUT_MS).await { + let response = match link.request_until_close(prompt).await { Ok(response) => response, Err(error) => { producer.finish_prompt(turn_id).await; @@ -988,15 +1015,12 @@ async fn fetch_autonomous_status( active_session_id: &str, ) -> Option { let response = link - .request( - DaemonCommand::WaitForHeadlessCompletion { - id: None, - active_session_id: active_session_id.to_string(), - wait_for_rlm_quiescence: None, - rest: Map::default(), - }, - TURN_TIMEOUT_MS, - ) + .request_until_close(DaemonCommand::WaitForHeadlessCompletion { + id: None, + active_session_id: active_session_id.to_string(), + wait_for_rlm_quiescence: None, + rest: Map::default(), + }) .await .ok()?; if !response.success { @@ -1031,7 +1055,7 @@ async fn session_cancel( rest: Map::default(), }; drop(guard); - let _ = link.request(abort, REQUEST_TIMEOUT_MS).await; + let _ = link.request(abort).await; Ok(()) } @@ -1067,32 +1091,26 @@ async fn handle_session_close( return; }; let _ = link - .request( - DaemonCommand::Abort { - id: None, - active_session_id: hosted.daemon_active_session_id.clone(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::Abort { + id: None, + active_session_id: hosted.daemon_active_session_id.clone(), + rest: Map::default(), + }) .await; if !hosted.mcp_server_names.is_empty() { let _ = replace_session_servers(link, &hosted, &[]).await; } // The worker dies before the config queue drains: a stalled // `set_model`/`set_thinking_level` (holding the queue on a wire - // request) fails fast once the worker is gone instead of parking the - // close for the turn timeout. Then the serialized config work settles - // before the producer fences (TS `await configTask`). + // request) fails fast once the worker is gone instead of parking + // the close on the request timeout. Then the serialized config work + // settles before the producer fences (TS `await configTask`). let _ = link - .request( - DaemonCommand::Kill { - id: None, - active_session_id: hosted.daemon_active_session_id.clone(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::Kill { + id: None, + active_session_id: hosted.daemon_active_session_id.clone(), + rest: Map::default(), + }) .await; let _ = hosted.config.queue.lock().await; hosted.producer.close().await; @@ -1111,24 +1129,18 @@ async fn teardown(link: &Arc, state: &Arc>) { // operation holding the queue releases once the wire peer is gone), // then the serialized config work settles before the producer fences. let _ = link - .request( - DaemonCommand::Abort { - id: None, - active_session_id: hosted.daemon_active_session_id.clone(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::Abort { + id: None, + active_session_id: hosted.daemon_active_session_id.clone(), + rest: Map::default(), + }) .await; let _ = link - .request( - DaemonCommand::Kill { - id: None, - active_session_id: hosted.daemon_active_session_id.clone(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::Kill { + id: None, + active_session_id: hosted.daemon_active_session_id.clone(), + rest: Map::default(), + }) .await; let _ = hosted.config.queue.lock().await; hosted.producer.close().await; diff --git a/crates/pa-daemon/src/acp/wire_config.rs b/crates/pa-daemon/src/acp/wire_config.rs index a0606aad6..abd552c7b 100644 --- a/crates/pa-daemon/src/acp/wire_config.rs +++ b/crates/pa-daemon/src/acp/wire_config.rs @@ -17,7 +17,7 @@ use super::config_options::{ config_options_value, model_value, publish_config_options, session_config_options, PickerModel, SessionConfigOption, }; -use super::daemon::{DaemonAcpState, DaemonLink, REQUEST_TIMEOUT_MS, TURN_TIMEOUT_MS}; +use super::daemon::{DaemonAcpState, DaemonLink}; use super::jsonrpc; use super::producer::{self, UpdateProducer}; use super::types; @@ -189,16 +189,13 @@ async fn apply_wire_config( WireConfigError::invalid_params(format!("Unavailable model: {value}")) })?; let response = link - .request( - DaemonCommand::SetModel { - id: None, - active_session_id: daemon_session_id.to_string(), - provider: model.provider.clone(), - model_id: model.id.clone(), - rest: Map::default(), - }, - TURN_TIMEOUT_MS, - ) + .request(DaemonCommand::SetModel { + id: None, + active_session_id: daemon_session_id.to_string(), + provider: model.provider.clone(), + model_id: model.id.clone(), + rest: Map::default(), + }) .await .map_err(|error| WireConfigError::internal(error.to_string()))?; if !response.success { @@ -250,15 +247,12 @@ async fn apply_wire_config( ))); } let response = link - .request( - DaemonCommand::SetThinkingLevel { - id: None, - active_session_id: daemon_session_id.to_string(), - level: value.to_string(), - rest: Map::default(), - }, - TURN_TIMEOUT_MS, - ) + .request(DaemonCommand::SetThinkingLevel { + id: None, + active_session_id: daemon_session_id.to_string(), + level: value.to_string(), + rest: Map::default(), + }) .await .map_err(|error| WireConfigError::internal(error.to_string()))?; if !response.success { @@ -283,14 +277,11 @@ pub(super) async fn fetch_connection_state( active_session_id: &str, ) -> Option { let response = link - .request( - DaemonCommand::GetConnectionState { - id: None, - active_session_id: active_session_id.to_string(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::GetConnectionState { + id: None, + active_session_id: active_session_id.to_string(), + rest: Map::default(), + }) .await .ok()?; if !response.success { @@ -305,14 +296,11 @@ pub(super) async fn fetch_available_models( active_session_id: &str, ) -> anyhow::Result> { let response = link - .request( - DaemonCommand::GetAvailableModels { - id: None, - active_session_id: active_session_id.to_string(), - rest: Map::default(), - }, - REQUEST_TIMEOUT_MS, - ) + .request(DaemonCommand::GetAvailableModels { + id: None, + active_session_id: active_session_id.to_string(), + rest: Map::default(), + }) .await?; if !response.success { anyhow::bail!(response -- knotmirror 0.0.1