diff --git a/crates/fff-mcp/src/main.rs b/crates/fff-mcp/src/main.rs index 497c5ee95..e14671f72 100644 --- a/crates/fff-mcp/src/main.rs +++ b/crates/fff-mcp/src/main.rs @@ -152,6 +152,14 @@ pub(crate) struct Args { /// Run a health check and print diagnostic information, then exit. #[arg(long = "healthcheck")] pub(crate) healthcheck: bool, + + /// Exit after this many seconds of inactivity. 0 = never exit. + #[arg( + long = "idle-timeout-secs", + env = "FFF_MCP_IDLE_TIMEOUT_SECS", + default_value_t = 900 + )] + idle_timeout_secs: u64, } /// Resolve default paths for the log file. @@ -252,7 +260,7 @@ async fn main() -> Result<(), Box> { // Initialize file picker (spawns background scan + watcher) FilePicker::new_with_shared_state( shared_picker.clone(), - shared_frecency.clone(), + shared_frecency, fff::FilePickerOptions { base_path, enable_mmap_cache: !args.no_warmup, @@ -273,7 +281,9 @@ async fn main() -> Result<(), Box> { } // Create and start the MCP server - let server = FffServer::new(shared_picker.clone(), shared_frecency.clone()); + let server = FffServer::new(shared_picker.clone()); + let last_activity = server.last_activity(); + let idle_timeout_secs = args.idle_timeout_secs; // Wait for initial scan in background — don't block server startup let picker_clone_for_scan = shared_picker.clone(); @@ -294,10 +304,45 @@ async fn main() -> Result<(), Box> { } }); - let service = server - .serve(stdio()) - .await - .map_err(|e| format!("Failed to start MCP server: {}", e))?; + const STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + let service = match tokio::time::timeout(STARTUP_TIMEOUT, server.serve(stdio())).await { + Ok(res) => res.map_err(|e| format!("Failed to start MCP server: {}", e))?, + Err(_) => { + return Err("MCP initialize handshake did not complete within 60s".into()); + } + }; + + if idle_timeout_secs > 0 { + last_activity.store( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0), + std::sync::atomic::Ordering::Relaxed, + ); + + let last_activity_for_watchdog = last_activity.clone(); + tokio::spawn(async move { + let tick = std::time::Duration::from_secs(60); + loop { + tokio::time::sleep(tick).await; + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0); + + let last = last_activity_for_watchdog.load(std::sync::atomic::Ordering::Relaxed); + if now.saturating_sub(last) >= idle_timeout_secs { + tracing::info!( + "Exiting after {}s of inactivity (idle_timeout_secs={})", + now.saturating_sub(last), + idle_timeout_secs + ); + std::process::exit(0); + } + } + }); + } let picker_for_shutdown = shared_picker.clone(); tokio::spawn(async move { diff --git a/crates/fff-mcp/src/server.rs b/crates/fff-mcp/src/server.rs index 0fc6102ab..c4affa334 100644 --- a/crates/fff-mcp/src/server.rs +++ b/crates/fff-mcp/src/server.rs @@ -1,29 +1,19 @@ -//! FFF MCP server — tool definitions and handlers. -//! -//! Uses the `rmcp` crate's `#[tool_router]` / `#[tool_handler]` macros -//! for declarative tool registration. Each tool method directly calls -//! `fff-core` APIs (no C FFI overhead). - -use std::borrow::Cow; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; - use crate::cursor::CursorStore; use crate::output::{GrepFormatter, OutputMode, file_suffix}; use fff::grep::{GrepMode, GrepSearchOptions, has_regex_metacharacters}; use fff::types::{FileItem, PaginationArgs}; -use fff::{FuzzySearchOptions, QueryParser, SharedFilePicker, SharedFrecency}; +use fff::{FuzzySearchOptions, QueryParser, SharedFilePicker}; use fff_query_parser::AiGrepConfig; use rmcp::handler::server::wrapper::Parameters; use rmcp::model::*; use rmcp::{ServerHandler, schemars, tool, tool_handler, tool_router}; +use std::borrow::Cow; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{SystemTime, UNIX_EPOCH}; + +const SCAN_READY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); -/// Normalize the caller-supplied `maxResults`. -/// -/// `None`, `Some(0)`, and non-positive / non-finite values fall back to -/// `default`. Issue #400 reported that grep returned 0 items for -/// `maxResults: 0` while `find_files` returned the entire dataset; treating -/// 0 as "use the default" makes both tools behave consistently. fn normalize_max_results(raw: Option, default: usize) -> usize { match raw { None => default, @@ -181,34 +171,64 @@ pub struct MultiGrepParams { #[derive(Clone)] pub struct FffServer { picker: SharedFilePicker, - #[allow(dead_code)] - frecency: SharedFrecency, cursor_store: Arc>, update_notice_sent: Arc, + last_activity: Arc, + scan_ready: Arc, +} + +fn now_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) } impl FffServer { - pub fn new(picker: SharedFilePicker, frecency: SharedFrecency) -> Self { + pub fn new(picker: SharedFilePicker) -> Self { Self { picker, - frecency, cursor_store: Arc::new(Mutex::new(CursorStore::new())), update_notice_sent: Arc::new(AtomicBool::new(false)), + last_activity: Arc::new(AtomicU64::new(now_secs())), + scan_ready: Arc::new(AtomicBool::new(false)), } } - #[allow(dead_code)] - pub fn wait_for_scan(&self) { + pub fn last_activity(&self) -> Arc { + self.last_activity.clone() + } + + fn bump_activity(&self) { + self.last_activity.store(now_secs(), Ordering::Relaxed); + } + + fn wait_for_scan(&self, timeout: std::time::Duration) -> Result<(), ErrorData> { + if self.scan_ready.load(Ordering::Relaxed) { + return Ok(()); + } + + let deadline = std::time::Instant::now() + timeout; + loop { - let guard = self.picker.read().ok(); - let is_scanning = guard + let is_scanning = self + .picker + .read() + .ok() .as_ref() .and_then(|g| g.as_ref()) .map(|p| p.is_scan_active()) .unwrap_or(true); if !is_scanning { - break; + self.scan_ready.store(true, Ordering::Relaxed); + return Ok(()); + } + if std::time::Instant::now() >= deadline { + return Err(ErrorData::internal_error( + "Index is still building; retry shortly", + None, + )); } std::thread::sleep(std::time::Duration::from_millis(50)); } @@ -407,6 +427,9 @@ impl FffServer { &self, Parameters(params): Parameters, ) -> Result { + self.bump_activity(); + self.wait_for_scan(SCAN_READY_TIMEOUT)?; + let max_results = normalize_max_results(params.max_results, 20); let query = ¶ms.query; @@ -519,6 +542,9 @@ impl FffServer { &self, Parameters(params): Parameters, ) -> Result { + self.bump_activity(); + self.wait_for_scan(SCAN_READY_TIMEOUT)?; + let max_results = normalize_max_results(params.max_results, 20); let output_mode = OutputMode::new(params.output_mode.as_deref()); @@ -553,6 +579,9 @@ impl FffServer { &self, Parameters(params): Parameters, ) -> Result { + self.bump_activity(); + self.wait_for_scan(SCAN_READY_TIMEOUT)?; + let mut result = self.multi_grep_inner(params)?; self.maybe_append_update_notice(&mut result); Ok(result)