diff --git a/crates/maestro-acp/src/lib.rs b/crates/maestro-acp/src/lib.rs index d2772ff..aba9b1c 100644 --- a/crates/maestro-acp/src/lib.rs +++ b/crates/maestro-acp/src/lib.rs @@ -33,7 +33,7 @@ #![deny(missing_debug_implementations)] -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{HashMap, VecDeque}; use std::ffi::OsString; use std::path::{Path, PathBuf}; use std::process::{ExitStatus, Stdio}; @@ -179,7 +179,6 @@ enum ChildCommand { #[derive(Debug, Default)] struct PermissionTurnState { active: Option, - cancelled_unresolved: HashSet, } impl PermissionTurnState { @@ -188,23 +187,17 @@ impl PermissionTurnState { } fn cancel_active(&mut self) { - if let Some(turn) = self.active.take() { - self.cancelled_unresolved.insert(turn); - } + self.active = None; } fn finish(&mut self, turn: u64) { if self.active == Some(turn) { self.active = None; } - self.cancelled_unresolved.remove(&turn); } fn callback_turn(&self) -> Option { - self.cancelled_unresolved - .is_empty() - .then_some(self.active) - .flatten() + self.active } } @@ -537,9 +530,8 @@ async fn run_connection( // history the event log already holds (#135). let replaying = Arc::new(AtomicBool::new(false)); let replaying_seen = replaying.clone(); - // ACP callbacks have no turn id. A cancelled prompt remains ambiguous - // until its response arrives, so callbacks are rejected rather than - // risking attribution to a later prompt. + // ACP callbacks have no turn id. The daemon never queues a later prompt + // until a cancelled prompt settles, so this active turn is unambiguous. let permission_turn = Arc::new(Mutex::new(PermissionTurnState::default())); let permission_turn_seen = permission_turn.clone(); let result = Client diff --git a/crates/maestro-acp/tests/fake_acp_agent.py b/crates/maestro-acp/tests/fake_acp_agent.py index a377772..32b593c 100755 --- a/crates/maestro-acp/tests/fake_acp_agent.py +++ b/crates/maestro-acp/tests/fake_acp_agent.py @@ -24,6 +24,7 @@ Environment knobs, all optional: so the SIGTERM path is exercised) FAKE_ACP_IGNORE_EOF=1 keep running after stdin closes, until signalled FAKE_ACP_EMIT_ON_CANCEL=1 emit a final notification before handling cancel + FAKE_ACP_CANCEL_SETTLES=1 complete the outstanding prompt on cancel FAKE_ACP_NO_LOAD=1 advertise `loadSession: false` FAKE_ACP_LOAD_REPLAY=n on `session/load`, replay n `agent_message_chunk` notifications ("replayed ") before answering, @@ -151,6 +152,12 @@ for raw in sys.stdin: elif method == "session/cancel": if os.environ.get("FAKE_ACP_EMIT_ON_CANCEL") == "1": send(chunk("final on cancel")) + if (os.environ.get("FAKE_ACP_CANCEL_SETTLES") == "1" + and pending_prompt_id is not None): + send({"jsonrpc": "2.0", "id": pending_prompt_id, + "result": {"stopReason": "cancelled"}}) + pending_prompt_id = None + pending_permission_id = None if os.environ.get("FAKE_ACP_EXIT_ON_CANCEL") == "1": sys.exit(0) elif method == "session/prompt": diff --git a/crates/maestro-daemon/src/local_api.rs b/crates/maestro-daemon/src/local_api.rs index c11d37f..1cdc19c 100644 --- a/crates/maestro-daemon/src/local_api.rs +++ b/crates/maestro-daemon/src/local_api.rs @@ -26,8 +26,7 @@ use maestro_protocol::{ use serde::Deserialize; use uuid::Uuid; -use crate::projection::TurnLifecycle; -use crate::session::SessionManager; +use crate::session::{QueueAcpPrompt, SessionManager}; use crate::sidecar::EventCursor; #[derive(Debug, Clone)] @@ -436,45 +435,24 @@ async fn structured_prompt( Path(id): Path, Json(body): Json, ) -> Response { - let (acp, turn) = match s.manager.begin_acp_turn(id).await { - Ok(crate::session::BeginAcpTurn::Started(acp, turn)) => (acp, turn), - Ok(crate::session::BeginAcpTurn::InFlight) => { - return error_response( - StatusCode::CONFLICT, - "turn_in_flight", - "the structured session already has an active turn", - ) - } - Ok(crate::session::BeginAcpTurn::Unavailable) => { - return error_response( - StatusCode::CONFLICT, - "not_structured", - "session is not a live structured session", - ) + match s.manager.queue_acp_prompt(id, body.prompt).await { + Ok(QueueAcpPrompt::Queued) => StatusCode::NO_CONTENT.into_response(), + Ok(QueueAcpPrompt::InFlight) => error_response( + StatusCode::CONFLICT, + "turn_in_flight", + "the structured session already has an active turn", + ), + Ok(QueueAcpPrompt::Unavailable) => error_response( + StatusCode::CONFLICT, + "not_structured", + "session is not a live structured session", + ), + Ok(QueueAcpPrompt::Failed(message)) => { + error_response(StatusCode::BAD_GATEWAY, "acp_prompt_failed", &message) } Err(e) => { let info = e.to_error_info(); - return error_response(session_error_status(&info.code), &info.code, &info.message); - } - }; - match acp.prompt(body.prompt, turn) { - Ok(done) => { - let manager = s.manager.clone(); - tokio::spawn(async move { - let lifecycle = match done.await { - Ok(Ok(())) => TurnLifecycle::Finish, - Ok(Err(_)) | Err(_) => TurnLifecycle::Failure, - }; - let _ = manager.apply_acp_turn_lifecycle(id, lifecycle, turn).await; - }); - StatusCode::NO_CONTENT.into_response() - } - Err(message) => { - let _ = s - .manager - .apply_acp_turn_lifecycle(id, TurnLifecycle::Failure, turn) - .await; - error_response(StatusCode::BAD_GATEWAY, "acp_prompt_failed", &message) + error_response(session_error_status(&info.code), &info.code, &info.message) } } } diff --git a/crates/maestro-daemon/src/session.rs b/crates/maestro-daemon/src/session.rs index 0ef56aa..087878b 100644 --- a/crates/maestro-daemon/src/session.rs +++ b/crates/maestro-daemon/src/session.rs @@ -122,6 +122,18 @@ struct AcpEventStream { projection: Projection, projected_live: broadcast::Sender, pending_approvals: HashMap, + /// Serializes prompt dispatch with cancellation. A prompt is not merely + /// claimed: it is queued onto ACP before another operation can proceed. + operation: Arc>, + dispatch: AcpDispatchState, +} + +#[derive(Debug, Default)] +struct AcpDispatchState { + queued_turn: Option, + cancelled_turn: Option, + failed: bool, + settled: Arc, } #[derive(Debug)] @@ -174,6 +186,15 @@ pub enum BeginAcpTurn { Unavailable, } +/// Result of serializing a projected turn claim with its ACP prompt enqueue. +#[derive(Debug)] +pub enum QueueAcpPrompt { + Queued, + InFlight, + Unavailable, + Failed(String), +} + impl SessionError { pub fn to_error_info(&self) -> ErrorInfo { let code = match self { @@ -540,6 +561,150 @@ impl SessionManager { Ok(BeginAcpTurn::Started(acp, turn)) } + /// Claim a turn and enqueue its ACP prompt as one serialized operation. + /// + /// ACP permission callbacks do not carry a prompt id. Therefore a + /// cancelled prompt must settle before a later prompt is dispatched. If + /// it does not settle within the cancel grace, the session is failed + /// instead of letting a later callback become ambiguous. + pub async fn queue_acp_prompt( + &self, + id: Uuid, + text: String, + ) -> Result { + let (acp, stream) = { + let sessions = self.inner.sessions.lock().await; + let record = sessions.get(&id).ok_or(SessionError::NotFound(id))?; + if !record.is_live_structured() { + return Ok(QueueAcpPrompt::Unavailable); + } + (record.acp.clone(), record.acp_events.clone()) + }; + let (Some(acp), Some(stream)) = (acp, stream) else { + return Ok(QueueAcpPrompt::Unavailable); + }; + let operation = stream.lock().await.operation.clone(); + let _operation = operation.lock().await; + + loop { + let waiting_for = { + let mut stream = stream.lock().await; + if stream.dispatch.failed { + return Ok(QueueAcpPrompt::Unavailable); + } + if let Some(turn) = stream.dispatch.queued_turn { + if stream.dispatch.cancelled_turn != Some(turn) { + return Ok(QueueAcpPrompt::InFlight); + } + Some((turn, stream.dispatch.settled.clone())) + } else { + if stream.projection.turn_active() { + return Ok(QueueAcpPrompt::InFlight); + } + let (turn, update) = stream.projection.begin_turn(); + let _ = stream.projected_live.send(update); + match acp.prompt(text.clone(), turn) { + Ok(done) => { + stream.dispatch.queued_turn = Some(turn); + let manager = self.clone(); + tokio::spawn(async move { + let lifecycle = match done.await { + Ok(Ok(())) => TurnLifecycle::Finish, + Ok(Err(_)) | Err(_) => TurnLifecycle::Failure, + }; + manager.finish_queued_acp_turn(id, turn, lifecycle).await; + }); + return Ok(QueueAcpPrompt::Queued); + } + Err(message) => { + cancel_pending_approvals(&mut stream); + let projected = stream + .projection + .turn_lifecycle(TurnLifecycle::Failure, turn) + .expect( + "newly claimed turn remains current while stream is locked", + ); + let _ = stream.projected_live.send(projected); + return Ok(QueueAcpPrompt::Failed(message)); + } + } + } + }; + let Some((turn, settled)) = waiting_for else { + continue; + }; + if tokio::time::timeout(maestro_acp::CANCEL_GRACE, settled.notified()) + .await + .is_err() + { + self.fail_cancelled_acp_turn(id, &stream, turn).await; + } + } + } + + async fn finish_queued_acp_turn(&self, id: Uuid, turn: u64, lifecycle: TurnLifecycle) { + let stream = self + .inner + .sessions + .lock() + .await + .get(&id) + .and_then(|record| record.acp_events.clone()); + let Some(stream) = stream else { + return; + }; + let mut stream = stream.lock().await; + if stream.dispatch.queued_turn != Some(turn) { + return; + } + stream.dispatch.queued_turn = None; + if stream.dispatch.cancelled_turn == Some(turn) { + stream.dispatch.cancelled_turn = None; + } + stream.dispatch.settled.notify_waiters(); + if stream.projection.active_turn() != Some(turn) { + return; + } + cancel_pending_approvals(&mut stream); + if let Some(projected) = stream.projection.turn_lifecycle(lifecycle, turn) { + let _ = stream.projected_live.send(projected); + } + } + + async fn fail_cancelled_acp_turn( + &self, + id: Uuid, + stream: &Arc>, + turn: u64, + ) { + let acp = self + .inner + .sessions + .lock() + .await + .get(&id) + .and_then(|record| record.acp.clone()); + { + let mut stream = stream.lock().await; + if stream.dispatch.queued_turn != Some(turn) + || stream.dispatch.cancelled_turn != Some(turn) + { + return; + } + stream.dispatch.failed = true; + stream.dispatch.queued_turn = None; + stream.dispatch.cancelled_turn = None; + stream.dispatch.settled.notify_waiters(); + } + mark_session_dead(&self.inner, id).await; + if let Some(acp) = acp { + tokio::spawn(async move { + acp.terminate(maestro_acp::CANCEL_GRACE, maestro_acp::TERM_GRACE) + .await; + }); + } + } + /// Cancel an ACP turn and immediately clear its client-facing lifecycle. pub async fn cancel_acp_turn(&self, id: Uuid) -> Result { let (acp, stream) = { @@ -553,10 +718,21 @@ impl SessionManager { let (Some(acp), Some(stream)) = (acp, stream) else { return Ok(false); }; - let mut stream = stream.lock().await; + let event_stream = stream; + let operation = event_stream.lock().await.operation.clone(); + let _operation = operation.lock().await; + let mut stream = event_stream.lock().await; acp.cancel(); cancel_pending_approvals(&mut stream); + let cancelled_queued_turn = stream.dispatch.queued_turn; + if let Some(turn) = cancelled_queued_turn { + stream.dispatch.cancelled_turn = Some(turn); + } let Some(turn) = stream.projection.active_turn() else { + drop(stream); + if let Some(turn) = cancelled_queued_turn { + self.schedule_cancelled_acp_turn_failure(id, event_stream, turn); + } return Ok(true); }; let projected = stream @@ -564,9 +740,26 @@ impl SessionManager { .turn_lifecycle(TurnLifecycle::Cancel, turn) .expect("active turn remains current while stream is locked"); let _ = stream.projected_live.send(projected); + drop(stream); + if let Some(turn) = cancelled_queued_turn { + self.schedule_cancelled_acp_turn_failure(id, event_stream, turn); + } Ok(true) } + fn schedule_cancelled_acp_turn_failure( + &self, + id: Uuid, + stream: Arc>, + turn: u64, + ) { + let manager = self.clone(); + tokio::spawn(async move { + tokio::time::sleep(maestro_acp::CANCEL_GRACE).await; + manager.fail_cancelled_acp_turn(id, &stream, turn).await; + }); + } + /// Resolve a daemon-issued approval nonce. Missing or already-consumed /// nonces deliberately succeed: two clients racing to answer must not /// turn a harmless duplicate click into an error. @@ -897,6 +1090,8 @@ impl SessionManager { projection, projected_live: projected_tx, pending_approvals: HashMap::new(), + operation: Arc::new(Mutex::new(())), + dispatch: AcpDispatchState::default(), }))) } Err(e) => { @@ -1341,6 +1536,8 @@ impl SessionManager { projection, projected_live: projected_tx, pending_approvals: HashMap::new(), + operation: Arc::new(Mutex::new(())), + dispatch: AcpDispatchState::default(), }))) } Err(error) => { diff --git a/crates/maestro-daemon/tests/acp.rs b/crates/maestro-daemon/tests/acp.rs index 2b1c76e..7a2f9ed 100644 --- a/crates/maestro-daemon/tests/acp.rs +++ b/crates/maestro-daemon/tests/acp.rs @@ -325,7 +325,7 @@ async fn stale_turn_completion_cannot_end_the_next_turn() { } #[tokio::test] -async fn delayed_cancelled_permission_cannot_attach_to_next_turn() { +async fn cancelled_prompt_that_never_settles_fails_the_session() { let fx = DaemonFixture::new(); let fake = Fake::new(); let cwd = TempDir::new().unwrap(); @@ -337,50 +337,91 @@ async fn delayed_cancelled_permission_cannot_attach_to_next_turn() { .await .unwrap(); - let maestro_daemon::session::BeginAcpTurn::Started(first_acp, first) = - fx.manager.begin_acp_turn(session.id).await.unwrap() - else { - panic!("first turn should start"); - }; - let _first_done = first_acp.prompt("first".into(), first).unwrap(); + assert!(matches!( + fx.manager + .queue_acp_prompt(session.id, "first".into()) + .await + .unwrap(), + maestro_daemon::session::QueueAcpPrompt::Queued + )); assert!(fx.manager.cancel_acp_turn(session.id).await.unwrap()); + let second = fx + .manager + .queue_acp_prompt(session.id, "second".into()) + .await + .unwrap(); + assert!( + matches!(second, maestro_daemon::session::QueueAcpPrompt::Unavailable), + "{second:?}" + ); + assert_eq!( + fx.manager.status(session.id).await, + Some(SessionStatus::Dead) + ); + assert!(!fake.log().contains("permission response")); +} - let maestro_daemon::session::BeginAcpTurn::Started(second_acp, second) = - fx.manager.begin_acp_turn(session.id).await.unwrap() - else { - panic!("second turn should start"); - }; - let _second_done = second_acp.prompt("second".into(), second).unwrap(); +#[tokio::test] +async fn cancellation_settlement_allows_the_next_prompt_permission() { + let fx = DaemonFixture::new(); + let fake = Fake::new(); + let cwd = TempDir::new().unwrap(); + let mut env = fake.env(); + env.insert("FAKE_ACP_REQUEST_PERMISSION".into(), "1".into()); + env.insert("FAKE_ACP_CANCEL_SETTLES".into(), "1".into()); + let session = fx + .manager + .create(structured(cwd.path(), env)) + .await + .unwrap(); + assert!(matches!( + fx.manager + .queue_acp_prompt(session.id, "first".into()) + .await + .unwrap(), + maestro_daemon::session::QueueAcpPrompt::Queued + )); + assert!(fx.manager.cancel_acp_turn(session.id).await.unwrap()); + assert!(matches!( + fx.manager + .queue_acp_prompt(session.id, "second".into()) + .await + .unwrap(), + maestro_daemon::session::QueueAcpPrompt::Queued + )); + + let nonce = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let (_, state, _) = fx + .manager + .subscribe_acp_projection(session.id, EventCursor::Since(0)) + .await + .unwrap() + .unwrap(); + if let Some(pending) = state.pending_approval { + break pending["nonce"].as_str().unwrap().to_owned(); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert!(fx + .manager + .resolve_acp_approval( + session.id, + &nonce, + maestro_daemon::session::ApprovalDecision::AllowOnce, + ) + .await + .unwrap() + .is_some()); wait_for_async(Duration::from_secs(2), || { let log = fake.log(); async move { log.contains("permission response") } }) .await; - - let Some((snapshot, state, _)) = fx - .manager - .subscribe_acp_projection(session.id, EventCursor::Since(0)) - .await - .unwrap() - else { - panic!("structured projection should remain available"); - }; - assert!(state.pending_approval.is_none()); - assert!(snapshot - .rows - .iter() - .all(|row| row.kind != TranscriptRowKind::Approval)); - assert!(matches!( - fx.manager.begin_acp_turn(session.id).await.unwrap(), - maestro_daemon::session::BeginAcpTurn::InFlight - )); - - assert!(fx - .manager - .apply_acp_turn_lifecycle(session.id, TurnLifecycle::Finish, second) - .await - .unwrap()); fx.manager.kill(session.id, false).await.unwrap(); } diff --git a/crates/maestro-daemon/tests/local_api.rs b/crates/maestro-daemon/tests/local_api.rs index 67229a7..fadfac6 100644 --- a/crates/maestro-daemon/tests/local_api.rs +++ b/crates/maestro-daemon/tests/local_api.rs @@ -260,6 +260,21 @@ async fn structured_prompt_rejects_a_second_in_flight_turn_and_cancel_clears_it( .await .unwrap(); assert_eq!(next.status(), reqwest::StatusCode::NO_CONTENT); + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if fake.log().matches("recv session/prompt").count() == 2 { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("both prompt commands reach ACP"); + let log = fake.log(); + let first_prompt = log.find("recv session/prompt").unwrap(); + let cancel = log.find("recv session/cancel").unwrap(); + let second_prompt = log.rfind("recv session/prompt").unwrap(); + assert!(first_prompt < cancel && cancel < second_prompt, "{log}"); } #[tokio::test]