use std::collections::HashMap;
use std::path::{Path, PathBuf};
#[derive(Debug, Clone)]
pub struct CustomCommand {
pub name: String,
pub description: String,
pub args_requirement: ArgsRequirement,
pub template: String,
pub source: PathBuf,
pub namespace: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ArgsRequirement {
Required,
Optional,
None,
}
pub struct CustomCommandRegistry {
commands: HashMap<String, CustomCommand>,
}
impl CustomCommandRegistry {
pub fn load(project_root: &Path) -> Self {
let config_dir = crate::config::Config::config_dir();
let mut commands = HashMap::new();
Self::load_from_dir(&config_dir.join("commands"), None, &mut commands);
Self::load_from_dir(&project_root.join(".atomcode/commands"), None, &mut commands);
for assets in crate::plugin::loader::iter_installed_plugin_assets() {
Self::load_from_dir(&assets.commands_dir(), Some(&assets.plugin), &mut commands);
}
Self { commands }
}
pub fn empty() -> Self {
Self {
commands: HashMap::new(),
}
}
fn load_from_dir(
dir: &Path,
namespace: Option<&str>,
commands: &mut HashMap<String, CustomCommand>,
) {
let entries = match std::fs::read_dir(dir) {
Ok(e) => e,
Err(_) => return,
};
for entry in entries.flatten() {
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("md") {
continue;
}
if let Some(mut cmd) = Self::parse_command_file(&path) {
if let Some(ns) = namespace {
cmd.namespace = Some(ns.to_string());
}
let key = match &cmd.namespace {
Some(ns) => format!("{}:{}", ns, cmd.name),
None => cmd.name.clone(),
};
commands.insert(key, cmd);
}
}
}
fn parse_command_file(path: &Path) -> Option<CustomCommand> {
let content = std::fs::read_to_string(path).ok()?;
let (frontmatter, template) = Self::split_frontmatter(&content)?;
let name = Self::extract_field(&frontmatter, "name")?;
let description = Self::extract_field(&frontmatter, "description")
.unwrap_or_else(|| format!("Custom command: {}", name));
let args_str = Self::extract_field(&frontmatter, "args").unwrap_or_else(|| "none".into());
let args_requirement = match args_str.as_str() {
"required" => ArgsRequirement::Required,
"optional" => ArgsRequirement::Optional,
_ => ArgsRequirement::None,
};
Some(CustomCommand {
name,
description,
args_requirement,
template: template.trim().to_string(),
source: path.to_path_buf(),
namespace: None,
})
}
fn split_frontmatter(content: &str) -> Option<(String, String)> {
let content = content.trim();
if !content.starts_with("---") {
return None;
}
let rest = &content[3..];
let end = rest.find("---")?;
Some((rest[..end].trim().to_string(), rest[end + 3..].to_string()))
}
fn extract_field(frontmatter: &str, key: &str) -> Option<String> {
for line in frontmatter.lines() {
let line = line.trim();
if let Some(rest) = line.strip_prefix(key) {
if let Some(value) = rest.trim_start().strip_prefix(':') {
return Some(value.trim().to_string());
}
}
}
None
}
pub fn get(&self, name: &str) -> Option<&CustomCommand> {
self.commands.get(name)
}
pub fn render(&self, name: &str, args: &str) -> Option<String> {
self.commands.get(name).map(|cmd| {
cmd.template
.replace("$ARGUMENTS", args)
.replace("${ARGUMENTS}", args)
})
}
pub fn list(&self) -> Vec<&CustomCommand> {
let mut cmds: Vec<_> = self.commands.values().collect();
cmds.sort_by_key(|c| &c.name);
cmds
}
pub fn command_names_and_descriptions(&self) -> Vec<(String, String)> {
self.list()
.iter()
.map(|c| (c.name.clone(), c.description.clone()))
.collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn split_frontmatter_works() {
let content = "---\nname: review\ndescription: Code review\n---\nTemplate body here";
let (fm, body) = CustomCommandRegistry::split_frontmatter(content).unwrap();
assert!(fm.contains("name: review"));
assert!(fm.contains("description: Code review"));
assert_eq!(body.trim(), "Template body here");
}
#[test]
fn split_frontmatter_returns_none_without_delimiters() {
assert!(CustomCommandRegistry::split_frontmatter("no frontmatter here").is_none());
assert!(CustomCommandRegistry::split_frontmatter("--- only opening").is_none());
}
#[test]
fn extract_field_works() {
let fm = "name: review\ndescription: Code review\nargs: optional";
assert_eq!(
CustomCommandRegistry::extract_field(fm, "name"),
Some("review".into())
);
assert_eq!(
CustomCommandRegistry::extract_field(fm, "description"),
Some("Code review".into())
);
assert_eq!(
CustomCommandRegistry::extract_field(fm, "args"),
Some("optional".into())
);
assert_eq!(CustomCommandRegistry::extract_field(fm, "missing"), None);
}
#[test]
fn parse_command_file_works() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("review.md");
std::fs::write(
&path,
"---\nname: review\ndescription: Code review\nargs: optional\n---\nReview: $ARGUMENTS",
)
.unwrap();
let cmd = CustomCommandRegistry::parse_command_file(&path).unwrap();
assert_eq!(cmd.name, "review");
assert_eq!(cmd.description, "Code review");
assert_eq!(cmd.args_requirement, ArgsRequirement::Optional);
assert_eq!(cmd.template, "Review: $ARGUMENTS");
assert_eq!(cmd.source, path);
}
#[test]
fn render_replaces_arguments() {
let dir = tempfile::tempdir().unwrap();
let cmd_dir = dir.path().join(".atomcode/commands");
std::fs::create_dir_all(&cmd_dir).unwrap();
std::fs::write(
cmd_dir.join("review.md"),
"---\nname: review\ndescription: Review\nargs: optional\n---\nReview $ARGUMENTS and ${ARGUMENTS} done",
)
.unwrap();
let reg = CustomCommandRegistry::load(dir.path());
let rendered = reg.render("review", "main.rs").unwrap();
assert_eq!(rendered, "Review main.rs and main.rs done");
}
#[test]
fn render_empty_args() {
let dir = tempfile::tempdir().unwrap();
let cmd_dir = dir.path().join(".atomcode/commands");
std::fs::create_dir_all(&cmd_dir).unwrap();
std::fs::write(
cmd_dir.join("test.md"),
"---\nname: test\ndescription: Run tests\n---\nRun all tests. Focus: $ARGUMENTS",
)
.unwrap();
let reg = CustomCommandRegistry::load(dir.path());
let rendered = reg.render("test", "").unwrap();
assert_eq!(rendered, "Run all tests. Focus: ");
}
#[test]
fn load_from_dir_skips_non_md() {
let dir = tempfile::tempdir().unwrap();
let cmd_dir = dir.path().join("commands");
std::fs::create_dir_all(&cmd_dir).unwrap();
std::fs::write(
cmd_dir.join("valid.md"),
"---\nname: valid\ndescription: Valid cmd\n---\nTemplate",
)
.unwrap();
std::fs::write(
cmd_dir.join("skip.txt"),
"---\nname: skip\ndescription: Skip\n---\nNope",
)
.unwrap();
std::fs::write(
cmd_dir.join("noext"),
"---\nname: noext\ndescription: No ext\n---\nNope",
)
.unwrap();
let mut commands = HashMap::new();
CustomCommandRegistry::load_from_dir(&cmd_dir, None, &mut commands);
assert_eq!(commands.len(), 1);
assert!(commands.contains_key("valid"));
}
#[test]
fn project_overrides_global_same_name() {
let root = tempfile::tempdir().unwrap();
let global_dir = root.path().join("global_commands");
std::fs::create_dir_all(&global_dir).unwrap();
std::fs::write(
global_dir.join("review.md"),
"---\nname: review\ndescription: Global review\n---\nGlobal template",
)
.unwrap();
let project_dir = root.path().join("project_commands");
std::fs::create_dir_all(&project_dir).unwrap();
std::fs::write(
project_dir.join("review.md"),
"---\nname: review\ndescription: Project review\n---\nProject template",
)
.unwrap();
let mut commands = HashMap::new();
CustomCommandRegistry::load_from_dir(&global_dir, None, &mut commands);
CustomCommandRegistry::load_from_dir(&project_dir, None, &mut commands);
let cmd = commands.get("review").unwrap();
assert_eq!(cmd.description, "Project review");
assert_eq!(cmd.template, "Project template");
}
#[test]
fn list_returns_sorted_commands() {
let dir = tempfile::tempdir().unwrap();
let cmd_dir = dir.path().join(".atomcode/commands");
std::fs::create_dir_all(&cmd_dir).unwrap();
std::fs::write(
cmd_dir.join("zebra.md"),
"---\nname: zebra\ndescription: Z\n---\nZ",
)
.unwrap();
std::fs::write(
cmd_dir.join("alpha.md"),
"---\nname: alpha\ndescription: A\n---\nA",
)
.unwrap();
let reg = CustomCommandRegistry::load(dir.path());
let names: Vec<_> = reg.list().iter().map(|c| c.name.as_str()).collect();
assert_eq!(names, vec!["alpha", "zebra"]);
}
#[test]
fn empty_registry_has_no_commands() {
let reg = CustomCommandRegistry::empty();
assert!(reg.list().is_empty());
assert!(reg.get("anything").is_none());
assert!(reg.render("anything", "").is_none());
}
#[test]
fn parse_file_without_frontmatter_returns_none() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("bad.md");
std::fs::write(&path, "No frontmatter here, just text.").unwrap();
assert!(CustomCommandRegistry::parse_command_file(&path).is_none());
}
#[test]
fn parse_file_without_name_returns_none() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("noname.md");
std::fs::write(
&path,
"---\ndescription: Missing name field\n---\nTemplate",
)
.unwrap();
assert!(CustomCommandRegistry::parse_command_file(&path).is_none());
}
#[test]
fn default_description_when_missing() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("nodesc.md");
std::fs::write(&path, "---\nname: nodesc\n---\nTemplate").unwrap();
let cmd = CustomCommandRegistry::parse_command_file(&path).unwrap();
assert_eq!(cmd.description, "Custom command: nodesc");
}
#[test]
fn default_args_requirement_is_none() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("noargs.md");
std::fs::write(&path, "---\nname: noargs\n---\nTemplate").unwrap();
let cmd = CustomCommandRegistry::parse_command_file(&path).unwrap();
assert_eq!(cmd.args_requirement, ArgsRequirement::None);
}
#[test]
#[serial_test::serial]
fn load_plugin_layer_namespaces_commands() {
let tmp = tempfile::tempdir().unwrap();
std::env::set_var("ATOMCODE_HOME", tmp.path());
let plugin_dir = tmp.path().join("plugins/marketplaces/p");
let cmd_dir = plugin_dir.join("commands");
std::fs::create_dir_all(&cmd_dir).unwrap();
std::fs::write(
cmd_dir.join("greet.md"),
"---\nname: greet\ndescription: hi\n---\nhello $ARGUMENTS",
)
.unwrap();
std::fs::write(
tmp.path().join("plugins/installed_plugins.json"),
r#"{"version":1,"plugins":{"p@p":{"marketplace":"p","plugin":"p","plugin_dir":"marketplaces/p","installed_at":"x"}}}"#,
)
.unwrap();
let working = tempfile::tempdir().unwrap();
let reg = CustomCommandRegistry::load(working.path());
assert!(reg.get("p:greet").is_some());
std::env::remove_var("ATOMCODE_HOME");
}
}