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