Rafaล Krzywaลผnia and Ben Brandt created
Closes: #32668
- Add
[tool_call_timeout_millis](https://github.com/cline/cline/pull/1904)
field to ContextServerCommand, like in Cline
- Update ModelContextServerBinary to include timeout configuration
- Modify Client to store and use configurable request timeout
- Replace hardcoded REQUEST_TIMEOUT with self.request_timeout
- Rename REQUEST_TIMEOUT to DEFAULT_REQUEST_TIMEOUT for clarity
- Maintain backward compatibility with 60-second default
Release Notes:
- context_server: Add support for configurable timeout for MCP tool
calls
---------
Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>crates/agent2/src/tests/mod.rs | 1
crates/context_server/src/client.rs | 18 ++++++++++---
crates/context_server/src/context_server.rs | 4 +++
crates/project/src/context_server_store.rs | 7 +++++
crates/project/src/context_server_store/extension.rs | 1
crates/project/src/project_settings.rs | 1
6 files changed, 28 insertions(+), 4 deletions(-)
@@ -2477,6 +2477,7 @@ fn setup_context_server(
path: "somebinary".into(),
args: Vec::new(),
env: None,
+ timeout: None,
},
},
);
@@ -25,7 +25,7 @@ use crate::{
};
const JSON_RPC_VERSION: &str = "2.0";
-const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
+const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
// Standard JSON-RPC error codes
pub const PARSE_ERROR: i32 = -32700;
@@ -60,6 +60,7 @@ pub(crate) struct Client {
executor: BackgroundExecutor,
#[allow(dead_code)]
transport: Arc<dyn Transport>,
+ request_timeout: Option<Duration>,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -143,6 +144,7 @@ pub struct ModelContextServerBinary {
pub executable: PathBuf,
pub args: Vec<String>,
pub env: Option<HashMap<String, String>>,
+ pub timeout: Option<u64>,
}
impl Client {
@@ -169,8 +171,9 @@ impl Client {
.map(|name| name.to_string_lossy().to_string())
.unwrap_or_else(String::new);
+ let timeout = binary.timeout.map(Duration::from_millis);
let transport = Arc::new(StdioTransport::new(binary, working_directory, &cx)?);
- Self::new(server_id, server_name.into(), transport, cx)
+ Self::new(server_id, server_name.into(), transport, timeout, cx)
}
/// Creates a new Client instance for a context server.
@@ -178,6 +181,7 @@ impl Client {
server_id: ContextServerId,
server_name: Arc<str>,
transport: Arc<dyn Transport>,
+ request_timeout: Option<Duration>,
cx: AsyncApp,
) -> Result<Self> {
let (outbound_tx, outbound_rx) = channel::unbounded::<String>();
@@ -237,6 +241,7 @@ impl Client {
io_tasks: Mutex::new(Some((input_task, output_task))),
output_done_rx: Mutex::new(Some(output_done_rx)),
transport,
+ request_timeout,
})
}
@@ -327,8 +332,13 @@ impl Client {
method: &str,
params: impl Serialize,
) -> Result<T> {
- self.request_with(method, params, None, Some(REQUEST_TIMEOUT))
- .await
+ self.request_with(
+ method,
+ params,
+ None,
+ self.request_timeout.or(Some(DEFAULT_REQUEST_TIMEOUT)),
+ )
+ .await
}
pub async fn request_with<T: DeserializeOwned>(
@@ -34,6 +34,8 @@ pub struct ContextServerCommand {
pub path: PathBuf,
pub args: Vec<String>,
pub env: Option<HashMap<String, String>>,
+ /// Timeout for tool calls in milliseconds. Defaults to 60000 (60 seconds) if not specified.
+ pub timeout: Option<u64>,
}
impl std::fmt::Debug for ContextServerCommand {
@@ -123,6 +125,7 @@ impl ContextServer {
executable: Path::new(&command.path).to_path_buf(),
args: command.args.clone(),
env: command.env.clone(),
+ timeout: command.timeout,
},
working_directory,
cx.clone(),
@@ -131,6 +134,7 @@ impl ContextServer {
client::ContextServerId(self.id.0.clone()),
self.id().0,
transport.clone(),
+ None,
cx.clone(),
)?,
})
@@ -976,6 +976,7 @@ mod tests {
path: "somebinary".into(),
args: vec!["arg".to_string()],
env: None,
+ timeout: None,
},
},
),
@@ -1016,6 +1017,7 @@ mod tests {
path: "somebinary".into(),
args: vec!["anotherArg".to_string()],
env: None,
+ timeout: None,
},
},
),
@@ -1098,6 +1100,7 @@ mod tests {
path: "somebinary".into(),
args: vec!["arg".to_string()],
env: None,
+ timeout: None,
},
},
)],
@@ -1150,6 +1153,7 @@ mod tests {
path: "somebinary".into(),
args: vec!["arg".to_string()],
env: None,
+ timeout: None,
},
},
)],
@@ -1177,6 +1181,7 @@ mod tests {
command: ContextServerCommand {
path: "somebinary".into(),
args: vec!["arg".to_string()],
+ timeout: None,
env: None,
},
},
@@ -1230,6 +1235,7 @@ mod tests {
path: "somebinary".into(),
args: vec!["arg".to_string()],
env: None,
+ timeout: None,
},
}
}
@@ -1318,6 +1324,7 @@ mod tests {
path: self.path.clone(),
args: vec!["arg1".to_string(), "arg2".to_string()],
env: None,
+ timeout: None,
}))
}
@@ -69,6 +69,7 @@ impl registry::ContextServerDescriptor for ContextServerDescriptor {
path: command.command,
args: command.args,
env: Some(command.env.into_iter().collect()),
+ timeout: None,
})
})
}
@@ -604,6 +604,7 @@ impl Settings for ProjectSettings {
path: cmd.command,
args: cmd.args.unwrap_or_default(),
env: cmd.env,
+ timeout: None,
}
}
}