app-server: add fuzzy search sessions for streaming file search

codex-rs/app-server-protocol/src/protocol/common.rs

@@ -350,6 +350,18 @@ client_request_definitions! {
         params: FuzzyFileSearchParams,
         response: FuzzyFileSearchResponse,
     },
+    FuzzyFileSearchSessionStart => "fuzzyFileSearch/sessionStart" {
+        params: FuzzyFileSearchSessionStartParams,
+        response: FuzzyFileSearchSessionStartResponse,
+    },
+    FuzzyFileSearchSessionUpdate => "fuzzyFileSearch/sessionUpdate" {
+        params: FuzzyFileSearchSessionUpdateParams,
+        response: FuzzyFileSearchSessionUpdateResponse,
+    },
+    FuzzyFileSearchSessionStop => "fuzzyFileSearch/sessionStop" {
+        params: FuzzyFileSearchSessionStopParams,
+        response: FuzzyFileSearchSessionStopResponse,
+    },
     /// Execute a command (argv vector) under the server's sandbox.
     ExecOneOffCommand {
         params: v1::ExecOneOffCommandParams,
@@ -594,6 +606,47 @@ pub struct FuzzyFileSearchResponse {
     pub files: Vec<FuzzyFileSearchResult>,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
+#[serde(rename_all = "camelCase")]
+#[ts(rename_all = "camelCase")]
+pub struct FuzzyFileSearchSessionStartParams {
+    pub session_id: String,
+    pub roots: Vec<String>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, Default)]
+pub struct FuzzyFileSearchSessionStartResponse {}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
+#[serde(rename_all = "camelCase")]
+#[ts(rename_all = "camelCase")]
+pub struct FuzzyFileSearchSessionUpdateParams {
+    pub session_id: String,
+    pub query: String,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, Default)]
+pub struct FuzzyFileSearchSessionUpdateResponse {}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
+#[serde(rename_all = "camelCase")]
+#[ts(rename_all = "camelCase")]
+pub struct FuzzyFileSearchSessionStopParams {
+    pub session_id: String,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, Default)]
+pub struct FuzzyFileSearchSessionStopResponse {}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
+#[serde(rename_all = "camelCase")]
+#[ts(rename_all = "camelCase")]
+pub struct FuzzyFileSearchSessionUpdatedNotification {
+    pub session_id: String,
+    pub query: String,
+    pub files: Vec<FuzzyFileSearchResult>,
+}
+
 server_notification_definitions! {
     /// NEW NOTIFICATIONS
     Error => "error" (v2::ErrorNotification),
@@ -625,6 +678,7 @@ server_notification_definitions! {
     ContextCompacted => "thread/compacted" (v2::ContextCompactedNotification),
     DeprecationNotice => "deprecationNotice" (v2::DeprecationNoticeNotification),
     ConfigWarning => "configWarning" (v2::ConfigWarningNotification),
+    FuzzyFileSearchSessionUpdated => "fuzzyFileSearch/sessionUpdated" (FuzzyFileSearchSessionUpdatedNotification),
 
     /// Notifies the user of world-writable directories on Windows, which cannot be protected by the sandbox.
     WindowsWorldWritableWarning => "windows/worldWritableWarning" (v2::WindowsWorldWritableWarningNotification),

codex-rs/app-server/src/codex_message_processor.rs

@@ -1,7 +1,9 @@
 use crate::bespoke_event_handling::apply_bespoke_event_handling;
 use crate::error_code::INTERNAL_ERROR_CODE;
 use crate::error_code::INVALID_REQUEST_ERROR_CODE;
+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::models::supported_models;
 use crate::outgoing_message::OutgoingMessageSender;
 use crate::outgoing_message::OutgoingNotification;
@@ -38,6 +40,12 @@ use codex_app_server_protocol::ForkConversationParams;
 use codex_app_server_protocol::ForkConversationResponse;
 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::GetAccountParams;
 use codex_app_server_protocol::GetAccountRateLimitsResponse;
 use codex_app_server_protocol::GetAccountResponse;
@@ -272,6 +280,7 @@ pub(crate) struct CodexMessageProcessor {
     pending_rollbacks: PendingRollbacks,
     turn_summary_store: TurnSummaryStore,
     pending_fuzzy_searches: Arc<Mutex<HashMap<String, Arc<AtomicBool>>>>,
+    fuzzy_search_sessions: Arc<Mutex<HashMap<String, FuzzyFileSearchSession>>>,
     feedback: CodexFeedback,
 }
 
@@ -328,6 +337,7 @@ impl CodexMessageProcessor {
             pending_rollbacks: Arc::new(Mutex::new(HashMap::new())),
             turn_summary_store: Arc::new(Mutex::new(HashMap::new())),
             pending_fuzzy_searches: Arc::new(Mutex::new(HashMap::new())),
+            fuzzy_search_sessions: Arc::new(Mutex::new(HashMap::new())),
             feedback,
         }
     }
@@ -581,6 +591,18 @@ impl CodexMessageProcessor {
             ClientRequest::FuzzyFileSearch { request_id, params } => {
                 self.fuzzy_file_search(request_id, params).await;
             }
+            ClientRequest::FuzzyFileSearchSessionStart { request_id, params } => {
+                self.fuzzy_file_search_session_start(request_id, params)
+                    .await;
+            }
+            ClientRequest::FuzzyFileSearchSessionUpdate { request_id, params } => {
+                self.fuzzy_file_search_session_update(request_id, params)
+                    .await;
+            }
+            ClientRequest::FuzzyFileSearchSessionStop { request_id, params } => {
+                self.fuzzy_file_search_session_stop(request_id, params)
+                    .await;
+            }
             ClientRequest::OneOffCommandExec { request_id, params } => {
                 self.exec_one_off_command(request_id, params).await;
             }
@@ -4591,6 +4613,89 @@ impl CodexMessageProcessor {
         self.outgoing.send_response(request_id, response).await;
     }
 
+    async fn fuzzy_file_search_session_start(
+        &mut self,
+        request_id: RequestId,
+        params: FuzzyFileSearchSessionStartParams,
+    ) {
+        let FuzzyFileSearchSessionStartParams { session_id, roots } = params;
+        if session_id.is_empty() {
+            let error = JSONRPCErrorError {
+                code: INVALID_REQUEST_ERROR_CODE,
+                message: "sessionId must not be empty".to_string(),
+                data: None,
+            };
+            self.outgoing.send_error(request_id, error).await;
+            return;
+        }
+
+        let session =
+            start_fuzzy_file_search_session(session_id.clone(), roots, self.outgoing.clone());
+        match session {
+            Ok(session) => {
+                let mut sessions = self.fuzzy_search_sessions.lock().await;
+                sessions.insert(session_id, session);
+                self.outgoing
+                    .send_response(request_id, FuzzyFileSearchSessionStartResponse {})
+                    .await;
+            }
+            Err(err) => {
+                let error = JSONRPCErrorError {
+                    code: INTERNAL_ERROR_CODE,
+                    message: format!("failed to start fuzzy file search session: {err}"),
+                    data: None,
+                };
+                self.outgoing.send_error(request_id, error).await;
+            }
+        }
+    }
+
+    async fn fuzzy_file_search_session_update(
+        &mut self,
+        request_id: RequestId,
+        params: FuzzyFileSearchSessionUpdateParams,
+    ) {
+        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 {
+            let error = JSONRPCErrorError {
+                code: INVALID_REQUEST_ERROR_CODE,
+                message: format!("fuzzy file search session not found: {session_id}"),
+                data: None,
+            };
+            self.outgoing.send_error(request_id, error).await;
+            return;
+        }
+
+        self.outgoing
+            .send_response(request_id, FuzzyFileSearchSessionUpdateResponse {})
+            .await;
+    }
+
+    async fn fuzzy_file_search_session_stop(
+        &mut self,
+        request_id: RequestId,
+        params: FuzzyFileSearchSessionStopParams,
+    ) {
+        let FuzzyFileSearchSessionStopParams { session_id } = params;
+        {
+            let mut sessions = self.fuzzy_search_sessions.lock().await;
+            sessions.remove(&session_id);
+        }
+
+        self.outgoing
+            .send_response(request_id, FuzzyFileSearchSessionStopResponse {})
+            .await;
+    }
+
     async fn upload_feedback(&self, request_id: RequestId, params: FeedbackUploadParams) {
         if !self.config.feedback_enabled {
             let error = JSONRPCErrorError {

codex-rs/app-server/src/fuzzy_file_search.rs

@@ -1,12 +1,18 @@
 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::FuzzyFileSearchResult;
+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;
 
@@ -77,3 +83,144 @@ pub(crate) async fn run_fuzzy_file_search(
 
     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 query.is_empty() || snapshot.query != query {
+            return;
+        }
+
+        let files = 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;
+        });
+    }
+}
+
+impl file_search::SessionReporter for SessionReporterImpl {
+    fn on_update(&self, snapshot: &file_search::FileSearchSnapshot) {
+        self.send_snapshot(snapshot);
+    }
+
+    fn on_complete(&self) {}
+}
+
+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(),
+                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
+}