agentguard-local 0.1.0

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

@@ -0,0 +1,613 @@
+use crate::pressure::{MemoryReading, PressureHysteresis, classify_pressure_with_previous};
+use crate::process::{Signal, pid_alive, sample_process_group, send_signal_to_group};
+use crate::protocol::{DoctorReport, PROTOCOL_VERSION, Request, Response};
+use crate::scheduler::{AdmissionDecision, decide_admission, ordered_queue, pause_candidates};
+use crate::store::{Store, event_payload};
+use crate::types::{JobRecord, JobState, Policy, PressureState, Priority, StatusSnapshot};
+use serde_json::json;
+use std::fs;
+use std::io::{BufRead, BufReader, Write};
+use std::os::unix::fs::PermissionsExt;
+use std::os::unix::net::{UnixListener, UnixStream};
+use std::path::{Path, PathBuf};
+use std::process::Command;
+use uuid::Uuid;
+
+#[derive(Debug, Clone)]
+pub struct DaemonConfig {
+    pub socket_path: PathBuf,
+    pub db_path: PathBuf,
+    pub foreground: bool,
+}
+
+pub fn run_daemon(config: DaemonConfig) -> anyhow::Result<()> {
+    if config.foreground && std::env::var_os("AGENTGUARD_TEST_MODE").is_none() {
+        eprintln!(
+            "agentguardd: foreground mode, socket={}, db={}",
+            config.socket_path.display(),
+            config.db_path.display()
+        );
+    }
+    if let Some(parent) = config.socket_path.parent() {
+        fs::create_dir_all(parent)?;
+        fs::set_permissions(parent, fs::Permissions::from_mode(0o700))?;
+    }
+    if config.socket_path.exists() {
+        fs::remove_file(&config.socket_path)?;
+    }
+    let listener = UnixListener::bind(&config.socket_path)?;
+    fs::set_permissions(&config.socket_path, fs::Permissions::from_mode(0o600))?;
+
+    let mut daemon = Daemon::new(config)?;
+    daemon.reconcile_jobs()?;
+    for stream in listener.incoming() {
+        match stream {
+            Ok(stream) => {
+                if let Err(err) = daemon.handle_stream(stream) {
+                    eprintln!("agentguardd: request failed: {err:#}");
+                }
+            }
+            Err(err) => eprintln!("agentguardd: socket error: {err}"),
+        }
+    }
+    Ok(())
+}
+
+struct Daemon {
+    config: DaemonConfig,
+    store: Store,
+    forced_pressure: Option<PressureState>,
+    pressure_hysteresis: PressureHysteresis,
+}
+
+impl Daemon {
+    fn new(config: DaemonConfig) -> anyhow::Result<Self> {
+        let store = Store::open(&config.db_path)?;
+        let policy = store.load_policy()?;
+        Ok(Self {
+            config,
+            store,
+            forced_pressure: None,
+            pressure_hysteresis: PressureHysteresis::new(policy.recovery_window_seconds as u32),
+        })
+    }
+
+    fn handle_stream(&mut self, stream: UnixStream) -> anyhow::Result<()> {
+        let mut reader = BufReader::new(stream.try_clone()?);
+        let mut line = String::new();
+        reader.read_line(&mut line)?;
+        let response = match serde_json::from_str::<Request>(&line) {
+            Ok(request) => self.handle_request(request),
+            Err(err) => Ok(Response::Error {
+                protocol_version: PROTOCOL_VERSION,
+                message: format!("invalid request: {err}"),
+            }),
+        }?;
+        let mut stream = stream;
+        serde_json::to_writer(&mut stream, &response)?;
+        stream.write_all(b"\n")?;
+        stream.flush()?;
+        Ok(())
+    }
+
+    fn handle_request(&mut self, request: Request) -> anyhow::Result<Response> {
+        if request.version() != PROTOCOL_VERSION {
+            return Ok(Response::Error {
+                protocol_version: PROTOCOL_VERSION,
+                message: format!("unsupported protocol version {}", request.version()),
+            });
+        }
+
+        if let Request::TestSetPressure { pressure, .. } = request {
+            if std::env::var_os("AGENTGUARD_TEST_MODE").is_none() {
+                return Ok(Response::Error {
+                    protocol_version: PROTOCOL_VERSION,
+                    message: "test pressure override requires AGENTGUARD_TEST_MODE=1".to_string(),
+                });
+            }
+            self.forced_pressure = Some(pressure);
+            self.store.add_event(
+                None,
+                "test_pressure_override",
+                json!({ "pressure": pressure }),
+            )?;
+            return Ok(Response::Ack {
+                protocol_version: PROTOCOL_VERSION,
+                message: format!("pressure set to {pressure}"),
+            });
+        }
+
+        let pressure = self.tick()?;
+        match request {
+            Request::Admission {
+                adapter,
+                command,
+                args,
+                cwd,
+                priority,
+                tty: _,
+                ..
+            } => self.admission(adapter, command, args, cwd, priority, pressure),
+            Request::LeaseCheck { job_id, .. } => self.lease_check(&job_id, pressure),
+            Request::Register {
+                job_id,
+                root_pid,
+                pgid,
+                ..
+            } => {
+                self.store.register_job(&job_id, root_pid, pgid)?;
+                self.store.add_event(
+                    Some(&job_id),
+                    "job_registered",
+                    json!({ "root_pid": root_pid, "pgid": pgid }),
+                )?;
+                Ok(Response::Ack {
+                    protocol_version: PROTOCOL_VERSION,
+                    message: "registered".to_string(),
+                })
+            }
+            Request::Exit {
+                job_id, exit_code, ..
+            } => {
+                self.store.complete_job(&job_id, exit_code)?;
+                self.store.add_event(
+                    Some(&job_id),
+                    if exit_code == 0 {
+                        "job_exited"
+                    } else {
+                        "job_failed"
+                    },
+                    json!({ "exit_code": exit_code }),
+                )?;
+                Ok(Response::Ack {
+                    protocol_version: PROTOCOL_VERSION,
+                    message: "exit recorded".to_string(),
+                })
+            }
+            Request::Status {
+                include_processes, ..
+            } => Ok(Response::Status {
+                protocol_version: PROTOCOL_VERSION,
+                snapshot: self.status_snapshot(pressure, include_processes)?,
+            }),
+            Request::PolicyGet { .. } => Ok(Response::Policy {
+                protocol_version: PROTOCOL_VERSION,
+                policy: self.store.load_policy()?,
+            }),
+            Request::PolicySet { key, value, .. } => {
+                let mut policy = self.store.load_policy()?;
+                policy.set_key(&key, &value)?;
+                self.store.save_policy(&policy)?;
+                self.store.add_event(
+                    None,
+                    "policy_updated",
+                    json!({ "key": key, "value": value }),
+                )?;
+                Ok(Response::Policy {
+                    protocol_version: PROTOCOL_VERSION,
+                    policy,
+                })
+            }
+            Request::Pause { job_id, .. } => self.pause_job(&job_id, "user requested pause"),
+            Request::Resume { job_id, .. } => self.resume_job(&job_id, "user requested resume"),
+            Request::Cancel { job_id, .. } => self.cancel_job(&job_id),
+            Request::ToolPathSet {
+                tool, real_path, ..
+            } => {
+                self.store.set_tool_path(&tool, &real_path)?;
+                self.store.add_event(
+                    None,
+                    "tool_path_updated",
+                    json!({ "tool": tool, "real_path": real_path }),
+                )?;
+                Ok(Response::Ack {
+                    protocol_version: PROTOCOL_VERSION,
+                    message: "tool path recorded".to_string(),
+                })
+            }
+            Request::Doctor { .. } => Ok(Response::Doctor {
+                protocol_version: PROTOCOL_VERSION,
+                report: self.doctor_report()?,
+            }),
+            Request::TestSetPressure { .. } => unreachable!("handled before tick"),
+        }
+    }
+
+    fn admission(
+        &mut self,
+        adapter: String,
+        command: String,
+        args: Vec<String>,
+        cwd: String,
+        priority: Priority,
+        pressure: PressureState,
+    ) -> anyhow::Result<Response> {
+        let policy = self.store.load_policy()?;
+        let id = Uuid::new_v4().to_string();
+        let mut job =
+            JobRecord::new_admission(id.clone(), adapter, command, args, cwd, priority, &policy)?;
+        let decision = decide_admission(&policy, pressure, self.store.active_count()?);
+        match decision {
+            AdmissionDecision::Admit => {
+                job.status = JobState::Starting;
+                job.started_at = Some(chrono::Utc::now());
+                self.store.insert_job(&job)?;
+                self.store.add_event(
+                    Some(&id),
+                    "admission_granted",
+                    json!({ "pressure": pressure, "priority": priority }),
+                )?;
+                Ok(Response::Admission {
+                    protocol_version: PROTOCOL_VERSION,
+                    job_id: id,
+                    admitted: true,
+                    state: JobState::Starting.to_string(),
+                    reason: None,
+                    codex_max_threads: policy.default_codex_max_threads,
+                })
+            }
+            AdmissionDecision::Queue { reason } => {
+                job.status = JobState::Queued;
+                job.queued_reason = Some(reason.clone());
+                self.store.insert_job(&job)?;
+                self.store.add_event(
+                    Some(&id),
+                    "admission_queued",
+                    json!({ "pressure": pressure, "reason": reason }),
+                )?;
+                Ok(Response::Admission {
+                    protocol_version: PROTOCOL_VERSION,
+                    job_id: id,
+                    admitted: false,
+                    state: JobState::Queued.to_string(),
+                    reason: Some(reason),
+                    codex_max_threads: policy.default_codex_max_threads,
+                })
+            }
+        }
+    }
+
+    fn lease_check(&mut self, job_id: &str, pressure: PressureState) -> anyhow::Result<Response> {
+        let Some(job) = self.store.get_job(job_id)? else {
+            return Ok(Response::Error {
+                protocol_version: PROTOCOL_VERSION,
+                message: format!("unknown job {job_id}"),
+            });
+        };
+        if job.status != JobState::Queued {
+            return Ok(Response::Lease {
+                protocol_version: PROTOCOL_VERSION,
+                job_id: job_id.to_string(),
+                admitted: matches!(job.status, JobState::Starting | JobState::Running),
+                state: job.status.to_string(),
+                reason: job.queued_reason,
+            });
+        }
+
+        let policy = self.store.load_policy()?;
+        let queued = ordered_queue(
+            self.store.list_jobs_by_state(&[JobState::Queued])?,
+            chrono::Utc::now(),
+        );
+        let is_next = queued.first().is_some_and(|next| next.id == job_id);
+        let active_count = self.store.active_count()?;
+        let decision = if is_next {
+            decide_admission(&policy, pressure, active_count)
+        } else {
+            AdmissionDecision::Queue {
+                reason: "waiting behind higher-priority queued job".to_string(),
+            }
+        };
+        match decision {
+            AdmissionDecision::Admit => {
+                self.store
+                    .update_job_state(job_id, JobState::Starting, None)?;
+                self.store.add_event(
+                    Some(job_id),
+                    "queued_job_admitted",
+                    json!({ "pressure": pressure }),
+                )?;
+                Ok(Response::Lease {
+                    protocol_version: PROTOCOL_VERSION,
+                    job_id: job_id.to_string(),
+                    admitted: true,
+                    state: JobState::Starting.to_string(),
+                    reason: None,
+                })
+            }
+            AdmissionDecision::Queue { reason } => {
+                self.store
+                    .update_job_state(job_id, JobState::Queued, Some(&reason))?;
+                Ok(Response::Lease {
+                    protocol_version: PROTOCOL_VERSION,
+                    job_id: job_id.to_string(),
+                    admitted: false,
+                    state: JobState::Queued.to_string(),
+                    reason: Some(reason),
+                })
+            }
+        }
+    }
+
+    fn tick(&mut self) -> anyhow::Result<PressureState> {
+        let policy = self.store.load_policy()?;
+        self.pressure_hysteresis
+            .set_required_recovery_ticks(policy.recovery_window_seconds as u32);
+        let pressure = self.current_pressure(&policy);
+        if pressure == PressureState::Critical && policy.pause_enabled {
+            let jobs = self.store.list_jobs()?;
+            if let Some(job) = pause_candidates(&jobs).first()
+                && let Err(err) = self.pause_job(&job.id, "critical memory pressure")
+            {
+                self.store.add_event(
+                    Some(&job.id),
+                    "pause_failed",
+                    json!({ "error": err.to_string() }),
+                )?;
+            }
+        } else if pressure == PressureState::Normal {
+            let paused = self.store.list_jobs_by_state(&[JobState::PausedByGuard])?;
+            if let Some(job) = paused.first()
+                && let Err(err) = self.resume_job(&job.id, "memory pressure recovered")
+            {
+                self.store.add_event(
+                    Some(&job.id),
+                    "resume_failed",
+                    json!({ "error": err.to_string() }),
+                )?;
+            }
+        }
+        Ok(pressure)
+    }
+
+    fn current_pressure(&mut self, policy: &Policy) -> PressureState {
+        if let Some(pressure) = self.forced_pressure {
+            return self.pressure_hysteresis.observe(pressure);
+        }
+        if let Ok(path) = std::env::var("AGENTGUARD_FAKE_PRESSURE_FILE")
+            && let Ok(contents) = fs::read_to_string(path)
+            && let Ok(pressure) = contents.parse()
+        {
+            return self.pressure_hysteresis.observe(pressure);
+        }
+        if let Ok(value) = std::env::var("AGENTGUARD_FAKE_PRESSURE")
+            && let Ok(pressure) = value.parse()
+        {
+            return self.pressure_hysteresis.observe(pressure);
+        }
+        let reading = sample_system_memory_reading().unwrap_or_else(|| MemoryReading::quiet(0));
+        let classified =
+            classify_pressure_with_previous(policy, reading, self.pressure_hysteresis.state());
+        self.pressure_hysteresis.observe(classified)
+    }
+
+    fn pause_job(&mut self, job_id: &str, reason: &str) -> anyhow::Result<Response> {
+        let Some(job) = self.store.get_job(job_id)? else {
+            anyhow::bail!("unknown job {job_id}");
+        };
+        if let Some(pgid) = job.pgid {
+            send_signal_to_group(pgid, Signal::Stop)?;
+        }
+        self.store
+            .update_job_state(job_id, JobState::PausedByGuard, Some(reason))?;
+        self.store.add_event(
+            Some(job_id),
+            "job_paused",
+            json!({ "reason": reason, "pgid": job.pgid }),
+        )?;
+        Ok(Response::Ack {
+            protocol_version: PROTOCOL_VERSION,
+            message: "paused".to_string(),
+        })
+    }
+
+    fn resume_job(&mut self, job_id: &str, reason: &str) -> anyhow::Result<Response> {
+        let Some(job) = self.store.get_job(job_id)? else {
+            anyhow::bail!("unknown job {job_id}");
+        };
+        if let Some(pgid) = job.pgid {
+            send_signal_to_group(pgid, Signal::Continue)?;
+        }
+        self.store
+            .update_job_state(job_id, JobState::Running, None)?;
+        self.store.add_event(
+            Some(job_id),
+            "job_resumed",
+            json!({ "reason": reason, "pgid": job.pgid }),
+        )?;
+        Ok(Response::Ack {
+            protocol_version: PROTOCOL_VERSION,
+            message: "resumed".to_string(),
+        })
+    }
+
+    fn cancel_job(&mut self, job_id: &str) -> anyhow::Result<Response> {
+        let Some(job) = self.store.get_job(job_id)? else {
+            anyhow::bail!("unknown job {job_id}");
+        };
+        if let Some(pgid) = job.pgid
+            && matches!(job.status, JobState::Running | JobState::PausedByGuard)
+        {
+            send_signal_to_group(pgid, Signal::Terminate)?;
+        }
+        self.store
+            .update_job_state(job_id, JobState::Cancelled, Some("user requested cancel"))?;
+        self.store
+            .add_event(Some(job_id), "job_cancelled", event_payload("cancelled"))?;
+        Ok(Response::Ack {
+            protocol_version: PROTOCOL_VERSION,
+            message: "cancelled".to_string(),
+        })
+    }
+
+    fn status_snapshot(
+        &mut self,
+        pressure: PressureState,
+        include_processes: bool,
+    ) -> anyhow::Result<StatusSnapshot> {
+        if include_processes {
+            for job in self
+                .store
+                .list_jobs_by_state(&[JobState::Running, JobState::PausedByGuard])?
+            {
+                if let Some(pgid) = job.pgid {
+                    let sample = sample_process_group(&job.id, pgid, pressure);
+                    self.store.record_sample(&sample)?;
+                }
+            }
+        }
+        Ok(StatusSnapshot {
+            pressure,
+            policy: self.store.load_policy()?,
+            jobs: self.store.list_jobs()?,
+            events: self.store.recent_events(25)?,
+            tool_paths: self.store.list_tool_paths()?,
+            socket_path: self.config.socket_path.display().to_string(),
+            db_path: self.config.db_path.display().to_string(),
+        })
+    }
+
+    fn reconcile_jobs(&mut self) -> anyhow::Result<()> {
+        for job in self.store.list_jobs()? {
+            if job.status.is_terminal() || job.status == JobState::Queued {
+                continue;
+            }
+            let reconciled = job.root_pid.is_some_and(pid_alive);
+            if reconciled {
+                self.store.add_event(
+                    Some(&job.id),
+                    "daemon_restart_reconciled",
+                    json!({ "root_pid": job.root_pid, "pgid": job.pgid }),
+                )?;
+            } else {
+                self.store.update_job_state(
+                    &job.id,
+                    JobState::Lost,
+                    Some("daemon restarted and process could not be reconciled"),
+                )?;
+                self.store.add_event(
+                    Some(&job.id),
+                    "daemon_restart_lost",
+                    json!({ "root_pid": job.root_pid, "pgid": job.pgid }),
+                )?;
+            }
+        }
+        Ok(())
+    }
+
+    fn doctor_report(&self) -> anyhow::Result<DoctorReport> {
+        let plist = home_dir()
+            .map(|home| home.join("Library/LaunchAgents/dev.agentguard.agentguardd.plist"))
+            .unwrap_or_default();
+        Ok(DoctorReport {
+            os_supported: std::env::consts::OS == "macos",
+            os: format!("{} {}", std::env::consts::OS, std::env::consts::ARCH),
+            socket_path: self.config.socket_path.display().to_string(),
+            socket_exists: self.config.socket_path.exists(),
+            db_path: self.config.db_path.display().to_string(),
+            db_exists: self.config.db_path.exists(),
+            launchd_plist_exists: plist.exists(),
+            discovered_tools: self
+                .store
+                .list_tool_paths()?
+                .into_iter()
+                .map(|tool| (tool.tool, tool.real_path))
+                .collect(),
+        })
+    }
+}
+
+fn sample_system_memory_reading() -> Option<MemoryReading> {
+    let total_bytes = sysctl_u64("hw.memsize")?;
+    let Ok(output) = Command::new("/usr/bin/vm_stat").output() else {
+        return None;
+    };
+    if !output.status.success() {
+        return None;
+    }
+    let text = String::from_utf8_lossy(&output.stdout);
+    let (page_size, available_pages) = parse_vm_stat(&text)?;
+    let available_bytes = available_pages.saturating_mul(page_size);
+    let used_percent =
+        100_u64.saturating_sub(available_bytes.saturating_mul(100) / total_bytes) as u8;
+    Some(MemoryReading::quiet(used_percent))
+}
+
+fn sysctl_u64(name: &str) -> Option<u64> {
+    let output = Command::new("/usr/sbin/sysctl")
+        .args(["-n", name])
+        .output()
+        .ok()?;
+    if !output.status.success() {
+        return None;
+    }
+    String::from_utf8_lossy(&output.stdout).trim().parse().ok()
+}
+
+fn parse_vm_stat(text: &str) -> Option<(u64, u64)> {
+    let page_size = text
+        .lines()
+        .next()?
+        .split("page size of ")
+        .nth(1)?
+        .split_whitespace()
+        .next()?
+        .parse()
+        .ok()?;
+    let mut available = 0_u64;
+    for line in text.lines() {
+        if line.starts_with("Pages free:")
+            || line.starts_with("Pages speculative:")
+            || line.starts_with("Pages inactive:")
+            || line.starts_with("Pages purgeable:")
+        {
+            available += digits(line)?.parse::<u64>().ok()?;
+        }
+    }
+    Some((page_size, available))
+}
+
+fn digits(line: &str) -> Option<String> {
+    let digits: String = line.chars().filter(|ch| ch.is_ascii_digit()).collect();
+    if digits.is_empty() {
+        None
+    } else {
+        Some(digits)
+    }
+}
+
+fn home_dir() -> Option<PathBuf> {
+    std::env::var_os("HOME").map(PathBuf::from)
+}
+
+pub fn default_socket_path() -> PathBuf {
+    home_dir()
+        .unwrap_or_else(|| Path::new(".").to_path_buf())
+        .join(".agentguard/agentguard.sock")
+}
+
+pub fn default_db_path() -> PathBuf {
+    home_dir()
+        .unwrap_or_else(|| Path::new(".").to_path_buf())
+        .join(".agentguard/state.db")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn vm_stat_parser_counts_reclaimable_pages_as_available() {
+        let text = "\
+Mach Virtual Memory Statistics: (page size of 16384 bytes)
+Pages free:                               10.
+Pages active:                             90.
+Pages inactive:                           20.
+Pages speculative:                        30.
+Pages throttled:                          0.
+Pages wired down:                         40.
+Pages purgeable:                          5.
+";
+        assert_eq!(parse_vm_stat(text), Some((16_384, 65)));
+    }
+}