nyx_agent_core/store/
mod.rs

1//! SQLite persistence layer for the agent.
2//!
3//! `Store::open` resolves a SQLite file under the configured state
4//! directory (`~/.local/share/nyx-agent/state.db` by default), applies
5//! every migration shipped under `migrations/`, sets WAL + the tuning
6//! pragmas the plan specifies, and hands back a clonable pool wrapper.
7//!
8//! Per-table repository structs (`RepoStore`, `RunStore`, ...) are
9//! returned by accessor methods on `Store`. Every repository borrows a
10//! pool reference; cloning `Store` is cheap because `SqlitePool` is an
11//! `Arc` internally.
12
13use std::path::{Path, PathBuf};
14use std::str::FromStr;
15
16use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions, SqliteSynchronous};
17use sqlx::{ConnectOptions, SqlitePool};
18use thiserror::Error;
19
20pub mod attack_graph;
21pub mod budget;
22pub mod candidate;
23pub mod chain;
24pub mod exploration_memory;
25pub mod feedback;
26pub mod finding;
27pub mod integration;
28pub mod payload;
29pub mod product;
30pub mod project;
31pub mod repo;
32pub mod repro;
33pub mod run;
34pub mod run_repo_outcome;
35pub mod schedule;
36pub mod schema;
37pub mod spec;
38pub mod trace;
39pub mod webhook;
40
41#[cfg(test)]
42pub(crate) mod testutil;
43
44pub use attack_graph::{attack_graph_edge_id, attack_graph_node_id, AttackGraphStore};
45pub use budget::{BudgetKind, BudgetRecord, BudgetStore};
46pub use candidate::{CandidateFindingRecord, CandidateFindingStore, CandidateStatus};
47pub use chain::{ChainRecord, ChainStore};
48pub use exploration_memory::{
49    compact_memory_for_prompt, ExplorationMemoryInput, ExplorationMemoryRecord,
50    ExplorationMemoryStore,
51};
52pub use feedback::{FeedbackRecord, FeedbackStore, OperatorVerdict};
53pub use finding::{
54    finding_id_hash, FindingFilter, FindingOrigin, FindingRecord, FindingStatus, FindingStore,
55    TriageState,
56};
57pub use integration::{
58    ProjectIntegrationEvent, ProjectIntegrationInsert, ProjectIntegrationKind,
59    ProjectIntegrationPatch, ProjectIntegrationRecord, ProjectIntegrationStore,
60    ProjectIntegrationStoredRecord,
61};
62pub use payload::{PayloadRecord, PayloadStore};
63pub use product::{
64    canonical_risk_rating, clamp_risk_score, AuthzMatrixEntryRecord, AuthzMatrixStore,
65    BusinessLogicTemplateRunRecord, BusinessLogicTemplateRunStore, EnvironmentRunRecord,
66    EnvironmentRunStore, LaunchProfileStore, NyxSignalRecord, NyxSignalStore,
67    PentestCandidateRecord, PentestCandidateStore, ProjectLaunchProfile, ProjectLaunchProfileInput,
68    RouteModelRecord, RouteModelStore, VerificationAttemptRecord, VerificationAttemptStore,
69    VerifiedVulnerabilityRecord, VerifiedVulnerabilityStore,
70};
71pub use project::{
72    ProjectPatch, ProjectPatchOption, ProjectRecord, ProjectRuntimeProfile, ProjectStore,
73    DEFAULT_PROJECT_ID, DEFAULT_PROJECT_NAME,
74};
75pub use repo::{PatchOption, RepoPatch, RepoRecord, RepoStore, SourceKind};
76pub use repro::{ReproBundleRecord, ReproBundleStore};
77pub use run::{RunRecord, RunStatus, RunStore, TriggeredBy};
78pub use run_repo_outcome::{RepoOutcomeLabel, RunRepoOutcomeRecord, RunRepoOutcomeStore};
79pub use schedule::{ScheduleRecord, ScheduleStore};
80pub use schema::{schema_version, CURRENT_SCHEMA_VERSION, MIGRATOR};
81pub use spec::{HarnessSpecRecord, HarnessSpecStore};
82pub use trace::{AgentTraceRecord, AgentTraceStore, TaskKind};
83pub use webhook::{WebhookRecord, WebhookStore};
84
85const DB_FILE_NAME: &str = "state.db";
86
87#[derive(Debug, Error)]
88pub enum StoreError {
89    #[error("failed to open database at {path}: {source}")]
90    Open {
91        path: PathBuf,
92        #[source]
93        source: sqlx::Error,
94    },
95    #[error("failed to apply migrations: {0}")]
96    Migrate(#[from] sqlx::migrate::MigrateError),
97    #[error("database error: {0}")]
98    Sqlx(#[from] sqlx::Error),
99    #[error("invalid project runtime profile JSON: {0}")]
100    ProjectRuntimeProfileJson(#[from] serde_json::Error),
101    #[error("invalid integration JSON: {0}")]
102    IntegrationJson(serde_json::Error),
103    #[error("invalid integration kind: {0}")]
104    InvalidIntegrationKind(String),
105}
106
107/// Connected, migrated SQLite store. Cloning is cheap (pool is `Arc`).
108#[derive(Debug, Clone)]
109pub struct Store {
110    pool: SqlitePool,
111    path: PathBuf,
112}
113
114impl Store {
115    /// Open `<state_dir>/state.db`, creating the file if missing, and
116    /// apply every migration shipped under `migrations/`.
117    pub async fn open(state_dir: &Path) -> Result<Self, StoreError> {
118        let path = state_dir.join(DB_FILE_NAME);
119        Self::open_at(&path).await
120    }
121
122    /// Open an arbitrary SQLite file. Used by tests against a tempdir.
123    pub async fn open_at(path: &Path) -> Result<Self, StoreError> {
124        if let Some(parent) = path.parent() {
125            std::fs::create_dir_all(parent).map_err(|source| StoreError::Open {
126                path: path.to_path_buf(),
127                source: sqlx::Error::Io(source),
128            })?;
129        }
130
131        let url = format!("sqlite://{}", path.display());
132        let opts = SqliteConnectOptions::from_str(&url)
133            .map_err(|source| StoreError::Open { path: path.to_path_buf(), source })?
134            .create_if_missing(true)
135            .journal_mode(SqliteJournalMode::Wal)
136            .synchronous(SqliteSynchronous::Normal)
137            .foreign_keys(true)
138            .pragma("cache_size", "-8000")
139            .pragma("temp_store", "MEMORY")
140            .disable_statement_logging();
141
142        let pool = SqlitePoolOptions::new()
143            .max_connections(8)
144            .min_connections(1)
145            .connect_with(opts)
146            .await
147            .map_err(|source| StoreError::Open { path: path.to_path_buf(), source })?;
148
149        MIGRATOR.run(&pool).await?;
150        Self::populate_meta(&pool).await?;
151        Self::ensure_default_project(&pool).await?;
152        Self::warn_if_finding_id_collision_pressure(&pool).await?;
153
154        Ok(Self { pool, path: path.to_path_buf() })
155    }
156
157    /// `finding_id_hash` truncates BLAKE3 to 64 bits so finding ids fit
158    /// neatly in UI rows. Birthday-collision pressure is negligible for
159    /// any plausible single-deployment row count, but a deployment that
160    /// crosses ~2^28 rows is close enough to the 2^32 expected-collision
161    /// bound that the operator deserves a warning while there is still
162    /// time to plan a schema migration (a `hash_version` column on
163    /// `findings`, or a widened id). Uses `MAX(rowid)` instead of
164    /// `COUNT(*)` so the check is O(log n) and adds no measurable
165    /// open-time latency on small databases.
166    async fn warn_if_finding_id_collision_pressure(pool: &SqlitePool) -> Result<(), sqlx::Error> {
167        const WARN_THRESHOLD: i64 = 1 << 28;
168        let (max_rowid,): (Option<i64>,) =
169            sqlx::query_as("SELECT MAX(rowid) FROM findings").fetch_one(pool).await?;
170        if max_rowid.unwrap_or(0) >= WARN_THRESHOLD {
171            tracing::warn!(
172                target: "nyx_agent_core::store",
173                approx_findings_rowid = max_rowid,
174                threshold = WARN_THRESHOLD,
175                "findings table crossed 2^28 rows; finding_id_hash truncates BLAKE3 to 64 bits so birthday collisions become statistically meaningful near 2^32 rows. Plan a schema migration that widens finding ids or adds a hash_version column before that boundary."
176            );
177        }
178        Ok(())
179    }
180
181    /// Stamp the singleton `meta` row with the running binary's version,
182    /// the applied schema version (mirrored from `_sqlx_migrations`), and
183    /// a real `created_at` on first run. Subsequent opens leave
184    /// `created_at` untouched so it remains the install timestamp.
185    async fn populate_meta(pool: &SqlitePool) -> Result<(), sqlx::Error> {
186        let now_ms = crate::time::now_epoch_ms();
187        let agent_version = env!("CARGO_PKG_VERSION");
188        let max_migration: (Option<i64>,) =
189            sqlx::query_as("SELECT MAX(version) FROM _sqlx_migrations").fetch_one(pool).await?;
190        let schema_v = max_migration.0.unwrap_or(0);
191        sqlx::query(
192            "UPDATE meta SET \
193                 schema_version = ?1, \
194                 created_at = CASE WHEN created_at = 0 THEN ?2 ELSE created_at END, \
195                 agent_version = ?3 \
196             WHERE id = 1",
197        )
198        .bind(schema_v)
199        .bind(now_ms)
200        .bind(agent_version)
201        .execute(pool)
202        .await?;
203        Ok(())
204    }
205
206    /// Transitional bootstrap: seed the `DEFAULT_PROJECT_ID` row so
207    /// legacy callers without an explicit project context can still
208    /// attach repos via the FK.
209    async fn ensure_default_project(pool: &SqlitePool) -> Result<(), StoreError> {
210        let now_ms = crate::time::now_epoch_ms();
211        ProjectStore::new(pool).ensure_default(now_ms).await?;
212        Ok(())
213    }
214
215    pub fn pool(&self) -> &SqlitePool {
216        &self.pool
217    }
218
219    pub fn path(&self) -> &Path {
220        &self.path
221    }
222
223    pub async fn close(self) {
224        self.pool.close().await;
225    }
226
227    /// Current applied schema version. Reads `meta.schema_version`, which
228    /// is kept in sync with `MAX(_sqlx_migrations.version)` on every
229    /// `open` via `populate_meta`. In debug builds the two are
230    /// cross-checked; in release the `meta` row is authoritative.
231    pub async fn schema_version(&self) -> Result<i64, StoreError> {
232        let (meta_v,): (i64,) = sqlx::query_as("SELECT schema_version FROM meta WHERE id = 1")
233            .fetch_one(&self.pool)
234            .await?;
235        debug_assert_eq!(
236            meta_v,
237            schema_version(&self.pool).await?,
238            "meta.schema_version drift from _sqlx_migrations"
239        );
240        Ok(meta_v)
241    }
242
243    pub fn projects(&self) -> ProjectStore<'_> {
244        ProjectStore::new(&self.pool)
245    }
246    pub fn repos(&self) -> RepoStore<'_> {
247        RepoStore::new(&self.pool)
248    }
249    pub fn runs(&self) -> RunStore<'_> {
250        RunStore::new(&self.pool)
251    }
252    pub fn attack_graph(&self) -> AttackGraphStore<'_> {
253        AttackGraphStore::new(&self.pool)
254    }
255    pub fn exploration_memory(&self) -> ExplorationMemoryStore<'_> {
256        ExplorationMemoryStore::new(&self.pool)
257    }
258    pub fn run_repo_outcomes(&self) -> RunRepoOutcomeStore<'_> {
259        RunRepoOutcomeStore::new(&self.pool)
260    }
261    pub fn findings(&self) -> FindingStore<'_> {
262        FindingStore::new(&self.pool)
263    }
264    pub fn integrations(&self) -> ProjectIntegrationStore<'_> {
265        ProjectIntegrationStore::new(&self.pool)
266    }
267    pub fn chains(&self) -> ChainStore<'_> {
268        ChainStore::new(&self.pool)
269    }
270    pub fn payloads(&self) -> PayloadStore<'_> {
271        PayloadStore::new(&self.pool)
272    }
273    pub fn launch_profiles(&self) -> LaunchProfileStore<'_> {
274        LaunchProfileStore::new(&self.pool)
275    }
276    pub fn environment_runs(&self) -> EnvironmentRunStore<'_> {
277        EnvironmentRunStore::new(&self.pool)
278    }
279    pub fn nyx_signals(&self) -> NyxSignalStore<'_> {
280        NyxSignalStore::new(&self.pool)
281    }
282    pub fn pentest_candidates(&self) -> PentestCandidateStore<'_> {
283        PentestCandidateStore::new(&self.pool)
284    }
285    pub fn business_logic_template_runs(&self) -> BusinessLogicTemplateRunStore<'_> {
286        BusinessLogicTemplateRunStore::new(&self.pool)
287    }
288    pub fn verification_attempts(&self) -> VerificationAttemptStore<'_> {
289        VerificationAttemptStore::new(&self.pool)
290    }
291    pub fn authz_matrix(&self) -> AuthzMatrixStore<'_> {
292        AuthzMatrixStore::new(&self.pool)
293    }
294    pub fn verified_vulnerabilities(&self) -> VerifiedVulnerabilityStore<'_> {
295        VerifiedVulnerabilityStore::new(&self.pool)
296    }
297    pub fn route_models(&self) -> RouteModelStore<'_> {
298        RouteModelStore::new(&self.pool)
299    }
300    pub fn candidate_findings(&self) -> CandidateFindingStore<'_> {
301        CandidateFindingStore::new(&self.pool)
302    }
303    pub fn agent_traces(&self) -> AgentTraceStore<'_> {
304        AgentTraceStore::new(&self.pool)
305    }
306    pub fn budgets(&self) -> BudgetStore<'_> {
307        BudgetStore::new(&self.pool)
308    }
309    pub fn repro_bundles(&self) -> ReproBundleStore<'_> {
310        ReproBundleStore::new(&self.pool)
311    }
312    pub fn schedules(&self) -> ScheduleStore<'_> {
313        ScheduleStore::new(&self.pool)
314    }
315    pub fn webhooks(&self) -> WebhookStore<'_> {
316        WebhookStore::new(&self.pool)
317    }
318    pub fn feedback(&self) -> FeedbackStore<'_> {
319        FeedbackStore::new(&self.pool)
320    }
321    pub fn harness_specs(&self) -> HarnessSpecStore<'_> {
322        HarnessSpecStore::new(&self.pool)
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329
330    async fn open_tmp() -> (tempfile::TempDir, Store) {
331        let tmp = tempfile::tempdir().expect("tempdir");
332        let store = Store::open(tmp.path()).await.expect("open");
333        (tmp, store)
334    }
335
336    #[tokio::test]
337    async fn open_creates_db_and_runs_migrations() {
338        let (tmp, store) = open_tmp().await;
339        let db_path = tmp.path().join("state.db");
340        assert!(db_path.exists(), "state.db should be created");
341        let v = store.schema_version().await.expect("schema_version");
342        assert_eq!(v, CURRENT_SCHEMA_VERSION);
343    }
344
345    #[tokio::test]
346    async fn meta_row_populated_with_real_values() {
347        let (_tmp, store) = open_tmp().await;
348        let (schema_v, created_at, agent_version): (i64, i64, String) = sqlx::query_as(
349            "SELECT schema_version, created_at, agent_version FROM meta WHERE id = 1",
350        )
351        .fetch_one(store.pool())
352        .await
353        .expect("meta row");
354        assert_eq!(schema_v, CURRENT_SCHEMA_VERSION);
355        assert!(created_at > 0, "created_at must be real epoch ms, got {created_at}");
356        assert_eq!(agent_version, env!("CARGO_PKG_VERSION"));
357    }
358
359    #[tokio::test]
360    async fn meta_created_at_stable_across_reopens() {
361        let tmp = tempfile::tempdir().expect("tempdir");
362        let store = Store::open(tmp.path()).await.expect("open");
363        let (first,): (i64,) = sqlx::query_as("SELECT created_at FROM meta WHERE id = 1")
364            .fetch_one(store.pool())
365            .await
366            .expect("created_at");
367        store.close().await;
368        assert!(first > 0);
369        // Sleep briefly so that any naive "always overwrite" would shift.
370        tokio::time::sleep(std::time::Duration::from_millis(5)).await;
371        let store = Store::open(tmp.path()).await.expect("reopen");
372        let (second,): (i64,) = sqlx::query_as("SELECT created_at FROM meta WHERE id = 1")
373            .fetch_one(store.pool())
374            .await
375            .expect("created_at");
376        assert_eq!(first, second, "created_at must be preserved across reopens");
377    }
378
379    #[tokio::test]
380    async fn migrations_idempotent_across_reopens() {
381        let tmp = tempfile::tempdir().expect("tempdir");
382        for _ in 0..3 {
383            let store = Store::open(tmp.path()).await.expect("open");
384            assert_eq!(
385                store.schema_version().await.expect("schema_version"),
386                CURRENT_SCHEMA_VERSION
387            );
388            store.close().await;
389        }
390    }
391
392    #[tokio::test]
393    async fn cross_table_finding_payload_chain_trace_roundtrip() {
394        // Acceptance test: insert a finding, payload, chain, and agent trace,
395        // query them back, then verify FK cascades.
396        let (_tmp, s) = open_tmp().await;
397        let testutil = &super::testutil::sample_repo;
398        s.repos().upsert(&testutil("repo")).await.expect("repo");
399        s.runs().insert(&super::testutil::sample_run("run-1")).await.expect("run");
400        let f = super::testutil::sample_finding("run-1", "repo", "src/a.rs", "rule-1");
401        let fid = f.id.clone();
402        s.findings().upsert(&f).await.expect("finding");
403        let p = super::testutil::sample_payload("payload-1", &fid);
404        s.payloads().insert(&p).await.expect("payload");
405        let c = super::testutil::sample_chain("chain-1", "run-1", &[&fid]);
406        s.chains().insert(&c).await.expect("chain");
407        s.findings().set_chain(&fid, "chain-1").await.expect("link");
408        let t = super::trace::AgentTraceRecord {
409            id: "trace-1".to_string(),
410            finding_id: Some(fid.clone()),
411            task_kind: "PayloadSynthesis".to_string(),
412            runtime_name: "anthropic".to_string(),
413            model: "claude-opus-4-7".to_string(),
414            prompt_version: Some("payload/v1".to_string()),
415            conversation_jsonl_path: None,
416            tokens_in: 10,
417            tokens_out: 20,
418            cost_usd_micros: 100,
419            cache_hits: 0,
420            cache_misses: 0,
421            duration_ms: Some(5_000),
422            started_at: 5_000,
423            finished_at: Some(10_000),
424            verifier_blob: None,
425        };
426        s.agent_traces().insert(&t).await.expect("trace");
427
428        // Read everything back.
429        assert_eq!(s.findings().get(&fid).await.expect("f").expect("row").id, fid);
430        assert_eq!(s.payloads().get("payload-1").await.expect("p").expect("row").finding_id, fid);
431        assert_eq!(s.chains().get("chain-1").await.expect("c").expect("row").run_id, "run-1");
432        assert_eq!(
433            s.agent_traces().get("trace-1").await.expect("t").expect("row").finding_id.as_deref(),
434            Some(fid.as_str())
435        );
436
437        // Deleting the run cascades to findings, payloads, chains, and
438        // SET-NULLs the trace's finding_id.
439        s.runs().delete("run-1").await.expect("del run");
440        assert!(s.findings().get(&fid).await.expect("f").is_none());
441        assert!(s.payloads().get("payload-1").await.expect("p").is_none());
442        assert!(s.chains().get("chain-1").await.expect("c").is_none());
443        let trace = s.agent_traces().get("trace-1").await.expect("t").expect("row");
444        assert!(
445            trace.finding_id.is_none(),
446            "trace.finding_id must be SET NULL after finding cascade-deleted"
447        );
448    }
449
450    #[tokio::test]
451    async fn pragmas_are_set() {
452        let (_tmp, store) = open_tmp().await;
453        let (jmode,): (String,) = sqlx::query_as("PRAGMA journal_mode")
454            .fetch_one(store.pool())
455            .await
456            .expect("journal_mode");
457        assert_eq!(jmode.to_lowercase(), "wal");
458
459        let (sync,): (i64,) = sqlx::query_as("PRAGMA synchronous")
460            .fetch_one(store.pool())
461            .await
462            .expect("synchronous");
463        // NORMAL = 1
464        assert_eq!(sync, 1);
465
466        let (fk,): (i64,) = sqlx::query_as("PRAGMA foreign_keys")
467            .fetch_one(store.pool())
468            .await
469            .expect("foreign_keys");
470        assert_eq!(fk, 1, "foreign keys must be enabled");
471    }
472}
nyx_agent_core/store/mod.rs

nyx_agent_core/store/
mod.rs
