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    // Dropped after `Drop` unlocked `file`: the next thread of this process
120    // in line only starts once the file lock is free.
121    _turn: crate::core::file_lock::Turn,
122}
123
124/// How long [`FileLock::acquire`] waits without progress. Threads of this
125/// process queue for a [`crate::core::file_lock::Turn`] first, and that wait
126/// restarts at every handover, so a moving queue never times out; the file
127/// lock itself is then polled (25 ms steps) for at most this long, which
128/// bounds other processes. Before (GitLab #1295) a waiter behind writers that
129/// took it back to back could pass the former 750 ms and fail its tool call.
130/// A dead holder's lock is released by the OS, so waiting longer never makes
131/// a stale lock permanent; it only lets a contended call through.
132pub const ACQUIRE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);
133
134impl FileLock {
135    pub fn acquire(path: &std::path::Path) -> Result<Self, String> {
136        Self::acquire_with_timeout(path, ACQUIRE_TIMEOUT)
137    }
138
139    pub(crate) fn acquire_with_timeout(
140        path: &std::path::Path,
141        timeout: std::time::Duration,
142    ) -> Result<Self, String> {
143        use fs2::FileExt;
144
145        const RETRY: std::time::Duration = std::time::Duration::from_millis(25);
146
147        // Threads of this process queue in arrival order and wait only as
148        // long as the queue stands still; other processes are bounded by
149        // the polled file lock below.
150        let turn = crate::core::file_lock::take_turn(path, timeout).ok_or_else(|| {
151            format!(
152                "lock {} timed out after {}ms",
153                path.display(),
154                timeout.as_millis()
155            )
156        })?;
157        let mut options = std::fs::OpenOptions::new();
158        options.create(true).truncate(false).read(true).write(true);
159        #[cfg(unix)]
160        {
161            use std::os::unix::fs::OpenOptionsExt;
162            options
163                .custom_flags(libc::O_NOFOLLOW | libc::O_CLOEXEC)
164                .mode(0o600);
165        }
166        let file = options
167            .open(path)
168            .map_err(|error| format!("lock {}: {error}", path.display()))?;
169        let deadline = std::time::Instant::now() + timeout;
170        loop {
171            match file.try_lock_exclusive() {
172                Ok(()) => {
173                    return Ok(Self { file, _turn: turn });
174                }
175                Err(error)
176                    if crate::core::file_lock::is_contended(&error)
177                        && std::time::Instant::now() < deadline =>
178                {
179                    std::thread::sleep(RETRY);
180                }
181                Err(error) if crate::core::file_lock::is_contended(&error) => {
182                    return Err(format!(
183                        "lock {} timed out after {}ms",
184                        path.display(),
185                        timeout.as_millis()
186                    ));
187                }
188                Err(error) => return Err(format!("lock {}: {error}", path.display())),
189            }
190        }
191    }
192}
193
194impl Drop for FileLock {
195    fn drop(&mut self) {
196        let _ = self.file.unlock();
197    }
198}
199
200#[cfg(test)]
201mod lock_tests {
202    use super::{ACQUIRE_TIMEOUT, FileLock};
203    use std::time::{Duration, Instant};
204
205    /// Holds the lock at `path` for `hold`, signalling once it holds it.
206    fn hold_for(path: &std::path::Path, hold: Duration) -> std::thread::JoinHandle<()> {
207        let path = path.to_path_buf();
208        let (held, wait) = std::sync::mpsc::channel();
209        let holder = std::thread::spawn(move || {
210            let _lock = FileLock::acquire_with_timeout(&path, Duration::from_secs(1)).unwrap();
211            held.send(()).unwrap();
212            std::thread::sleep(hold);
213        });
214        wait.recv().unwrap();
215        holder
216    }
217
218    /// A waiter behind two seconds of contention gets the lock (it failed
219    /// after 750 ms before, GitLab #1295).
220    #[test]
221    fn a_waiter_behind_two_seconds_of_contention_gets_the_lock() {
222        let dir = tempfile::tempdir().unwrap();
223        let path = dir.path().join("registry.lock");
224        let holder = hold_for(&path, Duration::from_secs(2));
225        let started = Instant::now();
226        let lock = FileLock::acquire(&path);
227        assert!(lock.is_ok(), "{:?}", lock.err());
228        assert!(started.elapsed() >= Duration::from_millis(1500));
229        holder.join().unwrap();
230    }
231
232    /// A waiter behind a holder that keeps the lock past the timeout gets a
233    /// clear error naming the lock and the wait.
234    #[test]
235    fn a_waiter_behind_a_longer_holder_gets_a_clear_error() {
236        let dir = tempfile::tempdir().unwrap();
237        let path = dir.path().join("registry.lock");
238        let holder = hold_for(&path, ACQUIRE_TIMEOUT + Duration::from_secs(1));
239        let error = FileLock::acquire(&path).err().expect("timed out");
240        assert_eq!(
241            error,
242            format!(
243                "lock {} timed out after {}ms",
244                path.display(),
245                ACQUIRE_TIMEOUT.as_millis()
246            )
247        );
248        assert_eq!(ACQUIRE_TIMEOUT, Duration::from_secs(5));
249        holder.join().unwrap();
250    }
251
252    /// Threads that take the lock back to back fail none of their waiters as
253    /// long as the lock changes hands within the timeout: the polled file lock
254    /// counted every holder ahead of a waiter against its timeout, so the
255    /// eighth of eight publishers whose writes took a second each failed after
256    /// five seconds (#2214, `knowledge_bridge` on Windows).
257    #[test]
258    fn a_queue_that_keeps_moving_fails_no_waiter() {
259        const THREADS: usize = 8;
260        let dir = tempfile::tempdir().unwrap();
261        let path = dir.path().join("bridge.lock");
262        let timeout = Duration::from_millis(400);
263        let start = std::sync::Arc::new(std::sync::Barrier::new(THREADS));
264        let waiters: Vec<_> = (0..THREADS)
265            .map(|_| {
266                let (path, start) = (path.clone(), std::sync::Arc::clone(&start));
267                std::thread::spawn(move || {
268                    start.wait();
269                    let lock = FileLock::acquire_with_timeout(&path, timeout)?;
270                    std::thread::sleep(Duration::from_millis(100));
271                    drop(lock);
272                    Ok::<(), String>(())
273                })
274            })
275            .collect();
276        for waiter in waiters {
277            assert_eq!(waiter.join().unwrap(), Ok(()));
278        }
279    }
280}
lean_ctx/core/agents/persistence.rs

lean_ctx/core/agents/
persistence.rs
