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            CREATE TABLE IF NOT EXISTS sessions (
17                id TEXT PRIMARY KEY, orchestrator_id TEXT,
18                name TEXT NOT NULL, repo TEXT NOT NULL,
19                status TEXT NOT NULL, agent_type TEXT NOT NULL,
20                cost_usd REAL NOT NULL DEFAULT 0, started_at INTEGER NOT NULL,
21                pr_number INTEGER, pr_id INTEGER,
22                workspace_path TEXT, pid INTEGER,
23                model TEXT, context_tokens INTEGER
24            );
25            CREATE TABLE IF NOT EXISTS orchestrators (
26                id TEXT PRIMARY KEY, name TEXT NOT NULL, created_at INTEGER NOT NULL
27            );
28            CREATE TABLE IF NOT EXISTS prs (
29                id INTEGER PRIMARY KEY, number INTEGER NOT NULL,
30                title TEXT NOT NULL, url TEXT NOT NULL,
31                body TEXT NOT NULL, session_id TEXT NOT NULL
32            );
33            CREATE TABLE IF NOT EXISTS ci_status (
34                pr_id INTEGER PRIMARY KEY, total INTEGER NOT NULL,
35                passing INTEGER NOT NULL, failing INTEGER NOT NULL,
36                pending INTEGER NOT NULL
37            );
38            CREATE TABLE IF NOT EXISTS review_comments (
39                id INTEGER PRIMARY KEY, pr_id INTEGER NOT NULL,
40                author TEXT NOT NULL, body TEXT NOT NULL,
41                path TEXT, line INTEGER, created_at INTEGER NOT NULL
42            );
43        ")?;
44        // Migrations for columns added after initial release — idempotent so
45        // both fresh and pre-existing databases end up with the same schema.
46        for (col, ddl) in [
47            ("model",             "ALTER TABLE sessions ADD COLUMN model TEXT"),
48            ("context_tokens",    "ALTER TABLE sessions ADD COLUMN context_tokens INTEGER"),
49            ("catalogue_path",    "ALTER TABLE sessions ADD COLUMN catalogue_path TEXT"),
50            ("claude_session_id", "ALTER TABLE sessions ADD COLUMN claude_session_id TEXT"),
51        ] {
52            if !Self::column_exists(&conn, "sessions", col)? {
53                conn.execute(ddl, [])?;
54            }
55        }
56        Ok(Self { conn: Mutex::new(conn) })
57    }
58
59    fn column_exists(conn: &Connection, table: &str, column: &str) -> Result<bool> {
60        let mut stmt = conn.prepare(&format!("PRAGMA table_info({table})"))?;
61        let exists = stmt
62            .query_map([], |r| r.get::<_, String>(1))?
63            .filter_map(|r| r.ok())
64            .any(|c| c == column);
65        Ok(exists)
66    }
67
68    pub fn upsert_session(&self, s: &Session) -> Result<()> {
69        let status = serde_json::to_string(&s.status)?.replace('"', "");
70        let conn = self.conn.lock().unwrap();
71        conn.execute(
72            "INSERT INTO sessions (id,orchestrator_id,name,repo,status,agent_type,
73             cost_usd,started_at,pr_number,pr_id,workspace_path,pid,model,context_tokens,
74             catalogue_path,claude_session_id)
75             VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,?15,?16)
76             ON CONFLICT(id) DO UPDATE SET
77             status=excluded.status,cost_usd=excluded.cost_usd,
78             started_at=excluded.started_at,
79             pr_number=excluded.pr_number,pr_id=excluded.pr_id,
80             workspace_path=excluded.workspace_path,pid=excluded.pid,
81             model=excluded.model,context_tokens=excluded.context_tokens,
82             catalogue_path=excluded.catalogue_path,
83             claude_session_id=excluded.claude_session_id",
84            params![
85                s.id, s.orchestrator_id, s.name, s.repo, status, s.agent_type,
86                s.cost_usd, s.started_at, s.pr_number, s.pr_id,
87                s.workspace_path, s.pid, s.model, s.context_tokens,
88                s.catalogue_path, s.claude_session_id
89            ],
90        )?;
91        Ok(())
92    }
93
94    pub fn list_sessions(&self) -> Result<Vec<Session>> {
95        let conn = self.conn.lock().unwrap();
96        let mut stmt = conn.prepare(
97            "SELECT id,orchestrator_id,name,repo,status,agent_type,cost_usd,
98             started_at,pr_number,pr_id,workspace_path,pid,model,context_tokens,
99             catalogue_path,claude_session_id
100             FROM sessions ORDER BY started_at DESC",
101        )?;
102        let rows = stmt.query_map([], |r| {
103            Ok((
104                r.get::<_, String>(0)?,
105                r.get::<_, Option<String>>(1)?,
106                r.get::<_, String>(2)?,
107                r.get::<_, String>(3)?,
108                r.get::<_, String>(4)?,
109                r.get::<_, String>(5)?,
110                r.get::<_, f64>(6)?,
111                r.get::<_, i64>(7)?,
112                r.get::<_, Option<u64>>(8)?,
113                r.get::<_, Option<i64>>(9)?,
114                r.get::<_, Option<String>>(10)?,
115                r.get::<_, Option<u32>>(11)?,
116                r.get::<_, Option<String>>(12)?,
117                r.get::<_, Option<i64>>(13)?,
118                r.get::<_, Option<String>>(14)?,
119                r.get::<_, Option<String>>(15)?,
120            ))
121        })?;
122        rows.map(|r| {
123            let (id, orchestrator_id, name, repo, status_str, agent_type,
124                 cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
125                 model, context_tokens, catalogue_path, claude_session_id) = r?;
126            let status = serde_json::from_str(&format!("\"{status_str}\""))
127                .unwrap_or(SessionStatus::Working);
128            Ok(Session {
129                id, orchestrator_id, name, repo, status, agent_type,
130                cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
131                model, context_tokens: context_tokens.map(|v| v.max(0) as u64),
132                catalogue_path, claude_session_id,
133            })
134        })
135        .collect()
136    }
137
138    pub fn get_session(&self, id: &str) -> Result<Option<Session>> {
139        let conn = self.conn.lock().unwrap();
140        let mut stmt = conn.prepare(
141            "SELECT id,orchestrator_id,name,repo,status,agent_type,cost_usd,
142             started_at,pr_number,pr_id,workspace_path,pid,model,context_tokens,
143             catalogue_path,claude_session_id
144             FROM sessions WHERE id = ?1",
145        )?;
146        let mut rows = stmt.query_map([id], |r| {
147            Ok((
148                r.get::<_, String>(0)?,
149                r.get::<_, Option<String>>(1)?,
150                r.get::<_, String>(2)?,
151                r.get::<_, String>(3)?,
152                r.get::<_, String>(4)?,
153                r.get::<_, String>(5)?,
154                r.get::<_, f64>(6)?,
155                r.get::<_, i64>(7)?,
156                r.get::<_, Option<u64>>(8)?,
157                r.get::<_, Option<i64>>(9)?,
158                r.get::<_, Option<String>>(10)?,
159                r.get::<_, Option<u32>>(11)?,
160                r.get::<_, Option<String>>(12)?,
161                r.get::<_, Option<i64>>(13)?,
162                r.get::<_, Option<String>>(14)?,
163                r.get::<_, Option<String>>(15)?,
164            ))
165        })?;
166        match rows.next() {
167            None => Ok(None),
168            Some(r) => {
169                let (id, orchestrator_id, name, repo, status_str, agent_type,
170                     cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
171                     model, context_tokens, catalogue_path, claude_session_id) = r?;
172                let status = serde_json::from_str(&format!("\"{status_str}\""))
173                    .unwrap_or(SessionStatus::Working);
174                Ok(Some(Session {
175                    id, orchestrator_id, name, repo, status, agent_type,
176                    cost_usd, started_at, pr_number, pr_id, workspace_path, pid,
177                    model, context_tokens: context_tokens.map(|v| v.max(0) as u64),
178                    catalogue_path, claude_session_id,
179                }))
180            }
181        }
182    }
183
184    /// Non-zero `cost_usd` samples recorded for sessions matching the given
185    /// agent harness (`agent_type`) and model — used to compute a
186    /// data-driven spawn-modal cost estimate once enough history exists for
187    /// a given preset. Read-only; built on `list_sessions` like
188    /// `sessions_by_orchestrator`.
189    pub fn cost_samples(&self, agent_type: &str, model: Option<&str>) -> Result<Vec<f64>> {
190        let sessions = self.list_sessions()?;
191        Ok(sessions
192            .into_iter()
193            .filter(|s| {
194                s.agent_type == agent_type
195                    && s.model.as_deref() == model
196                    && s.cost_usd > 0.0
197            })
198            .map(|s| s.cost_usd)
199            .collect())
200    }
201
202    pub fn upsert_orchestrator(&self, o: &Orchestrator) -> Result<()> {
203        let conn = self.conn.lock().unwrap();
204        conn.execute(
205            "INSERT INTO orchestrators(id,name,created_at) VALUES(?1,?2,?3)
206             ON CONFLICT(id) DO UPDATE SET name=excluded.name",
207            params![o.id, o.name, o.created_at],
208        )?;
209        Ok(())
210    }
211
212    pub fn sessions_by_orchestrator(&self, orchestrator_id: &str) -> Result<Vec<Session>> {
213        let sessions = self.list_sessions()?;
214        Ok(sessions.into_iter().filter(|s| s.orchestrator_id.as_deref() == Some(orchestrator_id)).collect())
215    }
216
217    pub fn delete_session(&self, id: &str) -> Result<()> {
218        let conn = self.conn.lock().unwrap();
219        conn.execute("DELETE FROM sessions WHERE id = ?1", [id])?;
220        Ok(())
221    }
222
223    pub fn delete_orchestrator(&self, id: &str) -> Result<()> {
224        let conn = self.conn.lock().unwrap();
225        conn.execute("DELETE FROM sessions WHERE orchestrator_id = ?1", [id])?;
226        conn.execute("DELETE FROM sessions WHERE id = ?1", [id])?;
227        conn.execute("DELETE FROM orchestrators WHERE id = ?1", [id])?;
228        Ok(())
229    }
230
231    pub fn upsert_pr(&self, pr: &PR) -> Result<()> {
232        let conn = self.conn.lock().unwrap();
233        conn.execute(
234            "INSERT OR REPLACE INTO prs
235             (id, number, title, url, body, session_id)
236             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
237            params![pr.id, pr.number, pr.title, pr.url, pr.body, pr.session_id],
238        )?;
239        Ok(())
240    }
241
242    pub fn get_pr(&self, id: PrId) -> Result<Option<PR>> {
243        let conn = self.conn.lock().unwrap();
244        let mut stmt = conn.prepare(
245            "SELECT id, number, title, url, body, session_id FROM prs WHERE id = ?1",
246        )?;
247        let mut rows = stmt.query_map([id], |r| {
248            Ok(PR {
249                id:         r.get(0)?,
250                number:     r.get(1)?,
251                title:      r.get(2)?,
252                url:        r.get(3)?,
253                body:       r.get(4)?,
254                session_id: r.get(5)?,
255            })
256        })?;
257        rows.next().transpose().map_err(Into::into)
258    }
259
260    pub fn upsert_ci_status(&self, ci: &CIStatus) -> Result<()> {
261        let conn = self.conn.lock().unwrap();
262        conn.execute(
263            "INSERT OR REPLACE INTO ci_status
264             (pr_id, total, passing, failing, pending)
265             VALUES (?1, ?2, ?3, ?4, ?5)",
266            params![ci.pr_id, ci.total, ci.passing, ci.failing, ci.pending],
267        )?;
268        Ok(())
269    }
270
271    pub fn upsert_comment(&self, c: &Comment) -> Result<()> {
272        let conn = self.conn.lock().unwrap();
273        conn.execute(
274            "INSERT OR REPLACE INTO review_comments
275             (id, pr_id, author, body, path, line, created_at)
276             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
277            params![c.id, c.pr_id, c.author, c.body, c.path, c.line, c.created_at],
278        )?;
279        Ok(())
280    }
281
282    pub fn list_orchestrators(&self) -> Result<Vec<Orchestrator>> {
283        let conn = self.conn.lock().unwrap();
284        let mut stmt = conn.prepare(
285            "SELECT id,name,created_at FROM orchestrators ORDER BY created_at DESC",
286        )?;
287        let rows = stmt.query_map([], |r| {
288            Ok(Orchestrator {
289                id: r.get(0)?,
290                name: r.get(1)?,
291                created_at: r.get(2)?,
292            })
293        })?;
294        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use tempfile::tempdir;
302
303    fn test_store() -> Store {
304        let dir = tempdir().unwrap();
305        let path = dir.path().join("t.db");
306        // keep dir alive for the lifetime of the test by leaking it
307        std::mem::forget(dir);
308        Store::open(path).unwrap()
309    }
310
311    #[test]
312    fn upsert_and_list_session() {
313        let store = test_store();
314        let session = Session {
315            id: "s1".into(), orchestrator_id: None, name: "worker-1".into(),
316            repo: "slievr/Athene".into(), status: SessionStatus::Working,
317            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
318            pr_number: None, pr_id: None, workspace_path: None, pid: None,
319            model: None, context_tokens: None, catalogue_path: None, claude_session_id: None,
320        };
321        store.upsert_session(&session).unwrap();
322        let list = store.list_sessions().unwrap();
323        assert_eq!(list.len(), 1);
324        assert_eq!(list[0].id, "s1");
325    }
326
327    #[test]
328    fn upsert_updates_status() {
329        let store = test_store();
330        let mut s = Session {
331            id: "s1".into(), orchestrator_id: None, name: "w".into(),
332            repo: "r".into(), status: SessionStatus::Working,
333            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
334            pr_number: None, pr_id: None, workspace_path: None, pid: None,
335            model: None, context_tokens: None, catalogue_path: None, claude_session_id: None,
336        };
337        store.upsert_session(&s).unwrap();
338        s.status = SessionStatus::Done;
339        store.upsert_session(&s).unwrap();
340        let list = store.list_sessions().unwrap();
341        assert_eq!(list.len(), 1);
342        assert!(matches!(list[0].status, SessionStatus::Done));
343    }
344
345    #[test]
346    fn get_session_by_id() {
347        let store = test_store();
348        let s = Session {
349            id: "s1".into(), orchestrator_id: None, name: "w".into(),
350            repo: "r".into(), status: SessionStatus::Working,
351            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
352            pr_number: None, pr_id: None, workspace_path: None, pid: None,
353            model: None, context_tokens: None, catalogue_path: None, claude_session_id: None,
354        };
355        store.upsert_session(&s).unwrap();
356        let found = store.get_session("s1").unwrap();
357        assert!(found.is_some());
358        assert_eq!(found.unwrap().name, "w");
359        assert!(store.get_session("missing").unwrap().is_none());
360    }
361
362    #[test]
363    fn model_and_context_tokens_round_trip() {
364        let store = test_store();
365        let s = Session {
366            id: "s1".into(), orchestrator_id: None, name: "w".into(),
367            repo: "r".into(), status: SessionStatus::Working,
368            agent_type: "claude-code".into(), cost_usd: 1.5, started_at: 0,
369            pr_number: None, pr_id: None, workspace_path: None, pid: None,
370            model: Some("claude-fable-5".into()), context_tokens: Some(214_000), catalogue_path: None,
371            claude_session_id: None,
372        };
373        store.upsert_session(&s).unwrap();
374        let found = store.get_session("s1").unwrap().unwrap();
375        assert_eq!(found.model.as_deref(), Some("claude-fable-5"));
376        assert_eq!(found.context_tokens, Some(214_000));
377    }
378
379    #[test]
380    fn catalogue_path_round_trips() {
381        let store = test_store();
382        let s = Session {
383            id: "s2".into(), orchestrator_id: None, name: "w".into(),
384            repo: "r".into(), status: SessionStatus::Working,
385            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
386            pr_number: None, pr_id: None, workspace_path: None, pid: None,
387            model: None, context_tokens: None,
388            catalogue_path: Some("/brains/x".into()),
389            claude_session_id: None,
390        };
391        store.upsert_session(&s).unwrap();
392        let found = store.get_session("s2").unwrap().unwrap();
393        assert_eq!(found.catalogue_path.as_deref(), Some("/brains/x"));
394        // list path decodes it too
395        assert_eq!(store.list_sessions().unwrap()[0].catalogue_path.as_deref(), Some("/brains/x"));
396    }
397
398    #[test]
399    fn claude_session_id_round_trips() {
400        let store = test_store();
401        let s = Session {
402            id: "s3".into(), orchestrator_id: None, name: "w".into(),
403            repo: "r".into(), status: SessionStatus::Working,
404            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
405            pr_number: None, pr_id: None, workspace_path: None, pid: None,
406            model: None, context_tokens: None, catalogue_path: None,
407            claude_session_id: Some("b7e0b3a0-0000-4000-8000-000000000001".into()),
408        };
409        store.upsert_session(&s).unwrap();
410        let found = store.get_session("s3").unwrap().unwrap();
411        assert_eq!(found.claude_session_id.as_deref(), Some("b7e0b3a0-0000-4000-8000-000000000001"));
412        // list path decodes it too
413        assert_eq!(
414            store.list_sessions().unwrap().iter().find(|x| x.id == "s3").unwrap().claude_session_id.as_deref(),
415            Some("b7e0b3a0-0000-4000-8000-000000000001"),
416        );
417
418        // None round-trips as None, not "" or "none"
419        let mut s2 = s.clone();
420        s2.id = "s4".into();
421        s2.claude_session_id = None;
422        store.upsert_session(&s2).unwrap();
423        assert_eq!(store.get_session("s4").unwrap().unwrap().claude_session_id, None);
424    }
425
426    /// A Re-file respawns the same session id with a fresh `started_at` —
427    /// the conflict-update path must persist it, or the in-memory time
428    /// silently reverts to the original spawn time on app restart.
429    #[test]
430    fn upsert_conflict_updates_started_at() {
431        let store = test_store();
432        let mut s = Session {
433            id: "s3".into(), orchestrator_id: None, name: "w".into(),
434            repo: "r".into(), status: SessionStatus::Working,
435            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 100,
436            pr_number: None, pr_id: None, workspace_path: None, pid: None,
437            model: None, context_tokens: None, catalogue_path: None, claude_session_id: None,
438        };
439        store.upsert_session(&s).unwrap();
440        s.started_at = 200;
441        store.upsert_session(&s).unwrap();
442        assert_eq!(store.get_session("s3").unwrap().unwrap().started_at, 200);
443    }
444
445    #[test]
446    fn get_pr_round_trips_and_misses_cleanly() {
447        let store = test_store();
448        assert!(store.get_pr(9).unwrap().is_none());
449        let pr = PR {
450            id: 9, number: 9, title: "t".into(),
451            url: "https://github.com/org/repo/pull/9".into(),
452            body: String::new(), session_id: "s1".into(),
453        };
454        store.upsert_pr(&pr).unwrap();
455        let found = store.get_pr(9).unwrap().unwrap();
456        assert_eq!(found.number, 9);
457        assert_eq!(found.session_id, "s1");
458        assert_eq!(found.url, pr.url);
459    }
460
461    #[test]
462    fn cost_samples_filters_by_agent_and_model_and_excludes_zero() {
463        let store = test_store();
464        for (id, agent_type, model, cost) in [
465            ("a", "claude-code", Some("claude-fable-5"), 3.0),
466            ("b", "claude-code", Some("claude-fable-5"), 5.0),
467            ("c", "claude-code", Some("claude-fable-5"), 0.0), // excluded: zero cost
468            ("d", "claude-code", Some("claude-opus-4-8"), 2.0), // excluded: different model
469            ("e", "codex",       Some("claude-fable-5"), 4.0),  // excluded: different harness
470        ] {
471            store.upsert_session(&Session {
472                id: id.into(), orchestrator_id: None, name: id.into(),
473                repo: "r".into(), status: SessionStatus::Working,
474                agent_type: agent_type.into(), cost_usd: cost, started_at: 0,
475                pr_number: None, pr_id: None, workspace_path: None, pid: None,
476                model: model.map(String::from), context_tokens: None, catalogue_path: None,
477                claude_session_id: None,
478            }).unwrap();
479        }
480        let samples = store.cost_samples("claude-code", Some("claude-fable-5")).unwrap();
481        assert_eq!(samples.len(), 2);
482        assert!((samples.iter().sum::<f64>() - 8.0).abs() < f64::EPSILON);
483    }
484}
ninox_core/store.rs

ninox_core/
store.rs
