| use std::collections::{HashMap, HashSet}; |
|
|
| use forge_domain::{Agent, ToolDefinition, ToolName}; |
| use glob::Pattern; |
|
|
| |
| |
| pub struct ToolResolver { |
| all_tool_definitions: Vec<ToolDefinition>, |
| } |
|
|
| |
| fn deprecated_tool_aliases() -> HashMap<&'static str, ToolName> { |
| HashMap::from([ |
| ("search", ToolName::new("fs_search")), |
| ("Read", ToolName::new("read")), |
| ("Write", ToolName::new("write")), |
| ("Task", ToolName::new("task")), |
| ]) |
| } |
|
|
| impl ToolResolver { |
| |
| pub fn new(all_tool_definitions: Vec<ToolDefinition>) -> Self { |
| Self { all_tool_definitions } |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| pub fn resolve<'a>(&'a self, agent: &Agent) -> Vec<&'a ToolDefinition> { |
| let patterns = Self::build_patterns(agent); |
| let mut resolved = self.match_tools(&patterns); |
| self.dedupe_tools(&mut resolved); |
| agent.tool_order().sort_refs(&mut resolved); |
| resolved |
| } |
|
|
| fn is_allowed_pattern(patterns: &[Pattern], tool_name: &ToolName) -> bool { |
| patterns |
| .iter() |
| .any(|pattern| pattern.matches(tool_name.as_str())) |
| } |
|
|
| pub fn is_allowed(agent: &Agent, tool_name: &ToolName) -> bool { |
| let aliases = deprecated_tool_aliases(); |
| let normalized_tool_name = aliases.get(tool_name.as_str()).unwrap_or(tool_name); |
| let legacy_mcp_tool_name = normalized_tool_name.to_legacy_mcp_name(); |
| let patterns = Self::build_patterns(agent); |
|
|
| Self::is_allowed_pattern(&patterns, normalized_tool_name) |
| || legacy_mcp_tool_name |
| .as_ref() |
| .is_some_and(|legacy_tool_name| { |
| Self::is_allowed_pattern(&patterns, legacy_tool_name) |
| }) |
| } |
|
|
| |
| |
| |
| fn build_patterns(agent: &Agent) -> Vec<Pattern> { |
| let aliases = deprecated_tool_aliases(); |
| let tool_names = agent |
| .tools |
| .iter() |
| .flatten() |
| .map(|name| { |
| |
| aliases.get(name.as_str()).unwrap_or(name) |
| }) |
| .collect::<HashSet<_>>(); |
|
|
| tool_names |
| .into_iter() |
| .filter_map(|pattern| Pattern::new(pattern.as_str()).ok()) |
| .collect() |
| } |
|
|
| |
| fn match_tools<'a>(&'a self, patterns: &[Pattern]) -> Vec<&'a ToolDefinition> { |
| self.all_tool_definitions |
| .iter() |
| .filter(|tool| Self::is_allowed_pattern(patterns, &tool.name)) |
| .collect() |
| } |
|
|
| |
| fn dedupe_tools(&self, resolved: &mut Vec<&ToolDefinition>) { |
| let mut seen = HashSet::new(); |
| resolved.retain(|tool| seen.insert(&tool.name)); |
| } |
| } |
|
|
| #[cfg(test)] |
| mod tests { |
| use forge_domain::{Agent, AgentId, ModelId, ProviderId, ToolDefinition, ToolName}; |
| use pretty_assertions::assert_eq; |
|
|
| use super::ToolResolver; |
|
|
| #[test] |
| fn test_resolve_filters_agent_tools() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("write").description("Write Tool"), |
| ToolDefinition::new("fs_search").description("Search Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("read"), ToolName::new("fs_search")]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| |
| let expected = vec![ |
| &tool_resolver.all_tool_definitions[0], |
| &tool_resolver.all_tool_definitions[2], |
| ]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_no_agent_tools() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("write").description("Write Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| let expected: Vec<&ToolDefinition> = vec![]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_nonexistent_tools() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("write").description("Write Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ |
| ToolName::new("nonexistent1"), |
| ToolName::new("nonexistent2"), |
| ]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| let expected: Vec<&ToolDefinition> = vec![]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_duplicate_agent_tools() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("write").description("Write Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ |
| ToolName::new("read"), |
| ToolName::new("read"), |
| ToolName::new("write"), |
| ]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| let expected = vec![ |
| &tool_resolver.all_tool_definitions[0], |
| &tool_resolver.all_tool_definitions[1], |
| ]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_glob_pattern_wildcard() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("fs_read").description("Read Tool"), |
| ToolDefinition::new("fs_write").description("Write Tool"), |
| ToolDefinition::new("fs_search").description("Search Tool"), |
| ToolDefinition::new("net_fetch").description("Fetch Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("fs_*")]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| let expected = vec![ |
| &tool_resolver.all_tool_definitions[0], |
| &tool_resolver.all_tool_definitions[2], |
| &tool_resolver.all_tool_definitions[1], |
| ]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_glob_pattern_no_matches() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("write").description("Write Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("fs_*")]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| let expected: Vec<&ToolDefinition> = vec![]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_mixed_exact_and_glob() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("fs_read").description("FS Read Tool"), |
| ToolDefinition::new("fs_write").description("FS Write Tool"), |
| ToolDefinition::new("net_fetch").description("Net Fetch Tool"), |
| ToolDefinition::new("shell").description("Shell Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("fs_*"), ToolName::new("shell")]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| let expected = vec![ |
| &tool_resolver.all_tool_definitions[0], |
| &tool_resolver.all_tool_definitions[1], |
| &tool_resolver.all_tool_definitions[3], |
| ]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_question_mark_wildcard() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read1").description("Read 1 Tool"), |
| ToolDefinition::new("read2").description("Read 2 Tool"), |
| ToolDefinition::new("read10").description("Read 10 Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("read?")]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| let expected = vec![ |
| &tool_resolver.all_tool_definitions[0], |
| &tool_resolver.all_tool_definitions[1], |
| ]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_resolve_with_overlapping_glob_patterns() { |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("fs_read").description("FS Read Tool"), |
| ToolDefinition::new("fs_write").description("FS Write Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ |
| ToolName::new("fs_*"), |
| ToolName::new("fs_read"), |
| ToolName::new("*_read"), |
| ]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| |
| |
| |
| let expected = vec![ |
| &tool_resolver.all_tool_definitions[1], |
| &tool_resolver.all_tool_definitions[0], |
| ]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_exact_legacy_mcp_tool_allows_claude_code_name() { |
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("mcp_github_tool_create_issue")]); |
|
|
| assert!(ToolResolver::is_allowed( |
| &fixture, |
| &ToolName::new("mcp__github__create_issue"), |
| )); |
| } |
|
|
| #[test] |
| fn test_glob_legacy_mcp_tool_allows_claude_code_name() { |
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("mcp_github_tool_*")]); |
|
|
| assert!(ToolResolver::is_allowed( |
| &fixture, |
| &ToolName::new("mcp__github__create_issue"), |
| )); |
| } |
|
|
| #[test] |
| fn test_backward_compatibility_search_alias() { |
| |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("fs_search").description("Search Tool"), |
| ]; |
|
|
| let tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| |
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("read"), ToolName::new("search")]); |
|
|
| let actual = tool_resolver.resolve(&fixture); |
| |
| |
| let expected = vec![ |
| &tool_resolver.all_tool_definitions[0], |
| &tool_resolver.all_tool_definitions[1], |
| ]; |
|
|
| assert_eq!(actual, expected); |
| } |
|
|
| #[test] |
| fn test_capitalized_read_alias() { |
| |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("write").description("Write Tool"), |
| ]; |
|
|
| let _tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| |
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("read"), ToolName::new("write")]); |
|
|
| |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("read"))); |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("Read"))); |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("write"))); |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("Write"))); |
| } |
|
|
| #[test] |
| fn test_capitalized_write_alias() { |
| |
| let all_tool_definitions = vec![ |
| ToolDefinition::new("read").description("Read Tool"), |
| ToolDefinition::new("write").description("Write Tool"), |
| ]; |
|
|
| let _tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("write")]); |
|
|
| |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("write"))); |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("Write"))); |
| } |
|
|
| #[test] |
| fn test_capitalized_task_alias() { |
| |
| let all_tool_definitions = vec![ToolDefinition::new("task").description("Task Tool")]; |
|
|
| let _tool_resolver = ToolResolver::new(all_tool_definitions); |
|
|
| let fixture = Agent::new( |
| AgentId::new("test-agent"), |
| ProviderId::ANTHROPIC, |
| ModelId::new("claude-3-5-sonnet-20241022"), |
| ) |
| .tools(vec![ToolName::new("task")]); |
|
|
| |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("task"))); |
| assert!(ToolResolver::is_allowed(&fixture, &ToolName::new("Task"))); |
| } |
| } |
|
|