1use std::{
2 collections::{BTreeMap, BTreeSet},
3 env,
4 fs::{self, File},
5 io::Write,
6 path::{Path, PathBuf},
7 sync::atomic::{AtomicU64, Ordering},
8 time::{SystemTime, UNIX_EPOCH},
9};
10
11use anyhow::{anyhow, bail, Context, Result};
12use chrono::{DateTime, Duration, Utc};
13use serde::{Deserialize, Serialize};
14use sha2::{Digest, Sha256};
15
16use crate::{
17 activity::{ActivityState, ActivityStateValue},
18 cli::{AnnotateArgs, CurrentArgs, DiscoverArgs, LookupArgs, PruneArgs, RegisterArgs},
19 names,
20};
21
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
23pub struct ActivitySummary {
24 pub text: String,
25 pub updated_at: DateTime<Utc>,
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
29pub struct ExtensionMetadata {
30 pub data: serde_json::Value,
31 pub updated_at: DateTime<Utc>,
32}
33
34#[derive(Debug, Clone)]
35pub struct ActivityUpdate {
36 pub summary: Option<String>,
37 pub clear_summary: bool,
38 pub state: Option<ActivityStateValue>,
39 pub clear_state: bool,
40 pub cwd: Option<String>,
41 pub clear_cwd: bool,
42 pub extensions: BTreeMap<String, serde_json::Value>,
43 pub clear_extensions: BTreeSet<String>,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
47pub struct Assignment {
48 pub version: u8,
49 pub session_id: String,
50 pub name: String,
51 pub slug: String,
52 pub first_name: String,
53 pub family_name: String,
54 pub realm: String,
55 #[serde(default)]
56 pub summary: Option<ActivitySummary>,
57 #[serde(default)]
58 pub state: Option<ActivityState>,
59 #[serde(default)]
60 pub cwd: Option<String>,
61 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
62 pub extensions: BTreeMap<String, ExtensionMetadata>,
63 pub created_at: DateTime<Utc>,
64 pub updated_at: DateTime<Utc>,
65}
66
67#[derive(Debug, Clone, Serialize)]
68pub struct PrunedIdentity {
69 pub session_id: String,
70 pub name: String,
71 pub slug: String,
72 pub updated_at: DateTime<Utc>,
73 pub claim_removed: bool,
74}
75
76#[derive(Debug, Clone, Serialize)]
77pub struct PruneReport {
78 pub cutoff: DateTime<Utc>,
79 pub dry_run: bool,
80 pub candidates: Vec<Assignment>,
81 pub removed: Vec<PrunedIdentity>,
82 pub errors: Vec<String>,
83}
84
85#[derive(Debug, Clone)]
86pub struct Registry {
87 root: PathBuf,
88}
89
90impl Registry {
91 pub fn from_env() -> Result<Self> {
92 let root = if let Some(path) = env::var_os("AGENT_ID_HOME") {
93 PathBuf::from(path)
94 } else if let Some(path) = env::var_os("XDG_DATA_HOME") {
95 PathBuf::from(path).join("agent-id")
96 } else {
97 home_dir()?.join(".local/share/agent-id")
98 };
99
100 Ok(Self::new(root))
101 }
102
103 pub fn new(root: PathBuf) -> Self {
104 Self { root }
105 }
106
107 pub fn register(
108 &self,
109 session_id: &str,
110 family_name: Option<&str>,
111 realm: &str,
112 ) -> Result<Assignment> {
113 let session_id = validate_session_id(session_id)?;
114 let session_path = self.session_path(&session_id);
115 if session_path.exists() {
116 let mut existing = read_assignment(&session_path)
117 .with_context(|| format!("read existing assignment {}", session_path.display()))?;
118 existing.updated_at = Utc::now();
119 replace_assignment(&session_path, &existing)?;
120 return Ok(existing);
121 }
122
123 let realm = normalize_component(realm, "realm")?;
124 let requested_family = family_name
125 .map(|value| normalize_component(value, "family name"))
126 .transpose()?;
127 let first_names = names::first_names();
128 let family_names = names::family_names();
129 if first_names.is_empty() || family_names.is_empty() {
130 bail!("name lists are empty");
131 }
132
133 fs::create_dir_all(self.root.join("by-session"))
134 .with_context(|| format!("create {}", self.root.display()))?;
135 fs::create_dir_all(self.root.join("by-name"))
136 .with_context(|| format!("create {}", self.root.display()))?;
137
138 for attempt in 0..100_000_u64 {
139 let (first, family) = candidate(
140 &session_id,
141 attempt,
142 &first_names,
143 &family_names,
144 requested_family.as_deref(),
145 );
146 if first == family {
147 continue;
148 }
149
150 let first_name = title_word(first);
151 let family_name = title_word(&family);
152
153 let name = format!("{first_name} {family_name} of {realm}");
154 let slug = slug(&first_name, &family_name, &realm);
155 let claim_path = self.root.join("by-name").join(&slug);
156
157 match claim_name(&claim_path, &session_id)? {
158 Claim::Owned => {}
159 Claim::Claimed => {}
160 Claim::Other => continue,
161 }
162
163 let now = Utc::now();
164 let assignment = Assignment {
165 version: 1,
166 session_id: session_id.clone(),
167 name,
168 slug,
169 first_name,
170 family_name,
171 realm: realm.clone(),
172 state: None,
173 cwd: None,
174 summary: None,
175 extensions: BTreeMap::new(),
176 created_at: now,
177 updated_at: now,
178 };
179 write_assignment(&session_path, &assignment)?;
180 return Ok(assignment);
181 }
182
183 bail!("exhausted available names for realm {realm}")
184 }
185
186 pub fn lookup(&self, input: &str) -> Result<Assignment> {
187 let input = require_nonempty(input, "lookup identifier")?;
188 if let Ok(session_id) = validate_session_id(&input) {
189 let session_path = self.session_path(&session_id);
190 if session_path.exists() {
191 return self.lookup_session(&session_id);
192 }
193 }
194
195 for slug in lookup_slugs(&input) {
196 let claim_path = self.root.join("by-name").join(&slug);
197 if !claim_path.is_file() {
198 continue;
199 }
200
201 let session_id = require_nonempty(
202 &fs::read_to_string(&claim_path)
203 .with_context(|| format!("read name claim {}", claim_path.display()))?,
204 "claimed session ID",
205 )?;
206 return self
207 .lookup_session(&session_id)
208 .with_context(|| format!("resolve name claim {slug}"));
209 }
210
211 bail!(
212 "no identity found for '{input}'; lookup accepts a session ID, canonical name, or slug"
213 )
214 }
215
216 fn lookup_session(&self, session_id: &str) -> Result<Assignment> {
217 let session_id = validate_session_id(session_id)?;
218 let path = self.session_path(&session_id);
219 if !path.exists() {
220 bail!(
221 "no identity registered for session {session_id}; run `agent-id register {session_id}`"
222 );
223 }
224
225 let assignment = read_assignment(&path)
226 .with_context(|| format!("read assignment for session {session_id}"))?;
227 if assignment.session_id != session_id {
228 bail!("identity registry entry does not belong to session {session_id}");
229 }
230 Ok(assignment)
231 }
232
233 pub fn annotate(&self, session_id: &str, update: ActivityUpdate) -> Result<Assignment> {
234 let session_id = validate_session_id(session_id)?;
235 if update.summary.is_none()
236 && !update.clear_summary
237 && update.state.is_none()
238 && !update.clear_state
239 && update.cwd.is_none()
240 && !update.clear_cwd
241 && update.extensions.is_empty()
242 && update.clear_extensions.is_empty()
243 {
244 bail!("pass at least one activity update");
245 }
246 if update.summary.is_some() && update.clear_summary {
247 bail!("summary and clear_summary are mutually exclusive");
248 }
249 if update.state.is_some() && update.clear_state {
250 bail!("state and clear_state are mutually exclusive");
251 }
252 if update.cwd.is_some() && update.clear_cwd {
253 bail!("cwd and clear_cwd are mutually exclusive");
254 }
255 if let Some(owner) = update
256 .extensions
257 .keys()
258 .find(|owner| update.clear_extensions.contains(*owner))
259 {
260 bail!("extension {owner} cannot be set and cleared together");
261 }
262
263 let mut assignment = self.lookup_session(&session_id)?;
264 let summary = update
265 .summary
266 .as_deref()
267 .map(normalize_summary)
268 .transpose()?;
269 let cwd = update.cwd.as_deref().map(normalize_cwd).transpose()?;
270 let now = Utc::now();
271 if update.summary.is_some() {
272 assignment.summary = summary.map(|text| ActivitySummary {
273 text,
274 updated_at: now,
275 });
276 } else if update.clear_summary {
277 assignment.summary = None;
278 }
279 if let Some(value) = update.state {
280 assignment.state = Some(ActivityState {
281 value,
282 updated_at: now,
283 });
284 } else if update.clear_state {
285 assignment.state = None;
286 }
287 if update.cwd.is_some() {
288 assignment.cwd = cwd;
289 } else if update.clear_cwd {
290 assignment.cwd = None;
291 }
292 for (owner, data) in update.extensions {
293 assignment.extensions.insert(
294 owner,
295 ExtensionMetadata {
296 data,
297 updated_at: now,
298 },
299 );
300 }
301 for owner in update.clear_extensions {
302 assignment.extensions.remove(&owner);
303 }
304 assignment.updated_at = now;
305 replace_assignment(&self.session_path(&session_id), &assignment)?;
306 Ok(assignment)
307 }
308
309 pub fn discover(
310 &self,
311 limit: usize,
312 recent_hours: Option<i64>,
313 realm: Option<&str>,
314 ) -> Result<Vec<Assignment>> {
315 if recent_hours.is_some_and(|hours| hours < 0) {
316 bail!("--recent must be non-negative");
317 }
318 let realm = realm
319 .map(|value| normalize_component(value, "realm"))
320 .transpose()?;
321 let cutoff = recent_hours.map(|hours| Utc::now() - Duration::hours(hours));
322 let path = self.root.join("by-session");
323 let entries = match fs::read_dir(&path) {
324 Ok(entries) => entries,
325 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
326 Err(error) => return Err(error).with_context(|| format!("read {}", path.display())),
327 };
328 let mut assignments = Vec::new();
329 for entry in entries {
330 let path = entry?.path();
331 if path.extension().and_then(|value| value.to_str()) != Some("json") {
332 continue;
333 }
334 let assignment = read_assignment(&path)
335 .with_context(|| format!("read assignment {}", path.display()))?;
336 if realm
337 .as_deref()
338 .is_some_and(|value| value != assignment.realm)
339 {
340 continue;
341 }
342 if cutoff.is_some_and(|value| assignment.updated_at < value) {
343 continue;
344 }
345 assignments.push(assignment);
346 }
347 assignments.sort_by(|left, right| right.updated_at.cmp(&left.updated_at));
348 if limit > 0 {
349 assignments.truncate(limit);
350 }
351 Ok(assignments)
352 }
353
354 pub fn prune(&self, cutoff: DateTime<Utc>, dry_run: bool) -> Result<PruneReport> {
355 let path = self.root.join("by-session");
356 let entries = match fs::read_dir(&path) {
357 Ok(entries) => entries,
358 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
359 return Ok(PruneReport {
360 cutoff,
361 dry_run,
362 candidates: Vec::new(),
363 removed: Vec::new(),
364 errors: Vec::new(),
365 })
366 }
367 Err(error) => return Err(error).with_context(|| format!("read {}", path.display())),
368 };
369 let mut candidates = Vec::new();
370 for entry in entries {
371 let path = entry?.path();
372 if path.extension().and_then(|value| value.to_str()) != Some("json") {
373 continue;
374 }
375 let assignment = read_assignment(&path)
376 .with_context(|| format!("read assignment {}", path.display()))?;
377 if assignment.updated_at < cutoff {
378 candidates.push(assignment);
379 }
380 }
381 candidates.sort_by(|left, right| left.updated_at.cmp(&right.updated_at));
382 let mut report = PruneReport {
383 cutoff,
384 dry_run,
385 candidates,
386 removed: Vec::new(),
387 errors: Vec::new(),
388 };
389 if dry_run {
390 return Ok(report);
391 }
392
393 for assignment in &report.candidates {
394 let session_path = self.session_path(&assignment.session_id);
395 if let Err(error) = fs::remove_file(&session_path) {
396 report
397 .errors
398 .push(format!("remove {}: {error}", session_path.display()));
399 continue;
400 }
401
402 let claim_path = self.root.join("by-name").join(&assignment.slug);
403 let claim_removed = match fs::read_to_string(&claim_path) {
404 Ok(owner) if owner.trim() == assignment.session_id => {
405 if let Err(error) = fs::remove_file(&claim_path) {
406 report.errors.push(format!(
407 "remove name claim {}: {error}",
408 claim_path.display()
409 ));
410 false
411 } else {
412 true
413 }
414 }
415 Ok(owner) => {
416 report.errors.push(format!(
417 "name claim {} belongs to {}, not {}",
418 claim_path.display(),
419 owner.trim(),
420 assignment.session_id
421 ));
422 false
423 }
424 Err(error) if error.kind() == std::io::ErrorKind::NotFound => false,
425 Err(error) => {
426 report
427 .errors
428 .push(format!("read name claim {}: {error}", claim_path.display()));
429 false
430 }
431 };
432 report.removed.push(PrunedIdentity {
433 session_id: assignment.session_id.clone(),
434 name: assignment.name.clone(),
435 slug: assignment.slug.clone(),
436 updated_at: assignment.updated_at,
437 claim_removed,
438 });
439 }
440 Ok(report)
441 }
442
443 fn session_path(&self, session_id: &str) -> PathBuf {
444 self.root
445 .join("by-session")
446 .join(format!("{session_id}.json"))
447 }
448}
449
450pub fn execute_register(args: &RegisterArgs) -> Result<()> {
451 let session_id = resolve_session(args.explicit_session())?;
452 let realm = resolve_realm(args.realm.as_deref())?;
453 let assignment = Registry::from_env()?.register(&session_id, args.family.as_deref(), &realm)?;
454 print_assignment(&assignment, args.json)
455}
456
457pub fn execute_lookup(args: &LookupArgs) -> Result<()> {
458 let input = resolve_session(args.explicit_input())?;
459 let assignment = Registry::from_env()?.lookup(&input)?;
460 print_assignment(&assignment, args.json)
461}
462
463pub fn execute_current(args: &CurrentArgs) -> Result<()> {
464 let session_id = resolve_session(None)?;
465 let assignment = Registry::from_env()?.lookup(&session_id)?;
466 print_assignment(&assignment, args.json)
467}
468
469pub fn execute_annotate(args: &AnnotateArgs) -> Result<()> {
470 let session_id = resolve_session(args.explicit_session())?;
471 let update = ActivityUpdate {
472 summary: args.summary.clone(),
473 clear_summary: args.clear_summary,
474 state: args.state,
475 clear_state: args.clear_state,
476 cwd: args.cwd.clone(),
477 clear_cwd: args.clear_cwd,
478 extensions: parse_extension_updates(&args.extensions)?,
479 clear_extensions: parse_extension_owners(&args.clear_extensions)?,
480 };
481 let assignment = Registry::from_env()?.annotate(&session_id, update)?;
482 print_assignment(&assignment, args.json)
483}
484
485pub fn execute_discover(args: &DiscoverArgs) -> Result<()> {
486 let assignments =
487 Registry::from_env()?.discover(args.limit, args.recent, args.realm.as_deref())?;
488 let records = crate::herdr::augment_discovery(assignments);
489 if args.json {
490 println!("{}", serde_json::to_string_pretty(&records)?);
491 } else if records.is_empty() {
492 println!("(no identities)");
493 } else {
494 for record in records {
495 let assignment = &record.assignment;
496 let mut annotations = Vec::new();
497 if let Some(state) = assignment.state.as_ref() {
498 annotations.push(format!("state:{}", state.value));
499 }
500 if let Some(summary) = assignment.summary.as_ref() {
501 annotations.push(format!("summary:{}", summary.text));
502 }
503 if let Some(cwd) = assignment.cwd.as_ref() {
504 annotations.push(format!("cwd:{cwd}"));
505 }
506 if let Some(runtime) = record.runtime {
507 for location in runtime.locations {
508 let workspace = location
509 .workspace_label
510 .as_deref()
511 .unwrap_or(&location.workspace_id);
512 annotations.push(format!(
513 "herdr:{} pane:{} workspace:{}",
514 location.agent_status, location.pane_id, workspace
515 ));
516 }
517 }
518 if annotations.is_empty() {
519 println!("{}\t{}", assignment.name, assignment.session_id);
520 } else {
521 println!(
522 "{}\t{}\t{}",
523 assignment.name,
524 assignment.session_id,
525 annotations.join("\t")
526 );
527 }
528 }
529 }
530 Ok(())
531}
532
533pub fn execute_prune(args: &PruneArgs) -> Result<()> {
534 let cutoff = DateTime::parse_from_rfc3339(&args.before)
535 .with_context(|| format!("parse --before timestamp {}", args.before))?
536 .with_timezone(&Utc);
537 let report = Registry::from_env()?.prune(cutoff, args.dry_run)?;
538 if args.json {
539 println!("{}", serde_json::to_string_pretty(&report)?);
540 } else {
541 let action = if args.dry_run {
542 "would prune"
543 } else {
544 "pruned"
545 };
546 println!(
547 "{action} {} identities before {}",
548 report.candidates.len(),
549 cutoff
550 );
551 for assignment in &report.candidates {
552 println!(
553 "{}\t{}\tupdated:{}",
554 assignment.name, assignment.session_id, assignment.updated_at
555 );
556 }
557 for error in &report.errors {
558 eprintln!("agent-id: {error}");
559 }
560 }
561 if report.errors.is_empty() {
562 Ok(())
563 } else {
564 bail!("prune completed with {} errors", report.errors.len())
565 }
566}
567
568pub fn resolve_session(explicit: Option<&str>) -> Result<String> {
569 if let Some(session_id) = explicit.filter(|value| !value.trim().is_empty()) {
570 return Ok(session_id.trim().to_string());
571 }
572
573 if let Some(value) = env::var_os("AGENT_ID_SESSION_ID") {
574 let value = value.to_string_lossy();
575 if !value.trim().is_empty() {
576 return Ok(value.trim().to_string());
577 }
578 }
579
580 bail!("no session ID found; pass SESSION_ID or --session-id, or set AGENT_ID_SESSION_ID")
581}
582
583fn resolve_realm(explicit: Option<&str>) -> Result<String> {
584 if let Some(realm) = explicit.filter(|value| !value.trim().is_empty()) {
585 return normalize_component(realm, "realm");
586 }
587 if let Some(realm) = env::var_os("AGENT_REALM") {
588 let realm = realm.to_string_lossy();
589 if !realm.trim().is_empty() {
590 return normalize_component(&realm, "realm");
591 }
592 }
593
594 let home = home_dir()?;
595 let config_home = env::var_os("XDG_CONFIG_HOME")
596 .map(PathBuf::from)
597 .unwrap_or_else(|| home.join(".config"));
598 let realm_path = config_home.join("agent-id/realm");
599
600 if realm_path.is_file() {
601 let value = fs::read_to_string(&realm_path)
602 .with_context(|| format!("read realm from {}", realm_path.display()))?;
603 return normalize_component(&value, "realm");
604 }
605
606 auto_create_realm(&realm_path)
607}
608
609fn auto_create_realm(path: &Path) -> Result<String> {
610 let candidates = names::candidate_realms();
611 if candidates.is_empty() {
612 bail!("bundled realm candidate list is empty");
613 }
614
615 let nanos = SystemTime::now()
616 .duration_since(UNIX_EPOCH)
617 .map(|duration| duration.as_nanos())
618 .unwrap_or_default();
619 let pid = std::process::id();
620 let hostname = env::var("HOSTNAME").unwrap_or_default();
621 let digest = Sha256::digest(format!("{nanos}:{pid}:{hostname}").as_bytes());
622 let index = usize::try_from(u64::from_be_bytes(digest[0..8].try_into().unwrap())).unwrap_or(0)
623 % candidates.len();
624 let realm = normalize_component(candidates[index], "realm")?;
625
626 let parent = path
627 .parent()
628 .ok_or_else(|| anyhow!("realm path has no parent directory"))?;
629 fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
630
631 let temp = temporary_path(path);
632 let mut file = File::create(&temp)
633 .with_context(|| format!("create temporary realm {}", temp.display()))?;
634 file.write_all(format!("{realm}\n").as_bytes())?;
635 file.sync_all()?;
636
637 match fs::hard_link(&temp, path) {
638 Ok(()) => {
639 let _ = fs::remove_file(temp);
640 Ok(realm)
641 }
642 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
643 let _ = fs::remove_file(temp);
644 let value = fs::read_to_string(path)
645 .with_context(|| format!("read existing realm {}", path.display()))?;
646 normalize_component(&value, "realm")
647 }
648 Err(error) => {
649 let _ = fs::remove_file(temp);
650 Err(error).with_context(|| format!("create realm file {}", path.display()))
651 }
652 }
653}
654
655fn print_assignment(assignment: &Assignment, json: bool) -> Result<()> {
656 if json {
657 println!("{}", serde_json::to_string_pretty(assignment)?);
658 } else {
659 println!("{}", assignment.name);
660 }
661 Ok(())
662}
663
664fn lookup_slugs(input: &str) -> Vec<String> {
665 let mut slugs = Vec::new();
666 let input = input.trim();
667 if is_slug(input) {
668 slugs.push(input.to_ascii_lowercase());
669 }
670 if let Some(slug) = canonical_name_slug(input) {
671 if !slugs.contains(&slug) {
672 slugs.push(slug);
673 }
674 }
675 slugs
676}
677
678fn canonical_name_slug(input: &str) -> Option<String> {
679 let mut parts = input.split_whitespace();
680 let first = normalize_component(parts.next()?, "first name").ok()?;
681 let family = normalize_component(parts.next()?, "family name").ok()?;
682 if !parts.next()?.eq_ignore_ascii_case("of") {
683 return None;
684 }
685 let realm = normalize_component(parts.next()?, "realm").ok()?;
686 if parts.next().is_some() {
687 return None;
688 }
689 Some(slug(&first, &family, &realm))
690}
691
692fn is_slug(input: &str) -> bool {
693 !input.is_empty()
694 && !input.starts_with('-')
695 && !input.ends_with('-')
696 && input
697 .chars()
698 .all(|character| character.is_ascii_alphanumeric() || character == '-')
699}
700
701fn candidate<'a>(
702 session_id: &str,
703 attempt: u64,
704 first_names: &'a [&'a str],
705 family_names: &'a [&'a str],
706 requested_family: Option<&str>,
707) -> (&'a str, String) {
708 let digest = Sha256::digest(format!("{session_id}:{attempt}").as_bytes());
709 let first_index = usize::try_from(u64::from_be_bytes(digest[0..8].try_into().unwrap()))
710 .unwrap_or(0)
711 % first_names.len();
712 let family_index = usize::try_from(u64::from_be_bytes(digest[8..16].try_into().unwrap()))
713 .unwrap_or(0)
714 % family_names.len();
715 let family = requested_family
716 .map(ToOwned::to_owned)
717 .unwrap_or_else(|| family_names[family_index].to_string());
718 (first_names[first_index], family)
719}
720
721#[derive(Debug, Clone, Copy, PartialEq, Eq)]
722enum Claim {
723 Claimed,
724 Owned,
725 Other,
726}
727
728fn claim_name(path: &Path, session_id: &str) -> Result<Claim> {
729 let parent = path
730 .parent()
731 .ok_or_else(|| anyhow!("name claim has no parent directory"))?;
732 fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
733
734 let temp = temporary_path(path);
735 let mut file = File::create(&temp)
736 .with_context(|| format!("create temporary claim {}", temp.display()))?;
737 file.write_all(session_id.as_bytes())?;
738 file.write_all(b"\n")?;
739 file.sync_all()?;
740
741 let claim = match fs::hard_link(&temp, path) {
742 Ok(()) => Claim::Claimed,
743 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
744 let owner = fs::read_to_string(path).unwrap_or_default();
745 if owner.trim() == session_id {
746 Claim::Owned
747 } else {
748 Claim::Other
749 }
750 }
751 Err(error) => return Err(error).with_context(|| format!("claim name {}", path.display())),
752 };
753 let _ = fs::remove_file(temp);
754 Ok(claim)
755}
756
757fn write_assignment(path: &Path, assignment: &Assignment) -> Result<()> {
758 let parent = path
759 .parent()
760 .ok_or_else(|| anyhow!("assignment has no parent directory"))?;
761 fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
762 let contents = format!("{}\n", serde_json::to_string_pretty(assignment)?);
763 let temp = temporary_path(path);
764 let mut file = File::create(&temp)
765 .with_context(|| format!("create temporary assignment {}", temp.display()))?;
766 file.write_all(contents.as_bytes())?;
767 file.sync_all()?;
768
769 let result = match fs::hard_link(&temp, path) {
770 Ok(()) => Ok(()),
771 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
772 bail!(
773 "session {} already has a registered identity",
774 assignment.session_id
775 )
776 }
777 Err(error) => Err(error).with_context(|| format!("write assignment {}", path.display())),
778 };
779 let _ = fs::remove_file(temp);
780 result
781}
782
783fn replace_assignment(path: &Path, assignment: &Assignment) -> Result<()> {
784 let parent = path
785 .parent()
786 .ok_or_else(|| anyhow!("assignment has no parent directory"))?;
787 fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
788 let contents = format!("{}\n", serde_json::to_string_pretty(assignment)?);
789 let temp = temporary_path(path);
790 let mut file = File::create(&temp)
791 .with_context(|| format!("create temporary assignment {}", temp.display()))?;
792 file.write_all(contents.as_bytes())?;
793 file.sync_all()?;
794
795 let result =
796 fs::rename(&temp, path).with_context(|| format!("replace assignment {}", path.display()));
797 if result.is_err() {
798 let _ = fs::remove_file(temp);
799 }
800 result
801}
802
803fn read_assignment(path: &Path) -> Result<Assignment> {
804 let contents = fs::read_to_string(path)?;
805 Ok(serde_json::from_str(&contents)?)
806}
807
808static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
809
810fn temporary_path(path: &Path) -> PathBuf {
811 let nanos = SystemTime::now()
812 .duration_since(UNIX_EPOCH)
813 .map(|duration| duration.as_nanos())
814 .unwrap_or_default();
815 let counter = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
816 path.with_extension(format!("tmp-{}-{nanos}-{counter}", std::process::id()))
817}
818
819fn slug(first_name: &str, family_name: &str, realm: &str) -> String {
820 [first_name, family_name, realm]
821 .iter()
822 .map(|part| part.to_ascii_lowercase().replace([' ', '\'', '_'], "-"))
823 .collect::<Vec<_>>()
824 .join("-")
825}
826
827fn normalize_component(value: &str, kind: &str) -> Result<String> {
828 let value = value.trim();
829 let valid = !value.is_empty()
830 && value.chars().all(|character| {
831 character.is_ascii_alphabetic() || character == '-' || character == '\''
832 })
833 && value
834 .chars()
835 .next()
836 .is_some_and(|character| character.is_ascii_alphabetic())
837 && value
838 .chars()
839 .last()
840 .is_some_and(|character| character.is_ascii_alphabetic());
841 if !valid {
842 bail!("{kind} must contain only letters, apostrophes, or hyphens")
843 }
844 Ok(title_word(value))
845}
846
847fn title_word(value: &str) -> String {
848 let mut characters = value.chars();
849 let Some(first) = characters.next() else {
850 return String::new();
851 };
852 first.to_uppercase().collect::<String>() + &characters.as_str().to_ascii_lowercase()
853}
854
855fn validate_session_id(value: &str) -> Result<String> {
856 let value = require_nonempty(value, "session ID")?;
857 let safe = value
858 .chars()
859 .all(|character| character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.'));
860 if !safe || value == "." || value == ".." {
861 bail!("session ID must be a filename-safe value using letters, digits, '.', '_' or '-'");
862 }
863 Ok(value)
864}
865
866const MAX_EXTENSION_OWNER_CHARS: usize = 64;
867const MAX_EXTENSION_JSON_BYTES: usize = 16 * 1024;
868
869fn parse_extension_updates(values: &[String]) -> Result<BTreeMap<String, serde_json::Value>> {
870 let mut extensions = BTreeMap::new();
871 for value in values {
872 let (owner, json) = value
873 .split_once('=')
874 .ok_or_else(|| anyhow!("--extension must use OWNER=JSON"))?;
875 let owner = normalize_extension_owner(owner)?;
876 if json.len() > MAX_EXTENSION_JSON_BYTES {
877 bail!("extension {owner} JSON must be at most {MAX_EXTENSION_JSON_BYTES} bytes");
878 }
879 let data = serde_json::from_str(json)
880 .with_context(|| format!("parse JSON for extension {owner}"))?;
881 if extensions.insert(owner.clone(), data).is_some() {
882 bail!("extension {owner} was provided more than once");
883 }
884 }
885 Ok(extensions)
886}
887
888fn parse_extension_owners(values: &[String]) -> Result<BTreeSet<String>> {
889 values
890 .iter()
891 .map(|owner| normalize_extension_owner(owner))
892 .collect()
893}
894
895fn normalize_extension_owner(value: &str) -> Result<String> {
896 let owner = require_nonempty(value, "extension owner")?;
897 let valid = owner.chars().count() <= MAX_EXTENSION_OWNER_CHARS
898 && owner
899 .chars()
900 .next()
901 .is_some_and(|character| character.is_ascii_lowercase() || character.is_ascii_digit())
902 && owner.chars().all(|character| {
903 character.is_ascii_lowercase()
904 || character.is_ascii_digit()
905 || matches!(character, '.' | '_' | '-')
906 });
907 if !valid {
908 bail!(
909 "extension owner must be at most {MAX_EXTENSION_OWNER_CHARS} characters using lowercase letters, digits, '.', '_' or '-'"
910 );
911 }
912 Ok(owner)
913}
914
915const MAX_SUMMARY_CHARS: usize = 240;
916
917fn normalize_summary(value: &str) -> Result<String> {
918 let mut normalized = String::new();
919 for word in value.split_whitespace() {
920 if !normalized.is_empty() {
921 normalized.push(' ');
922 }
923 normalized.push_str(word);
924 }
925 if normalized.is_empty() {
926 bail!("summary cannot be empty");
927 }
928 if normalized.chars().count() > MAX_SUMMARY_CHARS {
929 bail!("summary must be at most {MAX_SUMMARY_CHARS} characters");
930 }
931 Ok(normalized)
932}
933
934const MAX_CWD_CHARS: usize = 4096;
935
936fn normalize_cwd(value: &str) -> Result<String> {
937 let value = require_nonempty(value, "working directory")?;
938 if value.chars().any(char::is_control) {
939 bail!("working directory cannot contain control characters");
940 }
941 if value.chars().count() > MAX_CWD_CHARS {
942 bail!("working directory must be at most {MAX_CWD_CHARS} characters");
943 }
944 Ok(value)
945}
946
947fn require_nonempty(value: &str, kind: &str) -> Result<String> {
948 let value = value.trim();
949 if value.is_empty() {
950 bail!("{kind} cannot be empty")
951 }
952 Ok(value.to_string())
953}
954
955fn home_dir() -> Result<PathBuf> {
956 env::var_os("HOME")
957 .map(PathBuf::from)
958 .ok_or_else(|| anyhow!("HOME is not set; set AGENT_ID_HOME explicitly"))
959}
960
961#[cfg(test)]
962mod tests {
963 use super::*;
964
965 #[test]
966 fn generated_identity_has_canonical_parts() {
967 let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
968 let assignment = registry
969 .register("session-1", Some("Oak"), "Darkwood")
970 .unwrap();
971
972 assert_eq!(assignment.family_name, "Oak");
973 assert_eq!(assignment.realm, "Darkwood");
974 assert_eq!(
975 assignment.slug,
976 format!(
977 "{}-oak-darkwood",
978 assignment.first_name.to_ascii_lowercase()
979 )
980 );
981 assert_eq!(
982 assignment.name,
983 format!("{} Oak of Darkwood", assignment.first_name)
984 );
985 }
986
987 #[test]
988 fn session_id_is_the_registry_filename() {
989 let root = tempfile::tempdir().unwrap();
990 let registry = Registry::new(root.path().to_path_buf());
991 registry
992 .register("session-visible", None, "Darkwood")
993 .unwrap();
994
995 assert!(root
996 .path()
997 .join("by-session/session-visible.json")
998 .is_file());
999 }
1000
1001 #[test]
1002 fn unsafe_session_ids_are_rejected() {
1003 let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
1004 let error = registry
1005 .register("../escape", None, "Darkwood")
1006 .unwrap_err();
1007 assert!(error.to_string().contains("filename-safe"));
1008 }
1009
1010 #[test]
1011 fn lookup_accepts_session_name_and_slug() {
1012 let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
1013 let assignment = registry
1014 .register("session-lookup", Some("Oak"), "Darkwood")
1015 .unwrap();
1016
1017 assert_eq!(registry.lookup(&assignment.session_id).unwrap(), assignment);
1018 assert_eq!(registry.lookup(&assignment.name).unwrap(), assignment);
1019 assert_eq!(registry.lookup(&assignment.slug).unwrap(), assignment);
1020 }
1021
1022 #[test]
1023 fn missing_realm_is_auto_created_and_reused() {
1024 let config_dir = tempfile::tempdir().unwrap();
1025 let realm_file = config_dir.path().join("agent-id/realm");
1026 assert!(!realm_file.exists());
1027
1028 let first = auto_create_realm(&realm_file).unwrap();
1029 assert!(realm_file.is_file());
1030 let contents = fs::read_to_string(&realm_file).unwrap();
1031 assert_eq!(contents.trim(), first);
1032
1033 let second = auto_create_realm(&realm_file).unwrap();
1034 assert_eq!(second, first);
1035 }
1036
1037 #[test]
1038 fn concurrent_realm_creation_keeps_one_value() {
1039 let config_dir = tempfile::tempdir().unwrap();
1040 let realm_file = std::sync::Arc::new(config_dir.path().join("agent-id/realm"));
1041 let handles = (0..8)
1042 .map(|_| {
1043 let realm_file = std::sync::Arc::clone(&realm_file);
1044 std::thread::spawn(move || auto_create_realm(&realm_file).unwrap())
1045 })
1046 .collect::<Vec<_>>();
1047 let mut handles = handles.into_iter();
1048 let first = handles.next().unwrap().join().unwrap();
1049 for handle in handles {
1050 assert_eq!(handle.join().unwrap(), first);
1051 }
1052 }
1053
1054 #[test]
1055 fn cwd_metadata_is_bounded_and_single_line() {
1056 assert_eq!(normalize_cwd(" /tmp/agent-id ").unwrap(), "/tmp/agent-id");
1057 assert!(normalize_cwd("/tmp/agent\nid").is_err());
1058 assert!(normalize_cwd(&"x".repeat(MAX_CWD_CHARS + 1)).is_err());
1059 }
1060
1061 #[test]
1062 fn summaries_are_single_line_and_bounded() {
1063 assert_eq!(
1064 normalize_summary(" Implementing\n activity summaries ").unwrap(),
1065 "Implementing activity summaries"
1066 );
1067 assert!(normalize_summary(" \n\t ").is_err());
1068 assert!(normalize_summary(&"x".repeat(MAX_SUMMARY_CHARS + 1)).is_err());
1069 }
1070
1071 #[test]
1072 fn registering_a_session_updates_existing_identity() {
1073 let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
1074 let first = registry.register("session-1", None, "Darkwood").unwrap();
1075 let second = registry.register("session-1", None, "Darkwood").unwrap();
1076
1077 assert_eq!(second.name, first.name);
1078 assert_eq!(second.session_id, first.session_id);
1079 assert_eq!(second.created_at, first.created_at);
1080 assert!(second.updated_at >= first.updated_at);
1081 }
1082
1083 #[test]
1084 fn lookup_requires_a_registered_session() {
1085 let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
1086 let error = registry.lookup("missing").unwrap_err();
1087 assert!(error.to_string().contains("no identity found"));
1088 }
1089}