Add configurable timeout for context server tool calls (#33348)

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

@@ -2477,6 +2477,7 @@ fn setup_context_server(
                     path: "somebinary".into(),
                     args: Vec::new(),
                     env: None,
+                    timeout: None,
                 },
             },
         );

crates/context_server/src/client.rs

@@ -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>(

crates/context_server/src/context_server.rs

@@ -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(),
             )?,
         })

crates/project/src/context_server_store.rs

@@ -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,
             }))
         }
 

crates/project/src/context_server_store/extension.rs

@@ -69,6 +69,7 @@ impl registry::ContextServerDescriptor for ContextServerDescriptor {
                 path: command.command,
                 args: command.args,
                 env: Some(command.env.into_iter().collect()),
+                timeout: None,
             })
         })
     }

crates/project/src/project_settings.rs

@@ -604,6 +604,7 @@ impl Settings for ProjectSettings {
                     path: cmd.command,
                     args: cmd.args.unwrap_or_default(),
                     env: cmd.env,
+                    timeout: None,
                 }
             }
         }