Skip to main content

rz_cli/
registry.rs

1//! Agent registry for discovery and routing.
2//!
3//! Persists agent entries to `~/.rz/registry.json` so any process
4//! can discover peers by name, transport, or capability.
5
6use eyre::{Context, Result};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::fs;
10use std::path::PathBuf;
11use std::time::{SystemTime, UNIX_EPOCH};
12
13/// A single registered agent.
14#[derive(Serialize, Deserialize, Clone, Debug)]
15pub struct AgentEntry {
16    pub name: String,
17    /// UUID or cmux surface ID.
18    pub id: String,
19    /// One of: `cmux`, `http`, `file`, `stdio`.
20    pub transport: String,
21    /// Surface ID for cmux, URL for http, mailbox path for file.
22    pub endpoint: String,
23    /// Optional tags like `["code","review","search"]`.
24    #[serde(default)]
25    pub capabilities: Vec<String>,
26    /// Unix epoch milliseconds when the agent first registered.
27    pub registered_at: u64,
28    /// Unix epoch milliseconds, updated by [`touch`].
29    pub last_seen: u64,
30}
31
32/// Return the path to the registry file (`~/.rz/registry.json`).
33pub fn registry_path() -> PathBuf {
34    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
35    PathBuf::from(home).join(".rz").join("registry.json")
36}
37
38/// Load the registry from disk. Returns an empty map if the file does not exist.
39pub fn load() -> Result<HashMap<String, AgentEntry>> {
40    let path = registry_path();
41    if !path.exists() {
42        return Ok(HashMap::new());
43    }
44    let data = fs::read_to_string(&path)
45        .wrap_err_with(|| format!("failed to read {}", path.display()))?;
46    let map: HashMap<String, AgentEntry> =
47        serde_json::from_str(&data).wrap_err("failed to parse registry.json")?;
48    Ok(map)
49}
50
51/// Atomically write the registry to disk (write-tmp then rename).
52pub fn save(registry: &HashMap<String, AgentEntry>) -> Result<()> {
53    let path = registry_path();
54    if let Some(parent) = path.parent() {
55        fs::create_dir_all(parent)
56            .wrap_err_with(|| format!("failed to create {}", parent.display()))?;
57    }
58    let json = serde_json::to_string_pretty(registry)
59        .wrap_err("failed to serialize registry")?;
60
61    // Atomic write: temp file in same dir, then rename.
62    let tmp = path.with_extension("json.tmp");
63    fs::write(&tmp, json.as_bytes())
64        .wrap_err_with(|| format!("failed to write {}", tmp.display()))?;
65    fs::rename(&tmp, &path)
66        .wrap_err_with(|| format!("failed to rename {} -> {}", tmp.display(), path.display()))?;
67    Ok(())
68}
69
70/// Register (or update) an agent entry.
71pub fn register(entry: AgentEntry) -> Result<()> {
72    let mut reg = load()?;
73    reg.insert(entry.name.clone(), entry);
74    save(&reg)
75}
76
77/// Remove an agent by name.
78pub fn deregister(name: &str) -> Result<()> {
79    let mut reg = load()?;
80    reg.remove(name);
81    save(&reg)
82}
83
84/// Look up an agent by name.
85pub fn lookup(name: &str) -> Result<Option<AgentEntry>> {
86    let reg = load()?;
87    Ok(reg.get(name).cloned())
88}
89
90/// Return all registered agents.
91pub fn list_all() -> Result<Vec<AgentEntry>> {
92    let reg = load()?;
93    Ok(reg.into_values().collect())
94}
95
96/// Remove entries whose `last_seen` is older than `max_age_secs` seconds ago.
97/// Returns the number of entries removed.
98pub fn cleanup_stale(max_age_secs: u64) -> Result<usize> {
99    let mut reg = load()?;
100    let now_ms = SystemTime::now()
101        .duration_since(UNIX_EPOCH)
102        .unwrap_or_default()
103        .as_millis() as u64;
104    let cutoff = now_ms.saturating_sub(max_age_secs * 1000);
105
106    let before = reg.len();
107    reg.retain(|_, entry| entry.last_seen >= cutoff);
108    let removed = before - reg.len();
109
110    if removed > 0 {
111        save(&reg)?;
112    }
113    Ok(removed)
114}
115
116/// Update `last_seen` to now for the given agent name.
117pub fn touch(name: &str) -> Result<()> {
118    let mut reg = load()?;
119    if let Some(entry) = reg.get_mut(name) {
120        entry.last_seen = SystemTime::now()
121            .duration_since(UNIX_EPOCH)
122            .unwrap_or_default()
123            .as_millis() as u64;
124        save(&reg)?;
125    }
126    Ok(())
127}