1use 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_MESSAGE_BYTES;
18use crate::protocol::MessageAuthor;
19use crate::protocol::MessageDelivery;
20use crate::protocol::MessageEvent;
21use crate::protocol::MessageTarget;
22use crate::protocol::ModelStepContentPhase;
23use crate::protocol::SessionContext;
24use crate::protocol::SessionFileReference;
25use crate::protocol::TokenUsage;
26
27pub mod sqlite;
28
29pub(crate) const CHECKPOINT_VERSION: u32 = 13;
30pub(crate) const MAX_QUEUED_MESSAGES: usize = 1_024;
31const TURN_PAGE_BATCH_SIZE: usize = 100;
32const MAX_QUEUED_OWNER_BYTES: usize = 256;
33const MAX_QUEUED_ID_BYTES: usize = 4 * 1024;
34const MAX_QUEUED_TURN_ID_BYTES: usize = 4 * 1024;
35const MAX_QUEUED_MESSAGE_BYTES: usize = MAX_MESSAGE_BYTES * 2;
36
37#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(tag = "type", rename_all = "snake_case")]
40pub enum ExecutionPhase {
41 Model,
42 Completion {
43 last_assistant_message: Option<String>,
44 },
45}
46
47#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49pub struct ActiveExecution {
50 pub submission_id: String,
51 pub turn_id: String,
52 pub started_at_ms: i64,
53 pub model_calls: u64,
54 pub tool_calls: u64,
55 pub failed_tool_calls: u64,
56 pub usage: TokenUsage,
57 pub next_model_step: usize,
58 pub stop_hook_active: bool,
59 pub phase: ExecutionPhase,
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
64pub struct ActiveModelStep {
65 pub model_step_id: String,
66 pub step_index: usize,
67 pub started_at_ms: i64,
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
72#[serde(rename_all = "snake_case")]
73pub enum ExecutionOutcome {
74 Completed,
75 Aborted,
76 Failed,
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
81pub struct ExecutionRecord {
82 pub session_id: String,
83 pub submission_id: String,
84 pub turn_id: String,
85 pub started_at_ms: i64,
86 pub finished_at_ms: i64,
87 pub elapsed_ms: u64,
88 pub outcome: ExecutionOutcome,
89 pub model_calls: u64,
90 pub tool_calls: u64,
91 pub failed_tool_calls: u64,
92 pub usage: TokenUsage,
93}
94
95#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
97pub struct ExecutionStats {
98 pub run_count: u64,
99 pub failed_run_count: u64,
100 pub aborted_run_count: u64,
101 pub model_calls: u64,
102 pub tool_calls: u64,
103 pub failed_tool_calls: u64,
104 pub elapsed_ms: u64,
105 pub usage: TokenUsage,
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
110#[serde(rename_all = "snake_case")]
111pub enum ContextRewriteReason {
112 ContextOffloading,
113 Compaction,
114 Scratchpad,
115}
116
117impl ContextRewriteReason {
118 pub(crate) const fn as_str(self) -> &'static str {
119 match self {
120 Self::ContextOffloading => "context_offloading",
121 Self::Compaction => "compaction",
122 Self::Scratchpad => "scratchpad",
123 }
124 }
125}
126
127#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
129pub struct ContextRewrite {
130 pub epoch: u64,
131 pub reasons: Vec<ContextRewriteReason>,
132}
133
134impl ExecutionStats {
135 pub(crate) fn checked_record(&mut self, record: &ExecutionRecord) -> Option<()> {
136 let run_count = self.run_count.checked_add(1)?;
137 let failed_run_count = self
138 .failed_run_count
139 .checked_add(u64::from(record.outcome == ExecutionOutcome::Failed))?;
140 let aborted_run_count = self
141 .aborted_run_count
142 .checked_add(u64::from(record.outcome == ExecutionOutcome::Aborted))?;
143 let model_calls = self.model_calls.checked_add(record.model_calls)?;
144 let tool_calls = self.tool_calls.checked_add(record.tool_calls)?;
145 let failed_tool_calls = self
146 .failed_tool_calls
147 .checked_add(record.failed_tool_calls)?;
148 let elapsed_ms = self.elapsed_ms.checked_add(record.elapsed_ms)?;
149 let mut usage = self.usage.clone();
150 usage.checked_add(&record.usage)?;
151 *self = Self {
152 run_count,
153 failed_run_count,
154 aborted_run_count,
155 model_calls,
156 tool_calls,
157 failed_tool_calls,
158 elapsed_ms,
159 usage,
160 };
161 Some(())
162 }
163}
164
165#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
167pub struct PendingApproval {
168 pub submission_id: String,
169 pub turn_id: String,
170 pub request_id: String,
171 pub approval_call_ids: Vec<String>,
172 pub authorized_call_ids: Vec<String>,
173 pub calls: Vec<ToolCall>,
174 pub reason: String,
175 pub sandbox_mode: SandboxMode,
176 pub network_access: NetworkAccess,
177 pub decision_received: bool,
178}
179
180#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
182#[serde(tag = "type", rename_all = "snake_case")]
183pub enum QueuedMessageBoundary {
184 Turn,
185 Steer { turn_id: String },
186 Queue,
187}
188
189impl QueuedMessageBoundary {
190 pub(crate) const fn delivery(&self) -> MessageDelivery {
191 match self {
192 Self::Turn => MessageDelivery::Turn,
193 Self::Steer { .. } => MessageDelivery::Steer,
194 Self::Queue => MessageDelivery::Queue,
195 }
196 }
197
198 pub(crate) const fn starts_turn(&self) -> bool {
199 matches!(self, Self::Turn | Self::Queue)
200 }
201}
202
203#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
205pub struct QueuedMessage {
206 owner: String,
207 id: String,
208 boundary: QueuedMessageBoundary,
209 author: MessageAuthor,
210 message: String,
211 attachments: Vec<SessionFileReference>,
212}
213
214impl QueuedMessage {
215 pub(crate) fn new(
216 owner: &str,
217 id: &str,
218 boundary: QueuedMessageBoundary,
219 event: MessageEvent,
220 ) -> Result<Self> {
221 if event.delivery != boundary.delivery() || event.message_target.is_some() {
222 return Err(Error::Config("queued message event is inconsistent".into()));
223 }
224 let queued = Self {
225 owner: owner.into(),
226 id: id.into(),
227 boundary,
228 author: event.author,
229 message: event.text,
230 attachments: event.attachments,
231 };
232 queued.validate()?;
233 Ok(queued)
234 }
235
236 pub(crate) fn validate(&self) -> Result<()> {
237 validate_queued_message(self)
238 }
239
240 pub(crate) fn owner(&self) -> &str {
241 &self.owner
242 }
243
244 #[must_use]
246 pub fn id(&self) -> &str {
247 &self.id
248 }
249
250 pub(crate) fn boundary(&self) -> &QueuedMessageBoundary {
251 &self.boundary
252 }
253
254 pub(crate) fn event(&self) -> MessageEvent {
255 MessageEvent {
256 author: self.author.clone(),
257 delivery: self.boundary.delivery(),
258 text: self.message.clone(),
259 attachments: self.attachments.clone(),
260 message_target: None,
261 }
262 }
263
264 pub(crate) fn replace(&mut self, id: &str, event: MessageEvent) -> Result<()> {
265 let replacement = Self::new(&self.owner, id, self.boundary.clone(), event)?;
266 *self = replacement;
267 Ok(())
268 }
269
270 pub(crate) fn promote_to_next_turn(&mut self) -> Result<()> {
271 self.boundary = QueuedMessageBoundary::Queue;
272 Ok(())
273 }
274
275 pub(crate) fn into_parts(self) -> (String, MessageEvent) {
276 let event = self.event();
277 (self.id, event)
278 }
279}
280
281fn validate_queued_message(message: &QueuedMessage) -> Result<()> {
282 if message.owner.trim().is_empty() || message.owner.len() > MAX_QUEUED_OWNER_BYTES {
283 return Err(Error::Config("queued message owner is invalid".into()));
284 }
285 if message.id.trim().is_empty() || message.id.len() > MAX_QUEUED_ID_BYTES {
286 return Err(Error::Config("queued message ID is invalid".into()));
287 }
288 if matches!(
289 &message.boundary,
290 QueuedMessageBoundary::Steer { turn_id }
291 if turn_id.trim().is_empty() || turn_id.len() > MAX_QUEUED_TURN_ID_BYTES
292 ) {
293 return Err(Error::Config("queued message turn ID is invalid".into()));
294 }
295 crate::protocol::validate_message_content(
296 &message.author,
297 &message.message,
298 &message.attachments,
299 )?;
300 if serde_json::to_vec(&message.event())
301 .map_or(true, |value| value.len() > MAX_QUEUED_MESSAGE_BYTES)
302 {
303 return Err(Error::Config("queued message is invalid".into()));
304 }
305 Ok(())
306}
307
308#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
310pub struct Checkpoint {
311 pub version: u32,
312 pub session_id: String,
313 pub session_context: SessionContext,
314 pub metadata: BTreeMap<String, Value>,
315 pub catalog_visible: bool,
316 pub first_user_message: Option<String>,
317 pub model_route: Option<String>,
318 pub sequence: u64,
319 pub context: Vec<Value>,
320 pub context_epoch: u64,
321 pub compaction_count: u64,
322 pub last_context_rewrite: Option<ContextRewrite>,
323 pub total_usage: TokenUsage,
324 pub last_usage: Option<TokenUsage>,
325 pub pending_messages: Vec<QueuedMessage>,
326 pub active_execution: Option<ActiveExecution>,
327 pub active_model_step: Option<ActiveModelStep>,
328 pub execution_stats: ExecutionStats,
329 pub pending_tools: Vec<ToolCall>,
330 pub pending_approval: Option<PendingApproval>,
331}
332
333impl Checkpoint {
334 #[must_use]
336 pub fn empty(session_id: impl Into<String>) -> Self {
337 Self {
338 version: CHECKPOINT_VERSION,
339 session_id: session_id.into(),
340 session_context: SessionContext::default(),
341 metadata: BTreeMap::new(),
342 catalog_visible: true,
343 first_user_message: None,
344 model_route: None,
345 sequence: 0,
346 context: Vec::new(),
347 context_epoch: 0,
348 compaction_count: 0,
349 last_context_rewrite: None,
350 total_usage: TokenUsage::default(),
351 last_usage: None,
352 pending_messages: Vec::new(),
353 active_execution: None,
354 active_model_step: None,
355 execution_stats: ExecutionStats::default(),
356 pending_tools: Vec::new(),
357 pending_approval: None,
358 }
359 }
360
361 pub(crate) fn finish_execution(
362 &mut self,
363 outcome: ExecutionOutcome,
364 finished_at_ms: i64,
365 ) -> Result<ExecutionRecord> {
366 if self.active_model_step.is_some() {
367 return Err(Error::Checkpoint(
368 "turn ended with an active model step".into(),
369 ));
370 }
371 let active = self
372 .active_execution
373 .as_ref()
374 .ok_or_else(|| Error::Checkpoint("turn ended without an active execution".into()))?;
375 let finished_at_ms = finished_at_ms.max(active.started_at_ms);
376 let elapsed_ms = u64::try_from(finished_at_ms - active.started_at_ms)
377 .map_err(|_| Error::Checkpoint("execution elapsed time is unsupported".into()))?;
378 let record = ExecutionRecord {
379 session_id: self.session_id.clone(),
380 submission_id: active.submission_id.clone(),
381 turn_id: active.turn_id.clone(),
382 started_at_ms: active.started_at_ms,
383 finished_at_ms,
384 elapsed_ms,
385 outcome,
386 model_calls: active.model_calls,
387 tool_calls: active.tool_calls,
388 failed_tool_calls: active.failed_tool_calls,
389 usage: active.usage.clone(),
390 };
391 let mut stats = self.execution_stats.clone();
392 stats.checked_record(&record).ok_or_else(|| {
393 Error::Checkpoint("execution statistics exceed the supported range".into())
394 })?;
395 self.active_execution = None;
396 self.execution_stats = stats;
397 Ok(record)
398 }
399}
400
401#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
403pub struct SessionSummary {
404 pub session_id: String,
405 pub session_context: SessionContext,
406 pub parent_session_id: Option<String>,
407 pub parent_sequence: Option<u64>,
408 pub sequence: u64,
409 pub catalog_visible: bool,
410 pub first_user_message: Option<String>,
411 pub execution_stats: ExecutionStats,
412 pub created_at: i64,
413 pub updated_at: i64,
414}
415
416#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
418pub struct SessionCursor {
419 pub updated_at: i64,
420 pub sequence: u64,
421 pub session_id: String,
422}
423
424#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
426pub struct SessionPageRequest {
427 pub cursor: Option<SessionCursor>,
428 pub limit: usize,
429}
430
431#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
433pub struct SessionPage {
434 pub sessions: Vec<SessionSummary>,
435 pub next_cursor: Option<SessionCursor>,
436}
437
438#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
440pub struct TranscriptBatch {
441 pub sequence: u64,
442 pub created_at: i64,
443 pub items: Vec<Value>,
444}
445
446#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
448pub struct ExecutionPageRequest {
449 pub before_sequence: Option<u64>,
450 pub limit: usize,
451}
452
453#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
455pub struct ExecutionPage {
456 pub executions: Vec<ExecutionRecord>,
457 pub next_before_sequence: Option<u64>,
458}
459
460#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
462pub struct TranscriptPageRequest {
463 pub before_sequence: Option<u64>,
464 pub max_batches: usize,
465}
466
467#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
469pub struct TranscriptPage {
470 pub batches: Vec<TranscriptBatch>,
471 pub next_before_sequence: Option<u64>,
472}
473
474#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
476pub struct JournalEvent {
477 pub sequence: u64,
479 pub recorded_at_ms: i64,
481 pub event: Event,
483 pub stream_metrics: Vec<StreamMetrics>,
485}
486
487#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
489pub struct TimestampedEvent {
490 pub recorded_at_ms: i64,
491 pub event: Event,
492}
493
494#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
496pub struct StreamMetrics {
497 pub phase: ModelStepContentPhase,
498 pub first_delta_at_ms: i64,
499 pub last_delta_at_ms: i64,
500 pub chunk_count: u64,
501 pub utf8_bytes: u64,
502 pub longest_gap_ms: u64,
503}
504
505#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
507pub struct EventPageRequest {
508 pub before_sequence: Option<u64>,
509 pub limit: usize,
510}
511
512#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
514pub struct EventPage {
515 pub latest_sequence: u64,
517 pub events: Vec<JournalEvent>,
518 pub next_before_sequence: Option<u64>,
519}
520
521impl EventPage {
522 #[must_use]
524 pub fn into_chronological(mut self) -> Vec<JournalEvent> {
525 self.events.reverse();
526 self.events
527 }
528}
529
530pub async fn event_turn_page(
532 checkpoints: &dyn CheckpointStore,
533 session_id: &str,
534 before_sequence: Option<u64>,
535) -> Result<EventPage> {
536 let mut cursor = before_sequence;
537 let mut latest_sequence = 0;
538 let mut events = Vec::new();
539 let mut found_start = false;
540 let mut has_earlier_turn = false;
541
542 loop {
543 let page = checkpoints
544 .event_page(
545 session_id,
546 EventPageRequest {
547 before_sequence: cursor,
548 limit: TURN_PAGE_BATCH_SIZE,
549 },
550 )
551 .await?;
552 if events.is_empty() {
553 latest_sequence = page.latest_sequence;
554 }
555 for event in page.events {
556 if found_start {
557 if matches!(&event.event.msg, EventMsg::TurnStarted(_)) {
558 has_earlier_turn = true;
559 break;
560 }
561 } else {
562 found_start = matches!(&event.event.msg, EventMsg::TurnStarted(_));
563 events.push(event);
564 }
565 }
566 if has_earlier_turn {
567 break;
568 }
569 let Some(next) = page.next_before_sequence else {
570 break;
571 };
572 cursor = Some(next);
573 }
574
575 let Some((start_index, turn_id)) = events.iter().enumerate().find_map(|(index, event)| {
576 let EventMsg::TurnStarted(started) = &event.event.msg else {
577 return None;
578 };
579 Some((index, started.turn_id.as_str()))
580 }) else {
581 return Ok(EventPage {
582 latest_sequence,
583 events: Vec::new(),
584 next_before_sequence: None,
585 });
586 };
587 let page_start = events[..start_index]
588 .iter()
589 .position(|event| match &event.event.msg {
590 EventMsg::TurnComplete(completed) => completed.turn_id == turn_id,
591 EventMsg::TurnAborted(aborted) => aborted.turn_id == turn_id,
592 _ => false,
593 })
594 .unwrap_or(0);
595 let next_before_sequence = has_earlier_turn.then_some(events[start_index].sequence);
596 let events = events.drain(page_start..=start_index).collect();
597
598 Ok(EventPage {
599 latest_sequence,
600 events,
601 next_before_sequence,
602 })
603}
604
605impl TranscriptPage {
606 #[must_use]
608 pub fn into_positioned_items_chronological(self) -> Vec<(MessageTarget, Value)> {
609 self.batches
610 .into_iter()
611 .rev()
612 .flat_map(|batch| {
613 batch
614 .items
615 .into_iter()
616 .enumerate()
617 .map(move |(index, item)| {
618 (
619 MessageTarget {
620 checkpoint_sequence: batch.sequence,
621 batch_item_count: index + 1,
622 },
623 item,
624 )
625 })
626 })
627 .collect()
628 }
629}
630
631pub trait CheckpointStore: Send + Sync {
633 fn load<'a>(&'a self, session_id: &'a str) -> BoxFuture<'a, Result<Option<Checkpoint>>>;
635
636 fn delete_session<'a>(&'a self, session_id: &'a str) -> BoxFuture<'a, Result<bool>>;
640
641 fn save<'a>(
643 &'a self,
644 checkpoint: &'a Checkpoint,
645 transcript_delta: &'a [Value],
646 execution: Option<&'a ExecutionRecord>,
647 ) -> BoxFuture<'a, Result<()>>;
648
649 fn save_with_events<'a>(
651 &'a self,
652 checkpoint: &'a Checkpoint,
653 transcript_delta: &'a [Value],
654 execution: Option<&'a ExecutionRecord>,
655 events: &'a [TimestampedEvent],
656 ) -> BoxFuture<'a, Result<Vec<JournalEvent>>>;
657
658 fn append_event<'a>(
660 &'a self,
661 session_id: &'a str,
662 recorded_at_ms: i64,
663 event: &'a Event,
664 ) -> BoxFuture<'a, Result<JournalEvent>>;
665
666 fn event_page<'a>(
668 &'a self,
669 session_id: &'a str,
670 request: EventPageRequest,
671 ) -> BoxFuture<'a, Result<EventPage>>;
672
673 fn list_sessions_page(
675 &self,
676 _request: SessionPageRequest,
677 ) -> BoxFuture<'_, Result<SessionPage>> {
678 Box::pin(async {
679 Err(Error::Checkpoint(
680 "this checkpoint backend has no session catalog".into(),
681 ))
682 })
683 }
684
685 fn transcript_page<'a>(
687 &'a self,
688 session_id: &'a str,
689 request: TranscriptPageRequest,
690 ) -> BoxFuture<'a, Result<TranscriptPage>> {
691 Box::pin(async move {
692 if request.max_batches == 0 {
693 return Err(Error::Checkpoint(
694 "transcript page limit must be positive".into(),
695 ));
696 }
697 let Some(checkpoint) = self.load(session_id).await? else {
698 return Ok(TranscriptPage::default());
699 };
700 if checkpoint.context.is_empty()
701 || request
702 .before_sequence
703 .is_some_and(|before| checkpoint.sequence >= before)
704 {
705 return Ok(TranscriptPage::default());
706 }
707 Ok(TranscriptPage {
708 batches: vec![TranscriptBatch {
709 sequence: checkpoint.sequence,
710 created_at: 0,
711 items: checkpoint.context,
712 }],
713 next_before_sequence: None,
714 })
715 })
716 }
717
718 fn execution_page<'a>(
720 &'a self,
721 _session_id: &'a str,
722 _request: ExecutionPageRequest,
723 ) -> BoxFuture<'a, Result<ExecutionPage>> {
724 Box::pin(async {
725 Err(Error::Checkpoint(
726 "this checkpoint backend has no execution journal".into(),
727 ))
728 })
729 }
730
731 fn recent_executions(&self, _limit: usize) -> BoxFuture<'_, Result<Vec<ExecutionRecord>>> {
733 Box::pin(async {
734 Err(Error::Checkpoint(
735 "this checkpoint backend has no execution journal".into(),
736 ))
737 })
738 }
739
740 fn fork<'a>(
742 &'a self,
743 _parent_session_id: &'a str,
744 _parent_sequence: u64,
745 _checkpoint: &'a Checkpoint,
746 ) -> BoxFuture<'a, Result<SessionSummary>> {
747 Box::pin(async {
748 Err(Error::Checkpoint(
749 "this checkpoint backend cannot fork sessions".into(),
750 ))
751 })
752 }
753
754 fn load_state<'a>(
756 &'a self,
757 scope: &'a str,
758 key: &'a str,
759 ) -> BoxFuture<'a, Result<Option<Value>>>;
760
761 fn save_state<'a>(
763 &'a self,
764 scope: &'a str,
765 key: &'a str,
766 value: &'a Value,
767 ) -> BoxFuture<'a, Result<()>>;
768}