jamjet_agents/
memory_registry.rs1use crate::card::AgentCard;
5use crate::lifecycle::AgentStatus;
6use crate::registry::{Agent, AgentFilter, AgentId, AgentRegistry};
7use async_trait::async_trait;
8use chrono::Utc;
9use dashmap::DashMap;
10use uuid::Uuid;
11
12pub struct InMemoryAgentRegistry {
13 agents: DashMap<AgentId, Agent>,
14}
15
16impl InMemoryAgentRegistry {
17 pub fn new() -> Self {
18 Self {
19 agents: DashMap::new(),
20 }
21 }
22}
23
24impl Default for InMemoryAgentRegistry {
25 fn default() -> Self {
26 Self::new()
27 }
28}
29
30#[async_trait]
31impl AgentRegistry for InMemoryAgentRegistry {
32 async fn register(&self, card: AgentCard) -> Result<AgentId, String> {
33 let id = Uuid::new_v4();
34 let now = Utc::now();
35 let agent = Agent {
36 id,
37 card,
38 status: AgentStatus::Registered,
39 registered_at: now,
40 updated_at: now,
41 last_heartbeat: None,
42 };
43 self.agents.insert(id, agent);
44 Ok(id)
45 }
46
47 async fn get(&self, id: AgentId) -> Result<Option<Agent>, String> {
48 Ok(self.agents.get(&id).map(|r| r.value().clone()))
49 }
50
51 async fn get_by_uri(&self, uri: &str) -> Result<Option<Agent>, String> {
52 Ok(self
53 .agents
54 .iter()
55 .find(|r| r.value().card.uri == uri)
56 .map(|r| r.value().clone()))
57 }
58
59 async fn find(&self, filter: AgentFilter) -> Result<Vec<Agent>, String> {
60 let results: Vec<Agent> = self
61 .agents
62 .iter()
63 .filter(|r| {
64 let a = r.value();
65 if let Some(ref status) = filter.status {
66 if &a.status != status {
67 return false;
68 }
69 }
70 if let Some(ref skill) = filter.skill {
71 let has_skill = a
72 .card
73 .capabilities
74 .skills
75 .iter()
76 .any(|s| s.name.eq_ignore_ascii_case(skill));
77 if !has_skill {
78 return false;
79 }
80 }
81 if let Some(ref protocol) = filter.protocol {
82 let has_proto = a
83 .card
84 .capabilities
85 .protocols
86 .iter()
87 .any(|p| p.eq_ignore_ascii_case(protocol));
88 if !has_proto {
89 return false;
90 }
91 }
92 true
93 })
94 .map(|r| r.value().clone())
95 .collect();
96 Ok(results)
97 }
98
99 async fn update_status(&self, id: AgentId, status: AgentStatus) -> Result<(), String> {
100 match self.agents.get_mut(&id) {
101 Some(mut entry) => {
102 entry.status = status;
103 entry.updated_at = Utc::now();
104 Ok(())
105 }
106 None => Err(format!("agent not found: {id}")),
107 }
108 }
109
110 async fn heartbeat(&self, id: AgentId) -> Result<(), String> {
111 match self.agents.get_mut(&id) {
112 Some(mut entry) => {
113 entry.last_heartbeat = Some(Utc::now());
114 Ok(())
115 }
116 None => Err(format!("agent not found: {id}")),
117 }
118 }
119
120 async fn discover_remote(&self, url: &str) -> Result<Agent, String> {
121 let card_url = if url.ends_with("/.well-known/agent.json") {
122 url.to_string()
123 } else {
124 format!("{}/.well-known/agent.json", url.trim_end_matches('/'))
125 };
126 let resp = reqwest::get(&card_url)
127 .await
128 .map_err(|e| format!("failed to fetch agent card: {e}"))?;
129 let card: AgentCard = resp
130 .json()
131 .await
132 .map_err(|e| format!("failed to parse agent card: {e}"))?;
133 let id = self.register(card.clone()).await?;
134 self.get(id)
135 .await?
136 .ok_or_else(|| "agent registered but not found".to_string())
137 }
138}