Make task state store-owned and signal its changes to the lead

internal/agents/snapshots.go

@@ -43,10 +43,12 @@ func (m *Manager) snapshotSoon(l *live, trigger api.SnapshotTrigger, seq int64,
 
 	st := m.snapStateLocked(l.taskID)
 	if st.closed {
+		m.dropRequest(l.taskID, &snapRequest{review: review})
 		return
 	}
 	req := &snapRequest{trigger: trigger, seq: seq, review: review}
 	if st.running {
+		m.dropRequest(l.taskID, st.pending)
 		st.pending = req
 		return
 	}
@@ -66,8 +68,13 @@ func (m *Manager) runSnapshots(l *live, st *snapState, req *snapRequest) {
 				slog.Error("take snapshot", "task", l.taskID, "trigger", req.trigger, "err", err)
 			}
 			if req.review {
-				go m.afterSnapshot(snap)
+				go func() {
+					m.afterSnapshot(snap)
+					m.endReviewRequest(l.taskID)
+				}()
 			}
+		} else {
+			m.dropRequest(l.taskID, req)
 		}
 		st.mu.Unlock()
 
@@ -75,6 +82,7 @@ func (m *Manager) runSnapshots(l *live, st *snapState, req *snapRequest) {
 		req = st.pending
 		st.pending = nil
 		if req == nil || st.closed {
+			m.dropRequest(l.taskID, req)
 			st.running = false
 			m.snapMu.Unlock()
 			return
@@ -83,6 +91,19 @@ func (m *Manager) runSnapshots(l *live, st *snapState, req *snapRequest) {
 	}
 }
 
+func (m *Manager) dropRequest(taskID string, req *snapRequest) {
+	if req != nil && req.review {
+		m.endReviewRequest(taskID)
+	}
+}
+
+func (m *Manager) endReviewRequest(taskID string) {
+	err := m.store.EndReviewRequest(context.Background(), taskID)
+	if err != nil {
+		slog.Error("end review request", "task", taskID, "err", err)
+	}
+}
+
 func (m *Manager) takeSnapshot(ctx context.Context, l *live, trigger api.SnapshotTrigger, seq int64) (*api.Snapshot, error) {
 	if seq == 0 {
 		var err error