Skip to main content

agent_graph_mcp/
daemon.rs

1//! Lock-owning daemon primitives.
2use crate::migrations;
3use fs2::FileExt;
4use rusqlite::{Connection, OptionalExtension};
5use std::{
6    fs::{self, File, OpenOptions},
7    io,
8    os::unix::fs::OpenOptionsExt,
9    path::{Path, PathBuf},
10    time::{SystemTime, UNIX_EPOCH},
11};
12#[allow(dead_code)]
13pub const MAX_FRAME: usize = 1024 * 1024;
14#[derive(Debug)]
15pub enum DaemonError {
16    AlreadyOwned,
17    Io(io::Error),
18    Sql(rusqlite::Error),
19}
20impl From<io::Error> for DaemonError {
21    fn from(e: io::Error) -> Self {
22        Self::Io(e)
23    }
24}
25impl From<rusqlite::Error> for DaemonError {
26    fn from(e: rusqlite::Error) -> Self {
27        Self::Sql(e)
28    }
29}
30
31impl std::fmt::Display for DaemonError {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        match self {
34            Self::AlreadyOwned => write!(f, "data directory already owned by another daemon"),
35            Self::Io(e) => write!(f, "daemon io error: {e}"),
36            Self::Sql(e) => write!(f, "daemon sql error: {e}"),
37        }
38    }
39}
40
41impl std::error::Error for DaemonError {}
42impl DaemonError {
43    pub fn code(&self) -> &'static str {
44        match self {
45            Self::AlreadyOwned => "DATA_DIR_ALREADY_OWNED",
46            Self::Io(_) => "DAEMON_IO",
47            Self::Sql(_) => "DAEMON_SQL",
48        }
49    }
50}
51#[derive(Debug)]
52pub struct DaemonLock {
53    file: File,
54    #[allow(dead_code)]
55    pub path: PathBuf,
56}
57impl DaemonLock {
58    pub fn acquire(data_dir: &Path) -> Result<Self, DaemonError> {
59        fs::create_dir_all(data_dir)?;
60        let path = data_dir.join("daemon.lock");
61        let file = OpenOptions::new()
62            .create(true)
63            .read(true)
64            .write(true)
65            .mode(0o600)
66            .open(&path)?;
67        file.try_lock_exclusive().map_err(|e| {
68            if e.kind() == io::ErrorKind::WouldBlock {
69                DaemonError::AlreadyOwned
70            } else {
71                DaemonError::Io(e)
72            }
73        })?;
74        Ok(Self { file, path })
75    }
76}
77impl Drop for DaemonLock {
78    fn drop(&mut self) {
79        let _ = self.file.unlock();
80        let _ = self.file.sync_all();
81    }
82}
83
84#[allow(dead_code)]
85#[derive(Debug, Clone)]
86pub struct DaemonIdentity {
87    pub instance_id: String,
88    pub generation: i64,
89    pub pid: u32,
90    pub started_at: String,
91}
92
93#[allow(dead_code)]
94pub fn identity(conn: &Connection) -> rusqlite::Result<DaemonIdentity> {
95    use std::sync::atomic::{AtomicU64, Ordering};
96    static COUNTER: AtomicU64 = AtomicU64::new(0);
97    let now = SystemTime::now()
98        .duration_since(UNIX_EPOCH)
99        .unwrap_or_default()
100        .as_secs();
101    let instance_id = format!(
102        "{}-{}-{}",
103        std::process::id(),
104        now,
105        COUNTER.fetch_add(1, Ordering::SeqCst)
106    );
107    let generation: i64 = conn.query_row(
108        "SELECT COALESCE(MAX(generation),0)+1 FROM daemon_instances",
109        [],
110        |r| r.get(0),
111    )?;
112    let started_at = now.to_string();
113    conn.execute(
114        "INSERT INTO daemon_instances(instance_id,generation,pid,started_at,heartbeat_at) VALUES (?1,?2,?3,?4,?4)",
115        rusqlite::params![instance_id, generation, std::process::id(), started_at],
116    )?;
117    Ok(DaemonIdentity {
118        instance_id,
119        generation,
120        pid: std::process::id(),
121        started_at,
122    })
123}
124
125#[allow(dead_code)]
126pub fn recover_owned_state(
127    conn: &Connection,
128    instance_id: &str,
129    generation: i64,
130) -> rusqlite::Result<usize> {
131    let has_executions: bool = conn
132        .query_row(
133            "SELECT 1 FROM sqlite_master WHERE type='table' AND name='executions'",
134            [],
135            |row| row.get::<_, i64>(0),
136        )
137        .optional()?
138        .is_some();
139
140    let has_owner_column: bool = if has_executions {
141        conn.query_row(
142            "SELECT 1 FROM pragma_table_info('executions') WHERE name='owner_instance_id'",
143            [],
144            |row| row.get::<_, i64>(0),
145        )
146        .optional()?
147        .is_some()
148    } else {
149        false
150    };
151
152    let changed = if has_executions && has_owner_column {
153        conn.execute(
154            "UPDATE executions SET status='legacy_unverified' WHERE owner_instance_id IS NULL AND status IN ('accepted','running')",
155            [],
156        )?
157    } else {
158        0
159    };
160
161    conn.execute(
162        "UPDATE daemon_instances SET heartbeat_at=CURRENT_TIMESTAMP WHERE instance_id=?1 AND generation=?2",
163        rusqlite::params![instance_id, generation],
164    )?;
165    Ok(changed)
166}
167pub fn open_owned(
168    data_dir: &Path,
169    binary_digest: &str,
170) -> Result<(DaemonLock, Connection), DaemonError> {
171    let lock = DaemonLock::acquire(data_dir)?;
172    let mut c = Connection::open(data_dir.join("agent-graph.db"))?;
173    migrations::apply(&mut c, binary_digest)?;
174    Ok((lock, c))
175}
176
177/// Persist the integrity mode and reject mixed keyless/key-enabled restarts.
178pub fn enforce_startup_mode(conn: &Connection, key_enabled: bool) -> rusqlite::Result<()> {
179    conn.execute_batch("CREATE TABLE IF NOT EXISTS daemon_startup_mode (mode TEXT PRIMARY KEY CHECK(mode IN ('keyless','key-enabled')));")?;
180    let expected = if key_enabled {
181        "key-enabled"
182    } else {
183        "keyless"
184    };
185    let current: Option<String> = conn
186        .query_row("SELECT mode FROM daemon_startup_mode LIMIT 1", [], |r| {
187            r.get(0)
188        })
189        .optional()?;
190    match current {
191        Some(mode) if mode != expected => Err(rusqlite::Error::InvalidParameterName(
192            "STARTUP_MODE_MISMATCH".into(),
193        )),
194        Some(_) => Ok(()),
195        None => {
196            conn.execute(
197                "INSERT INTO daemon_startup_mode(mode) VALUES (?1)",
198                [expected],
199            )?;
200            Ok(())
201        }
202    }
203}
204#[allow(dead_code)]
205pub fn socket_path(runtime_dir: &Path, instance: &str) -> PathBuf {
206    runtime_dir
207        .join("agent-graph")
208        .join(instance)
209        .join("daemon.sock")
210}