Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 121 additions & 6 deletions crates/jp_cli/src/cmd/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ use jp_config::{
},
fs::{expand_tilde, load_partial},
model::parameters::{PartialCustomReasoningConfig, PartialReasoningConfig, ReasoningConfig},
style::reasoning::ReasoningDisplayConfig,
style::{mcp_startup::McpStartupConfig, reasoning::ReasoningDisplayConfig},
};
use jp_conversation::{
Conversation, ConversationEvent, ConversationId, ConversationStream,
Expand All @@ -97,9 +97,11 @@ use jp_llm::{
tool_definitions,
},
};
use jp_mcp::{StartupSet, id::McpServerId};
use jp_printer::Printer;
use jp_storage::backend::Projection;
use jp_task::task::TitleGeneratorTask;
use jp_term::width::{display_width, truncate_to_width};
use jp_workspace::{ConversationHandle, ConversationLock, Workspace};
use minijinja::{Environment, UndefinedBehavior};
use tool::{TerminalExecutorSource, ToolCoordinator};
Expand Down Expand Up @@ -131,6 +133,7 @@ use crate::{
parser::AttachmentUrlOrPath,
render::TurnView,
signals::SignalRouter,
timer::spawn_line_timer,
};

type BoxedResult<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
Expand Down Expand Up @@ -403,7 +406,7 @@ impl Query {
self.apply_pre_query_compaction(&lock, &cfg).await?;
}

let mut mcp_servers_handle = ctx.configure_active_mcp_servers().await?;
let mcp_servers_handle = ctx.configure_active_mcp_servers().await?;

let conv_title = lock.metadata().title.clone();

Expand Down Expand Up @@ -570,10 +573,16 @@ impl Query {
}
}

// Wait for all MCP servers to finish loading.
while let Some(result) = mcp_servers_handle.join_next().await {
result??;
}
// Wait for all MCP servers to finish loading, showing a timer line
// when the wait takes long enough to be noticeable.
await_mcp_servers(
mcp_servers_handle,
cfg.style.mcp_startup.clone(),
ctx.printer.clone(),
ctx.term.is_tty,
ctx.term.width,
)
.await?;

let forced_tool = cfg.assistant.tool_choice.function_name();
let tools =
Expand Down Expand Up @@ -1102,6 +1111,112 @@ impl Query {
}
}

/// Wait for background MCP server startups to complete.
///
/// Shows a single aggregate timer line on stderr once the wait exceeds the
/// configured delay, updating the listed server names as startups finish.
/// Servers that finish within the delay never trigger the line.
///
/// Returns the first startup error, after clearing the timer line so the error
/// renders on a clean row.
///
/// `width` bounds the rendered line to the terminal so a long server list wraps
/// no further than one row, keeping the timer's single-row clear on finish
/// sufficient.
/// `None` leaves the line unbounded (unknown width).
async fn await_mcp_servers(
mut startup: StartupSet,
config: McpStartupConfig,
printer: Arc<Printer>,
is_tty: bool,
width: Option<u16>,
) -> std::result::Result<(), cmd::Error> {
if startup.joins.is_empty() {
return Ok(());
}

let timer = spawn_line_timer(
printer,
config.show && is_tty,
Duration::from_secs(config.delay_secs.into()),
Duration::from_millis(config.interval_ms.into()),
move |secs, status| mcp_startup_line(secs, status, width),
);
if let Some(timer) = &timer {
timer.set_status(mcp_startup_status(&startup.pending));
}

let result = loop {
match startup.joins.join_next().await {
None => break Ok(()),
Some(Err(error)) => break Err(error.into()),
Some(Ok(Err(error))) => break Err(error.into()),
Some(Ok(Ok(id))) => {
startup.pending.retain(|pending| pending != &id);
if let Some(timer) = &timer
&& !startup.pending.is_empty()
{
timer.set_status(mcp_startup_status(&startup.pending));
}
}
}
};

if let Some(timer) = timer {
timer.finish().await;
}

result
}

/// Render the MCP startup timer line for `secs` elapsed and `status`, bounding
/// the visible text to `width` columns when known.
///
/// Truncation falls on the server-list fragment only: the ` ⏱ Starting `
/// prefix and the ` {secs:.1}s ` timer suffix are always preserved, so the
/// elapsed time keeps moving even when a long list overflows.
/// A terminal too narrow for even the prefix and suffix falls back to a bounded
/// `⏱ {secs:.1}s`.
/// The leading `\r\x1b[K` control prefix stays outside the width budget.
fn mcp_startup_line(secs: f64, status: Option<&str>, width: Option<u16>) -> String {
let status = status.unwrap_or("MCP servers");
let full = format!("⏱ Starting {status}… {secs:.1}s");
let line = match width {
Some(w) if display_width(&full) > usize::from(w) => {
let w = usize::from(w);
let prefix = "⏱ Starting ";
let suffix = format!(" {secs:.1}s");
let reserved = display_width(prefix) + display_width(&suffix);
if w <= reserved {
truncate_to_width(&format!("⏱ {secs:.1}s"), w)
} else {
let status = truncate_to_width(status, w - reserved);
format!("{prefix}{status}{suffix}")
}
}
_ => full,
};
format!("\r\x1b[K{line}")
}

/// Render the pending-server fragment for the MCP startup timer line.
///
/// One server renders as `MCP server bookworm`; several render as `2 MCP
/// servers (bookworm, grizzly)`.
fn mcp_startup_status(pending: &[McpServerId]) -> String {
match pending {
[id] => format!("MCP server {id}"),
ids => format!(
"{} MCP servers ({})",
ids.len(),
ids.iter()
.map(McpServerId::as_str)
.collect::<Vec<_>>()
.join(", ")
),
}
}

/// Return the most recent assistant message text in the stream.
///
/// Walks the stream in reverse and returns the first `ChatResponse::Message` it
Expand Down
Loading
Loading