1use std::{
5 collections::{HashMap, HashSet},
6 path::{Path, PathBuf},
7};
8
9use chrono::{DateTime, Utc};
10use serde::{Deserialize, Serialize};
11
12use crate::{
13 fs_atomic::write_file_atomic,
14 lock::RepoLock,
15 store::{HeddleError, Result},
16};
17
18const STALE_AGENT_TTL_DAYS: i64 = 7;
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct ContextQueryEntry {
23 pub path: String,
25 pub scope: Option<String>,
27 pub queried_at: DateTime<Utc>,
29}
30
31#[derive(Debug, Clone, Default, Serialize, Deserialize)]
32pub struct AgentUsageSummary {
33 #[serde(default)]
34 pub input_tokens: Option<u64>,
35 #[serde(default)]
36 pub output_tokens: Option<u64>,
37 #[serde(default)]
38 pub reasoning_tokens: Option<u64>,
39 #[serde(default)]
40 pub tool_calls: Option<u32>,
41 #[serde(default)]
42 pub cost_micros_usd: Option<u64>,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
47#[serde(rename_all = "snake_case")]
48pub enum ActorPresenceStatus {
49 Active,
51 Abandoned,
53 Complete,
55 Merged,
57}
58
59impl std::fmt::Display for ActorPresenceStatus {
60 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61 match self {
62 ActorPresenceStatus::Active => write!(f, "active"),
63 ActorPresenceStatus::Abandoned => write!(f, "abandoned"),
64 ActorPresenceStatus::Complete => write!(f, "complete"),
65 ActorPresenceStatus::Merged => write!(f, "merged"),
66 }
67 }
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct ActorPresence {
73 pub session_id: String,
75 #[serde(default)]
78 pub client_instance_id: Option<String>,
79 #[serde(default)]
81 pub native_actor_key: Option<String>,
82 #[serde(default)]
84 pub native_parent_actor_key: Option<String>,
85 #[serde(default)]
87 pub native_instance_key: Option<String>,
88 #[serde(default)]
91 pub heddle_session_id: Option<String>,
92 #[serde(default)]
94 pub thread_id: Option<String>,
95 pub thread: String,
97 #[serde(default)]
99 pub anchor_state: Option<String>,
100 #[serde(default)]
102 pub anchor_root: Option<String>,
103 #[serde(default)]
105 pub path: Option<PathBuf>,
106 pub base_state: String,
108 pub started_at: DateTime<Utc>,
110 pub provider: Option<String>,
112 pub model: Option<String>,
114 #[serde(default)]
116 pub harness: Option<String>,
117 #[serde(default)]
119 pub thinking_level: Option<String>,
120 #[serde(default)]
122 pub usage_summary: AgentUsageSummary,
123 #[serde(default)]
125 pub last_progress_at: Option<DateTime<Utc>>,
126 #[serde(default)]
128 pub report_flush_state: Option<String>,
129 #[serde(default)]
132 pub attach_reason: Option<String>,
133 #[serde(default)]
135 pub task_assignment_id: Option<String>,
136 #[serde(default)]
138 pub attach_precedence: Vec<String>,
139 #[serde(default)]
141 pub winning_attach_rule: Option<String>,
142 #[serde(default)]
144 pub probe_source: Option<String>,
145 #[serde(default)]
147 pub probe_confidence: Option<f32>,
148 pub status: ActorPresenceStatus,
150 #[serde(default)]
152 pub completed_at: Option<DateTime<Utc>>,
153 #[serde(default)]
156 pub context_queries: Vec<ContextQueryEntry>,
157}
158
159#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
161pub struct ActorChainNode {
162 pub session_id: String,
163 #[serde(default)]
164 pub native_actor_key: Option<String>,
165 #[serde(default)]
166 pub native_parent_actor_key: Option<String>,
167 pub thread: String,
168 pub status: ActorPresenceStatus,
169 #[serde(default)]
170 pub provider: Option<String>,
171 #[serde(default)]
172 pub model: Option<String>,
173 #[serde(default)]
174 pub harness: Option<String>,
175}
176
177impl From<&ActorPresence> for ActorChainNode {
178 fn from(entry: &ActorPresence) -> Self {
179 Self {
180 session_id: entry.session_id.clone(),
181 native_actor_key: entry.native_actor_key.clone(),
182 native_parent_actor_key: entry.native_parent_actor_key.clone(),
183 thread: entry.thread.clone(),
184 status: entry.status.clone(),
185 provider: entry.provider.clone(),
186 model: entry.model.clone(),
187 harness: entry.harness.clone(),
188 }
189 }
190}
191
192pub struct ActorPresenceStore {
194 presence_dir: PathBuf,
195}
196
197impl ActorPresenceStore {
198 pub fn new(heddle_dir: &Path) -> Self {
200 Self {
201 presence_dir: heddle_dir.join("actor-presence"),
202 }
203 }
204
205 fn entry_path(&self, session_id: &str) -> Result<PathBuf> {
206 if session_id.is_empty()
210 || !session_id
211 .bytes()
212 .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-')
213 {
214 return Err(HeddleError::Config(format!(
215 "invalid session ID '{}': only lowercase alphanumeric and hyphens allowed",
216 session_id
217 )));
218 }
219 Ok(self.presence_dir.join(format!("{}.toml", session_id)))
220 }
221
222 fn lock_path(&self) -> PathBuf {
223 self.presence_dir.join(".lock")
224 }
225
226 fn write_lock(&self) -> Result<crate::lock::WriteLockGuard> {
227 RepoLock::at(self.lock_path()).write().map_err(|err| {
228 HeddleError::Config(format!("failed to acquire agent registry lock: {err}"))
229 })
230 }
231
232 fn write_entry_file(&self, entry: &ActorPresence) -> Result<()> {
233 crate::fs_atomic::create_dir_all_durable(&self.presence_dir)?;
234 let path = self.entry_path(&entry.session_id)?;
235 let content =
236 toml::to_string_pretty(entry).map_err(|e| HeddleError::Config(e.to_string()))?;
237 Ok(write_file_atomic(&path, content.as_bytes())?)
238 }
239
240 fn load_entry_from_path(&self, path: &Path) -> Result<Option<ActorPresence>> {
241 if !path.exists() {
242 return Ok(None);
243 }
244
245 let content = std::fs::read_to_string(path)?;
246 let entry = toml::from_str(&content).map_err(|e| HeddleError::Config(e.to_string()))?;
247 Ok(Some(entry))
248 }
249
250 fn is_stale_terminal_entry(&self, entry: &ActorPresence) -> bool {
251 if matches!(entry.status, ActorPresenceStatus::Active) {
252 return false;
253 }
254
255 let terminal_at = entry.completed_at.unwrap_or(entry.started_at);
256 terminal_at <= Utc::now() - chrono::Duration::days(STALE_AGENT_TTL_DAYS)
257 }
258
259 fn prune_stale_entry_path(&self, path: &Path) -> Result<()> {
260 if path.exists() {
261 std::fs::remove_file(path)?;
262 }
263 Ok(())
264 }
265
266 pub fn current_entries(&self) -> Result<Vec<ActorPresence>> {
267 self.list()
268 }
269
270 pub fn active_entries(&self) -> Result<Vec<ActorPresence>> {
271 Ok(self
272 .current_entries()?
273 .into_iter()
274 .filter(|entry| entry.status == ActorPresenceStatus::Active)
275 .collect())
276 }
277
278 fn create_generated_entry_with<F, G>(
279 &self,
280 mut generate_id: G,
281 mut build_entry: F,
282 ) -> Result<ActorPresence>
283 where
284 F: FnMut(&str) -> Result<ActorPresence>,
285 G: FnMut() -> String,
286 {
287 let _lock = self.write_lock()?;
288
289 loop {
290 let session_id = generate_id();
291 let path = self.entry_path(&session_id)?;
292 if path.exists() {
293 continue;
294 }
295
296 let entry = build_entry(&session_id)?;
297 self.write_entry_file(&entry)?;
298 return Ok(entry);
299 }
300 }
301
302 pub fn create_generated_entry<F>(&self, build_entry: F) -> Result<ActorPresence>
304 where
305 F: FnMut(&str) -> Result<ActorPresence>,
306 {
307 self.create_generated_entry_with(generate_actor_session_id, build_entry)
308 }
309
310 pub fn save(&self, entry: &ActorPresence) -> Result<()> {
315 let _lock = self.write_lock()?;
316 self.write_entry_file(entry)
317 }
318
319 pub fn load(&self, session_id: &str) -> Result<Option<ActorPresence>> {
321 let path = self.entry_path(session_id)?;
322 let Some(entry) = self.load_entry_from_path(&path)? else {
323 return Ok(None);
324 };
325
326 if self.is_stale_terminal_entry(&entry) {
327 let _lock = self.write_lock()?;
328 if let Some(latest) = self.load_entry_from_path(&path)?
329 && self.is_stale_terminal_entry(&latest)
330 {
331 self.prune_stale_entry_path(&path)?;
332 return Ok(None);
333 }
334 }
335
336 Ok(Some(entry))
337 }
338
339 pub fn list(&self) -> Result<Vec<ActorPresence>> {
341 if !self.presence_dir.exists() {
342 return Ok(Vec::new());
343 }
344
345 let mut stale_paths = Vec::new();
346 let mut entries = Vec::new();
347 for dir_entry in std::fs::read_dir(&self.presence_dir)? {
348 let dir_entry = dir_entry?;
349 let path = dir_entry.path();
350 if path.extension().map(|e| e == "toml").unwrap_or(false) {
351 let content = std::fs::read_to_string(&path)?;
352 let entry = toml::from_str::<ActorPresence>(&content).map_err(|err| {
353 HeddleError::Config(format!(
354 "failed to parse agent registry entry '{}': {err}",
355 path.display()
356 ))
357 })?;
358 if self.is_stale_terminal_entry(&entry) {
359 stale_paths.push(path);
360 } else {
361 entries.push(entry);
362 }
363 }
364 }
365
366 if !stale_paths.is_empty() {
367 let _lock = self.write_lock()?;
368 for path in stale_paths {
369 if let Some(entry) = self.load_entry_from_path(&path)?
370 && self.is_stale_terminal_entry(&entry)
371 {
372 self.prune_stale_entry_path(&path)?;
373 }
374 }
375 }
376
377 entries.sort_by_key(|a| std::cmp::Reverse(a.started_at));
378 Ok(entries)
379 }
380
381 pub fn update_status(&self, session_id: &str, status: ActorPresenceStatus) -> Result<()> {
383 let _lock = self.write_lock()?;
384 let path = self.entry_path(session_id)?;
385 if let Some(mut entry) = self.load_entry_from_path(&path)? {
386 entry.status = status;
387 entry.completed_at = match entry.status {
388 ActorPresenceStatus::Active => None,
389 ActorPresenceStatus::Abandoned
390 | ActorPresenceStatus::Complete
391 | ActorPresenceStatus::Merged => Some(Utc::now()),
392 };
393 self.write_entry_file(&entry)?;
394 }
395 Ok(())
396 }
397
398 pub fn update_entry<F>(&self, session_id: &str, mut update: F) -> Result<Option<ActorPresence>>
400 where
401 F: FnMut(&mut ActorPresence),
402 {
403 let _lock = self.write_lock()?;
404 let path = self.entry_path(session_id)?;
405 let Some(mut entry) = self.load_entry_from_path(&path)? else {
406 return Ok(None);
407 };
408 update(&mut entry);
409 self.write_entry_file(&entry)?;
410 Ok(Some(entry))
411 }
412
413 pub fn find_or_create_active_entry<FMatch, FUpdate, FBuild>(
416 &self,
417 mut matches: FMatch,
418 mut update_existing: FUpdate,
419 mut build_entry: FBuild,
420 ) -> Result<(ActorPresence, bool)>
421 where
422 FMatch: FnMut(&ActorPresence) -> bool,
423 FUpdate: FnMut(&mut ActorPresence),
424 FBuild: FnMut(&str) -> Result<ActorPresence>,
425 {
426 let _lock = self.write_lock()?;
427 crate::fs_atomic::create_dir_all_durable(&self.presence_dir)?;
428
429 for dir_entry in std::fs::read_dir(&self.presence_dir)? {
430 let dir_entry = dir_entry?;
431 let path = dir_entry.path();
432 if !path.extension().map(|e| e == "toml").unwrap_or(false) {
433 continue;
434 }
435 let Some(mut entry) = self.load_entry_from_path(&path)? else {
436 continue;
437 };
438 if self.is_stale_terminal_entry(&entry) {
439 self.prune_stale_entry_path(&path)?;
440 continue;
441 }
442 if entry.status == ActorPresenceStatus::Active && matches(&entry) {
443 update_existing(&mut entry);
444 self.write_entry_file(&entry)?;
445 return Ok((entry, false));
446 }
447 }
448
449 loop {
450 let session_id = generate_actor_session_id();
451 let path = self.entry_path(&session_id)?;
452 if path.exists() {
453 continue;
454 }
455
456 let entry = build_entry(&session_id)?;
457 self.write_entry_file(&entry)?;
458 return Ok((entry, true));
459 }
460 }
461
462 pub fn find_active_by_path(&self, worktree_root: &Path) -> Result<Option<ActorPresence>> {
465 let canonical = worktree_root
466 .canonicalize()
467 .unwrap_or_else(|_| worktree_root.to_path_buf());
468 let entries = self.active_entries()?;
469 Ok(entries
470 .into_iter()
471 .find(|entry| entry_matches_root(entry, &canonical)))
472 }
473
474 pub fn find_active_by_heddle_session_id(
476 &self,
477 heddle_session_id: &str,
478 ) -> Result<Option<ActorPresence>> {
479 let entries = self.active_entries()?;
480 Ok(entries
481 .into_iter()
482 .find(|entry| entry.heddle_session_id.as_deref() == Some(heddle_session_id)))
483 }
484
485 pub fn find_active_by_client_instance_id(
488 &self,
489 client_instance_id: &str,
490 ) -> Result<Option<ActorPresence>> {
491 let entries = self.active_entries()?;
492 Ok(entries
493 .into_iter()
494 .find(|entry| entry.client_instance_id.as_deref() == Some(client_instance_id)))
495 }
496
497 pub fn find_active_by_native_actor_key(
499 &self,
500 native_actor_key: &str,
501 ) -> Result<Option<ActorPresence>> {
502 let entries = self.active_entries()?;
503 Ok(entries
504 .into_iter()
505 .find(|entry| entry.native_actor_key.as_deref() == Some(native_actor_key)))
506 }
507
508 pub fn actor_chain_for_session(&self, session_id: &str) -> Result<Vec<ActorChainNode>> {
515 let entries = self.current_entries()?;
516 let by_session: HashMap<&str, &ActorPresence> = entries
517 .iter()
518 .map(|entry| (entry.session_id.as_str(), entry))
519 .collect();
520 let by_native_key: HashMap<&str, &ActorPresence> = entries
521 .iter()
522 .filter_map(|entry| entry.native_actor_key.as_deref().map(|key| (key, entry)))
523 .collect();
524
525 let Some(mut current) = by_session.get(session_id).copied() else {
526 return Ok(Vec::new());
527 };
528 let mut leaf_to_root = vec![ActorChainNode::from(current)];
529 let mut seen = HashSet::from([current.session_id.as_str()]);
530
531 while let Some(parent_key) = current.native_parent_actor_key.as_deref() {
532 let Some(parent) = by_native_key.get(parent_key).copied() else {
533 break;
534 };
535 if !seen.insert(parent.session_id.as_str()) {
536 break;
537 }
538 leaf_to_root.push(ActorChainNode::from(parent));
539 current = parent;
540 }
541
542 leaf_to_root.reverse();
543 Ok(leaf_to_root)
544 }
545
546 pub fn find_active_by_native_instance_key_at_path(
549 &self,
550 native_instance_key: &str,
551 worktree_root: &Path,
552 ) -> Result<Option<ActorPresence>> {
553 let canonical = worktree_root
554 .canonicalize()
555 .unwrap_or_else(|_| worktree_root.to_path_buf());
556 let entries = self.active_entries()?;
557 Ok(entries.into_iter().find(|entry| {
558 entry.native_instance_key.as_deref() == Some(native_instance_key)
559 && entry_matches_root(entry, &canonical)
560 }))
561 }
562
563 pub fn log_context_query(&self, session_id: &str, query: ContextQueryEntry) -> Result<()> {
567 let _lock = self.write_lock()?;
568 let path = self.entry_path(session_id)?;
569 if let Some(mut entry) = self.load_entry_from_path(&path)?
570 && entry.status == ActorPresenceStatus::Active
571 {
572 entry.context_queries.push(query);
573 self.write_entry_file(&entry)?;
574 }
575 Ok(())
576 }
577
578 pub fn delete(&self, session_id: &str) -> Result<()> {
580 let path = self.entry_path(session_id)?;
581 if path.exists() {
582 std::fs::remove_file(path)?;
583 }
584 Ok(())
585 }
586}
587
588pub fn generate_actor_session_id() -> String {
593 let random_bytes: [u8; 12] = rand::random();
594 format!(
595 "agent-{}",
596 base32::encode(base32::Alphabet::Rfc4648 { padding: false }, &random_bytes).to_lowercase()
597 )
598}
599
600fn entry_matches_root(entry: &ActorPresence, canonical: &Path) -> bool {
601 entry
602 .path
603 .as_ref()
604 .map(|p| p.canonicalize().unwrap_or_else(|_| p.clone()) == canonical)
605 .unwrap_or(false)
606}
607
608#[cfg(test)]
609mod tests {
610 use tempfile::TempDir;
611
612 use super::*;
613
614 fn presence(session_id: &str, status: ActorPresenceStatus) -> ActorPresence {
615 ActorPresence {
616 session_id: session_id.to_string(),
617 client_instance_id: None,
618 native_actor_key: None,
619 native_parent_actor_key: None,
620 native_instance_key: None,
621 heddle_session_id: None,
622 thread_id: None,
623 thread: "feature/test".to_string(),
624 anchor_state: Some("hd-state".to_string()),
625 anchor_root: Some("root".to_string()),
626 path: None,
627 base_state: "hd-state".to_string(),
628 started_at: Utc::now(),
629 provider: Some("openai".to_string()),
630 model: Some("gpt-5".to_string()),
631 harness: Some("codex".to_string()),
632 thinking_level: None,
633 usage_summary: AgentUsageSummary::default(),
634 last_progress_at: None,
635 report_flush_state: None,
636 attach_reason: Some("test".to_string()),
637 task_assignment_id: None,
638 attach_precedence: vec!["explicit".to_string()],
639 winning_attach_rule: Some("explicit".to_string()),
640 probe_source: None,
641 probe_confidence: None,
642 status,
643 completed_at: None,
644 context_queries: vec![],
645 }
646 }
647
648 #[test]
649 fn active_presence_is_independent_of_writer_liveness() {
650 let temp = TempDir::new().unwrap();
651 let store = ActorPresenceStore::new(temp.path());
652 store
653 .save(&presence("agent-one", ActorPresenceStatus::Active))
654 .unwrap();
655
656 let active = store.active_entries().unwrap();
657 assert_eq!(active.len(), 1);
658 assert_eq!(active[0].session_id, "agent-one");
659 }
660
661 #[test]
662 fn active_presence_can_be_reused_by_native_identity() {
663 let temp = TempDir::new().unwrap();
664 let store = ActorPresenceStore::new(temp.path());
665 let (first, created) = store
666 .find_or_create_active_entry(
667 |_| false,
668 |_| {},
669 |session_id| {
670 let mut entry = presence(session_id, ActorPresenceStatus::Active);
671 entry.native_actor_key = Some("codex:thread:one".to_string());
672 Ok(entry)
673 },
674 )
675 .unwrap();
676 assert!(created);
677
678 let (second, created) = store
679 .find_or_create_active_entry(
680 |entry| entry.native_actor_key.as_deref() == Some("codex:thread:one"),
681 |_| {},
682 |_| panic!("matching presence should be reused"),
683 )
684 .unwrap();
685 assert!(!created);
686 assert_eq!(first.session_id, second.session_id);
687 }
688
689 #[test]
690 fn terminal_presence_is_retained_for_recent_provenance() {
691 let temp = TempDir::new().unwrap();
692 let store = ActorPresenceStore::new(temp.path());
693 let mut complete = presence("agent-done", ActorPresenceStatus::Complete);
694 complete.completed_at = Some(Utc::now());
695 store.save(&complete).unwrap();
696
697 let loaded = store.load("agent-done").unwrap().unwrap();
698 assert_eq!(loaded.status, ActorPresenceStatus::Complete);
699 }
700}