Move tool search RPC to dedicated processor

This commit is contained in:
Steven Lee
2026-07-06 20:02:27 +00:00
parent f1825e8741
commit 266e55bcd7
4 changed files with 105 additions and 89 deletions
+5 -1
View File
@@ -38,6 +38,7 @@ use crate::request_processors::RemoteControlRequestProcessor;
use crate::request_processors::SearchRequestProcessor;
use crate::request_processors::ThreadGoalRequestProcessor;
use crate::request_processors::ThreadRequestProcessor;
use crate::request_processors::ToolSearchRequestProcessor;
use crate::request_processors::TurnRequestProcessor;
use crate::request_processors::WindowsSandboxRequestProcessor;
use crate::request_serialization::QueuedInitializedRequest;
@@ -206,6 +207,7 @@ pub(crate) struct MessageProcessor {
search_processor: SearchRequestProcessor,
thread_goal_processor: ThreadGoalRequestProcessor,
thread_processor: ThreadRequestProcessor,
tool_search_processor: ToolSearchRequestProcessor,
turn_processor: TurnRequestProcessor,
windows_sandbox_processor: WindowsSandboxRequestProcessor,
request_serialization_queues: RequestSerializationQueues,
@@ -471,6 +473,7 @@ impl MessageProcessor {
);
let remote_control_processor = RemoteControlRequestProcessor::new(remote_control_handle);
let search_processor = SearchRequestProcessor::new(outgoing.clone());
let tool_search_processor = ToolSearchRequestProcessor::new(Arc::clone(&thread_manager));
let thread_goal_processor = ThreadGoalRequestProcessor::new(
Arc::clone(&thread_manager),
outgoing.clone(),
@@ -576,6 +579,7 @@ impl MessageProcessor {
search_processor,
thread_goal_processor,
thread_processor,
tool_search_processor,
turn_processor,
windows_sandbox_processor,
request_serialization_queues: RequestSerializationQueues::default(),
@@ -1387,7 +1391,7 @@ impl MessageProcessor {
.await
}
ClientRequest::ToolSearchSearch { params, .. } => {
self.mcp_processor.tool_search_search(params).await
self.tool_search_processor.tool_search_search(params).await
}
ClientRequest::McpResourceRead { params, .. } => {
self.mcp_processor
@@ -517,6 +517,7 @@ mod remote_control_processor;
mod search;
mod thread_processor;
mod token_usage_replay;
mod tool_search_processor;
mod turn_processor;
mod windows_sandbox_processor;
@@ -540,6 +541,7 @@ pub(crate) use remote_control_processor::RemoteControlRequestProcessor;
pub(crate) use search::SearchRequestProcessor;
pub(crate) use thread_goal_processor::ThreadGoalRequestProcessor;
pub(crate) use thread_processor::ThreadRequestProcessor;
pub(crate) use tool_search_processor::ToolSearchRequestProcessor;
pub(crate) use turn_processor::TurnRequestProcessor;
pub(crate) use windows_sandbox_processor::WindowsSandboxRequestProcessor;
@@ -53,15 +53,6 @@ impl McpRequestProcessor {
.map(|()| None)
}
pub(crate) async fn tool_search_search(
&self,
params: ToolSearchSearchParams,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
self.tool_search_search_response(params)
.await
.map(|response| Some(response.into()))
}
pub(crate) async fn mcp_resource_read(
&self,
request_id: &ConnectionRequestId,
@@ -118,32 +109,6 @@ impl McpRequestProcessor {
Ok((thread_id, thread))
}
async fn tool_search_search_response(
&self,
params: ToolSearchSearchParams,
) -> Result<ToolSearchSearchResponse, JSONRPCErrorError> {
let ToolSearchSearchParams {
thread_id,
query,
limit,
} = params;
let query = query.trim();
if query.is_empty() {
return Err(invalid_request("query must not be empty"));
}
let limit = limit.unwrap_or(codex_tools::TOOL_SEARCH_DEFAULT_LIMIT as u32);
if limit == 0 {
return Err(invalid_request("limit must be greater than zero"));
}
let (_, thread) = self.load_thread(&thread_id).await?;
let result = thread
.search_tool_search(query, limit as usize)
.await
.map_err(|err| internal_error(format!("failed to search tool_search: {err}")))?;
tool_search_result_to_response(result)
}
async fn mcp_server_oauth_login_response(
&self,
params: McpServerOauthLoginParams,
@@ -510,59 +475,6 @@ impl McpRequestProcessor {
}
}
fn tool_search_result_to_response(
result: codex_core::ToolSearchDebugResult,
) -> Result<ToolSearchSearchResponse, JSONRPCErrorError> {
Ok(ToolSearchSearchResponse {
indexed_tool_count: usize_to_u32(result.indexed_tool_count, "indexedToolCount")?,
matching_tool_count: usize_to_u32(result.matching_tool_count, "matchingToolCount")?,
requested_limit: usize_to_u32(result.requested_limit, "requestedLimit")?,
effective_limit: usize_to_u32(result.effective_limit, "effectiveLimit")?,
top_k_truncated: result.top_k_truncated,
tools: serialize_tool_specs(result.tools)?,
results: result
.results
.into_iter()
.map(tool_search_result_entry_to_response)
.collect::<Result<Vec<_>, _>>()?,
})
}
fn tool_search_result_entry_to_response(
result: codex_core::ToolSearchDebugResultEntry,
) -> Result<ToolSearchSearchResult, JSONRPCErrorError> {
Ok(ToolSearchSearchResult {
rank: usize_to_u32(result.rank, "rank")?,
index: usize_to_u32(result.index, "index")?,
score: result.score,
source_name: result.source_name,
source_description: result.source_description,
searchable_text: result.searchable_text,
tools: serialize_tool_specs(result.tools)?,
})
}
fn serialize_tool_specs(
tools: Vec<codex_tools::LoadableToolSpec>,
) -> Result<Vec<serde_json::Value>, JSONRPCErrorError> {
tools
.into_iter()
.map(|tool| {
serde_json::to_value(tool).map_err(|err| {
internal_error(format!("failed to serialize tool_search tool: {err}"))
})
})
.collect()
}
fn usize_to_u32(value: usize, field_name: &str) -> Result<u32, JSONRPCErrorError> {
u32::try_from(value).map_err(|err| {
internal_error(format!(
"tool_search {field_name} value {value} does not fit in u32: {err}"
))
})
}
fn with_mcp_tool_call_thread_id_meta(
meta: Option<serde_json::Value>,
thread_id: &str,
@@ -0,0 +1,98 @@
use super::*;
#[derive(Clone)]
pub(crate) struct ToolSearchRequestProcessor {
thread_manager: Arc<ThreadManager>,
}
impl ToolSearchRequestProcessor {
pub(crate) fn new(thread_manager: Arc<ThreadManager>) -> Self {
Self { thread_manager }
}
pub(crate) async fn tool_search_search(
&self,
params: ToolSearchSearchParams,
) -> Result<Option<ClientResponsePayload>, JSONRPCErrorError> {
let ToolSearchSearchParams {
thread_id,
query,
limit,
} = params;
let query = query.trim();
if query.is_empty() {
return Err(invalid_request("query must not be empty"));
}
let limit = limit.unwrap_or(codex_tools::TOOL_SEARCH_DEFAULT_LIMIT as u32);
if limit == 0 {
return Err(invalid_request("limit must be greater than zero"));
}
let thread_id = ThreadId::from_string(&thread_id)
.map_err(|err| invalid_request(format!("invalid thread id: {err}")))?;
let thread = self
.thread_manager
.get_thread(thread_id)
.await
.map_err(|_| invalid_request(format!("thread not found: {thread_id}")))?;
let result = thread
.search_tool_search(query, limit as usize)
.await
.map_err(|err| internal_error(format!("failed to search tool_search: {err}")))?;
tool_search_result_to_response(result).map(|response| Some(response.into()))
}
}
fn tool_search_result_to_response(
result: codex_core::ToolSearchDebugResult,
) -> Result<ToolSearchSearchResponse, JSONRPCErrorError> {
Ok(ToolSearchSearchResponse {
indexed_tool_count: usize_to_u32(result.indexed_tool_count, "indexedToolCount")?,
matching_tool_count: usize_to_u32(result.matching_tool_count, "matchingToolCount")?,
requested_limit: usize_to_u32(result.requested_limit, "requestedLimit")?,
effective_limit: usize_to_u32(result.effective_limit, "effectiveLimit")?,
top_k_truncated: result.top_k_truncated,
tools: serialize_tool_specs(result.tools)?,
results: result
.results
.into_iter()
.map(tool_search_result_entry_to_response)
.collect::<Result<Vec<_>, _>>()?,
})
}
fn tool_search_result_entry_to_response(
result: codex_core::ToolSearchDebugResultEntry,
) -> Result<ToolSearchSearchResult, JSONRPCErrorError> {
Ok(ToolSearchSearchResult {
rank: usize_to_u32(result.rank, "rank")?,
index: usize_to_u32(result.index, "index")?,
score: result.score,
source_name: result.source_name,
source_description: result.source_description,
searchable_text: result.searchable_text,
tools: serialize_tool_specs(result.tools)?,
})
}
fn serialize_tool_specs(
tools: Vec<codex_tools::LoadableToolSpec>,
) -> Result<Vec<serde_json::Value>, JSONRPCErrorError> {
tools
.into_iter()
.map(|tool| {
serde_json::to_value(tool).map_err(|err| {
internal_error(format!("failed to serialize tool_search tool: {err}"))
})
})
.collect()
}
fn usize_to_u32(value: usize, field_name: &str) -> Result<u32, JSONRPCErrorError> {
u32::try_from(value).map_err(|err| {
internal_error(format!(
"tool_search {field_name} value {value} does not fit in u32: {err}"
))
})
}