fix(maestro): key activity watermark by stream

crates/maestro-daemon/src/session.rs

@@ -10,7 +10,7 @@ use std::future::Future;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::task::{Context, Poll};
 use std::time::Duration;
 
@@ -102,8 +102,12 @@ pub struct SessionRecord {
     /// outside the sessions lock. Daemon-local; never rides the wire.
     pub acp: Option<Arc<AcpSession>>,
     /// Open log plus its post-append live sender. Their shared lock is the
-    /// replay/live seam used by the structured WebSocket in #138.
+    /// replay/live seam used by the structured WebSocket in #138. Callers may
+    /// clone this handle under `sessions`, but must never hold either guard
+    /// while acquiring the other.
     acp_events: Option<Arc<Mutex<AcpEventStream>>>,
+    /// Latest activity snapshot applied, keyed to its originating stream.
+    structured_activity_watermark: Option<StructuredActivityWatermark>,
     /// Completion signal for the task draining ACP notifications into the
     /// event log. Shutdown waits for it before returning.
     acp_event_pump: Option<watch::Receiver<bool>>,
@@ -118,6 +122,12 @@ impl SessionRecord {
     }
 }
 
+#[derive(Debug)]
+struct StructuredActivityWatermark {
+    stream: Weak<Mutex<AcpEventStream>>,
+    sequence: u64,
+}
+
 #[derive(Debug)]
 struct AcpEventStream {
     log: EventLog,
@@ -129,6 +139,14 @@ struct AcpEventStream {
     /// claimed: it is queued onto ACP before another operation can proceed.
     operation: Arc<Mutex<()>>,
     dispatch: AcpDispatchState,
+    activity_sequence: u64,
+}
+
+impl AcpEventStream {
+    fn activity_snapshot(&mut self) -> (u64, ActivityState) {
+        self.activity_sequence += 1;
+        (self.activity_sequence, self.projection.activity())
+    }
 }
 
 #[derive(Debug, Default)]
@@ -1220,6 +1238,7 @@ impl SessionManager {
                         pending_approvals: HashMap::new(),
                         operation: Arc::new(Mutex::new(())),
                         dispatch: AcpDispatchState::default(),
+                        activity_sequence: 0,
                     })))
                 }
                 Err(e) => {
@@ -1323,6 +1342,7 @@ impl SessionManager {
                     awaiting_suppressed: false,
                     acp,
                     acp_events,
+                    structured_activity_watermark: None,
                     acp_event_pump: event_pump,
                 },
             );
@@ -1534,6 +1554,7 @@ impl SessionManager {
                     // daemon; only tmux-backed sessions can be adopted.
                     acp: None,
                     acp_events: None,
+                    structured_activity_watermark: None,
                     acp_event_pump: None,
                 },
             );
@@ -1655,6 +1676,7 @@ impl SessionManager {
                             pending_approvals: HashMap::new(),
                             operation: Arc::new(Mutex::new(())),
                             dispatch: AcpDispatchState::default(),
+                            activity_sequence: 0,
                         })))
                     }
                     Err(error) => {
@@ -1702,6 +1724,7 @@ impl SessionManager {
                     awaiting_suppressed: false,
                     acp,
                     acp_events,
+                    structured_activity_watermark: None,
                     acp_event_pump: None,
                 },
             );
