chore: add timeout to tool use requests (#5143)

This commit is contained in:
Louis 2025-05-29 23:45:55 +07:00 committed by GitHub
parent ec7d0c6abf
commit 2b7b9c2dc6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,11 +1,11 @@
use std::{collections::HashMap, env, sync::Arc}; use std::{collections::HashMap, env, sync::Arc, time::Duration};
use rmcp::model::{CallToolRequestParam, CallToolResult, Tool}; use rmcp::model::{CallToolRequestParam, CallToolResult, Tool};
use rmcp::{service::RunningService, transport::TokioChildProcess, RoleClient, ServiceExt}; use rmcp::{service::RunningService, transport::TokioChildProcess, RoleClient, ServiceExt};
use serde_json::{Map, Value}; use serde_json::{Map, Value};
use std::fs;
use tauri::{AppHandle, Emitter, Runtime, State}; use tauri::{AppHandle, Emitter, Runtime, State};
use tokio::{process::Command, sync::Mutex}; use tokio::{process::Command, sync::Mutex, time::timeout};
use std::{fs};
use super::{cmd::get_jan_data_folder_path, state::AppState}; use super::{cmd::get_jan_data_folder_path, state::AppState};
@ -13,6 +13,8 @@ const DEFAULT_MCP_CONFIG: &str = r#"{
"mcpServers": {} "mcpServers": {}
}"#; }"#;
// Timeout for MCP tool calls (30 seconds)
const MCP_TOOL_CALL_TIMEOUT: Duration = Duration::from_secs(30);
/// Runs MCP commands by reading configuration from a JSON file and initializing servers /// Runs MCP commands by reading configuration from a JSON file and initializing servers
/// ///
@ -202,8 +204,18 @@ pub async fn get_tools(state: State<'_, AppState>) -> Result<Vec<Tool>, String>
let mut all_tools: Vec<Tool> = Vec::new(); let mut all_tools: Vec<Tool> = Vec::new();
for (_, service) in servers.iter() { for (_, service) in servers.iter() {
// List tools // List tools with timeout
let tools = service.list_all_tools().await.map_err(|e| e.to_string())?; let tools_future = service.list_all_tools();
let tools = match timeout(MCP_TOOL_CALL_TIMEOUT, tools_future).await {
Ok(result) => result.map_err(|e| e.to_string())?,
Err(_) => {
log::warn!(
"Listing tools timed out after {} seconds",
MCP_TOOL_CALL_TIMEOUT.as_secs()
);
continue; // Skip this server and continue with others
}
};
for tool in tools { for tool in tools {
all_tools.push(tool); all_tools.push(tool);
@ -240,13 +252,22 @@ pub async fn call_tool(
for (_, service) in servers.iter() { for (_, service) in servers.iter() {
if let Ok(tools) = service.list_all_tools().await { if let Ok(tools) = service.list_all_tools().await {
if tools.iter().any(|t| t.name == tool_name) { if tools.iter().any(|t| t.name == tool_name) {
return service println!("Found tool {} in server", tool_name);
.call_tool(CallToolRequestParam {
name: tool_name.into(), // Call the tool with timeout
let tool_call = service.call_tool(CallToolRequestParam {
name: tool_name.clone().into(),
arguments, arguments,
}) });
.await
.map_err(|e| e.to_string()); return match timeout(MCP_TOOL_CALL_TIMEOUT, tool_call).await {
Ok(result) => result.map_err(|e| e.to_string()),
Err(_) => Err(format!(
"Tool call '{}' timed out after {} seconds",
tool_name,
MCP_TOOL_CALL_TIMEOUT.as_secs()
)),
};
} }
} }
} }