timeouts for mcp tool calls (#3959)

defaults to 60sec, overridable with MCP_TOOL_TIMEOUT or on a per-server
basis in the config.
This commit is contained in:
Jeremy Rose
2025-09-22 10:30:59 -07:00
committed by GitHub
parent e258ca61b4
commit 19f46439ae
7 changed files with 166 additions and 50 deletions

View File

@@ -40,6 +40,9 @@ const MAX_TOOL_NAME_LENGTH: usize = 64;
/// Default timeout for initializing MCP server & initially listing tools.
const DEFAULT_STARTUP_TIMEOUT: Duration = Duration::from_secs(10);
/// Default timeout for individual tool calls.
const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(60);
/// Map that holds a startup error for every MCP server that could **not** be
/// spawned successfully.
pub type ClientStartErrors = HashMap<String, anyhow::Error>;
@@ -85,6 +88,7 @@ struct ToolInfo {
struct ManagedClient {
client: Arc<McpClient>,
startup_timeout: Duration,
tool_timeout: Option<Duration>,
}
/// A thin wrapper around a set of running [`McpClient`] instances.
@@ -132,10 +136,9 @@ impl McpConnectionManager {
continue;
}
let startup_timeout = cfg
.startup_timeout_ms
.map(Duration::from_millis)
.unwrap_or(DEFAULT_STARTUP_TIMEOUT);
let startup_timeout = cfg.startup_timeout_sec.unwrap_or(DEFAULT_STARTUP_TIMEOUT);
let tool_timeout = cfg.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT);
join_set.spawn(async move {
let McpServerConfig {
@@ -171,19 +174,19 @@ impl McpConnectionManager {
protocol_version: mcp_types::MCP_SCHEMA_VERSION.to_owned(),
};
let initialize_notification_params = None;
match client
let init_result = client
.initialize(
params,
initialize_notification_params,
Some(startup_timeout),
)
.await
{
Ok(_response) => (server_name, Ok((client, startup_timeout))),
Err(e) => (server_name, Err(e)),
}
.await;
(
(server_name, tool_timeout),
init_result.map(|_| (client, startup_timeout)),
)
}
Err(e) => (server_name, Err(e.into())),
Err(e) => ((server_name, tool_timeout), Err(e.into())),
}
});
}
@@ -191,8 +194,8 @@ impl McpConnectionManager {
let mut clients: HashMap<String, ManagedClient> = HashMap::with_capacity(join_set.len());
while let Some(res) = join_set.join_next().await {
let (server_name, client_res) = match res {
Ok((server_name, client_res)) => (server_name, client_res),
let ((server_name, tool_timeout), client_res) = match res {
Ok(result) => result,
Err(e) => {
warn!("Task panic when starting MCP server: {e:#}");
continue;
@@ -206,6 +209,7 @@ impl McpConnectionManager {
ManagedClient {
client: Arc::new(client),
startup_timeout,
tool_timeout: Some(tool_timeout),
},
);
}
@@ -243,14 +247,13 @@ impl McpConnectionManager {
server: &str,
tool: &str,
arguments: Option<serde_json::Value>,
timeout: Option<Duration>,
) -> Result<mcp_types::CallToolResult> {
let client = self
let managed = self
.clients
.get(server)
.ok_or_else(|| anyhow!("unknown MCP server '{server}'"))?
.client
.clone();
.ok_or_else(|| anyhow!("unknown MCP server '{server}'"))?;
let client = managed.client.clone();
let timeout = managed.tool_timeout;
client
.call_tool(tool.to_string(), arguments, timeout)