@@ -2122,57 +2145,59 @@ async fn sync_structured_activity_from_stream(
     id: Uuid,
     events: Arc<Mutex<AcpEventStream>>,
 ) {
-    let event_stream = events.lock().await;
-    let raw = event_stream.projection.activity();
-    // Stream state must stay fixed until its catalog transition is applied;
-    // every other path releases `sessions` before awaiting this stream lock.
-    apply_structured_activity(inner, id, Some(&events), raw).await;
+    let (sequence, raw) = events.lock().await.activity_snapshot();
+    apply_structured_activity(inner, id, Some((&events, sequence)), raw).await;
 }
 
 async fn apply_structured_activity(
     inner: &Arc<Inner>,
     id: Uuid,
-    expected_events: Option<&Arc<Mutex<AcpEventStream>>>,
+    expected_events: Option<(&Arc<Mutex<AcpEventStream>>, u64)>,
     raw: ActivityState,
 ) {
-    let updated = {
-        let mut sessions = inner.sessions.lock().await;
-        let Some(record) = sessions.get_mut(&id) else {
-            return;
-        };
-        if record.session.view != SessionView::Structured {
-            return;
-        }
-        let current_events = record.acp_events.as_ref();
-        let stream_matches = match (expected_events, current_events) {
-            (Some(expected), Some(current)) => Arc::ptr_eq(expected, current),
-            (None, None) => true,
-            _ => false,
-        };
-        if !stream_matches {
-            return;
-        }
-        let raw = if record.session.status == SessionStatus::Dead {
-            ActivityState::Error
-        } else {
-            raw
-        };
-        let activity = crate::activity::apply_focus_suppression(
-            raw,
-            record.focused_since,
-            std::time::Instant::now(),
-            &mut record.awaiting_suppressed,
-        );
-        if record.session.activity == activity {
-            None
-        } else {
-            record.session.activity = activity;
-            record.session.activity_since = Utc::now();
-            Some(record.session.clone())
-        }
+    let mut sessions = inner.sessions.lock().await;
+    let Some(record) = sessions.get_mut(&id) else {
+        return;
     };
-    if let Some(session) = updated {
-        let _ = inner.events.send(CatalogEvent::SessionUpdated { session });
+    if record.session.view != SessionView::Structured {
+        return;
+    }
+    let current_events = record.acp_events.as_ref();
+    match (expected_events, current_events) {
+        (Some((expected, sequence)), Some(current)) if Arc::ptr_eq(expected, current) => {
+            let stream = Arc::downgrade(current);
+            if record
+                .structured_activity_watermark
+                .as_ref()
+                .is_some_and(|watermark| {
+                    watermark.stream.ptr_eq(&stream) && sequence < watermark.sequence
+                })
+            {
+                return;
+            }
+            record.structured_activity_watermark =
+                Some(StructuredActivityWatermark { stream, sequence });
+        }
+        (None, None) => {}
+        _ => return,
+    }
+    let raw = if record.session.status == SessionStatus::Dead {
+        ActivityState::Error
+    } else {
+        raw
+    };
+    let activity = crate::activity::apply_focus_suppression(
+        raw,
+        record.focused_since,
+        std::time::Instant::now(),
+        &mut record.awaiting_suppressed,
+    );
+    if record.session.activity != activity {
+        record.session.activity = activity;
+        record.session.activity_since = Utc::now();
+        let _ = inner.events.send(CatalogEvent::SessionUpdated {
+            session: record.session.clone(),
+        });
     }
 }
 
@@ -2849,16 +2874,17 @@ mod tests {
         PermissionOption, PermissionOptionKind, RequestPermissionOutcome, RequestPermissionRequest,
         ToolCallUpdate, ToolCallUpdateFields,
     };
-    use maestro_protocol::{ActivityState, CreateSessionRequest, SessionView};
+    use maestro_protocol::{ActivityState, CatalogEvent, CreateSessionRequest, SessionView};
 
     use super::{
-        cancel_pending_approvals, permission_outcome, record_permission_request,
-        run_structured_activity_pass, sync_structured_activity_from_stream, AcpDispatchState,
-        AcpEventStream, ApprovalDecision, RecordOutcome, SessionManager,
+        apply_structured_activity, cancel_pending_approvals, permission_outcome,
+        record_permission_request, run_structured_activity_pass,
+        sync_structured_activity_from_stream, AcpDispatchState, AcpEventStream, ApprovalDecision,
+        RecordOutcome, SessionManager,
     };
     use crate::activity::FOCUS_CLEAR_HOLD;
     use crate::panes::FakePanes;
-    use crate::projection::Projection;
+    use crate::projection::{Projection, TurnLifecycle};
     use crate::sidecar::{
         EventCursor, EventLog, FaultyLogSink, PersistedEvent, SharedBuf, SidecarStore,
     };
@@ -2887,6 +2913,7 @@ mod tests {
             pending_approvals: std::collections::HashMap::new(),
             operation: std::sync::Arc::new(tokio::sync::Mutex::new(())),
             dispatch: AcpDispatchState::default(),
+            activity_sequence: 0,
         }
     }
 
@@ -3043,6 +3070,87 @@ mod tests {
             .expect("terminal resize remains available");
     }
 
+    #[tokio::test]
+    async fn structured_activity_releases_stream_before_waiting_for_catalog() {
+        let (_state, manager, _panes, structured_id, events, _terminal_id) =
+            manager_with_structured_stream().await;
+        let sessions = manager.inner.sessions.lock().await;
+        let sync =
+            sync_structured_activity_from_stream(&manager.inner, structured_id, events.clone());
+        tokio::pin!(sync);
+        assert!(matches!(futures_util::poll!(sync.as_mut()), Poll::Pending));
+
+        assert!(
+            events.try_lock().is_ok(),
+            "activity sync held the event stream while waiting for the catalog"
+        );
+
+        drop(sessions);
+        tokio::time::timeout(Duration::from_secs(1), sync)
+            .await
+            .expect("activity sync did not finish");
+    }
+
+    #[tokio::test]
+    async fn newer_structured_activity_snapshot_cannot_be_overtaken() {
+        let (_state, manager, _panes, structured_id, events, _terminal_id) =
+            manager_with_structured_stream().await;
+        let mut catalog = manager.events();
+        manager
+            .inner
+            .sessions
+            .lock()
+            .await
+            .get_mut(&structured_id)
+            .unwrap()
+            .focused_since = None;
+
+        let (older_sequence, older_activity) = events.lock().await.activity_snapshot();
+        let (newer_sequence, newer_activity) = {
+            let mut event_stream = events.lock().await;
+            let turn = event_stream.projection.active_turn().unwrap();
+            event_stream.projection.clear_approval();
+            event_stream
+                .projection
+                .turn_lifecycle(TurnLifecycle::Finish, turn);
+            event_stream.activity_snapshot()
+        };
+
+        apply_structured_activity(
+            &manager.inner,
+            structured_id,
+            Some((&events, newer_sequence)),
+            newer_activity,
+        )
+        .await;
+        apply_structured_activity(
+            &manager.inner,
+            structured_id,
+            Some((&events, older_sequence)),
+            older_activity,
+        )
+        .await;
+
+        let CatalogEvent::SessionUpdated { session } = catalog.try_recv().unwrap() else {
+            panic!("unexpected catalog event");
+        };
+        assert_eq!(session.activity, ActivityState::Idle);
+        assert!(matches!(
+            catalog.try_recv(),
+            Err(tokio::sync::broadcast::error::TryRecvError::Empty)
+        ));
+        assert_eq!(
+            manager
+                .snapshot()
+                .await
+                .into_iter()
+                .find(|session| session.id == structured_id)
+                .unwrap()
+                .activity,
+            ActivityState::Idle
+        );
+    }
+
     #[tokio::test]
     async fn delayed_activity_from_a_replaced_stream_is_discarded() {
         let (_state, manager, _panes, structured_id, old_events, _terminal_id) =
@@ -3076,6 +3184,73 @@ mod tests {
         ));
     }
 
