Skip to main content

agent_memory/
lib.rs

1mod cli;
2mod completion;
3
4use std::env;
5use std::ffi::OsString;
6use std::fs;
7use std::path::{Path, PathBuf};
8
9use clap::error::ErrorKind;
10use clap::{CommandFactory, Parser};
11use serde_json::json;
12
13use cli::{Cli, Command, IdArgs, ScopeArgs};
14
15use nils_common::cli_contract::exit;
16use nils_common::fs::display_path;
17
18const EXIT_OK: i32 = exit::SUCCESS;
19const EXIT_RUNTIME: i32 = exit::RUNTIME;
20const EXIT_USAGE: i32 = exit::USAGE;
21
22pub fn run() -> i32 {
23    run_with_args(env::args_os())
24}
25
26pub fn run_with_args<I, T>(args: I) -> i32
27where
28    I: IntoIterator<Item = T>,
29    T: Into<OsString>,
30{
31    let args: Vec<OsString> = args.into_iter().map(Into::into).collect();
32    if args.len() == 1 {
33        return match print_help() {
34            Ok(code) => code,
35            Err(err) => {
36                eprintln!("agent-memory: {}", err.message);
37                err.exit_code
38            }
39        };
40    }
41
42    let cli = match Cli::try_parse_from(args) {
43        Ok(cli) => cli,
44        Err(err) => {
45            let code = match err.kind() {
46                ErrorKind::DisplayHelp | ErrorKind::DisplayVersion => err.exit_code(),
47                _ => EXIT_USAGE,
48            };
49            let _ = err.print();
50            return code;
51        }
52    };
53
54    match dispatch(cli) {
55        Ok(code) => code,
56        Err(err) => {
57            eprintln!("agent-memory: {}", err.message);
58            err.exit_code
59        }
60    }
61}
62
63fn dispatch(cli: Cli) -> Result<i32, CliError> {
64    let layout = Layout::from_env()?;
65    match cli.command {
66        Command::Path(args) => {
67            println!(
68                "{}",
69                display_path(&layout.resolve_scope(args.scope.as_deref()))
70            );
71            Ok(EXIT_OK)
72        }
73        Command::List(args) => list_scope(&layout, &args),
74        Command::Index(args) => index_scope(&layout, &args),
75        Command::Agents => list_named_dirs(&layout.agents_dir()),
76        Command::Personas => list_named_dirs(&layout.personas_dir()),
77        Command::InitAgent(args) => init_agent(&layout, &args),
78        Command::InitPersona(args) => init_persona(&layout, &args),
79        Command::Resolve(args) => resolve_agent(&layout, &args),
80        Command::Env => {
81            print_env(&layout);
82            Ok(EXIT_OK)
83        }
84        Command::Doctor => doctor(&layout),
85        Command::Completion(args) => Ok(completion::run(args.shell)),
86        Command::Help => print_help(),
87    }
88}
89
90fn print_help() -> Result<i32, CliError> {
91    let mut command = Cli::command();
92    command
93        .print_long_help()
94        .map_err(|err| CliError::runtime(format!("failed to print help: {err}")))?;
95    println!();
96    Ok(EXIT_OK)
97}
98
99#[derive(Debug)]
100struct CliError {
101    message: String,
102    exit_code: i32,
103}
104
105impl CliError {
106    fn runtime(message: impl Into<String>) -> Self {
107        Self {
108            message: message.into(),
109            exit_code: EXIT_RUNTIME,
110        }
111    }
112
113    fn usage(message: impl Into<String>) -> Self {
114        Self {
115            message: message.into(),
116            exit_code: EXIT_USAGE,
117        }
118    }
119}
120
121#[derive(Debug)]
122struct Layout {
123    root: PathBuf,
124}
125
126impl Layout {
127    fn from_env() -> Result<Self, CliError> {
128        if let Some(value) = non_empty_env("AGENT_MEMORY_HOME") {
129            return Ok(Self {
130                root: PathBuf::from(value),
131            });
132        }
133
134        if let Some(value) = non_empty_env("XDG_CONFIG_HOME") {
135            return Ok(Self {
136                root: PathBuf::from(value).join("agent-memory"),
137            });
138        }
139
140        let home = non_empty_env("HOME").ok_or_else(|| {
141            CliError::runtime("HOME is not set and AGENT_MEMORY_HOME was omitted")
142        })?;
143
144        Ok(Self {
145            root: PathBuf::from(home).join(".config").join("agent-memory"),
146        })
147    }
148
149    fn global_dir(&self) -> PathBuf {
150        self.root.join("global")
151    }
152
153    fn agents_dir(&self) -> PathBuf {
154        self.root.join("agents")
155    }
156
157    fn personas_dir(&self) -> PathBuf {
158        self.root.join("personas")
159    }
160
161    fn resolve_scope(&self, scope: Option<&str>) -> PathBuf {
162        match scope.unwrap_or("root") {
163            "" | "root" => self.root.clone(),
164            "global" => self.global_dir(),
165            value if value.starts_with("agents/") || value.starts_with("personas/") => {
166                self.root.join(value)
167            }
168            value => self.agents_dir().join(value),
169        }
170    }
171}
172
173fn non_empty_env(name: &str) -> Option<OsString> {
174    env::var_os(name).filter(|value| !value.is_empty())
175}
176
177fn list_scope(layout: &Layout, args: &ScopeArgs) -> Result<i32, CliError> {
178    let path = layout.resolve_scope(args.scope.as_deref().or(Some("global")));
179    require_dir(&path)?;
180
181    for file in markdown_files(&path)? {
182        if let Some(name) = file.file_name() {
183            println!("{}", name.to_string_lossy());
184        }
185    }
186    Ok(EXIT_OK)
187}
188
189fn index_scope(layout: &Layout, args: &ScopeArgs) -> Result<i32, CliError> {
190    let path = layout.resolve_scope(args.scope.as_deref().or(Some("global")));
191    require_dir(&path)?;
192
193    let index = path.join("MEMORY.md");
194    if !index.is_file() {
195        return Err(CliError::runtime(format!(
196            "no MEMORY.md in {}",
197            display_path(&path)
198        )));
199    }
200
201    let contents = fs::read_to_string(&index).map_err(|err| {
202        CliError::runtime(format!("failed to read {}: {err}", display_path(&index)))
203    })?;
204    print!("{contents}");
205    Ok(EXIT_OK)
206}
207
208fn list_named_dirs(path: &Path) -> Result<i32, CliError> {
209    if !path.is_dir() {
210        return Ok(EXIT_OK);
211    }
212
213    for dir in child_dirs(path)? {
214        if let Some(name) = dir.file_name() {
215            println!("{}", name.to_string_lossy());
216        }
217    }
218    Ok(EXIT_OK)
219}
220
221fn init_agent(layout: &Layout, args: &IdArgs) -> Result<i32, CliError> {
222    validate_id(&args.id)?;
223    let path = layout.agents_dir().join(&args.id);
224    if path.exists() {
225        return Err(CliError::runtime(format!(
226            "already exists: {}",
227            display_path(&path)
228        )));
229    }
230
231    fs::create_dir_all(&path).map_err(|err| {
232        CliError::runtime(format!("failed to create {}: {err}", display_path(&path)))
233    })?;
234    fs::write(
235        path.join("MEMORY.md"),
236        format!("# Memory index ({})\n\n", args.id),
237    )
238    .map_err(|err| CliError::runtime(format!("failed to write MEMORY.md: {err}")))?;
239
240    println!("created: {}", display_path(&path));
241    Ok(EXIT_OK)
242}
243
244fn init_persona(layout: &Layout, args: &IdArgs) -> Result<i32, CliError> {
245    validate_id(&args.id)?;
246    let path = layout.personas_dir().join(&args.id);
247    if path.exists() {
248        return Err(CliError::runtime(format!(
249            "already exists: {}",
250            display_path(&path)
251        )));
252    }
253
254    let memory_dir = path.join("memory");
255    let claude_dir = path.join(".claude");
256    fs::create_dir_all(&memory_dir).map_err(|err| {
257        CliError::runtime(format!(
258            "failed to create {}: {err}",
259            display_path(&memory_dir)
260        ))
261    })?;
262    fs::create_dir_all(&claude_dir).map_err(|err| {
263        CliError::runtime(format!(
264            "failed to create {}: {err}",
265            display_path(&claude_dir)
266        ))
267    })?;
268
269    fs::write(path.join("CLAUDE.md"), persona_claude_template(&args.id)).map_err(|err| {
270        CliError::runtime(format!(
271            "failed to write {}: {err}",
272            display_path(&path.join("CLAUDE.md"))
273        ))
274    })?;
275    fs::write(
276        memory_dir.join("MEMORY.md"),
277        format!("# Memory index ({} persona)\n\n", args.id),
278    )
279    .map_err(|err| CliError::runtime(format!("failed to write persona MEMORY.md: {err}")))?;
280
281    let settings = json!({
282        "autoMemoryDirectory": to_tilde(&memory_dir),
283    });
284    fs::write(
285        claude_dir.join("settings.local.json"),
286        format!(
287            "{}\n",
288            serde_json::to_string_pretty(&settings).expect("settings json should serialize")
289        ),
290    )
291    .map_err(|err| CliError::runtime(format!("failed to write persona settings: {err}")))?;
292
293    println!("created: {}", display_path(&path));
294    println!("  next: $EDITOR {}/CLAUDE.md", display_path(&path));
295    println!(
296        "  launch: claude-{}  (after sourcing shell/agent-memory.zsh)",
297        args.id
298    );
299    Ok(EXIT_OK)
300}
301
302fn resolve_agent(layout: &Layout, args: &IdArgs) -> Result<i32, CliError> {
303    validate_id(&args.id)?;
304    println!("global\t{}", display_path(&layout.global_dir()));
305    println!(
306        "agent\t{}",
307        display_path(&layout.agents_dir().join(&args.id))
308    );
309    Ok(EXIT_OK)
310}
311
312fn print_env(layout: &Layout) {
313    println!(
314        "export AGENT_MEMORY_HOME={}",
315        shell_escape(&display_path(&layout.root))
316    );
317    println!(
318        "export AGENT_MEMORY_GLOBAL={}",
319        shell_escape(&display_path(&layout.global_dir()))
320    );
321    println!(
322        "export AGENT_MEMORY_AGENTS={}",
323        shell_escape(&display_path(&layout.agents_dir()))
324    );
325    println!(
326        "export AGENT_MEMORY_PERSONAS={}",
327        shell_escape(&display_path(&layout.personas_dir()))
328    );
329}
330
331fn doctor(layout: &Layout) -> Result<i32, CliError> {
332    let mut failed = false;
333
334    println!("AGENT_MEMORY_HOME={}", display_path(&layout.root));
335    if layout.root.is_dir() {
336        println!("  [ok]      root present");
337    } else {
338        eprintln!("  [missing] root");
339        failed = true;
340    }
341
342    let global = layout.global_dir();
343    match fs::symlink_metadata(&global) {
344        Ok(metadata) if metadata.file_type().is_symlink() => {
345            let target = fs::read_link(&global)
346                .map(|path| display_path(&path))
347                .unwrap_or_else(|_| "<unreadable>".to_string());
348            if global.is_dir() {
349                println!("  [ok]      global -> {target}");
350            } else {
351                eprintln!("  [broken]  global -> {target}");
352                failed = true;
353            }
354        }
355        Ok(metadata) if metadata.is_dir() => {
356            println!("  [ok]      global (real dir)");
357        }
358        _ => {
359            eprintln!("  [missing] global");
360            failed = true;
361        }
362    }
363
364    print_dir_count("agents/", &layout.agents_dir())?;
365    print_dir_count("personas/", &layout.personas_dir())?;
366
367    if failed {
368        Ok(EXIT_RUNTIME)
369    } else {
370        Ok(EXIT_OK)
371    }
372}
373
374fn print_dir_count(label: &str, path: &Path) -> Result<(), CliError> {
375    if path.is_dir() {
376        println!(
377            "  [ok]      {label:<10}({} entries)",
378            child_dirs(path)?.len()
379        );
380    } else if label == "agents/" {
381        println!("  [empty]   agents/   (run 'agent-memory init-agent <id>')");
382    } else if label == "personas/" {
383        println!("  [empty]   personas/ (run 'agent-memory init-persona <id>')");
384    } else {
385        println!("  [empty]   {label:<10}");
386    }
387    Ok(())
388}
389
390fn require_dir(path: &Path) -> Result<(), CliError> {
391    if path.is_dir() {
392        Ok(())
393    } else {
394        Err(CliError::runtime(format!(
395            "not found: {}",
396            display_path(path)
397        )))
398    }
399}
400
401fn validate_id(id: &str) -> Result<(), CliError> {
402    if id.is_empty() || id.contains('/') || id == "." || id == ".." {
403        return Err(CliError::usage(format!("invalid id: '{id}'")));
404    }
405    Ok(())
406}
407
408fn markdown_files(path: &Path) -> Result<Vec<PathBuf>, CliError> {
409    let mut files = Vec::new();
410    for entry in fs::read_dir(path)
411        .map_err(|err| CliError::runtime(format!("failed to read {}: {err}", display_path(path))))?
412    {
413        let entry =
414            entry.map_err(|err| CliError::runtime(format!("failed to read entry: {err}")))?;
415        let path = entry.path();
416        if path.is_file() && path.extension().is_some_and(|ext| ext == "md") {
417            files.push(path);
418        }
419    }
420    files.sort();
421    Ok(files)
422}
423
424fn child_dirs(path: &Path) -> Result<Vec<PathBuf>, CliError> {
425    let mut dirs = Vec::new();
426    for entry in fs::read_dir(path)
427        .map_err(|err| CliError::runtime(format!("failed to read {}: {err}", display_path(path))))?
428    {
429        let entry =
430            entry.map_err(|err| CliError::runtime(format!("failed to read entry: {err}")))?;
431        let path = entry.path();
432        if path.is_dir() {
433            dirs.push(path);
434        }
435    }
436    dirs.sort();
437    Ok(dirs)
438}
439
440fn persona_claude_template(id: &str) -> String {
441    format!(
442        r#"# Persona: {id}
443
444Claude Code session scoped to "{id}" persona. Loads on top of the base
445`~/.claude/CLAUDE.md` policies (this file is additive, not a replacement).
446
447## Scope
448
449- In scope: <fill in what this persona handles>
450- Out of scope: anything outside the persona's domain - recommend exiting
451  to base `claude` or another persona.
452
453## Memory
454
455- Auto-memory store: `./memory/` (this persona's isolated scope, wired via
456  `.claude/settings.local.json`).
457- Cross-persona facts (shell, git identity, host) belong in global memory,
458  not here.
459"#
460    )
461}
462
463fn to_tilde(path: &Path) -> String {
464    let Some(home) = env::var_os("HOME") else {
465        return display_path(path);
466    };
467    let home = PathBuf::from(home);
468    if path == home {
469        return "~".to_string();
470    }
471    if let Ok(rest) = path.strip_prefix(&home) {
472        if rest.as_os_str().is_empty() {
473            "~".to_string()
474        } else {
475            format!("~/{}", rest.to_string_lossy())
476        }
477    } else {
478        display_path(path)
479    }
480}
481
482fn shell_escape(value: &str) -> String {
483    if value.is_empty() {
484        return "''".to_string();
485    }
486
487    if value.bytes().all(is_shell_safe_byte) {
488        return value.to_string();
489    }
490
491    format!("'{}'", value.replace('\'', "'\\''"))
492}
493
494fn is_shell_safe_byte(byte: u8) -> bool {
495    byte.is_ascii_alphanumeric()
496        || matches!(
497            byte,
498            b'_' | b'@' | b'%' | b'+' | b'=' | b':' | b',' | b'.' | b'/' | b'-'
499        )
500}