Skip to main content

agentbridge/
local_registry.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// SPDX-License-Identifier: Apache-2.0
3
4//! On-host registry of agentbridge listeners started by `register`.
5//!
6//! `list --local` cannot ask a dataplane SLIM client for "who is listening"
7//! (that needs the controller channel). Instead each `register --slim-endpoint`
8//! or `--a2a-listen` writes a lease under `$SHADI_TMP_DIR/agentbridge-local`
9//! (or the process temp dir). `list --local` reads those files and drops any
10//! whose pid is gone, so a killed listener does not stay listed.
11
12use std::fs;
13use std::path::{Path, PathBuf};
14
15use serde::{Deserialize, Serialize};
16use shadi_a2a::A2ABinding;
17
18use crate::member_source::{CandidateMember, MemberSource};
19
20const REGISTRY_DIRNAME: &str = "agentbridge-local";
21
22/// One locally registered listener (SLIM and/or official A2A unicast).
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24pub struct LocalAdapterRecord {
25    pub name: String,
26    pub did: String,
27    #[serde(default)]
28    pub slim_endpoint: String,
29    /// Unicast A2A URL (`http://127.0.0.1:port`), empty when SLIM-only.
30    #[serde(default)]
31    pub a2a_url: String,
32    /// Official binding for [`Self::a2a_url`]. Missing on old leases → gRPC.
33    #[serde(default)]
34    pub a2a_binding: A2ABinding,
35    pub pid: u32,
36}
37
38impl LocalAdapterRecord {
39    pub fn to_candidate(&self) -> CandidateMember {
40        CandidateMember {
41            name: self.name.clone(),
42            did: self.did.clone(),
43            slim_endpoint: if self.slim_endpoint.is_empty() {
44                None
45            } else {
46                Some(self.slim_endpoint.clone())
47            },
48            a2a_url: if self.a2a_url.is_empty() {
49                None
50            } else {
51                Some(self.a2a_url.clone())
52            },
53            a2a_binding: if self.a2a_url.is_empty() {
54                None
55            } else {
56                Some(self.a2a_binding)
57            },
58        }
59    }
60}
61
62/// Directory of listener lease files.
63#[derive(Debug, Clone)]
64pub struct LocalAdapterRegistry {
65    dir: PathBuf,
66}
67
68/// Removes the lease file when the registering process exits cleanly.
69pub struct LocalAdapterLease {
70    path: PathBuf,
71}
72
73impl Drop for LocalAdapterLease {
74    fn drop(&mut self) {
75        let _ = fs::remove_file(&self.path);
76    }
77}
78
79impl LocalAdapterRegistry {
80    /// `$SHADI_TMP_DIR/agentbridge-local`, else `<temp>/agentbridge-local`.
81    /// Prefer `SHADI_TMP_DIR` so a sandboxed `register` can write the lease
82    /// under the same `--write` root as its mTLS material.
83    pub fn from_env() -> Self {
84        let root = std::env::var_os("SHADI_TMP_DIR")
85            .map(PathBuf::from)
86            .unwrap_or_else(std::env::temp_dir);
87        Self::with_dir(root.join(REGISTRY_DIRNAME))
88    }
89
90    pub fn with_dir(dir: PathBuf) -> Self {
91        Self { dir }
92    }
93
94    pub fn dir(&self) -> &Path {
95        &self.dir
96    }
97
98    /// Write a lease for `record`. Replaces a previous lease for the same
99    /// name+pid. The returned guard deletes the file on drop.
100    pub fn publish(&self, record: &LocalAdapterRecord) -> Result<LocalAdapterLease, String> {
101        if !is_safe_agent_name(&record.name) {
102            return Err(format!("unsafe agent name '{}'", record.name));
103        }
104        if record.did.is_empty() || (record.slim_endpoint.is_empty() && record.a2a_url.is_empty()) {
105            return Err("local adapter record needs a DID and slim_endpoint or a2a_url".to_string());
106        }
107        fs::create_dir_all(&self.dir).map_err(|err| {
108            format!(
109                "create local adapter registry {}: {err}",
110                self.dir.display()
111            )
112        })?;
113        #[cfg(unix)]
114        {
115            use std::os::unix::fs::PermissionsExt;
116            let _ = fs::set_permissions(&self.dir, fs::Permissions::from_mode(0o700));
117        }
118        let path = self.dir.join(record_filename(&record.name, record.pid)?);
119        let body = serde_json::to_vec_pretty(record)
120            .map_err(|err| format!("serialize local adapter record: {err}"))?;
121        fs::write(&path, body)
122            .map_err(|err| format!("write local adapter lease {}: {err}", path.display()))?;
123        Ok(LocalAdapterLease { path })
124    }
125
126    /// Live listeners only: skip unreadable files and delete leases whose
127    /// process is gone.
128    pub fn list_live(&self) -> Result<Vec<LocalAdapterRecord>, String> {
129        let mut records = Vec::new();
130        let entries = match fs::read_dir(&self.dir) {
131            Ok(entries) => entries,
132            Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(records),
133            Err(err) => {
134                return Err(format!(
135                    "read local adapter registry {}: {err}",
136                    self.dir.display()
137                ));
138            }
139        };
140        for entry in entries {
141            let entry = entry.map_err(|err| format!("read registry entry: {err}"))?;
142            let path = entry.path();
143            if !path.is_file() {
144                continue;
145            }
146            let Some(name) = path.file_name().and_then(|n| n.to_str()) else {
147                continue;
148            };
149            if !looks_like_lease_filename(name) {
150                continue;
151            }
152            let Ok(bytes) = fs::read(&path) else {
153                continue;
154            };
155            let Ok(record) = serde_json::from_slice::<LocalAdapterRecord>(&bytes) else {
156                continue;
157            };
158            if !is_safe_agent_name(&record.name) || record.did.is_empty() {
159                let _ = fs::remove_file(&path);
160                continue;
161            }
162            if !pid_is_alive(record.pid) {
163                let _ = fs::remove_file(&path);
164                continue;
165            }
166            records.push(record);
167        }
168        records.sort_by(|a, b| a.name.cmp(&b.name).then(a.pid.cmp(&b.pid)));
169        Ok(records)
170    }
171
172    /// Live listeners whose DID or adapter name matches `query`.
173    ///
174    /// DID is the portable name. A tool name (`copilot`) is a local alias for
175    /// whatever DID currently holds that lease. Two live records with the
176    /// same DID or the same name are returned so the caller can treat them
177    /// as ambiguous.
178    pub fn find_live(&self, query: &str) -> Result<Vec<LocalAdapterRecord>, String> {
179        let records = self.list_live()?;
180        if let Some(did) = crate::member_source::parse_peer_did(query) {
181            Ok(records.into_iter().filter(|r| r.did == did).collect())
182        } else {
183            Ok(records
184                .into_iter()
185                .filter(|r| crate::member_source::matches_local_alias(&r.name, query))
186                .collect())
187        }
188    }
189
190    /// Unique live listener for `query`, or an error if none / more than one.
191    pub fn resolve_live(&self, query: &str) -> Result<LocalAdapterRecord, String> {
192        let matches = self.find_live(query)?;
193        match matches.len() {
194            0 => Err(format!(
195                "no live agentbridge adapter matching '{query}'. \
196                 Address the agent by DID (`did:key:…`) after `list --local`, \
197                 or pass --a2a-url only as a locator override with --to <did>"
198            )),
199            1 => Ok(matches.into_iter().next().expect("len 1")),
200            n => Err(format!(
201                "ambiguous: {n} live adapters match '{query}'. \
202                 Two agents can share a URL; use the DID as the name"
203            )),
204        }
205    }
206}
207
208/// [`MemberSource`] over the on-host lease directory.
209pub struct LocalRegistrySource {
210    pub registry: LocalAdapterRegistry,
211}
212
213impl MemberSource for LocalRegistrySource {
214    fn resolve(&self) -> Result<Vec<CandidateMember>, String> {
215        Ok(self
216            .registry
217            .list_live()?
218            .into_iter()
219            .map(|record| record.to_candidate())
220            .collect())
221    }
222}
223
224fn is_safe_agent_name(name: &str) -> bool {
225    !name.is_empty()
226        && name.len() <= 64
227        && !name.starts_with('-')
228        && name
229            .chars()
230            .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
231}
232
233fn record_filename(name: &str, pid: u32) -> Result<String, String> {
234    if !is_safe_agent_name(name) {
235        return Err(format!("unsafe agent name '{name}'"));
236    }
237    Ok(format!("{name}-{pid}.json"))
238}
239
240fn looks_like_lease_filename(name: &str) -> bool {
241    let Some(stem) = name.strip_suffix(".json") else {
242        return false;
243    };
244    let Some((agent, pid)) = stem.rsplit_once('-') else {
245        return false;
246    };
247    is_safe_agent_name(agent) && pid.chars().all(|c| c.is_ascii_digit())
248}
249
250#[cfg(unix)]
251fn pid_is_alive(pid: u32) -> bool {
252    if pid == 0 {
253        return false;
254    }
255    // SAFETY: signal 0 only probes existence; ESRCH means gone, EPERM means
256    // the pid exists but we cannot signal it (still live).
257    let rc = unsafe { libc::kill(pid as libc::pid_t, 0) };
258    rc == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
259}
260
261#[cfg(windows)]
262fn pid_is_alive(pid: u32) -> bool {
263    use windows_sys::Win32::Foundation::{
264        CloseHandle, GetLastError, ERROR_ACCESS_DENIED, STILL_ACTIVE,
265    };
266    use windows_sys::Win32::System::Threading::{
267        GetExitCodeProcess, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION,
268    };
269
270    if pid == 0 {
271        return false;
272    }
273    // SAFETY: OpenProcess can succeed after exit while any handle (including
274    // this process's Child) is still open. GetExitCodeProcess distinguishes
275    // STILL_ACTIVE from a real exit. The handle is closed exactly once.
276    unsafe {
277        let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid);
278        if handle.is_null() {
279            return GetLastError() == ERROR_ACCESS_DENIED;
280        }
281        let mut code = 0u32;
282        let ok = GetExitCodeProcess(handle, &mut code);
283        CloseHandle(handle);
284        ok != 0 && code == STILL_ACTIVE as u32
285    }
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291    use std::sync::Mutex;
292
293    static ENV_LOCK: Mutex<()> = Mutex::new(());
294
295    fn temp_registry() -> (tempfile::TempDir, LocalAdapterRegistry) {
296        let dir = tempfile::tempdir().unwrap();
297        let registry = LocalAdapterRegistry::with_dir(dir.path().to_path_buf());
298        (dir, registry)
299    }
300
301    fn sample(name: &str, pid: u32) -> LocalAdapterRecord {
302        LocalAdapterRecord {
303            name: name.to_string(),
304            did: "did:key:zTest".to_string(),
305            slim_endpoint: "127.0.0.1:47357".to_string(),
306            a2a_url: String::new(),
307            a2a_binding: A2ABinding::Grpc,
308            pid,
309        }
310    }
311
312    #[test]
313    fn publish_then_list_sees_this_process() {
314        let (_dir, registry) = temp_registry();
315        let record = sample("copilot", std::process::id());
316        let _lease = registry.publish(&record).unwrap();
317        assert_eq!(registry.list_live().unwrap(), vec![record]);
318    }
319
320    #[test]
321    fn list_live_drops_dead_pids() {
322        let (_dir, registry) = temp_registry();
323        let mut child = dead_child();
324        let pid = child.id();
325        let _ = child.wait();
326        drop(child);
327        let path = registry.dir.join(record_filename("codex", pid).unwrap());
328        fs::create_dir_all(&registry.dir).unwrap();
329        fs::write(&path, serde_json::to_vec(&sample("codex", pid)).unwrap()).unwrap();
330        assert!(registry.list_live().unwrap().is_empty());
331        assert!(!path.exists());
332    }
333
334    #[test]
335    fn lease_drop_removes_file() {
336        let (_dir, registry) = temp_registry();
337        let record = sample("claude-code", std::process::id());
338        {
339            let _lease = registry.publish(&record).unwrap();
340            assert_eq!(registry.list_live().unwrap().len(), 1);
341        }
342        assert!(registry.list_live().unwrap().is_empty());
343    }
344
345    #[test]
346    fn publish_rejects_path_traversal_name() {
347        let (_dir, registry) = temp_registry();
348        let err = match registry.publish(&sample("../evil", std::process::id())) {
349            Ok(_) => panic!("path traversal name must be rejected"),
350            Err(err) => err,
351        };
352        assert!(err.contains("unsafe"), "{err}");
353    }
354
355    #[test]
356    fn publish_rejects_empty_did_or_endpoint() {
357        let (_dir, registry) = temp_registry();
358        let pid = std::process::id();
359        let mut no_did = sample("copilot", pid);
360        no_did.did.clear();
361        let err = match registry.publish(&no_did) {
362            Ok(_) => panic!("empty DID must be rejected"),
363            Err(err) => err,
364        };
365        assert!(err.contains("DID"), "{err}");
366        let mut no_ep = sample("copilot", pid);
367        no_ep.slim_endpoint.clear();
368        let err = match registry.publish(&no_ep) {
369            Ok(_) => panic!("empty endpoint must be rejected"),
370            Err(err) => err,
371        };
372        assert!(err.contains("a2a_url") || err.contains("slim_endpoint"), "{err}");
373        no_ep.a2a_url = "http://127.0.0.1:9".to_string();
374        registry
375            .publish(&no_ep)
376            .expect("gRPC-only lease with DID and a2a_url must be accepted");
377    }
378
379    #[test]
380    fn publish_rejects_dash_prefix_and_empty_name() {
381        let (_dir, registry) = temp_registry();
382        assert!(registry
383            .publish(&sample("-copilot", std::process::id()))
384            .is_err());
385        assert!(registry.publish(&sample("", std::process::id())).is_err());
386    }
387
388    #[test]
389    fn publish_fails_when_registry_parent_is_a_file() {
390        let dir = tempfile::tempdir().unwrap();
391        let blocked = dir.path().join("blocked");
392        fs::write(&blocked, b"x").unwrap();
393        let registry = LocalAdapterRegistry::with_dir(blocked.join("nested"));
394        let err = match registry.publish(&sample("copilot", std::process::id())) {
395            Ok(_) => panic!("file-as-parent must fail"),
396            Err(err) => err,
397        };
398        assert!(err.contains("create local adapter registry"), "{err}");
399    }
400
401    #[test]
402    fn from_env_uses_shadi_tmp_dir() {
403        let _guard = ENV_LOCK.lock().expect("env lock");
404        let tmp = tempfile::tempdir().unwrap();
405        let prev = std::env::var_os("SHADI_TMP_DIR");
406        std::env::set_var("SHADI_TMP_DIR", tmp.path());
407        let registry = LocalAdapterRegistry::from_env();
408        if let Some(value) = prev {
409            std::env::set_var("SHADI_TMP_DIR", value);
410        } else {
411            std::env::remove_var("SHADI_TMP_DIR");
412        }
413        assert_eq!(registry.dir(), tmp.path().join("agentbridge-local"));
414    }
415
416    #[test]
417    fn list_live_missing_dir_is_empty() {
418        let registry = LocalAdapterRegistry::with_dir(std::path::PathBuf::from(
419            "/no/such/agentbridge-local-test-dir",
420        ));
421        assert!(registry.list_live().unwrap().is_empty());
422    }
423
424    #[test]
425    fn list_live_errors_when_registry_path_is_a_file() {
426        let dir = tempfile::tempdir().unwrap();
427        let path = dir.path().join("not-a-dir");
428        fs::write(&path, b"x").unwrap();
429        let registry = LocalAdapterRegistry::with_dir(path);
430        let err = registry.list_live().unwrap_err();
431        assert!(err.contains("read local adapter registry"), "{err}");
432    }
433
434    #[test]
435    fn list_live_skips_junk_and_deletes_empty_did() {
436        let (_dir, registry) = temp_registry();
437        fs::create_dir_all(registry.dir()).unwrap();
438        fs::write(registry.dir().join("readme.txt"), b"ignore").unwrap();
439        fs::write(registry.dir().join("copilot.json"), b"{}").unwrap();
440        fs::write(registry.dir().join("copilot-abc.json"), b"{}").unwrap();
441        fs::create_dir(registry.dir().join("subdir")).unwrap();
442        fs::write(registry.dir().join("copilot-1.json"), b"not-json").unwrap();
443        let empty_did = registry.dir().join("ghost-1.json");
444        fs::write(
445            &empty_did,
446            serde_json::to_vec(&LocalAdapterRecord {
447                name: "ghost".to_string(),
448                did: String::new(),
449                slim_endpoint: "127.0.0.1:1".to_string(),
450                a2a_url: String::new(),
451                a2a_binding: A2ABinding::Grpc,
452                pid: 1,
453            })
454            .unwrap(),
455        )
456        .unwrap();
457        let zero = registry.dir().join("zero-0.json");
458        fs::write(&zero, serde_json::to_vec(&sample("zero", 0)).unwrap()).unwrap();
459        assert!(registry.list_live().unwrap().is_empty());
460        assert!(!empty_did.exists());
461        assert!(!zero.exists());
462    }
463
464    #[test]
465    fn lease_filename_helpers_reject_junk() {
466        assert!(record_filename("../evil", 1).is_err());
467        assert!(!looks_like_lease_filename("readme.md"));
468        assert!(!looks_like_lease_filename("copilot.json"));
469        assert!(!looks_like_lease_filename("copilot-abc.json"));
470        assert!(looks_like_lease_filename("copilot-12.json"));
471        assert!(!pid_is_alive(0));
472    }
473
474    #[test]
475    fn local_registry_source_maps_to_candidates() {
476        let (_dir, registry) = temp_registry();
477        let record = sample("cursor-agent", std::process::id());
478        let _lease = registry.publish(&record).unwrap();
479        let source = LocalRegistrySource { registry };
480        let members = source.resolve().unwrap();
481        assert_eq!(
482            members,
483            vec![CandidateMember {
484                name: "cursor-agent".to_string(),
485                did: "did:key:zTest".to_string(),
486                slim_endpoint: Some("127.0.0.1:47357".to_string()),
487                a2a_url: None,
488                a2a_binding: None,
489            }]
490        );
491    }
492
493    #[test]
494    fn resolve_live_follows_did_when_url_changes() {
495        let (_dir, registry) = temp_registry();
496        let pid = std::process::id();
497        let mut first = sample("copilot", pid);
498        first.did = "did:key:zSame".to_string();
499        first.slim_endpoint.clear();
500        first.a2a_url = "http://127.0.0.1:50051".to_string();
501        let _lease = registry.publish(&first).unwrap();
502        assert_eq!(
503            registry.resolve_live("did:key:zSame").unwrap().a2a_url,
504            "http://127.0.0.1:50051"
505        );
506        first.a2a_url = "http://127.0.0.1:50052".to_string();
507        let _lease = registry.publish(&first).unwrap();
508        let found = registry.resolve_live("did:key:zSame").unwrap();
509        assert_eq!(found.a2a_url, "http://127.0.0.1:50052");
510        assert_eq!(found.a2a_binding, A2ABinding::Grpc);
511        assert_eq!(found.name, "copilot");
512    }
513
514    #[test]
515    fn resolve_live_accepts_slim_channel_alias() {
516        let (_dir, registry) = temp_registry();
517        let mut record = sample("copilot", std::process::id());
518        record.did = "did:key:zChan".to_string();
519        let _lease = registry.publish(&record).unwrap();
520        let found = registry
521            .resolve_live("agntcy/shadi/copilot-a2a")
522            .expect("channel");
523        assert_eq!(found.did, "did:key:zChan");
524        assert_eq!(found.name, "copilot");
525    }
526
527    #[test]
528    fn old_lease_without_binding_defaults_to_grpc() {
529        let raw = serde_json::json!({
530            "name": "copilot",
531            "did": "did:key:zOld",
532            "slim_endpoint": "",
533            "a2a_url": "http://127.0.0.1:50051",
534            "pid": 1
535        });
536        let record: LocalAdapterRecord = serde_json::from_value(raw).unwrap();
537        assert_eq!(record.a2a_binding, A2ABinding::Grpc);
538        assert_eq!(
539            record.to_candidate().unicast_locator().unwrap().display_uri(),
540            "grpc://127.0.0.1:50051"
541        );
542    }
543
544    #[test]
545    fn resolve_live_treats_shared_did_as_ambiguous() {
546        let (_dir, registry) = temp_registry();
547        let pid = std::process::id();
548        let mut a = sample("copilot", pid);
549        a.did = "did:key:zShared".to_string();
550        let mut b = sample("codex", pid);
551        b.did = "did:key:zShared".to_string();
552        let _la = registry.publish(&a).unwrap();
553        let _lb = registry.publish(&b).unwrap();
554        let err = registry.resolve_live("did:key:zShared").unwrap_err();
555        assert!(err.contains("ambiguous"), "{err}");
556        assert!(registry.resolve_live("copilot").is_ok());
557    }
558
559    fn dead_child() -> std::process::Child {
560        #[cfg(unix)]
561        {
562            std::process::Command::new("true").spawn().unwrap()
563        }
564        #[cfg(windows)]
565        {
566            let mut cmd = std::process::Command::new("cmd");
567            cmd.args(["/C", "exit", "0"]);
568            cmd.spawn().unwrap()
569        }
570    }
571}