vv_agent/runtime/
state.rs

1use std::collections::BTreeMap;
2use std::io::{Error, ErrorKind, Result};
3use std::sync::{Arc, Mutex};
4
5use serde::{Deserialize, Serialize};
6
7use crate::types::{AgentStatus, CycleRecord, Message, Metadata};
8
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10pub struct Checkpoint {
11    pub task_id: String,
12    pub cycle_index: u32,
13    pub status: AgentStatus,
14    pub messages: Vec<Message>,
15    pub cycles: Vec<CycleRecord>,
16    pub shared_state: Metadata,
17}
18
19pub trait StateStore: Send + Sync {
20    fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()>;
21    fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>>;
22    fn delete_checkpoint(&self, task_id: &str) -> Result<()>;
23    fn list_checkpoints(&self) -> Result<Vec<String>>;
24}
25
26#[derive(Debug, Clone, Default)]
27pub struct InMemoryStateStore {
28    checkpoints: Arc<Mutex<BTreeMap<String, Checkpoint>>>,
29}
30
31impl InMemoryStateStore {
32    pub fn new() -> Self {
33        Self::default()
34    }
35}
36
37impl StateStore for InMemoryStateStore {
38    fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()> {
39        self.checkpoints
40            .lock()
41            .map_err(|_| poisoned("state store"))?
42            .insert(checkpoint.task_id.clone(), checkpoint);
43        Ok(())
44    }
45
46    fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>> {
47        Ok(self
48            .checkpoints
49            .lock()
50            .map_err(|_| poisoned("state store"))?
51            .get(task_id)
52            .cloned())
53    }
54
55    fn delete_checkpoint(&self, task_id: &str) -> Result<()> {
56        self.checkpoints
57            .lock()
58            .map_err(|_| poisoned("state store"))?
59            .remove(task_id);
60        Ok(())
61    }
62
63    fn list_checkpoints(&self) -> Result<Vec<String>> {
64        Ok(self
65            .checkpoints
66            .lock()
67            .map_err(|_| poisoned("state store"))?
68            .keys()
69            .cloned()
70            .collect())
71    }
72}
73
74pub(crate) fn checkpoint_status_value(status: AgentStatus) -> &'static str {
75    match status {
76        AgentStatus::Pending => "pending",
77        AgentStatus::Running => "running",
78        AgentStatus::WaitUser => "wait_user",
79        AgentStatus::Completed => "completed",
80        AgentStatus::Failed => "failed",
81        AgentStatus::MaxCycles => "max_cycles",
82    }
83}
84
85pub(crate) fn checkpoint_status_from_value(value: &str) -> Result<AgentStatus> {
86    match value {
87        "pending" => Ok(AgentStatus::Pending),
88        "running" => Ok(AgentStatus::Running),
89        "wait_user" => Ok(AgentStatus::WaitUser),
90        "completed" => Ok(AgentStatus::Completed),
91        "failed" => Ok(AgentStatus::Failed),
92        "max_cycles" => Ok(AgentStatus::MaxCycles),
93        other => Err(Error::new(
94            ErrorKind::InvalidData,
95            format!("unknown checkpoint status: {other}"),
96        )),
97    }
98}
99
100fn poisoned(name: &str) -> Error {
101    Error::other(format!("{name} lock is poisoned"))
102}
vv_agent/runtime/state.rs

vv_agent/runtime/
state.rs
