Skip to main content

jamjet_agents/
memory_registry.rs

1//! In-memory agent registry — no persistence, suitable for testing and
2//! ephemeral sandbox deployments (e.g. Glama).
3
4use 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}