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

feat: permission before tool call (#1313)

Browse Source
This commit is contained in:
Wendy Tang
2025-02-24 11:54:41 -08:00
committed by GitHub
parent bef7551459
commit 3723c64cc5
10 changed files with 226 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+3
View File
@@ -32,6 +32,9 @@ pub trait Agent: Send + Sync {
/// Add custom text to be included in the system prompt
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);
/// Override the system prompt with custom text
async fn override_system_prompt(&mut self, template: String);
}
crates/goose/src/agents/reference.rs
+4
View File
@@ -61,6 +61,10 @@ impl Agent for ReferenceAgent {
Ok(Value::Null)
}
async fn handle_confirmation(&self, _request_id: String, _confirmed: bool) {
// TODO implement
}
#[instrument(skip(self, messages), fields(user_message))]
async fn reply(
&self,
crates/goose/src/agents/truncate.rs
+100 -19
View File
@@ -2,12 +2,14 @@
/// It makes no attempt to handle context limits, and cannot read resources
use async_trait::async_trait;
use futures::stream::BoxStream;
use tokio::sync::mpsc;
use tokio::sync::Mutex;
use tracing::{debug, error, instrument, warn};
use super::Agent;
use crate::agents::capabilities::Capabilities;
use crate::agents::extension::{ExtensionConfig, ExtensionResult};
use crate::config::Config;
use crate::message::{Message, ToolRequest};
use crate::providers::base::Provider;
use crate::providers::base::ProviderUsage;
@@ -16,7 +18,7 @@ use crate::register_agent;
use crate::token_counter::TokenCounter;
use crate::truncate::{truncate_messages, OldestFirstTruncation};
use indoc::indoc;
use mcp_core::tool::Tool;
use mcp_core::{tool::Tool, Content};
use serde_json::{json, Value};
const MAX_TRUNCATION_ATTEMPTS: usize = 3;
@@ -26,14 +28,21 @@ 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)>>,
}
impl TruncateAgent {
pub fn new(provider: Box<dyn Provider>) -> Self {
let token_counter = TokenCounter::new(provider.get_model_config().tokenizer_name());
// Create channel with buffer size 32 (adjust if needed)
let (tx, rx) = mpsc::channel(32);
Self {
capabilities: Mutex::new(Capabilities::new(provider)),
token_counter,
confirmation_tx: tx,
confirmation_rx: Mutex::new(rx),
}
}
@@ -121,6 +130,13 @@ impl Agent for TruncateAgent {
Ok(Value::Null)
}
/// 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 {
error!("Failed to send confirmation: {}", e);
}
}
#[instrument(skip(self, messages), fields(user_message))]
async fn reply(
&self,
@@ -132,6 +148,10 @@ impl Agent for TruncateAgent {
let mut tools = capabilities.get_prefixed_tools().await?;
let mut truncation_attempt: usize = 0;
// Load settings from config
let config = Config::global();
let goose_mode = config.get("GOOSE_MODE").unwrap_or("auto".to_string());
// we add in the 2 resource tools if any extensions support resources
// TODO: make sure there is no collision with another extension's tool name
let read_resource_tool = Tool::new(
@@ -191,7 +211,6 @@ impl Agent for TruncateAgent {
Ok(Box::pin(async_stream::try_stream! {
let _reply_guard = reply_span.enter();
loop {
// Attempt to get completion from provider
match capabilities.provider().complete(
&system_prompt,
&messages,
@@ -218,24 +237,86 @@ impl Agent for TruncateAgent {
break;
}
// Then dispatch each in parallel
let futures: Vec<_> = tool_requests
.iter()
.filter_map(|request| request.tool_call.clone().ok())
.map(|tool_call| capabilities.dispatch_tool_call(tool_call))
.collect();
// Process all the futures in parallel but wait until all are finished
let outputs = futures::future::join_all(futures).await;
// Create a message with the responses
// Process tool requests depending on goose_mode
let mut message_tool_response = Message::user();
// Now combine these into MessageContent::ToolResponse using the original ID
for (request, output) in tool_requests.iter().zip(outputs.into_iter()) {
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
output,
);
// Clone goose_mode once before the match to avoid move issues
let mode = goose_mode.clone();
match mode.as_str() {
"approve" => {
// Process each tool request sequentially with confirmation
for request in &tool_requests {
if let Ok(tool_call) = request.tool_call.clone() {
let confirmation = Message::user().with_tool_confirmation_request(
request.id.clone(),
tool_call.name.clone(),
tool_call.arguments.clone(),
Some("Goose would like to call the tool: {}\nAllow? (y/n): ".to_string()),
);
yield confirmation;
// Wait for confirmation response through the channel
let mut rx = self.confirmation_rx.lock().await;
if let Some((req_id, confirmed)) = rx.recv().await {
if req_id == request.id {
if confirmed {
// User approved - dispatch the tool call
let output = capabilities.dispatch_tool_call(tool_call).await;
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
output,
);
} else {
// User declined - add declined response
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
Ok(vec![Content::text("User declined to run this tool.")]),
);
}
}
}
}
}
},
"chat" => {
// Skip all tool calls in chat mode
for request in &tool_requests {
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
Ok(vec![Content::text(
"The following tool call was skipped in Goose chat mode. \
In chat mode, you cannot run tool calls, instead, you can \
only provide a detailed plan to the user. Provide an \
explanation of the proposed tool call as if it were a plan. \
Only if the user asks, provide a short explanation to the \
user that they could consider running the tool above on \
their own or with a different goose mode."
)]),
);
}
},
_ => {
if mode != "auto" {
warn!("Unknown GOOSE_MODE: {mode:?}. Defaulting to 'auto' mode.");
}
// Process tool requests in parallel
let mut tool_futures = Vec::new();
for request in &tool_requests {
if let Ok(tool_call) = request.tool_call.clone() {
tool_futures.push(async {
let output = capabilities.dispatch_tool_call(tool_call).await;
(request.id.clone(), output)
});
}
}
// Wait for all tool calls to complete
let results = futures::future::join_all(tool_futures).await;
for (request_id, output) in results {
message_tool_response = message_tool_response.with_tool_response(
request_id,
output,
);
}
}
}
yield message_tool_response.clone();
crates/goose/src/message.rs
+44
View File
@@ -12,6 +12,7 @@ use mcp_core::content::{Content, ImageContent, TextContent};
use mcp_core::handler::ToolResult;
use mcp_core::role::Role;
use mcp_core::tool::ToolCall;
use serde_json::Value;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ToolRequest {
@@ -25,6 +26,14 @@ pub struct ToolResponse {
pub tool_result: ToolResult<Vec<Content>>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ToolConfirmationRequest {
pub id: String,
pub tool_name: String,
pub arguments: Value,
pub prompt: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
/// Content passed inside a message, which can be both simple content and tool content
pub enum MessageContent {
@@ -32,6 +41,7 @@ pub enum MessageContent {
Image(ImageContent),
ToolRequest(ToolRequest),
ToolResponse(ToolResponse),
ToolConfirmationRequest(ToolConfirmationRequest),
}
impl MessageContent {
@@ -64,6 +74,19 @@ impl MessageContent {
})
}
pub fn tool_confirmation_request<S: Into<String>>(
id: S,
tool_name: String,
arguments: Value,
prompt: Option<String>,
) -> Self {
MessageContent::ToolConfirmationRequest(ToolConfirmationRequest {
id: id.into(),
tool_name,
arguments,
prompt,
})
}
pub fn as_tool_request(&self) -> Option<&ToolRequest> {
if let MessageContent::ToolRequest(ref tool_request) = self {
Some(tool_request)
@@ -80,6 +103,14 @@ impl MessageContent {
}
}
pub fn as_tool_confirmation_request(&self) -> Option<&ToolConfirmationRequest> {
if let MessageContent::ToolConfirmationRequest(ref tool_confirmation_request) = self {
Some(tool_confirmation_request)
} else {
None
}
}
pub fn as_tool_response_text(&self) -> Option<String> {
if let Some(tool_response) = self.as_tool_response() {
if let Ok(contents) = &tool_response.tool_result {
@@ -178,6 +209,19 @@ impl Message {
self.with_content(MessageContent::tool_response(id, result))
}
/// Add a tool confirmation request to the message
pub fn with_tool_confirmation_request<S: Into<String>>(
self,
id: S,
tool_name: String,
arguments: Value,
prompt: Option<String>,
) -> Self {
self.with_content(MessageContent::tool_confirmation_request(
id, tool_name, arguments, prompt,
))
}
/// Get the concatenated text content of the message, separated by newlines
pub fn as_concat_text(&self) -> String {
self.content
crates/goose/src/providers/formats/anthropic.rs
+3
View File
@@ -57,6 +57,9 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
}));
}
}
MessageContent::ToolConfirmationRequest(_tool_confirmation_request) => {
// Skip tool confirmation requests
}
MessageContent::Image(_) => continue, // Anthropic doesn't support image content yet
}
}
crates/goose/src/providers/formats/bedrock.rs
+3
View File
@@ -28,6 +28,9 @@ pub fn to_bedrock_message(message: &Message) -> Result<bedrock::Message> {
pub fn to_bedrock_message_content(content: &MessageContent) -> Result<bedrock::ContentBlock> {
Ok(match content {
MessageContent::Text(text) => bedrock::ContentBlock::Text(text.text.to_string()),
MessageContent::ToolConfirmationRequest(_tool_confirmation_request) => {
bedrock::ContentBlock::Text("".to_string())
}
MessageContent::Image(_) => {
bail!("Image content is not supported by Bedrock provider yet")
}
crates/goose/src/providers/formats/openai.rs
+3
View File
@@ -136,6 +136,9 @@ pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<
}
}
}
MessageContent::ToolConfirmationRequest(_) => {
// Skip tool confirmation requests
}
MessageContent::Image(image) => {
// Handle direct image content
converted["content"] = json!([convert_image(image, image_format)]);