+    #[tokio::test]
+    async fn replacement_accepts_new_activity_after_a_high_old_watermark() {
+        let (_state, manager, _panes, structured_id, old_events, _terminal_id) =
+            manager_with_structured_stream().await;
+        let replacement = Arc::new(tokio::sync::Mutex::new(stream(
+            manager
+                .inner
+                .sidecar
+                .open_events(uuid::Uuid::new_v4())
+                .expect("open replacement event log"),
+        )));
+        manager
+            .inner
+            .sessions
+            .lock()
+            .await
+            .get_mut(&structured_id)
+            .unwrap()
+            .focused_since = None;
+        let _ = old_events.lock().await.activity_snapshot();
+        let (old_high_sequence, old_high_activity) = old_events.lock().await.activity_snapshot();
+        apply_structured_activity(
+            &manager.inner,
+            structured_id,
+            Some((&old_events, old_high_sequence)),
+            old_high_activity,
+        )
+        .await;
+        let (delayed_old_sequence, delayed_old_activity) =
+            old_events.lock().await.activity_snapshot();
+
+        {
+            let mut sessions = manager.inner.sessions.lock().await;
+            let record = sessions.get_mut(&structured_id).unwrap();
+            record.acp_events = Some(replacement.clone());
+            record.session.activity = ActivityState::Error;
+        }
+
+        let (new_sequence, new_activity) = replacement.lock().await.activity_snapshot();
+        assert!(new_sequence < old_high_sequence);
+        apply_structured_activity(
+            &manager.inner,
+            structured_id,
+            Some((&replacement, new_sequence)),
+            new_activity,
+        )
+        .await;
+        apply_structured_activity(
+            &manager.inner,
+            structured_id,
+            Some((&old_events, delayed_old_sequence)),
+            delayed_old_activity,
+        )
+        .await;
+
+        let sessions = manager.inner.sessions.lock().await;
+        let record = sessions.get(&structured_id).unwrap();
+        assert_eq!(record.session.activity, ActivityState::Idle);
+        let watermark = record.structured_activity_watermark.as_ref().unwrap();
+        assert_eq!(watermark.sequence, new_sequence);
+        assert!(watermark.stream.ptr_eq(&Arc::downgrade(&replacement)));
+        assert!(Arc::ptr_eq(
+            record.acp_events.as_ref().unwrap(),
+            &replacement
+        ));
+    }
+
     /// A handle that can never append *or* roll back: opened read-only, so
     /// both `write_all` and the rollback `truncate` fail with the OS's own
     /// error (`ftruncate` requires a writable fd too) instead of a