use std::collections::HashMap; use std::sync::Arc; use anyhow::{Context, Result}; use forge_app::domain::Command; use forge_app::{ DirectoryReaderInfra, EnvironmentInfra, FileInfoInfra, FileReaderInfra, FileWriterInfra, }; use gray_matter::Matter; use gray_matter::engine::YAML; pub struct CommandLoaderService { infra: Arc, // Cache is used to maintain the loaded commands // for this service instance. // So that they could live till user starts a new session. cache: tokio::sync::OnceCell>, } impl CommandLoaderService { pub fn new(infra: Arc) -> Self { Self { infra, cache: Default::default() } } /// Load built-in commands that are embedded in the application binary. fn init_default(&self) -> anyhow::Result> { parse_command_iter( [( "github-pr-description", include_str!("../../../commands/github-pr-description.md"), )] .into_iter() .map(|(name, content)| (name.to_string(), content.to_string())), ) } } #[async_trait::async_trait] impl forge_app::CommandLoaderService for CommandLoaderService { async fn get_commands(&self) -> anyhow::Result> { self.cache_or_init().await } } impl CommandLoaderService { /// Load all command definitions with caching support async fn cache_or_init(&self) -> anyhow::Result> { self.cache.get_or_try_init(|| self.init()).await.cloned() } async fn init(&self) -> anyhow::Result> { // Load built-in commands first (lowest precedence) let mut commands = self.init_default()?; // Load custom commands from global directory let dir = self.infra.get_environment().command_path(); let custom_commands = self.init_command_dir(&dir).await?; commands.extend(custom_commands); // Load custom commands from CWD let dir = self.infra.get_environment().command_path_local(); let cwd_commands = self.init_command_dir(&dir).await?; commands.extend(cwd_commands); // Handle command name conflicts by keeping the last occurrence // This gives precedence order: CWD > Global Custom > Built-in Ok(resolve_command_conflicts(commands)) } async fn init_command_dir(&self, dir: &std::path::Path) -> anyhow::Result> { if !self.infra.exists(dir).await? { return Ok(vec![]); } // Use DirectoryReaderInfra to read all .md files in parallel let files = self .infra .read_directory_files(dir, Some("*.md")) .await .with_context(|| format!("Failed to read commands from: {}", dir.display()))?; parse_command_iter(files.into_iter().map(|(path, content)| { let name = path .file_stem() .and_then(|s| s.to_str()) .unwrap_or("unknown") .to_string(); (name, content) })) } } /// Implementation function for resolving command name conflicts by keeping the /// last occurrence. This implements the precedence order: CWD Custom > Global /// Custom /// > Built-in fn resolve_command_conflicts(commands: Vec) -> Vec { // Use HashMap to deduplicate by command name, keeping the last occurrence let mut command_map: HashMap = HashMap::new(); for command in commands { command_map.insert(command.name.clone(), command); } // Convert back to vector (order is not guaranteed but doesn't matter for the // service) command_map.into_values().collect() } fn parse_command_iter, Content: AsRef>( contents: I, ) -> anyhow::Result> where I: Iterator, { let mut commands = vec![]; for (name, content) in contents { let command = parse_command_file(content.as_ref()) .with_context(|| format!("Failed to parse command: {}", name.as_ref()))?; commands.push(command); } Ok(commands) } /// Parse raw content into a Command with YAML frontmatter fn parse_command_file(content: &str) -> Result { // Parse the frontmatter using gray_matter with type-safe deserialization let gray_matter = Matter::::new(); let result = gray_matter.parse::(content)?; // Extract the frontmatter let command = result .data .context("Empty command frontmatter")? .prompt(result.content); Ok(command) } #[cfg(test)] mod tests { use pretty_assertions::assert_eq; use super::*; #[tokio::test] async fn test_parse_basic_command() { let content = forge_test_kit::fixture!("src/fixtures/commands/basic.md").await; let actual = parse_command_file(&content).unwrap(); assert_eq!(actual.name.as_str(), "test-basic"); assert_eq!(actual.description.as_str(), "A basic test command"); assert_eq!( actual.prompt.as_ref().unwrap(), "This is the prompt content for the basic test command." ); } #[tokio::test] async fn test_parse_command_with_multiline_prompt() { let content = forge_test_kit::fixture!("src/fixtures/commands/multiline.md").await; let actual = parse_command_file(&content).unwrap(); assert_eq!(actual.name.as_str(), "test-multiline"); assert_eq!(actual.description.as_str(), "Command with multiline prompt"); assert!(actual.prompt.as_ref().unwrap().contains("Step 1")); assert!(actual.prompt.as_ref().unwrap().contains("Step 2")); } #[tokio::test] async fn test_parse_invalid_frontmatter() { let content = forge_test_kit::fixture!("src/fixtures/commands/invalid.md").await; let result = parse_command_file(&content); assert!(result.is_err()); } #[tokio::test] async fn test_parse_builtin_commands() { // Test that all built-in commands parse correctly let builtin_commands = [ ("fixme", "../../.forge/commands/fixme.md"), ("check", "../../.forge/commands/check.md"), ]; for (name, path) in builtin_commands { let content = forge_test_kit::fixture!(path).await; let command = parse_command_file(&content) .with_context(|| format!("Failed to parse built-in command: {}", name)) .unwrap(); assert_eq!(command.name.as_str(), name); assert!(!command.description.is_empty()); assert!(command.prompt.is_some()); } } #[test] fn test_init_default_contains_builtin_commands() { // Fixture let service = CommandLoaderService::<()> { infra: Arc::new(()), cache: Default::default() }; // Execute let actual = service.init_default().unwrap(); // Verify github-pr-description let command = actual .iter() .find(|c| c.name.as_str() == "github-pr-description") .expect("github-pr-description should be a built-in command"); assert_eq!(command.name.as_str(), "github-pr-description"); assert!(!command.description.is_empty()); assert!(command.prompt.is_some()); } #[test] fn test_resolve_command_conflicts_no_duplicates() { let fixture = vec![ Command::default().name("command1").description("Command 1"), Command::default().name("command2").description("Command 2"), Command::default().name("command3").description("Command 3"), ]; let actual = resolve_command_conflicts(fixture.clone()); // Should return all commands when no conflicts assert_eq!(actual.len(), 3); let names: std::collections::HashSet<_> = actual.iter().map(|c| c.name.as_str()).collect(); assert!(names.contains("command1")); assert!(names.contains("command2")); assert!(names.contains("command3")); } #[test] fn test_resolve_command_conflicts_with_duplicates() { let fixture = vec![ Command::default() .name("command1") .description("Global Command 1"), Command::default() .name("command2") .description("Global Command 2"), Command::default() .name("command1") .description("CWD Command 1 - Override"), // Duplicate name, should override Command::default() .name("command3") .description("CWD Command 3"), ]; let actual = resolve_command_conflicts(fixture); // Should have 3 commands: command1 (CWD version), command2 (global), command3 // (CWD) assert_eq!(actual.len(), 3); let command1 = actual .iter() .find(|c| c.name.as_str() == "command1") .expect("Should have command1"); let expected_description = "CWD Command 1 - Override"; assert_eq!(command1.description.as_str(), expected_description); } #[test] fn test_resolve_command_conflicts_multiple_duplicates() { // Test scenario: Built-in -> Global -> CWD (CWD should win) let fixture = vec![ Command::default() .name("common") .description("Built-in Common Command"), Command::default() .name("unique1") .description("Built-in Unique 1"), Command::default() .name("common") .description("Global Common Command"), // Override built-in Command::default() .name("unique2") .description("Global Unique 2"), Command::default() .name("common") .description("CWD Common Command"), // Override global Command::default() .name("unique3") .description("CWD Unique 3"), ]; let actual = resolve_command_conflicts(fixture); // Should have 4 commands: common (CWD version), unique1, unique2, unique3 assert_eq!(actual.len(), 4); let common = actual .iter() .find(|c| c.name.as_str() == "common") .expect("Should have common command"); let expected_description = "CWD Common Command"; assert_eq!(common.description.as_str(), expected_description); // Verify all unique commands are present let names: std::collections::HashSet<_> = actual.iter().map(|c| c.name.as_str()).collect(); assert!(names.contains("common")); assert!(names.contains("unique1")); assert!(names.contains("unique2")); assert!(names.contains("unique3")); } #[test] fn test_resolve_command_conflicts_empty_input() { let fixture: Vec = vec![]; let actual = resolve_command_conflicts(fixture); assert_eq!(actual.len(), 0); } }