From f2b9e7a76c2014404bd007a33571b8ddd6f1e226 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 25 Nov 2025 21:19:17 -0800 Subject: [PATCH] Extract turn event helpers --- codex-rs/core/src/codex.rs | 186 +++----------------------- codex-rs/core/src/codex/turn_event.rs | 180 +++++++++++++++++++++++++ codex-rs/core/src/tools/parallel.rs | 1 + 3 files changed, 198 insertions(+), 169 deletions(-) create mode 100644 codex-rs/core/src/codex/turn_event.rs diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 2a9145d9e9..524be63af5 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -12,9 +12,11 @@ use crate::compact::run_inline_auto_compact_task; use crate::compact::should_use_remote_compact_task; use crate::compact_remote::run_inline_remote_auto_compact_task; use crate::features::Feature; +#[cfg(test)] use crate::function_tool::FunctionCallError; use crate::parse_command::parse_command; use crate::parse_turn_item; +mod turn_event; use crate::terminal; use crate::truncate::TruncationPolicy; use crate::user_notification::UserNotifier; @@ -127,6 +129,7 @@ use codex_otel::otel_event_manager::OtelEventManager; use codex_protocol::config_types::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; use codex_protocol::models::ContentItem; +#[cfg(test)] use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; @@ -135,6 +138,8 @@ use codex_protocol::protocol::InitialHistory; use codex_protocol::user_input::UserInput; use codex_utils_readiness::Readiness; use codex_utils_readiness::ReadinessFlag; +use turn_event::handle_non_tool_response_item; +use turn_event::handle_output_item_done; /// The high-level interface to the Codex system. /// It operates as a queue pair where you send submissions and receive events. @@ -2217,117 +2222,18 @@ async fn try_run_turn( ResponseEvent::Created => {} ResponseEvent::OutputItemDone(item) => { let previously_active_item = active_item.take(); - match ToolRouter::build_tool_call(sess.as_ref(), item.clone()).await { - Ok(Some(call)) => { - let payload_preview = call.payload.log_payload().into_owned(); - tracing::info!("ToolCall: {} {}", call.tool_name, payload_preview); - - sess.record_conversation_items( - &turn_context, - std::slice::from_ref(&item), - ) - .await; - - let sess_for_output = Arc::clone(&sess); - let turn_for_output = Arc::clone(&turn_context); - let response = - tool_runtime.handle_tool_call(call, cancellation_token.child_token()); - - in_flight.push(async move { - let response_input = response.await?; - if let Some(response_item) = - response_input_to_response_item(&response_input) - { - sess_for_output - .record_conversation_items( - turn_for_output.as_ref(), - std::slice::from_ref(&response_item), - ) - .await; - } - Ok(response_input) - } - .boxed()); - } - Ok(None) => { - if let Some(turn_item) = handle_non_tool_response_item(&item).await { - if previously_active_item.is_none() { - sess.emit_turn_item_started(&turn_context, &turn_item).await; - } - - sess.emit_turn_item_completed(&turn_context, turn_item) - .await; - } - - sess.record_conversation_items( - &turn_context, - std::slice::from_ref(&item), - ) - .await; - if let Some(agent_message) = last_assistant_message_from_item(&item) { - last_agent_message = Some(agent_message); - } - } - Err(FunctionCallError::MissingLocalShellCallId) => { - let msg = "LocalShellCall without call_id or id"; - turn_context - .client - .get_otel_event_manager() - .log_tool_failed("local_shell", msg); - error!(msg); - - let response = ResponseInputItem::FunctionCallOutput { - call_id: String::new(), - output: FunctionCallOutputPayload { - content: msg.to_string(), - ..Default::default() - }, - }; - sess.record_conversation_items( - &turn_context, - std::slice::from_ref(&item), - ) - .await; - if let Some(response_item) = - response_input_to_response_item(&response) - { - sess.record_conversation_items( - &turn_context, - std::slice::from_ref(&response_item), - ) - .await; - } - responses.push(response); - } - Err(FunctionCallError::RespondToModel(message)) - | Err(FunctionCallError::Denied(message)) => { - let response = ResponseInputItem::FunctionCallOutput { - call_id: String::new(), - output: FunctionCallOutputPayload { - content: message, - ..Default::default() - }, - }; - sess.record_conversation_items( - &turn_context, - std::slice::from_ref(&item), - ) - .await; - if let Some(response_item) = - response_input_to_response_item(&response) - { - sess.record_conversation_items( - &turn_context, - std::slice::from_ref(&response_item), - ) - .await; - } - responses.push(response); - } - Err(FunctionCallError::Fatal(message)) => { - return Err(CodexErr::Fatal(message)); - } - } + handle_output_item_done( + &sess, + &turn_context, + &tool_runtime, + item, + previously_active_item, + &mut in_flight, + &mut responses, + &mut last_agent_message, + cancellation_token.child_token(), + ) + .await?; } ResponseEvent::OutputItemAdded(item) => { if let Some(turn_item) = handle_non_tool_response_item(&item).await { @@ -2435,64 +2341,6 @@ async fn try_run_turn( } } -fn last_assistant_message_from_item(item: &ResponseItem) -> Option { - if let ResponseItem::Message { role, content, .. } = item - && role == "assistant" { - return content.iter().rev().find_map(|ci| match ci { - ContentItem::OutputText { text } => Some(text.clone()), - _ => None, - }); - } - None -} - -fn response_input_to_response_item(input: &ResponseInputItem) -> Option { - match input { - ResponseInputItem::FunctionCallOutput { call_id, output } => { - Some(ResponseItem::FunctionCallOutput { - call_id: call_id.clone(), - output: output.clone(), - }) - } - ResponseInputItem::CustomToolCallOutput { call_id, output } => { - Some(ResponseItem::CustomToolCallOutput { - call_id: call_id.clone(), - output: output.clone(), - }) - } - ResponseInputItem::McpToolCallOutput { call_id, result } => { - let output = match result { - Ok(call_tool_result) => FunctionCallOutputPayload::from(call_tool_result), - Err(err) => FunctionCallOutputPayload { - content: err.clone(), - success: Some(false), - ..Default::default() - }, - }; - Some(ResponseItem::FunctionCallOutput { - call_id: call_id.clone(), - output, - }) - } - _ => None, - } -} - -async fn handle_non_tool_response_item(item: &ResponseItem) -> Option { - debug!(?item, "Output item"); - - match item { - ResponseItem::Message { .. } - | ResponseItem::Reasoning { .. } - | ResponseItem::WebSearchCall { .. } => parse_turn_item(item), - ResponseItem::FunctionCallOutput { .. } | ResponseItem::CustomToolCallOutput { .. } => { - debug!("unexpected tool output from stream"); - None - } - _ => None, - } -} - pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option { responses.iter().rev().find_map(|item| { if let ResponseItem::Message { role, content, .. } = item { diff --git a/codex-rs/core/src/codex/turn_event.rs b/codex-rs/core/src/codex/turn_event.rs new file mode 100644 index 0000000000..ce4b996f86 --- /dev/null +++ b/codex-rs/core/src/codex/turn_event.rs @@ -0,0 +1,180 @@ +use std::sync::Arc; + +use futures::future::BoxFuture; +use futures::stream::FuturesUnordered; +use tokio_util::sync::CancellationToken; + +use super::ToolRouter; +use super::{CodexErr, CodexResult, Session, TurnContext, TurnItem}; +use crate::function_tool::FunctionCallError; +use crate::parse_turn_item; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ResponseItem; +use futures::FutureExt; +use tracing::debug; + +/// Handle a completed output item from the model stream, recording it and +/// queuing any tool execution futures. This records items immediately so +/// history and rollout stay in sync even if the turn is later cancelled. +#[allow(clippy::too_many_arguments)] +pub(super) async fn handle_output_item_done( + sess: &Arc, + turn_context: &Arc, + tool_runtime: &super::ToolCallRuntime, + item: ResponseItem, + previously_active_item: Option, + in_flight: &mut FuturesUnordered>>, + responses: &mut Vec, + last_agent_message: &mut Option, + cancellation_token: CancellationToken, +) -> CodexResult<()> { + match ToolRouter::build_tool_call(sess.as_ref(), item.clone()).await { + Ok(Some(call)) => { + let payload_preview = call.payload.log_payload().into_owned(); + tracing::info!("ToolCall: {} {}", call.tool_name, payload_preview); + + sess.record_conversation_items(turn_context, std::slice::from_ref(&item)) + .await; + + let sess_for_output: Arc = Arc::clone(sess); + let turn_for_output: Arc = Arc::clone(turn_context); + let runtime = tool_runtime.clone(); + + in_flight.push( + async move { + let response_input = runtime.handle_tool_call(call, cancellation_token).await?; + if let Some(response_item) = response_input_to_response_item(&response_input) { + sess_for_output + .record_conversation_items( + turn_for_output.as_ref(), + std::slice::from_ref(&response_item), + ) + .await; + } + Ok(response_input) + } + .boxed(), + ); + } + Ok(None) => { + if let Some(turn_item) = handle_non_tool_response_item(&item).await { + if previously_active_item.is_none() { + sess.emit_turn_item_started(turn_context, &turn_item).await; + } + + sess.emit_turn_item_completed(turn_context, turn_item).await; + } + + sess.record_conversation_items(turn_context, std::slice::from_ref(&item)) + .await; + if let Some(agent_message) = last_assistant_message_from_item(&item) { + *last_agent_message = Some(agent_message); + } + } + Err(FunctionCallError::MissingLocalShellCallId) => { + let msg = "LocalShellCall without call_id or id"; + turn_context + .client + .get_otel_event_manager() + .log_tool_failed("local_shell", msg); + tracing::error!(msg); + + let response = ResponseInputItem::FunctionCallOutput { + call_id: String::new(), + output: FunctionCallOutputPayload { + content: msg.to_string(), + ..Default::default() + }, + }; + sess.record_conversation_items(turn_context, std::slice::from_ref(&item)) + .await; + if let Some(response_item) = response_input_to_response_item(&response) { + sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) + .await; + } + responses.push(response); + } + Err(FunctionCallError::RespondToModel(message)) + | Err(FunctionCallError::Denied(message)) => { + let response = ResponseInputItem::FunctionCallOutput { + call_id: String::new(), + output: FunctionCallOutputPayload { + content: message, + ..Default::default() + }, + }; + sess.record_conversation_items(turn_context, std::slice::from_ref(&item)) + .await; + if let Some(response_item) = response_input_to_response_item(&response) { + sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) + .await; + } + responses.push(response); + } + Err(FunctionCallError::Fatal(message)) => { + return Err(CodexErr::Fatal(message)); + } + } + + Ok(()) +} + +pub(super) async fn handle_non_tool_response_item(item: &ResponseItem) -> Option { + debug!(?item, "Output item"); + + match item { + ResponseItem::Message { .. } + | ResponseItem::Reasoning { .. } + | ResponseItem::WebSearchCall { .. } => parse_turn_item(item), + ResponseItem::FunctionCallOutput { .. } | ResponseItem::CustomToolCallOutput { .. } => { + debug!("unexpected tool output from stream"); + None + } + _ => None, + } +} + +pub(super) fn last_assistant_message_from_item(item: &ResponseItem) -> Option { + if let ResponseItem::Message { role, content, .. } = item + && role == "assistant" + { + return content.iter().rev().find_map(|ci| match ci { + codex_protocol::models::ContentItem::OutputText { text } => Some(text.clone()), + _ => None, + }); + } + None +} + +pub(super) fn response_input_to_response_item(input: &ResponseInputItem) -> Option { + match input { + ResponseInputItem::FunctionCallOutput { call_id, output } => { + Some(ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output: output.clone(), + }) + } + ResponseInputItem::CustomToolCallOutput { call_id, output } => { + Some(ResponseItem::CustomToolCallOutput { + call_id: call_id.clone(), + output: output.clone(), + }) + } + ResponseInputItem::McpToolCallOutput { call_id, result } => { + let output = match result { + Ok(call_tool_result) => FunctionCallOutputPayload::from(call_tool_result), + Err(err) => FunctionCallOutputPayload { + content: err.clone(), + success: Some(false), + ..Default::default() + }, + }; + Some(ResponseItem::FunctionCallOutput { + call_id: call_id.clone(), + output, + }) + } + _ => None, + } +} diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 33dc42b936..971ea934d8 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -17,6 +17,7 @@ use crate::tools::router::ToolRouter; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; +#[derive(Clone)] pub(crate) struct ToolCallRuntime { router: Arc, session: Arc,