fix parallel tool calls

codex-rs/core/src/codex.rs

@@ -2476,6 +2476,7 @@ async fn try_run_turn(
     let mut needs_follow_up = false;
     let mut last_agent_message: Option<String> = None;
     let mut active_item: Option<TurnItem> = None;
+    let mut should_emit_turn_diff = false;
     let receiving_span = info_span!("receiving_stream");
     let outcome: CodexResult<TurnRunResult> = loop {
         let handle_responses = info_span!(
@@ -2551,14 +2552,7 @@ async fn try_run_turn(
             } => {
                 sess.update_token_usage_info(&turn_context, token_usage.as_ref())
                     .await;
-                let unified_diff = {
-                    let mut tracker = turn_diff_tracker.lock().await;
-                    tracker.get_unified_diff()
-                };
-                if let Ok(Some(unified_diff)) = unified_diff {
-                    let msg = EventMsg::TurnDiff(TurnDiffEvent { unified_diff });
-                    sess.send_event(&turn_context, msg).await;
-                }
+                should_emit_turn_diff = true;
 
                 break Ok(TurnRunResult {
                     needs_follow_up,
@@ -2632,7 +2626,18 @@ async fn try_run_turn(
         }
     };
 
-    drain_in_flight(&mut in_flight, sess, turn_context).await?;
+    drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?;
+
+    if should_emit_turn_diff {
+        let unified_diff = {
+            let mut tracker = turn_diff_tracker.lock().await;
+            tracker.get_unified_diff()
+        };
+        if let Ok(Some(unified_diff)) = unified_diff {
+            let msg = EventMsg::TurnDiff(TurnDiffEvent { unified_diff });
+            sess.clone().send_event(&turn_context, msg).await;
+        }
+    }
 
     outcome
 }

codex-rs/core/src/stream_events_utils.rs

@@ -16,7 +16,6 @@ use codex_protocol::models::FunctionCallOutputPayload;
 use codex_protocol::models::ResponseInputItem;
 use codex_protocol::models::ResponseItem;
 use futures::Future;
-use tracing::Instrument;
 use tracing::debug;
 use tracing::instrument;
 
@@ -59,16 +58,10 @@ pub(crate) async fn handle_output_item_done(
                 .await;
 
             let cancellation_token = ctx.cancellation_token.child_token();
-            let tool_runtime = ctx.tool_runtime.clone();
-
             let tool_future: InFlightFuture<'static> = Box::pin(
-                async move {
-                    let response_input = tool_runtime
-                        .handle_tool_call(call, cancellation_token)
-                        .await?;
-                    Ok(response_input)
-                }
-                .in_current_span(),
+                ctx.tool_runtime
+                    .clone()
+                    .handle_tool_call(call, cancellation_token),
             );
 
             output.needs_follow_up = true;

codex-rs/core/src/tools/parallel.rs

@@ -47,7 +47,7 @@ impl ToolCallRuntime {
 
     #[instrument(skip_all, fields(call = ?call))]
     pub(crate) fn handle_tool_call(
-        &self,
+        self,
         call: ToolCall,
         cancellation_token: CancellationToken,
     ) -> impl std::future::Future<Output = Result<ResponseInputItem, CodexErr>> {

codex-rs/core/tests/common/Cargo.toml

@@ -22,3 +22,7 @@ tokio = { workspace = true, features = ["time"] }
 walkdir = { workspace = true }
 wiremock = { workspace = true }
 shlex = { workspace = true }
+
+[dev-dependencies]
+pretty_assertions = { workspace = true }
+reqwest = { workspace = true }

codex-rs/core/tests/common/lib.rs

@@ -14,6 +14,7 @@ use std::path::PathBuf;
 use assert_cmd::cargo::cargo_bin;
 
 pub mod responses;
+pub mod streaming_sse;
 pub mod test_codex;
 pub mod test_codex_exec;