a3s_code_core/loop_checkpoint.rs
1//! Per-tool-round loop checkpoints for crash-tolerant runs (P3 cut 1).
2//!
3//! The agent loop persists a [`LoopCheckpoint`] after each completed tool
4//! round. The checkpoint captures the minimum state needed to recreate
5//! the loop's position so a future process — typically on a different
6//! node, dispatched by the host after a crash or planned migration — can
7//! resume from the last consistent boundary.
8//!
9//! Boundary policy: checkpoints are taken **only** between tool rounds,
10//! never mid-tool. If a process dies while a tool is executing, the
11//! work of that round is lost on resume; the LLM re-deliberates from
12//! the previous checkpoint. This trades retry cost for correctness —
13//! re-executing a non-idempotent tool (write, bash) on the wrong side
14//! of the boundary is worse than re-asking the LLM.
15//!
16//! [`crate::AgentSession::resume_run`] restores the checkpoint on a fresh run
17//! while preserving cumulative usage, turn budgets, and convergence guards.
18
19use crate::llm::{Message, TokenUsage};
20use crate::verification::VerificationReport;
21use async_trait::async_trait;
22use serde::{Deserialize, Serialize};
23
24/// Schema version. Bumped on incompatible format changes; impls of
25/// [`LoopCheckpointSink`] should reject loads from a future version.
26pub const LOOP_CHECKPOINT_SCHEMA_VERSION: u32 = 1;
27
28/// Loop state that must survive crash recovery to preserve convergence limits.
29#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
30pub struct LoopConvergenceState {
31 #[serde(default)]
32 pub parse_error_count: u32,
33 #[serde(default)]
34 pub continuation_count: u32,
35 #[serde(default)]
36 pub reasoning_only_repair_count: u32,
37 /// Tool name, SHA-256 argument fingerprint, and outcome; never raw args.
38 #[serde(default)]
39 pub recent_tool_signatures: Vec<String>,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
41 pub guarded_duplicate_signature: Option<String>,
42 #[serde(default)]
43 pub guarded_duplicate_count: u32,
44 #[serde(default, skip_serializing_if = "Option::is_none")]
45 pub last_incomplete_response_hash: Option<String>,
46 #[serde(default)]
47 pub incomplete_response_stalled: bool,
48}
49
50/// Snapshot of the agent loop at the boundary between tool rounds.
51///
52/// Stored under `run_id` so resume tooling can address the correct run
53/// without scanning all checkpoints of a session.
54#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct LoopCheckpoint {
56 /// Schema version — see [`LOOP_CHECKPOINT_SCHEMA_VERSION`].
57 #[serde(default)]
58 pub schema_version: u32,
59
60 /// Logical run identifier. Matches the `run_id` carried by
61 /// [`crate::run::RunSnapshot`] and `AgentEvent`s.
62 pub run_id: String,
63
64 /// Parent session id — redundant with `run_id` lookup but useful
65 /// for store layouts that key by `(session_id, run_id)`.
66 pub session_id: String,
67
68 /// 1-based tool round counter at checkpoint time.
69 /// `0` is reserved for "no rounds completed yet".
70 pub turn: usize,
71
72 /// Conversation history including the just-returned tool results.
73 /// On resume, the new agent loop starts from this exact message list.
74 pub messages: Vec<Message>,
75
76 /// Running token usage at checkpoint time. Lets resume re-emit
77 /// progress metrics without re-querying the LLM provider.
78 pub total_usage: TokenUsage,
79
80 /// How many tool calls have been executed total in this run.
81 pub tool_calls_count: usize,
82
83 /// Verification reports collected so far in this run.
84 #[serde(default)]
85 pub verification_reports: Vec<VerificationReport>,
86
87 /// Counters and fingerprints that prevent convergence budgets from
88 /// resetting when a run resumes on another process.
89 #[serde(default)]
90 pub convergence: LoopConvergenceState,
91
92 /// Wall-clock timestamp when the checkpoint was written
93 /// (Unix epoch ms — sourced from the session's
94 /// [`HostEnv`](crate::host_env::HostEnv)).
95 pub checkpoint_ms: u64,
96}
97
98impl LoopCheckpoint {
99 /// Reject a checkpoint written by a *newer*, incompatible schema
100 /// version than this build understands — honoring the contract on
101 /// [`LOOP_CHECKPOINT_SCHEMA_VERSION`].
102 ///
103 /// Field *additions* are absorbed transparently by `#[serde(default)]`,
104 /// so an older checkpoint (lower `schema_version`, including a pre-v1
105 /// `0`) always remains loadable. A *future* version, however, may have
106 /// changed the meaning of existing fields or the tool-round boundary
107 /// semantics; resuming from one risks silent corruption (e.g.
108 /// re-running a non-idempotent tool on the wrong side of the boundary).
109 ///
110 /// [`SessionStore`](crate::store::SessionStore) impls call this right
111 /// after deserialization, so both `resume_run` (which surfaces the
112 /// error to the caller) and the live-run [`LoopCheckpointSink`] (which
113 /// logs and starts fresh) refuse to act on an unreadable checkpoint.
114 pub fn ensure_loadable(&self) -> anyhow::Result<()> {
115 if self.schema_version > LOOP_CHECKPOINT_SCHEMA_VERSION {
116 anyhow::bail!(
117 "loop checkpoint for run {} has schema version {} but this build supports at \
118 most {}; refusing to resume from an incompatible future checkpoint",
119 self.run_id,
120 self.schema_version,
121 LOOP_CHECKPOINT_SCHEMA_VERSION
122 );
123 }
124 Ok(())
125 }
126
127 /// Verify that this value is stored under the run key it claims to own.
128 ///
129 /// Store keys are caller-controlled. Checking the redundant identifier in
130 /// the payload prevents a record written under one key from being replayed
131 /// as a different run.
132 pub fn ensure_addressed_by(&self, run_id: &str) -> anyhow::Result<()> {
133 if self.run_id != run_id {
134 anyhow::bail!(
135 "loop checkpoint key mismatch: requested run {:?}, payload belongs to {:?}",
136 run_id,
137 self.run_id
138 );
139 }
140 Ok(())
141 }
142
143 /// Verify both the addressed run and its owning session before replay.
144 pub fn ensure_owned_by(&self, run_id: &str, session_id: &str) -> anyhow::Result<()> {
145 self.ensure_addressed_by(run_id)?;
146 if self.session_id != session_id {
147 anyhow::bail!(
148 "loop checkpoint ownership mismatch for run {:?}: current session is {:?}, payload belongs to {:?}",
149 run_id,
150 session_id,
151 self.session_id
152 );
153 }
154 Ok(())
155 }
156}
157
158/// Receiver of per-tool-round checkpoints.
159///
160/// The framework ships one adapter:
161/// [`SessionStoreCheckpointSink`] which forwards to a
162/// [`crate::store::SessionStore`]. Hosts can implement custom sinks
163/// (e.g. push directly to Redis) by implementing this trait.
164#[async_trait]
165pub trait LoopCheckpointSink: Send + Sync {
166 /// Persist a checkpoint. Called from inside the agent loop after a
167 /// successful tool round. Errors are logged at warn level and
168 /// otherwise swallowed — losing a checkpoint must not halt the
169 /// live run.
170 async fn save_checkpoint(&self, checkpoint: &LoopCheckpoint);
171
172 /// Load the latest checkpoint for `run_id`, if any. Returns `None`
173 /// when no checkpoint has been recorded.
174 async fn load_latest(&self, run_id: &str) -> Option<LoopCheckpoint>;
175}
176
177/// Default adapter that forwards checkpoints to a
178/// [`SessionStore`](crate::store::SessionStore). Construct via
179/// [`SessionStoreCheckpointSink::new`].
180pub struct SessionStoreCheckpointSink {
181 inner: std::sync::Arc<dyn crate::store::SessionStore>,
182}
183
184impl SessionStoreCheckpointSink {
185 pub fn new(store: std::sync::Arc<dyn crate::store::SessionStore>) -> Self {
186 Self { inner: store }
187 }
188}
189
190#[async_trait]
191impl LoopCheckpointSink for SessionStoreCheckpointSink {
192 async fn save_checkpoint(&self, checkpoint: &LoopCheckpoint) {
193 if let Err(e) = self
194 .inner
195 .save_loop_checkpoint(&checkpoint.run_id, checkpoint)
196 .await
197 {
198 tracing::warn!(
199 run_id = %checkpoint.run_id,
200 error = %e,
201 "Loop checkpoint save failed; live run continues"
202 );
203 }
204 }
205
206 async fn load_latest(&self, run_id: &str) -> Option<LoopCheckpoint> {
207 match self.inner.load_loop_checkpoint(run_id).await {
208 Ok(opt) => opt,
209 Err(e) => {
210 tracing::warn!(
211 run_id = %run_id,
212 error = %e,
213 "Loop checkpoint load failed"
214 );
215 None
216 }
217 }
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 fn sample(run_id: &str, turn: usize) -> LoopCheckpoint {
226 LoopCheckpoint {
227 schema_version: LOOP_CHECKPOINT_SCHEMA_VERSION,
228 run_id: run_id.to_string(),
229 session_id: "session-1".to_string(),
230 turn,
231 messages: vec![Message::user("hi")],
232 total_usage: TokenUsage::default(),
233 tool_calls_count: 0,
234 verification_reports: Vec::new(),
235 convergence: LoopConvergenceState::default(),
236 checkpoint_ms: 1_700_000_000_000,
237 }
238 }
239
240 #[test]
241 fn checkpoint_round_trips_through_json() {
242 let mut cp = sample("run-1", 3);
243 cp.convergence.continuation_count = 2;
244 cp.convergence.recent_tool_signatures = vec!["read:deadbeef => ok".to_string()];
245 let json = serde_json::to_string(&cp).unwrap();
246 let back: LoopCheckpoint = serde_json::from_str(&json).unwrap();
247 assert_eq!(back.run_id, "run-1");
248 assert_eq!(back.turn, 3);
249 assert_eq!(back.schema_version, LOOP_CHECKPOINT_SCHEMA_VERSION);
250 assert_eq!(back.convergence, cp.convergence);
251 }
252
253 #[test]
254 fn missing_schema_version_defaults_to_zero() {
255 // Older payloads without the field must still load — they'll
256 // be interpreted as a pre-v1 snapshot.
257 let json = r#"{
258 "run_id": "run-1",
259 "session_id": "s",
260 "turn": 1,
261 "messages": [],
262 "total_usage": {"prompt_tokens":0,"completion_tokens":0,"total_tokens":0},
263 "tool_calls_count": 0,
264 "checkpoint_ms": 0
265 }"#;
266 let cp: LoopCheckpoint = serde_json::from_str(json).unwrap();
267 assert_eq!(cp.schema_version, 0);
268 assert_eq!(cp.convergence, LoopConvergenceState::default());
269 }
270
271 #[test]
272 fn checkpoint_rejects_run_key_and_session_owner_mismatches() {
273 let cp = sample("run-1", 1);
274 assert!(cp.ensure_owned_by("run-1", "session-1").is_ok());
275
276 let run_error = cp.ensure_owned_by("run-2", "session-1").unwrap_err();
277 assert!(run_error.to_string().contains("key mismatch"));
278
279 let session_error = cp.ensure_owned_by("run-1", "session-2").unwrap_err();
280 assert!(session_error.to_string().contains("ownership mismatch"));
281 }
282}