fix: harden structured replay

This commit is contained in:
Miguel Palhas
2026-09-05 01:56:07 +01:00
parent b98f5cb3e5
commit be23a2b539
6 changed files with 369 additions and 68 deletions
+24 -13
View File
@@ -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<Uuid>,
Json(body): Json<StructuredPromptRequest>,
) -> 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)
}
+54 -31
View File
@@ -35,7 +35,8 @@ pub struct Projection {
tools: HashMap<String, ToolCall>,
pending_tool_updates: HashMap<String, Vec<maestro_acp::schema::ToolCallUpdateFields>>,
state: ReducedState,
turn_active: bool,
active_turn: Option<u64>,
next_turn: u64,
turn_tool_calls: HashSet<String>,
}
@@ -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<ProjectionUpdate> {
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<u64> {
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!({
+78 -10
View File
@@ -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<Option<Vec<StructuredReplayEvent>>, 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<bool, SessionError> {
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<Option<Arc<AcpSession>>, SessionError> {
pub async fn begin_acp_turn(
&self,
id: Uuid,
) -> Result<Option<(Arc<AcpSession>, 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<bool, SessionError> {
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
+50 -4
View File
@@ -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();
+150 -10
View File
@@ -110,24 +110,49 @@ impl FakeAcp {
}
fn env(&self) -> HashMap<String, String> {
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<String, String>) -> 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<serde_json::Value> = tokio::time::timeout(Duration::from_secs(2), async {
loop {
let events: Vec<serde_json::Value> = 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<_>>(),
vec![1, 2]
);
let boundary: Vec<serde_json::Value> = 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;
+13
View File
@@ -794,6 +794,19 @@ pub struct TranscriptSnapshot {
pub rows: Vec<TranscriptRow>,
}
/// 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<TranscriptDelta>,
pub state: ReducedState,
}
/// A change to a transcript row.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]