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                    // Axel's backend transaction/tombstone is authoritative;
113                    // provider ID alone cannot key an in-memory receipt across
114                    // two different selected brains in this process.
115                    let result =
116                        if submitted.contains(&id) && job.provider_id.as_str() != "axel-host" {
117                            Ok(())
118                        } else if let Err(error) = job.payload.dispatch(job.provider.as_ref()) {
119                            // The call may have committed before its acknowledgement
120                            // was lost. Reconcile by idempotency key before accepting
121                            // an explicit retry; never blindly dispatch it again.
122                            match job.payload.query_committed(job.provider.as_ref()) {
123                                Ok(CaptureCommitState::Committed) => {
124                                    submitted.insert(id);
125                                    Ok(())
126                                }
127                                Ok(CaptureCommitState::Absent) | Err(_) => {
128                                    worker_failures.fetch_add(1, Ordering::Relaxed);
129                                    Err(error)
130                                }
131                            }
132                        } else {
133                            submitted.insert(id);
134                            Ok(())
135                        };
136                    if let Some(committed) = job.committed {
137                        let _ = committed.send(result);
138                    }
139                }
140            })
141            .expect("capture worker spawn");
142        Self {
143            sender,
144            dropped,
145            failures,
146        }
147    }
148
149    /// Apply lease capability and capture gates, build the bounded payload, and
150    /// attempt a nonblocking enqueue. `Ok(false)` means capture was gated off.
151    pub fn submit_terminal(
152        &self,
153        lease: &MemoryContextLease,
154        history: TerminalTurnHistory,
155        retention: RetentionClass,
156        provider: Arc<dyn CaptureProvider>,
157    ) -> Result<bool, CaptureBuildError> {
158        if lease.mode.turn_capture() != TurnCapture::Enabled
159            || lease.project_id != history.project_id
160            || lease.session_id != history.session_id
161        {
162            return Ok(false);
163        }
164        let capture = build_chat_turn_capture(&lease.project_id, history, retention)?;
165        let job = CaptureJob {
166            provider_id: lease.provider_id.clone(),
167            provider,
168            payload: CapturePayload::Turn(Box::new(capture)),
169            committed: None,
170        };
171        match self.sender.try_send(job) {
172            Ok(()) => Ok(true),
173            Err(mpsc::TrySendError::Full(_)) | Err(mpsc::TrySendError::Disconnected(_)) => {
174                self.dropped.fetch_add(1, Ordering::Relaxed);
175                Ok(true)
176            }
177        }
178    }
179
180    /// Submit a capture already built by the C1 builder. Historical import
181    /// uses this narrow seam so it shares the same bounded C3 queue and exact
182    /// leased provider dispatch as live terminal capture.
183    pub(crate) fn submit_built(
184        &self,
185        lease: &MemoryContextLease,
186        capture: ChatTurnCapture,
187        provider: Arc<dyn CaptureProvider>,
188    ) -> bool {
189        self.submit_built_job(lease, capture, provider, None)
190    }
191
192    /// Submit one historical capture with a completion receipt. Waiting for the
193    /// receipt happens only on the explicit import path, never turn completion.
194    pub(crate) fn submit_built_with_receipt(
195        &self,
196        lease: &MemoryContextLease,
197        capture: ChatTurnCapture,
198        provider: Arc<dyn CaptureProvider>,
199        committed: mpsc::Sender<Result<(), CaptureFailure>>,
200    ) -> bool {
201        self.submit_built_job(lease, capture, provider, Some(committed))
202    }
203
204    fn submit_built_job(
205        &self,
206        lease: &MemoryContextLease,
207        capture: ChatTurnCapture,
208        provider: Arc<dyn CaptureProvider>,
209        committed: Option<mpsc::Sender<Result<(), CaptureFailure>>>,
210    ) -> bool {
211        if lease.mode.turn_capture() != TurnCapture::Enabled
212            || lease.project_id != capture.project_id
213        {
214            return false;
215        }
216        let job = CaptureJob {
217            provider_id: lease.provider_id.clone(),
218            provider,
219            payload: CapturePayload::Turn(Box::new(capture)),
220            committed,
221        };
222        match self.sender.try_send(job) {
223            Ok(()) => true,
224            Err(mpsc::TrySendError::Full(_)) | Err(mpsc::TrySendError::Disconnected(_)) => {
225                self.dropped.fetch_add(1, Ordering::Relaxed);
226                false
227            }
228        }
229    }
230
231    /// Apply the same exact lease gate as turn capture, build a bounded
232    /// first-class compaction memory, and attempt a nonblocking enqueue.
233    #[allow(clippy::too_many_arguments)] // first-class summary metadata stays explicit at this boundary
234    pub fn submit_summary(
235        &self,
236        lease: &MemoryContextLease,
237        source: super::chat_capture::CompactionSource,
238        source_message_count: usize,
239        summary_text: &str,
240        redaction_policy: agent_core::compaction::RedactionPolicy,
241        retention: RetentionClass,
242        provider: Arc<dyn CaptureProvider>,
243    ) -> Result<bool, SummaryCaptureBuildError> {
244        if lease.mode.turn_capture() != TurnCapture::Enabled
245            || lease.project_id != source.project_id
246        {
247            return Ok(false);
248        }
249        let capture = build_conversation_summary_capture(
250            &lease.project_id,
251            source,
252            source_message_count,
253            summary_text,
254            redaction_policy,
255            retention,
256        )?;
257        let job = CaptureJob {
258            provider_id: lease.provider_id.clone(),
259            provider,
260            payload: CapturePayload::Summary(Box::new(capture)),
261            committed: None,
262        };
263        match self.sender.try_send(job) {
264            Ok(()) => Ok(true),
265            Err(mpsc::TrySendError::Full(_)) | Err(mpsc::TrySendError::Disconnected(_)) => {
266                self.dropped.fetch_add(1, Ordering::Relaxed);
267                Ok(true)
268            }
269        }
270    }
271
272    pub fn overflow_drops(&self) -> u64 {
273        self.dropped.load(Ordering::Relaxed)
274    }
275    pub fn provider_failures(&self) -> u64 {
276        self.failures.load(Ordering::Relaxed)
277    }
278}
279
280#[cfg(test)]
281mod tests {
282    use super::*;
283    use crate::runtime::chat_capture::{CanonicalCaptureItem, CaptureContentClass};
284    use crate::runtime::memory_context::{
285        mint_explicit_command_proof, CapturePolicy, MemoryContextMode, MemoryLeaseId, ProjectId,
286        RecallPolicy, SessionId,
287    };
288    use agent_core::core::disclosure::DisclosureClass;
289    use agent_core::TurnOutcome;
290    use std::sync::atomic::AtomicBool;
291    use std::sync::Mutex;
292    use std::time::{Duration, Instant, SystemTime};
293
294    struct ProviderDouble {
295        calls: AtomicU64,
296        queries: AtomicU64,
297        fail: AtomicBool,
298        committed: AtomicBool,
299        block: Mutex<Option<mpsc::Receiver<()>>>,
300    }
301
302    impl CaptureProvider for ProviderDouble {
303        fn capture(&self, _capture: ChatTurnCapture) -> Result<(), CaptureFailure> {
304            self.calls.fetch_add(1, Ordering::SeqCst);
305            if let Some(receiver) = self.block.lock().expect("lock").take() {
306                let _ = receiver.recv();
307            }
308            if self.fail.load(Ordering::SeqCst) {
309                Err(CaptureFailure {
310                    code: "fixture_failure",
311                })
312            } else {
313                self.committed.store(true, Ordering::SeqCst);
314                Ok(())
315            }
316        }
317
318        fn contains_capture(
319            &self,
320            _capture_id: &[u8; 32],
321        ) -> Result<CaptureCommitState, CaptureFailure> {
322            self.queries.fetch_add(1, Ordering::SeqCst);
323            Ok(if self.committed.load(Ordering::SeqCst) {
324                CaptureCommitState::Committed
325            } else {
326                CaptureCommitState::Absent
327            })
328        }
329    }
330
331    fn lease() -> MemoryContextLease {
332        let now = SystemTime::now();
333        MemoryContextLease::grant(
334            MemoryLeaseId::parse("lease-c3").unwrap(),
335            SessionId::parse("session-c3").unwrap(),
336            ProjectId::parse("project-c3").unwrap(),
337            ContextProviderId::parse("provider-c3").unwrap(),
338            MemoryContextMode::CaptureOnly,
339            CapturePolicy::EligibleTurnsOnly,
340            RecallPolicy::BoundedPerPrompt,
341            mint_explicit_command_proof(),
342            now,
343            None,
344        )
345        .unwrap()
346    }
347
348    fn history(turn: u64) -> TerminalTurnHistory {
349        let project_id = ProjectId::parse("project-c3").unwrap();
350        let started_at = SystemTime::UNIX_EPOCH + Duration::from_secs(turn);
351        TerminalTurnHistory {
352            project_id: project_id.clone(),
353            session_id: SessionId::parse("session-c3").unwrap(),
354            turn_id: super::super::memory_context::TurnId::parse(&format!("turn-{turn}")).unwrap(),
355            turn_ordinal: turn,
356            started_at,
357            completed_at: started_at + Duration::from_millis(1),
358            outcome: TurnOutcome::Completed,
359            items: vec![
360                CanonicalCaptureItem {
361                    project_id: project_id.clone(),
362                    class: CaptureContentClass::UserMessage,
363                    disclosure: DisclosureClass::ModelVisible,
364                    sensitivity: super::super::chat_capture::Sensitivity::Normal,
365                    text: "user".into(),
366                    tool_name: None,
367                },
368                CanonicalCaptureItem {
369                    project_id,
370                    class: CaptureContentClass::AssistantFinal,
371                    disclosure: DisclosureClass::ModelVisible,
372                    sensitivity: super::super::chat_capture::Sensitivity::Normal,
373                    text: "assistant".into(),
374                    tool_name: None,
375                },
376            ],
377            compaction: None,
378        }
379    }
380
381    fn wait_until(predicate: impl Fn() -> bool) {
382        let deadline = Instant::now() + Duration::from_secs(1);
383        while !predicate() && Instant::now() < deadline {
384            thread::sleep(Duration::from_millis(2));
385        }
386        assert!(predicate());
387    }
388
389    #[test]
390    fn memory_duplicate_capture_id_idempotent_at_engine_seam() {
391        let provider = Arc::new(ProviderDouble {
392            calls: AtomicU64::new(0),
393            queries: AtomicU64::new(0),
394            fail: AtomicBool::new(false),
395            committed: AtomicBool::new(false),
396            block: Mutex::new(None),
397        });
398        let worker = CaptureWorker::new(4);
399        worker
400            .submit_terminal(
401                &lease(),
402                history(1),
403                RetentionClass::Standard,
404                provider.clone(),
405            )
406            .unwrap();
407        worker
408            .submit_terminal(
409                &lease(),
410                history(1),
411                RetentionClass::Standard,
412                provider.clone(),
413            )
414            .unwrap();
415        wait_until(|| provider.calls.load(Ordering::SeqCst) == 1);
416        thread::sleep(Duration::from_millis(10));
417        assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
418    }
419
420    #[test]
421    fn memory_overflow_drops_exactly_and_never_blocks_completion() {
422        let (release_tx, release_rx) = mpsc::channel();
423        let provider = Arc::new(ProviderDouble {
424            calls: AtomicU64::new(0),
425            queries: AtomicU64::new(0),
426            fail: AtomicBool::new(false),
427            committed: AtomicBool::new(false),
428            block: Mutex::new(Some(release_rx)),
429        });
430        let worker = CaptureWorker::new(1);
431        worker
432            .submit_terminal(
433                &lease(),
434                history(1),
435                RetentionClass::Standard,
436                provider.clone(),
437            )
438            .unwrap();
439        wait_until(|| provider.calls.load(Ordering::SeqCst) == 1);
440        worker
441            .submit_terminal(
442                &lease(),
443                history(2),
444                RetentionClass::Standard,
445                provider.clone(),
446            )
447            .unwrap();
448        let started = Instant::now();
449        for turn in 3..103 {
450            worker
451                .submit_terminal(
452                    &lease(),
453                    history(turn),
454                    RetentionClass::Standard,
455                    provider.clone(),
456                )
457                .unwrap();
458        }
459        assert!(started.elapsed() < Duration::from_millis(50));
460        assert_eq!(worker.overflow_drops(), 100);
461        release_tx.send(()).unwrap();
462    }
463
464    #[test]
465    fn memory_possibly_committed_capture_is_queried_before_retry() {
466        struct AckLostAfterCommit {
467            calls: AtomicU64,
468            queries: AtomicU64,
469            committed: AtomicBool,
470        }
471        impl CaptureProvider for AckLostAfterCommit {
472            fn capture(&self, _capture: ChatTurnCapture) -> Result<(), CaptureFailure> {
473                self.calls.fetch_add(1, Ordering::SeqCst);
474                self.committed.store(true, Ordering::SeqCst);
475                Err(CaptureFailure { code: "ack_lost" })
476            }
477
478            fn contains_capture(
479                &self,
480                _capture_id: &[u8; 32],
481            ) -> Result<CaptureCommitState, CaptureFailure> {
482                self.queries.fetch_add(1, Ordering::SeqCst);
483                Ok(if self.committed.load(Ordering::SeqCst) {
484                    CaptureCommitState::Committed
485                } else {
486                    CaptureCommitState::Absent
487                })
488            }
489        }
490
491        let provider = Arc::new(AckLostAfterCommit {
492            calls: AtomicU64::new(0),
493            queries: AtomicU64::new(0),
494            committed: AtomicBool::new(false),
495        });
496        let worker = CaptureWorker::new(2);
497        for _ in 0..2 {
498            worker
499                .submit_terminal(
500                    &lease(),
501                    history(99),
502                    RetentionClass::Standard,
503                    provider.clone(),
504                )
505                .unwrap();
506            wait_until(|| provider.queries.load(Ordering::SeqCst) == 1);
507        }
508        thread::sleep(Duration::from_millis(10));
509        assert_eq!(
510            provider.calls.load(Ordering::SeqCst),
511            1,
512            "must not blindly retry"
513        );
514        assert_eq!(provider.queries.load(Ordering::SeqCst), 1);
515        assert_eq!(worker.provider_failures(), 0, "query proved the commit");
516    }
517
518    #[test]
519    fn memory_failed_capture_id_can_be_explicitly_retried() {
520        let provider = Arc::new(ProviderDouble {
521            calls: AtomicU64::new(0),
522            queries: AtomicU64::new(0),
523            fail: AtomicBool::new(true),
524            committed: AtomicBool::new(false),
525            block: Mutex::new(None),
526        });
527        let worker = CaptureWorker::new(2);
528        worker
529            .submit_terminal(
530                &lease(),
531                history(7),
532                RetentionClass::Standard,
533                provider.clone(),
534            )
535            .unwrap();
536        wait_until(|| worker.provider_failures() == 1);
537        provider.fail.store(false, Ordering::SeqCst);
538        worker
539            .submit_terminal(
540                &lease(),
541                history(7),
542                RetentionClass::Standard,
543                provider.clone(),
544            )
545            .unwrap();
546        wait_until(|| provider.calls.load(Ordering::SeqCst) == 2);
547    }
548
549    #[test]
550    fn memory_provider_capture_failure_leaves_typed_terminal_turn_completed() {
551        let provider = Arc::new(ProviderDouble {
552            calls: AtomicU64::new(0),
553            queries: AtomicU64::new(0),
554            fail: AtomicBool::new(true),
555            committed: AtomicBool::new(false),
556            block: Mutex::new(None),
557        });
558        let worker = CaptureWorker::new(1);
559        let terminal = history(1);
560        assert_eq!(terminal.outcome, TurnOutcome::Completed);
561        assert!(worker
562            .submit_terminal(
563                &lease(),
564                terminal,
565                RetentionClass::Standard,
566                provider.clone()
567            )
568            .unwrap());
569        wait_until(|| worker.provider_failures() == 1);
570    }
571}
agent_engine/runtime/capture_worker.rs

agent_engine/runtime/
capture_worker.rs
