codex: simplify tool search recovery (#30618)

This commit is contained in:
Anton Panasenko
2026-06-29 17:10:15 -07:00
parent b9f7dbf3c5
commit 4f20cfe962
9 changed files with 93 additions and 125 deletions

View File

@@ -5,8 +5,8 @@ use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseItem;
use codex_protocol::models::SearchToolCallParams;
use codex_protocol::openai_models::InputModality;
use codex_tools::INVALID_TOOL_SEARCH_QUERY;
use serde::Deserialize;
use std::borrow::Cow;
use std::collections::HashSet;
use tracing::info;
use tracing::warn;
@@ -16,6 +16,7 @@ use crate::util::error_or_panic;
const IMAGE_CONTENT_OMITTED_PLACEHOLDER: &str =
"image content omitted because you do not support image input";
const INVALID_TOOL_SEARCH_QUERY: &str = "[invalid tool_search arguments omitted]";
// Changing this value would change model-visible IDs and invalidate prompt caches.
const SYNTHETIC_OUTPUT_ID_NAMESPACE: Uuid = Uuid::from_u128(0x90d38d3e_6a5b_4d52_bfe2_2f1e634bfac4);
@@ -24,84 +25,67 @@ const SYNTHETIC_OUTPUT_ID_NAMESPACE: Uuid = Uuid::from_u128(0x90d38d3e_6a5b_4d52
/// Client-executed tool search is the only case handled here because its arguments have a fixed
/// schema that can be safely deserialized and reserialized. Function-call arguments and custom
/// tool input are intentionally preserved as schema-specific or freeform strings.
pub(crate) fn canonicalize_for_durable_history(
pub(crate) fn canonicalize_for_durable_history<'a>(
thread_id: &ThreadId,
items: &[ResponseItem],
) -> Option<Vec<ResponseItem>> {
items: &'a [ResponseItem],
) -> Cow<'a, [ResponseItem]> {
if !items.iter().any(|item| {
matches!(
item,
ResponseItem::ToolSearchCall { execution, .. } if execution == "client"
)
}) {
return None;
return Cow::Borrowed(items);
}
Some(
items
.iter()
.filter_map(|item| {
let ResponseItem::ToolSearchCall {
id,
let mut canonical_items = items.to_vec();
canonical_items.retain_mut(|item| {
let ResponseItem::ToolSearchCall {
call_id,
execution,
arguments,
..
} = item
else {
return true;
};
if execution != "client" {
return true;
}
let Some(call_id) = call_id.as_deref().filter(|call_id| !call_id.is_empty()) else {
warn!(
%thread_id,
"dropping client tool_search call with missing call_id from durable history"
);
return false;
};
*arguments = match SearchToolCallParams::deserialize(&*arguments) {
Ok(params) => serde_json::to_value(params).unwrap_or_else(|err| {
warn!(
%thread_id,
call_id,
status,
execution,
arguments,
internal_chat_message_metadata_passthrough,
} = item
else {
return Some(item.clone());
};
if execution != "client" {
return Some(item.clone());
}
let Some(call_id) = call_id.as_deref().filter(|call_id| !call_id.is_empty()) else {
warn!(
%thread_id,
"dropping client tool_search call with missing call_id from durable history"
);
return None;
};
let canonical_arguments = match SearchToolCallParams::deserialize(arguments) {
Ok(params) => match serde_json::to_value(params) {
Ok(canonical_arguments) => canonical_arguments,
Err(err) => {
warn!(
%thread_id,
call_id,
%err,
"failed to serialize canonical client tool_search arguments"
);
invalid_tool_search_arguments()
}
},
Err(err) => {
warn!(
%thread_id,
call_id,
error_category = ?err.classify(),
error_line = err.line(),
error_column = err.column(),
"replacing malformed client tool_search arguments in durable history"
);
invalid_tool_search_arguments()
}
};
Some(ResponseItem::ToolSearchCall {
id: id.clone(),
call_id: Some(call_id.to_string()),
status: status.clone(),
execution: execution.clone(),
arguments: canonical_arguments,
internal_chat_message_metadata_passthrough:
internal_chat_message_metadata_passthrough.clone(),
})
})
.collect(),
)
%err,
"failed to serialize canonical client tool_search arguments"
);
invalid_tool_search_arguments()
}),
Err(err) => {
warn!(
%thread_id,
call_id,
error_category = ?err.classify(),
error_line = err.line(),
error_column = err.column(),
"replacing malformed client tool_search arguments in durable history"
);
invalid_tool_search_arguments()
}
};
true
});
Cow::Owned(canonical_items)
}
fn invalid_tool_search_arguments() -> serde_json::Value {

View File

@@ -2709,28 +2709,24 @@ impl Session {
}
}
/// Prepares the history and raw-response views used when recording conversation items.
fn prepare_conversation_items_for_recording<'a>(
/// Records conversation items: append to history, persist to rollout, and
/// notify clients observing raw response items.
pub(crate) fn prepare_conversation_items_for_history<'a>(
&self,
turn_context: &TurnContext,
items: &'a [ResponseItem],
) -> (Cow<'a, [ResponseItem]>, Option<Vec<ResponseItem>>) {
) -> Cow<'a, [ResponseItem]> {
let mut items = Cow::Borrowed(items);
prepare_response_items(items.to_mut());
// Most response items get their passthrough turn ID at the durable history boundary.
for item in items.to_mut() {
item.set_turn_id_if_missing(&turn_context.sub_id);
}
let raw_response_items = if turn_context.config.features.enabled(Feature::ItemIds) {
if turn_context.config.features.enabled(Feature::ItemIds) {
Self::assign_missing_response_item_ids(items)
} else {
items
};
let history_items = crate::context_manager::canonicalize_for_durable_history(
&self.thread_id,
raw_response_items.as_ref(),
);
(raw_response_items, history_items)
}
}
fn assign_missing_response_item_ids(items: Cow<'_, [ResponseItem]>) -> Cow<'_, [ResponseItem]> {
@@ -2781,9 +2777,12 @@ impl Session {
turn_context: &TurnContext,
items: &[ResponseItem],
) {
let (raw_items, history_items) =
self.prepare_conversation_items_for_recording(turn_context, items);
let history_items = history_items.as_deref().unwrap_or(raw_items.as_ref());
let raw_items = self.prepare_conversation_items_for_history(turn_context, items);
let history_items = crate::context_manager::canonicalize_for_durable_history(
&self.thread_id,
raw_items.as_ref(),
);
let history_items = history_items.as_ref();
{
let mut state = self.state.lock().await;
state
@@ -2886,11 +2885,11 @@ impl Session {
) {
communication.set_turn_id_if_missing(&turn_context.sub_id);
let response_item = communication.to_model_input_item();
let (raw_items, history_items) = self.prepare_conversation_items_for_recording(
let items = self.prepare_conversation_items_for_history(
turn_context,
std::slice::from_ref(&response_item),
);
let items = history_items.as_deref().unwrap_or(raw_items.as_ref());
let items = items.as_ref();
let response_item = items[0].clone();
{
let mut state = self.state.lock().await;
@@ -2907,8 +2906,7 @@ impl Session {
RolloutItem::ResponseItem(response_item),
])
.await;
self.send_raw_response_items(turn_context, raw_items.as_ref())
.await;
self.send_raw_response_items(turn_context, items).await;
}
async fn maybe_warn_on_server_model_mismatch(

View File

@@ -1,5 +1,6 @@
use super::*;
use codex_protocol::models::DEFAULT_IMAGE_DETAIL;
use codex_protocol::models::SearchToolCallParams;
use core_test_support::assert_regex_match;
use pretty_assertions::assert_eq;
use serde_json::json;
@@ -320,7 +321,10 @@ fn custom_tool_calls_can_derive_text_from_content_items() {
#[test]
fn tool_search_payloads_roundtrip_as_tool_search_outputs() {
let payload = ToolPayload::ToolSearch {
arguments: json!({ "query": "calendar" }),
arguments: SearchToolCallParams {
query: "calendar".to_string(),
limit: None,
},
};
let response = ToolSearchOutput {
tools: vec![LoadableToolSpec::Function(codex_tools::ResponsesApiTool {

View File

@@ -10,7 +10,6 @@ use bm25::Document;
use bm25::Language;
use bm25::SearchEngine;
use bm25::SearchEngineBuilder;
use codex_protocol::models::SearchToolCallParams;
use codex_tools::LoadableToolSpec;
use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT;
use codex_tools::TOOL_SEARCH_TOOL_NAME;
@@ -19,7 +18,6 @@ use codex_tools::ToolSearchEntry;
use codex_tools::ToolSearchInfo;
use codex_tools::ToolSpec;
use codex_tools::coalesce_loadable_tool_specs;
use serde::Deserialize;
use std::sync::Arc;
use std::sync::Mutex;
use tracing::instrument;
@@ -121,7 +119,7 @@ impl ToolSearchHandler {
) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
let ToolInvocation { payload, .. } = invocation;
let arguments = match payload {
let args = match payload {
ToolPayload::ToolSearch { arguments } => arguments,
_ => {
return Err(FunctionCallError::Fatal(format!(
@@ -129,14 +127,6 @@ impl ToolSearchHandler {
)));
}
};
let args = SearchToolCallParams::deserialize(&arguments).map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to parse tool_search arguments: {:?} validation at line {}, column {}",
err.classify(),
err.line(),
err.column()
))
})?;
let query = args.query.trim();
if query.is_empty() {

View File

@@ -12,6 +12,7 @@ use crate::tools::spec_plan::build_tool_router;
use codex_mcp::ToolInfo;
use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_protocol::models::ResponseItem;
use codex_protocol::models::SearchToolCallParams;
use codex_tools::DiscoverableTool;
use codex_tools::ToolCall as ExtensionToolCall;
use codex_tools::ToolExecutor;
@@ -130,11 +131,19 @@ impl ToolRouter {
execution,
arguments,
..
} if execution == "client" && !call_id.is_empty() => Some(ToolCall {
tool_name: ToolName::plain("tool_search"),
call_id,
payload: ToolPayload::ToolSearch { arguments },
}),
} if execution == "client" && !call_id.is_empty() => {
let arguments = serde_json::from_value(arguments).unwrap_or(SearchToolCallParams {
// Enter the handler's normal validation failure path without retaining
// or logging the malformed payload.
query: String::new(),
limit: None,
});
Some(ToolCall {
tool_name: ToolName::plain("tool_search"),
call_id,
payload: ToolPayload::ToolSearch { arguments },
})
}
ResponseItem::ToolSearchCall { .. } => None,
ResponseItem::CustomToolCall {
name,

View File

@@ -9,15 +9,12 @@ use crate::tools::context::ToolCallSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
use codex_protocol::models::SearchToolCallParams;
use codex_rollout_trace::ExecutionStatus;
use codex_rollout_trace::ToolDispatchInvocation;
use codex_rollout_trace::ToolDispatchPayload;
use codex_rollout_trace::ToolDispatchRequester;
use codex_rollout_trace::ToolDispatchResult;
use codex_rollout_trace::ToolDispatchTraceContext;
use codex_tools::INVALID_TOOL_SEARCH_QUERY;
use serde::Deserialize;
/// Keeps registry early-return paths paired with trace end events.
pub(crate) struct ToolDispatchTrace {
@@ -108,15 +105,9 @@ fn tool_dispatch_payload(payload: &ToolPayload) -> ToolDispatchPayload {
ToolPayload::Function { arguments } => ToolDispatchPayload::Function {
arguments: arguments.clone(),
},
ToolPayload::ToolSearch { arguments } => {
let arguments = SearchToolCallParams::deserialize(arguments).unwrap_or_else(|_| {
SearchToolCallParams {
query: INVALID_TOOL_SEARCH_QUERY.to_string(),
limit: None,
}
});
ToolDispatchPayload::ToolSearch { arguments }
}
ToolPayload::ToolSearch { arguments } => ToolDispatchPayload::ToolSearch {
arguments: arguments.clone(),
},
ToolPayload::Custom { input } => ToolDispatchPayload::Custom {
input: input.clone(),
},

View File

@@ -1,5 +1,3 @@
#![allow(clippy::unwrap_used)]
use anyhow::Result;
use codex_protocol::models::InternalChatMessageMetadataPassthrough;
use codex_protocol::models::ResponseItem;

View File

@@ -98,7 +98,6 @@ pub use tool_executor::ToolExecutorFuture;
pub use tool_executor::ToolExposure;
pub use tool_output::JsonToolOutput;
pub use tool_output::ToolOutput;
pub use tool_payload::INVALID_TOOL_SEARCH_QUERY;
pub use tool_payload::ToolPayload;
pub use tool_search::ToolSearchEntry;
pub use tool_search::ToolSearchInfo;

View File

@@ -1,15 +1,12 @@
use std::borrow::Cow;
use codex_protocol::models::SearchToolCallParams;
use serde::Deserialize;
pub const INVALID_TOOL_SEARCH_QUERY: &str = "[invalid tool_search arguments omitted]";
/// Canonical payload shapes accepted by model-visible tool runtimes.
#[derive(Clone, Debug, PartialEq)]
pub enum ToolPayload {
Function { arguments: String },
ToolSearch { arguments: serde_json::Value },
ToolSearch { arguments: SearchToolCallParams },
Custom { input: String },
}
@@ -17,9 +14,7 @@ impl ToolPayload {
pub fn log_payload(&self) -> Cow<'_, str> {
match self {
ToolPayload::Function { arguments } => Cow::Borrowed(arguments),
ToolPayload::ToolSearch { arguments } => SearchToolCallParams::deserialize(arguments)
.map(|arguments| Cow::Owned(arguments.query))
.unwrap_or(Cow::Borrowed(INVALID_TOOL_SEARCH_QUERY)),
ToolPayload::ToolSearch { arguments } => Cow::Owned(arguments.query.clone()),
ToolPayload::Custom { input } => Cow::Borrowed(input),
}
}