Skip to main content

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