Skip to main content

agent_id_cli/
registry.rs

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