diff --git a/docs/semcode-mcp.md b/docs/semcode-mcp.md index 565a62b..52776b4 100644 --- a/docs/semcode-mcp.md +++ b/docs/semcode-mcp.md @@ -157,6 +157,21 @@ All `*_patterns` arrays below are OR'd within the array. **indexing_status**: show background indexing progress, errors, and timing. No parameters. +## Startup and readiness + +At startup the server checks whether the index covers the current commit +of the repository and indexes it when it does not. Lookups that arrive +meanwhile wait for the check (up to 30 seconds) rather than being told +to retry; only an indexing pass that is still running after that gets +the "currently being indexed" answer. + +The server speaks the sd_notify(3) protocol: when NOTIFY_SOCKET is set +it sends `STATUS=` lines for the phases and `READY=1` once lookups are +answered, and `STOPPING=1` when it exits. So `systemd-notify --fork -- +semcode-mcp ...` returns once the server is ready, and a unit with +`Type=notify` becomes active at that point. A check or indexing failure +sends a `STATUS=` line and no `READY=1`. + ## Lazy Loading Start the server with `--lazy` to cut initial context ~96%. The diff --git a/src/bin/semcode-mcp.rs b/src/bin/semcode-mcp.rs index bb25a11..0799801 100644 --- a/src/bin/semcode-mcp.rs +++ b/src/bin/semcode-mcp.rs @@ -1975,6 +1975,9 @@ struct Args { lazy: bool, } +/// How long a lookup waits for the startup pass before it is told to retry. +const STARTUP_WAIT_SECONDS: u64 = 30; + #[derive(Debug, Clone)] enum IndexingStatus { NotStarted, @@ -2912,17 +2915,34 @@ impl McpServer { } async fn check_database_status(&self) -> Option { - let state = self.indexing_state.lock().await; - - // Check if indexing is currently in progress - if matches!(state.status, IndexingStatus::InProgress { .. }) { - if let IndexingStatus::InProgress { ref phase, .. } = state.status { + // The startup pass first checks whether HEAD is indexed, which takes a + // moment, and indexes it when it is not, which can take long. Wait for + // the first case instead of sending the caller away; only a pass that + // is still running after that is reported. + let wait_until = + tokio::time::Instant::now() + std::time::Duration::from_secs(STARTUP_WAIT_SECONDS); + let state = loop { + let state = self.indexing_state.lock().await; + let busy = matches!( + state.status, + IndexingStatus::NotStarted | IndexingStatus::InProgress { .. } + ); + if !busy { + break state; + } + if tokio::time::Instant::now() >= wait_until { + let phase = match &state.status { + IndexingStatus::InProgress { phase, .. } => phase.clone(), + _ => "starting".to_string(), + }; return Some(format!( "Database is currently being indexed ({}). Please wait for indexing to complete and try again. Use the indexing_status tool to check progress.", phase )); } - } + drop(state); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + }; // Check if indexing failed if let IndexingStatus::Failed { ref error } = state.status { @@ -5704,6 +5724,8 @@ async fn index_current_commit_background( eprintln!("[Background] Background indexing task started"); send_notification("Semcode: Checking if indexing is needed...".to_string()); + semcode::notify::status("Checking the index against the current commit"); + // Get current git SHA let _git_sha = match semcode::git::get_git_sha(&git_repo) { Ok(Some(sha)) => { @@ -5725,6 +5747,7 @@ async fn index_current_commit_background( send_notification( "Semcode: Not in a git repository, skipping auto-indexing".to_string(), ); + semcode::notify::status("Not in a git repository, no index check"); return; } Err(e) => { @@ -5735,6 +5758,7 @@ async fn index_current_commit_background( error: error_msg.clone(), }; send_notification(format!("Semcode: {}", error_msg)); + semcode::notify::status(&error_msg); return; } }; @@ -5768,6 +5792,16 @@ async fn index_current_commit_background( changed_file.path, &new_hash[..8] ); + { + let mut state = indexing_state.lock().await; + state.status = + IndexingStatus::Completed { files_processed: 0 }; + state.completed_at = Some(std::time::SystemTime::now()); + } + semcode::notify::ready(&format!( + "Index up to date for {}", + &_git_sha[..12.min(_git_sha.len())] + )); return; } else { eprintln!( @@ -5782,6 +5816,11 @@ async fn index_current_commit_background( "[Background] Warning: Failed to get processed files: {}", e ); + let mut state = indexing_state.lock().await; + state.status = IndexingStatus::Failed { + error: format!("Failed to get processed files: {}", e), + }; + semcode::notify::status(&format!("Index check failed: {}", e)); return; } } @@ -5817,6 +5856,7 @@ async fn index_current_commit_background( }; } send_notification("Semcode: Indexing current commit...".to_string()); + semcode::notify::status(&format!("Indexing {}", &_git_sha[..12.min(_git_sha.len())])); // Create synthetic range for current commit: HEAD^..HEAD let git_range = format!("{}^..{}", _git_sha, _git_sha); @@ -5878,6 +5918,10 @@ async fn index_current_commit_background( ); } } + semcode::notify::ready(&format!( + "Index up to date for {}", + &_git_sha[..12.min(_git_sha.len())] + )); } Err(e) => { eprintln!("[Background] Warning: Auto-indexing failed: {}", e); @@ -5887,6 +5931,7 @@ async fn index_current_commit_background( }; state.completed_at = Some(std::time::SystemTime::now()); send_notification(format!("Semcode: Indexing failed: {}", e)); + semcode::notify::status(&format!("Indexing failed: {}", e)); } } } @@ -5933,6 +5978,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 +6010,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(""); @@ -6074,13 +6152,14 @@ async fn main() -> Result<()> { tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; // Run MCP server on stdio - run_stdio_server(server).await?; + let served = run_stdio_server(server).await; + semcode::notify::stopping(); // Gracefully shutdown the background indexing task indexing_handle.abort(); let _ = indexing_handle.await; - Ok(()) + served } #[cfg(test)] diff --git a/src/lib.rs b/src/lib.rs index 7ed091e..19b0552 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,7 @@ pub mod git; pub mod git_range; pub mod hash; pub mod indexer; +pub mod notify; pub mod perf_monitor; pub mod symbol_walkback; pub mod text_utils; diff --git a/src/notify.rs b/src/notify.rs new file mode 100644 index 0000000..6bc161f --- /dev/null +++ b/src/notify.rs @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +//! The sd_notify(3) protocol: a datagram of newline separated assignments to +//! the socket named by NOTIFY_SOCKET. A service manager that started us with +//! Type=notify, or systemd-notify --fork, learns from READY=1 that we answer +//! lookups. Without the variable nothing is sent. + +use std::os::linux::net::SocketAddrExt; +use std::os::unix::net::{SocketAddr, UnixDatagram}; + +/// Send one notification. A leading `@` names an abstract socket. +pub fn notify(message: &str) { + let Ok(socket) = std::env::var("NOTIFY_SOCKET") else { + return; + }; + if socket.is_empty() { + return; + } + let Ok(sender) = UnixDatagram::unbound() else { + return; + }; + let result = match socket.strip_prefix('@') { + Some(name) => SocketAddr::from_abstract_name(name.as_bytes()) + .and_then(|addr| sender.send_to_addr(message.as_bytes(), &addr)), + None => sender.send_to(message.as_bytes(), &socket), + }; + if let Err(e) = result { + tracing::debug!("sd_notify to {socket} failed: {e}"); + } +} + +/// Startup is over: lookups are answered from here on. +pub fn ready(status: &str) { + notify(&format!("READY=1\nSTATUS={status}\n")); +} + +pub fn status(status: &str) { + notify(&format!("STATUS={status}\n")); +} + +pub fn stopping() { + notify("STOPPING=1\n"); +}