diff --git a/crates/agent2/src/tests/mod.rs b/crates/agent2/src/tests/mod.rs index 9132c9a316e7e5122d24a0b413ae700a7adbaba9..884580ed69009d168b3266870acf4f698a2f5450 100644 --- a/crates/agent2/src/tests/mod.rs +++ b/crates/agent2/src/tests/mod.rs @@ -2477,6 +2477,7 @@ fn setup_context_server( path: "somebinary".into(), args: Vec::new(), env: None, + timeout: None, }, }, ); diff --git a/crates/context_server/src/client.rs b/crates/context_server/src/client.rs index 03cf047ac5273071354756119864ab6914e524c6..b3b44dbde67d92ce620d85a39a0925f27a4e2086 100644 --- a/crates/context_server/src/client.rs +++ b/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, + request_timeout: Option, } #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -143,6 +144,7 @@ pub struct ModelContextServerBinary { pub executable: PathBuf, pub args: Vec, pub env: Option>, + pub timeout: Option, } 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, transport: Arc, + request_timeout: Option, cx: AsyncApp, ) -> Result { let (outbound_tx, outbound_rx) = channel::unbounded::(); @@ -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 { - 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( diff --git a/crates/context_server/src/context_server.rs b/crates/context_server/src/context_server.rs index 9ca78138dbc229a6aedd5c53c460d9205502df94..b126bb393784664692b5de39fee5ed7f66e9948a 100644 --- a/crates/context_server/src/context_server.rs +++ b/crates/context_server/src/context_server.rs @@ -34,6 +34,8 @@ pub struct ContextServerCommand { pub path: PathBuf, pub args: Vec, pub env: Option>, + /// Timeout for tool calls in milliseconds. Defaults to 60000 (60 seconds) if not specified. + pub timeout: Option, } 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(), )?, }) diff --git a/crates/project/src/context_server_store.rs b/crates/project/src/context_server_store.rs index 49a430c26110fc58a9494d414ecbcf45a6c76c49..20188df5c4ae38b2ae305daee5b3eecc25319951 100644 --- a/crates/project/src/context_server_store.rs +++ b/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, })) } diff --git a/crates/project/src/context_server_store/extension.rs b/crates/project/src/context_server_store/extension.rs index 2a3a0c2e4b99e56c66993d0db1fbec5b3fb9ef29..ca5cacf3b549523dee8b85242bea86653eecbf7a 100644 --- a/crates/project/src/context_server_store/extension.rs +++ b/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, }) }) } diff --git a/crates/project/src/project_settings.rs b/crates/project/src/project_settings.rs index 4a97130f15d582df392c25d6d64482bc4ca17834..40874638111eb3b85d4f65ad0b531072ab082624 100644 --- a/crates/project/src/project_settings.rs +++ b/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, } } }