File size: 5,181 Bytes
afa0cbf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | 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 {})
}
}
|