use std::{
collections::{BTreeMap, BTreeSet},
env,
fs::{self, File},
io::Write,
path::{Path, PathBuf},
sync::atomic::{AtomicU64, Ordering},
time::{SystemTime, UNIX_EPOCH},
};
use anyhow::{anyhow, bail, Context, Result};
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use crate::{
activity::{ActivityState, ActivityStateValue},
cli::{AnnotateArgs, CurrentArgs, DiscoverArgs, LookupArgs, PruneArgs, RegisterArgs},
names,
};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ActivitySummary {
pub text: String,
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ExtensionMetadata {
pub data: serde_json::Value,
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone)]
pub struct ActivityUpdate {
pub summary: Option<String>,
pub clear_summary: bool,
pub state: Option<ActivityStateValue>,
pub clear_state: bool,
pub cwd: Option<String>,
pub clear_cwd: bool,
pub extensions: BTreeMap<String, serde_json::Value>,
pub clear_extensions: BTreeSet<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct Assignment {
pub version: u8,
pub session_id: String,
pub name: String,
pub slug: String,
pub first_name: String,
pub family_name: String,
pub realm: String,
pub summary: Option<ActivitySummary>,
pub state: ActivityState,
pub cwd: Option<String>,
#[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
pub extensions: BTreeMap<String, ExtensionMetadata>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Serialize)]
pub struct PrunedIdentity {
pub session_id: String,
pub name: String,
pub slug: String,
pub updated_at: DateTime<Utc>,
pub claim_removed: bool,
}
#[derive(Debug, Clone, Serialize)]
pub struct PruneReport {
pub cutoff: DateTime<Utc>,
pub dry_run: bool,
pub candidates: Vec<Assignment>,
pub removed: Vec<PrunedIdentity>,
pub errors: Vec<String>,
}
#[derive(Debug, Clone)]
pub struct Registry {
root: PathBuf,
}
impl Registry {
pub fn from_env() -> Result<Self> {
let root = if let Some(path) = env::var_os("AGENT_ID_HOME") {
PathBuf::from(path)
} else if let Some(path) = env::var_os("XDG_DATA_HOME") {
PathBuf::from(path).join("agent-id")
} else {
home_dir()?.join(".local/share/agent-id")
};
Ok(Self::new(root))
}
pub fn new(root: PathBuf) -> Self {
Self { root }
}
pub fn register(
&self,
session_id: &str,
family_name: Option<&str>,
realm: &str,
) -> Result<Assignment> {
let session_id = validate_session_id(session_id)?;
let session_path = self.session_path(&session_id);
if session_path.exists() {
let mut existing = read_assignment(&session_path)
.with_context(|| format!("read existing assignment {}", session_path.display()))?;
existing.updated_at = Utc::now();
replace_assignment(&session_path, &existing)?;
return Ok(existing);
}
let realm = normalize_component(realm, "realm")?;
let requested_family = family_name
.map(|value| normalize_component(value, "family name"))
.transpose()?;
let first_names = names::first_names();
let family_names = names::family_names();
if first_names.is_empty() || family_names.is_empty() {
bail!("name lists are empty");
}
fs::create_dir_all(self.root.join("by-session"))
.with_context(|| format!("create {}", self.root.display()))?;
fs::create_dir_all(self.root.join("by-name"))
.with_context(|| format!("create {}", self.root.display()))?;
for attempt in 0..100_000_u64 {
let (first, family) = candidate(
&session_id,
attempt,
&first_names,
&family_names,
requested_family.as_deref(),
);
if first == family {
continue;
}
let first_name = title_word(first);
let family_name = title_word(&family);
let name = format!("{first_name} {family_name} of {realm}");
let slug = slug(&first_name, &family_name, &realm);
let claim_path = self.root.join("by-name").join(&slug);
match claim_name(&claim_path, &session_id)? {
Claim::Owned => {}
Claim::Claimed => {}
Claim::Other => continue,
}
let now = Utc::now();
let assignment = Assignment {
version: 1,
session_id: session_id.clone(),
name,
slug,
first_name,
family_name,
realm: realm.clone(),
state: ActivityState::unknown(now),
cwd: None,
summary: None,
extensions: BTreeMap::new(),
created_at: now,
updated_at: now,
};
write_assignment(&session_path, &assignment)?;
return Ok(assignment);
}
bail!("exhausted available names for realm {realm}")
}
pub fn lookup(&self, input: &str) -> Result<Assignment> {
let input = require_nonempty(input, "lookup identifier")?;
if let Ok(session_id) = validate_session_id(&input) {
let session_path = self.session_path(&session_id);
if session_path.exists() {
return self.lookup_session(&session_id);
}
}
for slug in lookup_slugs(&input) {
let claim_path = self.root.join("by-name").join(&slug);
if !claim_path.is_file() {
continue;
}
let session_id = require_nonempty(
&fs::read_to_string(&claim_path)
.with_context(|| format!("read name claim {}", claim_path.display()))?,
"claimed session ID",
)?;
return self
.lookup_session(&session_id)
.with_context(|| format!("resolve name claim {slug}"));
}
bail!(
"no identity found for '{input}'; lookup accepts a session ID, canonical name, or slug"
)
}
fn lookup_session(&self, session_id: &str) -> Result<Assignment> {
let session_id = validate_session_id(session_id)?;
let path = self.session_path(&session_id);
if !path.exists() {
bail!(
"no identity registered for session {session_id}; run `agent-id register {session_id}`"
);
}
let assignment = read_assignment(&path)
.with_context(|| format!("read assignment for session {session_id}"))?;
if assignment.session_id != session_id {
bail!("identity registry entry does not belong to session {session_id}");
}
Ok(assignment)
}
pub fn annotate(&self, session_id: &str, update: ActivityUpdate) -> Result<Assignment> {
let session_id = validate_session_id(session_id)?;
if update.summary.is_none()
&& !update.clear_summary
&& update.state.is_none()
&& !update.clear_state
&& update.cwd.is_none()
&& !update.clear_cwd
&& update.extensions.is_empty()
&& update.clear_extensions.is_empty()
{
bail!("pass at least one activity update");
}
if update.summary.is_some() && update.clear_summary {
bail!("summary and clear_summary are mutually exclusive");
}
if update.state.is_some() && update.clear_state {
bail!("state and clear_state are mutually exclusive");
}
if update.cwd.is_some() && update.clear_cwd {
bail!("cwd and clear_cwd are mutually exclusive");
}
if let Some(owner) = update
.extensions
.keys()
.find(|owner| update.clear_extensions.contains(*owner))
{
bail!("extension {owner} cannot be set and cleared together");
}
let mut assignment = self.lookup_session(&session_id)?;
let summary = update
.summary
.as_deref()
.map(normalize_summary)
.transpose()?;
let cwd = update.cwd.as_deref().map(normalize_cwd).transpose()?;
let now = Utc::now();
if update.summary.is_some() {
assignment.summary = summary.map(|text| ActivitySummary {
text,
updated_at: now,
});
} else if update.clear_summary {
assignment.summary = None;
}
if update.cwd.is_some() {
assignment.cwd = cwd;
} else if update.clear_cwd {
assignment.cwd = None;
}
for (owner, data) in update.extensions {
assignment.extensions.insert(
owner,
ExtensionMetadata {
data,
updated_at: now,
},
);
}
for owner in update.clear_extensions {
assignment.extensions.remove(&owner);
}
if let Some(value) = update.state {
set_omp_activity_state(&mut assignment.extensions, value, now);
} else if update.clear_state {
clear_omp_activity_state(&mut assignment.extensions, now);
}
normalize_omp_activity_state(&mut assignment.extensions, now);
assignment.updated_at = now;
assignment.state = materialize_omp_activity_state(&assignment.extensions, now);
replace_assignment(&self.session_path(&session_id), &assignment)?;
Ok(assignment)
}
pub fn discover(
&self,
recent_hours: Option<i64>,
realm: Option<&str>,
) -> Result<Vec<Assignment>> {
if recent_hours.is_some_and(|hours| hours < 0) {
bail!("--recent must be non-negative");
}
let realm = realm
.map(|value| normalize_component(value, "realm"))
.transpose()?;
let cutoff = recent_hours.map(|hours| Utc::now() - Duration::hours(hours));
let path = self.root.join("by-session");
let entries = match fs::read_dir(&path) {
Ok(entries) => entries,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(error) => return Err(error).with_context(|| format!("read {}", path.display())),
};
let mut assignments = Vec::new();
for entry in entries {
let path = entry?.path();
if path.extension().and_then(|value| value.to_str()) != Some("json") {
continue;
}
let assignment = read_assignment(&path)
.with_context(|| format!("read assignment {}", path.display()))?;
if realm
.as_deref()
.is_some_and(|value| value != assignment.realm)
{
continue;
}
if cutoff.is_some_and(|value| assignment.updated_at < value) {
continue;
}
assignments.push(assignment);
}
assignments.sort_by(|left, right| right.updated_at.cmp(&left.updated_at));
Ok(assignments)
}
pub fn prune(&self, cutoff: DateTime<Utc>, dry_run: bool) -> Result<PruneReport> {
let path = self.root.join("by-session");
let entries = match fs::read_dir(&path) {
Ok(entries) => entries,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
return Ok(PruneReport {
cutoff,
dry_run,
candidates: Vec::new(),
removed: Vec::new(),
errors: Vec::new(),
})
}
Err(error) => return Err(error).with_context(|| format!("read {}", path.display())),
};
let mut candidates = Vec::new();
for entry in entries {
let path = entry?.path();
if path.extension().and_then(|value| value.to_str()) != Some("json") {
continue;
}
let assignment = read_assignment(&path)
.with_context(|| format!("read assignment {}", path.display()))?;
if assignment.updated_at < cutoff {
candidates.push(assignment);
}
}
candidates.sort_by(|left, right| left.updated_at.cmp(&right.updated_at));
let mut report = PruneReport {
cutoff,
dry_run,
candidates,
removed: Vec::new(),
errors: Vec::new(),
};
if dry_run {
return Ok(report);
}
for assignment in &report.candidates {
let session_path = self.session_path(&assignment.session_id);
if let Err(error) = fs::remove_file(&session_path) {
report
.errors
.push(format!("remove {}: {error}", session_path.display()));
continue;
}
let claim_path = self.root.join("by-name").join(&assignment.slug);
let claim_removed = match fs::read_to_string(&claim_path) {
Ok(owner) if owner.trim() == assignment.session_id => {
if let Err(error) = fs::remove_file(&claim_path) {
report.errors.push(format!(
"remove name claim {}: {error}",
claim_path.display()
));
false
} else {
true
}
}
Ok(owner) => {
report.errors.push(format!(
"name claim {} belongs to {}, not {}",
claim_path.display(),
owner.trim(),
assignment.session_id
));
false
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => false,
Err(error) => {
report
.errors
.push(format!("read name claim {}: {error}", claim_path.display()));
false
}
};
report.removed.push(PrunedIdentity {
session_id: assignment.session_id.clone(),
name: assignment.name.clone(),
slug: assignment.slug.clone(),
updated_at: assignment.updated_at,
claim_removed,
});
}
Ok(report)
}
fn session_path(&self, session_id: &str) -> PathBuf {
self.root
.join("by-session")
.join(format!("{session_id}.json"))
}
}
pub fn execute_register(args: &RegisterArgs) -> Result<()> {
let session_id = resolve_session(args.explicit_session())?;
let realm = resolve_realm(args.realm.as_deref())?;
let assignment = Registry::from_env()?.register(&session_id, args.family.as_deref(), &realm)?;
print_assignment(&assignment, args.json)
}
pub fn execute_lookup(args: &LookupArgs) -> Result<()> {
let input = resolve_session(args.explicit_input())?;
let assignment = Registry::from_env()?.lookup(&input)?;
print_assignment(&assignment, args.json)
}
pub fn execute_current(args: &CurrentArgs) -> Result<()> {
let session_id = resolve_session(None)?;
let assignment = Registry::from_env()?.lookup(&session_id)?;
print_assignment(&assignment, args.json)
}
pub fn execute_annotate(args: &AnnotateArgs) -> Result<()> {
let session_id = resolve_session(args.explicit_session())?;
let update = ActivityUpdate {
summary: args.summary.clone(),
clear_summary: args.clear_summary,
state: args.state,
clear_state: args.clear_state,
cwd: args.cwd.clone(),
clear_cwd: args.clear_cwd,
extensions: parse_extension_updates(&args.extensions)?,
clear_extensions: parse_extension_owners(&args.clear_extensions)?,
};
let assignment = Registry::from_env()?.annotate(&session_id, update)?;
print_assignment(&assignment, args.json)
}
pub fn execute_discover(args: &DiscoverArgs) -> Result<()> {
let assignments = Registry::from_env()?.discover(args.recent, args.realm.as_deref())?;
let mut records = crate::herdr::augment_discovery(assignments, args.all);
if !args.all {
records.retain(|record| record.assignment.state.value != ActivityStateValue::Stopped);
}
if args.limit > 0 {
records.truncate(args.limit);
}
if args.json {
println!("{}", serde_json::to_string_pretty(&records)?);
} else if records.is_empty() {
println!("(no identities)");
} else {
for record in records {
let assignment = &record.assignment;
let mut annotations = Vec::new();
annotations.push(format!("state:{}", assignment.state.value));
if let Some(summary) = assignment.summary.as_ref() {
annotations.push(format!("summary:{}", summary.text));
}
if let Some(cwd) = assignment.cwd.as_ref() {
annotations.push(format!("cwd:{cwd}"));
}
if let Some(runtime) = record.runtime {
for location in runtime.locations {
let workspace = location
.workspace_label
.as_deref()
.unwrap_or(&location.workspace_id);
annotations.push(format!(
"herdr:{} pane:{} workspace:{}",
location.agent_status, location.pane_id, workspace
));
}
}
println!(
"{}\t{}\t{}",
assignment.name,
assignment.session_id,
annotations.join("\t")
);
}
}
Ok(())
}
pub fn execute_prune(args: &PruneArgs) -> Result<()> {
let cutoff = DateTime::parse_from_rfc3339(&args.before)
.with_context(|| format!("parse --before timestamp {}", args.before))?
.with_timezone(&Utc);
let report = Registry::from_env()?.prune(cutoff, args.dry_run)?;
if args.json {
println!("{}", serde_json::to_string_pretty(&report)?);
} else {
let action = if args.dry_run {
"would prune"
} else {
"pruned"
};
println!(
"{action} {} identities before {}",
report.candidates.len(),
cutoff
);
for assignment in &report.candidates {
println!(
"{}\t{}\tupdated:{}",
assignment.name, assignment.session_id, assignment.updated_at
);
}
for error in &report.errors {
eprintln!("agent-id: {error}");
}
}
if report.errors.is_empty() {
Ok(())
} else {
bail!("prune completed with {} errors", report.errors.len())
}
}
pub fn resolve_session(explicit: Option<&str>) -> Result<String> {
if let Some(session_id) = explicit.filter(|value| !value.trim().is_empty()) {
return Ok(session_id.trim().to_string());
}
if let Some(value) = env::var_os("AGENT_ID_SESSION_ID") {
let value = value.to_string_lossy();
if !value.trim().is_empty() {
return Ok(value.trim().to_string());
}
}
bail!("no session ID found; pass SESSION_ID or --session-id, or set AGENT_ID_SESSION_ID")
}
fn resolve_realm(explicit: Option<&str>) -> Result<String> {
if let Some(realm) = explicit.filter(|value| !value.trim().is_empty()) {
return normalize_component(realm, "realm");
}
if let Some(realm) = env::var_os("AGENT_REALM") {
let realm = realm.to_string_lossy();
if !realm.trim().is_empty() {
return normalize_component(&realm, "realm");
}
}
let home = home_dir()?;
let config_home = env::var_os("XDG_CONFIG_HOME")
.map(PathBuf::from)
.unwrap_or_else(|| home.join(".config"));
let realm_path = config_home.join("agent-id/realm");
if realm_path.is_file() {
let value = fs::read_to_string(&realm_path)
.with_context(|| format!("read realm from {}", realm_path.display()))?;
return normalize_component(&value, "realm");
}
auto_create_realm(&realm_path)
}
fn auto_create_realm(path: &Path) -> Result<String> {
let candidates = names::candidate_realms();
if candidates.is_empty() {
bail!("bundled realm candidate list is empty");
}
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or_default();
let pid = std::process::id();
let hostname = env::var("HOSTNAME").unwrap_or_default();
let digest = Sha256::digest(format!("{nanos}:{pid}:{hostname}").as_bytes());
let index = usize::try_from(u64::from_be_bytes(digest[0..8].try_into().unwrap())).unwrap_or(0)
% candidates.len();
let realm = normalize_component(candidates[index], "realm")?;
let parent = path
.parent()
.ok_or_else(|| anyhow!("realm path has no parent directory"))?;
fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
let temp = temporary_path(path);
let mut file = File::create(&temp)
.with_context(|| format!("create temporary realm {}", temp.display()))?;
file.write_all(format!("{realm}\n").as_bytes())?;
file.sync_all()?;
match fs::hard_link(&temp, path) {
Ok(()) => {
let _ = fs::remove_file(temp);
Ok(realm)
}
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
let _ = fs::remove_file(temp);
let value = fs::read_to_string(path)
.with_context(|| format!("read existing realm {}", path.display()))?;
normalize_component(&value, "realm")
}
Err(error) => {
let _ = fs::remove_file(temp);
Err(error).with_context(|| format!("create realm file {}", path.display()))
}
}
}
fn print_assignment(assignment: &Assignment, json: bool) -> Result<()> {
if json {
println!("{}", serde_json::to_string_pretty(assignment)?);
} else {
println!("{}", assignment.name);
}
Ok(())
}
fn lookup_slugs(input: &str) -> Vec<String> {
let mut slugs = Vec::new();
let input = input.trim();
if is_slug(input) {
slugs.push(input.to_ascii_lowercase());
}
if let Some(slug) = canonical_name_slug(input) {
if !slugs.contains(&slug) {
slugs.push(slug);
}
}
slugs
}
fn canonical_name_slug(input: &str) -> Option<String> {
let mut parts = input.split_whitespace();
let first = normalize_component(parts.next()?, "first name").ok()?;
let family = normalize_component(parts.next()?, "family name").ok()?;
if !parts.next()?.eq_ignore_ascii_case("of") {
return None;
}
let realm = normalize_component(parts.next()?, "realm").ok()?;
if parts.next().is_some() {
return None;
}
Some(slug(&first, &family, &realm))
}
fn is_slug(input: &str) -> bool {
!input.is_empty()
&& !input.starts_with('-')
&& !input.ends_with('-')
&& input
.chars()
.all(|character| character.is_ascii_alphanumeric() || character == '-')
}
fn candidate<'a>(
session_id: &str,
attempt: u64,
first_names: &'a [&'a str],
family_names: &'a [&'a str],
requested_family: Option<&str>,
) -> (&'a str, String) {
let digest = Sha256::digest(format!("{session_id}:{attempt}").as_bytes());
let first_index = usize::try_from(u64::from_be_bytes(digest[0..8].try_into().unwrap()))
.unwrap_or(0)
% first_names.len();
let family_index = usize::try_from(u64::from_be_bytes(digest[8..16].try_into().unwrap()))
.unwrap_or(0)
% family_names.len();
let family = requested_family
.map(ToOwned::to_owned)
.unwrap_or_else(|| family_names[family_index].to_string());
(first_names[first_index], family)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Claim {
Claimed,
Owned,
Other,
}
fn claim_name(path: &Path, session_id: &str) -> Result<Claim> {
let parent = path
.parent()
.ok_or_else(|| anyhow!("name claim has no parent directory"))?;
fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
let temp = temporary_path(path);
let mut file = File::create(&temp)
.with_context(|| format!("create temporary claim {}", temp.display()))?;
file.write_all(session_id.as_bytes())?;
file.write_all(b"\n")?;
file.sync_all()?;
let claim = match fs::hard_link(&temp, path) {
Ok(()) => Claim::Claimed,
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
let owner = fs::read_to_string(path).unwrap_or_default();
if owner.trim() == session_id {
Claim::Owned
} else {
Claim::Other
}
}
Err(error) => return Err(error).with_context(|| format!("claim name {}", path.display())),
};
let _ = fs::remove_file(temp);
Ok(claim)
}
#[derive(Debug, Serialize)]
struct PersistedAssignment<'a> {
version: u8,
session_id: &'a str,
name: &'a str,
slug: &'a str,
first_name: &'a str,
family_name: &'a str,
realm: &'a str,
summary: &'a Option<ActivitySummary>,
cwd: &'a Option<String>,
#[serde(skip_serializing_if = "BTreeMap::is_empty")]
extensions: &'a BTreeMap<String, ExtensionMetadata>,
created_at: DateTime<Utc>,
updated_at: DateTime<Utc>,
}
fn persisted_assignment(assignment: &Assignment) -> PersistedAssignment<'_> {
PersistedAssignment {
version: assignment.version,
session_id: &assignment.session_id,
name: &assignment.name,
slug: &assignment.slug,
first_name: &assignment.first_name,
family_name: &assignment.family_name,
realm: &assignment.realm,
summary: &assignment.summary,
cwd: &assignment.cwd,
extensions: &assignment.extensions,
created_at: assignment.created_at,
updated_at: assignment.updated_at,
}
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct StoredAssignment {
version: u8,
session_id: String,
name: String,
slug: String,
first_name: String,
family_name: String,
realm: String,
#[serde(default)]
summary: Option<ActivitySummary>,
#[serde(default)]
cwd: Option<String>,
#[serde(default)]
extensions: BTreeMap<String, ExtensionMetadata>,
created_at: DateTime<Utc>,
updated_at: DateTime<Utc>,
}
impl StoredAssignment {
fn into_assignment(self) -> Assignment {
let StoredAssignment {
version,
session_id,
name,
slug,
first_name,
family_name,
realm,
summary,
cwd,
extensions,
created_at,
updated_at,
} = self;
let state = materialize_omp_activity_state(&extensions, updated_at);
Assignment {
version,
session_id,
name,
slug,
first_name,
family_name,
realm,
summary,
state,
cwd,
extensions,
created_at,
updated_at,
}
}
}
fn write_assignment(path: &Path, assignment: &Assignment) -> Result<()> {
let parent = path
.parent()
.ok_or_else(|| anyhow!("assignment has no parent directory"))?;
fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
let contents = format!(
"{}\n",
serde_json::to_string_pretty(&persisted_assignment(assignment))?
);
let temp = temporary_path(path);
let mut file = File::create(&temp)
.with_context(|| format!("create temporary assignment {}", temp.display()))?;
file.write_all(contents.as_bytes())?;
file.sync_all()?;
let result = match fs::hard_link(&temp, path) {
Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
bail!(
"session {} already has a registered identity",
assignment.session_id
)
}
Err(error) => Err(error).with_context(|| format!("write assignment {}", path.display())),
};
let _ = fs::remove_file(temp);
result
}
fn replace_assignment(path: &Path, assignment: &Assignment) -> Result<()> {
let parent = path
.parent()
.ok_or_else(|| anyhow!("assignment has no parent directory"))?;
fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
let contents = format!(
"{}\n",
serde_json::to_string_pretty(&persisted_assignment(assignment))?
);
let temp = temporary_path(path);
let mut file = File::create(&temp)
.with_context(|| format!("create temporary assignment {}", temp.display()))?;
file.write_all(contents.as_bytes())?;
file.sync_all()?;
let result =
fs::rename(&temp, path).with_context(|| format!("replace assignment {}", path.display()));
if result.is_err() {
let _ = fs::remove_file(temp);
}
result
}
fn read_assignment(path: &Path) -> Result<Assignment> {
let contents = fs::read_to_string(path)?;
Ok(serde_json::from_str::<StoredAssignment>(&contents)?.into_assignment())
}
const OMP_EXTENSION_OWNER: &str = "omp";
const OMP_STATE_KEY: &str = "state";
fn omp_activity_state(extensions: &BTreeMap<String, ExtensionMetadata>) -> Option<ActivityState> {
let metadata = extensions.get(OMP_EXTENSION_OWNER)?;
let state = metadata.data.get(OMP_STATE_KEY)?;
if let Some(value) = state.as_str() {
return Some(ActivityState::from_external(value, metadata.updated_at));
}
serde_json::from_value(state.clone()).ok()
}
fn materialize_omp_activity_state(
extensions: &BTreeMap<String, ExtensionMetadata>,
fallback_updated_at: DateTime<Utc>,
) -> ActivityState {
omp_activity_state(extensions).unwrap_or_else(|| ActivityState::unknown(fallback_updated_at))
}
fn set_omp_activity_state(
extensions: &mut BTreeMap<String, ExtensionMetadata>,
value: ActivityStateValue,
updated_at: DateTime<Utc>,
) {
let metadata = extensions
.entry(OMP_EXTENSION_OWNER.to_string())
.or_insert_with(|| ExtensionMetadata {
data: serde_json::json!({}),
updated_at,
});
if !metadata.data.is_object() {
metadata.data = serde_json::json!({});
}
metadata
.data
.as_object_mut()
.expect("state extension data was just normalized to an object")
.insert(
OMP_STATE_KEY.to_string(),
serde_json::json!({
"value": value,
"updated_at": updated_at,
}),
);
metadata.updated_at = updated_at;
}
fn clear_omp_activity_state(
extensions: &mut BTreeMap<String, ExtensionMetadata>,
updated_at: DateTime<Utc>,
) {
let Some(metadata) = extensions.get_mut(OMP_EXTENSION_OWNER) else {
return;
};
if let Some(data) = metadata.data.as_object_mut() {
data.remove(OMP_STATE_KEY);
metadata.updated_at = updated_at;
}
}
fn normalize_omp_activity_state(
extensions: &mut BTreeMap<String, ExtensionMetadata>,
updated_at: DateTime<Utc>,
) {
let value = extensions
.get(OMP_EXTENSION_OWNER)
.and_then(|metadata| metadata.data.get(OMP_STATE_KEY))
.and_then(|state| state.as_str())
.map(ActivityStateValue::from_external);
if let Some(value) = value {
set_omp_activity_state(extensions, value, updated_at);
}
}
static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
fn temporary_path(path: &Path) -> PathBuf {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or_default();
let counter = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
path.with_extension(format!("tmp-{}-{nanos}-{counter}", std::process::id()))
}
fn slug(first_name: &str, family_name: &str, realm: &str) -> String {
[first_name, family_name, realm]
.iter()
.map(|part| part.to_ascii_lowercase().replace([' ', '\'', '_'], "-"))
.collect::<Vec<_>>()
.join("-")
}
fn normalize_component(value: &str, kind: &str) -> Result<String> {
let value = value.trim();
let valid = !value.is_empty()
&& value.chars().all(|character| {
character.is_ascii_alphabetic() || character == '-' || character == '\''
})
&& value
.chars()
.next()
.is_some_and(|character| character.is_ascii_alphabetic())
&& value
.chars()
.last()
.is_some_and(|character| character.is_ascii_alphabetic());
if !valid {
bail!("{kind} must contain only letters, apostrophes, or hyphens")
}
Ok(title_word(value))
}
fn title_word(value: &str) -> String {
let mut characters = value.chars();
let Some(first) = characters.next() else {
return String::new();
};
first.to_uppercase().collect::<String>() + &characters.as_str().to_ascii_lowercase()
}
fn validate_session_id(value: &str) -> Result<String> {
let value = require_nonempty(value, "session ID")?;
let safe = value
.chars()
.all(|character| character.is_ascii_alphanumeric() || matches!(character, '-' | '_' | '.'));
if !safe || value == "." || value == ".." {
bail!("session ID must be a filename-safe value using letters, digits, '.', '_' or '-'");
}
Ok(value)
}
const MAX_EXTENSION_OWNER_CHARS: usize = 64;
const MAX_EXTENSION_JSON_BYTES: usize = 16 * 1024;
fn parse_extension_updates(values: &[String]) -> Result<BTreeMap<String, serde_json::Value>> {
let mut extensions = BTreeMap::new();
for value in values {
let (owner, json) = value
.split_once('=')
.ok_or_else(|| anyhow!("--extension must use OWNER=JSON"))?;
let owner = normalize_extension_owner(owner)?;
if json.len() > MAX_EXTENSION_JSON_BYTES {
bail!("extension {owner} JSON must be at most {MAX_EXTENSION_JSON_BYTES} bytes");
}
let data = serde_json::from_str(json)
.with_context(|| format!("parse JSON for extension {owner}"))?;
if extensions.insert(owner.clone(), data).is_some() {
bail!("extension {owner} was provided more than once");
}
}
Ok(extensions)
}
fn parse_extension_owners(values: &[String]) -> Result<BTreeSet<String>> {
values
.iter()
.map(|owner| normalize_extension_owner(owner))
.collect()
}
fn normalize_extension_owner(value: &str) -> Result<String> {
let owner = require_nonempty(value, "extension owner")?;
let valid = owner.chars().count() <= MAX_EXTENSION_OWNER_CHARS
&& owner
.chars()
.next()
.is_some_and(|character| character.is_ascii_lowercase() || character.is_ascii_digit())
&& owner.chars().all(|character| {
character.is_ascii_lowercase()
|| character.is_ascii_digit()
|| matches!(character, '.' | '_' | '-')
});
if !valid {
bail!(
"extension owner must be at most {MAX_EXTENSION_OWNER_CHARS} characters using lowercase letters, digits, '.', '_' or '-'"
);
}
Ok(owner)
}
const MAX_SUMMARY_CHARS: usize = 240;
fn normalize_summary(value: &str) -> Result<String> {
let mut normalized = String::new();
for word in value.split_whitespace() {
if !normalized.is_empty() {
normalized.push(' ');
}
normalized.push_str(word);
}
if normalized.is_empty() {
bail!("summary cannot be empty");
}
if normalized.chars().count() > MAX_SUMMARY_CHARS {
bail!("summary must be at most {MAX_SUMMARY_CHARS} characters");
}
Ok(normalized)
}
const MAX_CWD_CHARS: usize = 4096;
fn normalize_cwd(value: &str) -> Result<String> {
let value = require_nonempty(value, "working directory")?;
if value.chars().any(char::is_control) {
bail!("working directory cannot contain control characters");
}
if value.chars().count() > MAX_CWD_CHARS {
bail!("working directory must be at most {MAX_CWD_CHARS} characters");
}
Ok(value)
}
fn require_nonempty(value: &str, kind: &str) -> Result<String> {
let value = value.trim();
if value.is_empty() {
bail!("{kind} cannot be empty")
}
Ok(value.to_string())
}
fn home_dir() -> Result<PathBuf> {
env::var_os("HOME")
.map(PathBuf::from)
.ok_or_else(|| anyhow!("HOME is not set; set AGENT_ID_HOME explicitly"))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn generated_identity_has_canonical_parts() {
let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
let assignment = registry
.register("session-1", Some("Oak"), "Darkwood")
.unwrap();
assert_eq!(assignment.family_name, "Oak");
assert_eq!(assignment.realm, "Darkwood");
assert_eq!(
assignment.slug,
format!(
"{}-oak-darkwood",
assignment.first_name.to_ascii_lowercase()
)
);
assert_eq!(
assignment.name,
format!("{} Oak of Darkwood", assignment.first_name)
);
}
#[test]
fn session_id_is_the_registry_filename() {
let root = tempfile::tempdir().unwrap();
let registry = Registry::new(root.path().to_path_buf());
registry
.register("session-visible", None, "Darkwood")
.unwrap();
assert!(root
.path()
.join("by-session/session-visible.json")
.is_file());
}
#[test]
fn unsafe_session_ids_are_rejected() {
let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
let error = registry
.register("../escape", None, "Darkwood")
.unwrap_err();
assert!(error.to_string().contains("filename-safe"));
}
#[test]
fn lookup_accepts_session_name_and_slug() {
let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
let assignment = registry
.register("session-lookup", Some("Oak"), "Darkwood")
.unwrap();
assert_eq!(registry.lookup(&assignment.session_id).unwrap(), assignment);
assert_eq!(registry.lookup(&assignment.name).unwrap(), assignment);
assert_eq!(registry.lookup(&assignment.slug).unwrap(), assignment);
}
#[test]
fn missing_realm_is_auto_created_and_reused() {
let config_dir = tempfile::tempdir().unwrap();
let realm_file = config_dir.path().join("agent-id/realm");
assert!(!realm_file.exists());
let first = auto_create_realm(&realm_file).unwrap();
assert!(realm_file.is_file());
let contents = fs::read_to_string(&realm_file).unwrap();
assert_eq!(contents.trim(), first);
let second = auto_create_realm(&realm_file).unwrap();
assert_eq!(second, first);
}
#[test]
fn concurrent_realm_creation_keeps_one_value() {
let config_dir = tempfile::tempdir().unwrap();
let realm_file = std::sync::Arc::new(config_dir.path().join("agent-id/realm"));
let handles = (0..8)
.map(|_| {
let realm_file = std::sync::Arc::clone(&realm_file);
std::thread::spawn(move || auto_create_realm(&realm_file).unwrap())
})
.collect::<Vec<_>>();
let mut handles = handles.into_iter();
let first = handles.next().unwrap().join().unwrap();
for handle in handles {
assert_eq!(handle.join().unwrap(), first);
}
}
#[test]
fn cwd_metadata_is_bounded_and_single_line() {
assert_eq!(normalize_cwd(" /tmp/agent-id ").unwrap(), "/tmp/agent-id");
assert!(normalize_cwd("/tmp/agent\nid").is_err());
assert!(normalize_cwd(&"x".repeat(MAX_CWD_CHARS + 1)).is_err());
}
#[test]
fn summaries_are_single_line_and_bounded() {
assert_eq!(
normalize_summary(" Implementing\n activity summaries ").unwrap(),
"Implementing activity summaries"
);
assert!(normalize_summary(" \n\t ").is_err());
assert!(normalize_summary(&"x".repeat(MAX_SUMMARY_CHARS + 1)).is_err());
}
#[test]
fn registering_a_session_updates_existing_identity() {
let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
let first = registry.register("session-1", None, "Darkwood").unwrap();
let second = registry.register("session-1", None, "Darkwood").unwrap();
assert_eq!(second.name, first.name);
assert_eq!(second.session_id, first.session_id);
assert_eq!(second.created_at, first.created_at);
assert!(second.updated_at >= first.updated_at);
}
#[test]
fn lookup_requires_a_registered_session() {
let registry = Registry::new(tempfile::tempdir().unwrap().path().to_path_buf());
let error = registry.lookup("missing").unwrap_err();
assert!(error.to_string().contains("no identity found"));
}
}