Skip to main content

atman_runtime/memory/
plan.rs

1use std::path::{Path, PathBuf};
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5
6use crate::error::RuntimeError;
7use crate::memory::{append_jsonl, read_jsonl};
8
9#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
10pub struct PlanStep {
11    pub index: usize,
12    pub text: String,
13    #[serde(default)]
14    pub done: bool,
15    #[serde(default, skip_serializing_if = "Option::is_none")]
16    pub done_at: Option<DateTime<Utc>>,
17}
18
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
20pub struct Plan {
21    pub id: String,
22    pub title: String,
23    pub steps: Vec<PlanStep>,
24    pub created_at: DateTime<Utc>,
25    pub updated_at: DateTime<Utc>,
26}
27
28impl Plan {
29    pub fn new(id: impl Into<String>, title: impl Into<String>, steps: Vec<String>) -> Self {
30        let now = Utc::now();
31        Self {
32            id: id.into(),
33            title: title.into(),
34            steps: steps
35                .into_iter()
36                .enumerate()
37                .map(|(i, text)| PlanStep {
38                    index: i,
39                    text,
40                    done: false,
41                    done_at: None,
42                })
43                .collect(),
44            created_at: now,
45            updated_at: now,
46        }
47    }
48
49    pub fn progress(&self) -> (usize, usize) {
50        let done = self.steps.iter().filter(|s| s.done).count();
51        (done, self.steps.len())
52    }
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
56#[serde(tag = "op", rename_all = "snake_case")]
57enum PlanEntry {
58    Upsert(Plan),
59    Tick {
60        plan_id: String,
61        step_index: usize,
62        at: DateTime<Utc>,
63    },
64}
65
66pub struct PlanStore {
67    path: PathBuf,
68    notify: Option<tokio::sync::watch::Sender<Vec<Plan>>>,
69}
70
71impl PlanStore {
72    pub fn at(session_dir: impl AsRef<Path>) -> Self {
73        Self {
74            path: session_dir.as_ref().join("plans.jsonl"),
75            notify: None,
76        }
77    }
78
79    pub fn with_notify(mut self, tx: tokio::sync::watch::Sender<Vec<Plan>>) -> Self {
80        self.notify = Some(tx);
81        self
82    }
83
84    async fn notify_if_needed(&self) {
85        if let Some(tx) = &self.notify {
86            if let Ok(list) = self.list().await {
87                let _ = tx.send(list);
88            }
89        }
90    }
91
92    pub async fn upsert(&self, plan: Plan) -> Result<(), RuntimeError> {
93        let mut plan = plan;
94        plan.updated_at = Utc::now();
95        append_jsonl(&self.path, &PlanEntry::Upsert(plan)).await?;
96        self.notify_if_needed().await;
97        Ok(())
98    }
99
100    pub async fn tick(&self, plan_id: &str, step_index: usize) -> Result<(), RuntimeError> {
101        append_jsonl(
102            &self.path,
103            &PlanEntry::Tick {
104                plan_id: plan_id.into(),
105                step_index,
106                at: Utc::now(),
107            },
108        )
109        .await?;
110        self.notify_if_needed().await;
111        Ok(())
112    }
113
114    pub async fn list(&self) -> Result<Vec<Plan>, RuntimeError> {
115        let entries: Vec<PlanEntry> = read_jsonl(&self.path).await?;
116        let mut plans: Vec<Plan> = Vec::new();
117        for entry in entries {
118            match entry {
119                PlanEntry::Upsert(mut p) => {
120                    if let Some(idx) = plans.iter().position(|x| x.id == p.id) {
121                        let old = &plans[idx];
122                        p.created_at = old.created_at;
123                        plans[idx] = p;
124                    } else {
125                        plans.push(p);
126                    }
127                }
128                PlanEntry::Tick {
129                    plan_id,
130                    step_index,
131                    at,
132                } => {
133                    if let Some(p) = plans.iter_mut().find(|p| p.id == plan_id)
134                        && let Some(step) = p.steps.iter_mut().find(|s| s.index == step_index)
135                    {
136                        step.done = true;
137                        step.done_at = Some(at);
138                        p.updated_at = at;
139                    }
140                }
141            }
142        }
143        Ok(plans)
144    }
145
146    pub async fn get(&self, id: &str) -> Result<Option<Plan>, RuntimeError> {
147        Ok(self.list().await?.into_iter().find(|p| p.id == id))
148    }
149
150    pub async fn latest(&self) -> Result<Option<Plan>, RuntimeError> {
151        Ok(self.list().await?.into_iter().max_by_key(|p| p.updated_at))
152    }
153
154    pub fn path(&self) -> &Path {
155        &self.path
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use tempfile::TempDir;
163
164    fn sample(id: &str, title: &str, steps: &[&str]) -> Plan {
165        Plan::new(id, title, steps.iter().map(|s| s.to_string()).collect())
166    }
167
168    #[tokio::test]
169    async fn upsert_then_get_round_trips() {
170        let dir = TempDir::new().unwrap();
171        let store = PlanStore::at(dir.path());
172        let plan = sample("p1", "ship endurance", &["design", "implement", "ship"]);
173        store.upsert(plan.clone()).await.unwrap();
174        let fetched = store.get("p1").await.unwrap().unwrap();
175        assert_eq!(fetched.title, "ship endurance");
176        assert_eq!(fetched.steps.len(), 3);
177    }
178
179    #[tokio::test]
180    async fn latest_returns_most_recent_by_updated_at() {
181        let dir = TempDir::new().unwrap();
182        let store = PlanStore::at(dir.path());
183        store.upsert(sample("p1", "first", &["a"])).await.unwrap();
184        tokio::time::sleep(std::time::Duration::from_millis(5)).await;
185        store.upsert(sample("p2", "second", &["b"])).await.unwrap();
186        assert_eq!(store.latest().await.unwrap().unwrap().id, "p2");
187    }
188
189    #[tokio::test]
190    async fn tick_flips_step_done() {
191        let dir = TempDir::new().unwrap();
192        let store = PlanStore::at(dir.path());
193        store
194            .upsert(sample("p1", "t", &["a", "b", "c"]))
195            .await
196            .unwrap();
197        store.tick("p1", 1).await.unwrap();
198        let plan = store.get("p1").await.unwrap().unwrap();
199        assert!(!plan.steps[0].done);
200        assert!(plan.steps[1].done);
201        assert!(!plan.steps[2].done);
202        assert_eq!(plan.progress(), (1, 3));
203    }
204
205    #[tokio::test]
206    async fn upsert_replaces_by_id_preserving_created_at() {
207        let dir = TempDir::new().unwrap();
208        let store = PlanStore::at(dir.path());
209        let mut plan = sample("p1", "v1", &["a"]);
210        store.upsert(plan.clone()).await.unwrap();
211        let original_created = store.get("p1").await.unwrap().unwrap().created_at;
212        plan.title = "v2".into();
213        plan.steps.push(PlanStep {
214            index: 1,
215            text: "b".into(),
216            done: false,
217            done_at: None,
218        });
219        tokio::time::sleep(std::time::Duration::from_millis(5)).await;
220        store.upsert(plan).await.unwrap();
221        let updated = store.get("p1").await.unwrap().unwrap();
222        assert_eq!(updated.title, "v2");
223        assert_eq!(updated.steps.len(), 2);
224        assert_eq!(updated.created_at, original_created);
225    }
226}