diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d43724e5db..3fd2b9451a 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -216,6 +216,7 @@ pub(crate) mod session; pub(crate) mod step_context; pub(crate) mod time_reminder; mod token_budget; +mod tool_search_history; pub(crate) mod turn; pub(crate) mod turn_context; mod world_state; @@ -2709,16 +2710,14 @@ impl Session { } } - /// 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>( + fn prepare_raw_response_items<'a>( &self, turn_context: &TurnContext, items: &'a [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. + // Most response items get their passthrough turn ID at the recording boundary. for item in items.to_mut() { item.set_turn_id_if_missing(&turn_context.sub_id); } @@ -2777,18 +2776,22 @@ impl Session { turn_context: &TurnContext, items: &[ResponseItem], ) { - let items = self.prepare_conversation_items_for_history(turn_context, items); - let items = items.as_ref(); + let raw_items = self.prepare_raw_response_items(turn_context, items); + let history_items = + tool_search_history::prepare_for_durable_history(&self.thread_id, raw_items.as_ref()); { let mut state = self.state.lock().await; - state.current_time_reminder.note_recorded_items(items); + state + .current_time_reminder + .note_recorded_items(&history_items); state.record_items( - items.iter(), + history_items.iter(), turn_context.model_info.truncation_policy.into(), ); } - self.persist_rollout_response_items(items).await; - self.send_raw_response_items(turn_context, items).await; + self.persist_rollout_response_items(&history_items).await; + self.send_raw_response_items(turn_context, raw_items.as_ref()) + .await; } pub(crate) async fn record_step_world_state_if_changed( @@ -2878,17 +2881,16 @@ impl Session { ) { communication.set_turn_id_if_missing(&turn_context.sub_id); let response_item = communication.to_model_input_item(); - let items = self.prepare_conversation_items_for_history( - turn_context, - std::slice::from_ref(&response_item), - ); - let items = items.as_ref(); - let response_item = items[0].clone(); + let raw_items = + self.prepare_raw_response_items(turn_context, std::slice::from_ref(&response_item)); + let response_item = raw_items[0].clone(); { let mut state = self.state.lock().await; - state.current_time_reminder.note_recorded_items(items); + state + .current_time_reminder + .note_recorded_items(raw_items.as_ref()); state.record_items( - items.iter(), + raw_items.iter(), turn_context.model_info.truncation_policy.into(), ); } @@ -2899,7 +2901,8 @@ impl Session { RolloutItem::ResponseItem(response_item), ]) .await; - self.send_raw_response_items(turn_context, items).await; + self.send_raw_response_items(turn_context, raw_items.as_ref()) + .await; } async fn maybe_warn_on_server_model_mismatch( diff --git a/codex-rs/core/src/session/tool_search_history.rs b/codex-rs/core/src/session/tool_search_history.rs new file mode 100644 index 0000000000..600f57dace --- /dev/null +++ b/codex-rs/core/src/session/tool_search_history.rs @@ -0,0 +1,92 @@ +use codex_protocol::ThreadId; +use codex_protocol::models::ResponseItem; +use codex_protocol::models::SearchToolCallParams; +use serde::Deserialize; +use std::borrow::Cow; +use tracing::warn; + +pub(super) const INVALID_TOOL_SEARCH_QUERY: &str = "[invalid tool_search arguments omitted]"; + +pub(super) fn prepare_for_durable_history<'a>( + thread_id: &ThreadId, + items: &'a [ResponseItem], +) -> Cow<'a, [ResponseItem]> { + if !items.iter().any(|item| { + matches!( + item, + ResponseItem::ToolSearchCall { execution, .. } if execution == "client" + ) + }) { + return Cow::Borrowed(items); + } + + Cow::Owned( + items + .iter() + .filter_map(|item| { + let ResponseItem::ToolSearchCall { + 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(), + ) +} + +fn invalid_tool_search_arguments() -> serde_json::Value { + serde_json::json!({ "query": INVALID_TOOL_SEARCH_QUERY }) +} diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index b07c180b16..60ea244161 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -485,11 +485,25 @@ pub(crate) async fn handle_output_item_done( } // The tool request should be answered directly (or was denied); push that response into the transcript. Err(FunctionCallError::RespondToModel(message)) => { - let response = ResponseInputItem::FunctionCallOutput { - call_id: String::new(), - output: FunctionCallOutputPayload { - body: FunctionCallOutputBody::Text(message), - ..Default::default() + let response = match &item { + ResponseItem::ToolSearchCall { + call_id: Some(call_id), + execution, + .. + } if execution == "client" && !call_id.is_empty() => { + ResponseInputItem::ToolSearchOutput { + call_id: call_id.clone(), + status: "completed".to_string(), + execution: "client".to_string(), + tools: Vec::new(), + } + } + _ => ResponseInputItem::FunctionCallOutput { + call_id: String::new(), + output: FunctionCallOutputPayload { + body: FunctionCallOutputBody::Text(message), + ..Default::default() + }, }, }; record_completed_response_item(ctx.sess.as_ref(), ctx.turn_context.as_ref(), &item) diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index aca86b124b..692801958d 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -18,6 +18,7 @@ use codex_tools::ToolCall as ExtensionToolCall; use codex_tools::ToolExecutor; use codex_tools::ToolName; use codex_tools::ToolSpec; +use serde::Deserialize; use std::sync::Arc; use std::sync::atomic::AtomicBool; use tokio_util::sync::CancellationToken; @@ -131,13 +132,15 @@ impl ToolRouter { execution, arguments, .. - } if execution == "client" => { - let arguments: SearchToolCallParams = - serde_json::from_value(arguments).map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to parse tool_search arguments: {err}" - )) - })?; + } if execution == "client" && !call_id.is_empty() => { + let arguments = 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() + )) + })?; Ok(Some(ToolCall { tool_name: ToolName::plain("tool_search"), call_id, diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 3e27463370..4afdca9b82 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -122,6 +122,7 @@ mod subagent_notifications; mod token_budget; mod tool_harness; mod tool_parallelism; +mod tool_search_recovery; mod tools; mod truncation; mod turn_state; diff --git a/codex-rs/core/tests/suite/tool_search_recovery.rs b/codex-rs/core/tests/suite/tool_search_recovery.rs new file mode 100644 index 0000000000..52b93033ca --- /dev/null +++ b/codex-rs/core/tests/suite/tool_search_recovery.rs @@ -0,0 +1,306 @@ +#![allow(clippy::unwrap_used)] + +use anyhow::Result; +use codex_protocol::models::InternalChatMessageMetadataPassthrough; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::user_input::UserInput; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once; +use core_test_support::responses::mount_sse_sequence; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::Map; +use serde_json::Value; +use serde_json::json; +use std::fs; +use std::path::Path; +use std::time::Duration; +use tokio::time::timeout; + +const CALL_ID: &str = "malformed-search-call"; +const VALID_CALL_ID: &str = "valid-search-call"; +const INVALID_TOOL_SEARCH_QUERY: &str = "[invalid tool_search arguments omitted]"; + +fn tool_search_item( + id: &str, + call_id: Option<&str>, + execution: &str, + arguments: Value, +) -> ResponseItem { + ResponseItem::ToolSearchCall { + id: Some(id.to_string()), + call_id: call_id.map(str::to_string), + status: Some("completed".to_string()), + execution: execution.to_string(), + arguments, + internal_chat_message_metadata_passthrough: Some(InternalChatMessageMetadataPassthrough { + turn_id: Some("source-turn".to_string()), + }), + } +} + +fn ev_response_item_done(item: &ResponseItem) -> Value { + json!({"type": "response.output_item.done", "item": item}) +} + +fn without_item_id(mut item: Value) -> Value { + item.as_object_mut() + .expect("response item should serialize as an object") + .remove("id"); + item +} + +fn rollout_response_items(path: &Path) -> Result> { + Ok(fs::read_to_string(path)? + .lines() + .map(serde_json::from_str::) + .collect::, _>>()? + .into_iter() + .filter_map(|line| match line.item { + RolloutItem::ResponseItem(item) => Some(item), + _ => None, + }) + .collect()) +} + +fn tool_search_call( + request: &core_test_support::responses::ResponsesRequest, + call_id: &str, +) -> Value { + request + .inputs_of_type("tool_search_call") + .into_iter() + .find(|item| item.get("call_id").and_then(Value::as_str) == Some(call_id)) + .expect("tool_search_call should be present") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn malformed_tool_search_is_repaired_before_persistence_and_resume() -> Result<()> { + let oversized_key = "multilingual-generated-gibberish-世界-".repeat(24); + let original_arguments = Value::Object(Map::from_iter([(oversized_key.clone(), json!([]))])); + let malformed_item = tool_search_item( + "malformed-item", + Some(CALL_ID), + "client", + original_arguments, + ); + let valid_item = tool_search_item( + "valid-item", + Some(VALID_CALL_ID), + "client", + json!({"query": "calendar", "limit": 3, "unknown": oversized_key.clone()}), + ); + let missing_id_item = tool_search_item( + "missing-call-id-item", + None, + "client", + json!({"query": "missing"}), + ); + let empty_id_item = tool_search_item( + "empty-call-id-item", + Some(""), + "client", + json!({"query": "empty"}), + ); + let server_item = tool_search_item( + "server-item", + None, + "server", + json!({"serverGenerated": {"nested": true}}), + ); + let canonical_malformed_item = tool_search_item( + "malformed-item", + Some(CALL_ID), + "client", + json!({"query": INVALID_TOOL_SEARCH_QUERY}), + ); + let canonical_valid_item = tool_search_item( + "valid-item", + Some(VALID_CALL_ID), + "client", + json!({"query": "calendar", "limit": 3}), + ); + let canonical_malformed_request = + without_item_id(serde_json::to_value(&canonical_malformed_item)?); + let canonical_valid_request = without_item_id(serde_json::to_value(&canonical_valid_item)?); + let server_request = without_item_id(serde_json::to_value(&server_item)?); + let original_items = vec![ + malformed_item.clone(), + valid_item, + missing_id_item, + empty_id_item, + server_item.clone(), + ]; + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let mut first_response = vec![ev_response_created("resp-1")]; + first_response.extend(original_items.iter().map(ev_response_item_done)); + first_response.push(ev_completed("resp-1")); + let turn_mock = mount_sse_sequence( + &server, + vec![ + sse(first_response), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-1", "recovered"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + let mut builder = test_codex(); + let test = builder.build_with_auto_env(&server).await?; + let rollout_path = test + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + let home = test.home.clone(); + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "trigger malformed tool search".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + + let mut raw_tool_search_items = Vec::new(); + loop { + let event = timeout(Duration::from_secs(30), test.codex.next_event()) + .await + .expect("timeout waiting for turn event") + .expect("event stream ended unexpectedly"); + match event.msg { + EventMsg::RawResponseItem(raw) + if matches!(&raw.item, ResponseItem::ToolSearchCall { .. }) => + { + raw_tool_search_items.push(raw.item); + } + EventMsg::TurnComplete(_) => break, + EventMsg::Error(error) => panic!("turn failed: {}", error.message), + _ => {} + } + } + assert_eq!(raw_tool_search_items, original_items); + + let requests = turn_mock.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + tool_search_call(&requests[1], CALL_ID), + canonical_malformed_request + ); + assert_eq!( + tool_search_call(&requests[1], VALID_CALL_ID), + canonical_valid_request + ); + let tool_search_output = requests[1].tool_search_output(CALL_ID); + assert_eq!( + tool_search_output.get("status").and_then(Value::as_str), + Some("completed") + ); + assert_eq!( + tool_search_output.get("execution").and_then(Value::as_str), + Some("client") + ); + assert_eq!(tool_search_output.get("tools"), Some(&json!([]))); + assert_eq!( + requests[1].tool_search_output(VALID_CALL_ID).get("tools"), + Some(&json!([])) + ); + let persisted_calls = requests[1].inputs_of_type("tool_search_call"); + assert_eq!(persisted_calls.len(), 3); + assert!(persisted_calls.contains(&server_request)); + assert_eq!( + persisted_calls + .iter() + .filter_map(|item| item.get("call_id").and_then(Value::as_str)) + .collect::>(), + vec![CALL_ID, VALID_CALL_ID] + ); + assert!( + requests[1] + .inputs_of_type("function_call_output") + .iter() + .all(|item| item.get("call_id").and_then(Value::as_str) != Some("")) + ); + + let rollout_text = fs::read_to_string(&rollout_path)?; + assert!(!rollout_text.contains(&oversized_key)); + let rollout_items = rollout_response_items(&rollout_path)?; + assert!(rollout_items.contains(&canonical_malformed_item)); + assert!(rollout_items.contains(&canonical_valid_item)); + assert!(rollout_items.contains(&server_item)); + assert!(rollout_items.iter().all(|item| { + !matches!( + item, + ResponseItem::ToolSearchCall { + id: Some(id), + .. + } if id == "missing-call-id-item" || id == "empty-call-id-item" + ) + })); + assert!(rollout_items.iter().any(|item| { + matches!( + item, + ResponseItem::ToolSearchOutput { + call_id: Some(call_id), + tools, + .. + } if call_id == CALL_ID && tools.is_empty() + ) + })); + + drop(test); + + let resume_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-2", "resumed"), + ev_completed("resp-3"), + ]), + ) + .await; + let mut resume_builder = test_codex(); + let resumed = resume_builder.resume(&server, home, rollout_path).await?; + + resumed.submit_turn("continue").await?; + + let resumed_request = resume_mock.single_request(); + assert!( + resumed_request + .inputs_of_type("tool_search_call") + .contains(&server_request) + ); + assert_eq!( + tool_search_call(&resumed_request, CALL_ID), + canonical_malformed_request + ); + assert_eq!( + tool_search_call(&resumed_request, VALID_CALL_ID), + canonical_valid_request + ); + assert_eq!( + resumed_request.tool_search_output(CALL_ID).get("tools"), + Some(&json!([])) + ); + + Ok(()) +}