1use std::collections::BTreeMap;
2use std::io::{Error, ErrorKind, Result};
3use std::sync::{Arc, Mutex};
4use std::time::Instant;
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8
9use crate::budget::BudgetUsageSnapshot;
10use crate::types::{AgentStatus, CycleRecord, Message, Metadata};
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13pub struct Checkpoint {
14 pub task_id: String,
15 pub cycle_index: u32,
16 pub status: AgentStatus,
17 pub messages: Vec<Message>,
18 pub cycles: Vec<CycleRecord>,
19 pub shared_state: Metadata,
20 #[serde(default)]
21 pub revision: u64,
22 #[serde(default)]
23 pub claim_token: Option<String>,
24 #[serde(default)]
25 pub claimed_cycle: Option<u32>,
26 #[serde(default)]
27 pub lease_expires_at_ms: Option<u64>,
28 #[serde(default)]
29 pub terminal_result: Option<crate::types::AgentResult>,
30 #[serde(default)]
31 pub budget_usage: Option<BudgetUsageSnapshot>,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(rename_all = "snake_case")]
36pub enum StateStoreKind {
37 Sqlite,
38 Redis,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
42pub struct StateStoreSpec {
43 pub kind: StateStoreKind,
44 pub location: String,
45}
46
47impl StateStoreSpec {
48 pub fn sqlite(location: impl Into<String>) -> Result<Self> {
49 Self::new(StateStoreKind::Sqlite, location)
50 }
51
52 pub fn redis(location: impl Into<String>) -> Result<Self> {
53 Self::new(StateStoreKind::Redis, location)
54 }
55
56 fn new(kind: StateStoreKind, location: impl Into<String>) -> Result<Self> {
57 let location = location.into();
58 if location.trim().is_empty() {
59 return Err(Error::new(
60 ErrorKind::InvalidInput,
61 "state store location must be a non-empty string",
62 ));
63 }
64 Ok(Self { kind, location })
65 }
66
67 pub fn to_dict(&self) -> Value {
68 serde_json::json!({
69 "kind": match self.kind {
70 StateStoreKind::Sqlite => "sqlite",
71 StateStoreKind::Redis => "redis",
72 },
73 "location": self.location,
74 })
75 }
76
77 pub fn from_dict(payload: &Value) -> Result<Self> {
78 let object = payload
79 .as_object()
80 .ok_or_else(|| Error::new(ErrorKind::InvalidData, "state_store must be an object"))?;
81 let kind = match object.get("kind").and_then(Value::as_str) {
82 Some("sqlite") => StateStoreKind::Sqlite,
83 Some("redis") => StateStoreKind::Redis,
84 _ => {
85 return Err(Error::new(
86 ErrorKind::InvalidData,
87 "state_store.kind must be 'sqlite' or 'redis'",
88 ))
89 }
90 };
91 let location = object
92 .get("location")
93 .and_then(Value::as_str)
94 .ok_or_else(|| {
95 Error::new(
96 ErrorKind::InvalidData,
97 "state_store.location must be a non-empty string",
98 )
99 })?;
100 Self::new(kind, location)
101 }
102
103 pub fn build(&self) -> Result<Arc<dyn StateStore>> {
104 match self.kind {
105 StateStoreKind::Sqlite => Ok(Arc::new(
106 crate::runtime::stores::sqlite::SqliteStateStore::new(&self.location)?,
107 )),
108 StateStoreKind::Redis => Ok(Arc::new(
109 crate::runtime::stores::redis::RedisStateStore::new(&self.location)?,
110 )),
111 }
112 }
113}
114
115pub trait StateStore: Send + Sync {
116 fn create_checkpoint(&self, checkpoint: Checkpoint) -> Result<bool>;
117 fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()>;
118 fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>>;
119 fn claim_checkpoint(
120 &self,
121 task_id: &str,
122 cycle_index: u32,
123 claim_token: &str,
124 lease_expires_at_ms: u64,
125 now_ms: u64,
126 ) -> Result<Option<Checkpoint>>;
127 fn commit_checkpoint(
128 &self,
129 checkpoint: Checkpoint,
130 claim_token: &str,
131 expected_revision: u64,
132 ) -> Result<bool>;
133 fn renew_checkpoint_claim(
134 &self,
135 task_id: &str,
136 claim_token: &str,
137 expected_revision: u64,
138 lease_expires_at_ms: u64,
139 now_ms: u64,
140 ) -> Result<bool>;
141 fn finalize_checkpoint(&self, checkpoint: Checkpoint, expected_revision: u64) -> Result<bool>;
142 fn delete_checkpoint(&self, task_id: &str) -> Result<()>;
143 fn acknowledge_terminal(&self, task_id: &str, expected_revision: u64) -> Result<bool>;
144 fn list_checkpoints(&self) -> Result<Vec<String>>;
145 fn state_store_spec(&self) -> Option<StateStoreSpec>;
146}
147
148pub(crate) struct LeaseOperationClock {
149 now_ms: u64,
150 started: Instant,
151}
152
153impl LeaseOperationClock {
154 pub(crate) fn new(now_ms: u64) -> Self {
155 Self {
156 now_ms,
157 started: Instant::now(),
158 }
159 }
160
161 pub(crate) fn now_ms(&self) -> u64 {
162 self.now_ms
163 .saturating_add(u64::try_from(self.started.elapsed().as_millis()).unwrap_or(u64::MAX))
164 }
165}
166
167#[derive(Debug, Clone, Default)]
168pub struct InMemoryStateStore {
169 checkpoints: Arc<Mutex<BTreeMap<String, Checkpoint>>>,
170}
171
172impl InMemoryStateStore {
173 pub fn new() -> Self {
174 Self::default()
175 }
176}
177
178impl StateStore for InMemoryStateStore {
179 fn create_checkpoint(&self, checkpoint: Checkpoint) -> Result<bool> {
180 crate::runtime::checkpoint_codec::validate_checkpoint(&checkpoint)?;
181 let mut checkpoints = self
182 .checkpoints
183 .lock()
184 .map_err(|_| poisoned("state store"))?;
185 if checkpoints.contains_key(&checkpoint.task_id) {
186 return Ok(false);
187 }
188 checkpoints.insert(checkpoint.task_id.clone(), checkpoint);
189 Ok(true)
190 }
191
192 fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()> {
193 crate::runtime::checkpoint_codec::validate_checkpoint(&checkpoint)?;
194 self.checkpoints
195 .lock()
196 .map_err(|_| poisoned("state store"))?
197 .insert(checkpoint.task_id.clone(), checkpoint);
198 Ok(())
199 }
200
201 fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>> {
202 Ok(self
203 .checkpoints
204 .lock()
205 .map_err(|_| poisoned("state store"))?
206 .get(task_id)
207 .cloned())
208 }
209
210 fn claim_checkpoint(
211 &self,
212 task_id: &str,
213 cycle_index: u32,
214 claim_token: &str,
215 lease_expires_at_ms: u64,
216 now_ms: u64,
217 ) -> Result<Option<Checkpoint>> {
218 validate_claim(cycle_index, claim_token, lease_expires_at_ms, now_ms)?;
219 let mut checkpoints = self
220 .checkpoints
221 .lock()
222 .map_err(|_| poisoned("state store"))?;
223 let Some(checkpoint) = checkpoints.get_mut(task_id) else {
224 return Ok(None);
225 };
226 check_claim(checkpoint, cycle_index, now_ms)?;
227 checkpoint.revision = checkpoint.revision.saturating_add(1);
228 checkpoint.claim_token = Some(claim_token.to_string());
229 checkpoint.claimed_cycle = Some(cycle_index);
230 checkpoint.lease_expires_at_ms = Some(lease_expires_at_ms);
231 Ok(Some(checkpoint.clone()))
232 }
233
234 fn commit_checkpoint(
235 &self,
236 mut checkpoint: Checkpoint,
237 claim_token: &str,
238 expected_revision: u64,
239 ) -> Result<bool> {
240 let mut checkpoints = self
241 .checkpoints
242 .lock()
243 .map_err(|_| poisoned("state store"))?;
244 let current = checkpoints.get(&checkpoint.task_id);
245 if !claim_matches(current, &checkpoint, claim_token, expected_revision) {
246 return Ok(false);
247 }
248 checkpoint.revision = expected_revision.saturating_add(1);
249 clear_claim(&mut checkpoint);
250 crate::runtime::checkpoint_codec::validate_checkpoint(&checkpoint)?;
251 checkpoints.insert(checkpoint.task_id.clone(), checkpoint);
252 Ok(true)
253 }
254
255 fn renew_checkpoint_claim(
256 &self,
257 task_id: &str,
258 claim_token: &str,
259 expected_revision: u64,
260 lease_expires_at_ms: u64,
261 now_ms: u64,
262 ) -> Result<bool> {
263 validate_renew(claim_token, expected_revision, lease_expires_at_ms, now_ms)?;
264 let clock = LeaseOperationClock::new(now_ms);
265 let mut checkpoints = self
266 .checkpoints
267 .lock()
268 .map_err(|_| poisoned("state store"))?;
269 let current_now_ms = clock.now_ms();
270 let Some(checkpoint) = checkpoints.get_mut(task_id) else {
271 return Ok(false);
272 };
273 if checkpoint.revision != expected_revision
274 || checkpoint.claim_token.as_deref() != Some(claim_token)
275 || checkpoint.lease_expires_at_ms.unwrap_or(0) <= current_now_ms
276 || lease_expires_at_ms <= current_now_ms
277 {
278 return Ok(false);
279 }
280 checkpoint.lease_expires_at_ms = Some(lease_expires_at_ms);
281 Ok(true)
282 }
283
284 fn finalize_checkpoint(
285 &self,
286 mut checkpoint: Checkpoint,
287 expected_revision: u64,
288 ) -> Result<bool> {
289 if checkpoint.terminal_result.is_none() {
290 return Err(Error::new(
291 ErrorKind::InvalidInput,
292 "finalized checkpoint must include terminal_result",
293 ));
294 }
295 crate::runtime::checkpoint_codec::validate_checkpoint(&checkpoint)?;
296 let mut checkpoints = self
297 .checkpoints
298 .lock()
299 .map_err(|_| poisoned("state store"))?;
300 let Some(current) = checkpoints.get(&checkpoint.task_id) else {
301 return Ok(false);
302 };
303 if current.revision != expected_revision
304 || current.claim_token.is_some()
305 || current.terminal_result.is_some()
306 {
307 return Ok(false);
308 }
309 checkpoint.revision = expected_revision.saturating_add(1);
310 clear_claim(&mut checkpoint);
311 checkpoints.insert(checkpoint.task_id.clone(), checkpoint);
312 Ok(true)
313 }
314
315 fn delete_checkpoint(&self, task_id: &str) -> Result<()> {
316 self.checkpoints
317 .lock()
318 .map_err(|_| poisoned("state store"))?
319 .remove(task_id);
320 Ok(())
321 }
322
323 fn acknowledge_terminal(&self, task_id: &str, expected_revision: u64) -> Result<bool> {
324 let mut checkpoints = self
325 .checkpoints
326 .lock()
327 .map_err(|_| poisoned("state store"))?;
328 if checkpoints.get(task_id).is_none_or(|checkpoint| {
329 checkpoint.revision != expected_revision || checkpoint.terminal_result.is_none()
330 }) {
331 return Ok(false);
332 }
333 checkpoints.remove(task_id);
334 Ok(true)
335 }
336
337 fn list_checkpoints(&self) -> Result<Vec<String>> {
338 Ok(self
339 .checkpoints
340 .lock()
341 .map_err(|_| poisoned("state store"))?
342 .keys()
343 .cloned()
344 .collect())
345 }
346
347 fn state_store_spec(&self) -> Option<StateStoreSpec> {
348 None
349 }
350}
351
352pub(crate) fn checkpoint_status_value(status: AgentStatus) -> &'static str {
353 match status {
354 AgentStatus::Pending => "pending",
355 AgentStatus::Running => "running",
356 AgentStatus::WaitUser => "wait_user",
357 AgentStatus::Completed => "completed",
358 AgentStatus::Failed => "failed",
359 AgentStatus::MaxCycles => "max_cycles",
360 AgentStatus::ReconciliationRequired => "reconciliation_required",
361 }
362}
363
364pub(crate) fn checkpoint_status_from_value(value: &str) -> Result<AgentStatus> {
365 match value {
366 "pending" => Ok(AgentStatus::Pending),
367 "running" => Ok(AgentStatus::Running),
368 "wait_user" => Ok(AgentStatus::WaitUser),
369 "completed" => Ok(AgentStatus::Completed),
370 "failed" => Ok(AgentStatus::Failed),
371 "max_cycles" => Ok(AgentStatus::MaxCycles),
372 "reconciliation_required" => Ok(AgentStatus::ReconciliationRequired),
373 other => Err(Error::new(
374 ErrorKind::InvalidData,
375 format!("unknown checkpoint status: {other}"),
376 )),
377 }
378}
379
380fn poisoned(name: &str) -> Error {
381 Error::other(format!("{name} lock is poisoned"))
382}
383
384pub(crate) fn validate_claim(
385 cycle_index: u32,
386 claim_token: &str,
387 lease_expires_at_ms: u64,
388 now_ms: u64,
389) -> Result<()> {
390 if cycle_index == 0 {
391 return Err(Error::new(
392 ErrorKind::InvalidInput,
393 "claimed cycle_index must be between 1 and 4294967295",
394 ));
395 }
396 if claim_token.is_empty() {
397 return Err(Error::new(
398 ErrorKind::InvalidInput,
399 "claim_token must be a non-empty string",
400 ));
401 }
402 if lease_expires_at_ms <= now_ms {
403 return Err(Error::new(
404 ErrorKind::InvalidInput,
405 "lease_expires_at_ms must be greater than now_ms",
406 ));
407 }
408 Ok(())
409}
410
411pub(crate) fn validate_renew(
412 claim_token: &str,
413 _expected_revision: u64,
414 lease_expires_at_ms: u64,
415 now_ms: u64,
416) -> Result<()> {
417 if claim_token.is_empty() {
418 return Err(Error::new(
419 ErrorKind::InvalidInput,
420 "claim_token must be a non-empty string",
421 ));
422 }
423 if lease_expires_at_ms <= now_ms {
424 return Err(Error::new(
425 ErrorKind::InvalidInput,
426 "lease_expires_at_ms must be greater than now_ms",
427 ));
428 }
429 Ok(())
430}
431
432pub(crate) fn check_claim(checkpoint: &Checkpoint, cycle_index: u32, now_ms: u64) -> Result<()> {
433 let expected_cycle = cycle_index - 1;
434 if checkpoint.terminal_result.is_some() || checkpoint.status != AgentStatus::Running {
435 return Err(Error::new(
436 ErrorKind::AlreadyExists,
437 format!("checkpoint for task {} is terminal", checkpoint.task_id),
438 ));
439 }
440 if checkpoint.cycle_index != expected_cycle {
441 return Err(Error::new(
442 ErrorKind::AlreadyExists,
443 format!(
444 "checkpoint cycle conflict for task {}: expected {}, found {}",
445 checkpoint.task_id, expected_cycle, checkpoint.cycle_index
446 ),
447 ));
448 }
449 if checkpoint.claim_token.is_some() && checkpoint.lease_expires_at_ms.unwrap_or(0) > now_ms {
450 return Err(Error::new(
451 ErrorKind::AlreadyExists,
452 format!(
453 "checkpoint cycle {cycle_index} for task {} is already claimed",
454 checkpoint.task_id
455 ),
456 ));
457 }
458 Ok(())
459}
460
461pub(crate) fn claim_matches(
462 current: Option<&Checkpoint>,
463 checkpoint: &Checkpoint,
464 claim_token: &str,
465 expected_revision: u64,
466) -> bool {
467 current.is_some_and(|current| {
468 current.revision == expected_revision
469 && current.claim_token.as_deref() == Some(claim_token)
470 && current.claimed_cycle.is_some()
471 && checkpoint.cycle_index == current.claimed_cycle.unwrap_or_default()
472 })
473}
474
475pub(crate) fn clear_claim(checkpoint: &mut Checkpoint) {
476 checkpoint.claim_token = None;
477 checkpoint.claimed_cycle = None;
478 checkpoint.lease_expires_at_ms = None;
479}