agent_engine/runtime/
capture_worker.rs

1//! Nonblocking, bounded dispatch of completed memory captures.
2
3use super::chat_capture::{
4    build_chat_turn_capture, build_conversation_summary_capture, CaptureBuildError,
5    ChatTurnCapture, ConversationSummaryCapture, SummaryCaptureBuildError, TerminalTurnHistory,
6};
7use super::memory_context::{ContextProviderId, MemoryContextLease, RetentionClass, TurnCapture};
8use std::collections::HashSet;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::{mpsc, Arc};
11use std::thread;
12
13/// Outcome of an idempotency-key query after an ambiguous capture call.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum CaptureCommitState {
16    Committed,
17    Absent,
18}
19
20/// Provider seam. Implementations must treat every capture id as an
21/// idempotency key. Summary capture is additive so existing turn providers
22/// keep source compatibility while adopting the C4 record class.
23pub trait CaptureProvider: Send + Sync + 'static {
24    fn capture(&self, capture: ChatTurnCapture) -> Result<(), CaptureFailure>;
25
26    /// Query whether this idempotency key is already durable after an
27    /// interrupted/ambiguous call. Providers that cannot query return false;
28    /// no automatic blind retry is performed in either case.
29    fn contains_capture(
30        &self,
31        _capture_id: &[u8; 32],
32    ) -> Result<CaptureCommitState, CaptureFailure> {
33        Ok(CaptureCommitState::Absent)
34    }
35
36    fn capture_summary(&self, _capture: ConversationSummaryCapture) -> Result<(), CaptureFailure> {
37        Err(CaptureFailure {
38            code: "summary_capture_unsupported",
39        })
40    }
41}
42
43/// Content-free failure metadata suitable for diagnostics.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct CaptureFailure {
46    pub code: &'static str,
47}
48
49enum CapturePayload {
50    Turn(Box<ChatTurnCapture>),
51    Summary(Box<ConversationSummaryCapture>),
52}
53
54impl CapturePayload {
55    fn id_bytes(&self) -> [u8; 32] {
56        match self {
57            Self::Turn(capture) => *capture.capture_id.as_bytes(),
58            Self::Summary(capture) => *capture.capture_id.as_bytes(),
59        }
60    }
61
62    fn dispatch(&self, provider: &dyn CaptureProvider) -> Result<(), CaptureFailure> {
63        match self {
64            Self::Turn(capture) => provider.capture(capture.as_ref().clone()),
65            Self::Summary(capture) => provider.capture_summary(capture.as_ref().clone()),
66        }
67    }
68
69    fn query_committed(
70        &self,
71        provider: &dyn CaptureProvider,
72    ) -> Result<CaptureCommitState, CaptureFailure> {
73        match self {
74            Self::Turn(capture) => provider.contains_capture(capture.capture_id.as_bytes()),
75            // Summary retry/query support is additive and not part of the C5
76            // terminal-turn crash gate.
77            Self::Summary(_) => Ok(CaptureCommitState::Absent),
78        }
79    }
80}
81
82struct CaptureJob {
83    provider_id: ContextProviderId,
84    provider: Arc<dyn CaptureProvider>,
85    payload: CapturePayload,
86    /// Present for imports that may advance a durable host checkpoint only
87    /// after the provider has acknowledged (or reconciled) the capture.
88    committed: Option<mpsc::Sender<Result<(), CaptureFailure>>>,
89}
90
91/// Fixed-capacity worker. Enqueue is always `try_send`; persistence never runs
92/// on the turn-completion or compaction-transition path.
93pub struct CaptureWorker {
94    sender: mpsc::SyncSender<CaptureJob>,
95    dropped: Arc<AtomicU64>,
96    failures: Arc<AtomicU64>,
97}
98
99impl CaptureWorker {
100    pub fn new(capacity: usize) -> Self {
101        assert!(capacity > 0, "capture queue capacity must be nonzero");
102        let (sender, receiver) = mpsc::sync_channel::<CaptureJob>(capacity);
103        let dropped = Arc::new(AtomicU64::new(0));
104        let failures = Arc::new(AtomicU64::new(0));
105        let worker_failures = Arc::clone(&failures);
106        thread::Builder::new()
107            .name("memory-capture".into())
108            .spawn(move || {
109                let mut submitted = HashSet::new();
110                while let Ok(job) = receiver.recv() {
111                    let id = (job.provider_id.clone(), job.payload.id_bytes());
112                    let result = if submitted.contains(&id) {
113                        Ok(())
114                    } else if let Err(error) = job.payload.dispatch(job.provider.as_ref()) {
115                        // The call may have committed before its acknowledgement
116                        // was lost. Reconcile by idempotency key before accepting
117                        // an explicit retry; never blindly dispatch it again.
118                        match job.payload.query_committed(job.provider.as_ref()) {
119                            Ok(CaptureCommitState::Committed) => {
120                                submitted.insert(id);
121                                Ok(())
122                            }
123                            Ok(CaptureCommitState::Absent) | Err(_) => {
124                                worker_failures.fetch_add(1, Ordering::Relaxed);
125                                Err(error)
126                            }
127                        }
128                    } else {
129                        submitted.insert(id);
130                        Ok(())
131                    };
132                    if let Some(committed) = job.committed {
133                        let _ = committed.send(result);
134                    }
135                }
136            })
137            .expect("capture worker spawn");
138        Self {
139            sender,
140            dropped,
141            failures,
142        }
143    }
144
145    /// Apply lease capability and capture gates, build the bounded payload, and
146    /// attempt a nonblocking enqueue. `Ok(false)` means capture was gated off.
147    pub fn submit_terminal(
148        &self,
149        lease: &MemoryContextLease,
150        history: TerminalTurnHistory,
151        retention: RetentionClass,
152        provider: Arc<dyn CaptureProvider>,
153    ) -> Result<bool, CaptureBuildError> {
154        if lease.mode.turn_capture() != TurnCapture::Enabled
155            || lease.project_id != history.project_id
156            || lease.session_id != history.session_id
157        {
158            return Ok(false);
159        }
160        let capture = build_chat_turn_capture(&lease.project_id, history, retention)?;
161        let job = CaptureJob {
162            provider_id: lease.provider_id.clone(),
163            provider,
164            payload: CapturePayload::Turn(Box::new(capture)),
165            committed: None,
166        };
167        match self.sender.try_send(job) {
168            Ok(()) => Ok(true),
169            Err(mpsc::TrySendError::Full(_)) | Err(mpsc::TrySendError::Disconnected(_)) => {
170                self.dropped.fetch_add(1, Ordering::Relaxed);
171                Ok(true)
172            }
173        }
174    }
175
176    /// Submit a capture already built by the C1 builder. Historical import
177    /// uses this narrow seam so it shares the same bounded C3 queue and exact
178    /// leased provider dispatch as live terminal capture.
179    pub(crate) fn submit_built(
180        &self,
181        lease: &MemoryContextLease,
182        capture: ChatTurnCapture,
183        provider: Arc<dyn CaptureProvider>,
184    ) -> bool {
185        self.submit_built_job(lease, capture, provider, None)
186    }
187
188    /// Submit one historical capture with a completion receipt. Waiting for the
189    /// receipt happens only on the explicit import path, never turn completion.
190    pub(crate) fn submit_built_with_receipt(
191        &self,
192        lease: &MemoryContextLease,
193        capture: ChatTurnCapture,
194        provider: Arc<dyn CaptureProvider>,
195        committed: mpsc::Sender<Result<(), CaptureFailure>>,
196    ) -> bool {
197        self.submit_built_job(lease, capture, provider, Some(committed))
198    }
199
200    fn submit_built_job(
201        &self,
202        lease: &MemoryContextLease,
203        capture: ChatTurnCapture,
204        provider: Arc<dyn CaptureProvider>,
205        committed: Option<mpsc::Sender<Result<(), CaptureFailure>>>,
206    ) -> bool {
207        if lease.mode.turn_capture() != TurnCapture::Enabled
208            || lease.project_id != capture.project_id
209        {
210            return false;
211        }
212        let job = CaptureJob {
213            provider_id: lease.provider_id.clone(),
214            provider,
215            payload: CapturePayload::Turn(Box::new(capture)),
216            committed,
217        };
218        match self.sender.try_send(job) {
219            Ok(()) => true,
220            Err(mpsc::TrySendError::Full(_)) | Err(mpsc::TrySendError::Disconnected(_)) => {
221                self.dropped.fetch_add(1, Ordering::Relaxed);
222                false
223            }
224        }
225    }
226
227    /// Apply the same exact lease gate as turn capture, build a bounded
228    /// first-class compaction memory, and attempt a nonblocking enqueue.
229    #[allow(clippy::too_many_arguments)] // first-class summary metadata stays explicit at this boundary
230    pub fn submit_summary(
231        &self,
232        lease: &MemoryContextLease,
233        source: super::chat_capture::CompactionSource,
234        source_message_count: usize,
235        summary_text: &str,
236        redaction_policy: agent_core::compaction::RedactionPolicy,
237        retention: RetentionClass,
238        provider: Arc<dyn CaptureProvider>,
239    ) -> Result<bool, SummaryCaptureBuildError> {
240        if lease.mode.turn_capture() != TurnCapture::Enabled
241            || lease.project_id != source.project_id
242        {
243            return Ok(false);
244        }
245        let capture = build_conversation_summary_capture(
246            &lease.project_id,
247            source,
248            source_message_count,
249            summary_text,
250            redaction_policy,
251            retention,
252        )?;
253        let job = CaptureJob {
254            provider_id: lease.provider_id.clone(),
255            provider,
256            payload: CapturePayload::Summary(Box::new(capture)),
257            committed: None,
258        };
259        match self.sender.try_send(job) {
260            Ok(()) => Ok(true),
261            Err(mpsc::TrySendError::Full(_)) | Err(mpsc::TrySendError::Disconnected(_)) => {
262                self.dropped.fetch_add(1, Ordering::Relaxed);
263                Ok(true)
264            }
265        }
266    }
267
268    pub fn overflow_drops(&self) -> u64 {
269        self.dropped.load(Ordering::Relaxed)
270    }
271    pub fn provider_failures(&self) -> u64 {
272        self.failures.load(Ordering::Relaxed)
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279    use crate::runtime::chat_capture::{CanonicalCaptureItem, CaptureContentClass};
280    use crate::runtime::memory_context::{
281        mint_explicit_command_proof, CapturePolicy, MemoryContextMode, MemoryLeaseId, ProjectId,
282        RecallPolicy, SessionId,
283    };
284    use agent_core::core::disclosure::DisclosureClass;
285    use agent_core::TurnOutcome;
286    use std::sync::atomic::AtomicBool;
287    use std::sync::Mutex;
288    use std::time::{Duration, Instant, SystemTime};
289
290    struct ProviderDouble {
291        calls: AtomicU64,
292        queries: AtomicU64,
293        fail: AtomicBool,
294        committed: AtomicBool,
295        block: Mutex<Option<mpsc::Receiver<()>>>,
296    }
297
298    impl CaptureProvider for ProviderDouble {
299        fn capture(&self, _capture: ChatTurnCapture) -> Result<(), CaptureFailure> {
300            self.calls.fetch_add(1, Ordering::SeqCst);
301            if let Some(receiver) = self.block.lock().expect("lock").take() {
302                let _ = receiver.recv();
303            }
304            if self.fail.load(Ordering::SeqCst) {
305                Err(CaptureFailure {
306                    code: "fixture_failure",
307                })
308            } else {
309                self.committed.store(true, Ordering::SeqCst);
310                Ok(())
311            }
312        }
313
314        fn contains_capture(
315            &self,
316            _capture_id: &[u8; 32],
317        ) -> Result<CaptureCommitState, CaptureFailure> {
318            self.queries.fetch_add(1, Ordering::SeqCst);
319            Ok(if self.committed.load(Ordering::SeqCst) {
320                CaptureCommitState::Committed
321            } else {
322                CaptureCommitState::Absent
323            })
324        }
325    }
326
327    fn lease() -> MemoryContextLease {
328        let now = SystemTime::now();
329        MemoryContextLease::grant(
330            MemoryLeaseId::parse("lease-c3").unwrap(),
331            SessionId::parse("session-c3").unwrap(),
332            ProjectId::parse("project-c3").unwrap(),
333            ContextProviderId::parse("provider-c3").unwrap(),
334            MemoryContextMode::CaptureOnly,
335            CapturePolicy::EligibleTurnsOnly,
336            RecallPolicy::BoundedPerPrompt,
337            mint_explicit_command_proof(),
338            now,
339            None,
340        )
341        .unwrap()
342    }
343
344    fn history(turn: u64) -> TerminalTurnHistory {
345        let project_id = ProjectId::parse("project-c3").unwrap();
346        let started_at = SystemTime::UNIX_EPOCH + Duration::from_secs(turn);
347        TerminalTurnHistory {
348            project_id: project_id.clone(),
349            session_id: SessionId::parse("session-c3").unwrap(),
350            turn_id: super::super::memory_context::TurnId::parse(&format!("turn-{turn}")).unwrap(),
351            turn_ordinal: turn,
352            started_at,
353            completed_at: started_at + Duration::from_millis(1),
354            outcome: TurnOutcome::Completed,
355            items: vec![
356                CanonicalCaptureItem {
357                    project_id: project_id.clone(),
358                    class: CaptureContentClass::UserMessage,
359                    disclosure: DisclosureClass::ModelVisible,
360                    sensitivity: super::super::chat_capture::Sensitivity::Normal,
361                    text: "user".into(),
362                    tool_name: None,
363                },
364                CanonicalCaptureItem {
365                    project_id,
366                    class: CaptureContentClass::AssistantFinal,
367                    disclosure: DisclosureClass::ModelVisible,
368                    sensitivity: super::super::chat_capture::Sensitivity::Normal,
369                    text: "assistant".into(),
370                    tool_name: None,
371                },
372            ],
373            compaction: None,
374        }
375    }
376
377    fn wait_until(predicate: impl Fn() -> bool) {
378        let deadline = Instant::now() + Duration::from_secs(1);
379        while !predicate() && Instant::now() < deadline {
380            thread::sleep(Duration::from_millis(2));
381        }
382        assert!(predicate());
383    }
384
385    #[test]
386    fn memory_duplicate_capture_id_idempotent_at_engine_seam() {
387        let provider = Arc::new(ProviderDouble {
388            calls: AtomicU64::new(0),
389            queries: AtomicU64::new(0),
390            fail: AtomicBool::new(false),
391            committed: AtomicBool::new(false),
392            block: Mutex::new(None),
393        });
394        let worker = CaptureWorker::new(4);
395        worker
396            .submit_terminal(
397                &lease(),
398                history(1),
399                RetentionClass::Standard,
400                provider.clone(),
401            )
402            .unwrap();
403        worker
404            .submit_terminal(
405                &lease(),
406                history(1),
407                RetentionClass::Standard,
408                provider.clone(),
409            )
410            .unwrap();
411        wait_until(|| provider.calls.load(Ordering::SeqCst) == 1);
412        thread::sleep(Duration::from_millis(10));
413        assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
414    }
415
416    #[test]
417    fn memory_overflow_drops_exactly_and_never_blocks_completion() {
418        let (release_tx, release_rx) = mpsc::channel();
419        let provider = Arc::new(ProviderDouble {
420            calls: AtomicU64::new(0),
421            queries: AtomicU64::new(0),
422            fail: AtomicBool::new(false),
423            committed: AtomicBool::new(false),
424            block: Mutex::new(Some(release_rx)),
425        });
426        let worker = CaptureWorker::new(1);
427        worker
428            .submit_terminal(
429                &lease(),
430                history(1),
431                RetentionClass::Standard,
432                provider.clone(),
433            )
434            .unwrap();
435        wait_until(|| provider.calls.load(Ordering::SeqCst) == 1);
436        worker
437            .submit_terminal(
438                &lease(),
439                history(2),
440                RetentionClass::Standard,
441                provider.clone(),
442            )
443            .unwrap();
444        let started = Instant::now();
445        for turn in 3..103 {
446            worker
447                .submit_terminal(
448                    &lease(),
449                    history(turn),
450                    RetentionClass::Standard,
451                    provider.clone(),
452                )
453                .unwrap();
454        }
455        assert!(started.elapsed() < Duration::from_millis(50));
456        assert_eq!(worker.overflow_drops(), 100);
457        release_tx.send(()).unwrap();
458    }
459
460    #[test]
461    fn memory_possibly_committed_capture_is_queried_before_retry() {
462        struct AckLostAfterCommit {
463            calls: AtomicU64,
464            queries: AtomicU64,
465            committed: AtomicBool,
466        }
467        impl CaptureProvider for AckLostAfterCommit {
468            fn capture(&self, _capture: ChatTurnCapture) -> Result<(), CaptureFailure> {
469                self.calls.fetch_add(1, Ordering::SeqCst);
470                self.committed.store(true, Ordering::SeqCst);
471                Err(CaptureFailure { code: "ack_lost" })
472            }
473
474            fn contains_capture(
475                &self,
476                _capture_id: &[u8; 32],
477            ) -> Result<CaptureCommitState, CaptureFailure> {
478                self.queries.fetch_add(1, Ordering::SeqCst);
479                Ok(if self.committed.load(Ordering::SeqCst) {
480                    CaptureCommitState::Committed
481                } else {
482                    CaptureCommitState::Absent
483                })
484            }
485        }
486
487        let provider = Arc::new(AckLostAfterCommit {
488            calls: AtomicU64::new(0),
489            queries: AtomicU64::new(0),
490            committed: AtomicBool::new(false),
491        });
492        let worker = CaptureWorker::new(2);
493        for _ in 0..2 {
494            worker
495                .submit_terminal(
496                    &lease(),
497                    history(99),
498                    RetentionClass::Standard,
499                    provider.clone(),
500                )
501                .unwrap();
502            wait_until(|| provider.queries.load(Ordering::SeqCst) == 1);
503        }
504        thread::sleep(Duration::from_millis(10));
505        assert_eq!(
506            provider.calls.load(Ordering::SeqCst),
507            1,
508            "must not blindly retry"
509        );
510        assert_eq!(provider.queries.load(Ordering::SeqCst), 1);
511        assert_eq!(worker.provider_failures(), 0, "query proved the commit");
512    }
513
514    #[test]
515    fn memory_failed_capture_id_can_be_explicitly_retried() {
516        let provider = Arc::new(ProviderDouble {
517            calls: AtomicU64::new(0),
518            queries: AtomicU64::new(0),
519            fail: AtomicBool::new(true),
520            committed: AtomicBool::new(false),
521            block: Mutex::new(None),
522        });
523        let worker = CaptureWorker::new(2);
524        worker
525            .submit_terminal(
526                &lease(),
527                history(7),
528                RetentionClass::Standard,
529                provider.clone(),
530            )
531            .unwrap();
532        wait_until(|| worker.provider_failures() == 1);
533        provider.fail.store(false, Ordering::SeqCst);
534        worker
535            .submit_terminal(
536                &lease(),
537                history(7),
538                RetentionClass::Standard,
539                provider.clone(),
540            )
541            .unwrap();
542        wait_until(|| provider.calls.load(Ordering::SeqCst) == 2);
543    }
544
545    #[test]
546    fn memory_provider_capture_failure_leaves_typed_terminal_turn_completed() {
547        let provider = Arc::new(ProviderDouble {
548            calls: AtomicU64::new(0),
549            queries: AtomicU64::new(0),
550            fail: AtomicBool::new(true),
551            committed: AtomicBool::new(false),
552            block: Mutex::new(None),
553        });
554        let worker = CaptureWorker::new(1);
555        let terminal = history(1);
556        assert_eq!(terminal.outcome, TurnOutcome::Completed);
557        assert!(worker
558            .submit_terminal(
559                &lease(),
560                terminal,
561                RetentionClass::Standard,
562                provider.clone()
563            )
564            .unwrap());
565        wait_until(|| worker.provider_failures() == 1);
566    }
567}
agent_engine/runtime/capture_worker.rs

agent_engine/runtime/
capture_worker.rs
