lean_ctx/core/agents/
persistence.rs

1use super::AgentRegistry;
2use crate::ipc::process::ProcessIdentity;
3use std::collections::HashMap;
4use std::path::PathBuf;
5
6/// Compatibility store for process identities.
7///
8/// Older MCP processes deserialize and re-serialize `registry.json` with an
9/// older `AgentEntry` schema. Serde then drops fields they do not know about.
10/// Keeping the PID-reuse proof in its own file makes a rolling upgrade safe:
11/// an old process can still update its legacy registry record without erasing
12/// the immutable identity captured by a new binary.
13#[derive(Debug, Default, serde::Serialize, serde::Deserialize)]
14pub struct ProcessIdentityIndex {
15    #[serde(default)]
16    identities: HashMap<String, ProcessIdentity>,
17}
18
19impl ProcessIdentityIndex {
20    pub fn load() -> Self {
21        let Ok(path) = identity_index_path() else {
22            return Self::default();
23        };
24        std::fs::read_to_string(path)
25            .ok()
26            .and_then(|content| serde_json::from_str(&content).ok())
27            .unwrap_or_default()
28    }
29
30    pub fn get(&self, agent_id: &str) -> Option<&ProcessIdentity> {
31        self.identities.get(agent_id)
32    }
33
34    pub(crate) fn insert(&mut self, agent_id: &str, identity: &ProcessIdentity) -> bool {
35        self.identities
36            .insert(agent_id.to_string(), identity.clone())
37            .as_ref()
38            != Some(identity)
39    }
40
41    pub(crate) fn retain_agents(&mut self, agents: &AgentRegistry) -> bool {
42        let before = self.identities.len();
43        self.identities.retain(|agent_id, _| {
44            agents
45                .agents
46                .iter()
47                .any(|agent| agent.agent_id == *agent_id)
48        });
49        self.identities.len() != before
50    }
51
52    pub(crate) fn save(&self) -> Result<(), String> {
53        let path = identity_index_path()?;
54        let json = serde_json::to_string_pretty(self).map_err(|error| error.to_string())?;
55        crate::config_io::write_atomic(&path, &json)
56            .map_err(|error| format!("persist process identity index {}: {error}", path.display()))
57    }
58}
59
60pub(super) fn agents_dir() -> Result<PathBuf, String> {
61    let dir = crate::core::data_dir::lean_ctx_data_dir()?;
62    Ok(dir.join("agents"))
63}
64
65fn identity_index_path() -> Result<PathBuf, String> {
66    Ok(agents_dir()?.join("process-identities.json"))
67}
68
69/// Create the agent-registry directory, naming the path and the operation on
70/// failure (#1619). A bare `File exists (os error 17)` or `Operation not
71/// permitted` tells a user nothing they can act on; sandboxes and filesystem
72/// policies deny exactly one path, and that path is the whole diagnosis.
73pub(super) fn create_agents_dir(dir: &std::path::Path) -> Result<(), String> {
74    std::fs::create_dir_all(dir)
75        .map_err(|error| format!("create agent registry directory {}: {error}", dir.display()))
76}
77
78pub(super) fn mutate_persistent<T>(
79    mutate: impl FnOnce(&mut AgentRegistry) -> T,
80) -> Result<T, String> {
81    let dir = agents_dir()?;
82    create_agents_dir(&dir)?;
83    let _lock = FileLock::acquire(&dir.join("registry.lock"))?;
84    let path = dir.join("registry.json");
85    let mut registry = load_registry_file(&path)?.unwrap_or_default();
86    let result = mutate(&mut registry);
87    save_registry_file(&path, &registry)?;
88    Ok(result)
89}
90
91/// A missing registry is the normal first-run case. A malformed registry is a
92/// data-integrity failure, not an empty bus: callers that mutate must fail
93/// closed instead of overwriting messages and presence with a blank snapshot.
94pub(super) fn load_registry_file(path: &std::path::Path) -> Result<Option<AgentRegistry>, String> {
95    match std::fs::read_to_string(path) {
96        Ok(content) => serde_json::from_str(&content)
97            .map(Some)
98            .map_err(|error| format!("agent registry is corrupt at {}: {error}", path.display())),
99        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
100        Err(error) => Err(format!("read agent registry {}: {error}", path.display())),
101    }
102}
103
104pub(super) fn save_registry_file(
105    path: &std::path::Path,
106    registry: &AgentRegistry,
107) -> Result<(), String> {
108    let json = serde_json::to_string_pretty(registry).map_err(|error| error.to_string())?;
109    crate::config_io::write_atomic(path, &json)
110        .map_err(|error| format!("persist agent registry {}: {error}", path.display()))
111}
112
113pub(super) fn generate_short_id() -> String {
114    uuid::Uuid::new_v4().simple().to_string()
115}
116
117pub struct FileLock {
118    file: std::fs::File,
119}
120
121/// How long [`FileLock::acquire`] waits. The lock is polled (25 ms steps),
122/// not a fair queue (GitLab #1295): under load a waiter behind writers that
123/// take it back to back could pass the former 750 ms and fail its tool call.
124/// A dead holder's lock is released by the OS, so waiting longer never makes
125/// a stale lock permanent; it only lets a contended call through.
126pub(crate) const ACQUIRE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
127
128impl FileLock {
129    pub fn acquire(path: &std::path::Path) -> Result<Self, String> {
130        Self::acquire_with_timeout(path, ACQUIRE_TIMEOUT)
131    }
132
133    pub(crate) fn acquire_with_timeout(
134        path: &std::path::Path,
135        timeout: std::time::Duration,
136    ) -> Result<Self, String> {
137        use fs2::FileExt;
138
139        const RETRY: std::time::Duration = std::time::Duration::from_millis(25);
140
141        let mut options = std::fs::OpenOptions::new();
142        options.create(true).truncate(false).read(true).write(true);
143        #[cfg(unix)]
144        {
145            use std::os::unix::fs::OpenOptionsExt;
146            options
147                .custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC)
148                .mode(0o600);
149        }
150        let file = options
151            .open(path)
152            .map_err(|error| format!("lock {}: {error}", path.display()))?;
153        let deadline = std::time::Instant::now() + timeout;
154        loop {
155            match file.try_lock_exclusive() {
156                Ok(()) => return Ok(Self { file }),
157                Err(error)
158                    if crate::core::file_lock::is_contended(&error)
159                        && std::time::Instant::now() < deadline =>
160                {
161                    std::thread::sleep(RETRY);
162                }
163                Err(error) if crate::core::file_lock::is_contended(&error) => {
164                    return Err(format!(
165                        "lock {} timed out after {}ms",
166                        path.display(),
167                        timeout.as_millis()
168                    ));
169                }
170                Err(error) => return Err(format!("lock {}: {error}", path.display())),
171            }
172        }
173    }
174}
175
176impl Drop for FileLock {
177    fn drop(&mut self) {
178        let _ = self.file.unlock();
179    }
180}
181
182#[cfg(test)]
183mod lock_tests {
184    use super::{ACQUIRE_TIMEOUT, FileLock};
185    use std::time::{Duration, Instant};
186
187    /// Holds the lock at `path` for `hold`, signalling once it holds it.
188    fn hold_for(path: &std::path::Path, hold: Duration) -> std::thread::JoinHandle<()> {
189        let path = path.to_path_buf();
190        let (held, wait) = std::sync::mpsc::channel();
191        let holder = std::thread::spawn(move || {
192            let _lock = FileLock::acquire_with_timeout(&path, Duration::from_secs(1)).unwrap();
193            held.send(()).unwrap();
194            std::thread::sleep(hold);
195        });
196        wait.recv().unwrap();
197        holder
198    }
199
200    /// A waiter behind two seconds of contention gets the lock (it failed
201    /// after 750 ms before, GitLab #1295).
202    #[test]
203    fn a_waiter_behind_two_seconds_of_contention_gets_the_lock() {
204        let dir = tempfile::tempdir().unwrap();
205        let path = dir.path().join("registry.lock");
206        let holder = hold_for(&path, Duration::from_secs(2));
207        let started = Instant::now();
208        let lock = FileLock::acquire(&path);
209        assert!(lock.is_ok(), "{:?}", lock.err());
210        assert!(started.elapsed() >= Duration::from_millis(1500));
211        holder.join().unwrap();
212    }
213
214    /// A waiter behind a holder that keeps the lock past the timeout gets a
215    /// clear error naming the lock and the wait.
216    #[test]
217    fn a_waiter_behind_a_longer_holder_gets_a_clear_error() {
218        let dir = tempfile::tempdir().unwrap();
219        let path = dir.path().join("registry.lock");
220        let holder = hold_for(&path, ACQUIRE_TIMEOUT + Duration::from_secs(1));
221        let error = FileLock::acquire(&path).err().expect("timed out");
222        assert_eq!(
223            error,
224            format!(
225                "lock {} timed out after {}ms",
226                path.display(),
227                ACQUIRE_TIMEOUT.as_millis()
228            )
229        );
230        assert_eq!(ACQUIRE_TIMEOUT, Duration::from_secs(5));
231        holder.join().unwrap();
232    }
233}
lean_ctx/core/agents/persistence.rs

lean_ctx/core/agents/
persistence.rs
