agent-manager/src/state.rs

//! The local installation database: a JSON file recording every agent
//! installed by agent-manager (name, version, paths, PID, dates).

use anyhow::{anyhow, Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

pub const STATE_VERSION: u32 = 1;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateFile {
    #[serde(default = "state_version")]
    pub version: u32,
    #[serde(default)]
    pub installed: BTreeMap<String, InstalledEntry>,
}

impl Default for StateFile {
    fn default() -> Self {
        Self {
            version: STATE_VERSION,
            installed: BTreeMap::new(),
        }
    }
}

fn state_version() -> u32 {
    STATE_VERSION
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InstalledEntry {
    pub name: String,
    #[serde(default)]
    pub version: Option<String>,
    /// Installer method that was used (npm, pip, binary, git, ...).
    pub method: String,
    /// Root directory of this agent's installation.
    pub install_dir: String,
    /// Executable files/shims that belong to this installation.
    #[serde(default)]
    pub bins: Vec<String>,
    /// Run command (as configured at install time).
    pub run: String,
    /// RFC 3339 timestamp of the installation.
    pub installed_at: String,
    /// RFC 3339 timestamp of the last update (if any).
    #[serde(default)]
    pub updated_at: Option<String>,
    /// PID of the background process (when running).
    #[serde(default)]
    pub pid: Option<u32>,
    /// RFC 3339 timestamp of the last background start.
    #[serde(default)]
    pub started_at: Option<String>,
}

#[derive(Debug, Clone)]
pub struct StateStore {
    path: PathBuf,
}

impl StateStore {
    pub fn new(path: PathBuf) -> Self {
        Self { path }
    }

    pub fn path(&self) -> &Path {
        &self.path
    }

    pub fn load(&self) -> Result<StateFile> {
        if !self.path.exists() {
            return Ok(StateFile::default());
        }
        let text = std::fs::read_to_string(&self.path)
            .with_context(|| format!("cannot read state file {}", self.path.display()))?;
        let sf: StateFile = serde_json::from_str(&text).map_err(|e| {
            anyhow!(
                "state file {} is corrupt ({e}); run 'am doctor --fix' to repair it",
                self.path.display()
            )
        })?;
        Ok(sf)
    }

    pub fn save(&self, sf: &StateFile) -> Result<()> {
        if let Some(parent) = self.path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let tmp = self.path.with_extension("json.tmp");
        let text = serde_json::to_string_pretty(sf)?;
        std::fs::write(&tmp, text)?;
        std::fs::rename(&tmp, &self.path)
            .with_context(|| format!("cannot write state file {}", self.path.display()))?;
        Ok(())
    }

    pub fn get(&self, name: &str) -> Result<Option<InstalledEntry>> {
        Ok(self.load()?.installed.get(name).cloned())
    }

    pub fn all(&self) -> Result<BTreeMap<String, InstalledEntry>> {
        Ok(self.load()?.installed)
    }

    pub fn set(&self, entry: &InstalledEntry) -> Result<()> {
        let mut sf = self.load()?;
        sf.installed.insert(entry.name.clone(), entry.clone());
        self.save(&sf)
    }

    pub fn remove(&self, name: &str) -> Result<bool> {
        let mut sf = self.load()?;
        let removed = sf.installed.remove(name).is_some();
        if removed {
            self.save(&sf)?;
        }
        Ok(removed)
    }

    /// Update the PID / start timestamp of a background agent.
    pub fn update_pid(&self, name: &str, pid: Option<u32>, started_at: Option<String>) -> Result<()> {
        let mut sf = self.load()?;
        if let Some(e) = sf.installed.get_mut(name) {
            e.pid = pid;
            e.started_at = started_at;
            self.save(&sf)?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn store(dir: &Path) -> StateStore {
        StateStore::new(dir.join("state.json"))
    }

    #[test]
    fn roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let st = store(dir.path());
        assert!(st.all().unwrap().is_empty());
        let entry = InstalledEntry {
            name: "claude-code".to_string(),
            version: Some("1.2.3".to_string()),
            method: "npm".to_string(),
            install_dir: "/tmp/cc".to_string(),
            bins: vec!["/tmp/cc/bin/claude".to_string()],
            run: "claude".to_string(),
            installed_at: "2026-01-01T00:00:00Z".to_string(),
            updated_at: None,
            pid: Some(42),
            started_at: None,
        };
        st.set(&entry).unwrap();
        let got = st.get("claude-code").unwrap().unwrap();
        assert_eq!(got.version.as_deref(), Some("1.2.3"));
        assert_eq!(got.pid, Some(42));
        st.update_pid("claude-code", None, None).unwrap();
        assert_eq!(st.get("claude-code").unwrap().unwrap().pid, None);
        assert!(st.remove("claude-code").unwrap());
        assert!(st.all().unwrap().is_empty());
    }

    #[test]
    fn corrupt_file_is_reported() {
        let dir = tempfile::tempdir().unwrap();
        let st = store(dir.path());
        std::fs::write(st.path(), "{ not json").unwrap();
        let err = st.load().unwrap_err().to_string();
        assert!(err.contains("corrupt"), "unexpected error: {err}");
    }
}