Skip to main content

codetether_agent/ralph/
store_memory.rs

1//! In-memory implementation of RalphStateStore
2//!
3//! Used for tests and as a fallback when no external store is configured.
4
5use async_trait::async_trait;
6use std::collections::HashMap;
7use tokio::sync::RwLock;
8
9use super::state_store::{RalphRunState, RalphRunSummary, RalphStateStore, StoryResultEntry};
10use super::types::{Prd, ProgressEntry, RalphStatus};
11
12/// In-memory state store backed by a RwLock<HashMap>
13pub struct InMemoryStore {
14    runs: RwLock<HashMap<String, RalphRunState>>,
15}
16
17impl InMemoryStore {
18    pub fn new() -> Self {
19        Self {
20            runs: RwLock::new(HashMap::new()),
21        }
22    }
23}
24
25impl Default for InMemoryStore {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31#[async_trait]
32impl RalphStateStore for InMemoryStore {
33    async fn create_run(&self, state: &RalphRunState) -> anyhow::Result<()> {
34        let mut runs = self.runs.write().await;
35        runs.insert(state.run_id.clone(), state.clone());
36        Ok(())
37    }
38
39    async fn update_status(&self, run_id: &str, status: RalphStatus) -> anyhow::Result<()> {
40        let mut runs = self.runs.write().await;
41        if let Some(run) = runs.get_mut(run_id) {
42            run.status = status;
43            if status == RalphStatus::Running && run.started_at.is_none() {
44                run.started_at = Some(chrono::Utc::now().to_rfc3339());
45            }
46        }
47        Ok(())
48    }
49
50    async fn update_iteration(&self, run_id: &str, iteration: usize) -> anyhow::Result<()> {
51        let mut runs = self.runs.write().await;
52        if let Some(run) = runs.get_mut(run_id) {
53            run.current_iteration = iteration;
54        }
55        Ok(())
56    }
57
58    async fn record_story_result(
59        &self,
60        run_id: &str,
61        result: &StoryResultEntry,
62    ) -> anyhow::Result<()> {
63        let mut runs = self.runs.write().await;
64        if let Some(run) = runs.get_mut(run_id) {
65            // Update or insert
66            if let Some(existing) = run
67                .story_results
68                .iter_mut()
69                .find(|r| r.story_id == result.story_id)
70            {
71                *existing = result.clone();
72            } else {
73                run.story_results.push(result.clone());
74            }
75        }
76        Ok(())
77    }
78
79    async fn append_progress(&self, run_id: &str, entry: &ProgressEntry) -> anyhow::Result<()> {
80        let mut runs = self.runs.write().await;
81        if let Some(run) = runs.get_mut(run_id) {
82            run.progress_log.push(entry.clone());
83        }
84        Ok(())
85    }
86
87    async fn update_prd(&self, run_id: &str, prd: &Prd) -> anyhow::Result<()> {
88        let mut runs = self.runs.write().await;
89        if let Some(run) = runs.get_mut(run_id) {
90            run.prd = prd.clone();
91        }
92        Ok(())
93    }
94
95    async fn set_error(&self, run_id: &str, error: &str) -> anyhow::Result<()> {
96        let mut runs = self.runs.write().await;
97        if let Some(run) = runs.get_mut(run_id) {
98            run.error = Some(error.to_string());
99        }
100        Ok(())
101    }
102
103    async fn complete_run(&self, run_id: &str, status: RalphStatus) -> anyhow::Result<()> {
104        let mut runs = self.runs.write().await;
105        if let Some(run) = runs.get_mut(run_id) {
106            run.status = status;
107            run.completed_at = Some(chrono::Utc::now().to_rfc3339());
108        }
109        Ok(())
110    }
111
112    async fn get_run(&self, run_id: &str) -> anyhow::Result<Option<RalphRunState>> {
113        let runs = self.runs.read().await;
114        Ok(runs.get(run_id).cloned())
115    }
116
117    async fn list_runs(&self) -> anyhow::Result<Vec<RalphRunSummary>> {
118        let runs = self.runs.read().await;
119        Ok(runs
120            .values()
121            .map(|run| RalphRunSummary {
122                run_id: run.run_id.clone(),
123                okr_id: run.okr_id.clone(),
124                status: run.status,
125                passed: run.prd.passed_count(),
126                total: run.prd.user_stories.len(),
127                current_iteration: run.current_iteration,
128                created_at: run.created_at.clone(),
129            })
130            .collect())
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use crate::ralph::types::{Prd, QualityChecks, RalphConfig, UserStory};
138
139    fn test_prd() -> Prd {
140        Prd {
141            project: "test".to_string(),
142            feature: "test-feature".to_string(),
143            branch_name: "feature/test".to_string(),
144            version: "1.0".to_string(),
145            user_stories: vec![UserStory {
146                id: "US-001".to_string(),
147                title: "Test story".to_string(),
148                description: "A test story".to_string(),
149                acceptance_criteria: vec![],
150                verification_steps: vec![],
151                passes: false,
152                priority: 1,
153                depends_on: vec![],
154                complexity: 2,
155            }],
156            technical_requirements: vec![],
157            quality_checks: QualityChecks::default(),
158            created_at: String::new(),
159            updated_at: String::new(),
160        }
161    }
162
163    fn test_run_state() -> RalphRunState {
164        RalphRunState {
165            run_id: "test-run-1".to_string(),
166            okr_id: Some("okr-1".to_string()),
167            prd: test_prd(),
168            config: RalphConfig::default(),
169            status: RalphStatus::Pending,
170            current_iteration: 0,
171            max_iterations: 10,
172            progress_log: vec![],
173            story_results: vec![],
174            error: None,
175            created_at: chrono::Utc::now().to_rfc3339(),
176            started_at: None,
177            completed_at: None,
178        }
179    }
180
181    #[tokio::test]
182    async fn test_create_and_get_run() {
183        let store = InMemoryStore::new();
184        let state = test_run_state();
185
186        store.create_run(&state).await.unwrap();
187
188        let fetched = store.get_run("test-run-1").await.unwrap();
189        assert!(fetched.is_some());
190        let fetched = fetched.unwrap();
191        assert_eq!(fetched.run_id, "test-run-1");
192        assert_eq!(fetched.status, RalphStatus::Pending);
193    }
194
195    #[tokio::test]
196    async fn test_update_status() {
197        let store = InMemoryStore::new();
198        store.create_run(&test_run_state()).await.unwrap();
199
200        store
201            .update_status("test-run-1", RalphStatus::Running)
202            .await
203            .unwrap();
204
205        let run = store.get_run("test-run-1").await.unwrap().unwrap();
206        assert_eq!(run.status, RalphStatus::Running);
207        assert!(run.started_at.is_some());
208    }
209
210    #[tokio::test]
211    async fn test_record_story_result() {
212        let store = InMemoryStore::new();
213        store.create_run(&test_run_state()).await.unwrap();
214
215        let result = StoryResultEntry {
216            story_id: "US-001".to_string(),
217            title: "Test story".to_string(),
218            passed: true,
219            iteration: 1,
220            error: None,
221        };
222
223        store
224            .record_story_result("test-run-1", &result)
225            .await
226            .unwrap();
227
228        let run = store.get_run("test-run-1").await.unwrap().unwrap();
229        assert_eq!(run.story_results.len(), 1);
230        assert!(run.story_results[0].passed);
231    }
232
233    #[tokio::test]
234    async fn test_complete_run() {
235        let store = InMemoryStore::new();
236        store.create_run(&test_run_state()).await.unwrap();
237
238        store
239            .complete_run("test-run-1", RalphStatus::Completed)
240            .await
241            .unwrap();
242
243        let run = store.get_run("test-run-1").await.unwrap().unwrap();
244        assert_eq!(run.status, RalphStatus::Completed);
245        assert!(run.completed_at.is_some());
246    }
247
248    #[tokio::test]
249    async fn test_list_runs() {
250        let store = InMemoryStore::new();
251        store.create_run(&test_run_state()).await.unwrap();
252
253        let summaries = store.list_runs().await.unwrap();
254        assert_eq!(summaries.len(), 1);
255        assert_eq!(summaries[0].run_id, "test-run-1");
256        assert_eq!(summaries[0].total, 1);
257    }
258
259    #[tokio::test]
260    async fn test_get_nonexistent_run() {
261        let store = InMemoryStore::new();
262        let result = store.get_run("nonexistent").await.unwrap();
263        assert!(result.is_none());
264    }
265}