Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion crates/goose-test-support/src/mcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ impl McpFixtureServer {
}
}

#[tool(description = "Get the code")]
#[tool(description = "Get the code", annotations(read_only_hint = true))]
fn get_code(&self) -> Result<CallToolResult, McpError> {
Ok(CallToolResult::success(vec![Content::text(FAKE_CODE)]))
}
Expand Down
14 changes: 11 additions & 3 deletions crates/goose/src/agents/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -245,14 +245,18 @@ impl Agent {
tool_result_tx: tool_tx,
tool_result_rx: Arc::new(Mutex::new(tool_rx)),
retry_manager: RetryManager::new(),
tool_inspection_manager: Self::create_tool_inspection_manager(permission_manager),
tool_inspection_manager: Self::create_tool_inspection_manager(
permission_manager,
provider.clone(),
),
container: Mutex::new(None),
}
}

/// Create a tool inspection manager with default inspectors
fn create_tool_inspection_manager(
permission_manager: Arc<PermissionManager>,
provider: SharedProvider,
) -> ToolInspectionManager {
let mut tool_inspection_manager = ToolInspectionManager::new();

Expand All @@ -261,9 +265,8 @@ impl Agent {

// Add permission inspector (medium-high priority)
tool_inspection_manager.add_inspector(Box::new(PermissionInspector::new(
std::collections::HashSet::new(), // readonly tools - will be populated from extension manager
std::collections::HashSet::new(), // regular tools - will be populated from extension manager
permission_manager,
provider,
)));

// Add repetition inspector (lower priority - basic repetition checking)
Expand Down Expand Up @@ -350,6 +353,10 @@ impl Agent {
.prepare_tools_and_prompt(session_id, working_dir)
.await?;

if self.config.goose_mode == GooseMode::SmartApprove {
self.tool_inspection_manager.apply_tool_annotations(&tools);
}

Ok(ReplyContext {
conversation,
tools,
Expand Down Expand Up @@ -1261,6 +1268,7 @@ impl Agent {
// Run all tool inspectors
let inspection_results = self.tool_inspection_manager
.inspect_tools(
&session_config.id,
&remaining_requests,
conversation.messages(),
goose_mode,
Expand Down
73 changes: 73 additions & 0 deletions crates/goose/src/config/permission.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::config::paths::Paths;
use rmcp::model::Tool;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
Expand Down Expand Up @@ -97,6 +98,51 @@ impl PermissionManager {
self.config_path.as_path()
}

pub fn apply_tool_annotations(&self, tools: &[Tool]) {
let mut write_annotated = Vec::new();
for tool in tools {
let Some(anns) = &tool.annotations else {
continue;
};
if anns.read_only_hint == Some(false) {
write_annotated.push(tool.name.to_string());
}
}
if !write_annotated.is_empty() {
self.bulk_update_smart_approve_permissions(
&write_annotated,
PermissionLevel::AskBefore,
);
}
}

fn bulk_update_smart_approve_permissions(&self, tool_names: &[String], level: PermissionLevel) {
let mut map = self.permission_map.write().unwrap();
let permission_config = map.entry(SMART_APPROVE_PERMISSION.to_string()).or_default();

for tool_name in tool_names {
// Remove from all lists to avoid duplicates
permission_config.always_allow.retain(|p| p != tool_name);
permission_config.ask_before.retain(|p| p != tool_name);
permission_config.never_allow.retain(|p| p != tool_name);

// Add to the appropriate list
match &level {
PermissionLevel::AlwaysAllow => {
permission_config.always_allow.push(tool_name.clone())
}
PermissionLevel::AskBefore => permission_config.ask_before.push(tool_name.clone()),
PermissionLevel::NeverAllow => {
permission_config.never_allow.push(tool_name.clone())
}
Comment thread
codefromthecrypt marked this conversation as resolved.
}
}

let yaml_content =
serde_yaml::to_string(&*map).expect("Failed to serialize permission config");
fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml");
}

/// Helper function to retrieve the permission level for a specific permission category and tool.
fn get_permission(&self, name: &str, principal_name: &str) -> Option<PermissionLevel> {
let map = self.permission_map.read().unwrap();
Expand Down Expand Up @@ -191,6 +237,8 @@ impl PermissionManager {
#[cfg(test)]
mod tests {
use super::*;
use rmcp::model::ToolAnnotations;
use rmcp::object;
use tempfile::TempDir;

// Helper function to create a test instance of PermissionManager with a temp dir
Expand Down Expand Up @@ -313,4 +361,29 @@ mod tests {
fs::write(&permission_path, "{{invalid yaml: [broken").unwrap();
PermissionManager::new(temp_dir.path().to_path_buf());
}

use test_case::test_case;

#[test_case(
vec![Tool::new("tool".to_string(), String::new(), object!({"type": "object"}))
.annotate(ToolAnnotations::new().read_only(false))],
Some(PermissionLevel::AskBefore);
"write_annotation_caches_ask"
)]
#[test_case(
vec![Tool::new("tool".to_string(), String::new(), object!({"type": "object"}))],
None;
"unannotated_left_uncached"
)]
#[test_case(
vec![Tool::new("tool".to_string(), String::new(), object!({"type": "object"}))
.annotate(ToolAnnotations::new().read_only(true))],
None;
"readonly_annotation_skipped"
)]
fn test_apply_tool_annotations(tools: Vec<Tool>, expect_cache: Option<PermissionLevel>) {
let (manager, _temp_dir) = create_test_permission_manager();
manager.apply_tool_annotations(&tools);
assert_eq!(manager.get_smart_approve_permission("tool"), expect_cache);
}
}
1 change: 0 additions & 1 deletion crates/goose/src/permission/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,5 +5,4 @@ pub mod permission_store;

pub use permission_confirmation::{Permission, PermissionConfirmation};
pub use permission_inspector::PermissionInspector;
pub use permission_judge::detect_read_only_tools;
pub use permission_store::ToolPermissionStore;
168 changes: 144 additions & 24 deletions crates/goose/src/permission/permission_inspector.rs
Original file line number Diff line number Diff line change
@@ -1,34 +1,52 @@
use crate::agents::platform_extensions::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE;
use crate::agents::types::SharedProvider;
use crate::config::permission::PermissionLevel;
use crate::config::{GooseMode, PermissionManager};
use crate::conversation::message::{Message, ToolRequest};
use crate::permission::permission_judge::PermissionCheckResult;
use crate::permission::permission_judge::{detect_read_only_tools, PermissionCheckResult};
use crate::tool_inspection::{InspectionAction, InspectionResult, ToolInspector};
use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::Tool;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::{Arc, RwLock};

/// Permission Inspector that handles tool permission checking
pub struct PermissionInspector {
readonly_tools: HashSet<String>,
regular_tools: HashSet<String>,
pub permission_manager: Arc<PermissionManager>,
provider: SharedProvider,
readonly_tools: RwLock<HashSet<String>>,
}

impl PermissionInspector {
pub fn new(
readonly_tools: HashSet<String>,
regular_tools: HashSet<String>,
permission_manager: Arc<PermissionManager>,
) -> Self {
pub fn new(permission_manager: Arc<PermissionManager>, provider: SharedProvider) -> Self {
Self {
readonly_tools,
regular_tools,
permission_manager,
provider,
readonly_tools: RwLock::new(HashSet::new()),
}
}

// readonly_tools is per-agent to avoid concurrent session clobbering; write-annotated
// tools are cached globally via PermissionManager.
pub fn apply_tool_annotations(&self, tools: &[Tool]) {
let mut readonly_annotated = HashSet::new();
for tool in tools {
let Some(anns) = &tool.annotations else {
continue;
};
if anns.read_only_hint == Some(true) {
readonly_annotated.insert(tool.name.to_string());
}
}
*self.readonly_tools.write().unwrap() = readonly_annotated;
self.permission_manager.apply_tool_annotations(tools);
}

pub fn is_readonly_annotated_tool(&self, tool_name: &str) -> bool {
self.readonly_tools.read().unwrap().contains(tool_name)
}

/// Process inspection results into permission decisions
/// This method takes all inspection results and converts them into a PermissionCheckResult
/// that can be used by the agent to determine which tools to approve, deny, or ask for approval
Expand Down Expand Up @@ -105,12 +123,14 @@ impl ToolInspector for PermissionInspector {

async fn inspect(
&self,
session_id: &str,
tool_requests: &[ToolRequest],
_messages: &[Message],
goose_mode: GooseMode,
) -> Result<Vec<InspectionResult>> {
let mut results = Vec::new();
let permission_manager = &self.permission_manager;
let mut llm_detect_candidates: Vec<&ToolRequest> = Vec::new();

for request in tool_requests {
if let Ok(tool_call) = &request.tool_call {
Expand All @@ -129,21 +149,28 @@ impl ToolInspector for PermissionInspector {
InspectionAction::RequireApproval(None)
}
}
}
// 2. Check if it's a readonly or regular tool (both pre-approved)
else if self.readonly_tools.contains(&**tool_name)
|| self.regular_tools.contains(&**tool_name)
// 2. Check if it's a smart-approved tool (annotation or cached LLM decision)
} else if self.is_readonly_annotated_tool(tool_name)
|| (goose_mode == GooseMode::SmartApprove
&& permission_manager.get_smart_approve_permission(tool_name)
== Some(PermissionLevel::AlwaysAllow))
{
InspectionAction::Allow
Comment thread
codefromthecrypt marked this conversation as resolved.
}
// 4. Special case for extension management
else if tool_name == MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE {
// 3. Special case for extension management
} else if tool_name == MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE {
InspectionAction::RequireApproval(Some(
"Extension management requires approval for security".to_string(),
))
}
// 4. Defer to LLM detection (SmartApprove, not yet cached)
} else if goose_mode == GooseMode::SmartApprove
&& permission_manager
.get_smart_approve_permission(tool_name)
.is_none()
{
llm_detect_candidates.push(request);
continue;
// 5. Default: require approval for unknown tools
else {
} else {
InspectionAction::RequireApproval(None)
}
}
Expand All @@ -153,10 +180,10 @@ impl ToolInspector for PermissionInspector {
InspectionAction::Allow => {
if goose_mode == GooseMode::Auto {
"Auto mode - all tools approved".to_string()
} else if self.readonly_tools.contains(&**tool_name) {
"Tool marked as read-only".to_string()
} else if self.regular_tools.contains(&**tool_name) {
"Tool pre-approved".to_string()
} else if self.is_readonly_annotated_tool(tool_name) {
"Tool annotated as read-only".to_string()
} else if goose_mode == GooseMode::SmartApprove {
"SmartApprove cached as read-only".to_string()
} else {
"User permission allows this tool".to_string()
}
Expand All @@ -182,6 +209,99 @@ impl ToolInspector for PermissionInspector {
}
}

// LLM-based read-only detection for deferred SmartApprove candidates
if !llm_detect_candidates.is_empty() {
let detected: HashSet<String> = match self.provider.lock().await.clone() {
Some(provider) => {
detect_read_only_tools(provider, session_id, llm_detect_candidates.to_vec())
.await
.into_iter()
.collect()
}
Comment thread
codefromthecrypt marked this conversation as resolved.
None => Default::default(),
};

for candidate in &llm_detect_candidates {
let is_readonly = candidate
.tool_call
.as_ref()
.map(|tc| detected.contains(&tc.name.to_string()))
.unwrap_or(false);
Comment thread
codefromthecrypt marked this conversation as resolved.

// Cache the LLM decision for future calls
if let Ok(tc) = &candidate.tool_call {
let level = if is_readonly {
PermissionLevel::AlwaysAllow
} else {
PermissionLevel::AskBefore
};
Comment thread
codefromthecrypt marked this conversation as resolved.
permission_manager.update_smart_approve_permission(&tc.name, level);
}

results.push(InspectionResult {
tool_request_id: candidate.id.clone(),
action: if is_readonly {
InspectionAction::Allow
} else {
InspectionAction::RequireApproval(None)
},
reason: if is_readonly {
"LLM detected as read-only".to_string()
} else {
"Tool requires user approval".to_string()
},
confidence: 1.0, // Permission decisions are definitive
inspector_name: self.name().to_string(),
finding_id: None,
});
}
}

Ok(results)
}
}

#[cfg(test)]
mod tests {
use super::*;
use rmcp::model::CallToolRequestParams;
use rmcp::object;
use std::sync::Arc;
use test_case::test_case;
use tokio::sync::Mutex;

#[test_case(GooseMode::Auto, false, None, InspectionAction::Allow; "auto_allows")]
#[test_case(GooseMode::SmartApprove, true, None, InspectionAction::Allow; "smart_approve_annotation_allows")]
#[test_case(GooseMode::SmartApprove, false, Some(PermissionLevel::AlwaysAllow), InspectionAction::Allow; "smart_approve_cached_allow")]
#[test_case(GooseMode::SmartApprove, false, Some(PermissionLevel::AskBefore), InspectionAction::RequireApproval(None); "smart_approve_cached_ask")]
#[test_case(GooseMode::SmartApprove, false, None, InspectionAction::RequireApproval(None); "smart_approve_unknown_defers")]
#[test_case(GooseMode::Approve, false, None, InspectionAction::RequireApproval(None); "approve_requires_approval")]
#[test_case(GooseMode::Approve, false, Some(PermissionLevel::AlwaysAllow), InspectionAction::RequireApproval(None); "approve_ignores_cache")]
#[tokio::test]
async fn test_inspect_action(
mode: GooseMode,
smart_approved: bool,
cache: Option<PermissionLevel>,
expected: InspectionAction,
) {
let pm = Arc::new(PermissionManager::new(tempfile::tempdir().unwrap().keep()));
if let Some(level) = cache {
pm.update_smart_approve_permission("tool", level);
}
let inspector = PermissionInspector::new(pm, Arc::new(Mutex::new(None)));
if smart_approved {
*inspector.readonly_tools.write().unwrap() = ["tool".to_string()].into_iter().collect();
}
let req = ToolRequest {
id: "req".into(),
tool_call: Ok(CallToolRequestParams::new("tool").with_arguments(object!({}))),
metadata: None,
tool_meta: None,
};
let results = inspector
.inspect(goose_test_support::TEST_SESSION_ID, &[req], &[], mode)
.await
.unwrap();
assert_eq!(results[0].action, expected);
}
}
Loading
Loading