agentguard-local 0.1.0

package/crates/agentguard-core/src/store.rs

@@ -0,0 +1,434 @@
+use crate::lifecycle::ensure_transition;
+use crate::types::{EventRecord, JobRecord, JobState, Policy, ProcessSample, ToolPath};
+use chrono::{DateTime, Utc};
+use rusqlite::{Connection, OptionalExtension, params};
+use serde_json::{Value, json};
+use std::path::{Path, PathBuf};
+
+pub struct Store {
+    conn: Connection,
+    path: PathBuf,
+}
+
+impl Store {
+    pub fn open(path: impl AsRef<Path>) -> anyhow::Result<Self> {
+        let path = path.as_ref().to_path_buf();
+        if let Some(parent) = path.parent() {
+            std::fs::create_dir_all(parent)?;
+        }
+        let conn = Connection::open(&path)?;
+        conn.busy_timeout(std::time::Duration::from_secs(2))?;
+        let store = Self { conn, path };
+        store.init_schema()?;
+        Ok(store)
+    }
+
+    pub fn path(&self) -> &Path {
+        &self.path
+    }
+
+    fn init_schema(&self) -> anyhow::Result<()> {
+        self.conn.execute_batch(
+            r#"
+            PRAGMA journal_mode = WAL;
+            PRAGMA foreign_keys = ON;
+
+            CREATE TABLE IF NOT EXISTS jobs (
+                id TEXT PRIMARY KEY,
+                adapter TEXT NOT NULL,
+                command TEXT NOT NULL,
+                args_json TEXT NOT NULL,
+                cwd TEXT NOT NULL,
+                status TEXT NOT NULL,
+                priority TEXT NOT NULL,
+                root_pid INTEGER,
+                pgid INTEGER,
+                queued_reason TEXT,
+                started_at TEXT,
+                ended_at TEXT,
+                exit_code INTEGER,
+                policy_snapshot_json TEXT NOT NULL,
+                created_at TEXT NOT NULL,
+                updated_at TEXT NOT NULL
+            );
+
+            CREATE TABLE IF NOT EXISTS process_samples (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                job_id TEXT NOT NULL,
+                sampled_at TEXT NOT NULL,
+                pid_count INTEGER NOT NULL,
+                rss_bytes INTEGER NOT NULL,
+                footprint_bytes INTEGER NOT NULL,
+                cpu_pct REAL NOT NULL,
+                pressure_state TEXT NOT NULL
+            );
+
+            CREATE TABLE IF NOT EXISTS events (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                job_id TEXT,
+                ts TEXT NOT NULL,
+                type TEXT NOT NULL,
+                payload_json TEXT NOT NULL
+            );
+
+            CREATE TABLE IF NOT EXISTS policies (
+                name TEXT PRIMARY KEY,
+                created_at TEXT NOT NULL,
+                updated_at TEXT NOT NULL,
+                profile_json TEXT NOT NULL
+            );
+
+            CREATE TABLE IF NOT EXISTS tool_paths (
+                tool TEXT PRIMARY KEY,
+                real_path TEXT NOT NULL,
+                discovered_at TEXT NOT NULL,
+                hash TEXT
+            );
+            "#,
+        )?;
+        Ok(())
+    }
+
+    pub fn load_policy(&self) -> anyhow::Result<Policy> {
+        let profile: Option<String> = self
+            .conn
+            .query_row(
+                "SELECT profile_json FROM policies WHERE name = 'default'",
+                [],
+                |row| row.get(0),
+            )
+            .optional()?;
+        if let Some(profile) = profile {
+            Ok(serde_json::from_str(&profile)?)
+        } else {
+            let policy = Policy::default();
+            self.save_policy(&policy)?;
+            Ok(policy)
+        }
+    }
+
+    pub fn save_policy(&self, policy: &Policy) -> anyhow::Result<()> {
+        policy.validate()?;
+        let now = Utc::now().to_rfc3339();
+        let json = serde_json::to_string(policy)?;
+        self.conn.execute(
+            r#"
+            INSERT INTO policies (name, created_at, updated_at, profile_json)
+            VALUES ('default', ?1, ?1, ?2)
+            ON CONFLICT(name) DO UPDATE SET updated_at = excluded.updated_at,
+                                            profile_json = excluded.profile_json
+            "#,
+            params![now, json],
+        )?;
+        Ok(())
+    }
+
+    pub fn insert_job(&self, job: &JobRecord) -> anyhow::Result<()> {
+        self.conn.execute(
+            r#"
+            INSERT INTO jobs (
+                id, adapter, command, args_json, cwd, status, priority, root_pid, pgid,
+                queued_reason, started_at, ended_at, exit_code, policy_snapshot_json,
+                created_at, updated_at
+            )
+            VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)
+            "#,
+            params![
+                job.id,
+                job.adapter,
+                job.command,
+                serde_json::to_string(&job.args)?,
+                job.cwd,
+                job.status.to_string(),
+                job.priority.to_string(),
+                job.root_pid,
+                job.pgid,
+                job.queued_reason,
+                opt_ts(job.started_at),
+                opt_ts(job.ended_at),
+                job.exit_code,
+                job.policy_snapshot_json,
+                job.created_at.to_rfc3339(),
+                job.updated_at.to_rfc3339(),
+            ],
+        )?;
+        Ok(())
+    }
+
+    pub fn get_job(&self, id: &str) -> anyhow::Result<Option<JobRecord>> {
+        self.conn
+            .query_row("SELECT * FROM jobs WHERE id = ?1", params![id], map_job)
+            .optional()
+            .map_err(Into::into)
+    }
+
+    pub fn list_jobs(&self) -> anyhow::Result<Vec<JobRecord>> {
+        let mut stmt = self
+            .conn
+            .prepare("SELECT * FROM jobs ORDER BY created_at DESC, id ASC")?;
+        let rows = stmt.query_map([], map_job)?;
+        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
+    }
+
+    pub fn list_jobs_by_state(&self, states: &[JobState]) -> anyhow::Result<Vec<JobRecord>> {
+        let jobs = self.list_jobs()?;
+        Ok(jobs
+            .into_iter()
+            .filter(|job| states.contains(&job.status))
+            .collect())
+    }
+
+    pub fn active_count(&self) -> anyhow::Result<usize> {
+        Ok(self
+            .list_jobs()?
+            .into_iter()
+            .filter(|job| job.status.counts_as_active())
+            .count())
+    }
+
+    pub fn update_job_state(
+        &self,
+        id: &str,
+        to: JobState,
+        queued_reason: Option<&str>,
+    ) -> anyhow::Result<()> {
+        if let Some(current) = self.get_job(id)? {
+            ensure_transition(current.status, to)?;
+            let now = Utc::now();
+            let started_at = if matches!(to, JobState::Starting | JobState::Running) {
+                current.started_at.or(Some(now)).map(|ts| ts.to_rfc3339())
+            } else {
+                opt_ts(current.started_at)
+            };
+            let ended_at = if to.is_terminal() {
+                Some(now.to_rfc3339())
+            } else {
+                opt_ts(current.ended_at)
+            };
+            self.conn.execute(
+                r#"
+                UPDATE jobs
+                SET status = ?2, queued_reason = ?3, started_at = ?4, ended_at = ?5, updated_at = ?6
+                WHERE id = ?1
+                "#,
+                params![
+                    id,
+                    to.to_string(),
+                    queued_reason,
+                    started_at,
+                    ended_at,
+                    now.to_rfc3339()
+                ],
+            )?;
+        }
+        Ok(())
+    }
+
+    pub fn register_job(&self, id: &str, root_pid: i64, pgid: i64) -> anyhow::Result<()> {
+        if let Some(current) = self.get_job(id)? {
+            ensure_transition(current.status, JobState::Running)?;
+        }
+        let now = Utc::now().to_rfc3339();
+        self.conn.execute(
+            r#"
+            UPDATE jobs
+            SET status = 'running',
+                root_pid = ?2,
+                pgid = ?3,
+                started_at = COALESCE(started_at, ?4),
+                updated_at = ?4
+            WHERE id = ?1
+            "#,
+            params![id, root_pid, pgid, now],
+        )?;
+        Ok(())
+    }
+
+    pub fn complete_job(&self, id: &str, exit_code: i32) -> anyhow::Result<()> {
+        let state = if exit_code == 0 {
+            JobState::Exited
+        } else {
+            JobState::Failed
+        };
+        if let Some(current) = self.get_job(id)?
+            && current.status.is_terminal()
+        {
+            return Ok(());
+        }
+        let now = Utc::now().to_rfc3339();
+        self.conn.execute(
+            r#"
+            UPDATE jobs
+            SET status = ?2, ended_at = ?3, exit_code = ?4, updated_at = ?3
+            WHERE id = ?1
+            "#,
+            params![id, state.to_string(), now, exit_code],
+        )?;
+        Ok(())
+    }
+
+    pub fn add_event(
+        &self,
+        job_id: Option<&str>,
+        event_type: &str,
+        payload: Value,
+    ) -> anyhow::Result<()> {
+        self.conn.execute(
+            "INSERT INTO events (job_id, ts, type, payload_json) VALUES (?1, ?2, ?3, ?4)",
+            params![
+                job_id,
+                Utc::now().to_rfc3339(),
+                event_type,
+                serde_json::to_string(&payload)?,
+            ],
+        )?;
+        Ok(())
+    }
+
+    pub fn recent_events(&self, limit: usize) -> anyhow::Result<Vec<EventRecord>> {
+        let mut stmt = self.conn.prepare(
+            "SELECT id, job_id, ts, type, payload_json FROM events ORDER BY id DESC LIMIT ?1",
+        )?;
+        let rows = stmt.query_map(params![limit as i64], |row| {
+            let ts: String = row.get(2)?;
+            let payload: String = row.get(4)?;
+            Ok(EventRecord {
+                id: row.get(0)?,
+                job_id: row.get(1)?,
+                ts: parse_ts_sql(&ts)?,
+                event_type: row.get(3)?,
+                payload_json: serde_json::from_str(&payload).map_err(to_sql_error)?,
+            })
+        })?;
+        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
+    }
+
+    pub fn record_sample(&self, sample: &ProcessSample) -> anyhow::Result<()> {
+        self.conn.execute(
+            r#"
+            INSERT INTO process_samples (
+                job_id, sampled_at, pid_count, rss_bytes, footprint_bytes, cpu_pct, pressure_state
+            )
+            VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
+            "#,
+            params![
+                sample.job_id,
+                sample.sampled_at.to_rfc3339(),
+                sample.pid_count,
+                sample.rss_bytes,
+                sample.footprint_bytes,
+                sample.cpu_pct,
+                sample.pressure_state.to_string(),
+            ],
+        )?;
+        Ok(())
+    }
+
+    pub fn set_tool_path(&self, tool: &str, real_path: &str) -> anyhow::Result<()> {
+        let now = Utc::now().to_rfc3339();
+        self.conn.execute(
+            r#"
+            INSERT INTO tool_paths (tool, real_path, discovered_at, hash)
+            VALUES (?1, ?2, ?3, NULL)
+            ON CONFLICT(tool) DO UPDATE SET real_path = excluded.real_path,
+                                           discovered_at = excluded.discovered_at
+            "#,
+            params![tool, real_path, now],
+        )?;
+        Ok(())
+    }
+
+    pub fn list_tool_paths(&self) -> anyhow::Result<Vec<ToolPath>> {
+        let mut stmt = self
+            .conn
+            .prepare("SELECT tool, real_path, discovered_at, hash FROM tool_paths ORDER BY tool")?;
+        let rows = stmt.query_map([], |row| {
+            let discovered_at: String = row.get(2)?;
+            Ok(ToolPath {
+                tool: row.get(0)?,
+                real_path: row.get(1)?,
+                discovered_at: parse_ts_sql(&discovered_at)?,
+                hash: row.get(3)?,
+            })
+        })?;
+        rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
+    }
+}
+
+fn map_job(row: &rusqlite::Row<'_>) -> rusqlite::Result<JobRecord> {
+    let args_json: String = row.get("args_json")?;
+    let status: String = row.get("status")?;
+    let priority: String = row.get("priority")?;
+    let started_at: Option<String> = row.get("started_at")?;
+    let ended_at: Option<String> = row.get("ended_at")?;
+    let created_at: String = row.get("created_at")?;
+    let updated_at: String = row.get("updated_at")?;
+    Ok(JobRecord {
+        id: row.get("id")?,
+        adapter: row.get("adapter")?,
+        command: row.get("command")?,
+        args: serde_json::from_str(&args_json).map_err(to_sql_error)?,
+        cwd: row.get("cwd")?,
+        status: status.parse().map_err(string_sql_error)?,
+        priority: priority.parse().map_err(string_sql_error)?,
+        root_pid: row.get("root_pid")?,
+        pgid: row.get("pgid")?,
+        queued_reason: row.get("queued_reason")?,
+        started_at: parse_opt_ts(started_at)?,
+        ended_at: parse_opt_ts(ended_at)?,
+        exit_code: row.get("exit_code")?,
+        policy_snapshot_json: row.get("policy_snapshot_json")?,
+        created_at: parse_ts_sql(&created_at)?,
+        updated_at: parse_ts_sql(&updated_at)?,
+    })
+}
+
+fn opt_ts(ts: Option<DateTime<Utc>>) -> Option<String> {
+    ts.map(|ts| ts.to_rfc3339())
+}
+
+fn parse_opt_ts(value: Option<String>) -> rusqlite::Result<Option<DateTime<Utc>>> {
+    value.as_deref().map(parse_ts_sql).transpose()
+}
+
+fn parse_ts_sql(value: &str) -> rusqlite::Result<DateTime<Utc>> {
+    DateTime::parse_from_rfc3339(value)
+        .map(|dt| dt.with_timezone(&Utc))
+        .map_err(to_sql_error)
+}
+
+fn to_sql_error<E>(err: E) -> rusqlite::Error
+where
+    E: std::error::Error + Send + Sync + 'static,
+{
+    rusqlite::Error::ToSqlConversionFailure(Box::new(err))
+}
+
+fn string_sql_error(err: String) -> rusqlite::Error {
+    rusqlite::Error::FromSqlConversionFailure(
+        0,
+        rusqlite::types::Type::Text,
+        Box::new(std::io::Error::new(std::io::ErrorKind::InvalidData, err)),
+    )
+}
+
+pub fn event_payload(message: impl Into<String>) -> Value {
+    json!({ "message": message.into() })
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::tempdir;
+
+    #[test]
+    fn policy_load_save_defaults_round_trip() {
+        let dir = tempdir().unwrap();
+        let store = Store::open(dir.path().join("state.db")).unwrap();
+        let mut policy = store.load_policy().unwrap();
+        assert_eq!(policy.max_active_jobs, 1);
+        policy.set_key("max_active_jobs", "3").unwrap();
+        store.save_policy(&policy).unwrap();
+        assert_eq!(store.load_policy().unwrap().max_active_jobs, 3);
+    }
+}

