Codex Handbook
app-server/src/fuzzy_file_search.rs 256 lines
use std::num::NonZero;use std::path::PathBuf;use std::sync::Arc;use std::sync::Mutex;use std::sync::atomic::AtomicBool;use std::sync::atomic::Ordering;use codex_app_server_protocol::FuzzyFileSearchMatchType;use codex_app_server_protocol::FuzzyFileSearchResult;use codex_app_server_protocol::FuzzyFileSearchSessionCompletedNotification;use codex_app_server_protocol::FuzzyFileSearchSessionUpdatedNotification;use codex_app_server_protocol::ServerNotification;use codex_file_search as file_search;use tracing::warn;use crate::outgoing_message::OutgoingMessageSender;const MATCH_LIMIT: usize = 50;const MAX_THREADS: usize = 12;pub(crate) async fn run_fuzzy_file_search(    query: String,    roots: Vec<String>,    cancellation_flag: Arc<AtomicBool>,) -> Vec<FuzzyFileSearchResult> {    if roots.is_empty() {        return Vec::new();    }    #[expect(clippy::expect_used)]    let limit = NonZero::new(MATCH_LIMIT).expect("MATCH_LIMIT should be a valid non-zero usize");    let cores = std::thread::available_parallelism()        .map(std::num::NonZero::get)        .unwrap_or(1);    let threads = cores.min(MAX_THREADS);    #[expect(clippy::expect_used)]    let threads = NonZero::new(threads.max(1)).expect("threads should be non-zero");    let search_dirs: Vec<PathBuf> = roots.iter().map(PathBuf::from).collect();    let mut files = match tokio::task::spawn_blocking(move || {        file_search::run(            query.as_str(),            search_dirs,            file_search::FileSearchOptions {                limit,                threads,                compute_indices: true,                ..Default::default()            },            Some(cancellation_flag),        )    })    .await    {        Ok(Ok(res)) => res            .matches            .into_iter()            .map(|m| {                let file_name = m.path.file_name().unwrap_or_default();                FuzzyFileSearchResult {                    root: m.root.to_string_lossy().to_string(),                    path: m.path.to_string_lossy().to_string(),                    match_type: match m.match_type {                        file_search::MatchType::File => FuzzyFileSearchMatchType::File,                        file_search::MatchType::Directory => FuzzyFileSearchMatchType::Directory,                    },                    file_name: file_name.to_string_lossy().to_string(),                    score: m.score,                    indices: m.indices,                }            })            .collect::<Vec<_>>(),        Ok(Err(err)) => {            warn!("fuzzy-file-search failed: {err}");            Vec::new()        }        Err(err) => {            warn!("fuzzy-file-search join failed: {err}");            Vec::new()        }    };    files.sort_by(file_search::cmp_by_score_desc_then_path_asc::<        FuzzyFileSearchResult,        _,        _,    >(|f| f.score, |f| f.path.as_str()));    files}pub(crate) struct FuzzyFileSearchSession {    session: file_search::FileSearchSession,    shared: Arc<SessionShared>,}impl FuzzyFileSearchSession {    pub(crate) fn update_query(&self, query: String) {        if self.shared.canceled.load(Ordering::Relaxed) {            return;        }        {            #[expect(clippy::unwrap_used)]            let mut latest_query = self.shared.latest_query.lock().unwrap();            *latest_query = query.clone();        }        self.session.update_query(&query);    }}impl Drop for FuzzyFileSearchSession {    fn drop(&mut self) {        self.shared.canceled.store(true, Ordering::Relaxed);    }}pub(crate) fn start_fuzzy_file_search_session(    session_id: String,    roots: Vec<String>,    outgoing: Arc<OutgoingMessageSender>,) -> anyhow::Result<FuzzyFileSearchSession> {    #[expect(clippy::expect_used)]    let limit = NonZero::new(MATCH_LIMIT).expect("MATCH_LIMIT should be a valid non-zero usize");    let cores = std::thread::available_parallelism()        .map(std::num::NonZero::get)        .unwrap_or(1);    let threads = cores.min(MAX_THREADS);    #[expect(clippy::expect_used)]    let threads = NonZero::new(threads.max(1)).expect("threads should be non-zero");    let search_dirs: Vec<PathBuf> = roots.iter().map(PathBuf::from).collect();    let canceled = Arc::new(AtomicBool::new(false));    let shared = Arc::new(SessionShared {        session_id,        latest_query: Mutex::new(String::new()),        outgoing,        runtime: tokio::runtime::Handle::current(),        canceled: canceled.clone(),    });    let reporter = Arc::new(SessionReporterImpl {        shared: shared.clone(),    });    let session = file_search::create_session(        search_dirs,        file_search::FileSearchOptions {            limit,            threads,            compute_indices: true,            ..Default::default()        },        reporter,        Some(canceled),    )?;    Ok(FuzzyFileSearchSession { session, shared })}struct SessionShared {    session_id: String,    latest_query: Mutex<String>,    outgoing: Arc<OutgoingMessageSender>,    runtime: tokio::runtime::Handle,    canceled: Arc<AtomicBool>,}struct SessionReporterImpl {    shared: Arc<SessionShared>,}impl SessionReporterImpl {    fn send_snapshot(&self, snapshot: &file_search::FileSearchSnapshot) {        if self.shared.canceled.load(Ordering::Relaxed) {            return;        }        let query = {            #[expect(clippy::unwrap_used)]            self.shared.latest_query.lock().unwrap().clone()        };        if snapshot.query != query {            return;        }        let files = if query.is_empty() {            Vec::new()        } else {            collect_files(snapshot)        };        let notification = ServerNotification::FuzzyFileSearchSessionUpdated(            FuzzyFileSearchSessionUpdatedNotification {                session_id: self.shared.session_id.clone(),                query,                files,            },        );        let outgoing = self.shared.outgoing.clone();        self.shared.runtime.spawn(async move {            outgoing.send_server_notification(notification).await;        });    }    fn send_complete(&self) {        if self.shared.canceled.load(Ordering::Relaxed) {            return;        }        let session_id = self.shared.session_id.clone();        let outgoing = self.shared.outgoing.clone();        self.shared.runtime.spawn(async move {            let notification = ServerNotification::FuzzyFileSearchSessionCompleted(                FuzzyFileSearchSessionCompletedNotification { session_id },            );            outgoing.send_server_notification(notification).await;        });    }}impl file_search::SessionReporter for SessionReporterImpl {    fn on_update(&self, snapshot: &file_search::FileSearchSnapshot) {        self.send_snapshot(snapshot);    }    fn on_complete(&self) {        self.send_complete();    }}fn collect_files(snapshot: &file_search::FileSearchSnapshot) -> Vec<FuzzyFileSearchResult> {    let mut files = snapshot        .matches        .iter()        .map(|m| {            let file_name = m.path.file_name().unwrap_or_default();            FuzzyFileSearchResult {                root: m.root.to_string_lossy().to_string(),                path: m.path.to_string_lossy().to_string(),                match_type: match m.match_type {                    file_search::MatchType::File => FuzzyFileSearchMatchType::File,                    file_search::MatchType::Directory => FuzzyFileSearchMatchType::Directory,                },                file_name: file_name.to_string_lossy().to_string(),                score: m.score,                indices: m.indices.clone(),            }        })        .collect::<Vec<_>>();    files.sort_by(file_search::cmp_by_score_desc_then_path_asc::<        FuzzyFileSearchResult,        _,        _,    >(|f| f.score, |f| f.path.as_str()));    files}