Skip to content

Commit 15a489c

Browse files
authored
adds server load timeout for no interactive mode (#1851)
1 parent 3508012 commit 15a489c

File tree

4 files changed

+33
-7
lines changed

4 files changed

+33
-7
lines changed

crates/chat-cli/src/cli/chat/conversation_state.rs

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1049,12 +1049,13 @@ mod tests {
10491049
#[tokio::test]
10501050
async fn test_conversation_state_history_handling_truncation() {
10511051
let mut database = Database::new().await.unwrap();
1052+
let mut output = SharedWriter::null();
10521053

10531054
let mut tool_manager = ToolManager::default();
10541055
let mut conversation_state = ConversationState::new(
10551056
Context::new(),
10561057
"fake_conv_id",
1057-
tool_manager.load_tools(&database).await.unwrap(),
1058+
tool_manager.load_tools(&database, &mut output).await.unwrap(),
10581059
None,
10591060
None,
10601061
tool_manager,
@@ -1076,10 +1077,11 @@ mod tests {
10761077
#[tokio::test]
10771078
async fn test_conversation_state_history_handling_with_tool_results() {
10781079
let mut database = Database::new().await.unwrap();
1080+
let mut output = SharedWriter::null();
10791081

10801082
// Build a long conversation history of tool use results.
10811083
let mut tool_manager = ToolManager::default();
1082-
let tool_config = tool_manager.load_tools(&database).await.unwrap();
1084+
let tool_config = tool_manager.load_tools(&database, &mut output).await.unwrap();
10831085
let mut conversation_state = ConversationState::new(
10841086
Context::new(),
10851087
"fake_conv_id",
@@ -1150,6 +1152,7 @@ mod tests {
11501152
#[tokio::test]
11511153
async fn test_conversation_state_with_context_files() {
11521154
let mut database = Database::new().await.unwrap();
1155+
let mut output = SharedWriter::null();
11531156

11541157
let ctx = Context::builder().with_test_home().await.unwrap().build_fake();
11551158
ctx.fs().write(AMAZONQ_FILENAME, "test context").await.unwrap();
@@ -1158,7 +1161,7 @@ mod tests {
11581161
let mut conversation_state = ConversationState::new(
11591162
ctx,
11601163
"fake_conv_id",
1161-
tool_manager.load_tools(&database).await.unwrap(),
1164+
tool_manager.load_tools(&database, &mut output).await.unwrap(),
11621165
None,
11631166
None,
11641167
tool_manager,
@@ -1199,6 +1202,7 @@ mod tests {
11991202
// tracing_subscriber::fmt::try_init().ok();
12001203

12011204
let mut database = Database::new().await.unwrap();
1205+
let mut output = SharedWriter::null();
12021206

12031207
let mut tool_manager = ToolManager::default();
12041208
let ctx = Context::builder().with_test_home().await.unwrap().build_fake();
@@ -1227,7 +1231,7 @@ mod tests {
12271231
let mut conversation_state = ConversationState::new(
12281232
ctx,
12291233
"fake_conv_id",
1230-
tool_manager.load_tools(&database).await.unwrap(),
1234+
tool_manager.load_tools(&database, &mut output).await.unwrap(),
12311235
None,
12321236
Some(SharedWriter::stdout()),
12331237
tool_manager,

crates/chat-cli/src/cli/chat/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -412,7 +412,7 @@ pub async fn chat(
412412
.interactive(interactive)
413413
.build(telemetry, tool_manager_output)
414414
.await?;
415-
let tool_config = tool_manager.load_tools(database).await?;
415+
let tool_config = tool_manager.load_tools(database, &mut output).await?;
416416
let mut tool_permissions = ToolPermissions::new(tool_config.len());
417417
if accept_all || trust_all_tools {
418418
tool_permissions.trust_all = true;

crates/chat-cli/src/cli/chat/tool_manager.rs

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ use tracing::{
5252
warn,
5353
};
5454

55+
use super::util::shared_writer::SharedWriter;
5556
use crate::api_client::model::{
5657
ToolResult,
5758
ToolResultContentBlock,
@@ -723,7 +724,11 @@ impl Clone for ToolManager {
723724
}
724725

725726
impl ToolManager {
726-
pub async fn load_tools(&mut self, database: &Database) -> eyre::Result<HashMap<String, ToolSpec>> {
727+
pub async fn load_tools(
728+
&mut self,
729+
database: &Database,
730+
output: &mut SharedWriter,
731+
) -> eyre::Result<HashMap<String, ToolSpec>> {
727732
let tx = self.loading_status_sender.take();
728733
let display_task = self.loading_display_task.take();
729734
self.schema = {
@@ -771,7 +776,21 @@ impl ToolManager {
771776
.map_or(5000_u64, |s| s as u64);
772777
Box::pin(tokio::time::sleep(std::time::Duration::from_millis(init_timeout)))
773778
} else {
774-
Box::pin(future::pending())
779+
// if it is non-interactive we will want to use the "mcp.noInteractiveTimeout"
780+
let init_timeout = database
781+
.settings
782+
.get_int(Setting::McpNoInteractiveTimeout)
783+
.map_or(30_000_u64, |s| s as u64);
784+
Box::pin(async move {
785+
tokio::time::sleep(std::time::Duration::from_millis(init_timeout)).await;
786+
let _ = queue!(
787+
output,
788+
style::Print(
789+
"Not all mcp servers loaded. Configure no-interactive timeout with q settings mcp.noInteractiveTimeout"
790+
),
791+
style::Print("\n")
792+
);
793+
})
775794
};
776795
tokio::select! {
777796
_ = display_fut => {},

crates/chat-cli/src/database/settings.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ pub enum Setting {
2929
ApiCodeWhispererService,
3030
ApiQService,
3131
McpInitTimeout,
32+
McpNoInteractiveTimeout,
3233
McpLoadedBefore,
3334
}
3435

@@ -47,6 +48,7 @@ impl AsRef<str> for Setting {
4748
Self::ApiCodeWhispererService => "api.codewhisperer.service",
4849
Self::ApiQService => "api.q.service",
4950
Self::McpInitTimeout => "mcp.initTimeout",
51+
Self::McpNoInteractiveTimeout => "mcp.noInteractiveTimeout",
5052
Self::McpLoadedBefore => "mcp.loadedBefore",
5153
}
5254
}
@@ -75,6 +77,7 @@ impl TryFrom<&str> for Setting {
7577
"api.codewhisperer.service" => Ok(Self::ApiCodeWhispererService),
7678
"api.q.service" => Ok(Self::ApiQService),
7779
"mcp.initTimeout" => Ok(Self::McpInitTimeout),
80+
"mcp.noInteractiveTimeout" => Ok(Self::McpNoInteractiveTimeout),
7881
"mcp.loadedBefore" => Ok(Self::McpLoadedBefore),
7982
_ => Err(DatabaseError::InvalidSetting(value.to_string())),
8083
}

0 commit comments

Comments
 (0)