1use 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#[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#[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#[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#[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}