package/crates/agentguard-core/src/types.rs

@@ -0,0 +1,248 @@
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::{fmt, str::FromStr};
+
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, PartialOrd, Ord)]
+#[serde(rename_all = "snake_case")]
+pub enum Priority {
+    Low,
+    #[default]
+    Normal,
+    High,
+}
+
+impl Priority {
+    pub fn scheduler_weight(self) -> i64 {
+        match self {
+            Priority::Low => 0,
+            Priority::Normal => 10,
+            Priority::High => 20,
+        }
+    }
+}
+
+impl fmt::Display for Priority {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            Priority::Low => "low",
+            Priority::Normal => "normal",
+            Priority::High => "high",
+        })
+    }
+}
+
+impl FromStr for Priority {
+    type Err = String;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        match value {
+            "low" => Ok(Self::Low),
+            "normal" => Ok(Self::Normal),
+            "high" => Ok(Self::High),
+            other => Err(format!("unknown priority {other:?}")),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum PressureState {
+    Normal,
+    Soft,
+    Critical,
+}
+
+impl fmt::Display for PressureState {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            PressureState::Normal => "normal",
+            PressureState::Soft => "soft",
+            PressureState::Critical => "critical",
+        })
+    }
+}
+
+impl FromStr for PressureState {
+    type Err = String;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        match value.trim().to_ascii_lowercase().as_str() {
+            "normal" | "green" => Ok(Self::Normal),
+            "soft" | "warning" | "yellow" => Ok(Self::Soft),
+            "critical" | "red" => Ok(Self::Critical),
+            other => Err(format!("unknown pressure state {other:?}")),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum JobState {
+    PendingAdmission,
+    Queued,
+    Starting,
+    Running,
+    PausedByGuard,
+    Exiting,
+    Exited,
+    Failed,
+    Lost,
+    Cancelled,
+}
+
+impl JobState {
+    pub fn is_terminal(self) -> bool {
+        matches!(
+            self,
+            JobState::Exited | JobState::Failed | JobState::Lost | JobState::Cancelled
+        )
+    }
+
+    pub fn counts_as_active(self) -> bool {
+        matches!(
+            self,
+            JobState::Starting | JobState::Running | JobState::PausedByGuard | JobState::Exiting
+        )
+    }
+}
+
+impl fmt::Display for JobState {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            JobState::PendingAdmission => "pending_admission",
+            JobState::Queued => "queued",
+            JobState::Starting => "starting",
+            JobState::Running => "running",
+            JobState::PausedByGuard => "paused_by_guard",
+            JobState::Exiting => "exiting",
+            JobState::Exited => "exited",
+            JobState::Failed => "failed",
+            JobState::Lost => "lost",
+            JobState::Cancelled => "cancelled",
+        })
+    }
+}
+
+impl FromStr for JobState {
+    type Err = String;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        match value {
+            "pending_admission" => Ok(Self::PendingAdmission),
+            "queued" => Ok(Self::Queued),
+            "starting" => Ok(Self::Starting),
+            "running" => Ok(Self::Running),
+            "paused_by_guard" | "paused" => Ok(Self::PausedByGuard),
+            "exiting" => Ok(Self::Exiting),
+            "exited" | "completed" => Ok(Self::Exited),
+            "failed" => Ok(Self::Failed),
+            "lost" => Ok(Self::Lost),
+            "cancelled" | "canceled" => Ok(Self::Cancelled),
+            other => Err(format!("unknown job state {other:?}")),
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Policy {
+    pub max_active_jobs: u32,
+    pub soft_memory_pressure_percent: u8,
+    pub critical_memory_pressure_percent: u8,
+    pub recovery_memory_pressure_percent: u8,
+    pub default_codex_max_threads: u32,
+    pub pause_enabled: bool,
+    pub queue_timeout_seconds: u64,
+    pub recovery_window_seconds: u64,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct JobRecord {
+    pub id: String,
+    pub adapter: String,
+    pub command: String,
+    pub args: Vec<String>,
+    pub cwd: String,
+    pub status: JobState,
+    pub priority: Priority,
+    pub root_pid: Option<i64>,
+    pub pgid: Option<i64>,
+    pub queued_reason: Option<String>,
+    pub started_at: Option<DateTime<Utc>>,
+    pub ended_at: Option<DateTime<Utc>>,
+    pub exit_code: Option<i32>,
+    pub policy_snapshot_json: String,
+    pub created_at: DateTime<Utc>,
+    pub updated_at: DateTime<Utc>,
+}
+
+impl JobRecord {
+    pub fn new_admission(
+        id: String,
+        adapter: String,
+        command: String,
+        args: Vec<String>,
+        cwd: String,
+        priority: Priority,
+        policy: &Policy,
+    ) -> anyhow::Result<Self> {
+        let now = Utc::now();
+        Ok(Self {
+            id,
+            adapter,
+            command,
+            args,
+            cwd,
+            status: JobState::PendingAdmission,
+            priority,
+            root_pid: None,
+            pgid: None,
+            queued_reason: None,
+            started_at: None,
+            ended_at: None,
+            exit_code: None,
+            policy_snapshot_json: serde_json::to_string(policy)?,
+            created_at: now,
+            updated_at: now,
+        })
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EventRecord {
+    pub id: i64,
+    pub job_id: Option<String>,
+    pub ts: DateTime<Utc>,
+    pub event_type: String,
+    pub payload_json: Value,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ToolPath {
+    pub tool: String,
+    pub real_path: String,
+    pub discovered_at: DateTime<Utc>,
+    pub hash: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ProcessSample {
+    pub job_id: String,
+    pub sampled_at: DateTime<Utc>,
+    pub pid_count: u32,
+    pub rss_bytes: u64,
+    pub footprint_bytes: u64,
+    pub cpu_pct: f64,
+    pub pressure_state: PressureState,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StatusSnapshot {
+    pub pressure: PressureState,
+    pub policy: Policy,
+    pub jobs: Vec<JobRecord>,
+    pub events: Vec<EventRecord>,
+    pub tool_paths: Vec<ToolPath>,
+    pub socket_path: String,
+    pub db_path: String,
+}