diff --git a/crates/csp/src/bin/csp/mcp_server.rs b/crates/csp/src/bin/csp/mcp_server.rs index e12707c..2e7d4ea 100644 --- a/crates/csp/src/bin/csp/mcp_server.rs +++ b/crates/csp/src/bin/csp/mcp_server.rs @@ -173,7 +173,7 @@ pub fn run_mcp( content: Vec, ) -> Result<()> { let rt = tokio::runtime::Runtime::new()?; - rt.block_on(async move { + block_then_abandon(rt, async move { let service = CspMcpServer::new(default_source, default_ref, content) .serve(stdio()) .await?; @@ -182,12 +182,91 @@ pub fn run_mcp( }) } +/// Drive `fut` to completion on `rt`, then abandon the runtime without joining +/// its workers. +/// +/// Tool handlers run synchronous index/model work inline on a runtime worker. +/// Dropping the runtime would join that worker, so `csp mcp` would stay alive +/// after stdin EOF until an in-flight index build finished. +/// `shutdown_background` returns immediately, so the process exits without +/// waiting for them. It runs from a drop guard so a panic in `fut` abandons the +/// runtime too, instead of joining it during unwinding. +fn block_then_abandon(rt: tokio::runtime::Runtime, fut: F) -> F::Output { + struct Abandon(Option); + impl Drop for Abandon { + fn drop(&mut self) { + if let Some(rt) = self.0.take() { + rt.shutdown_background(); + } + } + } + let guard = Abandon(Some(rt)); + guard.0.as_ref().expect("runtime is set").block_on(fut) +} + #[cfg(test)] mod tests { use super::*; use std::fs; + use std::time::{Duration, Instant}; use tempfile::tempdir; + /// Stand-in for a tool handler's synchronous index/model work, which runs + /// inline on a runtime worker exactly like this. + fn synchronous_index_build() { + std::thread::sleep(Duration::from_secs(30)); + } + + #[test] + fn block_then_abandon_does_not_join_a_blocked_worker() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .build() + .unwrap(); + let start = Instant::now(); + block_then_abandon(rt, async { + // Park a worker in synchronous work, like a handler mid index build. + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let _ = started_tx.send(()); + synchronous_index_build(); + }); + // Wait until the worker is actually inside the blocking section. + started_rx.await.unwrap(); + }); + assert!( + start.elapsed() < Duration::from_secs(5), + "runtime teardown joined the blocked worker ({:?})", + start.elapsed() + ); + } + + #[test] + fn block_then_abandon_does_not_join_a_blocked_worker_on_panic() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .build() + .unwrap(); + let start = Instant::now(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + block_then_abandon(rt, async { + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + tokio::spawn(async move { + let _ = started_tx.send(()); + synchronous_index_build(); + }); + started_rx.await.unwrap(); + panic!("serve loop failed"); + }) + })); + assert!(result.is_err()); + assert!( + start.elapsed() < Duration::from_secs(5), + "unwinding joined the blocked worker ({:?})", + start.elapsed() + ); + } + #[test] fn search_params_deserialize_with_and_without_optionals() { // Only `query` is required; repo/top_k default to None.