From be23a2b53918d6328f64e01740691bee88689eb2 Mon Sep 17 00:00:00 2001 From: Miguel Palhas Date: Sat, 5 Sep 2026 01:56:07 +0100 Subject: [PATCH] fix: harden structured replay --- crates/maestro-daemon/src/local_api.rs | 37 ++++-- crates/maestro-daemon/src/projection.rs | 85 +++++++----- crates/maestro-daemon/src/session.rs | 88 +++++++++++-- crates/maestro-daemon/tests/acp.rs | 54 +++++++- crates/maestro-daemon/tests/local_api.rs | 160 +++++++++++++++++++++-- crates/maestro-protocol/src/lib.rs | 13 ++ 6 files changed, 369 insertions(+), 68 deletions(-) diff --git a/crates/maestro-daemon/src/local_api.rs b/crates/maestro-daemon/src/local_api.rs index 8087ace..e7959c4 100644 --- a/crates/maestro-daemon/src/local_api.rs +++ b/crates/maestro-daemon/src/local_api.rs @@ -400,15 +400,26 @@ async fn structured_replay( Some(before) => EventCursor::Before(Some(before)), None => EventCursor::Since(query.since.unwrap_or(0)), }; - match s.manager.read_acp_events( - id, - cursor, - query - .limit - .unwrap_or(100) - .clamp(1, MAX_STRUCTURED_REPLAY_LIMIT), - ) { - Ok(events) => Json(events).into_response(), + match s + .manager + .read_acp_projection_replay( + id, + cursor, + query + .limit + .unwrap_or(100) + .clamp(1, MAX_STRUCTURED_REPLAY_LIMIT), + ) + .await + { + Ok(Some(events)) => Json(events).into_response(), + // A terminal session, or a structured session whose adapter is no + // longer live, cannot provide a coherent replay continuation. + Ok(None) => error_response( + StatusCode::CONFLICT, + "not_structured", + "session is not a live structured session", + ), Err(e) => { let info = e.to_error_info(); error_response(session_error_status(&info.code), &info.code, &info.message) @@ -421,8 +432,8 @@ async fn structured_prompt( Path(id): Path, Json(body): Json, ) -> Response { - let acp = match s.manager.begin_acp_turn(id).await { - Ok(Some(acp)) => acp, + let (acp, turn) = match s.manager.begin_acp_turn(id).await { + Ok(Some(turn)) => turn, Ok(None) => { return error_response( StatusCode::CONFLICT, @@ -443,14 +454,14 @@ async fn structured_prompt( Ok(Ok(())) => TurnLifecycle::Finish, Ok(Err(_)) | Err(_) => TurnLifecycle::Failure, }; - let _ = manager.apply_acp_turn_lifecycle(id, lifecycle).await; + 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) + .apply_acp_turn_lifecycle(id, TurnLifecycle::Failure, turn) .await; error_response(StatusCode::BAD_GATEWAY, "acp_prompt_failed", &message) } diff --git a/crates/maestro-daemon/src/projection.rs b/crates/maestro-daemon/src/projection.rs index 426b0fd..4057e5b 100644 --- a/crates/maestro-daemon/src/projection.rs +++ b/crates/maestro-daemon/src/projection.rs @@ -35,7 +35,8 @@ pub struct Projection { tools: HashMap, pending_tool_updates: HashMap>, state: ReducedState, - turn_active: bool, + active_turn: Option, + next_turn: u64, turn_tool_calls: HashSet, } @@ -105,31 +106,53 @@ impl Projection { } pub fn turn_active(&self) -> bool { - self.turn_active + self.active_turn.is_some() } - /// Apply a transport-owned prompt lifecycle transition. - pub fn turn_lifecycle(&mut self, lifecycle: TurnLifecycle) -> ProjectionUpdate { - match lifecycle { - TurnLifecycle::Begin => { - self.turn_active = true; - self.state.turn.responding = true; - self.state.turn.thinking = true; - self.turn_tool_calls.clear(); - self.state.turn.active_tool_calls.clear(); - } - TurnLifecycle::Finish | TurnLifecycle::Cancel | TurnLifecycle::Failure => { - self.turn_active = false; - self.state.turn.responding = false; - self.state.turn.thinking = false; - self.turn_tool_calls.clear(); - self.state.turn.active_tool_calls.clear(); - } + /// Start a prompt turn and return its monotonically increasing token. + pub fn begin_turn(&mut self) -> (u64, ProjectionUpdate) { + self.next_turn += 1; + let turn = self.next_turn; + self.active_turn = Some(turn); + self.state.turn.responding = true; + self.state.turn.thinking = true; + self.turn_tool_calls.clear(); + self.state.turn.active_tool_calls.clear(); + ( + turn, + ProjectionUpdate { + transcript: None, + state: self.state(), + }, + ) + } + + /// End `turn` only when it remains the active prompt turn. + pub fn turn_lifecycle( + &mut self, + lifecycle: TurnLifecycle, + turn: u64, + ) -> Option { + if !matches!( + lifecycle, + TurnLifecycle::Finish | TurnLifecycle::Cancel | TurnLifecycle::Failure + ) || self.active_turn != Some(turn) + { + return None; } - ProjectionUpdate { + self.active_turn = None; + self.state.turn.responding = false; + self.state.turn.thinking = false; + self.turn_tool_calls.clear(); + self.state.turn.active_tool_calls.clear(); + Some(ProjectionUpdate { transcript: None, state: self.state(), - } + }) + } + + pub fn active_turn(&self) -> Option { + self.active_turn } fn content( @@ -234,28 +257,28 @@ impl Projection { } fn agent_text(&mut self) { - if self.turn_active { + if self.turn_active() { self.state.turn.responding = true; self.state.turn.thinking = false; } } fn agent_thought(&mut self) { - if self.turn_active { + if self.turn_active() { self.state.turn.responding = true; self.state.turn.thinking = true; } } fn track_tool_call(&mut self, tool: &ToolCall) { - if self.turn_active { + if self.turn_active() { self.turn_tool_calls.insert(tool.tool_call_id.to_string()); self.sync_active_tool(&tool.tool_call_id.to_string()); } } fn sync_active_tool(&mut self, tool_id: &str) { - if !self.turn_active || !self.turn_tool_calls.contains(tool_id) { + if !self.turn_active() || !self.turn_tool_calls.contains(tool_id) { return; } let active = self.tools.get(tool_id).is_some_and(|tool| { @@ -542,7 +565,7 @@ mod tests { ); assert_eq!(projection.state().turn, Default::default()); - projection.turn_lifecycle(TurnLifecycle::Begin); + let (turn, _) = projection.begin_turn(); assert!(projection.state().turn.responding); assert!(projection.state().turn.thinking); projection.apply( @@ -562,8 +585,8 @@ mod tests { TurnLifecycle::Cancel, TurnLifecycle::Failure, ] { - projection.turn_lifecycle(TurnLifecycle::Begin); - projection.turn_lifecycle(terminal); + let (turn, _) = projection.begin_turn(); + projection.turn_lifecycle(terminal, turn); assert!(!projection.state().turn.responding); assert!(!projection.state().turn.thinking); } @@ -575,7 +598,7 @@ mod tests { next_id: &str, ) { let mut projection = Projection::default(); - projection.turn_lifecycle(TurnLifecycle::Begin); + let (turn, _) = projection.begin_turn(); projection.apply( 1, &event(serde_json::json!({ @@ -585,11 +608,11 @@ mod tests { ); assert_eq!(projection.state().turn.active_tool_calls, [unfinished_id]); - projection.turn_lifecycle(terminal); + projection.turn_lifecycle(terminal, turn); assert!(projection.state().turn.active_tool_calls.is_empty()); assert_eq!(projection.snapshot(1).rows.len(), 1); - projection.turn_lifecycle(TurnLifecycle::Begin); + projection.begin_turn(); projection.apply( 2, &event(serde_json::json!({ diff --git a/crates/maestro-daemon/src/session.rs b/crates/maestro-daemon/src/session.rs index 00f0656..a3a4340 100644 --- a/crates/maestro-daemon/src/session.rs +++ b/crates/maestro-daemon/src/session.rs @@ -15,7 +15,7 @@ use chrono::Utc; use maestro_acp::{AcpSession, SpawnConfig}; use maestro_protocol::{ ActivityState, CatalogEvent, CreateSessionRequest, ErrorInfo, NameSource, Session, - SessionStatus, SessionView, WorktreeSpec, NAME_KEY, NAME_SOURCE_KEY, + SessionStatus, SessionView, StructuredReplayEvent, WorktreeSpec, NAME_KEY, NAME_SOURCE_KEY, }; use tokio::sync::{broadcast, watch, Mutex}; use uuid::Uuid; @@ -357,6 +357,60 @@ impl SessionManager { .map_err(SessionError::Sidecar) } + /// Replay durable structured-session data as client-neutral projections. + /// + /// The complete log is folded before selecting the requested cursor range, + /// so each returned state is exactly the state at its cursor. + pub async fn read_acp_projection_replay( + &self, + id: Uuid, + cursor: EventCursor, + limit: usize, + ) -> Result>, SessionError> { + let stream = { + let sessions = self.inner.sessions.lock().await; + let record = sessions.get(&id).ok_or(SessionError::NotFound(id))?; + if record.session.view != SessionView::Structured { + return Ok(None); + } + record.acp_events.clone() + }; + let Some(stream) = stream else { + return Ok(None); + }; + let stream = stream.lock().await; + let events = self + .inner + .sidecar + .read_events(id, EventCursor::Since(0), usize::MAX)?; + let mut projection = Projection::default(); + let mut replay = Vec::new(); + for event in events { + let update = projection.apply(event.seq, &event.notification); + let include = match cursor { + EventCursor::Since(since) => event.seq > since, + EventCursor::Before(Some(before)) => event.seq < before, + EventCursor::Before(None) => true, + }; + if include { + replay.push(StructuredReplayEvent { + cursor: event.seq, + transcript: update.transcript, + state: update.state, + }); + } + } + drop(stream); + match cursor { + EventCursor::Since(_) => replay.truncate(limit), + EventCursor::Before(_) if replay.len() > limit => { + replay.drain(..replay.len() - limit); + } + _ => {} + } + Ok(Some(replay)) + } + /// Subscribe to the two projections of a structured session's ACP log. /// /// The transcript snapshot and whole reduced state are folded once by the @@ -402,6 +456,7 @@ impl SessionManager { &self, id: Uuid, lifecycle: TurnLifecycle, + turn: u64, ) -> Result { let stream = { let sessions = self.inner.sessions.lock().await; @@ -412,14 +467,19 @@ impl SessionManager { return Ok(false); }; let mut stream = stream.lock().await; - let projected = stream.projection.turn_lifecycle(lifecycle); + let Some(projected) = stream.projection.turn_lifecycle(lifecycle, turn) else { + return Ok(false); + }; let _ = stream.projected_live.send(projected); Ok(true) } /// Begin an ACP turn only when the previous one has ended. The projection /// lock makes the rejection and lifecycle update one indivisible action. - pub async fn begin_acp_turn(&self, id: Uuid) -> Result>, SessionError> { + pub async fn begin_acp_turn( + &self, + id: Uuid, + ) -> Result, u64)>, SessionError> { let (acp, stream) = { let sessions = self.inner.sessions.lock().await; let record = sessions.get(&id).ok_or(SessionError::NotFound(id))?; @@ -432,24 +492,32 @@ impl SessionManager { if stream.projection.turn_active() { return Ok(None); } - let update = stream.projection.turn_lifecycle(TurnLifecycle::Begin); + let (turn, update) = stream.projection.begin_turn(); let _ = stream.projected_live.send(update); - Ok(Some(acp)) + Ok(Some((acp, turn))) } /// Cancel an ACP turn and immediately clear its client-facing lifecycle. pub async fn cancel_acp_turn(&self, id: Uuid) -> Result { - let acp = { + let (acp, stream) = { let sessions = self.inner.sessions.lock().await; let record = sessions.get(&id).ok_or(SessionError::NotFound(id))?; - record.acp.clone() + (record.acp.clone(), record.acp_events.clone()) }; - let Some(acp) = acp else { + let (Some(acp), Some(stream)) = (acp, stream) else { return Ok(false); }; + let mut stream = stream.lock().await; + let Some(turn) = stream.projection.active_turn() else { + return Ok(true); + }; acp.cancel(); - self.apply_acp_turn_lifecycle(id, TurnLifecycle::Cancel) - .await + let projected = stream + .projection + .turn_lifecycle(TurnLifecycle::Cancel, turn) + .expect("active turn remains current while stream is locked"); + let _ = stream.projected_live.send(projected); + Ok(true) } /// End every adapter child before the daemon exits (#134): each gets diff --git a/crates/maestro-daemon/tests/acp.rs b/crates/maestro-daemon/tests/acp.rs index 31ad938..5340bb0 100644 --- a/crates/maestro-daemon/tests/acp.rs +++ b/crates/maestro-daemon/tests/acp.rs @@ -232,11 +232,12 @@ async fn projection_snapshot_and_live_stream_share_one_fold() { assert_eq!(snapshot.rows[0].kind, TranscriptRowKind::AgentText); assert_eq!(state.turn, Default::default()); - assert!(fx + let (_, turn) = fx .manager - .apply_acp_turn_lifecycle(session.id, TurnLifecycle::Begin) + .begin_acp_turn(session.id) .await - .expect("begin lifecycle")); + .expect("begin lifecycle") + .expect("structured turn"); let begin = tokio::time::timeout(Duration::from_secs(5), live.recv()) .await .expect("begin lifecycle broadcast") @@ -264,7 +265,7 @@ async fn projection_snapshot_and_live_stream_share_one_fold() { assert!(fx .manager - .apply_acp_turn_lifecycle(session.id, TurnLifecycle::Finish) + .apply_acp_turn_lifecycle(session.id, TurnLifecycle::Finish, turn) .await .expect("finish lifecycle")); let finish = tokio::time::timeout(Duration::from_secs(5), live.recv()) @@ -277,6 +278,51 @@ async fn projection_snapshot_and_live_stream_share_one_fold() { fx.manager.kill(session.id, false).await.expect("kill"); } +#[tokio::test] +async fn stale_turn_completion_cannot_end_the_next_turn() { + let fx = DaemonFixture::new(); + let fake = Fake::new(); + let cwd = TempDir::new().unwrap(); + let session = fx + .manager + .create(structured(cwd.path(), fake.env())) + .await + .unwrap(); + + let (_, first) = fx + .manager + .begin_acp_turn(session.id) + .await + .unwrap() + .unwrap(); + assert!(fx.manager.cancel_acp_turn(session.id).await.unwrap()); + let (_, second) = fx + .manager + .begin_acp_turn(session.id) + .await + .unwrap() + .unwrap(); + assert!(second > first); + assert!(!fx + .manager + .apply_acp_turn_lifecycle(session.id, TurnLifecycle::Finish, first) + .await + .unwrap()); + assert!(fx + .manager + .begin_acp_turn(session.id) + .await + .unwrap() + .is_none()); + assert!(fx + .manager + .apply_acp_turn_lifecycle(session.id, TurnLifecycle::Finish, second) + .await + .unwrap()); + + fx.manager.kill(session.id, false).await.unwrap(); +} + #[tokio::test] async fn recovery_loads_the_acp_session_without_duplicating_history() { let fx = DaemonFixture::new(); diff --git a/crates/maestro-daemon/tests/local_api.rs b/crates/maestro-daemon/tests/local_api.rs index 3355808..188911c 100644 --- a/crates/maestro-daemon/tests/local_api.rs +++ b/crates/maestro-daemon/tests/local_api.rs @@ -110,24 +110,49 @@ impl FakeAcp { } fn env(&self) -> HashMap { - HashMap::from([( - "PATH".into(), - format!( - "{}:{}", - self.dir.path().join("bin").display(), - std::env::var("PATH").unwrap() + HashMap::from([ + ( + "PATH".into(), + format!( + "{}:{}", + self.dir.path().join("bin").display(), + std::env::var("PATH").unwrap() + ), ), - )]) + ( + "FAKE_ACP_LOG".into(), + self.dir.path().join("agent.log").display().to_string(), + ), + ]) + } + + fn pid(&self) -> u32 { + std::fs::read_to_string(self.dir.path().join("agent.log")) + .unwrap_or_default() + .lines() + .filter_map(|line| { + line.strip_prefix("start pid=")? + .split(' ') + .next()? + .parse() + .ok() + }) + .next_back() + .expect("fake ACP started") } } async fn structured_session(h: &Harness, fake: &FakeAcp) -> Uuid { + structured_session_with_env(h, fake.env()).await +} + +async fn structured_session_with_env(h: &Harness, env: HashMap) -> Uuid { h.manager .create(CreateSessionRequest { view: SessionView::Structured, cwd: "/tmp".into(), command: "claude".into(), - env: fake.env(), + env, metadata: HashMap::new(), worktree: None, }) @@ -183,7 +208,9 @@ async fn structured_ws_sends_snapshot_then_live_projection() { async fn structured_prompt_rejects_a_second_in_flight_turn_and_cancel_clears_it() { let h = spawn("naps-laptop").await; let fake = FakeAcp::new(); - let id = structured_session(&h, &fake).await; + let mut env = fake.env(); + env.insert("FAKE_ACP_DELAY_PROMPT_MS".into(), "200".into()); + let id = structured_session_with_env(&h, env).await; let client = reqwest::Client::new(); let first = client .post(url( @@ -229,6 +256,32 @@ async fn structured_prompt_rejects_a_second_in_flight_turn_and_cancel_clears_it( assert_eq!(next.status(), reqwest::StatusCode::NO_CONTENT); } +#[tokio::test] +async fn structured_prompt_failure_clears_the_turn() { + let h = spawn("naps-laptop").await; + let fake = FakeAcp::new(); + let mut env = fake.env(); + env.insert("FAKE_ACP_PROMPT_FAIL".into(), "1".into()); + let id = structured_session_with_env(&h, env).await; + let client = reqwest::Client::new(); + let path = format!("/api/sessions/{id}/structured/prompt"); + let first = client + .post(url(h.addr, &path)) + .json(&serde_json::json!({"prompt":"fails"})) + .send() + .await + .unwrap(); + assert_eq!(first.status(), reqwest::StatusCode::NO_CONTENT); + tokio::time::sleep(Duration::from_millis(50)).await; + let second = client + .post(url(h.addr, &path)) + .json(&serde_json::json!({"prompt":"available"})) + .send() + .await + .unwrap(); + assert_ne!(second.status(), reqwest::StatusCode::CONFLICT); +} + #[tokio::test] async fn structured_replay_honors_cursors_and_limit() { let h = spawn("naps-laptop").await; @@ -258,7 +311,8 @@ async fn structured_replay_honors_cursors_and_limit() { .await .unwrap(); assert_eq!(events.len(), 1); - assert_eq!(events[0]["seq"], 1); + assert_eq!(events[0]["cursor"], 1); + assert!(events[0].get("notification").is_none()); let invalid = client .get(url( @@ -271,6 +325,92 @@ async fn structured_replay_honors_cursors_and_limit() { assert_eq!(invalid.status(), reqwest::StatusCode::BAD_REQUEST); } +#[tokio::test] +async fn structured_replay_pages_projected_events_in_cursor_order() { + let h = spawn("naps-laptop").await; + let fake = FakeAcp::new(); + let id = structured_session(&h, &fake).await; + let client = reqwest::Client::new(); + unsafe { libc::kill(fake.pid() as libc::pid_t, libc::SIGUSR1) }; + let events: Vec = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let events: Vec = client + .get(url( + h.addr, + &format!("/api/sessions/{id}/structured/replay?before=99&limit=2"), + )) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + if events.len() == 2 { + return events; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert_eq!( + events + .iter() + .map(|event| event["cursor"].as_u64().unwrap()) + .collect::>(), + vec![1, 2] + ); + let boundary: Vec = client + .get(url( + h.addr, + &format!("/api/sessions/{id}/structured/replay?before=2&limit=10"), + )) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(boundary.len(), 1); + assert_eq!(boundary[0]["cursor"], 1); +} + +#[tokio::test] +async fn structured_replay_rejects_unknown_and_terminal_sessions() { + let h = spawn("naps-laptop").await; + let client = reqwest::Client::new(); + let unknown = client + .get(url( + h.addr, + &format!("/api/sessions/{}/structured/replay", Uuid::new_v4()), + )) + .send() + .await + .unwrap(); + assert_eq!(unknown.status(), reqwest::StatusCode::NOT_FOUND); + let terminal = h + .manager + .create(CreateSessionRequest { + view: SessionView::Terminal, + cwd: "/tmp".into(), + command: "sh".into(), + env: HashMap::new(), + metadata: HashMap::new(), + worktree: None, + }) + .await + .unwrap(); + let response = client + .get(url( + h.addr, + &format!("/api/sessions/{}/structured/replay", terminal.id), + )) + .send() + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::CONFLICT); +} + #[tokio::test] async fn info_returns_daemon_identity() { let h = spawn("naps-laptop").await; diff --git a/crates/maestro-protocol/src/lib.rs b/crates/maestro-protocol/src/lib.rs index 3f06413..d7b62d2 100644 --- a/crates/maestro-protocol/src/lib.rs +++ b/crates/maestro-protocol/src/lib.rs @@ -794,6 +794,19 @@ pub struct TranscriptSnapshot { pub rows: Vec, } +/// One durable, client-neutral replay step from a structured session. +/// +/// Results are ordered by increasing `cursor`. `since` returns the first +/// matching steps; `before` returns the final matching steps, still increasing. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct StructuredReplayEvent { + /// Durable ACP-log position represented by this projected step. + pub cursor: u64, + #[serde(skip_serializing_if = "Option::is_none")] + pub transcript: Option, + pub state: ReducedState, +} + /// A change to a transcript row. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")]