From ac2eb01bda4d500cddd4f46daf404c73a7569639 Mon Sep 17 00:00:00 2001 From: Miguel Palhas Date: Sat, 5 Sep 2026 11:29:44 +0100 Subject: [PATCH] fix(maestro): allow concurrent ACP handshakes --- crates/maestro-acp/tests/fake_acp_agent.py | 5 + crates/maestro-daemon/src/session.rs | 192 +++++++++++++++------ crates/maestro-daemon/tests/acp.rs | 185 ++++++++++++++++++++ 3 files changed, 325 insertions(+), 57 deletions(-) diff --git a/crates/maestro-acp/tests/fake_acp_agent.py b/crates/maestro-acp/tests/fake_acp_agent.py index 32b593c..6730a9c 100755 --- a/crates/maestro-acp/tests/fake_acp_agent.py +++ b/crates/maestro-acp/tests/fake_acp_agent.py @@ -19,6 +19,8 @@ Environment knobs, all optional: FAKE_ACP_INIT_FAIL=1 answer `initialize` with a JSON-RPC error FAKE_ACP_DELAY_SESSION_NEW_MS=n wait n milliseconds before answering `session/new` + FAKE_ACP_SESSION_NEW_GATE=path wait for `path` to exist before answering + `session/new` FAKE_ACP_EXIT_AFTER_NEW=1 exit 0 right after answering `session/new` FAKE_ACP_EXIT_ON_CANCEL=1 exit 0 on `session/cancel` (default: keep running so the SIGTERM path is exercised) @@ -133,6 +135,9 @@ for raw in sys.stdin: "authMethods": [], }}) elif method == "session/new": + gate = os.environ.get("FAKE_ACP_SESSION_NEW_GATE") + while gate and not os.path.exists(gate): + time.sleep(0.01) time.sleep(int(os.environ.get("FAKE_ACP_DELAY_SESSION_NEW_MS", "0")) / 1000) send({"jsonrpc": "2.0", "id": rid, "result": {"sessionId": session_id}}) send(chunk("hello from fake")) diff --git a/crates/maestro-daemon/src/session.rs b/crates/maestro-daemon/src/session.rs index 06cd2c7..6154c4e 100644 --- a/crates/maestro-daemon/src/session.rs +++ b/crates/maestro-daemon/src/session.rs @@ -6,9 +6,12 @@ //! handler interact with tmux. use std::collections::{HashMap, HashSet}; +use std::future::Future; use std::path::PathBuf; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; +use std::task::{Context, Poll}; use std::time::Duration; use chrono::Utc; @@ -18,7 +21,7 @@ use maestro_protocol::{ SessionStatus, SessionView, StructuredReplayEvent, TranscriptDelta, WorktreeSpec, NAME_KEY, NAME_SOURCE_KEY, }; -use tokio::sync::{broadcast, oneshot, watch, Mutex}; +use tokio::sync::{broadcast, oneshot, watch, Mutex, Notify}; use uuid::Uuid; use crate::panes::Panes; @@ -263,9 +266,99 @@ struct Inner { /// Once daemon shutdown starts, it owns structured-session finalization. /// Background passes must not classify those deliberate exits as natural. shutting_down: AtomicBool, - /// Serializes structured creation with shutdown so every adapter child is - /// either rejected or included in shutdown's final drain. - structured_lifecycle: Mutex<()>, + /// Admission is serialized with shutdown, while admitted handshakes run + /// independently and are counted until registration or rollback finishes. + structured_lifecycle: StructuredLifecycle, +} + +#[derive(Debug, Default)] +struct StructuredLifecycle { + gate: std::sync::Mutex<()>, + in_flight: AtomicUsize, + settled: Notify, +} + +impl StructuredLifecycle { + fn admit<'a>( + &'a self, + shutting_down: &AtomicBool, + ) -> Result, SessionError> { + let _gate = self.gate.lock().expect("structured lifecycle gate"); + if shutting_down.load(Ordering::Acquire) { + return Err(SessionError::ShuttingDown); + } + self.in_flight.fetch_add(1, Ordering::AcqRel); + Ok(StructuredCreate { lifecycle: self }) + } + + async fn wait_until_settled(&self) { + loop { + let settled = self.settled.notified(); + if self.in_flight.load(Ordering::Acquire) == 0 { + return; + } + settled.await; + } + } +} + +struct StructuredCreate<'a> { + lifecycle: &'a StructuredLifecycle, +} + +impl Drop for StructuredCreate<'_> { + fn drop(&mut self) { + if self.lifecycle.in_flight.fetch_sub(1, Ordering::AcqRel) == 1 { + self.lifecycle.settled.notify_one(); + } + } +} + +struct AdmittedStructuredCreate<'a> { + create: Pin> + Send + 'a>>, + /// Field order matters: dropping a cancelled create rolls back its inner + /// future before this permit releases shutdown's in-flight barrier. + _permit: StructuredCreate<'a>, +} + +impl Future for AdmittedStructuredCreate<'_> { + type Output = Result; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.create.as_mut().poll(cx) + } +} + +struct PendingCreate<'a> { + inner: &'a Inner, + id: Uuid, + fifo_path: PathBuf, + worktree_path: Option, + worktree_base_repo: Option, + committed: bool, +} + +impl PendingCreate<'_> { + fn commit(&mut self) { + self.committed = true; + } +} + +impl Drop for PendingCreate<'_> { + fn drop(&mut self) { + if self.committed { + return; + } + let _ = std::fs::remove_file(&self.fifo_path); + if let Err(error) = self.inner.sidecar.delete(self.id) { + tracing::warn!(%error, id = %self.id, "rollback: failed to delete create artifacts"); + } + if let (Some(base), Some(path)) = (&self.worktree_base_repo, &self.worktree_path) { + if let Err(error) = self.inner.worktree_mgr.remove(base, path) { + tracing::warn!(%error, ?path, "rollback: failed to remove worktree after partial create"); + } + } + } } impl SessionManager { @@ -309,7 +402,7 @@ impl SessionManager { viewports: Mutex::new(Viewports::new()), next_viewport: AtomicU64::new(0), shutting_down: AtomicBool::new(false), - structured_lifecycle: Mutex::new(()), + structured_lifecycle: StructuredLifecycle::default(), }), } } @@ -827,10 +920,18 @@ impl SessionManager { /// is on disk after the restart is truthful. tmux sessions are left /// alone — surviving a restart is what they are for. pub async fn shutdown(&self) { - if self.inner.shutting_down.swap(true, Ordering::AcqRel) { - return; + { + let _gate = self + .inner + .structured_lifecycle + .gate + .lock() + .expect("structured lifecycle gate"); + if self.inner.shutting_down.swap(true, Ordering::AcqRel) { + return; + } } - let _lifecycle = self.inner.structured_lifecycle.lock().await; + self.inner.structured_lifecycle.wait_until_settled().await; let children: Vec<(Uuid, Arc, Option>)> = { let sessions = self.inner.sessions.lock().await; sessions @@ -980,11 +1081,15 @@ impl SessionManager { /// Spawn a new session. On success returns the freshly-created `Session`. pub async fn create(&self, req: CreateSessionRequest) -> Result { if req.view == SessionView::Structured { - let _lifecycle = self.inner.structured_lifecycle.lock().await; - if self.inner.shutting_down.load(Ordering::Acquire) { - return Err(SessionError::ShuttingDown); + let permit = self + .inner + .structured_lifecycle + .admit(&self.inner.shutting_down)?; + return AdmittedStructuredCreate { + create: Box::pin(self.create_inner(req)), + _permit: permit, } - return self.create_inner(req).await; + .await; } self.create_inner(req).await } @@ -1022,22 +1127,17 @@ impl SessionManager { // Each side effect below unwinds the ones before it on failure, so a // partial create never leaks a worktree or — worse — a ghost tmux // session with no in-memory record (NAP-260). - let bridge = - match PtyBridge::start_with_replay(fifo_path.clone(), self.inner.replay_bytes).await { - Ok(b) => b, - Err(e) => { - self.rollback_worktree(&worktree_base_repo, &worktree_path); - return Err(e.into()); - } - }; - let mut session_env = match sanitize_env(&req.env) { - Ok(env) => env, - Err(e) => { - let _ = std::fs::remove_file(&fifo_path); - self.rollback_worktree(&worktree_base_repo, &worktree_path); - return Err(e); - } + let mut pending = PendingCreate { + inner: &self.inner, + id, + fifo_path: fifo_path.clone(), + worktree_path: worktree_path.clone(), + worktree_base_repo: worktree_base_repo.clone(), + committed: false, }; + let bridge = + PtyBridge::start_with_replay(fifo_path.clone(), self.inner.replay_bytes).await?; + let mut session_env = sanitize_env(&req.env)?; // Identity + reporting target for the in-session hook (#106). Ours // win over whatever the request carried, so a client cannot spoof // another session's identity (#103). @@ -1084,10 +1184,6 @@ impl SessionManager { .await }; if let Err(e) = spawned { - // `bridge` drops at the end of this block, closing the - // FIFO reader. - let _ = std::fs::remove_file(&fifo_path); - self.rollback_worktree(&worktree_base_repo, &worktree_path); return Err(SessionError::Tmux(e)); } (None, None) @@ -1105,8 +1201,6 @@ impl SessionManager { Err(error) => { acp.terminate(maestro_acp::CANCEL_GRACE, maestro_acp::TERM_GRACE) .await; - let _ = std::fs::remove_file(&fifo_path); - self.rollback_worktree(&worktree_base_repo, &worktree_path); return Err(SessionError::Sidecar(error)); } }; @@ -1131,8 +1225,6 @@ impl SessionManager { Err(e) => { acp.terminate(maestro_acp::CANCEL_GRACE, maestro_acp::TERM_GRACE) .await; - let _ = std::fs::remove_file(&fifo_path); - self.rollback_worktree(&worktree_base_repo, &worktree_path); return Err(SessionError::Sidecar(e)); } } @@ -1201,9 +1293,6 @@ impl SessionManager { let _ = tmux_blocking(move || tmux.kill_session(id)).await; } } - let _ = std::fs::remove_file(&fifo_path); - let _ = self.inner.sidecar.delete(id); - self.rollback_worktree(&worktree_base_repo, &worktree_path); return Err(SessionError::Sidecar(e)); } @@ -1213,6 +1302,11 @@ impl SessionManager { .map(|(acp, events)| (acp.clone(), events.clone())); { let mut sessions = self.inner.sessions.lock().await; + let event_pump = event_pump.map(|(acp, events)| { + let done = start_acp_event_pump(&self.inner, id, &acp, events.clone()); + start_acp_permission_pump(&self.inner, id, &acp, events); + done + }); sessions.insert( id, SessionRecord { @@ -1229,17 +1323,10 @@ impl SessionManager { awaiting_suppressed: false, acp, acp_events, - acp_event_pump: None, + acp_event_pump: event_pump, }, ); } - if let Some((acp, events)) = event_pump { - let done = start_acp_event_pump(&self.inner, id, &acp, events.clone()); - start_acp_permission_pump(&self.inner, id, &acp, events); - if let Some(record) = self.inner.sessions.lock().await.get_mut(&id) { - record.acp_event_pump = Some(done); - } - } let _ = self.inner.events.send(CatalogEvent::SessionAdded { session: session.clone(), }); @@ -1264,19 +1351,10 @@ impl SessionManager { } } + pending.commit(); Ok(session) } - /// Best-effort removal of a worktree created earlier in `create` when a - /// later step fails. No-op when the request had no worktree. - fn rollback_worktree(&self, base: &Option, path: &Option) { - if let (Some(base), Some(path)) = (base, path) { - if let Err(e) = self.inner.worktree_mgr.remove(base, path) { - tracing::warn!(error = %e, ?path, "rollback: failed to remove worktree after partial create"); - } - } - } - /// Kill a session and remove its on-disk artifacts. /// /// Refused (#79) when the session's worktree holds unsaved work — diff --git a/crates/maestro-daemon/tests/acp.rs b/crates/maestro-daemon/tests/acp.rs index 10f51c5..d8461ff 100644 --- a/crates/maestro-daemon/tests/acp.rs +++ b/crates/maestro-daemon/tests/acp.rs @@ -950,6 +950,191 @@ async fn shutdown_waits_for_structured_creation_and_rejects_new_creates() { assert_eq!(fake.log().matches("signal SIGTERM").count(), 1); } +#[tokio::test] +async fn structured_handshakes_run_independently() { + let fx = DaemonFixture::new(); + let blocked = Fake::new(); + let ready = Fake::new(); + let cwd = TempDir::new().unwrap(); + let gate = cwd.path().join("release-first"); + let mut blocked_env = blocked.env(); + blocked_env.insert( + "FAKE_ACP_SESSION_NEW_GATE".into(), + gate.to_string_lossy().into_owned(), + ); + + let first_cwd = cwd.path().to_path_buf(); + let manager = fx.manager.clone(); + let first = + tokio::spawn(async move { manager.create(structured(&first_cwd, blocked_env)).await }); + wait_for_async(Duration::from_secs(5), || async { + blocked.log().contains("recv session/new") + }) + .await; + + let second = tokio::time::timeout( + Duration::from_secs(5), + fx.manager.create(structured(cwd.path(), ready.env())), + ) + .await + .expect("second handshake was blocked by the first") + .expect("second create"); + assert_eq!(second.status, SessionStatus::Running); + assert!(!first.is_finished(), "the gate must still block the first"); + + std::fs::write(&gate, b"release").unwrap(); + let first = first.await.expect("first task").expect("first create"); + fx.manager.kill(first.id, false).await.expect("kill first"); + fx.manager + .kill(second.id, false) + .await + .expect("kill second"); +} + +#[tokio::test] +async fn shutdown_waits_for_every_admitted_structured_create() { + let fx = DaemonFixture::new(); + let a = Fake::new(); + let b = Fake::new(); + let late = Fake::new(); + let cwd = TempDir::new().unwrap(); + let gate_a = cwd.path().join("release-a"); + let gate_b = cwd.path().join("release-b"); + let mut env_a = a.env(); + let mut env_b = b.env(); + env_a.insert( + "FAKE_ACP_SESSION_NEW_GATE".into(), + gate_a.to_string_lossy().into_owned(), + ); + env_b.insert( + "FAKE_ACP_SESSION_NEW_GATE".into(), + gate_b.to_string_lossy().into_owned(), + ); + + let create_a = { + let manager = fx.manager.clone(); + let cwd = cwd.path().to_path_buf(); + tokio::spawn(async move { manager.create(structured(&cwd, env_a)).await }) + }; + let create_b = { + let manager = fx.manager.clone(); + let cwd = cwd.path().to_path_buf(); + tokio::spawn(async move { manager.create(structured(&cwd, env_b)).await }) + }; + wait_for_async(Duration::from_secs(5), || async { + a.log().contains("recv session/new") && b.log().contains("recv session/new") + }) + .await; + + let shutdown = { + let manager = fx.manager.clone(); + tokio::spawn(async move { manager.shutdown().await }) + }; + wait_for_async(Duration::from_secs(5), || async { + matches!( + fx.manager.create(structured(cwd.path(), late.env())).await, + Err(maestro_daemon::SessionError::ShuttingDown) + ) + }) + .await; + assert!(!shutdown.is_finished()); + + std::fs::write(&gate_a, b"release").unwrap(); + let session_a = create_a.await.expect("create a task").expect("create a"); + assert!( + !shutdown.is_finished(), + "shutdown ignored the second create" + ); + + std::fs::write(&gate_b, b"release").unwrap(); + let session_b = create_b.await.expect("create b task").expect("create b"); + shutdown.await.expect("shutdown task"); + + wait_for_exit(a.pid()).await; + wait_for_exit(b.pid()).await; + assert_eq!( + fx.manager.status(session_a.id).await, + Some(SessionStatus::Dead) + ); + assert_eq!( + fx.manager.status(session_b.id).await, + Some(SessionStatus::Dead) + ); +} + +#[tokio::test] +async fn cancelling_a_structured_handshake_rolls_back_and_settles_admission() { + let fx = DaemonFixture::new(); + let fake = Fake::new(); + let repo = TempDir::new().unwrap(); + init_repo(repo.path()); + let gate = repo.path().join("release-create"); + let mut env = fake.env(); + env.insert("FAKE_ACP_ECHO_ENV".into(), "MAESTRO_SESSION_ID".into()); + env.insert( + "FAKE_ACP_SESSION_NEW_GATE".into(), + gate.to_string_lossy().into_owned(), + ); + let request = CreateSessionRequest { + view: SessionView::Structured, + cwd: repo.path().to_path_buf(), + command: "claude".into(), + env, + metadata: HashMap::new(), + worktree: Some(WorktreeSpec { + base_repo: repo.path().to_path_buf(), + branch: "cancelled-create".into(), + create_branch: true, + }), + }; + + let create = { + let manager = fx.manager.clone(); + tokio::spawn(async move { manager.create(request).await }) + }; + wait_for_async(Duration::from_secs(5), || async { + fake.log().contains("recv session/new") + }) + .await; + let id = fake + .log() + .lines() + .find_map(|line| line.strip_prefix("env MAESTRO_SESSION_ID=")) + .and_then(|id| uuid::Uuid::parse_str(id).ok()) + .expect("session id in adapter log"); + let pid = fake.pid(); + + create.abort(); + let shutdown = { + let manager = fx.manager.clone(); + tokio::spawn(async move { manager.shutdown().await }) + }; + assert!(create.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(5), shutdown) + .await + .expect("cancelled admission kept shutdown waiting") + .expect("shutdown task"); + wait_for_exit(pid).await; + wait_for_async(Duration::from_secs(5), || async { + !fx.fifo_dir().join(format!("{id}.fifo")).exists() + && !repo.path().join("cancelled-create").exists() + }) + .await; + assert!(fx.manager.snapshot().await.is_empty()); + let id = id.to_string(); + assert!( + std::fs::read_dir(fx.sidecar_dir()) + .into_iter() + .flatten() + .all(|entry| !entry + .unwrap() + .file_name() + .to_string_lossy() + .starts_with(&id)), + "cancelled create left a sidecar artifact" + ); +} + #[tokio::test] async fn shutdown_drains_final_notifications_to_the_log() { let fx = DaemonFixture::new();