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

feat: start use tool permission confirmation struct (#2044)

Browse Source
This commit is contained in:
Yingjie He
2025-04-04 16:39:04 -07:00
committed by GitHub
parent bfed33c5a4
commit 890f266e7b
12 changed files with 88 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+2 -2
View File
@@ -9,9 +9,9 @@ use serde_json::Value;
use std::sync::Arc;
use super::extension::{ExtensionConfig, ExtensionResult};
use crate::message::Message;
use crate::providers::base::Provider;
use crate::session;
use crate::{message::Message, permission::PermissionConfirmation};
use mcp_core::{prompt::Prompt, protocol::GetPromptResult, Content, ToolResult};
/// Session configuration for an agent
@@ -50,7 +50,7 @@ pub trait Agent: Send + Sync {
async fn extend_system_prompt(&mut self, extension: String);
/// Handle a confirmation response for a tool request
async fn handle_confirmation(&self, request_id: String, confirmed: bool);
async fn handle_confirmation(&self, request_id: String, confirmation: PermissionConfirmation);
/// Override the system prompt with custom text
async fn override_system_prompt(&mut self, template: String);
crates/goose/src/agents/mod.rs
+1 -5
View File
@@ -1,9 +1,7 @@
mod agent;
mod capabilities;
pub mod capabilities;
pub mod extension;
mod factory;
mod permission_judge;
mod permission_store;
mod reference;
mod summarize;
mod truncate;
@@ -13,5 +11,3 @@ pub use agent::{Agent, SessionConfig};
pub use capabilities::Capabilities;
pub use extension::ExtensionConfig;
pub use factory::{register_agent, AgentFactory};
pub use permission_judge::detect_read_only_tools;
pub use permission_store::ToolPermissionStore;
crates/goose/src/agents/permission_judge.rs
-273
View File
@@ -1,273 +0,0 @@
use crate::agents::capabilities::Capabilities;
use crate::message::{Message, MessageContent, ToolRequest};
use chrono::Utc;
use indoc::indoc;
use mcp_core::tool::ToolAnnotations;
use mcp_core::{tool::Tool, TextContent};
use serde_json::{json, Value};
/// Creates the tool definition for checking read-only permissions.
fn create_read_only_tool() -> Tool {
Tool::new(
"platform__tool_by_tool_permission".to_string(),
indoc! {r#"
Analyze the tool requests and determine which ones perform read-only operations.
What constitutes a read-only operation:
- A read-only operation retrieves information without modifying any data or state.
- Examples include:
- Reading a file without writing to it.
- Querying a database without making updates.
- Retrieving information from APIs without performing POST, PUT, or DELETE operations.
Examples of read vs. write operations:
- Read Operations:
- `SELECT` query in SQL.
- Reading file metadata or content.
- Listing directory contents.
- Write Operations:
- `INSERT`, `UPDATE`, or `DELETE` in SQL.
- Writing or appending to a file.
- Modifying system configurations.
How to analyze tool requests:
- Inspect each tool request to identify its purpose based on its name and arguments.
- Categorize the operation as read-only if it does not involve any state or data modification.
- Return a list of tool names that are strictly read-only.
Use this analysis to generate the list of tools performing read-only operations from the provided tool requests.
"#}
.to_string(),
json!({
"type": "object",
"properties": {
"read_only_tools": {
"type": "array",
"items": {
"type": "string"
},
"description": "Optional list of tool names which has read-only operations."
}
},
"required": []
}),
Some(ToolAnnotations {
title: Some("Check tool operation".to_string()),
read_only_hint: true,
destructive_hint: false,
idempotent_hint: false,
open_world_hint: false,
}),
)
}
/// Builds the message to be sent to the LLM for detecting read-only operations.
fn create_check_messages(tool_requests: Vec<&ToolRequest>) -> Vec<Message> {
let mut check_messages = vec![];
check_messages.push(Message {
role: mcp_core::Role::User,
created: Utc::now().timestamp(),
content: vec![MessageContent::Text(TextContent {
text: format!(
"Here are the tool requests: {:?}\n\nAnalyze the tool requests and list the tools that perform read-only operations. \
\n\nGuidelines for Read-Only Operations: \
\n- Read-only operations do not modify any data or state. \
\n- Examples include file reading, SELECT queries in SQL, and directory listing. \
\n- Write operations include INSERT, UPDATE, DELETE, and file writing. \
\n\nPlease provide a list of tool names that qualify as read-only:",
tool_requests,
),
annotations: None,
})],
});
check_messages
}
/// Processes the response to extract the list of tools with read-only operations.
fn extract_read_only_tools(response: &Message) -> Option<Vec<String>> {
for content in &response.content {
if let MessageContent::ToolRequest(tool_request) = content {
if let Ok(tool_call) = &tool_request.tool_call {
if tool_call.name == "platform__tool_by_tool_permission" {
if let Value::Object(arguments) = &tool_call.arguments {
if let Some(Value::Array(read_only_tools)) =
arguments.get("read_only_tools")
{
return Some(
read_only_tools
.iter()
.filter_map(|tool| tool.as_str().map(String::from))
.collect(),
);
}
}
}
}
}
}
None
}
/// Executes the read-only tools detection and returns the list of tools with read-only operations.
pub async fn detect_read_only_tools(
capabilities: &Capabilities,
tool_requests: Vec<&ToolRequest>,
) -> Vec<String> {
if tool_requests.is_empty() {
return vec![];
}
let tool = create_read_only_tool();
let check_messages = create_check_messages(tool_requests);
let res = capabilities
.provider()
.complete(
"You are a good analyst and can detect operations whether they have read-only operations.",
&check_messages,
&[tool.clone()],
)
.await;
// Process the response and return an empty vector if the response is invalid
if let Ok((message, _usage)) = res {
extract_read_only_tools(&message).unwrap_or_default()
} else {
vec![]
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::agents::capabilities::Capabilities;
use crate::message::{Message, MessageContent, ToolRequest};
use crate::model::ModelConfig;
use crate::providers::base::{Provider, ProviderMetadata, ProviderUsage, Usage};
use crate::providers::errors::ProviderError;
use chrono::Utc;
use mcp_core::ToolCall;
use mcp_core::{tool::Tool, Role, ToolResult};
use serde_json::json;
#[derive(Clone)]
struct MockProvider {
model_config: ModelConfig,
}
#[async_trait::async_trait]
impl Provider for MockProvider {
fn metadata() -> ProviderMetadata {
ProviderMetadata::empty()
}
fn get_model_config(&self) -> ModelConfig {
self.model_config.clone()
}
async fn complete(
&self,
_system: &str,
_messages: &[Message],
_tools: &[Tool],
) -> anyhow::Result<(Message, ProviderUsage), ProviderError> {
Ok((
Message {
role: Role::Assistant,
created: Utc::now().timestamp(),
content: vec![MessageContent::ToolRequest(ToolRequest {
id: "mock_tool_request".to_string(),
tool_call: ToolResult::Ok(ToolCall {
name: "platform__tool_by_tool_permission".to_string(),
arguments: json!({
"read_only_tools": ["file_reader", "data_fetcher"]
}),
}),
})],
},
ProviderUsage::new("mock".to_string(), Usage::default()),
))
}
}
fn create_mock_capabilities() -> Capabilities {
let mock_model_config =
ModelConfig::new("test-model".to_string()).with_context_limit(200_000.into());
Capabilities::new(Box::new(MockProvider {
model_config: mock_model_config,
}))
}
#[tokio::test]
async fn test_create_read_only_tool() {
let tool = create_read_only_tool();
assert_eq!(tool.name, "platform__tool_by_tool_permission");
assert!(tool.description.contains("read-only operation"));
}
#[test]
fn test_create_check_messages() {
let tool_request = ToolRequest {
id: "tool_1".to_string(),
tool_call: ToolResult::Ok(ToolCall {
name: "file_reader".to_string(),
arguments: json!({"path": "/path/to/file"}),
}),
};
let messages = create_check_messages(vec![&tool_request]);
assert_eq!(messages.len(), 1);
let content = &messages[0].content[0];
if let MessageContent::Text(text_content) = content {
assert!(text_content
.text
.contains("Analyze the tool requests and list the tools"));
assert!(text_content.text.contains("file_reader"));
} else {
panic!("Expected text content");
}
}
#[test]
fn test_extract_read_only_tools() {
let message = Message {
role: Role::Assistant,
created: Utc::now().timestamp(),
content: vec![MessageContent::ToolRequest(ToolRequest {
id: "tool_2".to_string(),
tool_call: ToolResult::Ok(ToolCall {
name: "platform__tool_by_tool_permission".to_string(),
arguments: json!({
"read_only_tools": ["file_reader", "data_fetcher"]
}),
}),
})],
};
let result = extract_read_only_tools(&message);
assert!(result.is_some());
let tools = result.unwrap();
assert_eq!(tools, vec!["file_reader", "data_fetcher"]);
}
#[tokio::test]
async fn test_detect_read_only_tools() {
let capabilities = create_mock_capabilities();
let tool_request = ToolRequest {
id: "tool_1".to_string(),
tool_call: ToolResult::Ok(ToolCall {
name: "file_reader".to_string(),
arguments: json!({"path": "/path/to/file"}),
}),
};
let result = detect_read_only_tools(&capabilities, vec![&tool_request]).await;
assert_eq!(result, vec!["file_reader", "data_fetcher"]);
}
#[tokio::test]
async fn test_detect_read_only_tools_empty_requests() {
let capabilities = create_mock_capabilities();
let result = detect_read_only_tools(&capabilities, vec![]).await;
assert!(result.is_empty());
}
}
crates/goose/src/agents/permission_store.rs
-149
View File
@@ -1,149 +0,0 @@
use crate::message::ToolRequest;
use anyhow::Result;
use blake3::Hasher;
use chrono::Utc;
use etcetera::{choose_app_strategy, AppStrategy};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::time::Duration;
use std::{fs::File, path::PathBuf};
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ToolPermissionRecord {
tool_name: String,
allowed: bool,
context_hash: String, // Hash of the tool's arguments/context to differentiate similar calls
#[serde(skip_serializing_if = "Option::is_none")] // Don't serialize if None
readable_context: Option<String>, // Add this field
timestamp: i64,
expiry: Option<i64>, // Optional expiry timestamp
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ToolPermissionStore {
permissions: HashMap<String, Vec<ToolPermissionRecord>>,
version: u32, // For future schema migrations
#[serde(skip)] // Don't serialize this field
permissions_dir: PathBuf,
}
impl Default for ToolPermissionStore {
fn default() -> Self {
Self::new()
}
}
impl ToolPermissionStore {
pub fn new() -> Self {
let permissions_dir = choose_app_strategy(crate::config::APP_STRATEGY.clone())
.map(|strategy| strategy.config_dir())
.unwrap_or_else(|_| PathBuf::from(".config/goose"));
Self {
permissions: HashMap::new(),
version: 1,
permissions_dir,
}
}
pub fn load() -> Result<Self> {
let store = Self::new();
let file_path = store.permissions_dir.join("tool_permissions.json");
if !file_path.exists() {
return Ok(store);
}
let file = File::open(file_path)?;
let mut permissions: ToolPermissionStore = serde_json::from_reader(file)?;
permissions.permissions_dir = store.permissions_dir;
// Clean up expired entries on load
permissions.cleanup_expired()?;
Ok(permissions)
}
pub fn save(&self) -> anyhow::Result<()> {
std::fs::create_dir_all(&self.permissions_dir)?;
let path = self.permissions_dir.join("tool_permissions.json");
let temp_path = path.with_extension("tmp");
// Write complete content to temporary file
let content = serde_json::to_string_pretty(self)?;
std::fs::write(&temp_path, &content)?;
// Atomically rename temp file to target file
std::fs::rename(temp_path, path)?;
Ok(())
}
pub fn check_permission(&self, tool_request: &ToolRequest) -> Option<bool> {
let context_hash = self.hash_tool_context(tool_request);
let tool_call = tool_request.tool_call.as_ref().unwrap();
let key = format!("{}:{}", tool_call.name, context_hash);
self.permissions.get(&key).and_then(|records| {
records
.iter()
.filter(|record| record.expiry.is_none_or(|exp| exp > Utc::now().timestamp()))
.next_back()
.map(|record| record.allowed)
})
}
pub fn record_permission(
&mut self,
tool_request: &ToolRequest,
allowed: bool,
expiry_duration: Option<Duration>,
) -> anyhow::Result<()> {
let context_hash = self.hash_tool_context(tool_request);
let tool_call = tool_request.tool_call.as_ref().unwrap();
let key = format!("{}:{}", tool_call.name, context_hash);
let record = ToolPermissionRecord {
tool_name: tool_call.name.clone(),
allowed,
context_hash,
readable_context: Some(tool_request.to_readable_string()),
timestamp: Utc::now().timestamp(),
expiry: expiry_duration.map(|d| Utc::now().timestamp() + d.as_secs() as i64),
};
self.permissions.entry(key).or_default().push(record);
self.save()?;
Ok(())
}
fn hash_tool_context(&self, tool_request: &ToolRequest) -> String {
// Create a hash of the tool's arguments to differentiate similar calls
// This helps identify when the same tool is being used in a different context
let mut hasher = Hasher::new();
hasher.update(
serde_json::to_string(&tool_request.tool_call.as_ref().unwrap().arguments)
.unwrap_or_default()
.as_bytes(),
);
hasher.finalize().to_hex().to_string()
}
pub fn cleanup_expired(&mut self) -> anyhow::Result<()> {
let now = Utc::now().timestamp();
let mut changed = false;
self.permissions.retain(|_, records| {
records.retain(|record| record.expiry.is_none_or(|exp| exp > now));
changed = changed || records.is_empty();
!records.is_empty()
});
if changed {
self.save()?;
}
Ok(())
}
}
crates/goose/src/agents/reference.rs
+6 -1
View File
@@ -15,6 +15,7 @@ use super::Agent;
use crate::agents::capabilities::Capabilities;
use crate::agents::extension::{ExtensionConfig, ExtensionResult};
use crate::message::{Message, ToolRequest};
use crate::permission::PermissionConfirmation;
use crate::providers::base::Provider;
use crate::token_counter::TokenCounter;
use crate::{register_agent, session};
@@ -73,7 +74,11 @@ impl Agent for ReferenceAgent {
Ok(Value::Null)
}
async fn handle_confirmation(&self, _request_id: String, _confirmed: bool) {
async fn handle_confirmation(
&self,
_request_id: String,
_confirmation: PermissionConfirmation,
) {
// TODO implement
}
crates/goose/src/agents/summarize.rs
+9 -7
View File
@@ -12,7 +12,6 @@ use tracing::{debug, error, instrument, warn};
use super::agent::SessionConfig;
use super::capabilities::get_parameter_names;
use super::detect_read_only_tools;
use super::extension::ToolInfo;
use super::Agent;
use crate::agents::capabilities::Capabilities;
@@ -20,6 +19,9 @@ use crate::agents::extension::{ExtensionConfig, ExtensionResult};
use crate::config::Config;
use crate::memory_condense::condense_messages;
use crate::message::{Message, ToolRequest};
use crate::permission::detect_read_only_tools;
use crate::permission::Permission;
use crate::permission::PermissionConfirmation;
use crate::providers::base::Provider;
use crate::providers::errors::ProviderError;
use crate::register_agent;
@@ -38,8 +40,8 @@ const ESTIMATE_FACTOR_DECAY: f32 = 0.9;
pub struct SummarizeAgent {
capabilities: Mutex<Capabilities>,
token_counter: TokenCounter,
confirmation_tx: mpsc::Sender<(String, bool)>, // (request_id, confirmed)
confirmation_rx: Mutex<mpsc::Receiver<(String, bool)>>,
confirmation_tx: mpsc::Sender<(String, PermissionConfirmation)>,
confirmation_rx: Mutex<mpsc::Receiver<(String, PermissionConfirmation)>>,
tool_result_tx: mpsc::Sender<(String, ToolResult<Vec<Content>>)>,
}
@@ -159,8 +161,8 @@ impl Agent for SummarizeAgent {
}
/// Handle a confirmation response for a tool request
async fn handle_confirmation(&self, request_id: String, confirmed: bool) {
if let Err(e) = self.confirmation_tx.send((request_id, confirmed)).await {
async fn handle_confirmation(&self, request_id: String, confirmation: PermissionConfirmation) {
if let Err(e) = self.confirmation_tx.send((request_id, confirmation)).await {
error!("Failed to send confirmation: {}", e);
}
}
@@ -321,9 +323,9 @@ impl Agent for SummarizeAgent {
// Wait for confirmation response through the channel
let mut rx = self.confirmation_rx.lock().await;
// Loop the recv until we have a matched req_id due to potential duplicate messages.
while let Some((req_id, confirmed)) = rx.recv().await {
while let Some((req_id, tool_confirmation)) = rx.recv().await {
if req_id == request.id {
if confirmed {
if tool_confirmation.permission == Permission::AllowOnce || tool_confirmation.permission == Permission::AlwaysAllow {
// User approved - dispatch the tool call
let output = capabilities.dispatch_tool_call(tool_call).await;
message_tool_response = message_tool_response.with_tool_response(
crates/goose/src/agents/truncate.rs
+10 -11
View File
@@ -10,15 +10,17 @@ use tokio::sync::Mutex;
use tracing::{debug, error, instrument, warn};
use super::agent::SessionConfig;
use super::detect_read_only_tools;
use super::extension::ToolInfo;
use super::types::ToolResultReceiver;
use super::Agent;
use crate::agents::capabilities::{get_parameter_names, Capabilities};
use crate::agents::extension::{ExtensionConfig, ExtensionResult};
use crate::agents::ToolPermissionStore;
use crate::config::Config;
use crate::message::{Message, MessageContent, ToolRequest};
use crate::permission::detect_read_only_tools;
use crate::permission::Permission;
use crate::permission::PermissionConfirmation;
use crate::permission::ToolPermissionStore;
use crate::providers::base::Provider;
use crate::providers::errors::ProviderError;
use crate::providers::toolshim::{
@@ -34,7 +36,6 @@ use mcp_core::{
prompt::Prompt, protocol::GetPromptResult, tool::Tool, Content, ToolError, ToolResult,
};
use serde_json::{json, Value};
use std::time::Duration;
const MAX_TRUNCATION_ATTEMPTS: usize = 3;
const ESTIMATE_FACTOR_DECAY: f32 = 0.9;
@@ -43,8 +44,8 @@ const ESTIMATE_FACTOR_DECAY: f32 = 0.9;
pub struct TruncateAgent {
capabilities: Mutex<Capabilities>,
token_counter: TokenCounter,
confirmation_tx: mpsc::Sender<(String, bool)>, // (request_id, confirmed)
confirmation_rx: Mutex<mpsc::Receiver<(String, bool)>>,
confirmation_tx: mpsc::Sender<(String, PermissionConfirmation)>,
confirmation_rx: Mutex<mpsc::Receiver<(String, PermissionConfirmation)>>,
tool_result_tx: mpsc::Sender<(String, ToolResult<Vec<Content>>)>,
tool_result_rx: ToolResultReceiver,
}
@@ -160,8 +161,8 @@ impl Agent for TruncateAgent {
}
/// Handle a confirmation response for a tool request
async fn handle_confirmation(&self, request_id: String, confirmed: bool) {
if let Err(e) = self.confirmation_tx.send((request_id, confirmed)).await {
async fn handle_confirmation(&self, request_id: String, confirmation: PermissionConfirmation) {
if let Err(e) = self.confirmation_tx.send((request_id, confirmation)).await {
error!("Failed to send confirmation: {}", e);
}
}
@@ -432,12 +433,10 @@ impl Agent for TruncateAgent {
// Wait for confirmation response through the channel
let mut rx = self.confirmation_rx.lock().await;
while let Some((req_id, confirmed)) = rx.recv().await {
while let Some((req_id, tool_confirmation)) = rx.recv().await {
if req_id == request.id {
// Store the user's response with 30-day expiration
let mut store = ToolPermissionStore::load()?;
store.record_permission(request, confirmed, Some(Duration::from_secs(30 * 24 * 60 * 60)))?;
let confirmed = tool_confirmation.permission == Permission::AllowOnce || tool_confirmation.permission == Permission::AlwaysAllow;
if confirmed {
// Add this tool call to the futures collection
let tool_future = Self::create_tool_future(&capabilities, tool_call, request.id.clone());