Simplify parallel (#4829)

make tool processing return a future and then collect futures.
handle cleanup on Drop

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

@@ -1,6 +1,8 @@
 use std::sync::Arc;
 
-use tokio::task::JoinHandle;
+use tokio::sync::RwLock;
+use tokio_util::either::Either;
+use tokio_util::task::AbortOnDropHandle;
 
 use crate::codex::Session;
 use crate::codex::TurnContext;
@@ -11,20 +13,13 @@ use crate::tools::router::ToolCall;
 use crate::tools::router::ToolRouter;
 use codex_protocol::models::ResponseInputItem;
 
-use crate::codex::ProcessedResponseItem;
-
-struct PendingToolCall {
-    index: usize,
-    handle: JoinHandle<Result<ResponseInputItem, FunctionCallError>>,
-}
-
 pub(crate) struct ToolCallRuntime {
     router: Arc<ToolRouter>,
     session: Arc<Session>,
     turn_context: Arc<TurnContext>,
     tracker: SharedTurnDiffTracker,
     sub_id: String,
-    pending_calls: Vec<PendingToolCall>,
+    parallel_execution: Arc<RwLock<()>>,
 }
 
 impl ToolCallRuntime {
@@ -41,97 +36,45 @@ impl ToolCallRuntime {
             turn_context,
             tracker,
             sub_id,
-            pending_calls: Vec::new(),
+            parallel_execution: Arc::new(RwLock::new(())),
         }
     }
 
-    pub(crate) async fn handle_tool_call(
-        &mut self,
+    pub(crate) fn handle_tool_call(
+        &self,
         call: ToolCall,
-        output_index: usize,
-        output: &mut [ProcessedResponseItem],
-    ) -> Result<(), CodexErr> {
+    ) -> impl std::future::Future<Output = Result<ResponseInputItem, CodexErr>> {
         let supports_parallel = self.router.tool_supports_parallel(&call.tool_name);
-        if supports_parallel {
-            self.spawn_parallel(call, output_index);
-        } else {
-            self.resolve_pending(output).await?;
-            let response = self.dispatch_serial(call).await?;
-            let slot = output.get_mut(output_index).ok_or_else(|| {
-                CodexErr::Fatal(format!("tool output index {output_index} out of bounds"))
-            })?;
-            slot.response = Some(response);
-        }
 
-        Ok(())
-    }
-
-    pub(crate) fn abort_all(&mut self) {
-        while let Some(pending) = self.pending_calls.pop() {
-            pending.handle.abort();
-        }
-    }
-
-    pub(crate) async fn resolve_pending(
-        &mut self,
-        output: &mut [ProcessedResponseItem],
-    ) -> Result<(), CodexErr> {
-        while let Some(PendingToolCall { index, handle }) = self.pending_calls.pop() {
-            match handle.await {
-                Ok(Ok(response)) => {
-                    if let Some(slot) = output.get_mut(index) {
-                        slot.response = Some(response);
-                    }
-                }
-                Ok(Err(FunctionCallError::Fatal(message))) => {
-                    self.abort_all();
-                    return Err(CodexErr::Fatal(message));
-                }
-                Ok(Err(other)) => {
-                    self.abort_all();
-                    return Err(CodexErr::Fatal(other.to_string()));
-                }
-                Err(join_err) => {
-                    self.abort_all();
-                    return Err(CodexErr::Fatal(format!(
-                        "tool task failed to join: {join_err}"
-                    )));
-                }
-            }
-        }
-
-        Ok(())
-    }
-
-    fn spawn_parallel(&mut self, call: ToolCall, index: usize) {
         let router = Arc::clone(&self.router);
         let session = Arc::clone(&self.session);
         let turn = Arc::clone(&self.turn_context);
         let tracker = Arc::clone(&self.tracker);
         let sub_id = self.sub_id.clone();
-        let handle = tokio::spawn(async move {
-            router
-                .dispatch_tool_call(session, turn, tracker, sub_id, call)
-                .await
-        });
-        self.pending_calls.push(PendingToolCall { index, handle });
-    }
+        let lock = Arc::clone(&self.parallel_execution);
 
-    async fn dispatch_serial(&self, call: ToolCall) -> Result<ResponseInputItem, CodexErr> {
-        match self
-            .router
-            .dispatch_tool_call(
-                Arc::clone(&self.session),
-                Arc::clone(&self.turn_context),
-                Arc::clone(&self.tracker),
-                self.sub_id.clone(),
-                call,
-            )
-            .await
-        {
-            Ok(response) => Ok(response),
-            Err(FunctionCallError::Fatal(message)) => Err(CodexErr::Fatal(message)),
-            Err(other) => Err(CodexErr::Fatal(other.to_string())),
+        let handle: AbortOnDropHandle<Result<ResponseInputItem, FunctionCallError>> =
+            AbortOnDropHandle::new(tokio::spawn(async move {
+                let _guard = if supports_parallel {
+                    Either::Left(lock.read().await)
+                } else {
+                    Either::Right(lock.write().await)
+                };
+
+                router
+                    .dispatch_tool_call(session, turn, tracker, sub_id, call)
+                    .await
+            }));
+
+        async move {
+            match handle.await {
+                Ok(Ok(response)) => Ok(response),
+                Ok(Err(FunctionCallError::Fatal(message))) => Err(CodexErr::Fatal(message)),
+                Ok(Err(other)) => Err(CodexErr::Fatal(other.to_string())),
+                Err(err) => Err(CodexErr::Fatal(format!(
+                    "tool task failed to receive: {err:?}"
+                ))),
+            }
         }
     }
 }