From 08a0feefd4f97fce906420ee25961c44727efac8 Mon Sep 17 00:00:00 2001 From: Christian Brauner Date: Wed, 4 Nov 2026 15:46:20 +0100 Subject: [PATCH] semcode-mcp: answer requests concurrently The MCP server reads one request from stdin, awaits handle_request() for it and only then reads the next one. So a query that scans the whole database, grep_functions over every function body or find_implementors, holds up every other client of the server until it is done, and they all run into the idle timeout of the MCP client at the same time. Twenty review agents on one server stalled for the better part of an hour behind a few such scans, while the CLI, which runs the same query code in a process of its own, kept answering in under a second. So hand each request to a task of its own and let the responses go out as they finish. JSON-RPC matches them by id, the order doesn't matter, and the stdout lock already serializes the writes with the notification task. The handlers borrow their arguments across awaits in a way rustc can't prove Send for, so the task drives the request on a thread of the blocking pool with Handle::block_on() instead of running it on a worker. That also keeps a CPU bound scan off the runtime's workers. The number of requests in flight is capped at the number of CPUs with a semaphore, so that any number of scans can't thrash the machine. Notifications are still handled inline, there is nothing to wait for there. A task that can't write its response gives up on it and logs that, the loop ends on its own when stdin is closed. A find_function sent right after a grep_functions over all function bodies of a kernel tree used to be answered after the scan, 26 seconds in. Now it comes back in 4 seconds while the scan is still running. Signed-off-by: Christian Brauner (Amutable) --- src/bin/semcode-mcp.rs | 47 +++++++++++++++++++++++++++++++++++------- 1 file changed, 40 insertions(+), 7 deletions(-) diff --git a/src/bin/semcode-mcp.rs b/src/bin/semcode-mcp.rs index bb25a11..13bb14d 100644 --- a/src/bin/semcode-mcp.rs +++ b/src/bin/semcode-mcp.rs @@ -5933,6 +5933,18 @@ async fn run_stdio_server(server: Arc) -> Result<()> { let stdout = stdout_clone; + // Requests used to be answered one after the other from the loop below, + // so a query that scans the whole database held up every other client + // of the server until it was done. Each request gets a task of its own + // now and the responses go out as they finish. JSON-RPC matches them by + // id, so the order doesn't matter. A full scan is CPU bound, so the + // number of requests in flight is capped at the number of CPUs. + let max_in_flight = std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(4) + .max(2); + let in_flight = Arc::new(tokio::sync::Semaphore::new(max_in_flight)); + let mut line = String::new(); loop { @@ -5953,25 +5965,46 @@ async fn run_stdio_server(server: Arc) -> Result<()> { Ok(request) => { // Check if this is a notification (no id) or a request (has id) if request.get("id").is_some() { - // It's a request, send a response - let response = server.handle_request(request).await; - if let Ok(response_str) = serde_json::to_string(&response) { + // It's a request, answer it from a task of its + // own, see max_in_flight above. + let server = server.clone(); + let stdout = stdout.clone(); + let in_flight = in_flight.clone(); + tokio::spawn(async move { + let Ok(_permit) = in_flight.acquire_owned().await else { + return; + }; + // The handlers borrow their arguments across + // awaits in a way rustc can't prove Send for, + // so drive the request on a thread of the + // blocking pool instead. That also keeps a + // CPU bound scan off the runtime's workers. + let handle = tokio::runtime::Handle::current(); + let Ok(response) = tokio::task::spawn_blocking(move || { + handle.block_on(server.handle_request(request)) + }) + .await + else { + return; + }; + let Ok(response_str) = serde_json::to_string(&response) else { + return; + }; let mut stdout_guard = stdout.lock().await; if let Err(e) = stdout_guard.write_all(response_str.as_bytes()).await { eprintln!("Failed to write response: {e}"); - break; + return; } if let Err(e) = stdout_guard.write_all(b"\n").await { eprintln!("Failed to write newline: {e}"); - break; + return; } if let Err(e) = stdout_guard.flush().await { eprintln!("Failed to flush stdout: {e}"); - break; } - } + }); } else { // It's a notification, handle without response let method = request["method"].as_str().unwrap_or("");