Skip to main content

codetether_agent/okr/
persistence.rs

1//! OKR Persistence Layer
2//!
3//! Provides file-based persistence for OKR entities using Config::data_dir().
4//! Supports CRUD operations for ID, run, and status queries.
5
6use crate::okr::{Okr, OkrRun, OkrRunStatus, OkrStatus};
7use anyhow::{Context, Result};
8use std::path::PathBuf;
9use std::sync::Arc;
10use tokio::fs;
11use tokio::sync::RwLock;
12use uuid::Uuid;
13
14/// Repository for OKR persistence operations
15pub struct OkrRepository {
16    base_path: PathBuf,
17    cache: Arc<RwLock<OkrCache>>,
18}
19
20/// In-memory cache for OKR data
21#[derive(Debug, Default)]
22struct OkrCache {
23    okrs: Vec<Okr>,
24    runs: Vec<OkrRun>,
25}
26
27impl OkrRepository {
28    /// Create a new OKR repository with the given base path
29    pub fn new(base_path: PathBuf) -> Self {
30        Self {
31            base_path,
32            cache: Arc::new(RwLock::new(OkrCache::default())),
33        }
34    }
35
36    /// Create repository from Config::data_dir() or fallback to temp
37    pub async fn from_config() -> Result<Self> {
38        let base_path = crate::config::Config::data_dir()
39            .map(|p| p.join("okr"))
40            .unwrap_or_else(|| std::env::temp_dir().join("codetether").join("okr"));
41
42        // Ensure directory exists
43        fs::create_dir_all(&base_path)
44            .await
45            .context("Failed to create OKR data directory")?;
46
47        tracing::info!(path = %base_path.display(), "OKR repository initialized");
48
49        Ok(Self::new(base_path))
50    }
51
52    /// Get the file path for an OKR
53    fn okr_path(&self, id: Uuid) -> PathBuf {
54        self.base_path
55            .join("objectives")
56            .join(format!("{}.json", id))
57    }
58
59    /// Get the file path for an OKR run
60    fn run_path(&self, id: Uuid) -> PathBuf {
61        self.base_path.join("runs").join(format!("{}.json", id))
62    }
63
64    /// Get the directory path for OKRs
65    fn okrs_dir(&self) -> PathBuf {
66        self.base_path.join("objectives")
67    }
68
69    /// Get the directory path for runs
70    fn runs_dir(&self) -> PathBuf {
71        self.base_path.join("runs")
72    }
73
74    // ============ OKR CRUD ============
75
76    /// Create a new OKR
77    pub async fn create_okr(&self, okr: Okr) -> Result<Okr> {
78        // Validate
79        okr.validate()?;
80
81        // Ensure directory exists
82        fs::create_dir_all(self.okrs_dir()).await?;
83
84        // Write to file (atomic: tmp + rename)
85        let path = self.okr_path(okr.id);
86        let tmp_path = path.with_extension("json.tmp");
87        let json = serde_json::to_string_pretty(&okr)?;
88        fs::write(&tmp_path, &json).await?;
89        fs::rename(&tmp_path, &path).await?;
90
91        // Update cache
92        let mut cache = self.cache.write().await;
93        cache.okrs.push(okr.clone());
94
95        tracing::info!(okr_id = %okr.id, title = %okr.title, "OKR created");
96        Ok(okr)
97    }
98
99    /// Get an OKR by ID
100    pub async fn get_okr(&self, id: Uuid) -> Result<Option<Okr>> {
101        // Check cache first
102        {
103            let cache = self.cache.read().await;
104            if let Some(okr) = cache.okrs.iter().find(|o| o.id == id) {
105                return Ok(Some(okr.clone()));
106            }
107        }
108
109        // Load from disk
110        let path = self.okr_path(id);
111        if !path.exists() {
112            return Ok(None);
113        }
114
115        let data = fs::read_to_string(&path).await?;
116        let okr: Okr = serde_json::from_str(&data).context("Failed to parse OKR")?;
117
118        // Update cache
119        let mut cache = self.cache.write().await;
120        if !cache.okrs.iter().any(|o| o.id == okr.id) {
121            cache.okrs.push(okr.clone());
122        }
123
124        Ok(Some(okr))
125    }
126
127    /// Update an existing OKR
128    pub async fn update_okr(&self, mut okr: Okr) -> Result<Okr> {
129        okr.updated_at = chrono::Utc::now();
130
131        // Validate
132        okr.validate()?;
133
134        // Write to file
135        let path = self.okr_path(okr.id);
136        fs::create_dir_all(self.okrs_dir()).await?;
137        let json = serde_json::to_string_pretty(&okr)?;
138        fs::write(&path, &json).await?;
139
140        // Update cache
141        let mut cache = self.cache.write().await;
142        if let Some(existing) = cache.okrs.iter_mut().find(|o| o.id == okr.id) {
143            *existing = okr.clone();
144        }
145
146        tracing::info!(okr_id = %okr.id, "OKR updated");
147        Ok(okr)
148    }
149
150    /// Delete an OKR by ID
151    pub async fn delete_okr(&self, id: Uuid) -> Result<bool> {
152        let path = self.okr_path(id);
153
154        if !path.exists() {
155            return Ok(false);
156        }
157
158        fs::remove_file(&path).await?;
159
160        // Remove from cache
161        let mut cache = self.cache.write().await;
162        cache.okrs.retain(|o| o.id != id);
163
164        tracing::info!(okr_id = %id, "OKR deleted");
165        Ok(true)
166    }
167
168    /// List all OKRs
169    pub async fn list_okrs(&self) -> Result<Vec<Okr>> {
170        // Check cache first
171        {
172            let cache = self.cache.read().await;
173            if !cache.okrs.is_empty() {
174                return Ok(cache.okrs.clone());
175            }
176        }
177
178        // Load from disk
179        let dir = self.okrs_dir();
180        if !dir.exists() {
181            return Ok(Vec::new());
182        }
183
184        let mut okrs = Vec::new();
185        let mut entries = fs::read_dir(&dir).await?;
186
187        while let Some(entry) = entries.next_entry().await? {
188            let path = entry.path();
189            if path.extension().is_some_and(|e| e == "json")
190                && let Ok(data) = fs::read_to_string(&path).await
191                && let Ok(okr) = serde_json::from_str::<Okr>(&data)
192            {
193                okrs.push(okr);
194            }
195        }
196
197        // Update cache
198        let mut cache = self.cache.write().await;
199        cache.okrs = okrs.clone();
200
201        Ok(okrs)
202    }
203
204    /// Query OKRs by status
205    pub async fn query_okrs_by_status(&self, status: OkrStatus) -> Result<Vec<Okr>> {
206        let all = self.list_okrs().await?;
207        Ok(all.into_iter().filter(|o| o.status == status).collect())
208    }
209
210    /// Query OKRs by owner
211    pub async fn query_okrs_by_owner(&self, owner: &str) -> Result<Vec<Okr>> {
212        let all = self.list_okrs().await?;
213        Ok(all
214            .into_iter()
215            .filter(|o| o.owner.as_deref() == Some(owner))
216            .collect())
217    }
218
219    /// Query OKRs by tenant
220    pub async fn query_okrs_by_tenant(&self, tenant_id: &str) -> Result<Vec<Okr>> {
221        let all = self.list_okrs().await?;
222        Ok(all
223            .into_iter()
224            .filter(|o| o.tenant_id.as_deref() == Some(tenant_id))
225            .collect())
226    }
227
228    // ============ OKR Run CRUD ============
229
230    /// Create a new OKR run
231    pub async fn create_run(&self, run: OkrRun) -> Result<OkrRun> {
232        // Validate
233        run.validate()?;
234
235        // Ensure directory exists
236        fs::create_dir_all(self.runs_dir()).await?;
237
238        // Write to file
239        let path = self.run_path(run.id);
240        let json = serde_json::to_string_pretty(&run)?;
241        fs::write(&path, &json).await?;
242
243        // Update cache
244        let mut cache = self.cache.write().await;
245        cache.runs.push(run.clone());
246
247        tracing::info!(run_id = %run.id, okr_id = %run.okr_id, name = %run.name, "OKR run created");
248        Ok(run)
249    }
250
251    /// Get a run by ID
252    pub async fn get_run(&self, id: Uuid) -> Result<Option<OkrRun>> {
253        // Check cache
254        {
255            let cache = self.cache.read().await;
256            if let Some(run) = cache.runs.iter().find(|r| r.id == id) {
257                return Ok(Some(run.clone()));
258            }
259        }
260
261        // Load from disk
262        let path = self.run_path(id);
263        if !path.exists() {
264            return Ok(None);
265        }
266
267        let data = fs::read_to_string(&path).await?;
268        let run: OkrRun = serde_json::from_str(&data).context("Failed to parse OKR run")?;
269
270        // Update cache
271        let mut cache = self.cache.write().await;
272        if !cache.runs.iter().any(|r| r.id == run.id) {
273            cache.runs.push(run.clone());
274        }
275
276        Ok(Some(run))
277    }
278
279    /// Update a run
280    pub async fn update_run(&self, mut run: OkrRun) -> Result<OkrRun> {
281        run.updated_at = chrono::Utc::now();
282
283        // Write to file
284        let path = self.run_path(run.id);
285        fs::create_dir_all(self.runs_dir()).await?;
286        let json = serde_json::to_string_pretty(&run)?;
287        fs::write(&path, &json).await?;
288
289        // Update cache
290        let mut cache = self.cache.write().await;
291        if let Some(existing) = cache.runs.iter_mut().find(|r| r.id == run.id) {
292            *existing = run.clone();
293        }
294
295        tracing::info!(run_id = %run.id, "OKR run updated");
296        Ok(run)
297    }
298
299    /// Delete a run
300    pub async fn delete_run(&self, id: Uuid) -> Result<bool> {
301        let path = self.run_path(id);
302
303        if !path.exists() {
304            return Ok(false);
305        }
306
307        fs::remove_file(&path).await?;
308
309        // Remove from cache
310        let mut cache = self.cache.write().await;
311        cache.runs.retain(|r| r.id != id);
312
313        tracing::info!(run_id = %id, "OKR run deleted");
314        Ok(true)
315    }
316
317    /// List all runs
318    pub async fn list_runs(&self) -> Result<Vec<OkrRun>> {
319        // Check cache
320        {
321            let cache = self.cache.read().await;
322            if !cache.runs.is_empty() {
323                return Ok(cache.runs.clone());
324            }
325        }
326
327        // Load from disk
328        let dir = self.runs_dir();
329        if !dir.exists() {
330            return Ok(Vec::new());
331        }
332
333        let mut runs = Vec::new();
334        let mut entries = fs::read_dir(&dir).await?;
335
336        while let Some(entry) = entries.next_entry().await? {
337            let path = entry.path();
338            if path.extension().is_some_and(|e| e == "json")
339                && let Ok(data) = fs::read_to_string(&path).await
340                && let Ok(run) = serde_json::from_str::<OkrRun>(&data)
341            {
342                runs.push(run);
343            }
344        }
345
346        // Update cache
347        let mut cache = self.cache.write().await;
348        cache.runs = runs.clone();
349
350        Ok(runs)
351    }
352
353    /// Query runs by OKR ID
354    pub async fn query_runs_by_okr(&self, okr_id: Uuid) -> Result<Vec<OkrRun>> {
355        let all = self.list_runs().await?;
356        Ok(all.into_iter().filter(|r| r.okr_id == okr_id).collect())
357    }
358
359    /// Query runs by status
360    pub async fn query_runs_by_status(&self, status: OkrRunStatus) -> Result<Vec<OkrRun>> {
361        let all = self.list_runs().await?;
362        Ok(all.into_iter().filter(|r| r.status == status).collect())
363    }
364
365    /// Query runs by correlation ID
366    pub async fn query_runs_by_correlation(&self, correlation_id: &str) -> Result<Vec<OkrRun>> {
367        let all = self.list_runs().await?;
368        Ok(all
369            .into_iter()
370            .filter(|r| r.correlation_id.as_deref() == Some(correlation_id))
371            .collect())
372    }
373
374    /// Query runs by relay checkpoint ID
375    pub async fn query_runs_by_checkpoint(&self, checkpoint_id: &str) -> Result<Vec<OkrRun>> {
376        let all = self.list_runs().await?;
377        Ok(all
378            .into_iter()
379            .filter(|r| r.relay_checkpoint_id.as_deref() == Some(checkpoint_id))
380            .collect())
381    }
382
383    /// Query runs by session ID
384    pub async fn query_runs_by_session(&self, session_id: &str) -> Result<Vec<OkrRun>> {
385        let all = self.list_runs().await?;
386        Ok(all
387            .into_iter()
388            .filter(|r| r.session_id.as_deref() == Some(session_id))
389            .collect())
390    }
391
392    // ============ Utility Methods ============
393
394    /// Clear the in-memory cache (useful for forcing reload from disk)
395    #[allow(dead_code)]
396    pub async fn clear_cache(&self) {
397        let mut cache = self.cache.write().await;
398        cache.okrs.clear();
399        cache.runs.clear();
400        tracing::debug!("OKR repository cache cleared");
401    }
402
403    /// Get statistics about the repository
404    pub async fn stats(&self) -> Result<OkrRepositoryStats> {
405        let okrs = self.list_okrs().await?;
406        let runs = self.list_runs().await?;
407
408        let okr_status_counts = okrs
409            .iter()
410            .fold(std::collections::HashMap::new(), |mut acc, o| {
411                *acc.entry(o.status).or_insert(0) += 1;
412                acc
413            });
414
415        let run_status_counts = runs
416            .iter()
417            .fold(std::collections::HashMap::new(), |mut acc, r| {
418                *acc.entry(r.status).or_insert(0) += 1;
419                acc
420            });
421
422        Ok(OkrRepositoryStats {
423            total_okrs: okrs.len(),
424            total_runs: runs.len(),
425            okr_status_counts,
426            run_status_counts,
427        })
428    }
429}
430
431/// Statistics about the OKR repository
432#[derive(Debug, serde::Serialize)]
433pub struct OkrRepositoryStats {
434    pub total_okrs: usize,
435    pub total_runs: usize,
436    pub okr_status_counts: std::collections::HashMap<OkrStatus, usize>,
437    pub run_status_counts: std::collections::HashMap<OkrRunStatus, usize>,
438}
439
440#[cfg(test)]
441mod tests {
442    use super::*;
443    use crate::okr::KeyResult;
444
445    fn temp_repo() -> OkrRepository {
446        let temp_dir = std::env::temp_dir().join(format!("okr_test_{}", Uuid::new_v4()));
447        OkrRepository::new(temp_dir)
448    }
449
450    #[tokio::test]
451    async fn test_create_and_get_okr() {
452        let repo = temp_repo();
453
454        let mut okr = Okr::new("Test Objective", "Description");
455        let kr = KeyResult::new(okr.id, "KR1", 100.0, "%");
456        okr.add_key_result(kr);
457
458        let created = repo.create_okr(okr).await.unwrap();
459
460        let fetched = repo.get_okr(created.id).await.unwrap();
461        assert!(fetched.is_some());
462        assert_eq!(fetched.unwrap().title, "Test Objective");
463    }
464
465    #[tokio::test]
466    async fn test_update_okr() {
467        let repo = temp_repo();
468
469        let mut okr = Okr::new("Test", "Desc");
470        let kr = KeyResult::new(okr.id, "KR1", 100.0, "%");
471        okr.add_key_result(kr);
472
473        let created = repo.create_okr(okr).await.unwrap();
474
475        let mut to_update = created.clone();
476        to_update.title = "Updated Title".to_string();
477        to_update.status = OkrStatus::Active;
478
479        let updated = repo.update_okr(to_update).await.unwrap();
480        assert_eq!(updated.title, "Updated Title");
481        assert_eq!(updated.status, OkrStatus::Active);
482    }
483
484    #[tokio::test]
485    async fn test_delete_okr() {
486        let repo = temp_repo();
487
488        let mut okr = Okr::new("Test", "Desc");
489        let kr = KeyResult::new(okr.id, "KR1", 100.0, "%");
490        okr.add_key_result(kr);
491
492        let created = repo.create_okr(okr).await.unwrap();
493        let deleted = repo.delete_okr(created.id).await.unwrap();
494        assert!(deleted);
495
496        let fetched = repo.get_okr(created.id).await.unwrap();
497        assert!(fetched.is_none());
498    }
499
500    #[tokio::test]
501    async fn test_query_okrs_by_status() {
502        let repo = temp_repo();
503
504        let mut okr1 = Okr::new("Active OKR", "Desc");
505        let kr1 = KeyResult::new(okr1.id, "KR1", 100.0, "%");
506        okr1.add_key_result(kr1);
507        okr1.status = OkrStatus::Active;
508
509        let mut okr2 = Okr::new("Draft OKR", "Desc");
510        let kr2 = KeyResult::new(okr2.id, "KR2", 100.0, "%");
511        okr2.add_key_result(kr2);
512        okr2.status = OkrStatus::Draft;
513
514        repo.create_okr(okr1).await.unwrap();
515        repo.create_okr(okr2).await.unwrap();
516
517        let active = repo.query_okrs_by_status(OkrStatus::Active).await.unwrap();
518        assert_eq!(active.len(), 1);
519        assert_eq!(active[0].title, "Active OKR");
520    }
521
522    #[tokio::test]
523    async fn test_crud_runs() {
524        let repo = temp_repo();
525
526        // First create an OKR
527        let mut okr = Okr::new("Test OKR", "Desc");
528        let okr_id = okr.id;
529        let kr = KeyResult::new(okr.id, "KR1", 100.0, "%");
530        okr.add_key_result(kr);
531        repo.create_okr(okr).await.unwrap();
532
533        // Create a run
534        let mut run = OkrRun::new(okr_id, "Q1 Run");
535        run.correlation_id = Some("corr-123".to_string());
536        run.session_id = Some("session-456".to_string());
537
538        let created = repo.create_run(run).await.unwrap();
539
540        // Fetch by ID
541        let fetched = repo.get_run(created.id).await.unwrap();
542        assert!(fetched.is_some());
543        assert_eq!(fetched.unwrap().name, "Q1 Run");
544
545        // Query by OKR
546        let runs_for_okr = repo.query_runs_by_okr(okr_id).await.unwrap();
547        assert_eq!(runs_for_okr.len(), 1);
548
549        // Query by correlation
550        let runs_by_corr = repo.query_runs_by_correlation("corr-123").await.unwrap();
551        assert_eq!(runs_by_corr.len(), 1);
552
553        // Query by status
554        let pending = repo
555            .query_runs_by_status(OkrRunStatus::Draft)
556            .await
557            .unwrap();
558        assert_eq!(pending.len(), 1);
559    }
560
561    #[tokio::test]
562    async fn test_list_and_stats() {
563        let repo = temp_repo();
564
565        // Create OKRs
566        let mut okr = Okr::new("Test", "Desc");
567        let kr = KeyResult::new(okr.id, "KR", 100.0, "%");
568        okr.add_key_result(kr);
569        repo.create_okr(okr).await.unwrap();
570
571        // List
572        let all_okrs = repo.list_okrs().await.unwrap();
573        assert_eq!(all_okrs.len(), 1);
574
575        let all_runs = repo.list_runs().await.unwrap();
576        assert_eq!(all_runs.len(), 0);
577
578        // Stats
579        let stats = repo.stats().await.unwrap();
580        assert_eq!(stats.total_okrs, 1);
581    }
582}