fix: type allowed tools validation

This commit is contained in:
bellman
2026-06-04 12:01:58 +09:00
parent 22fdaeae2c
commit ecd3e4ceb9
6 changed files with 400 additions and 57 deletions

View File

@@ -201,30 +201,20 @@ impl GlobalToolRegistry {
return Ok(None);
}
let builtin_specs = mvp_tool_specs();
let canonical_names = builtin_specs
.iter()
.map(|spec| spec.name.to_string())
.chain(
self.plugin_tools
.iter()
.map(|tool| tool.definition().name.clone()),
)
.chain(self.runtime_tools.iter().map(|tool| tool.name.clone()))
.collect::<Vec<_>>();
let mut name_map = canonical_names
.iter()
.map(|name| (normalize_tool_name(name), name.clone()))
.collect::<BTreeMap<_, _>>();
let actual_names = self.actual_tool_names();
let canonical_names = self.canonical_allowed_tool_names();
let canonical_name_set = canonical_names.iter().cloned().collect::<BTreeSet<_>>();
let mut name_map = BTreeMap::new();
for actual in &actual_names {
let canonical = canonical_allowed_tool_name(actual);
name_map.insert(allowed_tool_lookup_key(actual), canonical.clone());
name_map.insert(allowed_tool_lookup_key(&canonical), canonical);
}
for (alias, canonical) in [
("read", "read_file"),
("write", "write_file"),
("edit", "edit_file"),
("glob", "glob_search"),
("grep", "grep_search"),
] {
name_map.insert(alias.to_string(), canonical.to_string());
for (alias, canonical) in self.allowed_tool_aliases() {
if canonical_name_set.contains(&canonical) {
name_map.insert(allowed_tool_lookup_key(&alias), canonical);
}
}
let mut allowed = BTreeSet::new();
@@ -233,11 +223,11 @@ impl GlobalToolRegistry {
.split(|ch: char| ch == ',' || ch.is_whitespace())
.filter(|token| !token.is_empty())
{
let normalized = normalize_tool_name(token);
let canonical = name_map.get(&normalized).ok_or_else(|| {
let canonical = name_map.get(&allowed_tool_lookup_key(token)).ok_or_else(|| {
format!(
"unsupported tool in --allowedTools: {token} (expected one of: {})",
canonical_names.join(", ")
"invalid_tool_name: unsupported tool in --allowedTools: {token}\nAvailable: {}\nAliases: {}\nHint: Use canonical snake_case tool names from Available or aliases from Aliases.",
canonical_names.join(", "),
format_allowed_tool_aliases(&self.allowed_tool_aliases())
)
})?;
allowed.insert(canonical.clone());
@@ -258,7 +248,10 @@ impl GlobalToolRegistry {
pub fn definitions(&self, allowed_tools: Option<&BTreeSet<String>>) -> Vec<ToolDefinition> {
let builtin = mvp_tool_specs()
.into_iter()
.filter(|spec| allowed_tools.is_none_or(|allowed| allowed.contains(spec.name)))
.filter(|spec| {
allowed_tools
.is_none_or(|allowed| allowed.contains(&canonical_allowed_tool_name(spec.name)))
})
.map(|spec| ToolDefinition {
name: spec.name.to_string(),
description: Some(spec.description.to_string()),
@@ -267,7 +260,11 @@ impl GlobalToolRegistry {
let runtime = self
.runtime_tools
.iter()
.filter(|tool| allowed_tools.is_none_or(|allowed| allowed.contains(tool.name.as_str())))
.filter(|tool| {
allowed_tools.is_none_or(|allowed| {
allowed.contains(&canonical_allowed_tool_name(&tool.name))
})
})
.map(|tool| ToolDefinition {
name: tool.name.clone(),
description: tool.description.clone(),
@@ -277,8 +274,11 @@ impl GlobalToolRegistry {
.plugin_tools
.iter()
.filter(|tool| {
allowed_tools
.is_none_or(|allowed| allowed.contains(tool.definition().name.as_str()))
allowed_tools.is_none_or(|allowed| {
allowed.contains(&canonical_allowed_tool_name(
tool.definition().name.as_str(),
))
})
})
.map(|tool| ToolDefinition {
name: tool.definition().name.clone(),
@@ -294,19 +294,29 @@ impl GlobalToolRegistry {
) -> Result<Vec<(String, PermissionMode)>, String> {
let builtin = mvp_tool_specs()
.into_iter()
.filter(|spec| allowed_tools.is_none_or(|allowed| allowed.contains(spec.name)))
.filter(|spec| {
allowed_tools
.is_none_or(|allowed| allowed.contains(&canonical_allowed_tool_name(spec.name)))
})
.map(|spec| (spec.name.to_string(), spec.required_permission));
let runtime = self
.runtime_tools
.iter()
.filter(|tool| allowed_tools.is_none_or(|allowed| allowed.contains(tool.name.as_str())))
.filter(|tool| {
allowed_tools.is_none_or(|allowed| {
allowed.contains(&canonical_allowed_tool_name(&tool.name))
})
})
.map(|tool| (tool.name.clone(), tool.required_permission));
let plugin = self
.plugin_tools
.iter()
.filter(|tool| {
allowed_tools
.is_none_or(|allowed| allowed.contains(tool.definition().name.as_str()))
allowed_tools.is_none_or(|allowed| {
allowed.contains(&canonical_allowed_tool_name(
tool.definition().name.as_str(),
))
})
})
.map(|tool| {
permission_mode_from_plugin(tool.required_permission())
@@ -316,6 +326,52 @@ impl GlobalToolRegistry {
Ok(builtin.chain(runtime).chain(plugin).collect())
}
#[must_use]
pub fn actual_tool_names(&self) -> Vec<String> {
mvp_tool_specs()
.iter()
.map(|spec| spec.name.to_string())
.chain(
self.plugin_tools
.iter()
.map(|tool| tool.definition().name.clone()),
)
.chain(self.runtime_tools.iter().map(|tool| tool.name.clone()))
.collect()
}
#[must_use]
pub fn canonical_allowed_tool_names(&self) -> Vec<String> {
self.actual_tool_names()
.into_iter()
.map(|name| canonical_allowed_tool_name(&name))
.collect::<BTreeSet<_>>()
.into_iter()
.collect()
}
#[must_use]
pub fn allowed_tool_aliases(&self) -> BTreeMap<String, String> {
let mut aliases = BTreeMap::from([
("read".to_string(), "read_file".to_string()),
("Read".to_string(), "read_file".to_string()),
("write".to_string(), "write_file".to_string()),
("Write".to_string(), "write_file".to_string()),
("edit".to_string(), "edit_file".to_string()),
("Edit".to_string(), "edit_file".to_string()),
("glob".to_string(), "glob_search".to_string()),
("Glob".to_string(), "glob_search".to_string()),
("grep".to_string(), "grep_search".to_string()),
("Grep".to_string(), "grep_search".to_string()),
]);
for actual in self.actual_tool_names() {
let canonical = canonical_allowed_tool_name(&actual);
if actual != canonical {
aliases.insert(actual, canonical);
}
}
aliases
}
#[must_use]
pub fn has_runtime_tool(&self, name: &str) -> bool {
self.runtime_tools.iter().any(|tool| tool.name == name)
@@ -378,8 +434,40 @@ impl GlobalToolRegistry {
}
}
fn normalize_tool_name(value: &str) -> String {
value.trim().replace('-', "_").to_ascii_lowercase()
pub fn canonical_allowed_tool_name(value: &str) -> String {
let trimmed = value.trim().replace('-', "_");
let mut output = String::new();
let chars = trimmed.chars().collect::<Vec<_>>();
for (index, ch) in chars.iter().copied().enumerate() {
if ch == '_' || ch.is_whitespace() {
output.push('_');
continue;
}
let previous = index.checked_sub(1).and_then(|i| chars.get(i)).copied();
let next = chars.get(index + 1).copied();
if ch.is_ascii_uppercase()
&& index > 0
&& !output.ends_with('_')
&& (previous.is_some_and(|p| p.is_ascii_lowercase() || p.is_ascii_digit())
|| next.is_some_and(|n| n.is_ascii_lowercase()))
{
output.push('_');
}
output.push(ch.to_ascii_lowercase());
}
output.trim_matches('_').to_string()
}
fn allowed_tool_lookup_key(value: &str) -> String {
canonical_allowed_tool_name(value).replace('_', "")
}
fn format_allowed_tool_aliases(aliases: &BTreeMap<String, String>) -> String {
aliases
.iter()
.map(|(alias, canonical)| format!("{alias}={canonical}"))
.collect::<Vec<_>>()
.join(", ")
}
fn permission_mode_from_plugin(value: &str) -> Result<PermissionMode, String> {
@@ -4210,7 +4298,7 @@ fn allowed_tools_for_subagent(subagent_type: &str) -> BTreeSet<String> {
"PowerShell",
],
};
tools.into_iter().map(str::to_string).collect()
tools.into_iter().map(canonical_allowed_tool_name).collect()
}
fn agent_permission_policy() -> PermissionPolicy {
@@ -5238,7 +5326,10 @@ impl SubagentToolExecutor {
impl ToolExecutor for SubagentToolExecutor {
fn execute(&mut self, tool_name: &str, input: &str) -> Result<String, ToolError> {
if !self.allowed_tools.contains(tool_name) {
if !self
.allowed_tools
.contains(&canonical_allowed_tool_name(tool_name))
{
return Err(ToolError::new(format!(
"tool `{tool_name}` is not enabled for this sub-agent"
)));
@@ -5253,7 +5344,10 @@ impl ToolExecutor for SubagentToolExecutor {
fn tool_specs_for_allowed_tools(allowed_tools: Option<&BTreeSet<String>>) -> Vec<ToolSpec> {
mvp_tool_specs()
.into_iter()
.filter(|spec| allowed_tools.is_none_or(|allowed| allowed.contains(spec.name)))
.filter(|spec| {
allowed_tools
.is_none_or(|allowed| allowed.contains(&canonical_allowed_tool_name(spec.name)))
})
.collect()
}
@@ -7603,6 +7697,29 @@ mod tests {
}
}
#[test]
fn allowed_tools_normalize_to_canonical_snake_case_and_aliases_432() {
let registry = GlobalToolRegistry::builtin();
let allowed = registry
.normalize_allowed_tools(&["Read,WebFetch,MCP".to_string()])
.expect("aliases and legacy names should normalize")
.expect("allow-list should be populated");
assert!(allowed.contains("read_file"));
assert!(allowed.contains("web_fetch"));
assert!(allowed.contains("mcp"));
assert!(!allowed.contains("Read"));
assert!(!allowed.contains("WebFetch"));
let canonical = registry.canonical_allowed_tool_names();
assert!(canonical.contains(&"web_fetch".to_string()));
assert!(canonical.contains(&"todo_write".to_string()));
assert!(!canonical.contains(&"WebFetch".to_string()));
assert_eq!(
registry.allowed_tool_aliases().get("WebFetch"),
Some(&"web_fetch".to_string())
);
}
#[test]
fn runtime_tools_extend_registry_definitions_permissions_and_search() {
let registry = GlobalToolRegistry::builtin()
@@ -8584,7 +8701,7 @@ mod tests {
.expect("spawn job should be captured");
assert_eq!(captured_job.prompt, "Check tests and outstanding work.");
assert!(captured_job.allowed_tools.contains("read_file"));
assert!(!captured_job.allowed_tools.contains("Agent"));
assert!(!captured_job.allowed_tools.contains("agent"));
let normalized = execute_tool(
"Agent",
@@ -9184,7 +9301,7 @@ mod tests {
let general = allowed_tools_for_subagent("general-purpose");
assert!(general.contains("bash"));
assert!(general.contains("write_file"));
assert!(!general.contains("Agent"));
assert!(!general.contains("agent"));
let explore = allowed_tools_for_subagent("Explore");
assert!(explore.contains("read_file"));
@@ -9192,13 +9309,13 @@ mod tests {
assert!(!explore.contains("bash"));
let plan = allowed_tools_for_subagent("Plan");
assert!(plan.contains("TodoWrite"));
assert!(plan.contains("StructuredOutput"));
assert!(!plan.contains("Agent"));
assert!(plan.contains("todo_write"));
assert!(plan.contains("structured_output"));
assert!(!plan.contains("agent"));
let verification = allowed_tools_for_subagent("Verification");
assert!(verification.contains("bash"));
assert!(verification.contains("PowerShell"));
assert!(verification.contains("power_shell"));
assert!(!verification.contains("write_file"));
}