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 upsert_ci_status(&self, ci: &CIStatus) -> Result<()> {
239        let conn = self.conn.lock().unwrap();
240        conn.execute(
241            "INSERT OR REPLACE INTO ci_status
242             (pr_id, total, passing, failing, pending)
243             VALUES (?1, ?2, ?3, ?4, ?5)",
244            params![ci.pr_id, ci.total, ci.passing, ci.failing, ci.pending],
245        )?;
246        Ok(())
247    }
248
249    pub fn upsert_comment(&self, c: &Comment) -> Result<()> {
250        let conn = self.conn.lock().unwrap();
251        conn.execute(
252            "INSERT OR REPLACE INTO review_comments
253             (id, pr_id, author, body, path, line, created_at)
254             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
255            params![c.id, c.pr_id, c.author, c.body, c.path, c.line, c.created_at],
256        )?;
257        Ok(())
258    }
259
260    pub fn list_orchestrators(&self) -> Result<Vec<Orchestrator>> {
261        let conn = self.conn.lock().unwrap();
262        let mut stmt = conn.prepare(
263            "SELECT id,name,created_at FROM orchestrators ORDER BY created_at DESC",
264        )?;
265        let rows = stmt.query_map([], |r| {
266            Ok(Orchestrator {
267                id: r.get(0)?,
268                name: r.get(1)?,
269                created_at: r.get(2)?,
270            })
271        })?;
272        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279    use tempfile::tempdir;
280
281    fn test_store() -> Store {
282        let dir = tempdir().unwrap();
283        let path = dir.path().join("t.db");
284        // keep dir alive for the lifetime of the test by leaking it
285        std::mem::forget(dir);
286        Store::open(path).unwrap()
287    }
288
289    #[test]
290    fn upsert_and_list_session() {
291        let store = test_store();
292        let session = Session {
293            id: "s1".into(), orchestrator_id: None, name: "worker-1".into(),
294            repo: "slievr/Athene".into(), status: SessionStatus::Working,
295            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
296            pr_number: None, pr_id: None, workspace_path: None, pid: None,
297            model: None, context_tokens: None, catalogue_path: None,
298        };
299        store.upsert_session(&session).unwrap();
300        let list = store.list_sessions().unwrap();
301        assert_eq!(list.len(), 1);
302        assert_eq!(list[0].id, "s1");
303    }
304
305    #[test]
306    fn upsert_updates_status() {
307        let store = test_store();
308        let mut s = Session {
309            id: "s1".into(), orchestrator_id: None, name: "w".into(),
310            repo: "r".into(), status: SessionStatus::Working,
311            agent_type: "c".into(), cost_usd: 0.0, started_at: 0,
312            pr_number: None, pr_id: None, workspace_path: None, pid: None,
313            model: None, context_tokens: None, catalogue_path: None,
314        };
315        store.upsert_session(&s).unwrap();
316        s.status = SessionStatus::Done;
317        store.upsert_session(&s).unwrap();
318        let list = store.list_sessions().unwrap();
319        assert_eq!(list.len(), 1);
320        assert!(matches!(list[0].status, SessionStatus::Done));
321    }
322
323    #[test]
324    fn get_session_by_id() {
325        let store = test_store();
326        let 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        let found = store.get_session("s1").unwrap();
335        assert!(found.is_some());
336        assert_eq!(found.unwrap().name, "w");
337        assert!(store.get_session("missing").unwrap().is_none());
338    }
339
340    #[test]
341    fn model_and_context_tokens_round_trip() {
342        let store = test_store();
343        let s = Session {
344            id: "s1".into(), orchestrator_id: None, name: "w".into(),
345            repo: "r".into(), status: SessionStatus::Working,
346            agent_type: "claude-code".into(), cost_usd: 1.5, started_at: 0,
347            pr_number: None, pr_id: None, workspace_path: None, pid: None,
348            model: Some("claude-fable-5".into()), context_tokens: Some(214_000), catalogue_path: None,
349        };
350        store.upsert_session(&s).unwrap();
351        let found = store.get_session("s1").unwrap().unwrap();
352        assert_eq!(found.model.as_deref(), Some("claude-fable-5"));
353        assert_eq!(found.context_tokens, Some(214_000));
354    }
355
356    #[test]
357    fn catalogue_path_round_trips() {
358        let store = test_store();
359        let s = Session {
360            id: "s2".into(), orchestrator_id: None, name: "w".into(),
361            repo: "r".into(), status: SessionStatus::Working,
362            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 0,
363            pr_number: None, pr_id: None, workspace_path: None, pid: None,
364            model: None, context_tokens: None,
365            catalogue_path: Some("/brains/x".into()),
366        };
367        store.upsert_session(&s).unwrap();
368        let found = store.get_session("s2").unwrap().unwrap();
369        assert_eq!(found.catalogue_path.as_deref(), Some("/brains/x"));
370        // list path decodes it too
371        assert_eq!(store.list_sessions().unwrap()[0].catalogue_path.as_deref(), Some("/brains/x"));
372    }
373
374    /// A Re-file respawns the same session id with a fresh `started_at` —
375    /// the conflict-update path must persist it, or the in-memory time
376    /// silently reverts to the original spawn time on app restart.
377    #[test]
378    fn upsert_conflict_updates_started_at() {
379        let store = test_store();
380        let mut s = Session {
381            id: "s3".into(), orchestrator_id: None, name: "w".into(),
382            repo: "r".into(), status: SessionStatus::Working,
383            agent_type: "claude-code".into(), cost_usd: 0.0, started_at: 100,
384            pr_number: None, pr_id: None, workspace_path: None, pid: None,
385            model: None, context_tokens: None, catalogue_path: None,
386        };
387        store.upsert_session(&s).unwrap();
388        s.started_at = 200;
389        store.upsert_session(&s).unwrap();
390        assert_eq!(store.get_session("s3").unwrap().unwrap().started_at, 200);
391    }
392
393    #[test]
394    fn cost_samples_filters_by_agent_and_model_and_excludes_zero() {
395        let store = test_store();
396        for (id, agent_type, model, cost) in [
397            ("a", "claude-code", Some("claude-fable-5"), 3.0),
398            ("b", "claude-code", Some("claude-fable-5"), 5.0),
399            ("c", "claude-code", Some("claude-fable-5"), 0.0), // excluded: zero cost
400            ("d", "claude-code", Some("claude-opus-4-8"), 2.0), // excluded: different model
401            ("e", "codex",       Some("claude-fable-5"), 4.0),  // excluded: different harness
402        ] {
403            store.upsert_session(&Session {
404                id: id.into(), orchestrator_id: None, name: id.into(),
405                repo: "r".into(), status: SessionStatus::Working,
406                agent_type: agent_type.into(), cost_usd: cost, started_at: 0,
407                pr_number: None, pr_id: None, workspace_path: None, pid: None,
408                model: model.map(String::from), context_tokens: None, catalogue_path: None,
409            }).unwrap();
410        }
411        let samples = store.cost_samples("claude-code", Some("claude-fable-5")).unwrap();
412        assert_eq!(samples.len(), 2);
413        assert!((samples.iter().sum::<f64>() - 8.0).abs() < f64::EPSILON);
414    }
415}
ninox_core/store.rs

ninox_core/
store.rs
