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    /// How many questions it put to a person, for the same reason. A recorded
80    /// answer is matched by position like everything else, so a resumed run has
81    /// to start where the first half stopped.
82    #[serde(default)]
83    pub consultations: u32,
84}
85
86/// Why a snapshot could not be used.
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub enum SnapshotError {
89    Io(String),
90    Malformed(String),
91    /// A snapshot of the other kind, or of a version this build does not write.
92    WrongKind {
93        found: String,
94    },
95    UnsupportedVersion {
96        found: String,
97    },
98    /// The artifact is not the one the run stopped in.
99    DifferentArtifact {
100        agent: String,
101    },
102    /// The node to continue from is gone.
103    UnknownNode {
104        node: String,
105    },
106}
107
108impl std::fmt::Display for SnapshotError {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        match self {
111            SnapshotError::Io(reason) => write!(f, "{reason}"),
112            SnapshotError::Malformed(reason) => write!(f, "this is not a snapshot: {reason}"),
113            SnapshotError::WrongKind { found } => write!(
114                f,
115                "this is a `{found}` snapshot, not a `{KIND}` one\n  \
116                 a memory store is passed to `--memory`, not `--resume`"
117            ),
118            SnapshotError::UnsupportedVersion { found } => write!(
119                f,
120                "this snapshot declares version `{found}`; this build writes `{SNAPSHOT_VERSION}`"
121            ),
122            SnapshotError::DifferentArtifact { agent } => write!(
123                f,
124                "`{agent}` has changed since the run stopped\n  \
125                 continuing against a modified program produces a result that is neither \
126                 program's, and nothing in the record would say which parts came from which\n  \
127                 start the run again"
128            ),
129            SnapshotError::UnknownNode { node } => write!(
130                f,
131                "the snapshot continues from node `{node}`, which this artifact does not have"
132            ),
133        }
134    }
135}
136
137impl std::error::Error for SnapshotError {}
138
139/// The digest a snapshot identifies its artifact by.
140pub fn artifact_digest(ir: &AgentIr) -> String {
141    let mut hasher = Sha256::new();
142    hasher.update(ir.to_canonical_json().as_bytes());
143    format!("sha256:{:x}", hasher.finalize())
144}
145
146impl Resumption {
147    /// Read one, and refuse anything that is not one.
148    pub fn load(path: &Path) -> Result<Resumption, SnapshotError> {
149        let text = std::fs::read_to_string(path)
150            .map_err(|error| SnapshotError::Io(format!("reading {}: {error}", path.display())))?;
151        let snapshot: Resumption = serde_json::from_str(&text).map_err(|error| {
152            // A memory store parses far enough to have a `kind`, so say which
153            // file this is before complaining about its fields.
154            match serde_json::from_str::<serde_json::Value>(&text) {
155                Ok(value) if value.get("kind").and_then(Value::as_str) == Some("memory") => {
156                    SnapshotError::WrongKind {
157                        found: "memory".to_string(),
158                    }
159                }
160                _ => SnapshotError::Malformed(error.to_string()),
161            }
162        })?;
163
164        if snapshot.kind != KIND {
165            return Err(SnapshotError::WrongKind {
166                found: snapshot.kind,
167            });
168        }
169        if snapshot.ingot_snapshot != SNAPSHOT_VERSION {
170            return Err(SnapshotError::UnsupportedVersion {
171                found: snapshot.ingot_snapshot,
172            });
173        }
174        Ok(snapshot)
175    }
176
177    /// Write it, with sorted keys and a trailing newline.
178    ///
179    /// The same rule the IR follows: a file two identical runs produce
180    /// differently is a file nobody can diff.
181    pub fn save(&self, path: &Path) -> Result<(), SnapshotError> {
182        if let Some(parent) = path.parent() {
183            std::fs::create_dir_all(parent).map_err(|error| {
184                SnapshotError::Io(format!("creating {}: {error}", parent.display()))
185            })?;
186        }
187        let mut text = serde_json::to_string_pretty(self)
188            .map_err(|error| SnapshotError::Malformed(error.to_string()))?;
189        text.push('\n');
190        std::fs::write(path, text)
191            .map_err(|error| SnapshotError::Io(format!("writing {}: {error}", path.display())))
192    }
193
194    /// Whether this snapshot belongs to `ir`, and continues into a node it has.
195    ///
196    /// There is no override. See
197    /// [Runtime 0.5 §2.4](../../../specs/runtime/v0.5.md).
198    pub fn check(&self, ir: &AgentIr) -> Result<(), SnapshotError> {
199        if self.artifact != artifact_digest(ir) {
200            return Err(SnapshotError::DifferentArtifact {
201                agent: self.agent.clone(),
202            });
203        }
204        if ir.node(&self.resume_at).is_none() {
205            return Err(SnapshotError::UnknownNode {
206                node: self.resume_at.clone(),
207            });
208        }
209        Ok(())
210    }
211}
212
213/// Every label a run could be stopped at, in flow order.
214///
215/// Used to refuse `--stop-at` before the run starts, and to say what was
216/// available when the label does not match.
217pub fn resumable_labels(ir: &AgentIr) -> Vec<String> {
218    ir.nodes
219        .iter()
220        .filter(|node| node.resumable)
221        .filter_map(|node| node.label.clone())
222        .collect()
223}
224
225/// Every checkpoint label, resumable or not.
226///
227/// The difference between this and [`resumable_labels`] is what turns "no such
228/// checkpoint" into "that checkpoint is inside a loop".
229pub fn all_checkpoint_labels(ir: &AgentIr) -> Vec<String> {
230    ir.nodes
231        .iter()
232        .filter(|node| node.kind == ingot_ir::NodeKind::Checkpoint)
233        .filter_map(|node| node.label.clone())
234        .collect()
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use ingot_ir::{Budget, ModelRequirement, Node, NodeKind, Requirements, IR_VERSION};
241
242    fn artifact() -> AgentIr {
243        let mut checkpoint = Node::new("n0", NodeKind::Checkpoint);
244        checkpoint.label = Some("half-way".to_string());
245        checkpoint.resumable = true;
246        checkpoint.next = Some("n1".to_string());
247
248        let mut nested = Node::new("n1", NodeKind::Checkpoint);
249        nested.label = Some("inside".to_string());
250
251        AgentIr {
252            ir_version: IR_VERSION.to_string(),
253            language: "0.2".to_string(),
254            agent: "test.Stops".to_string(),
255            doc: None,
256            inputs: BTreeMap::new(),
257            outputs: BTreeMap::new(),
258            types: BTreeMap::new(),
259            requirements: Requirements {
260                model: ModelRequirement::Unspecified,
261            },
262            tools: Vec::new(),
263            state: BTreeMap::new(),
264            persistent: BTreeMap::new(),
265            budget: Budget::default(),
266            policy: BTreeMap::new(),
267            effects: Vec::new(),
268            entry: Some("n0".to_string()),
269            nodes: vec![checkpoint, nested],
270        }
271    }
272
273    fn snapshot(ir: &AgentIr) -> Resumption {
274        Resumption {
275            ingot_snapshot: SNAPSHOT_VERSION.to_string(),
276            kind: KIND.to_string(),
277            agent: ir.agent.clone(),
278            artifact: artifact_digest(ir),
279            label: "half-way".to_string(),
280            stopped_at: "n0".to_string(),
281            resume_at: "n1".to_string(),
282            inputs: BTreeMap::new(),
283            bindings: BTreeMap::new(),
284            state: BTreeMap::new(),
285            outputs: BTreeMap::new(),
286            steps: 3,
287            usage: Usage::default(),
288            spend: Spend::default(),
289            model_calls: 0,
290            tool_calls: 0,
291            consultations: 0,
292        }
293    }
294
295    #[test]
296    fn a_snapshot_round_trips_through_json() {
297        let ir = artifact();
298        let original = snapshot(&ir);
299        let text = serde_json::to_string(&original).unwrap();
300        let parsed: Resumption = serde_json::from_str(&text).unwrap();
301        assert_eq!(parsed, original);
302    }
303
304    #[test]
305    fn a_snapshot_belongs_to_exactly_one_artifact() {
306        let ir = artifact();
307        let snapshot = snapshot(&ir);
308        assert!(snapshot.check(&ir).is_ok());
309
310        // One node id renamed is a different program.
311        let mut edited = ir.clone();
312        edited.budget.steps = Some(9);
313        let error = snapshot.check(&edited).unwrap_err();
314        assert!(
315            matches!(error, SnapshotError::DifferentArtifact { .. }),
316            "{error}"
317        );
318        // The message says to start again, because there is no override.
319        assert!(error.to_string().contains("start the run again"));
320    }
321
322    #[test]
323    fn only_a_top_level_checkpoint_is_offered() {
324        let ir = artifact();
325        assert_eq!(resumable_labels(&ir), vec!["half-way".to_string()]);
326        assert_eq!(
327            all_checkpoint_labels(&ir),
328            vec!["half-way".to_string(), "inside".to_string()]
329        );
330    }
331
332    #[test]
333    fn a_memory_store_pointed_at_resume_says_which_file_it_is() {
334        let dir = std::env::temp_dir().join(format!("ingot-snapshot-{}", std::process::id()));
335        std::fs::create_dir_all(&dir).unwrap();
336        let path = dir.join("store.json");
337        std::fs::write(
338            &path,
339            r#"{"ingotSnapshot":"0.1","kind":"memory","agent":"a","shape":{},"fields":{}}"#,
340        )
341        .unwrap();
342        let error = Resumption::load(&path).unwrap_err();
343        assert!(
344            matches!(error, SnapshotError::WrongKind { ref found } if found == "memory"),
345            "{error}"
346        );
347        assert!(error.to_string().contains("--memory"));
348        let _ = std::fs::remove_dir_all(&dir);
349    }
350}
ingot_runtime/snapshot.rs

ingot_runtime/
snapshot.rs
