1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::path::PathBuf;
4
5const MAX_SCRATCHPAD_ENTRIES: usize = 200;
6const MAX_DIARY_ENTRIES: usize = 100;
7
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct AgentRegistry {
10 pub agents: Vec<AgentEntry>,
11 pub scratchpad: Vec<ScratchpadEntry>,
12 pub updated_at: DateTime<Utc>,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct AgentDiary {
17 pub agent_id: String,
18 pub agent_type: String,
19 pub project_root: String,
20 pub entries: Vec<DiaryEntry>,
21 pub created_at: DateTime<Utc>,
22 pub updated_at: DateTime<Utc>,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct DiaryEntry {
27 pub entry_type: DiaryEntryType,
28 pub content: String,
29 pub context: Option<String>,
30 pub timestamp: DateTime<Utc>,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
34pub enum DiaryEntryType {
35 Discovery,
36 Decision,
37 Blocker,
38 Progress,
39 Insight,
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct AgentEntry {
44 pub agent_id: String,
45 pub agent_type: String,
46 pub role: Option<String>,
47 pub project_root: String,
48 pub started_at: DateTime<Utc>,
49 pub last_active: DateTime<Utc>,
50 pub pid: u32,
51 pub status: AgentStatus,
52 pub status_message: Option<String>,
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
56pub enum AgentStatus {
57 Active,
58 Idle,
59 Finished,
60}
61
62impl std::fmt::Display for AgentStatus {
63 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
64 match self {
65 AgentStatus::Active => write!(f, "active"),
66 AgentStatus::Idle => write!(f, "idle"),
67 AgentStatus::Finished => write!(f, "finished"),
68 }
69 }
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct ScratchpadEntry {
74 pub id: String,
75 pub from_agent: String,
76 pub to_agent: Option<String>,
77 pub category: String,
78 pub message: String,
79 pub timestamp: DateTime<Utc>,
80 pub read_by: Vec<String>,
81}
82
83impl AgentRegistry {
84 pub fn new() -> Self {
85 Self {
86 agents: Vec::new(),
87 scratchpad: Vec::new(),
88 updated_at: Utc::now(),
89 }
90 }
91
92 pub fn register(&mut self, agent_type: &str, role: Option<&str>, project_root: &str) -> String {
93 let pid = std::process::id();
94 let agent_id = format!("{}-{}-{}", agent_type, pid, &generate_short_id());
95
96 if let Some(existing) = self.agents.iter_mut().find(|a| a.pid == pid) {
97 existing.last_active = Utc::now();
98 existing.status = AgentStatus::Active;
99 if let Some(r) = role {
100 existing.role = Some(r.to_string());
101 }
102 return existing.agent_id.clone();
103 }
104
105 self.agents.push(AgentEntry {
106 agent_id: agent_id.clone(),
107 agent_type: agent_type.to_string(),
108 role: role.map(|r| r.to_string()),
109 project_root: project_root.to_string(),
110 started_at: Utc::now(),
111 last_active: Utc::now(),
112 pid,
113 status: AgentStatus::Active,
114 status_message: None,
115 });
116
117 self.updated_at = Utc::now();
118 crate::core::events::emit_agent_action(&agent_id, "register", None);
119 agent_id
120 }
121
122 pub fn update_heartbeat(&mut self, agent_id: &str) {
123 if let Some(agent) = self.agents.iter_mut().find(|a| a.agent_id == agent_id) {
124 agent.last_active = Utc::now();
125 }
126 }
127
128 pub fn set_status(&mut self, agent_id: &str, status: AgentStatus, message: Option<&str>) {
129 if let Some(agent) = self.agents.iter_mut().find(|a| a.agent_id == agent_id) {
130 agent.status = status;
131 agent.status_message = message.map(|s| s.to_string());
132 agent.last_active = Utc::now();
133 }
134 self.updated_at = Utc::now();
135 }
136
137 pub fn list_active(&self, project_root: Option<&str>) -> Vec<&AgentEntry> {
138 self.agents
139 .iter()
140 .filter(|a| {
141 if let Some(root) = project_root {
142 a.project_root == root && a.status != AgentStatus::Finished
143 } else {
144 a.status != AgentStatus::Finished
145 }
146 })
147 .collect()
148 }
149
150 pub fn list_all(&self) -> &[AgentEntry] {
151 &self.agents
152 }
153
154 pub fn post_message(
155 &mut self,
156 from_agent: &str,
157 to_agent: Option<&str>,
158 category: &str,
159 message: &str,
160 ) -> String {
161 let id = generate_short_id();
162 self.scratchpad.push(ScratchpadEntry {
163 id: id.clone(),
164 from_agent: from_agent.to_string(),
165 to_agent: to_agent.map(|s| s.to_string()),
166 category: category.to_string(),
167 message: message.to_string(),
168 timestamp: Utc::now(),
169 read_by: vec![from_agent.to_string()],
170 });
171
172 if self.scratchpad.len() > MAX_SCRATCHPAD_ENTRIES {
173 self.scratchpad
174 .drain(0..self.scratchpad.len() - MAX_SCRATCHPAD_ENTRIES);
175 }
176
177 self.updated_at = Utc::now();
178 id
179 }
180
181 pub fn read_messages(&mut self, agent_id: &str) -> Vec<&ScratchpadEntry> {
182 let unread: Vec<usize> = self
183 .scratchpad
184 .iter()
185 .enumerate()
186 .filter(|(_, e)| {
187 !e.read_by.contains(&agent_id.to_string())
188 && (e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
189 })
190 .map(|(i, _)| i)
191 .collect();
192
193 for i in &unread {
194 self.scratchpad[*i].read_by.push(agent_id.to_string());
195 }
196
197 self.scratchpad
198 .iter()
199 .filter(|e| e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
200 .filter(|e| e.from_agent != agent_id)
201 .collect()
202 }
203
204 pub fn read_unread(&mut self, agent_id: &str) -> Vec<&ScratchpadEntry> {
205 let unread_indices: Vec<usize> = self
206 .scratchpad
207 .iter()
208 .enumerate()
209 .filter(|(_, e)| {
210 !e.read_by.contains(&agent_id.to_string())
211 && e.from_agent != agent_id
212 && (e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
213 })
214 .map(|(i, _)| i)
215 .collect();
216
217 for i in &unread_indices {
218 self.scratchpad[*i].read_by.push(agent_id.to_string());
219 }
220
221 self.updated_at = Utc::now();
222
223 self.scratchpad
224 .iter()
225 .filter(|e| {
226 e.from_agent != agent_id
227 && (e.to_agent.is_none() || e.to_agent.as_deref() == Some(agent_id))
228 && e.read_by.contains(&agent_id.to_string())
229 && e.read_by.iter().filter(|r| *r == agent_id).count() == 1
230 })
231 .collect()
232 }
233
234 pub fn cleanup_stale(&mut self, max_age_hours: u64) {
235 let cutoff = Utc::now() - chrono::Duration::hours(max_age_hours as i64);
236
237 for agent in &mut self.agents {
238 if agent.last_active < cutoff
239 && agent.status != AgentStatus::Finished
240 && !is_process_alive(agent.pid)
241 {
242 agent.status = AgentStatus::Finished;
243 }
244 }
245
246 self.agents
247 .retain(|a| !(a.status == AgentStatus::Finished && a.last_active < cutoff));
248
249 self.updated_at = Utc::now();
250 }
251
252 pub fn save(&self) -> Result<(), String> {
253 let dir = agents_dir()?;
254 std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
255
256 let path = dir.join("registry.json");
257 let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
258
259 let lock_path = dir.join("registry.lock");
260 let _lock = FileLock::acquire(&lock_path)?;
261
262 std::fs::write(&path, json).map_err(|e| e.to_string())
263 }
264
265 pub fn load() -> Option<Self> {
266 let dir = agents_dir().ok()?;
267 let path = dir.join("registry.json");
268 let content = std::fs::read_to_string(&path).ok()?;
269 serde_json::from_str(&content).ok()
270 }
271
272 pub fn load_or_create() -> Self {
273 Self::load().unwrap_or_default()
274 }
275}
276
277impl Default for AgentRegistry {
278 fn default() -> Self {
279 Self::new()
280 }
281}
282
283impl AgentDiary {
284 pub fn new(agent_id: &str, agent_type: &str, project_root: &str) -> Self {
285 let now = Utc::now();
286 Self {
287 agent_id: agent_id.to_string(),
288 agent_type: agent_type.to_string(),
289 project_root: project_root.to_string(),
290 entries: Vec::new(),
291 created_at: now,
292 updated_at: now,
293 }
294 }
295
296 pub fn add_entry(&mut self, entry_type: DiaryEntryType, content: &str, context: Option<&str>) {
297 self.entries.push(DiaryEntry {
298 entry_type,
299 content: content.to_string(),
300 context: context.map(|s| s.to_string()),
301 timestamp: Utc::now(),
302 });
303 if self.entries.len() > MAX_DIARY_ENTRIES {
304 self.entries
305 .drain(0..self.entries.len() - MAX_DIARY_ENTRIES);
306 }
307 self.updated_at = Utc::now();
308 }
309
310 pub fn format_summary(&self) -> String {
311 if self.entries.is_empty() {
312 return format!("Diary [{}]: empty", self.agent_id);
313 }
314 let mut out = format!(
315 "Diary [{}] ({} entries):\n",
316 self.agent_id,
317 self.entries.len()
318 );
319 for e in self.entries.iter().rev().take(10) {
320 let age = (Utc::now() - e.timestamp).num_minutes();
321 let prefix = match e.entry_type {
322 DiaryEntryType::Discovery => "FOUND",
323 DiaryEntryType::Decision => "DECIDED",
324 DiaryEntryType::Blocker => "BLOCKED",
325 DiaryEntryType::Progress => "DONE",
326 DiaryEntryType::Insight => "INSIGHT",
327 };
328 let ctx = e
329 .context
330 .as_deref()
331 .map(|c| format!(" [{c}]"))
332 .unwrap_or_default();
333 out.push_str(&format!(" [{prefix}] {}{ctx} ({age}m ago)\n", e.content));
334 }
335 out
336 }
337
338 pub fn format_compact(&self) -> String {
339 if self.entries.is_empty() {
340 return String::new();
341 }
342 let items: Vec<String> = self
343 .entries
344 .iter()
345 .rev()
346 .take(5)
347 .map(|e| {
348 let prefix = match e.entry_type {
349 DiaryEntryType::Discovery => "F",
350 DiaryEntryType::Decision => "D",
351 DiaryEntryType::Blocker => "B",
352 DiaryEntryType::Progress => "P",
353 DiaryEntryType::Insight => "I",
354 };
355 format!("{prefix}:{}", truncate(&e.content, 50))
356 })
357 .collect();
358 format!("diary:{}|{}", self.agent_id, items.join("|"))
359 }
360
361 pub fn save(&self) -> Result<(), String> {
362 let dir = diary_dir()?;
363 std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
364 let path = dir.join(format!("{}.json", sanitize_filename(&self.agent_id)));
365 let json = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
366 std::fs::write(&path, json).map_err(|e| e.to_string())
367 }
368
369 pub fn load(agent_id: &str) -> Option<Self> {
370 let dir = diary_dir().ok()?;
371 let path = dir.join(format!("{}.json", sanitize_filename(agent_id)));
372 let content = std::fs::read_to_string(&path).ok()?;
373 serde_json::from_str(&content).ok()
374 }
375
376 pub fn load_or_create(agent_id: &str, agent_type: &str, project_root: &str) -> Self {
377 Self::load(agent_id).unwrap_or_else(|| Self::new(agent_id, agent_type, project_root))
378 }
379
380 pub fn list_all() -> Vec<(String, usize, DateTime<Utc>)> {
381 let dir = match diary_dir() {
382 Ok(d) => d,
383 Err(_) => return Vec::new(),
384 };
385 if !dir.exists() {
386 return Vec::new();
387 }
388 let mut results = Vec::new();
389 if let Ok(entries) = std::fs::read_dir(&dir) {
390 for entry in entries.flatten() {
391 if entry.path().extension().and_then(|e| e.to_str()) == Some("json") {
392 if let Ok(content) = std::fs::read_to_string(entry.path()) {
393 if let Ok(diary) = serde_json::from_str::<AgentDiary>(&content) {
394 results.push((diary.agent_id, diary.entries.len(), diary.updated_at));
395 }
396 }
397 }
398 }
399 }
400 results.sort_by(|a, b| b.2.cmp(&a.2));
401 results
402 }
403}
404
405impl std::fmt::Display for DiaryEntryType {
406 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
407 match self {
408 DiaryEntryType::Discovery => write!(f, "discovery"),
409 DiaryEntryType::Decision => write!(f, "decision"),
410 DiaryEntryType::Blocker => write!(f, "blocker"),
411 DiaryEntryType::Progress => write!(f, "progress"),
412 DiaryEntryType::Insight => write!(f, "insight"),
413 }
414 }
415}
416
417fn diary_dir() -> Result<PathBuf, String> {
418 let home = dirs::home_dir().ok_or("Cannot determine home directory")?;
419 Ok(home.join(".lean-ctx").join("agents").join("diaries"))
420}
421
422fn sanitize_filename(name: &str) -> String {
423 name.chars()
424 .map(|c| {
425 if c.is_alphanumeric() || c == '-' || c == '_' {
426 c
427 } else {
428 '_'
429 }
430 })
431 .collect()
432}
433
434fn truncate(s: &str, max: usize) -> String {
435 if s.len() <= max {
436 s.to_string()
437 } else {
438 format!("{}...", &s[..max.saturating_sub(3)])
439 }
440}
441
442fn agents_dir() -> Result<PathBuf, String> {
443 let home = dirs::home_dir().ok_or("Cannot determine home directory")?;
444 Ok(home.join(".lean-ctx").join("agents"))
445}
446
447fn generate_short_id() -> String {
448 use std::collections::hash_map::DefaultHasher;
449 use std::hash::{Hash, Hasher};
450 use std::time::SystemTime;
451
452 let mut hasher = DefaultHasher::new();
453 SystemTime::now().hash(&mut hasher);
454 std::process::id().hash(&mut hasher);
455 format!("{:08x}", hasher.finish() as u32)
456}
457
458fn is_process_alive(pid: u32) -> bool {
459 #[cfg(unix)]
460 {
461 std::process::Command::new("kill")
462 .args(["-0", &pid.to_string()])
463 .output()
464 .map(|o| o.status.success())
465 .unwrap_or(false)
466 }
467 #[cfg(not(unix))]
468 {
469 let _ = pid;
470 true
471 }
472}
473
474struct FileLock {
475 path: PathBuf,
476}
477
478impl FileLock {
479 fn acquire(path: &std::path::Path) -> Result<Self, String> {
480 for _ in 0..50 {
481 match std::fs::OpenOptions::new()
482 .write(true)
483 .create_new(true)
484 .open(path)
485 {
486 Ok(_) => {
487 return Ok(Self {
488 path: path.to_path_buf(),
489 })
490 }
491 Err(_) => {
492 if let Ok(metadata) = std::fs::metadata(path) {
493 if let Ok(modified) = metadata.modified() {
494 if modified.elapsed().unwrap_or_default().as_secs() > 5 {
495 let _ = std::fs::remove_file(path);
496 continue;
497 }
498 }
499 }
500 std::thread::sleep(std::time::Duration::from_millis(100));
501 }
502 }
503 }
504 Err("Could not acquire lock after 5 seconds".to_string())
505 }
506}
507
508impl Drop for FileLock {
509 fn drop(&mut self) {
510 let _ = std::fs::remove_file(&self.path);
511 }
512}
513
514#[cfg(test)]
515mod tests {
516 use super::*;
517
518 #[test]
519 fn register_and_list() {
520 let mut reg = AgentRegistry::new();
521 let id = reg.register("cursor", Some("dev"), "/tmp/project");
522 assert!(!id.is_empty());
523 assert_eq!(reg.list_active(None).len(), 1);
524 assert_eq!(reg.list_active(None)[0].agent_type, "cursor");
525 }
526
527 #[test]
528 fn reregister_same_pid() {
529 let mut reg = AgentRegistry::new();
530 let id1 = reg.register("cursor", Some("dev"), "/tmp/project");
531 let id2 = reg.register("cursor", Some("review"), "/tmp/project");
532 assert_eq!(id1, id2);
533 assert_eq!(reg.agents.len(), 1);
534 assert_eq!(reg.agents[0].role, Some("review".to_string()));
535 }
536
537 #[test]
538 fn post_and_read_messages() {
539 let mut reg = AgentRegistry::new();
540 reg.post_message("agent-a", None, "finding", "Found a bug in auth.rs");
541 reg.post_message("agent-b", Some("agent-a"), "request", "Please review");
542
543 let msgs = reg.read_unread("agent-a");
544 assert_eq!(msgs.len(), 1);
545 assert_eq!(msgs[0].category, "request");
546 }
547
548 #[test]
549 fn set_status() {
550 let mut reg = AgentRegistry::new();
551 let id = reg.register("claude", None, "/tmp/project");
552 reg.set_status(&id, AgentStatus::Idle, Some("waiting for review"));
553 assert_eq!(reg.agents[0].status, AgentStatus::Idle);
554 assert_eq!(
555 reg.agents[0].status_message,
556 Some("waiting for review".to_string())
557 );
558 }
559
560 #[test]
561 fn broadcast_message() {
562 let mut reg = AgentRegistry::new();
563 reg.post_message("agent-a", None, "status", "Starting refactor");
564
565 let msgs_b = reg.read_unread("agent-b");
566 assert_eq!(msgs_b.len(), 1);
567 assert_eq!(msgs_b[0].message, "Starting refactor");
568
569 let msgs_a = reg.read_unread("agent-a");
570 assert!(msgs_a.is_empty());
571 }
572
573 #[test]
574 fn diary_add_and_format() {
575 let mut diary = AgentDiary::new("test-agent-001", "cursor", "/tmp/project");
576 diary.add_entry(
577 DiaryEntryType::Discovery,
578 "Found auth module at src/auth.rs",
579 Some("auth"),
580 );
581 diary.add_entry(
582 DiaryEntryType::Decision,
583 "Use JWT RS256 for token signing",
584 None,
585 );
586 diary.add_entry(
587 DiaryEntryType::Progress,
588 "Implemented login endpoint",
589 Some("auth"),
590 );
591
592 assert_eq!(diary.entries.len(), 3);
593
594 let summary = diary.format_summary();
595 assert!(summary.contains("test-agent-001"));
596 assert!(summary.contains("FOUND"));
597 assert!(summary.contains("DECIDED"));
598 assert!(summary.contains("DONE"));
599 }
600
601 #[test]
602 fn diary_compact_format() {
603 let mut diary = AgentDiary::new("test-agent-002", "claude", "/tmp/project");
604 diary.add_entry(DiaryEntryType::Insight, "DB queries are N+1", None);
605 diary.add_entry(
606 DiaryEntryType::Blocker,
607 "Missing API credentials",
608 Some("deploy"),
609 );
610
611 let compact = diary.format_compact();
612 assert!(compact.contains("diary:test-agent-002"));
613 assert!(compact.contains("B:Missing API credentials"));
614 assert!(compact.contains("I:DB queries are N+1"));
615 }
616
617 #[test]
618 fn diary_entry_types() {
619 let types = vec![
620 DiaryEntryType::Discovery,
621 DiaryEntryType::Decision,
622 DiaryEntryType::Blocker,
623 DiaryEntryType::Progress,
624 DiaryEntryType::Insight,
625 ];
626 for t in types {
627 assert!(!format!("{}", t).is_empty());
628 }
629 }
630
631 #[test]
632 fn diary_truncation() {
633 let mut diary = AgentDiary::new("test-agent", "cursor", "/tmp");
634 for i in 0..150 {
635 diary.add_entry(DiaryEntryType::Progress, &format!("Step {i}"), None);
636 }
637 assert!(diary.entries.len() <= 100);
638 }
639}