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

More acp tools (#7843)

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
Co-authored-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Jack Amadeo
2026-03-13 04:28:02 -04:00
committed by GitHub
parent b475ecfc09
commit c947620843
24 changed files with 443 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+7 -5
View File
@@ -542,22 +542,25 @@ impl Agent {
};
}
let ctx = super::tool_execution::ToolCallContext::new(
session.id.clone(),
Some(session.working_dir.clone()),
Some(request_id.clone()),
);
debug!("WAITING_TOOL_START: {}", tool_call.name);
let result: ToolCallResult = if self.is_frontend_tool(&tool_call.name).await {
// For frontend tools, return an error indicating we need frontend execution
ToolCallResult::from(Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
"Frontend tool execution required".to_string(),
None,
)))
} else {
// Clone the result to ensure no references to extension_manager are returned
let result = self
.extension_manager
.dispatch_tool_call(
&session.id,
&ctx,
tool_call.clone(),
Some(session.working_dir.as_path()),
cancellation_token.unwrap_or_default(),
)
.await;
@@ -566,7 +569,6 @@ impl Agent {
"tool_execution_failed",
&format!("{}: {}", tool_call.name, e),
);
// Try to downcast to ErrorData to avoid double wrapping
let error_data = e.downcast::<ErrorData>().unwrap_or_else(|e| {
ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), None)
});
crates/goose/src/agents/extension_manager.rs
+36 -65
View File
@@ -30,7 +30,7 @@ use super::extension::{
ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, PlatformExtensionContext,
ToolInfo, PLATFORM_EXTENSIONS,
};
use super::tool_execution::ToolCallResult;
use super::tool_execution::{ToolCallContext, ToolCallResult};
use super::types::SharedProvider;
use crate::agents::extension::{Envs, ProcessExit};
use crate::agents::extension_malware_check;
@@ -1385,13 +1385,12 @@ impl ExtensionManager {
pub async fn dispatch_tool_call(
&self,
session_id: &str,
ctx: &super::tool_execution::ToolCallContext,
tool_call: CallToolRequestParams,
working_dir: Option<&std::path::Path>,
cancellation_token: CancellationToken,
) -> Result<ToolCallResult> {
let tool_name_str = tool_call.name.to_string();
let resolved = self.resolve_tool(session_id, &tool_name_str).await?;
let resolved = self.resolve_tool(&ctx.session_id, &tool_name_str).await?;
if let Some(extension) = self.extensions.lock().await.get(&resolved.extension_name) {
if !extension
@@ -1413,25 +1412,22 @@ impl ExtensionManager {
let arguments = tool_call.arguments.clone();
let client = resolved.client.clone();
let notifications_receiver = client.subscribe().await;
let session_id = session_id.to_string();
let actual_tool_name = resolved.actual_tool_name;
let working_dir_str = working_dir.map(|p| p.to_string_lossy().to_string());
let owned_ctx = ToolCallContext::new(
ctx.session_id.clone(),
ctx.working_dir.clone(),
ctx.tool_call_request_id.clone(),
);
let fut = async move {
tracing::debug!(
"dispatch_tool_call: calling client.call_tool tool={} session_id={} working_dir={:?}",
actual_tool_name,
session_id,
working_dir_str
owned_ctx.session_id,
owned_ctx.working_dir,
);
client
.call_tool(
&session_id,
&actual_tool_name,
arguments,
working_dir_str.as_deref(),
cancellation_token,
)
.call_tool(&owned_ctx, &actual_tool_name, arguments, cancellation_token)
.await
.map_err(|e| match e {
ServiceError::McpError(error_data) => error_data,
@@ -1798,10 +1794,9 @@ mod tests {
async fn call_tool(
&self,
_session_id: &str,
_ctx: &ToolCallContext,
name: &str,
_arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
match name {
@@ -1838,6 +1833,8 @@ mod tests {
#[tokio::test]
async fn test_dispatch_tool_call() {
use super::super::tool_execution::ToolCallContext;
let temp_dir = tempfile::tempdir().unwrap();
let extension_manager =
ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
@@ -1855,16 +1852,17 @@ mod tests {
.add_mock_extension("client 🚀".to_string(), Arc::new(MockClient {}))
.await;
let ctx = ToolCallContext::new(
"test-session-id".to_string(),
None,
Some("test-req-id".to_string()),
);
let tool_call =
CallToolRequestParams::new("test_client__tool".to_string()).with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
@@ -1872,12 +1870,7 @@ mod tests {
.with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
@@ -1885,12 +1878,7 @@ mod tests {
.with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
@@ -1898,12 +1886,7 @@ mod tests {
CallToolRequestParams::new("client___tool".to_string()).with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
@@ -1911,12 +1894,7 @@ mod tests {
CallToolRequestParams::new("client___tools".to_string()).with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
invalid_tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, invalid_tool_call, CancellationToken::default())
.await;
if let Err(err) = result {
let tool_err = err.downcast_ref::<ErrorData>().expect("Expected ErrorData");
@@ -1929,12 +1907,7 @@ mod tests {
CallToolRequestParams::new("_client__tools".to_string()).with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
invalid_tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, invalid_tool_call, CancellationToken::default())
.await;
if let Err(err) = result {
let tool_err = err.downcast_ref::<ErrorData>().expect("Expected ErrorData");
@@ -2009,6 +1982,8 @@ mod tests {
#[tokio::test]
async fn test_dispatch_unavailable_tool_returns_error() {
use super::super::tool_execution::ToolCallContext;
let temp_dir = tempfile::tempdir().unwrap();
let extension_manager =
ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
@@ -2023,16 +1998,17 @@ mod tests {
)
.await;
let ctx = ToolCallContext::new(
"test-session-id".to_string(),
None,
Some("test-req-id".to_string()),
);
let unavailable_tool_call = CallToolRequestParams::new("test_extension__tool".to_string())
.with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
unavailable_tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, unavailable_tool_call, CancellationToken::default())
.await;
if let Err(err) = result {
@@ -2048,12 +2024,7 @@ mod tests {
.with_arguments(object!({}));
let result = extension_manager
.dispatch_tool_call(
"test-session-id",
available_tool_call,
None,
CancellationToken::default(),
)
.dispatch_tool_call(&ctx, available_tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
crates/goose/src/agents/mcp_client.rs
+9 -5
View File
@@ -1,4 +1,5 @@
use crate::action_required_manager::ActionRequiredManager;
use crate::agents::tool_execution::ToolCallContext;
use crate::agents::types::SharedProvider;
use crate::session_context::{SESSION_ID_HEADER, WORKING_DIR_HEADER};
use rmcp::model::{
@@ -47,10 +48,9 @@ pub trait McpClientTrait: Send + Sync {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
working_dir: Option<&str>,
cancel_token: CancellationToken,
) -> Result<CallToolResult, Error>;
@@ -594,10 +594,9 @@ impl McpClientTrait for McpClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
working_dir: Option<&str>,
cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let mut params = CallToolRequestParams::new(name.to_string());
@@ -607,7 +606,12 @@ impl McpClientTrait for McpClient {
let request = ClientRequest::CallToolRequest(Request::new(params));
let result = self
.send_request_with_context(session_id, working_dir, request, cancel_token)
.send_request_with_context(
&ctx.session_id,
ctx.working_dir_str(),
request,
cancel_token,
)
.await;
match result? {
crates/goose/src/agents/mod.rs
+1
View File
@@ -30,4 +30,5 @@ pub use extension_manager::ExtensionManager;
pub use prompt_manager::PromptManager;
pub use subagent_handler::SUBAGENT_TOOL_REQUEST_TYPE;
pub use subagent_task_config::TaskConfig;
pub use tool_execution::ToolCallContext;
pub use types::{FrontendTool, RetryConfig, SessionConfig, SuccessCheck};
crates/goose/src/agents/platform_extensions/analyze/mod.rs
+3 -3
View File
@@ -5,6 +5,7 @@ pub mod parser;
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use ignore::WalkBuilder;
@@ -234,13 +235,12 @@ impl McpClientTrait for AnalyzeClient {
async fn call_tool(
&self,
_session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let working_dir = working_dir.map(Path::new);
let working_dir = ctx.working_dir.as_deref();
match name {
"analyze" => match Self::parse_args::<AnalyzeParams>(arguments) {
Ok(params) => {
crates/goose/src/agents/platform_extensions/apps.rs
+3 -2
View File
@@ -1,5 +1,6 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use crate::config::paths::Paths;
use crate::conversation::message::Message;
use crate::goose_apps::McpAppResource;
@@ -526,12 +527,12 @@ impl McpClientTrait for AppsManagerClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let session_id = &ctx.session_id;
let result = match name {
"list_apps" => self.handle_list_apps(arguments).await,
"create_app" => self.handle_create_app(session_id, arguments).await,
crates/goose/src/agents/platform_extensions/chatrecall.rs
+3 -2
View File
@@ -1,5 +1,6 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use indoc::indoc;
@@ -273,12 +274,12 @@ impl McpClientTrait for ChatRecallClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let session_id = &ctx.session_id;
let content = match name {
"chatrecall" => self.handle_chatrecall(session_id, arguments).await,
_ => Err(format!("Unknown tool: {}", name)),
crates/goose/src/agents/platform_extensions/code_execution.rs
+11 -5
View File
@@ -1,6 +1,7 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::extension_manager::get_tool_owner;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use indoc::indoc;
@@ -252,8 +253,13 @@ fn create_tool_callback(
}
params
};
let ctx = crate::agents::ToolCallContext::new(
session_id,
None,
Some("tool-request-id".to_string()),
);
match manager
.dispatch_tool_call(&session_id, tool_call, None, CancellationToken::new())
.dispatch_tool_call(&ctx, tool_call, CancellationToken::new())
.await
{
Ok(dispatch_result) => match dispatch_result.result.await {
@@ -318,10 +324,10 @@ impl McpClientTrait for CodeExecutionClient {
"list_functions".to_string(),
indoc! {r#"
List all available functions across all namespaces.
This will not return function input and output types.
After determining which functions are needed use
get_function_details to get input and output type
get_function_details to get input and output type
information about specific functions.
"#}
.to_string(),
@@ -420,12 +426,12 @@ impl McpClientTrait for CodeExecutionClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let session_id = &ctx.session_id;
let result = match name {
"list_functions" => self.handle_list_functions(session_id).await,
"get_function_details" => {
crates/goose/src/agents/platform_extensions/developer/mod.rs
+9 -11
View File
@@ -4,6 +4,7 @@ pub mod tree;
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use edit::{EditTools, FileEditParams, FileWriteParams};
@@ -15,7 +16,6 @@ use rmcp::model::{
use schemars::{schema_for, JsonSchema};
use serde_json::Value;
use shell::{ShellOutput, ShellParams, ShellTool};
use std::path::Path;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use tree::{TreeParams, TreeTool};
@@ -146,13 +146,12 @@ impl McpClientTrait for DeveloperClient {
async fn call_tool(
&self,
_session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
working_dir: Option<&str>,
_cancellation_token: CancellationToken,
_cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let working_dir = working_dir.map(Path::new);
let working_dir = ctx.working_dir.as_deref();
match name {
"shell" => match Self::parse_args::<ShellParams>(arguments) {
Ok(params) => Ok(self.shell_tool.shell_with_cwd(params, working_dir).await),
@@ -231,15 +230,15 @@ mod tests {
let cwd = temp.path().join("workspace");
fs::create_dir_all(&cwd).unwrap();
let ctx = ToolCallContext::new("session".to_owned(), Some(cwd.clone()), None);
let write = client
.call_tool(
"session",
&ctx,
"write",
Some(object!({
"path": "notes.txt",
"content": "first line"
})),
Some(cwd.to_str().unwrap()),
CancellationToken::new(),
)
.await
@@ -252,14 +251,13 @@ mod tests {
let edit = client
.call_tool(
"session",
&ctx,
"edit",
Some(object!({
"path": "notes.txt",
"before": "first",
"after": "updated"
})),
Some(cwd.to_str().unwrap()),
CancellationToken::new(),
)
.await
@@ -279,14 +277,14 @@ mod tests {
let cwd = temp.path().join("workspace");
fs::create_dir_all(&cwd).unwrap();
let ctx = ToolCallContext::new("session".to_owned(), Some(cwd.clone()), None);
let result = client
.call_tool(
"session",
&ctx,
"shell",
Some(object!({
"command": "pwd"
})),
Some(cwd.to_str().unwrap()),
CancellationToken::new(),
)
.await
crates/goose/src/agents/platform_extensions/developer/shell.rs
+1 -1
View File
@@ -13,7 +13,7 @@ use tokio_stream::{wrappers::SplitStream, StreamExt};
use crate::subprocess::SubprocessExt;
const OUTPUT_LIMIT_LINES: usize = 2000;
const OUTPUT_LIMIT_BYTES: usize = 50_000;
pub const OUTPUT_LIMIT_BYTES: usize = 50_000;
const OUTPUT_PREVIEW_LINES: usize = 50;
const OUTPUT_SLOTS: usize = 8;
crates/goose/src/agents/platform_extensions/ext_manager.rs
+3 -2
View File
@@ -1,5 +1,6 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use crate::config::get_extension_by_name;
use anyhow::Result;
use async_trait::async_trait;
@@ -408,12 +409,12 @@ impl McpClientTrait for ExtensionManagerClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let session_id = &ctx.session_id;
let result = match name {
SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME => {
self.handle_search_available_extensions().await
crates/goose/src/agents/platform_extensions/summarize.rs
+4 -3
View File
@@ -1,6 +1,7 @@
use std::path::{Path, PathBuf};
use std::sync::Arc;
use crate::agents::tool_execution::ToolCallContext;
use async_trait::async_trait;
use ignore::gitignore::{Gitignore, GitignoreBuilder};
use rmcp::model::{
@@ -98,10 +99,9 @@ impl McpClientTrait for SummarizeClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
if name != "summarize" {
@@ -111,7 +111,7 @@ impl McpClientTrait for SummarizeClient {
))]));
}
let Some(working_dir) = working_dir else {
let Some(working_dir) = ctx.working_dir_str() else {
return Ok(CallToolResult::error(vec![Content::text(
"Error: working_dir is required for summarize",
)]));
@@ -148,6 +148,7 @@ impl McpClientTrait for SummarizeClient {
}
};
let session_id = &ctx.session_id;
match execute_summarize(provider, session_id, params, &working_dir).await {
Ok(result) => Ok(result),
Err(msg) => Ok(CallToolResult::error(vec![Content::text(format!(
crates/goose/src/agents/platform_extensions/summon.rs
+5 -3
View File
@@ -9,6 +9,7 @@ use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::subagent_handler::{run_subagent_task, OnMessageCallback, SubagentRunParams};
use crate::agents::subagent_task_config::{TaskConfig, DEFAULT_SUBAGENT_MAX_TURNS};
use crate::agents::tool_execution::ToolCallContext;
use crate::agents::AgentConfig;
use crate::config::paths::Paths;
use crate::config::Config;
@@ -1815,12 +1816,12 @@ impl McpClientTrait for SummonClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let session_id = &ctx.session_id;
let content = match name {
"load" => self.handle_load(session_id, arguments).await,
"delegate" => {
@@ -2252,8 +2253,9 @@ You review code."#;
let names: Vec<_> = result.tools.iter().map(|t| t.name.as_ref()).collect();
assert!(names.contains(&"load") && names.contains(&"delegate"));
let ctx = ToolCallContext::new("test".to_string(), None, None);
let result = client
.call_tool("test", "unknown", None, None, CancellationToken::new())
.call_tool(&ctx, "unknown", None, CancellationToken::new())
.await
.unwrap();
assert!(result.is_error.unwrap_or(false));
crates/goose/src/agents/platform_extensions/todo.rs
+3 -2
View File
@@ -1,5 +1,6 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use crate::session::extension_data;
use crate::session::extension_data::ExtensionState;
use anyhow::Result;
@@ -155,12 +156,12 @@ impl McpClientTrait for TodoClient {
async fn call_tool(
&self,
session_id: &str,
ctx: &ToolCallContext,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let session_id = &ctx.session_id;
let content = match name {
"todo_write" => self.handle_write_todo(session_id, arguments).await,
_ => Err(format!("Unknown tool: {}", name)),
crates/goose/src/agents/platform_extensions/tom.rs
+2 -2
View File
@@ -1,5 +1,6 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::mcp_client::{Error, McpClientTrait};
use crate::agents::tool_execution::ToolCallContext;
use anyhow::Result;
use async_trait::async_trait;
use rmcp::model::{
@@ -45,10 +46,9 @@ impl McpClientTrait for TomClient {
async fn call_tool(
&self,
_session_id: &str,
_ctx: &ToolCallContext,
name: &str,
_arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
Ok(CallToolResult::error(vec![Content::text(format!(
crates/goose/src/agents/tool_execution.rs
+27
View File
@@ -8,11 +8,38 @@ use futures::{Stream, StreamExt};
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use std::path::PathBuf;
use crate::config::permission::PermissionLevel;
use crate::mcp_utils::ToolResult;
use crate::permission::Permission;
use rmcp::model::{Content, ServerNotification};
/// Context passed through the tool call dispatch chain.
pub struct ToolCallContext {
pub session_id: String,
pub working_dir: Option<PathBuf>,
pub tool_call_request_id: Option<String>,
}
impl ToolCallContext {
pub fn new(
session_id: String,
working_dir: Option<PathBuf>,
tool_call_request_id: Option<String>,
) -> Self {
Self {
session_id,
working_dir,
tool_call_request_id,
}
}
pub fn working_dir_str(&self) -> Option<&str> {
self.working_dir.as_ref().and_then(|p| p.to_str())
}
}
// ToolCallResult combines the result of a tool call with an optional notification stream that
// can be used to receive notifications from the tool.
pub struct ToolCallResult {