ingot_runtime/
snapshot.rs

1//! A run that stopped, in a form that can be continued.
2//!
3//! Everything an interrupted run held that the rest of it needs: the inputs,
4//! every binding in scope, working memory, the outputs produced so far, and the
5//! counters. It is a JSON document a person can read, and that is a
6//! requirement rather than a convenience — it is what rules out serialising a
7//! continuation, which in turn is why only a top-level `checkpoint` is
8//! resumable. See
9//! [RFC-0018](../../../rfcs/0018-state-that-outlives-a-run.md) §4.
10//!
11//! # What is deliberately absent
12//!
13//! **Persistent memory.** It belongs to the agent, not to the interrupted run,
14//! and both halves read and write the agent's store as normal.
15//!
16//! **A cassette.** A resumed run is given one the same way the first half was.
17//! Copying the recording into the snapshot would make the two disagree the
18//! moment either was re-recorded.
19
20use std::collections::BTreeMap;
21use std::path::Path;
22
23use serde::{Deserialize, Serialize};
24use serde_json::Value;
25use sha2::{Digest, Sha256};
26
27use ingot_ir::AgentIr;
28
29use crate::events::Artifact;
30use crate::price::Spend;
31use crate::provider::Usage;
32
33/// Format version of this document.
34pub const SNAPSHOT_VERSION: &str = "0.1";
35
36/// What kind of snapshot this is.
37///
38/// A memory store is the other one, and pointing `--resume` at it is a
39/// plausible mistake. Naming both kinds turns "unexpected field" into a
40/// sentence that explains itself.
41pub const KIND: &str = "resumption";
42
43#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
44#[serde(rename_all = "camelCase")]
45pub struct Resumption {
46    pub ingot_snapshot: String,
47    pub kind: String,
48    pub agent: String,
49    /// SHA-256 of the artifact's canonical JSON, `sha256:…`.
50    pub artifact: String,
51    /// The label of the checkpoint the run stopped at.
52    pub label: String,
53    /// The checkpoint node itself, which the `runStopped` event names.
54    pub stopped_at: String,
55    /// The node to continue from — the one **after** the checkpoint, so a
56    /// resumed run does not re-emit the checkpoint's event.
57    pub resume_at: String,
58    pub inputs: BTreeMap<String, Value>,
59    pub bindings: BTreeMap<String, Value>,
60    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
61    pub state: BTreeMap<String, Value>,
62    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
63    pub outputs: BTreeMap<String, Artifact>,
64    pub steps: u32,
65    pub usage: Usage,
66    #[serde(default)]
67    pub spend: Spend,
68    /// How many model calls the first half made.
69    ///
70    /// A cassette is matched by position, so a resumed run replaying one has to
71    /// start where the first half stopped. Counted here rather than asked of
72    /// the provider because it is a property of the run: a live run records the
73    /// same number, and a resumed live run simply ignores it.
74    #[serde(default)]
75    pub model_calls: u32,
76    /// How many tool calls it made, for the same reason.
77    #[serde(default)]
78    pub tool_calls: u32,
79}
80
81/// Why a snapshot could not be used.
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub enum SnapshotError {
84    Io(String),
85    Malformed(String),
86    /// A snapshot of the other kind, or of a version this build does not write.
87    WrongKind {
88        found: String,
89    },
90    UnsupportedVersion {
91        found: String,
92    },
93    /// The artifact is not the one the run stopped in.
94    DifferentArtifact {
95        agent: String,
96    },
97    /// The node to continue from is gone.
98    UnknownNode {
99        node: String,
100    },
101}
102
103impl std::fmt::Display for SnapshotError {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        match self {
106            SnapshotError::Io(reason) => write!(f, "{reason}"),
107            SnapshotError::Malformed(reason) => write!(f, "this is not a snapshot: {reason}"),
108            SnapshotError::WrongKind { found } => write!(
109                f,
110                "this is a `{found}` snapshot, not a `{KIND}` one\n  \
111                 a memory store is passed to `--memory`, not `--resume`"
112            ),
113            SnapshotError::UnsupportedVersion { found } => write!(
114                f,
115                "this snapshot declares version `{found}`; this build writes `{SNAPSHOT_VERSION}`"
116            ),
117            SnapshotError::DifferentArtifact { agent } => write!(
118                f,
119                "`{agent}` has changed since the run stopped\n  \
120                 continuing against a modified program produces a result that is neither \
121                 program's, and nothing in the record would say which parts came from which\n  \
122                 start the run again"
123            ),
124            SnapshotError::UnknownNode { node } => write!(
125                f,
126                "the snapshot continues from node `{node}`, which this artifact does not have"
127            ),
128        }
129    }
130}
131
132impl std::error::Error for SnapshotError {}
133
134/// The digest a snapshot identifies its artifact by.
135pub fn artifact_digest(ir: &AgentIr) -> String {
136    let mut hasher = Sha256::new();
137    hasher.update(ir.to_canonical_json().as_bytes());
138    format!("sha256:{:x}", hasher.finalize())
139}
140
141impl Resumption {
142    /// Read one, and refuse anything that is not one.
143    pub fn load(path: &Path) -> Result<Resumption, SnapshotError> {
144        let text = std::fs::read_to_string(path)
145            .map_err(|error| SnapshotError::Io(format!("reading {}: {error}", path.display())))?;
146        let snapshot: Resumption = serde_json::from_str(&text).map_err(|error| {
147            // A memory store parses far enough to have a `kind`, so say which
148            // file this is before complaining about its fields.
149            match serde_json::from_str::<serde_json::Value>(&text) {
150                Ok(value) if value.get("kind").and_then(Value::as_str) == Some("memory") => {
151                    SnapshotError::WrongKind {
152                        found: "memory".to_string(),
153                    }
154                }
155                _ => SnapshotError::Malformed(error.to_string()),
156            }
157        })?;
158
159        if snapshot.kind != KIND {
160            return Err(SnapshotError::WrongKind {
161                found: snapshot.kind,
162            });
163        }
164        if snapshot.ingot_snapshot != SNAPSHOT_VERSION {
165            return Err(SnapshotError::UnsupportedVersion {
166                found: snapshot.ingot_snapshot,
167            });
168        }
169        Ok(snapshot)
170    }
171
172    /// Write it, with sorted keys and a trailing newline.
173    ///
174    /// The same rule the IR follows: a file two identical runs produce
175    /// differently is a file nobody can diff.
176    pub fn save(&self, path: &Path) -> Result<(), SnapshotError> {
177        if let Some(parent) = path.parent() {
178            std::fs::create_dir_all(parent).map_err(|error| {
179                SnapshotError::Io(format!("creating {}: {error}", parent.display()))
180            })?;
181        }
182        let mut text = serde_json::to_string_pretty(self)
183            .map_err(|error| SnapshotError::Malformed(error.to_string()))?;
184        text.push('\n');
185        std::fs::write(path, text)
186            .map_err(|error| SnapshotError::Io(format!("writing {}: {error}", path.display())))
187    }
188
189    /// Whether this snapshot belongs to `ir`, and continues into a node it has.
190    ///
191    /// There is no override. See
192    /// [Runtime 0.5 §2.4](../../../specs/runtime/v0.5.md).
193    pub fn check(&self, ir: &AgentIr) -> Result<(), SnapshotError> {
194        if self.artifact != artifact_digest(ir) {
195            return Err(SnapshotError::DifferentArtifact {
196                agent: self.agent.clone(),
197            });
198        }
199        if ir.node(&self.resume_at).is_none() {
200            return Err(SnapshotError::UnknownNode {
201                node: self.resume_at.clone(),
202            });
203        }
204        Ok(())
205    }
206}
207
208/// Every label a run could be stopped at, in flow order.
209///
210/// Used to refuse `--stop-at` before the run starts, and to say what was
211/// available when the label does not match.
212pub fn resumable_labels(ir: &AgentIr) -> Vec<String> {
213    ir.nodes
214        .iter()
215        .filter(|node| node.resumable)
216        .filter_map(|node| node.label.clone())
217        .collect()
218}
219
220/// Every checkpoint label, resumable or not.
221///
222/// The difference between this and [`resumable_labels`] is what turns "no such
223/// checkpoint" into "that checkpoint is inside a loop".
224pub fn all_checkpoint_labels(ir: &AgentIr) -> Vec<String> {
225    ir.nodes
226        .iter()
227        .filter(|node| node.kind == ingot_ir::NodeKind::Checkpoint)
228        .filter_map(|node| node.label.clone())
229        .collect()
230}
231
232#[cfg(test)]
233mod tests {
234    use super::*;
235    use ingot_ir::{Budget, ModelRequirement, Node, NodeKind, Requirements, IR_VERSION};
236
237    fn artifact() -> AgentIr {
238        let mut checkpoint = Node::new("n0", NodeKind::Checkpoint);
239        checkpoint.label = Some("half-way".to_string());
240        checkpoint.resumable = true;
241        checkpoint.next = Some("n1".to_string());
242
243        let mut nested = Node::new("n1", NodeKind::Checkpoint);
244        nested.label = Some("inside".to_string());
245
246        AgentIr {
247            ir_version: IR_VERSION.to_string(),
248            language: "0.2".to_string(),
249            agent: "test.Stops".to_string(),
250            doc: None,
251            inputs: BTreeMap::new(),
252            outputs: BTreeMap::new(),
253            types: BTreeMap::new(),
254            requirements: Requirements {
255                model: ModelRequirement::Unspecified,
256            },
257            tools: Vec::new(),
258            state: BTreeMap::new(),
259            persistent: BTreeMap::new(),
260            budget: Budget::default(),
261            policy: BTreeMap::new(),
262            effects: Vec::new(),
263            entry: Some("n0".to_string()),
264            nodes: vec![checkpoint, nested],
265        }
266    }
267
268    fn snapshot(ir: &AgentIr) -> Resumption {
269        Resumption {
270            ingot_snapshot: SNAPSHOT_VERSION.to_string(),
271            kind: KIND.to_string(),
272            agent: ir.agent.clone(),
273            artifact: artifact_digest(ir),
274            label: "half-way".to_string(),
275            stopped_at: "n0".to_string(),
276            resume_at: "n1".to_string(),
277            inputs: BTreeMap::new(),
278            bindings: BTreeMap::new(),
279            state: BTreeMap::new(),
280            outputs: BTreeMap::new(),
281            steps: 3,
282            usage: Usage::default(),
283            spend: Spend::default(),
284            model_calls: 0,
285            tool_calls: 0,
286        }
287    }
288
289    #[test]
290    fn a_snapshot_round_trips_through_json() {
291        let ir = artifact();
292        let original = snapshot(&ir);
293        let text = serde_json::to_string(&original).unwrap();
294        let parsed: Resumption = serde_json::from_str(&text).unwrap();
295        assert_eq!(parsed, original);
296    }
297
298    #[test]
299    fn a_snapshot_belongs_to_exactly_one_artifact() {
300        let ir = artifact();
301        let snapshot = snapshot(&ir);
302        assert!(snapshot.check(&ir).is_ok());
303
304        // One node id renamed is a different program.
305        let mut edited = ir.clone();
306        edited.budget.steps = Some(9);
307        let error = snapshot.check(&edited).unwrap_err();
308        assert!(
309            matches!(error, SnapshotError::DifferentArtifact { .. }),
310            "{error}"
311        );
312        // The message says to start again, because there is no override.
313        assert!(error.to_string().contains("start the run again"));
314    }
315
316    #[test]
317    fn only_a_top_level_checkpoint_is_offered() {
318        let ir = artifact();
319        assert_eq!(resumable_labels(&ir), vec!["half-way".to_string()]);
320        assert_eq!(
321            all_checkpoint_labels(&ir),
322            vec!["half-way".to_string(), "inside".to_string()]
323        );
324    }
325
326    #[test]
327    fn a_memory_store_pointed_at_resume_says_which_file_it_is() {
328        let dir = std::env::temp_dir().join(format!("ingot-snapshot-{}", std::process::id()));
329        std::fs::create_dir_all(&dir).unwrap();
330        let path = dir.join("store.json");
331        std::fs::write(
332            &path,
333            r#"{"ingotSnapshot":"0.1","kind":"memory","agent":"a","shape":{},"fields":{}}"#,
334        )
335        .unwrap();
336        let error = Resumption::load(&path).unwrap_err();
337        assert!(
338            matches!(error, SnapshotError::WrongKind { ref found } if found == "memory"),
339            "{error}"
340        );
341        assert!(error.to_string().contains("--memory"));
342        let _ = std::fs::remove_dir_all(&dir);
343    }
344}
ingot_runtime/snapshot.rs

ingot_runtime/
snapshot.rs
