1use chrono::Utc;
2use std::collections::HashSet;
3use std::path::Path;
4
5#[cfg(test)]
6use super::diary::{AgentDiary, DiaryEntryType};
7use super::persistence::{
8 FileLock, ProcessIdentityIndex, agents_dir, create_agents_dir, generate_short_id,
9 load_registry_file, mutate_persistent, save_registry_file,
10};
11use super::{
12 AgentEntry, AgentRegistry, AgentStatus, LogicalSessionPresence, REQUEST_DIGEST_METADATA_KEY,
13 REQUEST_ID_METADATA_KEY, RoleSource, ScratchpadEntry, TransportClient,
14};
15use crate::core::a2a::message::{MessagePriority, PrivacyLevel};
16
17const PLACEHOLDER_ROLE: &str = "context-engine";
19const WORKER_LEASE_EXPIRED: &str = "worker lease expired; register again before continuing";
20const TRANSPORT_IDLE: &str = "connected; no tool call within the worker lease";
23
24#[derive(Default)]
26struct RegistrationMeta {
27 role_source: Option<RoleSource>,
28 transport: Option<TransportClient>,
29}
30
31fn explicit_role_source(role: Option<&str>) -> Option<RoleSource> {
34 match role {
35 None => None,
36 Some(PLACEHOLDER_ROLE) => Some(RoleSource::Assumed),
37 Some(_) => Some(RoleSource::Declared),
38 }
39}
40
41static RENEWED: std::sync::Mutex<std::collections::BTreeSet<String>> =
44 std::sync::Mutex::new(std::collections::BTreeSet::new());
45
46fn note_renewed(agent_id: &str) {
47 if let Ok(mut renewed) = RENEWED.lock() {
48 renewed.insert(agent_id.to_string());
49 }
50 crate::core::agents::friction::record_for(
51 crate::core::agents::friction::REGISTRATION_RENEWED,
52 Some(agent_id),
53 None,
54 );
55}
56
57pub fn take_renewed(agent_id: &str) -> bool {
60 RENEWED
61 .lock()
62 .is_ok_and(|mut renewed| renewed.remove(agent_id))
63}
64
65const WAIT_POLL: std::time::Duration = std::time::Duration::from_millis(500);
67const WAIT_KEEPALIVE: std::time::Duration = std::time::Duration::from_secs(30);
69const WAIT_RACY_MTIME: std::time::Duration = std::time::Duration::from_secs(3);
73
74const BROADCAST_GRACE: chrono::Duration = chrono::Duration::minutes(10);
78
79fn unread_for(
85 e: &ScratchpadEntry,
86 agent_id: &str,
87 project_root: Option<&str>,
88 now: chrono::DateTime<Utc>,
89 floor: Option<chrono::DateTime<Utc>>,
90) -> bool {
91 !e.read_by.iter().any(|reader| reader == agent_id)
92 && e.from_agent != agent_id
93 && (e.to_agent.as_deref() == Some(agent_id)
94 || (e.to_agent.is_none() && floor.is_none_or(|floor| e.timestamp >= floor)))
95 && project_root.is_none_or(|root| {
96 e.project_root
97 .as_deref()
98 .is_some_and(|own| crate::core::pathutil::same_project(own, root))
99 })
100 && e.expires_at.is_none_or(|exp| exp > now)
101}
102
103fn owned_by_current_process(agent: &AgentEntry) -> bool {
105 let pid = std::process::id();
106 agent.pid == pid
107 && agent.process_identity.is_some()
108 && crate::ipc::process::identity(pid).as_ref() == agent.process_identity.as_ref()
109}
110
111const LOGICAL_SESSION_SOURCE_MAX_BYTES: usize = 64;
112const LOGICAL_SESSION_WORKSPACE_MAX_BYTES: usize = 4096;
113const LOGICAL_SESSION_ID_MAX_BYTES: usize = 256;
114const MAX_RETAINED_FINISHED_AGENTS: usize = 32;
115const DEFAULT_ACTIVE_WORKER_LEASE_SECONDS: i64 = 120;
119
120fn active_worker_lease_seconds() -> i64 {
121 let configured = crate::core::config::Config::load()
122 .agents
123 .active_worker_lease_seconds;
124 i64::try_from(configured)
125 .unwrap_or(DEFAULT_ACTIVE_WORKER_LEASE_SECONDS)
126 .max(1)
127}
128
129fn presence_ttl() -> u64 {
130 crate::core::config::Config::load()
131 .agents
132 .presence_ttl_hours
133}
134
135fn max_scratchpad() -> usize {
136 crate::core::config::Config::load()
137 .agents
138 .max_scratchpad_entries
139}
140
141pub(crate) fn process_identity_matches(
142 agent: &AgentEntry,
143 compatibility_identities: &ProcessIdentityIndex,
144) -> bool {
145 agent
146 .process_identity
147 .as_ref()
148 .or_else(|| compatibility_identities.get(&agent.agent_id))
149 .is_some_and(|identity| crate::ipc::process::matches_identity(agent.pid, identity))
150}
151
152fn safe_legacy_identity(agent: &AgentEntry) -> Option<crate::ipc::process::ProcessIdentity> {
153 if agent.process_identity.is_some() {
154 return None;
155 }
156 let identity = crate::ipc::process::identity(agent.pid)?;
157 legacy_identity_matches_registration(agent, &identity).then_some(identity)
158}
159
160fn legacy_identity_matches_registration(
161 agent: &AgentEntry,
162 identity: &crate::ipc::process::ProcessIdentity,
163) -> bool {
164 if Path::new(&identity.executable)
165 .file_name()
166 .and_then(|name| name.to_str())
167 != Some("lean-ctx")
168 {
169 return false;
170 }
171 let started_at = agent.started_at.timestamp_micros();
172 let Ok(process_started_at) = i64::try_from(identity.start_marker) else {
173 return false;
174 };
175 const REGISTRATION_GRACE_MICROS: i64 = 5 * 60 * 1_000_000;
176 process_started_at <= started_at
177 && started_at.saturating_sub(process_started_at) <= REGISTRATION_GRACE_MICROS
178}
179
180fn is_recoverable_legacy_finished(agent: &AgentEntry) -> bool {
181 agent.process_identity.is_none()
182 && agent.status == AgentStatus::Finished
183 && agent.status_message.as_deref() == Some("process identity no longer matches")
184}
185
186impl AgentRegistry {
187 pub fn new() -> Self {
188 Self {
189 agents: Vec::new(),
190 scratchpad: Vec::new(),
191 logical_sessions: Vec::new(),
192 logical_session_telemetry_seen: false,
193 updated_at: Utc::now(),
194 }
195 }
196
197 pub fn register(
198 &mut self,
199 agent_type: &str,
200 role: Option<&str>,
201 project_root: &str,
202 durable_identity_id: Option<&str>,
203 ) -> Result<String, String> {
204 self.register_process(
205 agent_type,
206 role,
207 project_root,
208 durable_identity_id,
209 std::process::id(),
210 )
211 }
212
213 fn register_process(
214 &mut self,
215 agent_type: &str,
216 role: Option<&str>,
217 project_root: &str,
218 durable_identity_id: Option<&str>,
219 pid: u32,
220 ) -> Result<String, String> {
221 self.register_process_with(
222 agent_type,
223 role,
224 project_root,
225 durable_identity_id,
226 pid,
227 RegistrationMeta {
228 role_source: explicit_role_source(role),
229 transport: None,
230 },
231 )
232 }
233
234 pub fn register_child(
239 &mut self,
240 agent_type: &str,
241 role: &str,
242 project_root: &str,
243 parent_id: &str,
244 ) -> Result<String, String> {
245 let pid = std::process::id();
246 let identity = crate::ipc::process::identity(pid)
247 .ok_or_else(|| format!("cannot establish immutable process identity for PID {pid}"))?;
248 if !self.held_by_this_process(parent_id) {
249 return Err(format!(
250 "agent {parent_id} is not registered by this process"
251 ));
252 }
253 let now = Utc::now();
254 if let Some(child) = self.agents.iter_mut().find(|a| {
255 a.parent.as_deref() == Some(parent_id)
256 && a.role.as_deref() == Some(role)
257 && a.status != AgentStatus::Finished
258 }) {
259 child.last_active = now;
260 child.status = AgentStatus::Active;
261 child.project_root = project_root.to_string();
262 return Ok(child.agent_id.clone());
263 }
264 let agent_id = format!("{}-{}-{}", agent_type, pid, generate_short_id());
265 self.agents.push(AgentEntry {
266 agent_id: agent_id.clone(),
267 agent_type: agent_type.to_string(),
268 role: Some(role.to_string()),
269 project_root: project_root.to_string(),
270 started_at: now,
271 last_active: now,
272 pid,
273 process_identity: Some(identity),
274 durable_identity_id: None,
275 status: AgentStatus::Active,
276 status_message: None,
277 transport: None,
278 role_source: Some(RoleSource::Declared),
279 parent: Some(parent_id.to_string()),
280 app_session_id: crate::core::activity::session_id_from_env(),
281 display_name: crate::core::activity::session_name_from_env(),
282 });
283 self.updated_at = now;
284 crate::core::events::emit_agent_action(&agent_id, "register", None);
285 Ok(agent_id)
286 }
287
288 fn broadcast_floor(&self, agent_id: &str) -> Option<chrono::DateTime<Utc>> {
291 self.agents
292 .iter()
293 .find(|a| a.agent_id == agent_id)
294 .map(|a| a.started_at - BROADCAST_GRACE)
295 }
296
297 pub fn held_by_this_process(&self, agent_id: &str) -> bool {
300 let pid = std::process::id();
301 let identity = crate::ipc::process::identity(pid);
302 self.agents.iter().any(|a| {
303 a.agent_id == agent_id
304 && a.pid == pid
305 && a.status != AgentStatus::Finished
306 && a.process_identity.is_some()
307 && a.process_identity == identity
308 })
309 }
310
311 pub fn set_display_name(&mut self, agent_id: &str, name: &str) -> bool {
314 match self.agents.iter_mut().find(|a| a.agent_id == agent_id) {
315 Some(agent) => {
316 agent.display_name = Some(name.to_owned());
317 true
318 }
319 None => false,
320 }
321 }
322
323 pub fn declared_role(&self, agent_id: &str) -> Option<String> {
325 self.agents
326 .iter()
327 .find(|a| a.agent_id == agent_id && a.role_source == Some(RoleSource::Declared))
328 .and_then(|a| a.role.clone())
329 }
330
331 pub fn register_transport(
334 &mut self,
335 role: &str,
336 role_source: RoleSource,
337 project_root: &str,
338 durable_identity_id: Option<&str>,
339 transport: TransportClient,
340 ) -> Result<String, String> {
341 self.register_process_with(
342 "mcp",
343 Some(role),
344 project_root,
345 durable_identity_id,
346 std::process::id(),
347 RegistrationMeta {
348 role_source: Some(role_source),
349 transport: Some(transport),
350 },
351 )
352 }
353
354 fn register_process_with(
355 &mut self,
356 agent_type: &str,
357 role: Option<&str>,
358 project_root: &str,
359 durable_identity_id: Option<&str>,
360 pid: u32,
361 meta: RegistrationMeta,
362 ) -> Result<String, String> {
363 if let Some(identity_id) = durable_identity_id {
364 crate::core::agent_registry::require_active(identity_id)?;
365 }
366 let identity = crate::ipc::process::identity(pid)
367 .ok_or_else(|| format!("cannot establish immutable process identity for PID {pid}"))?;
368 let agent_id = format!("{}-{}-{}", agent_type, pid, generate_short_id());
369
370 if let Some(existing_index) = self.agents.iter().position(|a| {
371 a.pid == pid
372 && a.parent.is_none()
373 && a.status != AgentStatus::Finished
374 && a.process_identity.as_ref() == Some(&identity)
375 }) {
376 let existing_binding = self.agents[existing_index].durable_identity_id.clone();
377 if let (Some(existing_id), Some(requested_id)) =
378 (existing_binding.as_deref(), durable_identity_id)
379 && existing_id != requested_id
380 {
381 return Err(format!(
382 "agent presence '{}' is already bound to durable identity '{}'; refusing rebind to '{}'",
383 self.agents[existing_index].agent_id, existing_id, requested_id
384 ));
385 }
386 if existing_binding.is_some() {
387 self.require_active_for_presence(
388 &self.agents[existing_index].agent_id,
389 "registration",
390 )?;
391 }
392 let existing = &mut self.agents[existing_index];
393 existing.last_active = Utc::now();
394 existing.status = AgentStatus::Active;
395 if existing.status_message.as_deref() == Some(TRANSPORT_IDLE) {
396 existing.status_message = None;
397 }
398 if meta.transport.is_none() || existing.agent_type == "mcp" {
401 existing.agent_type = agent_type.to_string();
402 }
403 existing.project_root = project_root.to_string();
404 match role {
410 Some(PLACEHOLDER_ROLE) if existing.role.is_some() => {}
411 Some(_)
412 if meta.role_source == Some(RoleSource::Assumed)
413 && existing.role_source == Some(RoleSource::Declared) => {}
414 Some(r) => {
415 existing.role = Some(r.to_string());
416 existing.role_source = meta.role_source;
417 }
418 None => {}
419 }
420 if let Some(transport) = meta.transport
422 && (transport.is_identified() || existing.transport.is_none())
423 {
424 existing.transport = Some(transport);
425 }
426 if existing.durable_identity_id.is_none() {
427 existing.durable_identity_id = durable_identity_id.map(str::to_string);
428 }
429 if existing.app_session_id.is_none() {
430 existing.app_session_id = crate::core::activity::session_id_from_env();
431 }
432 if let Some(name) = crate::core::activity::session_name_from_env() {
433 existing.display_name = Some(name);
434 }
435 return Ok(self.adopt_app_session_identity(existing_index));
436 }
437
438 for existing in self
442 .agents
443 .iter_mut()
444 .filter(|agent| agent.pid == pid && agent.status != AgentStatus::Finished)
445 {
446 existing.status = AgentStatus::Finished;
447 existing.status_message =
448 Some("superseded by a different process identity".to_string());
449 }
450
451 self.agents.push(AgentEntry {
456 agent_id: agent_id.clone(),
457 agent_type: agent_type.to_string(),
458 role: role.map(std::string::ToString::to_string),
459 project_root: project_root.to_string(),
460 started_at: Utc::now(),
461 last_active: Utc::now(),
462 pid,
463 process_identity: Some(identity),
464 durable_identity_id: durable_identity_id.map(str::to_string),
465 status: AgentStatus::Active,
466 status_message: None,
467 transport: meta.transport,
468 role_source: meta.role_source,
469 parent: None,
470 app_session_id: crate::core::activity::session_id_from_env(),
471 display_name: crate::core::activity::session_name_from_env(),
472 });
473 let agent_id = self.adopt_app_session_identity(self.agents.len() - 1);
474
475 self.updated_at = Utc::now();
476 crate::core::events::emit_agent_action(&agent_id, "register", None);
477 Ok(agent_id)
478 }
479
480 fn adopt_app_session_identity(&mut self, index: usize) -> String {
489 let current = &self.agents[index];
490 let (Some(app), Some(role)) = (current.app_session_id.clone(), current.role.clone()) else {
491 return current.agent_id.clone();
492 };
493 if role == PLACEHOLDER_ROLE || current.parent.is_some() {
494 return current.agent_id.clone();
495 }
496 let binding = current.durable_identity_id.clone();
497 let earlier = self.agents.iter().enumerate().position(|(i, a)| {
498 i != index
499 && a.parent.is_none()
500 && a.app_session_id.as_deref() == Some(app.as_str())
501 && a.role.as_deref() == Some(role.as_str())
502 && (a.durable_identity_id.is_none() || a.durable_identity_id == binding)
503 && (a.status == AgentStatus::Finished
504 || a.process_identity.as_ref().is_some_and(|identity| {
505 !crate::ipc::process::matches_identity(a.pid, identity)
506 }))
507 });
508 let Some(earlier) = earlier else {
509 return self.agents[index].agent_id.clone();
510 };
511 let old = self.agents.remove(earlier);
512 let index = if earlier < index { index - 1 } else { index };
513 let current = &mut self.agents[index];
514 current.agent_id = old.agent_id;
515 current.started_at = old.started_at;
516 if current.durable_identity_id.is_none() {
517 current.durable_identity_id = old.durable_identity_id;
518 }
519 current.agent_id.clone()
520 }
521
522 pub fn register_mcp_process(
524 project_root: &str,
525 durable_identity_id: Option<&str>,
526 ) -> Result<String, String> {
527 Self::register_mcp_process_as(
528 project_root,
529 PLACEHOLDER_ROLE,
530 RoleSource::Assumed,
531 durable_identity_id,
532 )
533 }
534
535 pub fn register_mcp_process_as(
540 project_root: &str,
541 role: &str,
542 role_source: RoleSource,
543 durable_identity_id: Option<&str>,
544 ) -> Result<String, String> {
545 let transport = TransportClient::current();
546 mutate_persistent(|registry| {
547 registry.cleanup_stale(presence_ttl());
548 registry.register_transport(
549 role,
550 role_source,
551 project_root,
552 durable_identity_id,
553 transport,
554 )
555 })
556 .and_then(|result| result)
557 }
558
559 pub fn record_transport_client_persistent(
563 transport: &TransportClient,
564 ) -> Result<usize, String> {
565 mutate_persistent(|registry| registry.record_transport_client(transport))
566 }
567
568 pub(crate) fn record_transport_client(&mut self, transport: &TransportClient) -> usize {
569 let mut updated = 0;
570 for agent in self.agents.iter_mut().filter(|agent| {
571 agent.transport.is_some()
572 && agent.status != AgentStatus::Finished
573 && owned_by_current_process(agent)
574 }) {
575 agent.transport = Some(transport.clone());
576 updated += 1;
577 }
578 if updated > 0 {
579 self.updated_at = Utc::now();
580 }
581 updated
582 }
583
584 pub fn transport_presence_is_live(agent_id: &str) -> bool {
588 Self::load().is_some_and(|registry| {
589 registry.agents.iter().any(|agent| {
590 agent.agent_id == agent_id
591 && agent.status != AgentStatus::Finished
592 && owned_by_current_process(agent)
593 })
594 })
595 }
596
597 pub fn any_live_mcp_transport() -> bool {
603 let Some(registry) = Self::load() else {
604 return false;
605 };
606 let identities = ProcessIdentityIndex::load();
607 registry.agents.iter().any(|agent| {
608 agent.transport.is_some()
609 && agent.status != AgentStatus::Finished
610 && process_identity_matches(agent, &identities)
611 })
612 }
613
614 pub fn heartbeat_persistent(agent_id: &str) -> Result<(), String> {
619 mutate_persistent(|registry| registry.update_transport_heartbeat(agent_id))?
620 }
621
622 pub(crate) fn update_transport_heartbeat(&mut self, agent_id: &str) -> Result<(), String> {
623 if self.revive_lease_expired(agent_id) {
624 note_renewed(agent_id);
625 }
626 self.update_heartbeat(agent_id)
627 }
628
629 fn revive_lease_expired(&mut self, agent_id: &str) -> bool {
633 let Some(agent) = self
634 .agents
635 .iter_mut()
636 .find(|agent| agent.agent_id == agent_id)
637 else {
638 return false;
639 };
640 if agent.status == AgentStatus::Finished
641 && agent.status_message.as_deref() == Some(WORKER_LEASE_EXPIRED)
642 && agent.durable_identity_id.is_none()
643 && owned_by_current_process(agent)
644 {
645 agent.status = AgentStatus::Active;
646 agent.status_message = None;
647 return true;
648 }
649 false
650 }
651
652 pub fn renew_worker(&mut self, agent_id: &str) -> Result<bool, String> {
657 let revived = self.revive_lease_expired(agent_id);
658 self.update_heartbeat(agent_id)?;
659 if revived {
660 note_renewed(agent_id);
661 request_announce();
662 }
663 Ok(revived)
664 }
665
666 pub fn renew_worker_persistent(agent_id: &str) -> Result<bool, String> {
668 mutate_persistent(|registry| registry.renew_worker(agent_id))?
669 }
670
671 pub fn finish_persistent(agent_id: &str) -> Result<(), String> {
673 mutate_persistent(|registry| {
674 registry.set_status(agent_id, AgentStatus::Finished, Some("connection closed"))
675 })?
676 }
677
678 pub fn update_heartbeat(&mut self, agent_id: &str) -> Result<(), String> {
679 let index = self
680 .agents
681 .iter()
682 .position(|agent| agent.agent_id == agent_id)
683 .ok_or_else(|| format!("agent presence '{agent_id}' was not found"))?;
684 if self.agents[index].status == AgentStatus::Finished {
685 return Err(format!(
686 "agent presence '{agent_id}' is finished; register a new process presence"
687 ));
688 }
689 self.require_active_for_presence(agent_id, "heartbeat")?;
690 let agent = &mut self.agents[index];
691 let pid = std::process::id();
692 if agent.pid != pid {
693 return Err(format!(
694 "agent presence '{agent_id}' belongs to PID {}; heartbeat came from PID {pid}",
695 agent.pid
696 ));
697 }
698 let identity = crate::ipc::process::identity(pid)
699 .ok_or_else(|| format!("cannot establish immutable process identity for PID {pid}"))?;
700 if let Some(expected) = &agent.process_identity
701 && expected != &identity
702 {
703 return Err(format!(
704 "agent presence '{agent_id}' process identity no longer matches"
705 ));
706 }
707 agent.process_identity = Some(identity);
710 agent.status = AgentStatus::Active;
711 if agent.status_message.as_deref() == Some(TRANSPORT_IDLE) {
712 agent.status_message = None;
713 }
714 agent.last_active = Utc::now();
715 Ok(())
716 }
717
718 pub fn set_status(
719 &mut self,
720 agent_id: &str,
721 status: AgentStatus,
722 message: Option<&str>,
723 ) -> Result<(), String> {
724 let index = self
725 .agents
726 .iter()
727 .position(|agent| agent.agent_id == agent_id)
728 .ok_or_else(|| format!("agent presence '{agent_id}' was not found"))?;
729 if status != AgentStatus::Finished {
732 self.require_active_for_presence(agent_id, "status transition")?;
733 }
734 let agent = &mut self.agents[index];
735 agent.status = status;
736 agent.status_message = message.map(std::string::ToString::to_string);
737 agent.last_active = Utc::now();
738 self.updated_at = Utc::now();
739 Ok(())
740 }
741 pub(crate) fn open_or_heartbeat_logical_session(
744 &mut self,
745 source: &str,
746 workspace: &str,
747 session_id: &str,
748 ) {
749 let now = Utc::now();
750 self.logical_session_telemetry_seen = true;
751 if let Some(session) = self.logical_sessions.iter_mut().find(|session| {
752 session.source == source
753 && session.workspace == workspace
754 && session.session_id == session_id
755 }) {
756 session.last_heartbeat = now;
757 } else {
758 self.logical_sessions.push(LogicalSessionPresence {
759 source: source.to_string(),
760 workspace: workspace.to_string(),
761 session_id: session_id.to_string(),
762 opened_at: now,
763 last_heartbeat: now,
764 });
765 }
766 self.updated_at = now;
767 }
768
769 pub(crate) fn close_logical_session(
770 &mut self,
771 source: &str,
772 workspace: &str,
773 session_id: &str,
774 ) -> bool {
775 self.logical_session_telemetry_seen = true;
776 let previous_len = self.logical_sessions.len();
777 self.logical_sessions.retain(|session| {
778 session.source != source
779 || session.workspace != workspace
780 || session.session_id != session_id
781 });
782 let removed = self.logical_sessions.len() != previous_len;
783 self.updated_at = Utc::now();
784 removed
785 }
786
787 pub fn cleanup_stale_logical_sessions(&mut self, max_age_seconds: u64) {
788 let seconds = i64::try_from(max_age_seconds).unwrap_or(i64::MAX);
789 let cutoff = Utc::now() - chrono::Duration::seconds(seconds);
790 self.logical_sessions
791 .retain(|session| session.last_heartbeat >= cutoff);
792 self.updated_at = Utc::now();
793 }
794
795 pub fn record_logical_session_presence(
796 event: &str,
797 source: &str,
798 workspace: &str,
799 session_id: &str,
800 ) -> Result<(), String> {
801 let valid_field = |value: &str, max_bytes: usize| {
802 !value.is_empty() && value.len() <= max_bytes && !value.chars().any(char::is_control)
803 };
804 if !valid_field(source, LOGICAL_SESSION_SOURCE_MAX_BYTES)
805 || !valid_field(workspace, LOGICAL_SESSION_WORKSPACE_MAX_BYTES)
806 || !valid_field(session_id, LOGICAL_SESSION_ID_MAX_BYTES)
807 {
808 return Err(
809 "presence fields are empty, too long, or contain control characters".to_string(),
810 );
811 }
812 if !matches!(event, "open" | "heartbeat" | "close") {
813 return Err("event must be open, heartbeat, or close".to_string());
814 }
815
816 let ttl = crate::core::config::Config::load()
817 .agents
818 .logical_session_ttl_seconds;
819 mutate_persistent(|registry| {
820 registry.cleanup_stale_logical_sessions(ttl);
821 match event {
822 "open" | "heartbeat" => {
823 registry.open_or_heartbeat_logical_session(source, workspace, session_id);
824 }
825 "close" => {
826 registry.close_logical_session(source, workspace, session_id);
827 }
828 _ => unreachable!("event validated above"),
829 }
830 })
831 }
832
833 pub fn require_active_for_presence(
837 &self,
838 agent_id: &str,
839 operation: &str,
840 ) -> Result<(), String> {
841 let Some(agent) = self.agents.iter().find(|agent| agent.agent_id == agent_id) else {
842 return Ok(());
843 };
844 let Some(durable_id) = agent.durable_identity_id.as_deref() else {
845 return Ok(());
846 };
847 if agent.status == AgentStatus::Finished {
848 return Err(format!(
849 "agent presence '{agent_id}' is finished; cannot {operation}"
850 ));
851 }
852 crate::core::agent_registry::require_active(durable_id).map_err(|error| {
853 format!("agent presence '{agent_id}' is not authorized for {operation}: {error}")
854 })
855 }
856
857 pub fn list_active(&self, project_root: Option<&str>) -> Vec<&AgentEntry> {
858 let compatibility_identities = ProcessIdentityIndex::load();
859 self.agents
860 .iter()
861 .filter(|a| {
862 if let Some(root) = project_root {
863 crate::core::pathutil::same_project(&a.project_root, root)
864 && a.status != AgentStatus::Finished
865 && process_identity_matches(a, &compatibility_identities)
866 && self
867 .require_active_for_presence(&a.agent_id, "listing")
868 .is_ok()
869 } else {
870 a.status != AgentStatus::Finished
871 && process_identity_matches(a, &compatibility_identities)
872 && self
873 .require_active_for_presence(&a.agent_id, "listing")
874 .is_ok()
875 }
876 })
877 .collect()
878 }
879
880 pub fn list_all(&self) -> &[AgentEntry] {
881 &self.agents
882 }
883
884 pub(crate) fn post_message(
885 &mut self,
886 from_agent: &str,
887 to_agent: Option<&str>,
888 category: &str,
889 message: &str,
890 ) -> Result<String, String> {
891 self.post_message_full(
892 from_agent,
893 to_agent,
894 category,
895 message,
896 PrivacyLevel::default(),
897 MessagePriority::default(),
898 None,
899 )
900 }
901
902 pub(crate) fn post_message_full(
903 &mut self,
904 from_agent: &str,
905 to_agent: Option<&str>,
906 category: &str,
907 message: &str,
908 privacy: PrivacyLevel,
909 priority: MessagePriority,
910 ttl_hours: Option<u64>,
911 ) -> Result<String, String> {
912 self.post_message_scoped(
913 None, from_agent, to_agent, category, message, privacy, priority, ttl_hours,
914 )
915 }
916
917 #[allow(clippy::too_many_arguments)]
918 pub fn post_message_scoped(
919 &mut self,
920 project_root: Option<&str>,
921 from_agent: &str,
922 to_agent: Option<&str>,
923 category: &str,
924 message: &str,
925 privacy: PrivacyLevel,
926 priority: MessagePriority,
927 ttl_hours: Option<u64>,
928 ) -> Result<String, String> {
929 self.post_message_keyed(
930 project_root,
931 from_agent,
932 to_agent,
933 category,
934 message,
935 privacy,
936 priority,
937 ttl_hours,
938 None,
939 )
940 }
941
942 #[allow(clippy::too_many_arguments)]
949 pub fn post_message_keyed(
950 &mut self,
951 project_root: Option<&str>,
952 from_agent: &str,
953 to_agent: Option<&str>,
954 category: &str,
955 message: &str,
956 privacy: PrivacyLevel,
957 priority: MessagePriority,
958 ttl_hours: Option<u64>,
959 request_id: Option<&str>,
960 ) -> Result<String, String> {
961 self.require_active_for_presence(from_agent, "posting a message")?;
962 if let Some(target) = to_agent {
963 self.require_active_for_presence(target, "receiving a message")?;
964 }
965 let request = match request_id {
966 Some(key) => {
967 super::check_request_id(key)?;
968 let digest = super::request_digest(
969 "leanctx/agent-post/v1",
970 &serde_json::json!([
971 project_root,
972 to_agent,
973 category,
974 message,
975 privacy,
976 priority,
977 ttl_hours
978 ]),
979 );
980 let now = Utc::now();
981 if let Some(entry) = self.scratchpad.iter().find(|entry| {
982 entry.from_agent == from_agent
983 && entry.project_root.as_deref() == project_root
984 && entry.expires_at.is_none_or(|expires| expires > now)
985 && entry
986 .metadata
987 .get(REQUEST_ID_METADATA_KEY)
988 .map(String::as_str)
989 == Some(key)
990 }) {
991 return if entry.metadata.get(REQUEST_DIGEST_METADATA_KEY) == Some(&digest) {
992 Ok(entry.id.clone())
993 } else {
994 Err(super::request_id_conflict())
995 };
996 }
997 Some((key, digest))
998 }
999 None => None,
1000 };
1001 let (message, mut metadata) =
1005 ScratchpadEntry::admit_message(from_agent, project_root, message)?;
1006 if let Some((key, digest)) = request {
1007 metadata.insert(REQUEST_ID_METADATA_KEY.to_owned(), key.to_owned());
1008 metadata.insert(REQUEST_DIGEST_METADATA_KEY.to_owned(), digest);
1009 }
1010 let id = generate_short_id();
1011 let default_ttl_hours = crate::core::config::Config::load()
1012 .agents
1013 .scratchpad_default_ttl_hours;
1014 let expires_at = Some(match ttl_hours {
1015 Some(hours) => Utc::now() + chrono::Duration::hours(hours as i64),
1016 None => Utc::now() + chrono::Duration::hours(default_ttl_hours as i64),
1017 });
1018 self.scratchpad.push(ScratchpadEntry {
1019 id: id.clone(),
1020 from_agent: from_agent.to_string(),
1021 to_agent: to_agent.map(std::string::ToString::to_string),
1022 task_id: None,
1023 category: category.to_string(),
1024 priority,
1025 privacy,
1026 message,
1027 metadata,
1028 project_root: project_root.map(std::string::ToString::to_string),
1029 timestamp: Utc::now(),
1030 read_by: vec![from_agent.to_string()],
1031 expires_at,
1032 });
1033
1034 let max_scratchpad_entries = max_scratchpad();
1035 if self.scratchpad.len() > max_scratchpad_entries {
1036 self.scratchpad
1037 .drain(0..self.scratchpad.len() - max_scratchpad_entries);
1038 }
1039
1040 self.updated_at = Utc::now();
1041 Ok(id)
1042 }
1043
1044 #[cfg(test)]
1045 pub(crate) fn read_unread(&mut self, agent_id: &str) -> Result<Vec<&ScratchpadEntry>, String> {
1046 self.read_unread_for_project(agent_id, None)
1047 }
1048
1049 #[cfg(test)]
1050 pub(crate) fn read_unread_scoped(
1051 &mut self,
1052 agent_id: &str,
1053 project_root: &str,
1054 ) -> Result<Vec<&ScratchpadEntry>, String> {
1055 self.read_unread_for_project(agent_id, Some(project_root))
1056 }
1057
1058 pub fn read_unread_page(
1064 &mut self,
1065 agent_id: &str,
1066 project_root: &str,
1067 budget: usize,
1068 render: impl Fn(&ScratchpadEntry) -> String,
1069 ) -> Result<(Vec<String>, usize), String> {
1070 self.require_active_for_presence(agent_id, "reading messages")?;
1071 let now = Utc::now();
1072 let floor = self.broadcast_floor(agent_id);
1073 let unread: Vec<usize> = self
1074 .scratchpad
1075 .iter()
1076 .enumerate()
1077 .filter(|(_, e)| unread_for(e, agent_id, Some(project_root), now, floor))
1078 .map(|(i, _)| i)
1079 .collect();
1080 let (mut shown, mut used) = (Vec::new(), 0_usize);
1081 for &index in &unread {
1082 let line = render(&self.scratchpad[index]);
1083 if !shown.is_empty() && used + line.len() > budget {
1084 break;
1085 }
1086 used += line.len();
1087 shown.push(line);
1088 self.scratchpad[index].read_by.push(agent_id.to_string());
1089 }
1090 if !shown.is_empty() {
1091 self.updated_at = now;
1092 }
1093 let left = unread.len() - shown.len();
1094 Ok((shown, left))
1095 }
1096
1097 pub(crate) fn has_unread_scoped(&self, agent_id: &str, project_root: &str) -> bool {
1100 let now = Utc::now();
1101 let floor = self.broadcast_floor(agent_id);
1102 self.scratchpad
1103 .iter()
1104 .any(|e| unread_for(e, agent_id, Some(project_root), now, floor))
1105 }
1106
1107 pub fn unread_scoped(&self, agent_id: &str, project_root: &str) -> Vec<&ScratchpadEntry> {
1110 let now = Utc::now();
1111 let floor = self.broadcast_floor(agent_id);
1112 self.scratchpad
1113 .iter()
1114 .filter(|e| unread_for(e, agent_id, Some(project_root), now, floor))
1115 .collect()
1116 }
1117
1118 pub fn mark_read(&mut self, agent_id: &str, ids: &[String]) -> usize {
1121 let mut marked = 0;
1122 for entry in self.scratchpad.iter_mut().filter(|e| ids.contains(&e.id)) {
1123 if !entry.read_by.iter().any(|reader| reader == agent_id) {
1124 entry.read_by.push(agent_id.to_owned());
1125 marked += 1;
1126 }
1127 }
1128 if marked > 0 {
1129 self.updated_at = Utc::now();
1130 }
1131 marked
1132 }
1133
1134 pub fn wait_for_unread(
1139 agent_id: &str,
1140 project_root: &str,
1141 timeout: std::time::Duration,
1142 mut on_poll: impl FnMut(),
1143 ) -> bool {
1144 let path = agents_dir().ok().map(|dir| dir.join("registry.json"));
1145 let modified = || {
1146 path.as_ref()
1147 .and_then(|p| std::fs::metadata(p).ok())
1148 .and_then(|m| m.modified().ok())
1149 };
1150 let deadline = std::time::Instant::now() + timeout;
1151 let mut last_poll = std::time::Instant::now();
1152 let mut seen = None;
1153 loop {
1154 let now_modified = modified();
1155 if seen.is_none() || now_modified != seen.flatten() {
1156 seen = Some(now_modified);
1157 if Self::load().is_some_and(|r| r.has_unread_scoped(agent_id, project_root)) {
1158 return true;
1159 }
1160 let racy = now_modified.is_some_and(|mtime| {
1164 !std::time::SystemTime::now()
1165 .duration_since(mtime)
1166 .is_ok_and(|age| age >= WAIT_RACY_MTIME)
1167 });
1168 if racy {
1169 seen = None;
1170 }
1171 }
1172 let now = std::time::Instant::now();
1173 if now >= deadline {
1174 return false;
1175 }
1176 if now.duration_since(last_poll) >= WAIT_KEEPALIVE {
1177 on_poll();
1178 last_poll = now;
1179 seen = None;
1181 }
1182 std::thread::sleep(WAIT_POLL.min(deadline - now));
1183 }
1184 }
1185
1186 #[cfg(test)]
1187 fn read_unread_for_project(
1188 &mut self,
1189 agent_id: &str,
1190 project_root: Option<&str>,
1191 ) -> Result<Vec<&ScratchpadEntry>, String> {
1192 self.require_active_for_presence(agent_id, "reading messages")?;
1193 let now = Utc::now();
1194 let floor = self.broadcast_floor(agent_id);
1195 let unread_indices: Vec<usize> = self
1196 .scratchpad
1197 .iter()
1198 .enumerate()
1199 .filter(|(_, e)| unread_for(e, agent_id, project_root, now, floor))
1200 .map(|(i, _)| i)
1201 .collect();
1202
1203 for i in &unread_indices {
1204 self.scratchpad[*i].read_by.push(agent_id.to_string());
1205 }
1206
1207 self.updated_at = Utc::now();
1208
1209 Ok(unread_indices
1210 .into_iter()
1211 .map(|index| &self.scratchpad[index])
1212 .collect())
1213 }
1214
1215 pub fn cleanup_stale(&mut self, max_age_hours: u64) {
1216 let now = Utc::now();
1217 let cutoff = now - chrono::Duration::hours(max_age_hours as i64);
1218 let worker_cutoff = now - chrono::Duration::seconds(active_worker_lease_seconds());
1219 let mut compatibility_identities = ProcessIdentityIndex::load();
1220 let mut identities_changed = false;
1221 let mut newly_certified = Vec::new();
1222
1223 for agent in &mut self.agents {
1224 if agent.status == AgentStatus::Finished {
1225 if is_recoverable_legacy_finished(agent) {
1226 if let Some(identity) = safe_legacy_identity(agent) {
1227 identities_changed |=
1228 compatibility_identities.insert(&agent.agent_id, &identity);
1229 newly_certified.push(agent.agent_id.clone());
1230 agent.status = AgentStatus::Active;
1231 agent.status_message =
1232 Some("legacy process identity recovered".to_string());
1233 }
1234 }
1235 continue;
1236 }
1237 if let Some(identity) = &agent.process_identity {
1238 identities_changed |= compatibility_identities.insert(&agent.agent_id, identity);
1239 }
1240 if process_identity_matches(agent, &compatibility_identities) {
1241 if agent.status == AgentStatus::Active && agent.last_active < worker_cutoff {
1250 if agent.transport.is_some() {
1251 agent.status = AgentStatus::Idle;
1252 agent.status_message = Some(TRANSPORT_IDLE.to_string());
1253 } else {
1254 agent.status = AgentStatus::Finished;
1255 agent.status_message = Some(WORKER_LEASE_EXPIRED.to_string());
1256 }
1257 }
1258 continue;
1259 }
1260 if let Some(identity) = safe_legacy_identity(agent) {
1261 identities_changed |= compatibility_identities.insert(&agent.agent_id, &identity);
1262 newly_certified.push(agent.agent_id.clone());
1263 continue;
1264 }
1265 {
1266 agent.status = AgentStatus::Finished;
1267 agent.status_message = Some("process identity no longer matches".to_string());
1268 }
1269 }
1270
1271 let mut finished_by_recency: Vec<_> = self
1275 .agents
1276 .iter()
1277 .filter(|agent| agent.status == AgentStatus::Finished && agent.last_active >= cutoff)
1278 .collect();
1279 finished_by_recency.sort_unstable_by(|left, right| {
1280 right
1281 .last_active
1282 .cmp(&left.last_active)
1283 .then_with(|| left.agent_id.cmp(&right.agent_id))
1284 });
1285 let retained_finished_ids: HashSet<String> = finished_by_recency
1286 .into_iter()
1287 .take(MAX_RETAINED_FINISHED_AGENTS)
1288 .map(|agent| agent.agent_id.clone())
1289 .collect();
1290
1291 self.agents.retain(|a| {
1294 let retire = a.status == AgentStatus::Finished
1295 && (a.last_active < cutoff || !retained_finished_ids.contains(&a.agent_id));
1296 if retire {
1297 crate::core::agent_budget::remove(&a.agent_id);
1298 }
1299 !retire
1300 });
1301 identities_changed |= compatibility_identities.retain_agents(self);
1302 if identities_changed && compatibility_identities.save().is_err() {
1303 for agent in &mut self.agents {
1304 if newly_certified.contains(&agent.agent_id) {
1305 agent.status = AgentStatus::Finished;
1306 agent.status_message =
1307 Some("legacy process identity could not be persisted".to_string());
1308 }
1309 }
1310 }
1311
1312 let now = Utc::now();
1314 self.scratchpad
1315 .retain(|entry| entry.expires_at.is_none_or(|exp| exp > now));
1316
1317 self.updated_at = Utc::now();
1318 }
1319
1320 #[cfg(any(test, feature = "test-support"))]
1321 pub fn save(&self) -> Result<(), String> {
1322 let dir = agents_dir()?;
1323 create_agents_dir(&dir)?;
1324
1325 let lock_path = dir.join("registry.lock");
1326 let _lock = FileLock::acquire(&lock_path)?;
1327
1328 self.save_locked(&dir)
1329 }
1330
1331 fn save_locked(&self, dir: &std::path::Path) -> Result<(), String> {
1332 let path = dir.join("registry.json");
1333 save_registry_file(&path, self)
1334 }
1335
1336 pub fn load() -> Option<Self> {
1337 let dir = agents_dir().ok()?;
1338 let path = dir.join("registry.json");
1339 load_registry_file(&path).ok().flatten()
1340 }
1341
1342 pub fn load_or_create() -> Self {
1343 Self::load().unwrap_or_default()
1344 }
1345
1346 pub fn load_checked() -> Result<Self, String> {
1351 let dir = agents_dir()?;
1352 Ok(load_registry_file(&dir.join("registry.json"))?.unwrap_or_default())
1353 }
1354
1355 pub fn mutate_locked<T>(f: impl FnOnce(&mut Self) -> T) -> Result<(Self, T), String> {
1365 let dir = agents_dir()?;
1366 create_agents_dir(&dir)?;
1367
1368 let lock_path = dir.join("registry.lock");
1369 let _lock = FileLock::acquire(&lock_path)?;
1370
1371 let path = dir.join("registry.json");
1372 let mut registry = load_registry_file(&path)?.unwrap_or_default();
1373 let out = f(&mut registry);
1374 registry.save_locked(&dir)?;
1375 Ok((registry, out))
1376 }
1377}
1378
1379impl Default for AgentRegistry {
1380 fn default() -> Self {
1381 Self::new()
1382 }
1383}
1384
1385static ANNOUNCE: std::sync::OnceLock<fn()> = std::sync::OnceLock::new();
1389
1390pub fn register_announce(request: fn()) {
1392 let _ = ANNOUNCE.set(request);
1393}
1394
1395fn request_announce() {
1396 if let Some(request) = ANNOUNCE.get() {
1397 request();
1398 }
1399}
1400
1401#[cfg(test)]
1402#[path = "registry_tests.rs"]
1403mod tests;
1404
1405#[cfg(test)]
1406#[path = "registry_logical_session_tests.rs"]
1407mod logical_session_tests;
1408#[cfg(test)]
1410#[path = "registry_presence_tests.rs"]
1411mod presence_tests;