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

ninox_core/
store.rs
