tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 8 Packages Projects Releases Wiki Activity

feat: check tool annotation before permission check (#1957)

Browse Source
This commit is contained in:
Yingjie He
2025-04-01 17:07:07 -07:00
committed by GitHub
parent 262d7f70d9
commit 3d41768b1f
1 changed files with 31 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/truncate.rs
+31 -6
View File
@@ -235,6 +235,21 @@ impl Agent for TruncateAgent {
tools.push(list_resources_tool); tools.push(list_resources_tool);
} }
let (tools_with_readonly_annotation, tools_without_annotation): (Vec<String>, Vec<String>) =
tools.iter().fold((vec![], vec![]), |mut acc, tool| {
match &tool.annotations {
Some(annotations) => {
if annotations.read_only_hint {
acc.0.push(tool.name.clone());
}
}
None => {
acc.1.push(tool.name.clone());
}
}
acc
});
let config = capabilities.provider().get_model_config(); let config = capabilities.provider().get_model_config();
let mut system_prompt = capabilities.get_system_prompt().await; let mut system_prompt = capabilities.get_system_prompt().await;
let mut toolshim_tools = vec![]; let mut toolshim_tools = vec![];
@@ -312,30 +327,40 @@ impl Agent for TruncateAgent {
let mode = goose_mode.clone(); let mode = goose_mode.clone();
match mode.as_str() { match mode.as_str() {
"approve" | "smart_approve" => { "approve" | "smart_approve" => {
let mut read_only_tools = Vec::new();
let mut needs_confirmation = Vec::<&ToolRequest>::new(); let mut needs_confirmation = Vec::<&ToolRequest>::new();
let mut approved_tools = Vec::new(); let mut approved_tools = Vec::new();
let mut llm_detect_candidates = Vec::<&ToolRequest>::new();
let mut detected_read_only_tools = Vec::new();
// First check permissions for all tools // First check permissions for all tools
let store = ToolPermissionStore::load()?; let store = ToolPermissionStore::load()?;
for request in tool_requests.iter() { for request in tool_requests.iter() {
if let Ok(tool_call) = request.tool_call.clone() { if let Ok(tool_call) = request.tool_call.clone() {
if let Some(allowed) = store.check_permission(request) { if tools_with_readonly_annotation.contains(&tool_call.name) {
approved_tools.push((request.id.clone(), tool_call));
} else if let Some(allowed) = store.check_permission(request) {
if allowed { if allowed {
// Instead of executing immediately, collect approved tools // Instead of executing immediately, collect approved tools
approved_tools.push((request.id.clone(), tool_call)); approved_tools.push((request.id.clone(), tool_call));
} else { } else {
// If the tool doesn't have any annotation, we can use llm-as-a-judge to check permission.
if tools_without_annotation.contains(&tool_call.name) {
llm_detect_candidates.push(request);
}
needs_confirmation.push(request); needs_confirmation.push(request);
} }
} else { } else {
if tools_without_annotation.contains(&tool_call.name) {
llm_detect_candidates.push(request);
}
needs_confirmation.push(request); needs_confirmation.push(request);
} }
} }
} }
// Only check read-only status for tools needing confirmation // Only check read-only status for tools without annotation
if !needs_confirmation.is_empty() && mode == "smart_approve" { if !llm_detect_candidates.is_empty() && mode == "smart_approve" {
read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.clone()).await; detected_read_only_tools = detect_read_only_tools(&capabilities, llm_detect_candidates.clone()).await;
} }
// Handle pre-approved and read-only tools in parallel // Handle pre-approved and read-only tools in parallel
@@ -351,7 +376,7 @@ impl Agent for TruncateAgent {
for request in &needs_confirmation { for request in &needs_confirmation {
if let Ok(tool_call) = request.tool_call.clone() { if let Ok(tool_call) = request.tool_call.clone() {
// Skip confirmation if the tool_call.name is in the read_only_tools list // Skip confirmation if the tool_call.name is in the read_only_tools list
if read_only_tools.contains(&tool_call.name) { if detected_read_only_tools.contains(&tool_call.name) {
let tool_future = Self::create_tool_future(&capabilities, tool_call, request.id.clone()); let tool_future = Self::create_tool_future(&capabilities, tool_call, request.id.clone());
tool_futures.push(tool_future); tool_futures.push(tool_future);
} else { } else {