ninox_core/
store.rs

1use crate::types::*;
2use anyhow::Result;
3use rusqlite::{params, Connection};
4use std::{path::Path, sync::Mutex};
5
6pub struct Store {
7    conn: Mutex<Connection>,
8}
9
10
11impl Store {
12    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
13        let conn = Connection::open(path)?;
14        conn.execute_batch("
15            PRAGMA journal_mode=WAL;
16            PRAGMA busy_timeout=5000;
17            CREATE TABLE IF NOT EXISTS sessions (
18                id TEXT PRIMARY KEY, orchestrator_id TEXT,
19                name TEXT NOT NULL, repo TEXT NOT NULL,
20                status TEXT NOT NULL, agent_type TEXT NOT NULL,
21                cost_usd REAL NOT NULL DEFAULT 0, started_at INTEGER NOT NULL,
22                pr_number INTEGER, pr_id INTEGER,
23                workspace_path TEXT, pid INTEGER,
24                model TEXT, context_tokens INTEGER
25            );
26            CREATE TABLE IF NOT EXISTS orchestrators (
27                id TEXT PRIMARY KEY, name TEXT NOT NULL, created_at INTEGER NOT NULL
28            );
29            CREATE TABLE IF NOT EXISTS prs (
30                id INTEGER PRIMARY KEY, number INTEGER NOT NULL,
31                title TEXT NOT NULL, url TEXT NOT NULL,
32                body TEXT NOT NULL, session_id TEXT NOT NULL
33            );
34            CREATE TABLE IF NOT EXISTS ci_status (
35                pr_id INTEGER PRIMARY KEY, total INTEGER NOT NULL,
36                passing INTEGER NOT NULL, failing INTEGER NOT NULL,
37                pending INTEGER NOT NULL
38            );
39            CREATE TABLE IF NOT EXISTS review_comments (
40                id INTEGER PRIMARY KEY, pr_id INTEGER NOT NULL,
41                author TEXT NOT NULL, body TEXT NOT NULL,
42                path TEXT, line INTEGER, created_at INTEGER NOT NULL
43            );
44        ")?;
45        // Migrations for columns added after initial release — idempotent so
46        // both fresh and pre-existing databases end up with the same schema.
47        for (col, ddl) in [
48            ("model",                "ALTER TABLE sessions ADD COLUMN model TEXT"),
49            ("context_tokens",       "ALTER TABLE sessions ADD COLUMN context_tokens INTEGER"),
50            ("catalogue_path",       "ALTER TABLE sessions ADD COLUMN catalogue_path TEXT"),
51            ("context_used_pct",     "ALTER TABLE sessions ADD COLUMN context_used_pct REAL"),
52            ("context_total_tokens", "ALTER TABLE sessions ADD COLUMN context_total_tokens INTEGER"),
53            ("context_window_size",  "ALTER TABLE sessions ADD COLUMN context_window_size INTEGER"),
54            ("claude_session_id",    "ALTER TABLE sessions ADD COLUMN claude_session_id TEXT"),
55            ("summary",              "ALTER TABLE sessions ADD COLUMN summary TEXT"),
56            ("terminal_at",          "ALTER TABLE sessions ADD COLUMN terminal_at INTEGER"),
57        ] {
58            if !Self::column_exists(&conn, "sessions", col)? {
59                conn.execute(ddl, [])?;
60            }
61        }
62        Ok(Self { conn: Mutex::new(conn) })
63    }
64
65    fn column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
66        let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
67        let exists = stmt
68            .query_map([], |r| r.get::<_, String>(1))?
69            .filter_map(|r| r.ok())
70            .any(|c| c == column);
71        Ok(exists)
72    }
73
74    pub fn upsert_session(&self, s: &Session) -> Result<()> {
75        let status = serde_json::to_string(&s.status)?.replace('"', "");
76        let conn = self.conn.lock().unwrap();
77        conn.execute(
78            "INSERT INTO sessions (id,orchestrator_id,name,repo,status,agent_type,
79             cost_usd,started_at,pr_number,pr_id,workspace_path,pid,model,context_tokens,
80             catalogue_path,context_used_pct,context_total_tokens,context_window_size,
81             claude_session_id,summary,terminal_at)
82             VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,?15,?16,?17,?18,?19,?20,?21)
83             ON CONFLICT(id) DO UPDATE SET
84             repo=excluded.repo,
85             status=excluded.status,cost_usd=excluded.cost_usd,
86             started_at=excluded.started_at,
87             pr_number=excluded.pr_number,pr_id=excluded.pr_id,
88             workspace_path=excluded.workspace_path,pid=excluded.pid,
89             model=excluded.model,context_tokens=excluded.context_tokens,
90             catalogue_path=excluded.catalogue_path,
91             context_used_pct=excluded.context_used_pct,
92             context_total_tokens=excluded.context_total_tokens,
93             context_window_size=excluded.context_window_size,
94             claude_session_id=excluded.claude_session_id,
95             summary=excluded.summary,
96             terminal_at=excluded.terminal_at",
97            params![
98                s.id, s.orchestrator_id, s.name, s.repo, status, s.agent_type,
99                s.cost_usd, s.started_at, s.pr_number, s.pr_id,
100                s.workspace_path, s.pid, s.model, s.context_tokens,
101                s.catalogue_path, s.context_used_pct, s.context_total_tokens,
102                s.context_window_size, s.claude_session_id, s.summary, s.terminal_at
103            ],
104        )?;
105        Ok(())
106    }
107
108    pub fn list_sessions(&self) -> Result<Vec<Session>> {
109        let conn = self.conn.lock().unwrap();
110        let mut stmt = conn.prepare(
111            "SELECT id,orchestrator_id,name,repo,status,agent_type,cost_usd,
112             started_at,pr_number,pr_id,workspace_path,pid,model,context_tokens,
113             catalogue_path,context_used_pct,context_total_tokens,context_window_size,
114             claude_session_id,summary,terminal_at
115             FROM sessions ORDER BY started_at DESC",
116        )?;
117        let rows = stmt.query_map([], |r| {
118            Ok((
119                r.get::<_, String>(0)?,
120                r.get::<_, Option<String>>(1)?,
121                r.get::<_, String>(2)?,
122                r.get::<_, String>(3)?,
123                r.get::<_, String>(4)?,
124                r.get::<_, String>(5)?,
125                r.get::<_, f64>(6)?,
126                r.get::<_, i64>(7)?,
127                r.get::<_, Option<u64>>(8)?,
128                r.get::<_, Option<i64>>(9)?,
129                r.get::<_, Option<String>>(10)?,
130                r.get::<_, Option<u32>>(11)?,
131                r.get::<_, Option<String>>(12)?,
132                r.get::<_, Option<i64>>(13)?,
133                r.get::<_, Option<String>>(14)?,
134                r.get::<_, Option<f64>>(15)?,
135                r.get::<_, Option<i64>>(16)?,
136                r.get::<_, Option<i64>>(17)?,
137                r.get::<_, Option<String>>(18)?,
138                r.get::<_, Option<String>>(19)?,
139                r.get::<_, Option<i64>>(20)?,
140            ))
141        })?;
142        rows.map(|r| {
143            let (id, orchestrator_id, name, repo, status_str, agent_type,
144                 cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
145                 model, context_tokens, catalogue_path, context_used_pct,
146                 context_total_tokens, context_window_size, claude_session_id,
147                 summary, terminal_at) = r?;
148            let status = serde_json::from_str(&format!("\"{status_str}\""))
149                .unwrap_or(SessionStatus::Working);
150            Ok(Session {
151                id, orchestrator_id, name, repo, status, agent_type,
152                cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
153                model, context_tokens: context_tokens.map(|v| v.max(0) as u64),
154                catalogue_path,
155                context_used_pct,
156                context_total_tokens: context_total_tokens.map(|v| v.max(0) as u64),
157                context_window_size: context_window_size.map(|v| v.max(0) as u64),
158                claude_session_id,
159                summary,
160                terminal_at,
161            })
162        })
163        .collect()
164    }
165
166    pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
167        let conn = self.conn.lock().unwrap();
168        let mut stmt = conn.prepare(
169            "SELECT id,orchestrator_id,name,repo,status,agent_type,cost_usd,
170             started_at,pr_number,pr_id,workspace_path,pid,model,context_tokens,
171             catalogue_path,context_used_pct,context_total_tokens,context_window_size,
172             claude_session_id,summary,terminal_at
173             FROM sessions WHERE id = ?1",
174        )?;
175        let mut rows = stmt.query_map([id], |r| {
176            Ok((
177                r.get::<_, String>(0)?,
178                r.get::<_, Option<String>>(1)?,
179                r.get::<_, String>(2)?,
180                r.get::<_, String>(3)?,
181                r.get::<_, String>(4)?,
182                r.get::<_, String>(5)?,
183                r.get::<_, f64>(6)?,
184                r.get::<_, i64>(7)?,
185                r.get::<_, Option<u64>>(8)?,
186                r.get::<_, Option<i64>>(9)?,
187                r.get::<_, Option<String>>(10)?,
188                r.get::<_, Option<u32>>(11)?,
189                r.get::<_, Option<String>>(12)?,
190                r.get::<_, Option<i64>>(13)?,
191                r.get::<_, Option<String>>(14)?,
192                r.get::<_, Option<f64>>(15)?,
193                r.get::<_, Option<i64>>(16)?,
194                r.get::<_, Option<i64>>(17)?,
195                r.get::<_, Option<String>>(18)?,
196                r.get::<_, Option<String>>(19)?,
197                r.get::<_, Option<i64>>(20)?,
198            ))
199        })?;
200        match rows.next() {
201            None => Ok(None),
202            Some(r) => {
203                let (id, orchestrator_id, name, repo, status_str, agent_type,
204                     cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
205                     model, context_tokens, catalogue_path, context_used_pct,
206                     context_total_tokens, context_window_size, claude_session_id,
207                     summary, terminal_at) = r?;
208                let status = serde_json::from_str(&format!("\"{status_str}\""))
209                    .unwrap_or(SessionStatus::Working);
210                Ok(Some(Session {
211                    id, orchestrator_id, name, repo, status, agent_type,
212                    cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
213                    model, context_tokens: context_tokens.map(|v| v.max(0) as u64),
214                    catalogue_path,
215                    context_used_pct,
216                    context_total_tokens: context_total_tokens.map(|v| v.max(0) as u64),
217                    context_window_size: context_window_size.map(|v| v.max(0) as u64),
218                    claude_session_id,
219                    summary,
220                    terminal_at,
221                }))
222            }
223        }
224    }
225
226    /// Non-zero `cost_usd` samples recorded for sessions matching the given
227    /// agent harness (`agent_type`) and model — used to compute a
228    /// data-driven spawn-modal cost estimate once enough history exists for
229    /// a given preset. Read-only; built on `list_sessions` like
230    /// `sessions_by_orchestrator`.
231    pub fn cost_samples(&self, agent_type: &str, model: Option<&str>) -> Result<Vec<f64>> {
232        let sessions = self.list_sessions()?;
233        Ok(sessions
234            .into_iter()
235            .filter(|s| {
236                s.agent_type == agent_type
237                    && s.model.as_deref() == model
238                    && s.cost_usd > 0.0
239            })
240            .map(|s| s.cost_usd)
241            .collect())
242    }
243
244    pub fn upsert_orchestrator(&self, o: &Orchestrator) -> Result<()> {
245        let conn = self.conn.lock().unwrap();
246        conn.execute(
247            "INSERT INTO orchestrators(id,name,created_at) VALUES(?1,?2,?3)
248             ON CONFLICT(id) DO UPDATE SET name=excluded.name",
249            params![o.id, o.name, o.created_at],
250        )?;
251        Ok(())
252    }
253
254    pub fn sessions_by_orchestrator(&self, orchestrator_id: &str) -> Result<Vec<Session>> {
255        let sessions = self.list_sessions()?;
256        Ok(sessions.into_iter().filter(|s| s.orchestrator_id.as_deref() == Some(orchestrator_id)).collect())
257    }
258
259    pub fn delete_session(&self, id: &str) -> Result<()> {
260        let conn = self.conn.lock().unwrap();
261        conn.execute("DELETE FROM sessions WHERE id = ?1", [id])?;
262        Ok(())
263    }
264
265    pub fn delete_orchestrator(&self, id: &str) -> Result<()> {
266        let conn = self.conn.lock().unwrap();
267        conn.execute("DELETE FROM sessions WHERE orchestrator_id = ?1", [id])?;
268        conn.execute("DELETE FROM sessions WHERE id = ?1", [id])?;
269        conn.execute("DELETE FROM orchestrators WHERE id = ?1", [id])?;
270        Ok(())
271    }
272
273    pub fn upsert_pr(&self, pr: &PR) -> Result<()> {
274        let conn = self.conn.lock().unwrap();
275        conn.execute(
276            "INSERT OR REPLACE INTO prs
277             (id, number, title, url, body, session_id)
278             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
279            params![pr.id, pr.number, pr.title, pr.url, pr.body, pr.session_id],
280        )?;
281        Ok(())
282    }
283
284    pub fn get_pr(&self, id: PrId) -> Result<Option<PR>> {
285        let conn = self.conn.lock().unwrap();
286        let mut stmt = conn.prepare(
287            "SELECT id, number, title, url, body, session_id FROM prs WHERE id = ?1",
288        )?;
289        let mut rows = stmt.query_map([id], |r| {
290            Ok(PR {
291                id:         r.get(0)?,
292                number:     r.get(1)?,
293                title:      r.get(2)?,
294                url:        r.get(3)?,
295                body:       r.get(4)?,
296                session_id: r.get(5)?,
297            })
298        })?;
299        rows.next().transpose().map_err(Into::into)
300    }
301
302    pub fn upsert_ci_status(&self, ci: &CIStatus) -> Result<()> {
303        let conn = self.conn.lock().unwrap();
304        conn.execute(
305            "INSERT OR REPLACE INTO ci_status
306             (pr_id, total, passing, failing, pending)
307             VALUES (?1, ?2, ?3, ?4, ?5)",
308            params![ci.pr_id, ci.total, ci.passing, ci.failing, ci.pending],
309        )?;
310        Ok(())
311    }
312
313    pub fn upsert_comment(&self, c: &Comment) -> Result<()> {
314        let conn = self.conn.lock().unwrap();
315        conn.execute(
316            "INSERT OR REPLACE INTO review_comments
317             (id, pr_id, author, body, path, line, created_at)
318             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
319            params![c.id, c.pr_id, c.author, c.body, c.path, c.line, c.created_at],
320        )?;
321        Ok(())
322    }
323
324    pub fn list_orchestrators(&self) -> Result<Vec<Orchestrator>> {
325        let conn = self.conn.lock().unwrap();
326        let mut stmt = conn.prepare(
327            "SELECT id,name,created_at FROM orchestrators ORDER BY created_at DESC",
328        )?;
329        let rows = stmt.query_map([], |r| {
330            Ok(Orchestrator {
331                id: r.get(0)?,
332                name: r.get(1)?,
333                created_at: r.get(2)?,
334            })
335        })?;
336        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343    use tempfile::tempdir;
344
345    fn test_store() -> Store {
346        let dir = tempdir().unwrap();
347        let path = dir.path().join("t.db");
348        // keep dir alive for the lifetime of the test by leaking it
349        std::mem::forget(dir);
350        Store::open(path).unwrap()
351    }
352
353    #[test]
354    fn upsert_and_list_session() {
355        let store = test_store();
356        let session = Session {
357            id: "s1".into(), orchestrator_id: None, name: "worker-1".into(),
358            repo: "slievr/Athene".into(), status: SessionStatus::Working,
359            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
360            pr_number: None, pr_id: None, workspace_path: None, pid: None,
361            model: None, context_tokens: None, catalogue_path: None,
362            context_used_pct: None, context_total_tokens: None, context_window_size: None,
363            claude_session_id: None,
364            summary: None,
365            terminal_at: None,
366        };
367        store.upsert_session(&session).unwrap();
368        let list = store.list_sessions().unwrap();
369        assert_eq!(list.len(), 1);
370        assert_eq!(list[0].id, "s1");
371    }
372
373    #[test]
374    fn upsert_updates_status() {
375        let store = test_store();
376        let mut s = Session {
377            id: "s1".into(), orchestrator_id: None, name: "w".into(),
378            repo: "r".into(), status: SessionStatus::Working,
379            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
380            pr_number: None, pr_id: None, workspace_path: None, pid: None,
381            model: None, context_tokens: None, catalogue_path: None,
382            context_used_pct: None, context_total_tokens: None, context_window_size: None,
383            claude_session_id: None,
384            summary: None,
385            terminal_at: None,
386        };
387        store.upsert_session(&s).unwrap();
388        s.status = SessionStatus::Done;
389        store.upsert_session(&s).unwrap();
390        let list = store.list_sessions().unwrap();
391        assert_eq!(list.len(), 1);
392        assert!(matches!(list[0].status, SessionStatus::Done));
393    }
394
395    /// The poller's dual-remote self-heal (`session.repo` corrected once a
396    /// PR is found against a different remote than the one on record)
397    /// depends on `repo` being an updatable column, not just a write-once
398    /// field set at insert time.
399    #[test]
400    fn upsert_updates_repo() {
401        let store = test_store();
402        let mut s = Session {
403            id: "s1".into(), orchestrator_id: None, name: "w".into(),
404            repo: "OwnerA/repoA".into(), status: SessionStatus::Working,
405            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
406            pr_number: None, pr_id: None, workspace_path: None, pid: None,
407            model: None, context_tokens: None, catalogue_path: None,
408            context_used_pct: None, context_total_tokens: None, context_window_size: None,
409            claude_session_id: None,
410            summary: None,
411            terminal_at: None,
412        };
413        store.upsert_session(&s).unwrap();
414        s.repo = "OwnerB/repoB".into();
415        store.upsert_session(&s).unwrap();
416        let updated = store.get_session("s1").unwrap().unwrap();
417        assert_eq!(updated.repo, "OwnerB/repoB");
418    }
419
420    #[test]
421    fn get_session_by_id() {
422        let store = test_store();
423        let s = Session {
424            id: "s1".into(), orchestrator_id: None, name: "w".into(),
425            repo: "r".into(), status: SessionStatus::Working,
426            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
427            pr_number: None, pr_id: None, workspace_path: None, pid: None,
428            model: None, context_tokens: None, catalogue_path: None,
429            context_used_pct: None, context_total_tokens: None, context_window_size: None,
430            claude_session_id: None,
431            summary: None,
432            terminal_at: None,
433        };
434        store.upsert_session(&s).unwrap();
435        let found = store.get_session("s1").unwrap();
436        assert!(found.is_some());
437        assert_eq!(found.unwrap().name, "w");
438        assert!(store.get_session("missing").unwrap().is_none());
439    }
440
441    #[test]
442    fn model_and_context_tokens_round_trip() {
443        let store = test_store();
444        let s = Session {
445            id: "s1".into(), orchestrator_id: None, name: "w".into(),
446            repo: "r".into(), status: SessionStatus::Working,
447            agent_type: "claude-code".into(), cost_usd: 1.5, started_at: 0,
448            pr_number: None, pr_id: None, workspace_path: None, pid: None,
449            model: Some("claude-fable-5".into()), context_tokens: Some(214_000), catalogue_path: None,
450            context_used_pct: None, context_total_tokens: None, context_window_size: None,
451            claude_session_id: None,
452            summary: None,
453            terminal_at: None,
454        };
455        store.upsert_session(&s).unwrap();
456        let found = store.get_session("s1").unwrap().unwrap();
457        assert_eq!(found.model.as_deref(), Some("claude-fable-5"));
458        assert_eq!(found.context_tokens, Some(214_000));
459    }
460
461    #[test]
462    fn catalogue_path_round_trips() {
463        let store = test_store();
464        let s = Session {
465            id: "s2".into(), orchestrator_id: None, name: "w".into(),
466            repo: "r".into(), status: SessionStatus::Working,
467            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
468            pr_number: None, pr_id: None, workspace_path: None, pid: None,
469            model: None, context_tokens: None,
470            catalogue_path: Some("/brains/x".into()),
471            context_used_pct: None, context_total_tokens: None, context_window_size: None,
472            claude_session_id: None,
473            summary: None,
474            terminal_at: None,
475        };
476        store.upsert_session(&s).unwrap();
477        let found = store.get_session("s2").unwrap().unwrap();
478        assert_eq!(found.catalogue_path.as_deref(), Some("/brains/x"));
479        // list path decodes it too
480        assert_eq!(store.list_sessions().unwrap()[0].catalogue_path.as_deref(), Some("/brains/x"));
481    }
482
483    #[test]
484    fn summary_round_trips() {
485        let store = test_store();
486        let s = Session {
487            id: "s2b".into(), orchestrator_id: None, name: "w".into(),
488            repo: "r".into(), status: SessionStatus::Working,
489            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
490            pr_number: None, pr_id: None, workspace_path: None, pid: None,
491            model: None, context_tokens: None, catalogue_path: None,
492            context_used_pct: None, context_total_tokens: None, context_window_size: None,
493            claude_session_id: None,
494            summary: Some("Fix flaky CI on the auth suite".into()),
495            terminal_at: None,
496        };
497        store.upsert_session(&s).unwrap();
498        let found = store.get_session("s2b").unwrap().unwrap();
499        assert_eq!(found.summary.as_deref(), Some("Fix flaky CI on the auth suite"));
500        // list path decodes it too
501        assert_eq!(
502            store.list_sessions().unwrap().iter().find(|x| x.id == "s2b").unwrap().summary.as_deref(),
503            Some("Fix flaky CI on the auth suite"),
504        );
505
506        // None round-trips as None, not "" or "none"
507        let mut s2 = s.clone();
508        s2.id = "s2c".into();
509        s2.summary = None;
510        store.upsert_session(&s2).unwrap();
511        assert_eq!(store.get_session("s2c").unwrap().unwrap().summary, None);
512    }
513
514    #[test]
515    fn claude_session_id_round_trips() {
516        let store = test_store();
517        let s = Session {
518            id: "s3".into(), orchestrator_id: None, name: "w".into(),
519            repo: "r".into(), status: SessionStatus::Working,
520            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
521            pr_number: None, pr_id: None, workspace_path: None, pid: None,
522            model: None, context_tokens: None, catalogue_path: None,
523            context_used_pct: None, context_total_tokens: None, context_window_size: None,
524            claude_session_id: Some("b7e0b3a0-0000-4000-8000-000000000001".into()),
525            summary: None,
526            terminal_at: None,
527        };
528        store.upsert_session(&s).unwrap();
529        let found = store.get_session("s3").unwrap().unwrap();
530        assert_eq!(found.claude_session_id.as_deref(), Some("b7e0b3a0-0000-4000-8000-000000000001"));
531        // list path decodes it too
532        assert_eq!(
533            store.list_sessions().unwrap().iter().find(|x| x.id == "s3").unwrap().claude_session_id.as_deref(),
534            Some("b7e0b3a0-0000-4000-8000-000000000001"),
535        );
536
537        // None round-trips as None, not "" or "none"
538        let mut s2 = s.clone();
539        s2.id = "s4".into();
540        s2.claude_session_id = None;
541        store.upsert_session(&s2).unwrap();
542        assert_eq!(store.get_session("s4").unwrap().unwrap().claude_session_id, None);
543    }
544
545    #[test]
546    fn terminal_at_round_trips() {
547        let store = test_store();
548        let s = Session {
549            id: "s5".into(), orchestrator_id: None, name: "w".into(),
550            repo: "r".into(), status: SessionStatus::Done,
551            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
552            pr_number: None, pr_id: None, workspace_path: None, pid: None,
553            model: None, context_tokens: None, catalogue_path: None,
554            context_used_pct: None, context_total_tokens: None, context_window_size: None,
555            claude_session_id: None,
556            summary: None,
557            terminal_at: Some(1_720_000_000_000),
558        };
559        store.upsert_session(&s).unwrap();
560        let found = store.get_session("s5").unwrap().unwrap();
561        assert_eq!(found.terminal_at, Some(1_720_000_000_000));
562        // list path decodes it too
563        assert_eq!(
564            store.list_sessions().unwrap().iter().find(|x| x.id == "s5").unwrap().terminal_at,
565            Some(1_720_000_000_000),
566        );
567
568        // None round-trips as None — non-terminal / not-yet-retired sessions.
569        let mut s2 = s.clone();
570        s2.id = "s6".into();
571        s2.terminal_at = None;
572        store.upsert_session(&s2).unwrap();
573        assert_eq!(store.get_session("s6").unwrap().unwrap().terminal_at, None);
574    }
575
576    /// A Re-file respawns the same session id with a fresh `started_at` —
577    /// the conflict-update path must persist it, or the in-memory time
578    /// silently reverts to the original spawn time on app restart.
579    #[test]
580    fn upsert_conflict_updates_started_at() {
581        let store = test_store();
582        let mut s = Session {
583            id: "s3".into(), orchestrator_id: None, name: "w".into(),
584            repo: "r".into(), status: SessionStatus::Working,
585            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 100,
586            pr_number: None, pr_id: None, workspace_path: None, pid: None,
587            model: None, context_tokens: None, catalogue_path: None,
588            context_used_pct: None, context_total_tokens: None, context_window_size: None,
589            claude_session_id: None,
590            summary: None,
591            terminal_at: None,
592        };
593        store.upsert_session(&s).unwrap();
594        s.started_at = 200;
595        store.upsert_session(&s).unwrap();
596        assert_eq!(store.get_session("s3").unwrap().unwrap().started_at, 200);
597    }
598
599    #[test]
600    fn get_pr_round_trips_and_misses_cleanly() {
601        let store = test_store();
602        assert!(store.get_pr(9).unwrap().is_none());
603        let pr = PR {
604            id: 9, number: 9, title: "t".into(),
605            url: "https://github.com/org/repo/pull/9".into(),
606            body: String::new(), session_id: "s1".into(),
607        };
608        store.upsert_pr(&pr).unwrap();
609        let found = store.get_pr(9).unwrap().unwrap();
610        assert_eq!(found.number, 9);
611        assert_eq!(found.session_id, "s1");
612        assert_eq!(found.url, pr.url);
613    }
614
615    #[test]
616    fn context_fields_round_trip() {
617        let store = test_store();
618        let s = Session {
619            id: "s1".into(), orchestrator_id: None, name: "w".into(),
620            repo: "r".into(), status: SessionStatus::Working,
621            agent_type: "claude-code".into(), cost_usd: 2.6, started_at: 0,
622            pr_number: None, pr_id: None, workspace_path: None, pid: None,
623            model: None, context_tokens: None, catalogue_path: None,
624            context_used_pct: Some(62.0),
625            context_total_tokens: Some(124_000),
626            context_window_size: Some(200_000),
627            claude_session_id: None,
628            summary: None,
629            terminal_at: None,
630        };
631        store.upsert_session(&s).unwrap();
632        let found = store.get_session("s1").unwrap().unwrap();
633        assert_eq!(found.context_used_pct, Some(62.0));
634        assert_eq!(found.context_total_tokens, Some(124_000));
635        assert_eq!(found.context_window_size, Some(200_000));
636        // list path decodes it too
637        assert_eq!(store.list_sessions().unwrap()[0].context_used_pct, Some(62.0));
638    }
639
640    #[test]
641    fn context_fields_default_to_none() {
642        let store = test_store();
643        let s = Session {
644            id: "s2".into(), orchestrator_id: None, name: "w".into(),
645            repo: "r".into(), status: SessionStatus::Working,
646            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
647            pr_number: None, pr_id: None, workspace_path: None, pid: None,
648            model: None, context_tokens: None, catalogue_path: None,
649            context_used_pct: None, context_total_tokens: None, context_window_size: None,
650            claude_session_id: None,
651            summary: None,
652            terminal_at: None,
653        };
654        store.upsert_session(&s).unwrap();
655        let found = store.get_session("s2").unwrap().unwrap();
656        assert_eq!(found.context_used_pct, None);
657        assert_eq!(found.context_total_tokens, None);
658        assert_eq!(found.context_window_size, None);
659    }
660
661    #[test]
662    fn cost_samples_filters_by_agent_and_model_and_excludes_zero() {
663        let store = test_store();
664        for (id, agent_type, model, cost) in [
665            ("a", "claude-code", Some("claude-fable-5"), 3.0),
666            ("b", "claude-code", Some("claude-fable-5"), 5.0),
667            ("c", "claude-code", Some("claude-fable-5"), 0.0), // excluded: zero cost
668            ("d", "claude-code", Some("claude-opus-4-8"), 2.0), // excluded: different model
669            ("e", "codex",       Some("claude-fable-5"), 4.0),  // excluded: different harness
670        ] {
671            store.upsert_session(&Session {
672                id: id.into(), orchestrator_id: None, name: id.into(),
673                repo: "r".into(), status: SessionStatus::Working,
674                agent_type: agent_type.into(), cost_usd: cost, started_at: 0,
675                pr_number: None, pr_id: None, workspace_path: None, pid: None,
676                model: model.map(String::from), context_tokens: None, catalogue_path: None,
677                context_used_pct: None, context_total_tokens: None, context_window_size: None,
678                claude_session_id: None, summary: None,
679                terminal_at: None,
680            }).unwrap();
681        }
682        let samples = store.cost_samples("claude-code", Some("claude-fable-5")).unwrap();
683        assert_eq!(samples.len(), 2);
684        assert!((samples.iter().sum::<f64>() - 8.0).abs() < f64::EPSILON);
685    }
686}
ninox_core/store.rs

ninox_core/
store.rs
