Codex Handbook
app-server/src/request_processors/search.rs 134 lines
use std::collections::HashMap;use std::sync::Arc;use std::sync::atomic::AtomicBool;use std::sync::atomic::Ordering;use crate::error_code::internal_error;use crate::error_code::invalid_request;use crate::fuzzy_file_search::FuzzyFileSearchSession;use crate::fuzzy_file_search::run_fuzzy_file_search;use crate::fuzzy_file_search::start_fuzzy_file_search_session;use crate::outgoing_message::OutgoingMessageSender;use codex_app_server_protocol::FuzzyFileSearchParams;use codex_app_server_protocol::FuzzyFileSearchResponse;use codex_app_server_protocol::FuzzyFileSearchSessionStartParams;use codex_app_server_protocol::FuzzyFileSearchSessionStartResponse;use codex_app_server_protocol::FuzzyFileSearchSessionStopParams;use codex_app_server_protocol::FuzzyFileSearchSessionStopResponse;use codex_app_server_protocol::FuzzyFileSearchSessionUpdateParams;use codex_app_server_protocol::FuzzyFileSearchSessionUpdateResponse;use codex_app_server_protocol::JSONRPCErrorError;use tokio::sync::Mutex;#[derive(Clone)]pub(crate) struct SearchRequestProcessor {    outgoing: Arc<OutgoingMessageSender>,    pending_fuzzy_searches: Arc<Mutex<HashMap<String, Arc<AtomicBool>>>>,    fuzzy_search_sessions: Arc<Mutex<HashMap<String, FuzzyFileSearchSession>>>,}impl SearchRequestProcessor {    pub(crate) fn new(outgoing: Arc<OutgoingMessageSender>) -> Self {        Self {            outgoing,            pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())),            fuzzy_search_sessions: Arc::new(Mutex::new(HashMap::new())),        }    }    pub(crate) async fn fuzzy_file_search(        &self,        params: FuzzyFileSearchParams,    ) -> Result<FuzzyFileSearchResponse, JSONRPCErrorError> {        let FuzzyFileSearchParams {            query,            roots,            cancellation_token,        } = params;        let cancel_flag = match cancellation_token.clone() {            Some(token) => {                let mut pending_fuzzy_searches = self.pending_fuzzy_searches.lock().await;                // if a cancellation_token is provided and a pending_request exists for                // that token, cancel it                if let Some(existing) = pending_fuzzy_searches.get(&token) {                    existing.store(true, Ordering::Relaxed);                }                let flag = Arc::new(AtomicBool::new(false));                pending_fuzzy_searches.insert(token.clone(), flag.clone());                flag            }            None => Arc::new(AtomicBool::new(false)),        };        let results = match query.as_str() {            "" => vec![],            _ => run_fuzzy_file_search(query, roots, cancel_flag.clone()).await,        };        if let Some(token) = cancellation_token {            let mut pending_fuzzy_searches = self.pending_fuzzy_searches.lock().await;            if let Some(current_flag) = pending_fuzzy_searches.get(&token)                && Arc::ptr_eq(current_flag, &cancel_flag)            {                pending_fuzzy_searches.remove(&token);            }        }        Ok(FuzzyFileSearchResponse { files: results })    }    pub(crate) async fn fuzzy_file_search_session_start_response(        &self,        params: FuzzyFileSearchSessionStartParams,    ) -> Result<FuzzyFileSearchSessionStartResponse, JSONRPCErrorError> {        let FuzzyFileSearchSessionStartParams { session_id, roots } = params;        if session_id.is_empty() {            return Err(invalid_request("sessionId must not be empty"));        }        let session =            start_fuzzy_file_search_session(session_id.clone(), roots, self.outgoing.clone())                .map_err(|err| {                    internal_error(format!("failed to start fuzzy file search session: {err}"))                })?;        self.fuzzy_search_sessions            .lock()            .await            .insert(session_id, session);        Ok(FuzzyFileSearchSessionStartResponse {})    }    pub(crate) async fn fuzzy_file_search_session_update_response(        &self,        params: FuzzyFileSearchSessionUpdateParams,    ) -> Result<FuzzyFileSearchSessionUpdateResponse, JSONRPCErrorError> {        let FuzzyFileSearchSessionUpdateParams { session_id, query } = params;        let found = {            let sessions = self.fuzzy_search_sessions.lock().await;            if let Some(session) = sessions.get(&session_id) {                session.update_query(query);                true            } else {                false            }        };        if !found {            return Err(invalid_request(format!(                "fuzzy file search session not found: {session_id}"            )));        }        Ok(FuzzyFileSearchSessionUpdateResponse {})    }    pub(crate) async fn fuzzy_file_search_session_stop(        &self,        params: FuzzyFileSearchSessionStopParams,    ) -> Result<FuzzyFileSearchSessionStopResponse, JSONRPCErrorError> {        let FuzzyFileSearchSessionStopParams { session_id } = params;        self.fuzzy_search_sessions.lock().await.remove(&session_id);        Ok(FuzzyFileSearchSessionStopResponse {})    }}