diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 524be63af5..3b43560d37 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2179,12 +2179,12 @@ async fn try_run_turn( .or_cancel(&cancellation_token) .await??; - let tool_runtime = ToolCallRuntime::new( + let tool_runtime = Arc::new(ToolCallRuntime::new( Arc::clone(&router), Arc::clone(&sess), Arc::clone(&turn_context), Arc::clone(&turn_diff_tracker), - ); + )); let mut in_flight: FuturesUnordered>> = FuturesUnordered::new(); let mut responses: Vec = Vec::new(); @@ -2225,7 +2225,7 @@ async fn try_run_turn( handle_output_item_done( &sess, &turn_context, - &tool_runtime, + Arc::clone(&tool_runtime), item, previously_active_item, &mut in_flight, diff --git a/codex-rs/core/src/codex/turn_event.rs b/codex-rs/core/src/codex/turn_event.rs index ce4b996f86..8b4524b2b5 100644 --- a/codex-rs/core/src/codex/turn_event.rs +++ b/codex-rs/core/src/codex/turn_event.rs @@ -1,30 +1,37 @@ +use std::pin::Pin; use std::sync::Arc; -use futures::future::BoxFuture; use futures::stream::FuturesUnordered; use tokio_util::sync::CancellationToken; +use super::CodexErr; +use super::CodexResult; +use super::Session; use super::ToolRouter; -use super::{CodexErr, CodexResult, Session, TurnContext, TurnItem}; +use super::TurnContext; +use super::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 futures::Future; 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. +pub(super) type InFlightFuture<'f> = + Pin> + Send + 'f>>; + #[allow(clippy::too_many_arguments)] pub(super) async fn handle_output_item_done( sess: &Arc, turn_context: &Arc, - tool_runtime: &super::ToolCallRuntime, + tool_runtime: Arc, item: ResponseItem, previously_active_item: Option, - in_flight: &mut FuturesUnordered>>, + in_flight: &mut FuturesUnordered>, responses: &mut Vec, last_agent_message: &mut Option, cancellation_token: CancellationToken, @@ -39,23 +46,22 @@ pub(super) async fn handle_output_item_done( let sess_for_output: Arc = Arc::clone(sess); let turn_for_output: Arc = Arc::clone(turn_context); - let runtime = tool_runtime.clone(); + let tool_runtime = Arc::clone(&tool_runtime); - 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) + in_flight.push(Box::pin(async move { + let response_input = tool_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; } - .boxed(), - ); + Ok(response_input) + })); } Ok(None) => { if let Some(turn_item) = handle_non_tool_response_item(&item).await { diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 971ea934d8..33dc42b936 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -17,7 +17,6 @@ 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,