Skip to main content

vv_agent/
checkpoint.rs

1//! Public building blocks for the opt-in durable checkpoint v2 protocol.
2//!
3//! The v1 checkpoint implementation deliberately remains in
4//! runtime::state and runtime::checkpoint_codec. This module owns the
5//! language-neutral values which are new in checkpoint v2 and the canonical
6//! JSON helpers used by definitions, journals, and event identities.
7
8use std::cmp::Ordering;
9use std::collections::{BTreeMap, BTreeSet};
10use std::fmt;
11use std::sync::{Arc, Mutex};
12
13use serde::{Deserialize, Serialize};
14use serde_json::{Map, Number, Value};
15use sha2::{Digest, Sha256};
16
17use crate::runtime::backends::CapabilityRef;
18use crate::runtime::state_v2::CheckpointStoreV2;
19
20mod canonical;
21
22pub use canonical::*;
23use canonical::{
24    require_non_empty, require_positive, utf16_cmp, validate_capability_ref,
25    validate_capability_slot, validate_i_json, validate_pointer,
26};
27
28pub const MAX_WIRE_INTEGER: u64 = (1_u64 << 53) - 1;
29pub const MAX_CHECKPOINT_KEY_BYTES: usize = 512;
30pub const MAX_EXTENSION_NAMESPACE_BYTES: usize = 128;
31pub const MAX_EXTENSION_ENTRY_BYTES: usize = 65_536;
32pub const DEFAULT_MAX_EXTENSION_STATE_BYTES: u64 = 262_144;
33pub const CHECKPOINT_V2_SCHEMA: &str = "vv-agent.checkpoint.v2";
34pub const RUN_DEFINITION_SCHEMA: &str = "vv-agent.run-definition.v1";
35pub const OPERATION_REQUEST_SCHEMA: &str = "vv-agent.operation-request.v1";
36pub const EVENT_CURSOR_SCHEMA: &str = "vv-agent.event-cursor.v1";
37pub const CREDENTIAL_REDACTED: &str = "<credential-redacted>";
38
39/// A stable, observable checkpoint error.
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct CheckpointError {
42    code: String,
43    message: String,
44}
45
46impl CheckpointError {
47    pub fn new(code: impl Into<String>, message: impl Into<String>) -> Self {
48        Self {
49            code: code.into(),
50            message: message.into(),
51        }
52    }
53
54    pub fn code(&self) -> &str {
55        &self.code
56    }
57
58    pub fn error_code(&self) -> &str {
59        self.code()
60    }
61
62    pub fn message(&self) -> &str {
63        &self.message
64    }
65}
66
67impl fmt::Display for CheckpointError {
68    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
69        write!(formatter, "{}: {}", self.code, self.message)
70    }
71}
72
73impl std::error::Error for CheckpointError {}
74
75impl From<serde_json::Error> for CheckpointError {
76    fn from(error: serde_json::Error) -> Self {
77        Self::new("checkpoint_json_invalid", error.to_string())
78    }
79}
80
81impl From<std::io::Error> for CheckpointError {
82    fn from(error: std::io::Error) -> Self {
83        Self::new("checkpoint_store_io", error.to_string())
84    }
85}
86
87pub type CheckpointResult<T> = Result<T, CheckpointError>;
88
89#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
90#[serde(rename_all = "snake_case")]
91pub enum ResumePolicy {
92    #[default]
93    New,
94    ResumeIfPresent,
95    RequireExisting,
96}
97
98#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
99#[serde(rename_all = "snake_case")]
100pub enum AmbiguousModelPolicy {
101    #[default]
102    RequireReconciliation,
103    RetryWithDuplicateRisk,
104}
105
106#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
107#[serde(rename_all = "snake_case")]
108pub enum AmbiguousToolPolicy {
109    #[default]
110    RequireReconciliation,
111    RetryIdempotentOnly,
112}
113
114#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
115#[serde(rename_all = "snake_case")]
116pub enum ToolIdempotency {
117    Supported,
118    Unsupported,
119    #[default]
120    Unknown,
121}
122
123#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
124#[serde(rename_all = "snake_case")]
125pub enum OperationKind {
126    Model,
127    Tool,
128}
129
130#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
131#[serde(rename_all = "snake_case")]
132pub enum OperationState {
133    Planned,
134    Started,
135    Succeeded,
136    Failed,
137    Ambiguous,
138}
139
140#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
141#[serde(rename_all = "snake_case")]
142pub enum ClaimMode {
143    Continue,
144    Recovery,
145}
146
147#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
148#[serde(rename_all = "snake_case")]
149pub enum CheckpointStatus {
150    Pending,
151    Running,
152    WaitUser,
153    Completed,
154    Failed,
155    MaxCycles,
156    ReconciliationRequired,
157}
158
159impl CheckpointStatus {
160    pub fn as_str(self) -> &'static str {
161        match self {
162            Self::Pending => "pending",
163            Self::Running => "running",
164            Self::WaitUser => "wait_user",
165            Self::Completed => "completed",
166            Self::Failed => "failed",
167            Self::MaxCycles => "max_cycles",
168            Self::ReconciliationRequired => "reconciliation_required",
169        }
170    }
171
172    pub fn is_terminal(self) -> bool {
173        matches!(
174            self,
175            Self::WaitUser | Self::Completed | Self::Failed | Self::MaxCycles
176        )
177    }
178}
179
180#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
181#[serde(rename_all = "snake_case")]
182pub enum ReconciliationDecisionKind {
183    Defer,
184    Retry,
185    ReplaySuccess,
186    RecordFailure,
187    Abort,
188}
189
190/// A typed observation retained when a started external operation has no
191/// durable receipt.
192#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
193pub struct ResumeObservation {
194    pub operation_id: String,
195    pub operation_kind: OperationKind,
196    pub cycle_index: u64,
197    #[serde(default = "ambiguous_state")]
198    pub state: OperationState,
199    pub risk: String,
200    pub idempotency_support: Option<ToolIdempotency>,
201}
202
203fn ambiguous_state() -> OperationState {
204    OperationState::Ambiguous
205}
206
207impl ResumeObservation {
208    pub fn validate(&self) -> CheckpointResult<()> {
209        require_non_empty(&self.operation_id, "resume observation operation_id")?;
210        require_positive(self.cycle_index, "resume observation cycle_index")?;
211        if self.state != OperationState::Ambiguous {
212            return Err(CheckpointError::new(
213                "checkpoint_resume_observation_invalid",
214                "resume observation state must be ambiguous",
215            ));
216        }
217        require_non_empty(&self.risk, "resume observation risk")
218    }
219}
220
221#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
222pub struct ReconciliationError {
223    pub code: String,
224    pub message: String,
225    #[serde(default)]
226    pub retryable: bool,
227}
228
229impl ReconciliationError {
230    pub fn new(code: impl Into<String>, message: impl Into<String>, retryable: bool) -> Self {
231        Self {
232            code: code.into(),
233            message: message.into(),
234            retryable,
235        }
236    }
237
238    pub fn validate(&self) -> CheckpointResult<()> {
239        require_non_empty(&self.code, "reconciliation error code")?;
240        require_non_empty(&self.message, "reconciliation error message")
241    }
242}
243
244#[derive(Debug, Clone, PartialEq)]
245pub struct ReconciliationDecision {
246    pub kind: ReconciliationDecisionKind,
247    pub response: Option<Value>,
248    pub result: Option<Value>,
249    pub error: Option<ReconciliationError>,
250}
251
252impl ReconciliationDecision {
253    pub fn defer() -> Self {
254        Self {
255            kind: ReconciliationDecisionKind::Defer,
256            response: None,
257            result: None,
258            error: None,
259        }
260    }
261
262    pub fn retry() -> Self {
263        Self {
264            kind: ReconciliationDecisionKind::Retry,
265            response: None,
266            result: None,
267            error: None,
268        }
269    }
270
271    pub fn replay_response(response: Value) -> Self {
272        Self {
273            kind: ReconciliationDecisionKind::ReplaySuccess,
274            response: Some(response),
275            result: None,
276            error: None,
277        }
278    }
279
280    pub fn replay_result(result: Value) -> Self {
281        Self {
282            kind: ReconciliationDecisionKind::ReplaySuccess,
283            response: None,
284            result: Some(result),
285            error: None,
286        }
287    }
288
289    pub fn record_failure(error: ReconciliationError) -> Self {
290        Self {
291            kind: ReconciliationDecisionKind::RecordFailure,
292            response: None,
293            result: None,
294            error: Some(error),
295        }
296    }
297
298    pub fn abort(error: ReconciliationError) -> Self {
299        Self {
300            kind: ReconciliationDecisionKind::Abort,
301            response: None,
302            result: None,
303            error: Some(error),
304        }
305    }
306
307    pub fn validate(&self) -> CheckpointResult<()> {
308        match self.kind {
309            ReconciliationDecisionKind::ReplaySuccess
310                if self.response.is_some() ^ self.result.is_some() => {}
311            ReconciliationDecisionKind::ReplaySuccess => {
312                return Err(CheckpointError::new(
313                    "checkpoint_reconciliation_decision_invalid",
314                    "replay_success requires exactly one response or result",
315                ));
316            }
317            ReconciliationDecisionKind::RecordFailure | ReconciliationDecisionKind::Abort => {
318                let Some(error) = &self.error else {
319                    return Err(CheckpointError::new(
320                        "checkpoint_reconciliation_decision_invalid",
321                        "failure and abort decisions require a typed error",
322                    ));
323                };
324                error.validate()?;
325            }
326            _ if self.response.is_none() && self.result.is_none() && self.error.is_none() => {}
327            _ => {
328                return Err(CheckpointError::new(
329                    "checkpoint_reconciliation_decision_invalid",
330                    "this reconciliation decision carries an unexpected payload",
331                ));
332            }
333        }
334        if let Some(response) = &self.response {
335            validate_i_json(response, "reconciliation response")?;
336        }
337        if let Some(result) = &self.result {
338            validate_i_json(result, "reconciliation result")?;
339        }
340        Ok(())
341    }
342}
343
344pub trait ReconciliationProvider: Send + Sync {
345    fn reconcile(
346        &self,
347        observation: &ResumeObservation,
348    ) -> CheckpointResult<ReconciliationDecision>;
349}
350
351pub trait CheckpointExtension: Send + Sync {
352    fn namespace(&self) -> &str;
353    fn version(&self) -> &str;
354    fn required(&self) -> bool;
355    fn snapshot(&self) -> CheckpointResult<Value>;
356    fn restore(&self, state: &Value) -> CheckpointResult<()>;
357}
358
359/// Runtime configuration for enabling checkpoint v2.
360#[derive(Clone)]
361pub struct CheckpointConfig {
362    pub store: Option<Arc<dyn CheckpointStoreV2>>,
363    pub store_ref: Option<CapabilityRef>,
364    pub key: Option<String>,
365    pub resume_policy: ResumePolicy,
366    pub ambiguous_model_policy: AmbiguousModelPolicy,
367    pub ambiguous_tool_policy: AmbiguousToolPolicy,
368    pub required_extension_namespaces: Vec<String>,
369    pub max_extension_state_bytes: u64,
370    pub credential_slots: Vec<String>,
371    pub capability_refs: BTreeMap<String, CapabilityRef>,
372}
373
374impl Default for CheckpointConfig {
375    fn default() -> Self {
376        Self {
377            store: None,
378            store_ref: None,
379            key: None,
380            resume_policy: ResumePolicy::New,
381            ambiguous_model_policy: AmbiguousModelPolicy::RequireReconciliation,
382            ambiguous_tool_policy: AmbiguousToolPolicy::RequireReconciliation,
383            required_extension_namespaces: Vec::new(),
384            max_extension_state_bytes: DEFAULT_MAX_EXTENSION_STATE_BYTES,
385            credential_slots: Vec::new(),
386            capability_refs: BTreeMap::new(),
387        }
388    }
389}
390
391impl fmt::Debug for CheckpointConfig {
392    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
393        formatter
394            .debug_struct("CheckpointConfig")
395            .field("store", &self.store.as_ref().map(|_| "configured"))
396            .field("store_ref", &self.store_ref)
397            .field("key", &self.key)
398            .field("resume_policy", &self.resume_policy)
399            .field("ambiguous_model_policy", &self.ambiguous_model_policy)
400            .field("ambiguous_tool_policy", &self.ambiguous_tool_policy)
401            .field(
402                "required_extension_namespaces",
403                &self.required_extension_namespaces,
404            )
405            .field("max_extension_state_bytes", &self.max_extension_state_bytes)
406            .field("credential_slots", &self.credential_slots)
407            .field("capability_refs", &self.capability_refs)
408            .finish()
409    }
410}
411
412impl CheckpointConfig {
413    pub fn new(store: Arc<dyn CheckpointStoreV2>) -> Self {
414        Self {
415            store: Some(store),
416            ..Self::default()
417        }
418    }
419
420    pub fn with_store<S>(store: S) -> Self
421    where
422        S: CheckpointStoreV2 + 'static,
423    {
424        Self::new(Arc::new(store))
425    }
426
427    pub fn validate(&self) -> CheckpointResult<()> {
428        if self.store.is_some() == self.store_ref.is_some() {
429            return Err(CheckpointError::new(
430                "checkpoint_store_selection_invalid",
431                "exactly one of store or store_ref is required",
432            ));
433        }
434        if let Some(reference) = &self.store_ref {
435            validate_capability_ref(reference, "CheckpointConfig.store_ref")?;
436        }
437        match (&self.key, self.resume_policy) {
438            (None, ResumePolicy::New) => {}
439            (None, _) => {
440                return Err(CheckpointError::new(
441                    "checkpoint_key_required",
442                    "resume_if_present and require_existing need an explicit key",
443                ));
444            }
445            (Some(key), _) => validate_checkpoint_key(key)?,
446        }
447        if self.max_extension_state_bytes > MAX_WIRE_INTEGER {
448            return Err(CheckpointError::new(
449                "checkpoint_extension_limit_invalid",
450                "max_extension_state_bytes exceeds the JSON-safe integer range",
451            ));
452        }
453        let mut previous_slot: Option<&str> = None;
454        for slot in &self.credential_slots {
455            validate_pointer(slot).map_err(|error| {
456                CheckpointError::new("checkpoint_credential_slots_invalid", error.message)
457            })?;
458            if previous_slot.is_some_and(|previous| utf16_cmp(previous, slot) != Ordering::Less) {
459                return Err(CheckpointError::new(
460                    "checkpoint_credential_slots_invalid",
461                    "credential slots must be sorted and unique",
462                ));
463            }
464            previous_slot = Some(slot);
465        }
466        let mut seen = BTreeSet::new();
467        for namespace in &self.required_extension_namespaces {
468            validate_extension_namespace(namespace)?;
469            if !seen.insert(namespace) {
470                return Err(CheckpointError::new(
471                    "checkpoint_extension_namespace_duplicate",
472                    format!("duplicate extension namespace {namespace}"),
473                ));
474            }
475        }
476        if self
477            .required_extension_namespaces
478            .windows(2)
479            .any(|window| window[0] > window[1])
480        {
481            return Err(CheckpointError::new(
482                "checkpoint_extension_namespace_invalid",
483                "required extension namespaces must be sorted",
484            ));
485        }
486        for (slot, reference) in &self.capability_refs {
487            validate_capability_slot(slot)?;
488            validate_capability_ref(reference, &format!("capability_refs.{slot}"))?;
489        }
490        Ok(())
491    }
492
493    pub fn capability_ref(&self, slot: &str) -> Option<&CapabilityRef> {
494        self.capability_refs.get(slot)
495    }
496}
497
498#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
499pub struct EventCursor {
500    pub schema_version: String,
501    pub store_ref: CapabilityRef,
502    pub value: Value,
503    pub last_event_id: Option<String>,
504}
505
506impl EventCursor {
507    pub fn new(store_ref: CapabilityRef, value: Value, last_event_id: Option<String>) -> Self {
508        Self {
509            schema_version: EVENT_CURSOR_SCHEMA.to_string(),
510            store_ref,
511            value,
512            last_event_id,
513        }
514    }
515
516    pub fn validate(&self) -> CheckpointResult<()> {
517        if self.schema_version != EVENT_CURSOR_SCHEMA {
518            return Err(CheckpointError::new(
519                "checkpoint_event_cursor_schema_unsupported",
520                "unsupported event cursor schema_version",
521            ));
522        }
523        validate_capability_ref(&self.store_ref, "event cursor store_ref")?;
524        validate_i_json(&self.value, "event cursor value")?;
525        if self
526            .last_event_id
527            .as_ref()
528            .is_some_and(|event_id| event_id.trim().is_empty())
529        {
530            return Err(CheckpointError::new(
531                "checkpoint_event_cursor_invalid",
532                "last_event_id must be non-empty when present",
533            ));
534        }
535        Ok(())
536    }
537}
538
539#[derive(Debug, Clone, PartialEq)]
540pub struct AppendOnceResult {
541    pub inserted: bool,
542    pub cursor: Value,
543}
544
545pub trait IdempotentRunEventStore: Send + Sync {
546    fn append_once(
547        &self,
548        event_id: &str,
549        payload_digest: &str,
550        event: &Value,
551    ) -> CheckpointResult<AppendOnceResult>;
552}
553
554/// A process-local accepted-once event store useful for deterministic tests
555/// and hosts that already own their process lifetime.
556#[derive(Debug, Default, Clone)]
557pub struct InMemoryRunEventStore {
558    entries: Arc<Mutex<RunEventEntries>>,
559    next_sequence: Arc<Mutex<u64>>,
560}
561
562type RunEventEntries = BTreeMap<String, (String, Value, Value)>;
563
564impl IdempotentRunEventStore for InMemoryRunEventStore {
565    fn append_once(
566        &self,
567        event_id: &str,
568        payload_digest: &str,
569        event: &Value,
570    ) -> CheckpointResult<AppendOnceResult> {
571        require_non_empty(event_id, "event_id")?;
572        validate_sha256(payload_digest, "payload_digest")?;
573        if !event.is_object() {
574            return Err(CheckpointError::new(
575                "event_payload_invalid",
576                "event payload must be an object",
577            ));
578        }
579        let calculated = event_payload_digest(event)?;
580        if calculated != payload_digest {
581            return Err(CheckpointError::new(
582                "event_payload_digest_mismatch",
583                "payload_digest does not match the canonical event",
584            ));
585        }
586        let mut entries = self.entries.lock().map_err(|_| {
587            CheckpointError::new(
588                "checkpoint_store_lock_poisoned",
589                "event store lock poisoned",
590            )
591        })?;
592        if let Some((existing_digest, _existing_event, cursor)) = entries.get(event_id) {
593            if existing_digest != payload_digest {
594                return Err(CheckpointError::new(
595                    "event_identity_conflict",
596                    format!("event id {event_id} was recorded with a different digest"),
597                ));
598            }
599            return Ok(AppendOnceResult {
600                inserted: false,
601                cursor: cursor.clone(),
602            });
603        }
604        let mut next = self.next_sequence.lock().map_err(|_| {
605            CheckpointError::new(
606                "checkpoint_store_lock_poisoned",
607                "event sequence lock poisoned",
608            )
609        })?;
610        *next = next.checked_add(1).ok_or_else(|| {
611            CheckpointError::new("event_cursor_overflow", "event sequence overflow")
612        })?;
613        let cursor = serde_json::json!({"sequence": *next});
614        entries.insert(
615            event_id.to_string(),
616            (payload_digest.to_string(), event.clone(), cursor.clone()),
617        );
618        Ok(AppendOnceResult {
619            inserted: true,
620            cursor,
621        })
622    }
623}
624
625impl crate::event_store::RunEventStore for InMemoryRunEventStore {
626    fn append(
627        &self,
628        event: &crate::events::RunEvent,
629    ) -> Result<(), crate::event_store::EventStoreError> {
630        let value = serde_json::to_value(event).map_err(|error| {
631            crate::event_store::EventStoreError::new(
632                "event_store_serialization_error",
633                error.to_string(),
634            )
635        })?;
636        let digest = event_payload_digest(&value).map_err(|error| {
637            crate::event_store::EventStoreError::new(
638                "event_store_checkpoint_error",
639                error.to_string(),
640            )
641        })?;
642        IdempotentRunEventStore::append_once(self, event.event_id().as_str(), &digest, &value)
643            .map_err(|error| {
644                crate::event_store::EventStoreError::new(
645                    "event_store_checkpoint_error",
646                    error.to_string(),
647                )
648            })?;
649        Ok(())
650    }
651
652    fn replay(
653        &self,
654        query: crate::event_store::RunEventReplayQuery,
655    ) -> Result<crate::event_store::RunEventIter, crate::event_store::EventStoreError> {
656        let entries = self.entries.lock().map_err(|_| {
657            crate::event_store::EventStoreError::new(
658                "event_store_lock_poisoned",
659                "event store lock poisoned",
660            )
661        })?;
662        let mut events = entries
663            .values()
664            .filter_map(|(_, event, cursor)| {
665                let sequence = cursor.get("sequence").and_then(Value::as_u64)?;
666                let event =
667                    serde_json::from_value::<crate::events::RunEvent>(event.clone()).ok()?;
668                let include = event.run_id() == query.run_id()
669                    || (query.should_include_children()
670                        && event.parent_run_id() == Some(query.run_id()));
671                include.then_some((sequence, event))
672            })
673            .collect::<Vec<_>>();
674        events.sort_by_key(|(sequence, _)| *sequence);
675        Ok(Box::new(events.into_iter().map(|(_, event)| Ok(event))))
676    }
677
678    fn append_once(
679        &self,
680        event_id: &str,
681        payload_digest: &str,
682        event: &crate::events::RunEvent,
683    ) -> Result<Option<EventCursor>, crate::event_store::EventStoreError> {
684        let value = serde_json::to_value(event).map_err(|error| {
685            crate::event_store::EventStoreError::new(
686                "event_store_serialization_error",
687                error.to_string(),
688            )
689        })?;
690        let result = IdempotentRunEventStore::append_once(self, event_id, payload_digest, &value)
691            .map_err(|error| {
692            crate::event_store::EventStoreError::new(
693                "event_store_checkpoint_error",
694                error.to_string(),
695            )
696        })?;
697        Ok(Some(EventCursor::new(
698            crate::runtime::backends::CapabilityRef {
699                id: "events.in-memory".to_string(),
700                version: "1".to_string(),
701            },
702            result.cursor,
703            Some(event_id.to_string()),
704        )))
705    }
706}