Skip to main content

agent_graph_mcp/
migrations.rs

1//! Versioned durable-daemon migrations. Kept independent so the store can call it.
2use rusqlite::{params, Connection, OptionalExtension};
3use sha2::{Digest, Sha256};
4
5pub const CURRENT_VERSION: i64 = 3;
6pub const LEGACY_OWNER_UNKNOWN: &str = "legacy_owner_unknown";
7
8#[allow(dead_code)]
9pub trait MigrationStore {
10    fn connection(&self) -> &Connection;
11}
12
13pub fn apply(conn: &mut Connection, binary_digest: &str) -> rusqlite::Result<()> {
14    let tx = conn.transaction()?;
15    tx.execute_batch("CREATE TABLE IF NOT EXISTS schema_migrations (version INTEGER PRIMARY KEY, applied_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, migration_digest TEXT NOT NULL);")?;
16    let exists: Option<i64> = tx
17        .query_row(
18            "SELECT version FROM schema_migrations WHERE version = ?1",
19            [CURRENT_VERSION],
20            |r| r.get(0),
21        )
22        .optional()?;
23    if exists.is_none() {
24        tx.execute_batch("CREATE TABLE IF NOT EXISTS daemon_instances (instance_id TEXT PRIMARY KEY, generation INTEGER NOT NULL UNIQUE, pid INTEGER NOT NULL, boot_id TEXT, executable_digest TEXT, started_at TEXT NOT NULL, heartbeat_at TEXT NOT NULL, clean_shutdown_at TEXT); CREATE TABLE IF NOT EXISTS run_publication_state (run_id TEXT PRIMARY KEY, state TEXT NOT NULL, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP, reason TEXT); CREATE TABLE IF NOT EXISTS operator_receipts (receipt_id TEXT PRIMARY KEY, request_digest TEXT NOT NULL, action TEXT NOT NULL, resource_kind TEXT NOT NULL, resource_id TEXT NOT NULL, state_digest TEXT NOT NULL, operator_uid INTEGER NOT NULL, daemon_instance_id TEXT NOT NULL, nonce TEXT NOT NULL UNIQUE, issued_at TEXT NOT NULL, expires_at TEXT NOT NULL, consumed_at TEXT); CREATE INDEX IF NOT EXISTS idx_operator_receipts_nonce ON operator_receipts(nonce);")?;
25        let has_owner: Option<String> = tx
26            .query_row(
27                "SELECT 1 FROM pragma_table_info('executions') WHERE name='owner_instance_id'",
28                [],
29                |r| r.get(0),
30            )
31            .optional()?;
32        if has_owner.is_none() {
33            // Only add owner_instance_id if executions table exists.
34            // The store's own migration creates it; if it doesn't exist yet,
35            // the store migration will need to include the column.
36            let has_executions: Option<i64> = tx
37                .query_row(
38                    "SELECT 1 FROM sqlite_master WHERE type='table' AND name='executions'",
39                    [],
40                    |r| r.get(0),
41                )
42                .optional()?;
43            if has_executions.is_some() {
44                tx.execute_batch("ALTER TABLE executions ADD COLUMN owner_instance_id TEXT;")?;
45                tx.execute("UPDATE executions SET owner_instance_id = ?1 WHERE owner_instance_id IS NULL AND status IN ('accepted','running')", [LEGACY_OWNER_UNKNOWN])?;
46            }
47        }
48        let digest = migration_digest(binary_digest);
49        tx.execute(
50            "INSERT INTO schema_migrations(version,migration_digest) VALUES (?1,?2)",
51            params![CURRENT_VERSION, digest],
52        )?;
53    }
54    tx.commit()
55}
56
57pub fn migration_digest(binary_digest: &str) -> String {
58    let mut h = Sha256::new();
59    h.update(format!(
60        "agent-graph-mcp:migration:{CURRENT_VERSION}:{binary_digest}"
61    ));
62    format!("{:x}", h.finalize())
63}
64
65#[allow(dead_code)]
66pub fn owner_for_new_run<'a>(owner: &'a str) -> rusqlite::Result<&'a str> {
67    if owner.is_empty() {
68        Err(rusqlite::Error::InvalidParameterName(
69            "owner_instance_id".into(),
70        ))
71    } else {
72        Ok(owner)
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    #[test]
80    fn legacy_active_rows_are_quarantined() {
81        let mut c = Connection::open_in_memory().unwrap();
82        c.execute_batch("CREATE TABLE executions(run_id TEXT PRIMARY KEY,status TEXT NOT NULL); INSERT INTO executions VALUES ('a','running'),('b','completed');").unwrap();
83        apply(&mut c, "bin").unwrap();
84        let owner: String = c
85            .query_row(
86                "SELECT owner_instance_id FROM executions WHERE run_id='a'",
87                [],
88                |r| r.get(0),
89            )
90            .unwrap();
91        assert_eq!(owner, LEGACY_OWNER_UNKNOWN);
92        apply(&mut c, "bin").unwrap();
93        assert_eq!(
94            c.query_row::<i64, _, _>("SELECT count(*) FROM schema_migrations", [], |r| r.get(0))
95                .unwrap(),
96            1
97        );
98    }
99}