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

ninox_core/
store.rs
