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

Detect client disconnects and cancel tool calls (#3782)

Browse Source
This commit is contained in:
Jack Amadeo
2025-08-05 05:48:15 -04:00
committed by GitHub
parent bef7622123
commit 7174fc74c3
8 changed files with 331 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+12 -6
View File
@@ -440,13 +440,19 @@ impl Agent {
// Check if the tool is read_resource and handle it separately
ToolCallResult::from(
extension_manager
.read_resource(tool_call.arguments.clone())
.read_resource(
tool_call.arguments.clone(),
cancellation_token.unwrap_or_default(),
)
.await,
)
} else if tool_call.name == PLATFORM_LIST_RESOURCES_TOOL_NAME {
ToolCallResult::from(
extension_manager
.list_resources(tool_call.arguments.clone())
.list_resources(
tool_call.arguments.clone(),
cancellation_token.unwrap_or_default(),
)
.await,
)
} else if tool_call.name == PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME {
@@ -470,7 +476,7 @@ impl Agent {
} else {
// Clone the result to ensure no references to extension_manager are returned
let result = extension_manager
.dispatch_tool_call(tool_call.clone())
.dispatch_tool_call(tool_call.clone(), cancellation_token.unwrap_or_default())
.await;
result.unwrap_or_else(|e| {
ToolCallResult::from(Err(ToolError::ExecutionError(e.to_string())))
@@ -1181,7 +1187,7 @@ impl Agent {
pub async fn list_extension_prompts(&self) -> HashMap<String, Vec<Prompt>> {
let extension_manager = self.extension_manager.read().await;
extension_manager
.list_prompts()
.list_prompts(CancellationToken::default())
.await
.expect("Failed to list prompts")
}
@@ -1191,7 +1197,7 @@ impl Agent {
// First find which extension has this prompt
let prompts = extension_manager
.list_prompts()
.list_prompts(CancellationToken::default())
.await
.map_err(|e| anyhow!("Failed to list prompts: {}", e))?;
@@ -1201,7 +1207,7 @@ impl Agent {
.map(|(extension, _)| extension)
{
return extension_manager
.get_prompt(extension, name, arguments)
.get_prompt(extension, name, arguments, CancellationToken::default())
.await
.map_err(|e| anyhow!("Failed to get prompt: {}", e));
}
crates/goose/src/agents/extension_manager.rs
+94 -77
View File
@@ -1,6 +1,6 @@
use anyhow::Result;
use axum::http::{HeaderMap, HeaderName};
use chrono::{DateTime, TimeZone, Utc};
use chrono::{DateTime, Utc};
use futures::stream::{FuturesUnordered, StreamExt};
use futures::{future, FutureExt};
use mcp_core::handler::require_str_parameter;
@@ -13,7 +13,6 @@ use rmcp::transport::{
use std::collections::{HashMap, HashSet};
use std::process::Stdio;
use std::sync::Arc;
use std::sync::LazyLock;
use std::time::Duration;
use tempfile::tempdir;
use tokio::io::AsyncReadExt;
@@ -21,6 +20,7 @@ use tokio::process::Command;
use tokio::sync::Mutex;
use tokio::task;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use tracing::{error, warn};
use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, ToolInfo};
@@ -29,14 +29,9 @@ use crate::agents::extension::{Envs, ProcessExit};
use crate::config::{Config, ExtensionConfigManager};
use crate::prompt_template;
use mcp_client::client::{McpClient, McpClientTrait};
use rmcp::model::{Content, GetPromptResult, Prompt, Resource, ResourceContents, Tool};
use rmcp::model::{Content, GetPromptResult, Prompt, ResourceContents, Tool};
use serde_json::Value;
// By default, we set it to Jan 1, 2020 if the resource does not have a timestamp
// This is to ensure that the resource is considered less important than resources with a more recent timestamp
static DEFAULT_TIMESTAMP: LazyLock<DateTime<Utc>> =
LazyLock::new(|| Utc.with_ymd_and_hms(2020, 1, 1, 0, 0, 0).unwrap());
type McpClientBox = Arc<Mutex<Box<dyn McpClientTrait>>>;
/// Manages Goose extensions / MCP clients and their interactions
@@ -457,7 +452,9 @@ impl ExtensionManager {
task::spawn(async move {
let mut tools = Vec::new();
let client_guard = client.lock().await;
let mut client_tools = client_guard.list_tools(None).await?;
let mut client_tools = client_guard
.list_tools(None, CancellationToken::default())
.await?;
loop {
for tool in client_tools.tools {
@@ -474,7 +471,9 @@ impl ExtensionManager {
break;
}
client_tools = client_guard.list_tools(client_tools.next_cursor).await?;
client_tools = client_guard
.list_tools(client_tools.next_cursor, CancellationToken::default())
.await?;
}
Ok::<Vec<Tool>, ExtensionError>(tools)
@@ -497,43 +496,6 @@ impl ExtensionManager {
Ok(tools)
}
/// Get client resources and their contents
pub async fn get_resources(&self) -> ExtensionResult<Vec<ResourceItem>> {
let mut result: Vec<ResourceItem> = Vec::new();
for (name, client) in &self.clients {
let client_guard = client.lock().await;
let resources = client_guard.list_resources(None).await?;
for resource in resources.resources {
// Skip reading the resource if it's not marked active
// This avoids blowing up the context with inactive resources
if !resource_is_active(&resource) {
continue;
}
if let Ok(contents) = client_guard.read_resource(&resource.uri).await {
for content in contents.contents {
let (uri, content_str) = match content {
ResourceContents::TextResourceContents { uri, text, .. } => (uri, text),
ResourceContents::BlobResourceContents { uri, blob, .. } => (uri, blob),
};
result.push(ResourceItem::new(
name.clone(),
uri,
resource.name.clone(),
content_str,
resource.timestamp().unwrap_or(*DEFAULT_TIMESTAMP),
resource.priority().unwrap_or(0.0),
));
}
}
}
}
Ok(result)
}
/// Get the extension prompt including client instructions
pub async fn get_planning_prompt(&self, tools_info: Vec<ToolInfo>) -> String {
let mut context: HashMap<&str, Value> = HashMap::new();
@@ -551,14 +513,22 @@ impl ExtensionManager {
}
// Function that gets executed for read_resource tool
pub async fn read_resource(&self, params: Value) -> Result<Vec<Content>, ToolError> {
pub async fn read_resource(
&self,
params: Value,
cancellation_token: CancellationToken,
) -> Result<Vec<Content>, ToolError> {
let uri = require_str_parameter(&params, "uri")?;
let extension_name = params.get("extension_name").and_then(|v| v.as_str());
// If extension name is provided, we can just look it up
if extension_name.is_some() {
let result = self
.read_resource_from_extension(uri, extension_name.unwrap())
.read_resource_from_extension(
uri,
extension_name.unwrap(),
cancellation_token.clone(),
)
.await?;
return Ok(result);
}
@@ -568,7 +538,9 @@ impl ExtensionManager {
// TODO: do we want to find if a provided uri is in multiple extensions?
// currently it will return the first match and skip any others
for extension_name in self.resource_capable_extensions.iter() {
let result = self.read_resource_from_extension(uri, extension_name).await;
let result = self
.read_resource_from_extension(uri, extension_name, cancellation_token.clone())
.await;
match result {
Ok(result) => return Ok(result),
Err(_) => continue,
@@ -594,6 +566,7 @@ impl ExtensionManager {
&self,
uri: &str,
extension_name: &str,
cancellation_token: CancellationToken,
) -> Result<Vec<Content>, ToolError> {
let available_extensions = self
.clients
@@ -612,9 +585,12 @@ impl ExtensionManager {
.ok_or(ToolError::InvalidParameters(error_msg))?;
let client_guard = client.lock().await;
let read_result = client_guard.read_resource(uri).await.map_err(|_| {
ToolError::ExecutionError(format!("Could not read resource with uri: {}", uri))
})?;
let read_result = client_guard
.read_resource(uri, cancellation_token)
.await
.map_err(|_| {
ToolError::ExecutionError(format!("Could not read resource with uri: {}", uri))
})?;
let mut result = Vec::new();
for content in read_result.contents {
@@ -631,6 +607,7 @@ impl ExtensionManager {
async fn list_resources_from_extension(
&self,
extension_name: &str,
cancellation_token: CancellationToken,
) -> Result<Vec<Content>, ToolError> {
let client = self.clients.get(extension_name).ok_or_else(|| {
ToolError::InvalidParameters(format!("Extension {} is not valid", extension_name))
@@ -638,7 +615,7 @@ impl ExtensionManager {
let client_guard = client.lock().await;
client_guard
.list_resources(None)
.list_resources(None, cancellation_token)
.await
.map_err(|e| {
ToolError::ExecutionError(format!(
@@ -658,13 +635,18 @@ impl ExtensionManager {
})
}
pub async fn list_resources(&self, params: Value) -> Result<Vec<Content>, ToolError> {
pub async fn list_resources(
&self,
params: Value,
cancellation_token: CancellationToken,
) -> Result<Vec<Content>, ToolError> {
let extension = params.get("extension").and_then(|v| v.as_str());
match extension {
Some(extension_name) => {
// Handle single extension case
self.list_resources_from_extension(extension_name).await
self.list_resources_from_extension(extension_name, cancellation_token)
.await
}
None => {
// Handle all extensions case using FuturesUnordered
@@ -672,8 +654,10 @@ impl ExtensionManager {
// Create futures for each resource_capable_extension
for extension_name in &self.resource_capable_extensions {
let token = cancellation_token.clone();
futures.push(async move {
self.list_resources_from_extension(extension_name).await
self.list_resources_from_extension(extension_name, token)
.await
});
}
@@ -708,7 +692,11 @@ impl ExtensionManager {
}
}
pub async fn dispatch_tool_call(&self, tool_call: ToolCall) -> Result<ToolCallResult> {
pub async fn dispatch_tool_call(
&self,
tool_call: ToolCall,
cancellation_token: CancellationToken,
) -> Result<ToolCallResult> {
// Dispatch tool call based on the prefix naming convention
let (client_name, client) = self
.get_client_for_tool(&tool_call.name)
@@ -729,7 +717,7 @@ impl ExtensionManager {
let fut = async move {
let client_guard = client.lock().await;
client_guard
.call_tool(&tool_name, arguments)
.call_tool(&tool_name, arguments, cancellation_token)
.await
.map(|call| call.content)
.map_err(|e| ToolError::ExecutionError(e.to_string()))
@@ -744,6 +732,7 @@ impl ExtensionManager {
pub async fn list_prompts_from_extension(
&self,
extension_name: &str,
cancellation_token: CancellationToken,
) -> Result<Vec<Prompt>, ToolError> {
let client = self.clients.get(extension_name).ok_or_else(|| {
ToolError::InvalidParameters(format!("Extension {} is not valid", extension_name))
@@ -751,7 +740,7 @@ impl ExtensionManager {
let client_guard = client.lock().await;
client_guard
.list_prompts(None)
.list_prompts(None, cancellation_token)
.await
.map_err(|e| {
ToolError::ExecutionError(format!(
@@ -762,14 +751,19 @@ impl ExtensionManager {
.map(|lp| lp.prompts)
}
pub async fn list_prompts(&self) -> Result<HashMap<String, Vec<Prompt>>, ToolError> {
pub async fn list_prompts(
&self,
cancellation_token: CancellationToken,
) -> Result<HashMap<String, Vec<Prompt>>, ToolError> {
let mut futures = FuturesUnordered::new();
for extension_name in self.clients.keys() {
let token = cancellation_token.clone();
futures.push(async move {
(
extension_name,
self.list_prompts_from_extension(extension_name).await,
self.list_prompts_from_extension(extension_name, token)
.await,
)
});
}
@@ -809,6 +803,7 @@ impl ExtensionManager {
extension_name: &str,
name: &str,
arguments: Value,
cancellation_token: CancellationToken,
) -> Result<GetPromptResult> {
let client = self
.clients
@@ -817,7 +812,7 @@ impl ExtensionManager {
let client_guard = client.lock().await;
client_guard
.get_prompt(name, arguments)
.get_prompt(name, arguments, cancellation_token)
.await
.map_err(|e| anyhow::anyhow!("Failed to get prompt: {}", e))
}
@@ -896,10 +891,6 @@ impl ExtensionManager {
}
}
fn resource_is_active(resource: &Resource) -> bool {
resource.priority().is_some_and(|p| (p - 1.0).abs() < 1e-6)
}
#[cfg(test)]
mod tests {
use super::*;
@@ -927,19 +918,33 @@ mod tests {
async fn list_resources(
&self,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListResourcesResult, Error> {
Err(Error::TransportClosed)
}
async fn read_resource(&self, _uri: &str) -> Result<ReadResourceResult, Error> {
async fn read_resource(
&self,
_uri: &str,
_cancellation_token: CancellationToken,
) -> Result<ReadResourceResult, Error> {
Err(Error::TransportClosed)
}
async fn list_tools(&self, _next_cursor: Option<String>) -> Result<ListToolsResult, Error> {
async fn list_tools(
&self,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
Err(Error::TransportClosed)
}
async fn call_tool(&self, name: &str, _arguments: Value) -> Result<CallToolResult, Error> {
async fn call_tool(
&self,
name: &str,
_arguments: Value,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
match name {
"tool" | "test__tool" => Ok(CallToolResult {
content: vec![],
@@ -952,6 +957,7 @@ mod tests {
async fn list_prompts(
&self,
_next_cursor: Option<String>,
_cancellation_token: CancellationToken,
) -> Result<ListPromptsResult, Error> {
Err(Error::TransportClosed)
}
@@ -960,6 +966,7 @@ mod tests {
&self,
_name: &str,
_arguments: Value,
_cancellation_token: CancellationToken,
) -> Result<GetPromptResult, Error> {
Err(Error::TransportClosed)
}
@@ -1043,7 +1050,9 @@ mod tests {
arguments: json!({}),
};
let result = extension_manager.dispatch_tool_call(tool_call).await;
let result = extension_manager
.dispatch_tool_call(tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
let tool_call = ToolCall {
@@ -1051,7 +1060,9 @@ mod tests {
arguments: json!({}),
};
let result = extension_manager.dispatch_tool_call(tool_call).await;
let result = extension_manager
.dispatch_tool_call(tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
// verify a multiple underscores dispatch
@@ -1060,7 +1071,9 @@ mod tests {
arguments: json!({}),
};
let result = extension_manager.dispatch_tool_call(tool_call).await;
let result = extension_manager
.dispatch_tool_call(tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
// Test unicode in tool name, "client 🚀" should become "client_"
@@ -1069,7 +1082,9 @@ mod tests {
arguments: json!({}),
};
let result = extension_manager.dispatch_tool_call(tool_call).await;
let result = extension_manager
.dispatch_tool_call(tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
let tool_call = ToolCall {
@@ -1077,7 +1092,9 @@ mod tests {
arguments: json!({}),
};
let result = extension_manager.dispatch_tool_call(tool_call).await;
let result = extension_manager
.dispatch_tool_call(tool_call, CancellationToken::default())
.await;
assert!(result.is_ok());
// this should error out, specifically for an ToolError::ExecutionError
@@ -1087,7 +1104,7 @@ mod tests {
};
let result = extension_manager
.dispatch_tool_call(invalid_tool_call)
.dispatch_tool_call(invalid_tool_call, CancellationToken::default())
.await
.unwrap()
.result
@@ -1105,7 +1122,7 @@ mod tests {
};
let result = extension_manager
.dispatch_tool_call(invalid_tool_call)
.dispatch_tool_call(invalid_tool_call, CancellationToken::default())
.await;
if let Err(err) = result {
let tool_err = err.downcast_ref::<ToolError>().expect("Expected ToolError");
crates/goose/src/agents/subagent.rs
+2 -1
View File
@@ -15,6 +15,7 @@ use serde::{Deserialize, Serialize};
// use serde_json::{self};
use std::{collections::HashMap, sync::Arc};
use tokio::sync::{Mutex, RwLock};
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, instrument};
/// Status of a subagent
@@ -197,7 +198,7 @@ impl SubAgent {
.extension_manager
.read()
.await
.dispatch_tool_call(tool_call.clone())
.dispatch_tool_call(tool_call.clone(), CancellationToken::default())
.await
{
Ok(result) => result.result.await,