From 4f20cfe96224eead05ae26eb69ec814e4217a192 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Mon, 29 Jun 2026 17:10:15 -0700 Subject: [PATCH] codex: simplify tool search recovery (#30618) --- .../core/src/context_manager/normalize.rs | 122 ++++++++---------- codex-rs/core/src/session/mod.rs | 32 +++-- codex-rs/core/src/tools/context_tests.rs | 6 +- .../core/src/tools/handlers/tool_search.rs | 12 +- codex-rs/core/src/tools/router.rs | 19 ++- .../core/src/tools/tool_dispatch_trace.rs | 15 +-- .../core/tests/suite/tool_search_recovery.rs | 2 - codex-rs/tools/src/lib.rs | 1 - codex-rs/tools/src/tool_payload.rs | 9 +- 9 files changed, 93 insertions(+), 125 deletions(-) diff --git a/codex-rs/core/src/context_manager/normalize.rs b/codex-rs/core/src/context_manager/normalize.rs index 2208563cde..5ba7382574 100644 --- a/codex-rs/core/src/context_manager/normalize.rs +++ b/codex-rs/core/src/context_manager/normalize.rs @@ -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> { + 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 { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 3ab78aa47c..6e08300e67 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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>) { + ) -> 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( diff --git a/codex-rs/core/src/tools/context_tests.rs b/codex-rs/core/src/tools/context_tests.rs index b9f90e08d4..6ac5b6978f 100644 --- a/codex-rs/core/src/tools/context_tests.rs +++ b/codex-rs/core/src/tools/context_tests.rs @@ -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 { diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 3b55fc7727..272da2f226 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -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, 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() { diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 2c66232169..9d6326b3ab 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -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, diff --git a/codex-rs/core/src/tools/tool_dispatch_trace.rs b/codex-rs/core/src/tools/tool_dispatch_trace.rs index 267639d951..098dc8ddef 100644 --- a/codex-rs/core/src/tools/tool_dispatch_trace.rs +++ b/codex-rs/core/src/tools/tool_dispatch_trace.rs @@ -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(), }, diff --git a/codex-rs/core/tests/suite/tool_search_recovery.rs b/codex-rs/core/tests/suite/tool_search_recovery.rs index 8ee80613b3..fac3fca6d4 100644 --- a/codex-rs/core/tests/suite/tool_search_recovery.rs +++ b/codex-rs/core/tests/suite/tool_search_recovery.rs @@ -1,5 +1,3 @@ -#![allow(clippy::unwrap_used)] - use anyhow::Result; use codex_protocol::models::InternalChatMessageMetadataPassthrough; use codex_protocol::models::ResponseItem; diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 8a18e0a304..8bcad84981 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -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; diff --git a/codex-rs/tools/src/tool_payload.rs b/codex-rs/tools/src/tool_payload.rs index 14ad241d87..101d95a278 100644 --- a/codex-rs/tools/src/tool_payload.rs +++ b/codex-rs/tools/src/tool_payload.rs @@ -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), } }