1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::path::PathBuf;
4
5const MAX_SCRATCHPAD_ENTRIES: usize = 200;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct AgentRegistry {
9 pub agents: Vec<AgentEntry>,
10 pub scratchpad: Vec<ScratchpadEntry>,
11 pub updated_at: DateTime<Utc>,
12}
13
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct AgentEntry {
16 pub agent_id: String,
17 pub agent_type: String,
18 pub role: Option<String>,
19 pub project_root: String,
20 pub started_at: DateTime<Utc>,
21 pub last_active: DateTime<Utc>,
22 pub pid: u32,
23 pub status: AgentStatus,
24 pub status_message: Option<String>,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
28pub enum AgentStatus {
29 Active,
30 Idle,
31 Finished,
32}
33
34impl std::fmt::Display for AgentStatus {
35 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36 match self {
37 AgentStatus::Active => write!(f, "active"),
38 AgentStatus::Idle => write!(f, "idle"),
39 AgentStatus::Finished => write!(f, "finished"),
40 }
41 }
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct ScratchpadEntry {
46 pub id: String,
47 pub from_agent: String,
48 pub to_agent: Option<String>,
49 pub category: String,
50 pub message: String,
51 pub timestamp: DateTime<Utc>,
52 pub read_by: Vec<String>,
53}
54
55impl AgentRegistry {
56 pub fn new() -> Self {
57 Self {
58 agents: Vec::new(),
59 scratchpad: Vec::new(),
60 updated_at: Utc::now(),
61 }
62 }
63
64 pub fn register(&mut self, agent_type: &str, role: Option<&str>, project_root: &str) -> String {
65 let pid = std::process::id();
66 let agent_id = format!("{}-{}-{}", agent_type, pid, &generate_short_id());
67
68 if let Some(existing) = self.agents.iter_mut().find(|a| a.pid == pid) {
69 existing.last_active = Utc::now();
70 existing.status = AgentStatus::Active;
71 if let Some(r) = role {
72 existing.role = Some(r.to_string());
73 }
74 return existing.agent_id.clone();
75 }
76
77 self.agents.push(AgentEntry {
78 agent_id: agent_id.clone(),
79 agent_type: agent_type.to_string(),
80 role: role.map(|r| r.to_string()),
81 project_root: project_root.to_string(),
82 started_at: Utc::now(),
83 last_active: Utc::now(),
84 pid,
85 status: AgentStatus::Active,
86 status_message: None,
87 });
88
89 self.updated_at = Utc::now();
90 agent_id
91 }
92
93 pub fn update_heartbeat(&mut self, agent_id: &str) {
94 if let Some(agent) = self.agents.iter_mut().find(|a| a.agent_id == agent_id) {
95 agent.last_active = Utc::now();
96 }
97 }
98
99 pub fn set_status(&mut self, agent_id: &str, status: AgentStatus, message: Option<&str>) {
100 if let Some(agent) = self.agents.iter_mut().find(|a| a.agent_id == agent_id) {
101 agent.status = status;
102 agent.status_message = message.map(|s| s.to_string());
103 agent.last_active = Utc::now();
104 }
105 self.updated_at = Utc::now();
106 }
107
108 pub fn list_active(&self, project_root: Option<&str>) -> Vec<&AgentEntry> {
109 self.agents
110 .iter()
111 .filter(|a| {
112 if let Some(root) = project_root {
113 a.project_root == root && a.status != AgentStatus::Finished
114 } else {
115 a.status != AgentStatus::Finished
116 }
117 })
118 .collect()
119 }
120
121 pub fn list_all(&self) -> &[AgentEntry] {
122 &self.agents
123 }
124
125 pub fn post_message(
126 &mut self,
127 from_agent: &str,
128 to_agent: Option<&str>,
129 category: &str,
130 message: &str,
131 ) -> String {
132 let id = generate_short_id();
133 self.scratchpad.push(ScratchpadEntry {
134 id: id.clone(),
135 from_agent: from_agent.to_string(),
136 to_agent: to_agent.map(|s| s.to_string()),
137 category: category.to_string(),
138 message: message.to_string(),
139 timestamp: Utc::now(),
140 read_by: vec![from_agent.to_string()],
141 });
142
143 if self.scratchpad.len() > MAX_SCRATCHPAD_ENTRIES {
144 self.scratchpad
145 .drain(0..self.scratchpad.len() - MAX_SCRATCHPAD_ENTRIES);
146 }
147
148 self.updated_at = Utc::now();
149 id
150 }
151
152 pub fn read_messages(&mut self, agent_id: &str) -> Vec<&ScratchpadEntry> {
153 let unread: Vec<usize> = self
154 .scratchpad
155 .iter()
156 .enumerate()
157 .filter(|(_, e)| {
158 !e.read_by.contains(&agent_id.to_string())
159 && (e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
160 })
161 .map(|(i, _)| i)
162 .collect();
163
164 for i in &unread {
165 self.scratchpad[*i].read_by.push(agent_id.to_string());
166 }
167
168 self.scratchpad
169 .iter()
170 .filter(|e| e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
171 .filter(|e| e.from_agent != agent_id)
172 .collect()
173 }
174
175 pub fn read_unread(&mut self, agent_id: &str) -> Vec<&ScratchpadEntry> {
176 let unread_indices: Vec<usize> = self
177 .scratchpad
178 .iter()
179 .enumerate()
180 .filter(|(_, e)| {
181 !e.read_by.contains(&agent_id.to_string())
182 && e.from_agent != agent_id
183 && (e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
184 })
185 .map(|(i, _)| i)
186 .collect();
187
188 for i in &unread_indices {
189 self.scratchpad[*i].read_by.push(agent_id.to_string());
190 }
191
192 self.updated_at = Utc::now();
193
194 self.scratchpad
195 .iter()
196 .filter(|e| {
197 e.from_agent != agent_id
198 && (e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
199 && e.read_by.contains(&agent_id.to_string())
200 && e.read_by.iter().filter(|r| *r == agent_id).count() == 1
201 })
202 .collect()
203 }
204
205 pub fn cleanup_stale(&mut self, max_age_hours: u64) {
206 let cutoff = Utc::now() - chrono::Duration::hours(max_age_hours as i64);
207
208 for agent in &mut self.agents {
209 if agent.last_active < cutoff
210 && agent.status != AgentStatus::Finished
211 && !is_process_alive(agent.pid)
212 {
213 agent.status = AgentStatus::Finished;
214 }
215 }
216
217 self.agents
218 .retain(|a| !(a.status == AgentStatus::Finished && a.last_active < cutoff));
219
220 self.updated_at = Utc::now();
221 }
222
223 pub fn save(&self) -> Result<(), String> {
224 let dir = agents_dir()?;
225 std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
226
227 let path = dir.join("registry.json");
228 let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
229
230 let lock_path = dir.join("registry.lock");
231 let _lock = FileLock::acquire(&lock_path)?;
232
233 std::fs::write(&path, json).map_err(|e| e.to_string())
234 }
235
236 pub fn load() -> Option<Self> {
237 let dir = agents_dir().ok()?;
238 let path = dir.join("registry.json");
239 let content = std::fs::read_to_string(&path).ok()?;
240 serde_json::from_str(&content).ok()
241 }
242
243 pub fn load_or_create() -> Self {
244 Self::load().unwrap_or_default()
245 }
246}
247
248impl Default for AgentRegistry {
249 fn default() -> Self {
250 Self::new()
251 }
252}
253
254fn agents_dir() -> Result<PathBuf, String> {
255 let home = dirs::home_dir().ok_or("Cannot determine home directory")?;
256 Ok(home.join(".lean-ctx").join("agents"))
257}
258
259fn generate_short_id() -> String {
260 use std::collections::hash_map::DefaultHasher;
261 use std::hash::{Hash, Hasher};
262 use std::time::SystemTime;
263
264 let mut hasher = DefaultHasher::new();
265 SystemTime::now().hash(&mut hasher);
266 std::process::id().hash(&mut hasher);
267 format!("{:08x}", hasher.finish() as u32)
268}
269
270fn is_process_alive(pid: u32) -> bool {
271 #[cfg(unix)]
272 {
273 std::process::Command::new("kill")
274 .args(["-0", &pid.to_string()])
275 .output()
276 .map(|o| o.status.success())
277 .unwrap_or(false)
278 }
279 #[cfg(not(unix))]
280 {
281 let _ = pid;
282 true
283 }
284}
285
286struct FileLock {
287 path: PathBuf,
288}
289
290impl FileLock {
291 fn acquire(path: &std::path::Path) -> Result<Self, String> {
292 for _ in 0..50 {
293 match std::fs::OpenOptions::new()
294 .write(true)
295 .create_new(true)
296 .open(path)
297 {
298 Ok(_) => {
299 return Ok(Self {
300 path: path.to_path_buf(),
301 })
302 }
303 Err(_) => {
304 if let Ok(metadata) = std::fs::metadata(path) {
305 if let Ok(modified) = metadata.modified() {
306 if modified.elapsed().unwrap_or_default().as_secs() > 5 {
307 let _ = std::fs::remove_file(path);
308 continue;
309 }
310 }
311 }
312 std::thread::sleep(std::time::Duration::from_millis(100));
313 }
314 }
315 }
316 Err("Could not acquire lock after 5 seconds".to_string())
317 }
318}
319
320impl Drop for FileLock {
321 fn drop(&mut self) {
322 let _ = std::fs::remove_file(&self.path);
323 }
324}
325
326#[cfg(test)]
327mod tests {
328 use super::*;
329
330 #[test]
331 fn register_and_list() {
332 let mut reg = AgentRegistry::new();
333 let id = reg.register("cursor", Some("dev"), "/tmp/project");
334 assert!(!id.is_empty());
335 assert_eq!(reg.list_active(None).len(), 1);
336 assert_eq!(reg.list_active(None)[0].agent_type, "cursor");
337 }
338
339 #[test]
340 fn reregister_same_pid() {
341 let mut reg = AgentRegistry::new();
342 let id1 = reg.register("cursor", Some("dev"), "/tmp/project");
343 let id2 = reg.register("cursor", Some("review"), "/tmp/project");
344 assert_eq!(id1, id2);
345 assert_eq!(reg.agents.len(), 1);
346 assert_eq!(reg.agents[0].role, Some("review".to_string()));
347 }
348
349 #[test]
350 fn post_and_read_messages() {
351 let mut reg = AgentRegistry::new();
352 reg.post_message("agent-a", None, "finding", "Found a bug in auth.rs");
353 reg.post_message("agent-b", Some("agent-a"), "request", "Please review");
354
355 let msgs = reg.read_unread("agent-a");
356 assert_eq!(msgs.len(), 1);
357 assert_eq!(msgs[0].category, "request");
358 }
359
360 #[test]
361 fn set_status() {
362 let mut reg = AgentRegistry::new();
363 let id = reg.register("claude", None, "/tmp/project");
364 reg.set_status(&id, AgentStatus::Idle, Some("waiting for review"));
365 assert_eq!(reg.agents[0].status, AgentStatus::Idle);
366 assert_eq!(
367 reg.agents[0].status_message,
368 Some("waiting for review".to_string())
369 );
370 }
371
372 #[test]
373 fn broadcast_message() {
374 let mut reg = AgentRegistry::new();
375 reg.post_message("agent-a", None, "status", "Starting refactor");
376
377 let msgs_b = reg.read_unread("agent-b");
378 assert_eq!(msgs_b.len(), 1);
379 assert_eq!(msgs_b[0].message, "Starting refactor");
380
381 let msgs_a = reg.read_unread("agent-a");
382 assert!(msgs_a.is_empty());
383 }
384}