mobius/backend/checkpoint/
mod.rs

1//! Durable agent checkpoints.
2
3use std::collections::BTreeMap;
4
5use serde::Deserialize;
6use serde::Serialize;
7use serde_json::Value;
8
9use crate::BoxFuture;
10use crate::Error;
11use crate::Result;
12use crate::backend::model::ToolCall;
13use crate::backend::sandbox::NetworkAccess;
14use crate::backend::sandbox::SandboxMode;
15use crate::protocol::Event;
16use crate::protocol::EventMsg;
17use crate::protocol::MAX_CAPABILITY_INPUT_BYTES;
18use crate::protocol::MessageTarget;
19use crate::protocol::ModelStepContentPhase;
20use crate::protocol::SessionContext;
21use crate::protocol::TokenUsage;
22
23pub mod sqlite;
24
25pub(crate) const CHECKPOINT_VERSION: u32 = 6;
26pub(crate) const MAX_QUEUED_INPUTS: usize = 1_024;
27const TURN_PAGE_BATCH_SIZE: usize = 100;
28const MAX_QUEUED_OWNER_BYTES: usize = 256;
29const MAX_QUEUED_ID_BYTES: usize = 4 * 1024;
30
31/// Mutable counters for the user turn currently running.
32#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
33pub struct ActiveExecution {
34    pub submission_id: String,
35    pub turn_id: String,
36    pub started_at_ms: i64,
37    pub model_calls: u64,
38    pub tool_calls: u64,
39    pub failed_tool_calls: u64,
40    pub usage: TokenUsage,
41}
42
43/// The model step currently in flight for an active execution.
44#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
45pub struct ActiveModelStep {
46    pub model_step_id: String,
47    pub step_index: usize,
48    pub started_at_ms: i64,
49}
50
51/// Terminal outcome of one user turn.
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
53#[serde(rename_all = "snake_case")]
54pub enum ExecutionOutcome {
55    Completed,
56    Aborted,
57    Failed,
58}
59
60/// Durable observability record for one completed user turn.
61#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
62pub struct ExecutionRecord {
63    pub session_id: String,
64    pub submission_id: String,
65    pub turn_id: String,
66    pub started_at_ms: i64,
67    pub finished_at_ms: i64,
68    pub elapsed_ms: u64,
69    pub outcome: ExecutionOutcome,
70    pub model_calls: u64,
71    pub tool_calls: u64,
72    pub failed_tool_calls: u64,
73    pub usage: TokenUsage,
74}
75
76/// Aggregate execution metrics for one durable session.
77#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
78pub struct ExecutionStats {
79    pub run_count: u64,
80    pub failed_run_count: u64,
81    pub aborted_run_count: u64,
82    pub model_calls: u64,
83    pub tool_calls: u64,
84    pub failed_tool_calls: u64,
85    pub elapsed_ms: u64,
86    pub usage: TokenUsage,
87}
88
89/// One intentional replacement of active model history.
90#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
91#[serde(rename_all = "snake_case")]
92pub enum ContextRewriteReason {
93    ContextOffloading,
94    Compaction,
95    Scratchpad,
96}
97
98impl ContextRewriteReason {
99    pub(crate) const fn as_str(self) -> &'static str {
100        match self {
101            Self::ContextOffloading => "context_offloading",
102            Self::Compaction => "compaction",
103            Self::Scratchpad => "scratchpad",
104        }
105    }
106}
107
108/// The latest deliberate active-context rewrite.
109#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
110pub struct ContextRewrite {
111    pub epoch: u64,
112    pub reasons: Vec<ContextRewriteReason>,
113}
114
115impl ExecutionStats {
116    pub(crate) fn checked_record(&mut self, record: &ExecutionRecord) -> Option<()> {
117        let run_count = self.run_count.checked_add(1)?;
118        let failed_run_count = self
119            .failed_run_count
120            .checked_add(u64::from(record.outcome == ExecutionOutcome::Failed))?;
121        let aborted_run_count = self
122            .aborted_run_count
123            .checked_add(u64::from(record.outcome == ExecutionOutcome::Aborted))?;
124        let model_calls = self.model_calls.checked_add(record.model_calls)?;
125        let tool_calls = self.tool_calls.checked_add(record.tool_calls)?;
126        let failed_tool_calls = self
127            .failed_tool_calls
128            .checked_add(record.failed_tool_calls)?;
129        let elapsed_ms = self.elapsed_ms.checked_add(record.elapsed_ms)?;
130        let mut usage = self.usage.clone();
131        usage.checked_add(&record.usage)?;
132        *self = Self {
133            run_count,
134            failed_run_count,
135            aborted_run_count,
136            model_calls,
137            tool_calls,
138            failed_tool_calls,
139            elapsed_ms,
140            usage,
141        };
142        Some(())
143    }
144}
145
146/// A tool batch waiting for a frontend decision.
147#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
148pub struct PendingApproval {
149    pub submission_id: String,
150    pub turn_id: String,
151    pub request_id: String,
152    pub approval_call_ids: Vec<String>,
153    pub authorized_call_ids: Vec<String>,
154    pub calls: Vec<ToolCall>,
155    pub reason: String,
156    pub sandbox_mode: SandboxMode,
157    pub network_access: NetworkAccess,
158    pub decision_received: bool,
159}
160
161/// One active-turn message waiting for its next model boundary.
162#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
163pub struct QueuedInput {
164    owner: String,
165    id: String,
166    text: String,
167}
168
169impl QueuedInput {
170    pub(crate) fn new(owner: &str, id: &str, text: &str) -> Result<Self> {
171        validate_queued_input(owner, id, text).map_err(|message| Error::Config(message.into()))?;
172        Ok(Self {
173            owner: owner.into(),
174            id: id.into(),
175            text: text.into(),
176        })
177    }
178
179    pub(crate) fn validate(&self) -> std::result::Result<(), &'static str> {
180        validate_queued_input(&self.owner, &self.id, &self.text)
181    }
182
183    pub(crate) fn owner(&self) -> &str {
184        &self.owner
185    }
186
187    /// Returns the submission that owns this queued input.
188    #[must_use]
189    pub fn id(&self) -> &str {
190        &self.id
191    }
192
193    /// Returns the exact queued text.
194    #[must_use]
195    pub fn text(&self) -> &str {
196        &self.text
197    }
198
199    pub(crate) fn into_text(self) -> String {
200        self.text
201    }
202
203    pub(crate) fn into_id_and_text(self) -> (String, String) {
204        (self.id, self.text)
205    }
206}
207
208fn validate_queued_input(
209    owner: &str,
210    id: &str,
211    text: &str,
212) -> std::result::Result<(), &'static str> {
213    if owner.trim().is_empty() || owner.len() > MAX_QUEUED_OWNER_BYTES {
214        return Err("queued input owner is invalid");
215    }
216    if id.trim().is_empty() || id.len() > MAX_QUEUED_ID_BYTES {
217        return Err("queued input ID is invalid");
218    }
219    if text.trim().is_empty() || text.len() > MAX_CAPABILITY_INPUT_BYTES {
220        return Err("queued input text is invalid");
221    }
222    Ok(())
223}
224
225/// Versioned state persisted at each durable loop boundary.
226#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
227pub struct Checkpoint {
228    pub version: u32,
229    pub session_id: String,
230    pub session_context: SessionContext,
231    pub metadata: BTreeMap<String, Value>,
232    pub catalog_visible: bool,
233    pub first_user_message: Option<String>,
234    pub model_route: Option<String>,
235    pub sequence: u64,
236    pub context: Vec<Value>,
237    pub context_epoch: u64,
238    pub compaction_count: u64,
239    pub last_context_rewrite: Option<ContextRewrite>,
240    pub total_usage: TokenUsage,
241    pub last_usage: Option<TokenUsage>,
242    pub pending_input: Vec<QueuedInput>,
243    pub active_execution: Option<ActiveExecution>,
244    pub active_model_step: Option<ActiveModelStep>,
245    pub execution_stats: ExecutionStats,
246    pub pending_tools: Vec<ToolCall>,
247    pub pending_approval: Option<PendingApproval>,
248}
249
250impl Checkpoint {
251    /// Creates an empty session checkpoint.
252    #[must_use]
253    pub fn empty(session_id: impl Into<String>) -> Self {
254        Self {
255            version: CHECKPOINT_VERSION,
256            session_id: session_id.into(),
257            session_context: SessionContext::default(),
258            metadata: BTreeMap::new(),
259            catalog_visible: true,
260            first_user_message: None,
261            model_route: None,
262            sequence: 0,
263            context: Vec::new(),
264            context_epoch: 0,
265            compaction_count: 0,
266            last_context_rewrite: None,
267            total_usage: TokenUsage::default(),
268            last_usage: None,
269            pending_input: Vec::new(),
270            active_execution: None,
271            active_model_step: None,
272            execution_stats: ExecutionStats::default(),
273            pending_tools: Vec::new(),
274            pending_approval: None,
275        }
276    }
277
278    pub(crate) fn finish_execution(
279        &mut self,
280        outcome: ExecutionOutcome,
281        finished_at_ms: i64,
282    ) -> Result<ExecutionRecord> {
283        if self.active_model_step.is_some() {
284            return Err(Error::Checkpoint(
285                "turn ended with an active model step".into(),
286            ));
287        }
288        let active = self
289            .active_execution
290            .as_ref()
291            .ok_or_else(|| Error::Checkpoint("turn ended without an active execution".into()))?;
292        let finished_at_ms = finished_at_ms.max(active.started_at_ms);
293        let elapsed_ms = u64::try_from(finished_at_ms - active.started_at_ms)
294            .map_err(|_| Error::Checkpoint("execution elapsed time is unsupported".into()))?;
295        let record = ExecutionRecord {
296            session_id: self.session_id.clone(),
297            submission_id: active.submission_id.clone(),
298            turn_id: active.turn_id.clone(),
299            started_at_ms: active.started_at_ms,
300            finished_at_ms,
301            elapsed_ms,
302            outcome,
303            model_calls: active.model_calls,
304            tool_calls: active.tool_calls,
305            failed_tool_calls: active.failed_tool_calls,
306            usage: active.usage.clone(),
307        };
308        let mut stats = self.execution_stats.clone();
309        stats.checked_record(&record).ok_or_else(|| {
310            Error::Checkpoint("execution statistics exceed the supported range".into())
311        })?;
312        self.active_execution = None;
313        self.execution_stats = stats;
314        Ok(record)
315    }
316}
317
318/// Catalog metadata for one durable session.
319#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
320pub struct SessionSummary {
321    pub session_id: String,
322    pub session_context: SessionContext,
323    pub parent_session_id: Option<String>,
324    pub parent_sequence: Option<u64>,
325    pub sequence: u64,
326    pub catalog_visible: bool,
327    pub first_user_message: Option<String>,
328    pub execution_stats: ExecutionStats,
329    pub created_at: i64,
330    pub updated_at: i64,
331}
332
333/// Stable key for continuing a newest-first session catalog query.
334#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
335pub struct SessionCursor {
336    pub updated_at: i64,
337    pub sequence: u64,
338    pub session_id: String,
339}
340
341/// Bounds one newest-first session catalog query.
342#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
343pub struct SessionPageRequest {
344    pub cursor: Option<SessionCursor>,
345    pub limit: usize,
346}
347
348/// One page of durable sessions.
349#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
350pub struct SessionPage {
351    pub sessions: Vec<SessionSummary>,
352    pub next_cursor: Option<SessionCursor>,
353}
354
355/// One append-only transcript delta at its durable checkpoint sequence.
356#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
357pub struct TranscriptBatch {
358    pub sequence: u64,
359    pub created_at: i64,
360    pub items: Vec<Value>,
361}
362
363/// Bounds one newest-first execution-journal query.
364#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
365pub struct ExecutionPageRequest {
366    pub before_sequence: Option<u64>,
367    pub limit: usize,
368}
369
370/// One newest-first page of terminal user-turn records.
371#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
372pub struct ExecutionPage {
373    pub executions: Vec<ExecutionRecord>,
374    pub next_before_sequence: Option<u64>,
375}
376
377/// Bounds one newest-first transcript query.
378#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
379pub struct TranscriptPageRequest {
380    pub before_sequence: Option<u64>,
381    pub max_batches: usize,
382}
383
384/// One newest-first page of transcript deltas.
385#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
386pub struct TranscriptPage {
387    pub batches: Vec<TranscriptBatch>,
388    pub next_before_sequence: Option<u64>,
389}
390
391/// One normalized frontend event in the durable session journal.
392#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
393pub struct JournalEvent {
394    /// Monotonic sequence within one session.
395    pub sequence: u64,
396    /// Framework record time in Unix milliseconds.
397    pub recorded_at_ms: i64,
398    /// Provider-neutral framework event.
399    pub event: Event,
400    /// Compact delivery characteristics retained after progressive deltas are removed.
401    pub stream_metrics: Vec<StreamMetrics>,
402}
403
404/// One normalized event paired with its framework receipt time.
405#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
406pub struct TimestampedEvent {
407    pub recorded_at_ms: i64,
408    pub event: Event,
409}
410
411/// Delivery metrics for one typed text stream within a completed model step.
412#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
413pub struct StreamMetrics {
414    pub phase: ModelStepContentPhase,
415    pub first_delta_at_ms: i64,
416    pub last_delta_at_ms: i64,
417    pub chunk_count: u64,
418    pub utf8_bytes: u64,
419    pub longest_gap_ms: u64,
420}
421
422/// Bounds one newest-first event-journal query.
423#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
424pub struct EventPageRequest {
425    pub before_sequence: Option<u64>,
426    pub limit: usize,
427}
428
429/// One newest-first page of normalized session events.
430#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
431pub struct EventPage {
432    /// Durable sequence high-water, including intentionally discarded transient events.
433    pub latest_sequence: u64,
434    pub events: Vec<JournalEvent>,
435    pub next_before_sequence: Option<u64>,
436}
437
438impl EventPage {
439    /// Returns this newest-first page in replay order.
440    #[must_use]
441    pub fn into_chronological(mut self) -> Vec<JournalEvent> {
442        self.events.reverse();
443        self.events
444    }
445}
446
447/// Loads the newest logical turn before a durable event cursor.
448pub async fn event_turn_page(
449    checkpoints: &dyn CheckpointStore,
450    session_id: &str,
451    before_sequence: Option<u64>,
452) -> Result<EventPage> {
453    let mut cursor = before_sequence;
454    let mut latest_sequence = 0;
455    let mut events = Vec::new();
456    let mut found_start = false;
457    let mut has_earlier_turn = false;
458
459    loop {
460        let page = checkpoints
461            .event_page(
462                session_id,
463                EventPageRequest {
464                    before_sequence: cursor,
465                    limit: TURN_PAGE_BATCH_SIZE,
466                },
467            )
468            .await?;
469        if events.is_empty() {
470            latest_sequence = page.latest_sequence;
471        }
472        for event in page.events {
473            if found_start {
474                if matches!(&event.event.msg, EventMsg::TurnStarted(_)) {
475                    has_earlier_turn = true;
476                    break;
477                }
478            } else {
479                found_start = matches!(&event.event.msg, EventMsg::TurnStarted(_));
480                events.push(event);
481            }
482        }
483        if has_earlier_turn {
484            break;
485        }
486        let Some(next) = page.next_before_sequence else {
487            break;
488        };
489        cursor = Some(next);
490    }
491
492    let Some((start_index, turn_id)) = events.iter().enumerate().find_map(|(index, event)| {
493        let EventMsg::TurnStarted(started) = &event.event.msg else {
494            return None;
495        };
496        Some((index, started.turn_id.as_str()))
497    }) else {
498        return Ok(EventPage {
499            latest_sequence,
500            events: Vec::new(),
501            next_before_sequence: None,
502        });
503    };
504    let page_start = events[..start_index]
505        .iter()
506        .position(|event| match &event.event.msg {
507            EventMsg::TurnComplete(completed) => completed.turn_id == turn_id,
508            EventMsg::TurnAborted(aborted) => aborted.turn_id == turn_id,
509            _ => false,
510        })
511        .unwrap_or(0);
512    let next_before_sequence = has_earlier_turn.then_some(events[start_index].sequence);
513    let events = events.drain(page_start..=start_index).collect();
514
515    Ok(EventPage {
516        latest_sequence,
517        events,
518        next_before_sequence,
519    })
520}
521
522impl TranscriptPage {
523    /// Flattens this newest-first page into chronological items with durable positions.
524    #[must_use]
525    pub fn into_positioned_items_chronological(self) -> Vec<(MessageTarget, Value)> {
526        self.batches
527            .into_iter()
528            .rev()
529            .flat_map(|batch| {
530                batch
531                    .items
532                    .into_iter()
533                    .enumerate()
534                    .map(move |(index, item)| {
535                        (
536                            MessageTarget {
537                                checkpoint_sequence: batch.sequence,
538                                batch_item_count: index + 1,
539                            },
540                            item,
541                        )
542                    })
543            })
544            .collect()
545    }
546}
547
548/// Stores durable session checkpoints and middleware state.
549pub trait CheckpointStore: Send + Sync {
550    /// Loads the latest checkpoint for a session.
551    fn load<'a>(&'a self, session_id: &'a str) -> BoxFuture<'a, Result<Option<Checkpoint>>>;
552
553    /// Permanently deletes a session, its descendants, and their session-scoped state.
554    ///
555    /// Returns whether the requested session existed.
556    fn delete_session<'a>(&'a self, session_id: &'a str) -> BoxFuture<'a, Result<bool>>;
557
558    /// Atomically replaces the checkpoint, appends transcript items, and records a finished turn.
559    fn save<'a>(
560        &'a self,
561        checkpoint: &'a Checkpoint,
562        transcript_delta: &'a [Value],
563        execution: Option<&'a ExecutionRecord>,
564    ) -> BoxFuture<'a, Result<()>>;
565
566    /// Atomically saves one checkpoint and appends its normalized event batch.
567    fn save_with_events<'a>(
568        &'a self,
569        checkpoint: &'a Checkpoint,
570        transcript_delta: &'a [Value],
571        execution: Option<&'a ExecutionRecord>,
572        events: &'a [TimestampedEvent],
573    ) -> BoxFuture<'a, Result<Vec<JournalEvent>>>;
574
575    /// Assigns a session-local sequence and appends one normalized event atomically.
576    fn append_event<'a>(
577        &'a self,
578        session_id: &'a str,
579        recorded_at_ms: i64,
580        event: &'a Event,
581    ) -> BoxFuture<'a, Result<JournalEvent>>;
582
583    /// Loads one newest-first page of normalized session events.
584    fn event_page<'a>(
585        &'a self,
586        session_id: &'a str,
587        request: EventPageRequest,
588    ) -> BoxFuture<'a, Result<EventPage>>;
589
590    /// Lists one page of the most recently updated sessions, newest first.
591    fn list_sessions_page(
592        &self,
593        _request: SessionPageRequest,
594    ) -> BoxFuture<'_, Result<SessionPage>> {
595        Box::pin(async {
596            Err(Error::Checkpoint(
597                "this checkpoint backend has no session catalog".into(),
598            ))
599        })
600    }
601
602    /// Loads one newest-first page of append-only transcript deltas.
603    fn transcript_page<'a>(
604        &'a self,
605        session_id: &'a str,
606        request: TranscriptPageRequest,
607    ) -> BoxFuture<'a, Result<TranscriptPage>> {
608        Box::pin(async move {
609            if request.max_batches == 0 {
610                return Err(Error::Checkpoint(
611                    "transcript page limit must be positive".into(),
612                ));
613            }
614            let Some(checkpoint) = self.load(session_id).await? else {
615                return Ok(TranscriptPage::default());
616            };
617            if checkpoint.context.is_empty()
618                || request
619                    .before_sequence
620                    .is_some_and(|before| checkpoint.sequence >= before)
621            {
622                return Ok(TranscriptPage::default());
623            }
624            Ok(TranscriptPage {
625                batches: vec![TranscriptBatch {
626                    sequence: checkpoint.sequence,
627                    created_at: 0,
628                    items: checkpoint.context,
629                }],
630                next_before_sequence: None,
631            })
632        })
633    }
634
635    /// Loads one newest-first page of terminal user-turn records.
636    fn execution_page<'a>(
637        &'a self,
638        _session_id: &'a str,
639        _request: ExecutionPageRequest,
640    ) -> BoxFuture<'a, Result<ExecutionPage>> {
641        Box::pin(async {
642            Err(Error::Checkpoint(
643                "this checkpoint backend has no execution journal".into(),
644            ))
645        })
646    }
647
648    /// Loads the most recently started terminal user turns across all sessions.
649    fn recent_executions(&self, _limit: usize) -> BoxFuture<'_, Result<Vec<ExecutionRecord>>> {
650        Box::pin(async {
651            Err(Error::Checkpoint(
652                "this checkpoint backend has no execution journal".into(),
653            ))
654        })
655    }
656
657    /// Creates a child session at an exact durable parent sequence.
658    fn fork<'a>(
659        &'a self,
660        _parent_session_id: &'a str,
661        _parent_sequence: u64,
662        _checkpoint: &'a Checkpoint,
663    ) -> BoxFuture<'a, Result<SessionSummary>> {
664        Box::pin(async {
665            Err(Error::Checkpoint(
666                "this checkpoint backend cannot fork sessions".into(),
667            ))
668        })
669    }
670
671    /// Loads the latest opaque state owned by one middleware namespace.
672    fn load_state<'a>(
673        &'a self,
674        scope: &'a str,
675        key: &'a str,
676    ) -> BoxFuture<'a, Result<Option<Value>>>;
677
678    /// Durably replaces opaque middleware state.
679    fn save_state<'a>(
680        &'a self,
681        scope: &'a str,
682        key: &'a str,
683        value: &'a Value,
684    ) -> BoxFuture<'a, Result<()>>;
685}
mobius/backend/checkpoint/mod.rs

mobius/backend/checkpoint/
mod.rs
