Skip to main content

lc_a2a/
discovery.rs

1//! P2-8: agent discovery.
2//!
3//! A caller that can't enumerate 1000 agent URLs needs a directory to ask:
4//! "who can do X?" [`AgentRegistry`] is a local in-memory catalog of
5//! [`AgentCard`]s with skill- and data-boundary-aware lookup.
6//! [`RegistryClient`] pulls such a catalog from a remote HTTP registry
7//! (`GET /registry.json`).
8//!
9//! DNS-SD / mDNS is a possible future transport for the same directory; the
10//! catalog shape (a `Vec<AgentCard>`) is transport-agnostic so a DNS-SD-backed
11//! provider could slot in behind the same lookup helpers.
12
13use std::collections::HashMap;
14use std::sync::Mutex;
15
16use crate::protocol::{AgentCard, AgentSkill};
17
18/// Errors raised by the discovery components.
19#[derive(Debug, thiserror::Error)]
20#[non_exhaustive]
21pub enum RegistryError {
22    /// The agent is not in the registry.
23    #[error("agent `{0}` is not registered")]
24    UnknownAgent(String),
25    /// `register` was called twice with the same URL (use `upsert` to replace).
26    #[error("agent `{0}` is already registered")]
27    AlreadyRegistered(String),
28    /// A remote registry request failed at the transport level.
29    #[error("registry request failed: {0}")]
30    Http(String),
31    /// A remote registry returned an unparseable catalog.
32    #[error("registry payload malformed: {0}")]
33    Parse(String),
34}
35
36/// In-memory agent directory keyed by agent URL (P2-8).
37///
38/// Register every agent's [`AgentCard`], then discover by skill
39/// ([`AgentRegistry::search_skill`]) or data boundary
40/// ([`AgentRegistry::filter_data_class`]) instead of hardcoding URLs. All
41/// operations are cheap hash lookups behind a short-lived mutex.
42#[derive(Debug, Default)]
43pub struct AgentRegistry {
44    by_url: Mutex<HashMap<String, AgentCard>>,
45}
46
47impl AgentRegistry {
48    /// An empty directory.
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    /// Register an agent. Fails if the URL is already registered — use
54    /// [`AgentRegistry::upsert`] to replace a card in place.
55    pub fn register(&self, card: AgentCard) -> Result<(), RegistryError> {
56        let mut by_url = self.by_url.lock().unwrap_or_else(|e| e.into_inner());
57        if by_url.contains_key(&card.url) {
58            return Err(RegistryError::AlreadyRegistered(card.url));
59        }
60        by_url.insert(card.url.clone(), card);
61        Ok(())
62    }
63
64    /// Register or replace an agent card in place (idempotent).
65    pub fn upsert(&self, card: AgentCard) {
66        self.by_url
67            .lock()
68            .unwrap_or_else(|e| e.into_inner())
69            .insert(card.url.clone(), card);
70    }
71
72    /// Remove an agent by URL. Returns `true` if it was present.
73    pub fn unregister(&self, url: &str) -> bool {
74        self.by_url
75            .lock()
76            .unwrap_or_else(|e| e.into_inner())
77            .remove(url)
78            .is_some()
79    }
80
81    /// Look up an agent card by URL.
82    pub fn lookup(&self, url: &str) -> Option<AgentCard> {
83        self.by_url
84            .lock()
85            .unwrap_or_else(|e| e.into_inner())
86            .get(url)
87            .cloned()
88    }
89
90    /// Every registered agent card, in arbitrary order.
91    pub fn agents(&self) -> Vec<AgentCard> {
92        self.by_url
93            .lock()
94            .unwrap_or_else(|e| e.into_inner())
95            .values()
96            .cloned()
97            .collect()
98    }
99
100    /// Discover agents advertising a skill whose id/name/description contains
101    /// `query` (case-insensitive substring match).
102    pub fn search_skill(&self, query: &str) -> Vec<AgentCard> {
103        let q = query.to_lowercase();
104        self.by_url
105            .lock()
106            .unwrap_or_else(|e| e.into_inner())
107            .values()
108            .filter(|card| card.skills.iter().any(|s| skill_matches(s, &q)))
109            .cloned()
110            .collect()
111    }
112
113    /// Discover agents whose card declares exactly `class` as its data class
114    /// (data boundary, P2-8).
115    pub fn filter_data_class(&self, class: &str) -> Vec<AgentCard> {
116        self.by_url
117            .lock()
118            .unwrap_or_else(|e| e.into_inner())
119            .values()
120            .filter(|card| card.data_class.as_deref() == Some(class))
121            .cloned()
122            .collect()
123    }
124
125    /// How many agents are registered.
126    pub fn len(&self) -> usize {
127        self.by_url.lock().unwrap_or_else(|e| e.into_inner()).len()
128    }
129
130    /// Whether the registry is empty.
131    pub fn is_empty(&self) -> bool {
132        self.len() == 0
133    }
134}
135
136/// Client that fetches an agent catalog from a remote HTTP registry (P2-8).
137///
138/// The registry is expected to serve a JSON array of [`AgentCard`]s at
139/// `GET {base}/registry.json`. Discovery helpers (`search_skill`) filter the
140/// catalog locally so a remote directory needs zero extra endpoints.
141pub struct RegistryClient {
142    base_url: String,
143    http: reqwest::Client,
144}
145
146impl RegistryClient {
147    /// A client for a remote registry at `base_url`.
148    ///
149    /// The client disables proxy usage: registries are typically on a private
150    /// network, and an environment proxy must not intercept catalog fetches.
151    pub fn new(base_url: impl Into<String>) -> Result<Self, RegistryError> {
152        Ok(Self {
153            base_url: base_url.into(),
154            http: reqwest::Client::builder()
155                .no_proxy()
156                .build()
157                .map_err(|e| RegistryError::Http(format!("failed to build HTTP client: {e}")))?,
158        })
159    }
160
161    /// Fetch the full catalog from the remote registry.
162    pub async fn fetch_catalog(&self) -> Result<Vec<AgentCard>, RegistryError> {
163        let url = format!("{}/registry.json", self.base_url);
164        let resp = self
165            .http
166            .get(&url)
167            .send()
168            .await
169            .map_err(|e| RegistryError::Http(e.to_string()))?;
170        if !resp.status().is_success() {
171            return Err(RegistryError::Http(format!(
172                "registry returned {}",
173                resp.status()
174            )));
175        }
176        resp.json::<Vec<AgentCard>>()
177            .await
178            .map_err(|e| RegistryError::Parse(e.to_string()))
179    }
180
181    /// Fetch the catalog and filter to agents advertising `query` as a skill.
182    pub async fn search_skill(&self, query: &str) -> Result<Vec<AgentCard>, RegistryError> {
183        let q = query.to_lowercase();
184        Ok(self
185            .fetch_catalog()
186            .await?
187            .into_iter()
188            .filter(|card| card.skills.iter().any(|s| skill_matches(s, &q)))
189            .collect())
190    }
191}
192
193/// Whether a skill's id/name/description matches a lowercased query.
194fn skill_matches(skill: &AgentSkill, query: &str) -> bool {
195    skill.id.to_lowercase().contains(query)
196        || skill.name.to_lowercase().contains(query)
197        || skill.description.to_lowercase().contains(query)
198}
199
200#[cfg(test)]
201mod tests {
202    use super::*;
203    use tokio::io::{AsyncReadExt, AsyncWriteExt};
204
205    fn card(url: &str, skill: &str, data_class: Option<&str>) -> AgentCard {
206        let mut card = AgentCard::new("agent", "a test agent", url).with_skill(AgentSkill::new(
207            skill,
208            skill,
209            format!("provides {skill}"),
210        ));
211        if let Some(class) = data_class {
212            card = card.with_data_class(class);
213        }
214        card
215    }
216
217    #[test]
218    fn registry_register_lookup_unregister() {
219        let reg = AgentRegistry::new();
220        assert!(reg.is_empty());
221        reg.register(card("http://a", "summarize", None)).unwrap();
222        assert_eq!(reg.len(), 1);
223        assert!(reg.lookup("http://a").is_some());
224        assert!(reg.unregister("http://a"));
225        assert!(!reg.unregister("http://a"));
226        assert!(reg.is_empty());
227    }
228
229    #[test]
230    fn registry_rejects_duplicate_register_but_upsert_replaces() {
231        let reg = AgentRegistry::new();
232        reg.register(card("http://a", "summarize", None)).unwrap();
233        assert!(matches!(
234            reg.register(card("http://a", "translate", None)),
235            Err(RegistryError::AlreadyRegistered(_))
236        ));
237        reg.upsert(card("http://a", "translate", None));
238        assert_eq!(reg.agents()[0].skills[0].id, "translate");
239    }
240
241    #[test]
242    fn registry_searches_skill_case_insensitively() {
243        let reg = AgentRegistry::new();
244        reg.upsert(card("http://sum", "summarize", None));
245        reg.upsert(card("http://translate", "translate", None));
246
247        let hits = reg.search_skill("Summ");
248        assert_eq!(hits.len(), 1);
249        assert_eq!(hits[0].url, "http://sum");
250
251        assert!(reg.search_skill("nothing").is_empty());
252    }
253
254    #[test]
255    fn registry_filters_by_data_class() {
256        let reg = AgentRegistry::new();
257        reg.upsert(card("http://a", "summarize", Some("public")));
258        reg.upsert(card("http://b", "translate", Some("confidential")));
259
260        let public = reg.filter_data_class("public");
261        assert_eq!(public.len(), 1);
262        assert_eq!(public[0].url, "http://a");
263        assert!(reg.filter_data_class("internal").is_empty());
264    }
265
266    #[tokio::test]
267    async fn registry_client_fetches_remote_catalog() {
268        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
269        let addr = listener.local_addr().unwrap();
270        // Serve the same catalog on every connection so a client that issues
271        // several fetches (e.g. `search_skill` after `fetch_catalog`) succeeds.
272        tokio::spawn(async move {
273            let body = serde_json::json!([
274                {
275                    "name": "summarizer",
276                    "description": "summarizes",
277                    "url": "http://sum",
278                    "skills": [
279                        { "id": "summarize", "name": "summarize", "description": "provides summarize" }
280                    ],
281                    "protocolVersion": "0.3.0",
282                    "interfaces": {},
283                    "securitySchemes": []
284                }
285            ])
286            .to_string();
287            let resp = format!(
288                "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nContent-Type: application/json\r\n\r\n{}",
289                body.len(),
290                body
291            );
292            while let Ok((mut stream, _)) = listener.accept().await {
293                let resp = resp.clone();
294                tokio::spawn(async move {
295                    // Best-effort drain of the request head: a GET has no body,
296                    // so we must not block waiting for bytes that will never
297                    // come. A brief timeout lets the write below proceed
298                    // regardless of how the OS chunks the small request.
299                    let mut chunk = [0u8; 4096];
300                    let _ = tokio::time::timeout(
301                        std::time::Duration::from_millis(200),
302                        stream.read(&mut chunk),
303                    )
304                    .await;
305                    let _ = stream.write_all(resp.as_bytes()).await;
306                });
307            }
308        });
309
310        let client = RegistryClient::new(format!("http://{addr}")).unwrap();
311        let catalog = client.fetch_catalog().await.unwrap();
312        assert_eq!(catalog.len(), 1);
313        assert_eq!(catalog[0].url, "http://sum");
314
315        let hits = client.search_skill("summ").await.unwrap();
316        assert_eq!(hits.len(), 1);
317        assert_eq!(hits[0].name, "summarizer");
318    }
319}