runifold_workflow/
checkpoint.rs

1use std::{collections::BTreeMap, fmt, num::NonZeroU16, sync::Arc};
2
3use runifold_core::{
4    Checkpoint, CheckpointError, CheckpointErrorKind, CheckpointId, CheckpointStore, RunContext,
5    RunId, Usage,
6};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9
10use crate::{StepId, WorkflowError, WorkflowOutcome, WorkflowRemediationCheckpoint, WorkflowWait};
11use crate::{WorkflowLease, WorkflowStore};
12
13const CHECKPOINT_KIND: &str = "runifold.workflow";
14const CHECKPOINT_SCHEMA_VERSION: u32 = 5;
15const MIN_CHECKPOINT_SCHEMA_VERSION: u32 = 3;
16
17/// Recovery behavior for a workflow interrupted inside one node.
18#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
19#[non_exhaustive]
20pub enum WorkflowResumePolicy {
21    /// Reject recovery that could duplicate model cost or external effects.
22    #[default]
23    RejectAmbiguous,
24    /// Explicitly retry only the interrupted workflow node or review substage.
25    RetryInterruptedStep,
26}
27
28/// Maximum number of immutable checkpoint revisions returned by one query.
29#[derive(Clone, Copy, Debug, Eq, PartialEq)]
30pub struct WorkflowCheckpointHistoryLimit(NonZeroU16);
31
32impl WorkflowCheckpointHistoryLimit {
33    /// Creates a bounded history page limit.
34    ///
35    /// # Errors
36    ///
37    /// Rejects zero and values above 256.
38    pub fn new(value: u16) -> Result<Self, CheckpointError> {
39        NonZeroU16::new(value)
40            .filter(|value| value.get() <= 256)
41            .map(Self)
42            .ok_or_else(|| {
43                CheckpointError::new(
44                    CheckpointErrorKind::InvalidPayload,
45                    "workflow checkpoint history limit must be in 1..=256",
46                )
47            })
48    }
49
50    /// Returns the validated page size.
51    pub const fn get(self) -> u16 {
52        self.0.get()
53    }
54}
55
56/// Immutable, typed view of one historical workflow checkpoint revision.
57#[derive(Clone, Debug, PartialEq)]
58pub struct WorkflowCheckpointRevision {
59    /// Workflow checkpoint whose history owns this revision.
60    pub checkpoint_id: CheckpointId,
61    /// Monotonic revision within the checkpoint.
62    pub revision: u64,
63    /// Run that produced the immutable revision.
64    pub run_id: RunId,
65    /// Store timestamp carried by the checkpoint envelope.
66    pub updated_at_ms: u64,
67    /// Validated workflow state at this revision.
68    pub state: WorkflowCheckpointState,
69}
70
71impl WorkflowCheckpointRevision {
72    #[doc(hidden)]
73    pub fn from_checkpoint(checkpoint: Checkpoint) -> Result<Self, CheckpointError> {
74        decode_revision(checkpoint)
75    }
76}
77
78/// Explicit safety policy for forking a historical checkpoint.
79#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
80#[serde(rename_all = "snake_case")]
81#[non_exhaustive]
82pub enum WorkflowForkPolicy {
83    /// Fork only from a stable boundary that cannot repeat an ambiguous node.
84    #[default]
85    RejectAmbiguous,
86    /// Re-run one serial node or review substage persisted as in-flight.
87    RetryInterruptedStep,
88}
89
90/// Idempotent command that creates a new execution from immutable history.
91#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
92pub struct WorkflowForkCommand {
93    /// Caller-stable identity of the new execution branch.
94    pub fork_checkpoint_id: CheckpointId,
95    /// Existing workflow whose history is being selected.
96    pub source_checkpoint_id: CheckpointId,
97    /// Exact immutable source revision.
98    pub source_revision: u64,
99    /// Explicit ambiguous-replay policy.
100    pub policy: WorkflowForkPolicy,
101}
102
103impl WorkflowForkCommand {
104    /// Creates a fork command with a generated target identity.
105    pub fn new(
106        source_checkpoint_id: CheckpointId,
107        source_revision: u64,
108        policy: WorkflowForkPolicy,
109    ) -> Self {
110        Self::with_id(
111            CheckpointId::new(),
112            source_checkpoint_id,
113            source_revision,
114            policy,
115        )
116    }
117
118    /// Creates a retryable fork command with a caller-owned target identity.
119    pub const fn with_id(
120        fork_checkpoint_id: CheckpointId,
121        source_checkpoint_id: CheckpointId,
122        source_revision: u64,
123        policy: WorkflowForkPolicy,
124    ) -> Self {
125        Self {
126            fork_checkpoint_id,
127            source_checkpoint_id,
128            source_revision,
129            policy,
130        }
131    }
132
133    #[doc(hidden)]
134    pub fn prepare_checkpoint(&self, source: Checkpoint) -> Result<Checkpoint, CheckpointError> {
135        fork_checkpoint(source, self.fork_checkpoint_id, self.policy)
136    }
137}
138
139/// Immutable parent relationship of one forked workflow execution.
140#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
141pub struct WorkflowLineage {
142    /// Source workflow checkpoint.
143    pub parent_checkpoint_id: CheckpointId,
144    /// Exact parent revision selected for the fork.
145    pub parent_revision: u64,
146    /// Safety policy used to create the child.
147    pub policy: WorkflowForkPolicy,
148}
149
150/// Result of an idempotent workflow fork command.
151#[derive(Clone, Copy, Debug, Eq, PartialEq)]
152#[non_exhaustive]
153pub enum WorkflowForkOutcome {
154    /// A new workflow branch was atomically created.
155    Created {
156        /// New branch identity.
157        checkpoint_id: CheckpointId,
158    },
159    /// The same target identity was already bound to the same source.
160    Duplicate {
161        /// Existing branch identity.
162        checkpoint_id: CheckpointId,
163    },
164}
165
166/// Persisted workflow execution phase.
167#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
168#[serde(tag = "state", rename_all = "snake_case")]
169#[non_exhaustive]
170pub enum WorkflowCheckpointPhase {
171    /// Stable output is ready for the next node.
172    Ready,
173    /// One node may have partially executed.
174    StepInFlight {
175        /// Stable interrupted node identity.
176        step: StepId,
177    },
178    /// One reviewable serial node is generating, reviewing, or committing.
179    Remediating {
180        /// Stable repairable node identity.
181        step: StepId,
182        /// One-based generation attempt.
183        attempt: u32,
184        /// Original stable value supplied to the node.
185        original_input: Value,
186        /// Durable remediation substate.
187        checkpoint: WorkflowRemediationCheckpoint,
188    },
189    /// The worker lease was released while this node awaits a durable wake.
190    Waiting {
191        /// Stable waiting node identity.
192        step: StepId,
193        /// Durable wake condition.
194        wait: WorkflowWait,
195    },
196    /// A parallel node has one or more incomplete branches.
197    ParallelInFlight {
198        /// Stable parallel node identity.
199        step: StepId,
200        /// Durable branch progress keyed independently of completion order.
201        branches: BTreeMap<StepId, ParallelBranchCheckpoint>,
202    },
203    /// A first-success race has no durable winner yet, or awaits commit.
204    RaceInFlight {
205        /// Stable race node identity.
206        step: StepId,
207        /// Durable branch progress keyed independently of completion order.
208        branches: BTreeMap<StepId, ParallelBranchCheckpoint>,
209    },
210    /// The workflow reached a terminal output.
211    Completed {
212        /// Complete canonical workflow result.
213        outcome: WorkflowOutcome,
214    },
215}
216
217/// Persisted state of one branch inside an in-flight parallel node.
218#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
219#[serde(tag = "state", rename_all = "snake_case")]
220#[non_exhaustive]
221pub enum ParallelBranchCheckpoint {
222    /// The branch may have partially executed.
223    InFlight,
224    /// The branch returned a stable canonical output.
225    Completed {
226        /// Canonical branch output.
227        output: Value,
228    },
229    /// The branch returned a known failure.
230    Failed {
231        /// Safe persisted failure explanation.
232        message: String,
233    },
234    /// The branch was abandoned after another branch won.
235    Cancelled,
236}
237
238/// Versioned workflow state stored in a domain-neutral checkpoint envelope.
239#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
240pub struct WorkflowCheckpointState {
241    /// Stable workflow definition name.
242    pub workflow: String,
243    /// Caller-managed definition version.
244    pub workflow_version: u32,
245    /// Ordered node layout used to reject incompatible definitions.
246    pub layout: Vec<StepId>,
247    /// Index of the next node to execute.
248    pub next_index: usize,
249    /// Canonical value presented to the next node.
250    pub value: Value,
251    /// Stable outputs of all completed nodes.
252    pub outputs: BTreeMap<StepId, Value>,
253    /// Shared usage snapshot at persistence time.
254    pub usage: Usage,
255    /// Current recovery phase.
256    pub phase: WorkflowCheckpointPhase,
257}
258
259impl WorkflowCheckpointState {
260    pub(crate) fn outcome(&self) -> Option<WorkflowOutcome> {
261        match &self.phase {
262            WorkflowCheckpointPhase::Completed { outcome } => Some(outcome.clone()),
263            _ => None,
264        }
265    }
266}
267
268/// Stable handle binding one workflow checkpoint identity to a store.
269#[derive(Clone)]
270pub struct WorkflowCheckpoint {
271    id: CheckpointId,
272    backend: WorkflowCheckpointBackend,
273}
274
275#[derive(Clone)]
276enum WorkflowCheckpointBackend {
277    Local(Arc<dyn CheckpointStore>),
278    Distributed {
279        store: Arc<dyn WorkflowStore>,
280        lease: WorkflowLease,
281    },
282}
283
284impl WorkflowCheckpoint {
285    /// Creates a new workflow checkpoint handle.
286    pub fn new(store: Arc<dyn CheckpointStore>) -> Self {
287        Self {
288            id: CheckpointId::new(),
289            backend: WorkflowCheckpointBackend::Local(store),
290        }
291    }
292
293    /// Reconnects to an existing workflow checkpoint.
294    pub fn existing(id: CheckpointId, store: Arc<dyn CheckpointStore>) -> Self {
295        Self {
296            id,
297            backend: WorkflowCheckpointBackend::Local(store),
298        }
299    }
300
301    /// Binds a distributed checkpoint to the current fenced workflow lease.
302    pub fn distributed(store: Arc<dyn WorkflowStore>, lease: WorkflowLease) -> Self {
303        Self {
304            id: lease.checkpoint_id,
305            backend: WorkflowCheckpointBackend::Distributed { store, lease },
306        }
307    }
308
309    /// Returns the stable checkpoint identity.
310    pub const fn id(&self) -> CheckpointId {
311        self.id
312    }
313
314    /// Loads and validates the latest typed workflow state.
315    ///
316    /// # Errors
317    ///
318    /// Returns [`CheckpointError`] for storage or payload failures.
319    pub fn load(&self) -> Result<(Checkpoint, WorkflowCheckpointState), CheckpointError> {
320        let WorkflowCheckpointBackend::Local(store) = &self.backend else {
321            return Err(CheckpointError::new(
322                CheckpointErrorKind::Storage,
323                "distributed workflow checkpoints must be loaded asynchronously",
324            ));
325        };
326        let checkpoint = store.load(self.id)?;
327        decode(checkpoint)
328    }
329
330    /// Loads and validates either a local or distributed checkpoint.
331    ///
332    /// # Errors
333    ///
334    /// Returns [`CheckpointError`] for ownership, storage, or payload failures.
335    pub async fn load_async(
336        &self,
337    ) -> Result<(Checkpoint, WorkflowCheckpointState), CheckpointError> {
338        let checkpoint = match &self.backend {
339            WorkflowCheckpointBackend::Local(store) => store.load(self.id)?,
340            WorkflowCheckpointBackend::Distributed { store, lease } => {
341                store.load_checkpoint(lease.clone()).await?
342            }
343        };
344        decode(checkpoint)
345    }
346
347    async fn compare_and_swap(
348        &self,
349        checkpoint: &Checkpoint,
350        expected_revision: Option<u64>,
351    ) -> Result<(), CheckpointError> {
352        match &self.backend {
353            WorkflowCheckpointBackend::Local(store) => {
354                store.compare_and_swap(checkpoint, expected_revision)
355            }
356            WorkflowCheckpointBackend::Distributed { store, lease } => {
357                store
358                    .compare_and_swap_checkpoint(
359                        lease.clone(),
360                        checkpoint.clone(),
361                        expected_revision,
362                    )
363                    .await
364            }
365        }
366    }
367}
368
369fn decode(
370    checkpoint: Checkpoint,
371) -> Result<(Checkpoint, WorkflowCheckpointState), CheckpointError> {
372    if checkpoint.kind != CHECKPOINT_KIND
373        || !(MIN_CHECKPOINT_SCHEMA_VERSION..=CHECKPOINT_SCHEMA_VERSION)
374            .contains(&checkpoint.schema_version)
375    {
376        return Err(CheckpointError::new(
377            CheckpointErrorKind::InvalidPayload,
378            "checkpoint kind or schema version is not supported",
379        ));
380    }
381    let state = serde_json::from_value(checkpoint.payload.clone()).map_err(|error| {
382        CheckpointError::new(CheckpointErrorKind::InvalidPayload, error.to_string())
383    })?;
384    Ok((checkpoint, state))
385}
386
387pub(crate) fn decode_revision(
388    checkpoint: Checkpoint,
389) -> Result<WorkflowCheckpointRevision, CheckpointError> {
390    let (checkpoint, state) = decode(checkpoint)?;
391    Ok(WorkflowCheckpointRevision {
392        checkpoint_id: checkpoint.id,
393        revision: checkpoint.revision,
394        run_id: checkpoint.run_id,
395        updated_at_ms: checkpoint.updated_at_ms,
396        state,
397    })
398}
399
400pub(crate) fn fork_checkpoint(
401    source: Checkpoint,
402    target: CheckpointId,
403    policy: WorkflowForkPolicy,
404) -> Result<Checkpoint, CheckpointError> {
405    let (_, mut state) = decode(source)?;
406    match state.phase.clone() {
407        WorkflowCheckpointPhase::StepInFlight { .. }
408            if policy == WorkflowForkPolicy::RetryInterruptedStep =>
409        {
410            state.phase = WorkflowCheckpointPhase::Ready;
411        }
412        WorkflowCheckpointPhase::Remediating {
413            step,
414            attempt,
415            original_input,
416            checkpoint: WorkflowRemediationCheckpoint::GenerationInFlight { input },
417        } if policy == WorkflowForkPolicy::RetryInterruptedStep => {
418            state.phase = WorkflowCheckpointPhase::Remediating {
419                step,
420                attempt,
421                original_input,
422                checkpoint: WorkflowRemediationCheckpoint::GenerationReady { input },
423            };
424        }
425        WorkflowCheckpointPhase::Remediating {
426            step,
427            attempt,
428            original_input,
429            checkpoint: WorkflowRemediationCheckpoint::ReviewInFlight { candidate },
430        } if policy == WorkflowForkPolicy::RetryInterruptedStep => {
431            state.phase = WorkflowCheckpointPhase::Remediating {
432                step,
433                attempt,
434                original_input,
435                checkpoint: WorkflowRemediationCheckpoint::ReviewReady { candidate },
436            };
437        }
438        WorkflowCheckpointPhase::StepInFlight { .. }
439        | WorkflowCheckpointPhase::Remediating {
440            checkpoint:
441                WorkflowRemediationCheckpoint::GenerationInFlight { .. }
442                | WorkflowRemediationCheckpoint::ReviewInFlight { .. },
443            ..
444        }
445        | WorkflowCheckpointPhase::ParallelInFlight { .. }
446        | WorkflowCheckpointPhase::RaceInFlight { .. } => {
447            return Err(CheckpointError::new(
448                CheckpointErrorKind::Conflict,
449                "workflow checkpoint is ambiguous and cannot be forked safely",
450            ));
451        }
452        _ => {}
453    }
454    Ok(Checkpoint::initial(
455        target,
456        RunId::new(),
457        CHECKPOINT_KIND,
458        CHECKPOINT_SCHEMA_VERSION,
459        serde_json::to_value(state).map_err(|error| {
460            CheckpointError::new(CheckpointErrorKind::InvalidPayload, error.to_string())
461        })?,
462    ))
463}
464
465impl fmt::Debug for WorkflowCheckpoint {
466    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
467        formatter
468            .debug_struct("WorkflowCheckpoint")
469            .field("id", &self.id)
470            .finish_non_exhaustive()
471    }
472}
473
474pub(crate) struct WorkflowCheckpointCursor {
475    handle: WorkflowCheckpoint,
476    envelope: Checkpoint,
477}
478
479impl WorkflowCheckpointCursor {
480    pub(crate) async fn create(
481        handle: &WorkflowCheckpoint,
482        run: &RunContext,
483        state: &WorkflowCheckpointState,
484    ) -> Result<Self, WorkflowError> {
485        let envelope = Checkpoint::initial(
486            handle.id,
487            run.run_id(),
488            CHECKPOINT_KIND,
489            CHECKPOINT_SCHEMA_VERSION,
490            serialize(state)?,
491        );
492        handle.compare_and_swap(&envelope, None).await?;
493        Ok(Self {
494            handle: handle.clone(),
495            envelope,
496        })
497    }
498
499    pub(crate) fn loaded(handle: &WorkflowCheckpoint, envelope: Checkpoint) -> Self {
500        Self {
501            handle: handle.clone(),
502            envelope,
503        }
504    }
505
506    pub(crate) async fn save(
507        &mut self,
508        state: &WorkflowCheckpointState,
509    ) -> Result<(), WorkflowError> {
510        let next = self.envelope.next(serialize(state)?)?;
511        self.handle
512            .compare_and_swap(&next, Some(self.envelope.revision))
513            .await?;
514        self.envelope = next;
515        Ok(())
516    }
517}
518
519fn serialize(state: &WorkflowCheckpointState) -> Result<Value, WorkflowError> {
520    serde_json::to_value(state).map_err(|error| {
521        CheckpointError::new(CheckpointErrorKind::InvalidPayload, error.to_string()).into()
522    })
523}
524
525#[cfg(test)]
526mod tests {
527    use runifold_core::{CheckpointId, RunId, Usage};
528    use serde_json::json;
529
530    use super::*;
531
532    fn in_flight_checkpoint() -> Checkpoint {
533        let state = WorkflowCheckpointState {
534            workflow: "fork-test".into(),
535            workflow_version: 1,
536            layout: vec![StepId::parse("charge").unwrap()],
537            next_index: 0,
538            value: json!({"amount": 42}),
539            outputs: BTreeMap::new(),
540            usage: Usage {
541                tokens: 7,
542                ..Usage::default()
543            },
544            phase: WorkflowCheckpointPhase::StepInFlight {
545                step: StepId::parse("charge").unwrap(),
546            },
547        };
548        Checkpoint::initial(
549            CheckpointId::new(),
550            RunId::new(),
551            CHECKPOINT_KIND,
552            CHECKPOINT_SCHEMA_VERSION,
553            serde_json::to_value(state).unwrap(),
554        )
555    }
556
557    #[test]
558    fn fork_rejects_ambiguous_replay_unless_explicitly_authorized() {
559        let source = in_flight_checkpoint();
560        let error = fork_checkpoint(
561            source.clone(),
562            CheckpointId::new(),
563            WorkflowForkPolicy::RejectAmbiguous,
564        )
565        .unwrap_err();
566        assert_eq!(error.kind, CheckpointErrorKind::Conflict);
567
568        let forked = fork_checkpoint(
569            source,
570            CheckpointId::new(),
571            WorkflowForkPolicy::RetryInterruptedStep,
572        )
573        .unwrap();
574        let revision = decode_revision(forked).unwrap();
575        assert!(matches!(
576            revision.state.phase,
577            WorkflowCheckpointPhase::Ready
578        ));
579        assert_eq!(revision.state.usage.tokens, 7);
580        assert_eq!(revision.state.next_index, 0);
581    }
582
583    #[test]
584    fn schema_v4_ready_checkpoint_remains_readable() {
585        let checkpoint = Checkpoint::initial(
586            CheckpointId::new(),
587            RunId::new(),
588            CHECKPOINT_KIND,
589            4,
590            json!({
591                "workflow": "v4-compatible",
592                "workflow_version": 7,
593                "layout": ["draft"],
594                "next_index": 0,
595                "value": {"request": "review this"},
596                "outputs": {},
597                "usage": {
598                    "tokens": 11,
599                    "cost_microusd": 12,
600                    "duration_micros": 13,
601                    "turns": 14,
602                    "tool_calls": 15,
603                    "delegations": 16
604                },
605                "phase": {"state": "ready"}
606            }),
607        );
608
609        let revision = decode_revision(checkpoint).unwrap();
610
611        assert_eq!(revision.state.workflow, "v4-compatible");
612        assert_eq!(revision.state.workflow_version, 7);
613        assert_eq!(revision.state.layout, [StepId::parse("draft").unwrap()]);
614        assert_eq!(revision.state.usage.tokens, 11);
615        assert!(matches!(
616            revision.state.phase,
617            WorkflowCheckpointPhase::Ready
618        ));
619    }
620
621    #[test]
622    fn future_checkpoint_schema_is_rejected() {
623        let mut checkpoint = in_flight_checkpoint();
624        checkpoint.schema_version = CHECKPOINT_SCHEMA_VERSION + 1;
625
626        let error = decode_revision(checkpoint).unwrap_err();
627
628        assert_eq!(error.kind, CheckpointErrorKind::InvalidPayload);
629    }
630}
runifold_workflow/checkpoint.rs

runifold_workflow/
checkpoint.rs
