feat(mcp): add toolCallTimeoutMs, timeout/reconnect/error handling

- Add toolCallTimeoutMs to stdio MCP config with 60s default
- tools/call runs under timeout with dedicated Timeout error
- Handle malformed JSON/broken protocol as InvalidResponse
- Reset/reconnect stdio state on child exit or transport drop
- Add tests: slow timeout, invalid JSON response, stdio reconnect
- Verified: cargo test -p runtime 113 passed, clippy clean
This commit is contained in:
YeonGyu-Kim
2026-04-02 18:24:30 +09:00
parent 6e4b0123a6
commit 3b18ce9f3f
4 changed files with 866 additions and 135 deletions

View File

@@ -106,6 +106,7 @@ pub struct McpStdioServerConfig {
pub command: String, pub command: String,
pub args: Vec<String>, pub args: Vec<String>,
pub env: BTreeMap<String, String>, pub env: BTreeMap<String, String>,
pub tool_call_timeout_ms: Option<u64>,
} }
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
@@ -791,6 +792,7 @@ fn parse_mcp_server_config(
command: expect_string(object, "command", context)?.to_string(), command: expect_string(object, "command", context)?.to_string(),
args: optional_string_array(object, "args", context)?.unwrap_or_default(), args: optional_string_array(object, "args", context)?.unwrap_or_default(),
env: optional_string_map(object, "env", context)?.unwrap_or_default(), env: optional_string_map(object, "env", context)?.unwrap_or_default(),
tool_call_timeout_ms: optional_u64(object, "toolCallTimeoutMs", context)?,
})), })),
"sse" => Ok(McpServerConfig::Sse(parse_mcp_remote_server_config( "sse" => Ok(McpServerConfig::Sse(parse_mcp_remote_server_config(
object, context, object, context,
@@ -914,6 +916,27 @@ fn optional_u16(
} }
} }
fn optional_u64(
object: &BTreeMap<String, JsonValue>,
key: &str,
context: &str,
) -> Result<Option<u64>, ConfigError> {
match object.get(key) {
Some(value) => {
let Some(number) = value.as_i64() else {
return Err(ConfigError::Parse(format!(
"{context}: field {key} must be a non-negative integer"
)));
};
let number = u64::try_from(number).map_err(|_| {
ConfigError::Parse(format!("{context}: field {key} is out of range"))
})?;
Ok(Some(number))
}
None => Ok(None),
}
}
fn parse_bool_map(value: &JsonValue, context: &str) -> Result<BTreeMap<String, bool>, ConfigError> { fn parse_bool_map(value: &JsonValue, context: &str) -> Result<BTreeMap<String, bool>, ConfigError> {
let Some(map) = value.as_object() else { let Some(map) = value.as_object() else {
return Err(ConfigError::Parse(format!( return Err(ConfigError::Parse(format!(

View File

@@ -84,10 +84,13 @@ pub fn mcp_server_signature(config: &McpServerConfig) -> Option<String> {
pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String { pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String {
let rendered = match &config.config { let rendered = match &config.config {
McpServerConfig::Stdio(stdio) => format!( McpServerConfig::Stdio(stdio) => format!(
"stdio|{}|{}|{}", "stdio|{}|{}|{}|{}",
stdio.command, stdio.command,
render_command_signature(&stdio.args), render_command_signature(&stdio.args),
render_env_signature(&stdio.env) render_env_signature(&stdio.env),
stdio
.tool_call_timeout_ms
.map_or_else(String::new, |timeout_ms| timeout_ms.to_string())
), ),
McpServerConfig::Sse(remote) => format!( McpServerConfig::Sse(remote) => format!(
"sse|{}|{}|{}|{}", "sse|{}|{}|{}|{}",
@@ -245,6 +248,7 @@ mod tests {
command: "uvx".to_string(), command: "uvx".to_string(),
args: vec!["mcp-server".to_string()], args: vec!["mcp-server".to_string()],
env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]), env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]),
tool_call_timeout_ms: None,
}); });
assert_eq!( assert_eq!(
mcp_server_signature(&stdio), mcp_server_signature(&stdio),

View File

@@ -3,6 +3,8 @@ use std::collections::BTreeMap;
use crate::config::{McpOAuthConfig, McpServerConfig, ScopedMcpServerConfig}; use crate::config::{McpOAuthConfig, McpServerConfig, ScopedMcpServerConfig};
use crate::mcp::{mcp_server_signature, mcp_tool_prefix, normalize_name_for_mcp}; use crate::mcp::{mcp_server_signature, mcp_tool_prefix, normalize_name_for_mcp};
pub const DEFAULT_MCP_TOOL_CALL_TIMEOUT_MS: u64 = 60_000;
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
pub enum McpClientTransport { pub enum McpClientTransport {
Stdio(McpStdioTransport), Stdio(McpStdioTransport),
@@ -18,6 +20,7 @@ pub struct McpStdioTransport {
pub command: String, pub command: String,
pub args: Vec<String>, pub args: Vec<String>,
pub env: BTreeMap<String, String>, pub env: BTreeMap<String, String>,
pub tool_call_timeout_ms: Option<u64>,
} }
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
@@ -75,6 +78,7 @@ impl McpClientTransport {
command: config.command.clone(), command: config.command.clone(),
args: config.args.clone(), args: config.args.clone(),
env: config.env.clone(), env: config.env.clone(),
tool_call_timeout_ms: config.tool_call_timeout_ms,
}), }),
McpServerConfig::Sse(config) => Self::Sse(McpRemoteTransport { McpServerConfig::Sse(config) => Self::Sse(McpRemoteTransport {
url: config.url.clone(), url: config.url.clone(),
@@ -105,6 +109,14 @@ impl McpClientTransport {
} }
} }
impl McpStdioTransport {
#[must_use]
pub fn resolved_tool_call_timeout_ms(&self) -> u64 {
self.tool_call_timeout_ms
.unwrap_or(DEFAULT_MCP_TOOL_CALL_TIMEOUT_MS)
}
}
impl McpClientAuth { impl McpClientAuth {
#[must_use] #[must_use]
pub fn from_oauth(oauth: Option<McpOAuthConfig>) -> Self { pub fn from_oauth(oauth: Option<McpOAuthConfig>) -> Self {
@@ -136,6 +148,7 @@ mod tests {
command: "uvx".to_string(), command: "uvx".to_string(),
args: vec!["mcp-server".to_string()], args: vec!["mcp-server".to_string()],
env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]), env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]),
tool_call_timeout_ms: Some(15_000),
}), }),
}; };
@@ -154,6 +167,7 @@ mod tests {
transport.env.get("TOKEN").map(String::as_str), transport.env.get("TOKEN").map(String::as_str),
Some("secret") Some("secret")
); );
assert_eq!(transport.tool_call_timeout_ms, Some(15_000));
} }
other => panic!("expected stdio transport, got {other:?}"), other => panic!("expected stdio transport, got {other:?}"),
} }

File diff suppressed because it is too large Load Diff