1use std::collections::HashMap;
14use std::sync::Mutex;
15
16use crate::protocol::{AgentCard, AgentSkill};
17
18#[derive(Debug, thiserror::Error)]
20#[non_exhaustive]
21pub enum RegistryError {
22 #[error("agent `{0}` is not registered")]
24 UnknownAgent(String),
25 #[error("agent `{0}` is already registered")]
27 AlreadyRegistered(String),
28 #[error("registry request failed: {0}")]
30 Http(String),
31 #[error("registry payload malformed: {0}")]
33 Parse(String),
34}
35
36#[derive(Debug, Default)]
43pub struct AgentRegistry {
44 by_url: Mutex<HashMap<String, AgentCard>>,
45}
46
47impl AgentRegistry {
48 pub fn new() -> Self {
50 Self::default()
51 }
52
53 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 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 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 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 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 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 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 pub fn len(&self) -> usize {
127 self.by_url.lock().unwrap_or_else(|e| e.into_inner()).len()
128 }
129
130 pub fn is_empty(&self) -> bool {
132 self.len() == 0
133 }
134}
135
136pub struct RegistryClient {
142 base_url: String,
143 http: reqwest::Client,
144}
145
146impl RegistryClient {
147 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 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 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
193fn 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 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 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}