From cacf8c91bb941b4160c197de128b6a08fd8c7c24 Mon Sep 17 00:00:00 2001 From: Jeremy Rose Date: Fri, 22 Aug 2025 15:23:25 -0700 Subject: [PATCH] wip --- codex-rs/core/src/client.rs | 2 +- codex-rs/core/src/codex.rs | 246 ++++++++++++++++++++++++++- codex-rs/tui/src/chatwidget.rs | 14 +- codex-rs/tui/src/chatwidget/tests.rs | 1 + codex-rs/tui/src/history_cell.rs | 55 +++++- 5 files changed, 304 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 6206b4264a..f12b1ac9c0 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -185,7 +185,7 @@ impl ModelClient { input: &input_with_instructions, tools: &tools_json, tool_choice: "auto", - parallel_tool_calls: false, + parallel_tool_calls: true, reasoning, store, stream: true, diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 79b73a3335..783f46af72 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1643,6 +1643,31 @@ async fn try_run_turn( let mut stream = turn_context.client.clone().stream(&prompt).await?; let mut output = Vec::new(); + // Stage tool calls to execute after the model signals completion. + // We do this so multiple tool calls can run in parallel (where possible) + // and so we don't block SSE processing. + enum StagedCall { + Exec { + item: ResponseItem, + name: String, + args: String, + call_id: String, + }, + LocalShell { + item: ResponseItem, + id: Option, + call_id: Option, + action: crate::models::LocalShellAction, + }, + Mcp { + item: ResponseItem, + server: String, + tool: String, + args: String, + call_id: String, + }, + } + let mut staged: Vec = Vec::new(); loop { // Poll the next item from the model stream. We must inspect *both* Ok and Err // cases so that transient stream failures (e.g., dropped SSE connection before @@ -1669,15 +1694,59 @@ async fn try_run_turn( match event { ResponseEvent::Created => {} ResponseEvent::OutputItemDone(item) => { - let response = handle_response_item( - sess, - turn_context, - turn_diff_tracker, - sub_id, - item.clone(), - ) - .await?; - output.push(ProcessedResponseItem { item, response }); + match &item { + ResponseItem::FunctionCall { + name, + arguments, + call_id, + .. + } => { + // Classify as MCP vs Exec and stage execution for after completion. + if let Some((server, tool)) = + sess.mcp_connection_manager.parse_tool_name(name) + { + staged.push(StagedCall::Mcp { + item: item.clone(), + server, + tool, + args: arguments.clone(), + call_id: call_id.clone(), + }); + } else { + staged.push(StagedCall::Exec { + item: item.clone(), + name: name.clone(), + args: arguments.clone(), + call_id: call_id.clone(), + }); + } + } + ResponseItem::LocalShellCall { + id, + call_id, + action, + .. + } => { + staged.push(StagedCall::LocalShell { + item: item.clone(), + id: id.clone(), + call_id: call_id.clone(), + action: action.clone(), + }); + } + _ => { + // Non-tool items are handled immediately (messages, reasoning, deltas). + let response = handle_response_item( + sess, + turn_context, + turn_diff_tracker, + sub_id, + item.clone(), + ) + .await?; + output.push(ProcessedResponseItem { item, response }); + } + } } ResponseEvent::Completed { response_id: _, @@ -1703,6 +1772,165 @@ async fn try_run_turn( let _ = sess.tx_event.send(event).await; } + // Execute any staged tool calls before returning the turn output. + use futures::future::join_all; + + // Run MCP, LocalShell, and non-apply exec tool calls fully in parallel. + let mut mcp_futs = Vec::new(); + let mut local_futs = Vec::new(); + let mut exec_parallel_futs = Vec::new(); + // Keep apply_patch exec calls sequential for diff tracking. + let mut exec_sequential_calls = Vec::new(); + + for sc in staged { + match sc { + StagedCall::Mcp { + item, + server, + tool, + args, + call_id, + } => { + let fut = handle_mcp_tool_call( + sess, + sub_id, + call_id.clone(), + server, + tool, + args, + None, + ); + mcp_futs.push(async move { + let response = fut.await; + ProcessedResponseItem { + item, + response: Some(response), + } + }); + } + StagedCall::LocalShell { + item, + id, + call_id, + action, + } => { + let effective_call_id = match (call_id, id) { + (Some(call_id), _) => call_id, + (None, Some(id)) => id, + (None, None) => String::new(), + }; + let crate::models::LocalShellAction::Exec(action) = action; + let params = ShellToolCallParams { + command: action.command, + workdir: action.working_directory, + timeout_ms: action.timeout_ms, + with_escalated_permissions: None, + justification: None, + }; + let exec_params = to_exec_params(params, turn_context); + let fut = async move { + // Use a fresh tracker – non-apply_patch calls won't emit diffs anyway. + let mut tdt = TurnDiffTracker::new(); + let response = handle_container_exec_with_params( + exec_params, + sess, + turn_context, + &mut tdt, + sub_id.to_string(), + effective_call_id, + ) + .await; + ProcessedResponseItem { + item, + response: Some(response), + } + }; + local_futs.push(fut); + } + StagedCall::Exec { + item, + name, + args, + call_id, + } => { + if name == "apply_patch" { + exec_sequential_calls.push(StagedCall::Exec { + item, + name, + args, + call_id, + }); + } else { + let fut = async move { + let response = match parse_container_exec_arguments( + args, + turn_context, + &call_id, + ) { + Ok(params) => { + let mut tdt = TurnDiffTracker::new(); + handle_container_exec_with_params( + params, + sess, + turn_context, + &mut tdt, + sub_id.to_string(), + call_id, + ) + .await + } + Err(output) => *output, + }; + ProcessedResponseItem { + item, + response: Some(response), + } + }; + exec_parallel_futs.push(fut); + } + } + other => exec_sequential_calls.push(other), + } + } + + let mut processed: Vec = join_all(mcp_futs).await; + processed.extend(join_all(local_futs).await); + processed.extend(join_all(exec_parallel_futs).await); + + // Now handle exec/shell calls one by one so we can use the turn_diff_tracker safely. + for sc in exec_sequential_calls { + match sc { + StagedCall::Exec { + item, + name, + args, + call_id, + } => { + // Reuse existing logic for exec tools. + let response = handle_function_call( + sess, + turn_context, + turn_diff_tracker, + sub_id.to_string(), + name, + args, + call_id, + ) + .await; + processed.push(ProcessedResponseItem { + item, + response: Some(response), + }); + } + StagedCall::LocalShell { .. } => {} + StagedCall::Mcp { .. } => {} + } + } + + // Preserve original ordering: non-tool items were already pushed to `output` as they arrived; + // now append the processed tool-call items in their staged order. + output.extend(processed); + return Ok(output); } ResponseEvent::OutputTextDelta(delta) => { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 0e4bd85664..a1d8a33746 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -54,6 +54,7 @@ use crate::history_cell::CommandOutput; use crate::history_cell::ExecCell; use crate::history_cell::HistoryCell; use crate::history_cell::PatchEventType; +use crate::history_cell::RunningMcpCell; use crate::slash_command::SlashCommand; use crate::tui::FrameRequester; // streaming internals are provided by crate::streaming and crate::markdown_stream @@ -73,6 +74,8 @@ use codex_core::protocol::AskForApproval; use codex_core::protocol::SandboxPolicy; use codex_core::protocol_config_types::ReasoningEffort as ReasoningEffortConfig; use codex_file_search::FileMatch; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use uuid::Uuid; // Track information about an in-flight exec command. @@ -93,6 +96,7 @@ pub(crate) struct ChatWidget { // Stream lifecycle controller stream: StreamController, running_commands: HashMap, + running_mcp: HashMap>, pending_exec_completions: Vec<(Vec, Vec, CommandOutput)>, task_complete_pending: bool, // Queue of interruptive UI events deferred during an active write cycle @@ -486,10 +490,17 @@ impl ChatWidget { pub(crate) fn handle_mcp_begin_now(&mut self, ev: McpToolCallBeginEvent) { self.flush_answer_stream_with_separator(); - self.add_to_history(history_cell::new_active_mcp_tool_call(ev.invocation)); + // Track running state so the spinner stops when the call ends. + let is_running = Arc::new(AtomicBool::new(true)); + self.running_mcp + .insert(ev.call_id.clone(), Arc::clone(&is_running)); + self.add_boxed_history(Box::new(RunningMcpCell::new(ev.invocation, is_running))); } pub(crate) fn handle_mcp_end_now(&mut self, ev: McpToolCallEndEvent) { self.flush_answer_stream_with_separator(); + if let Some(flag) = self.running_mcp.remove(&ev.call_id) { + flag.store(false, Ordering::Relaxed); + } self.add_boxed_history(history_cell::new_completed_mcp_tool_call( 80, ev.invocation, @@ -558,6 +569,7 @@ impl ChatWidget { last_token_usage: TokenUsage::default(), stream: StreamController::new(config), running_commands: HashMap::new(), + running_mcp: HashMap::new(), pending_exec_completions: Vec::new(), task_complete_pending: false, interrupts: InterruptManager::new(), diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index e43dfc7824..addc8672d0 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -149,6 +149,7 @@ fn make_chatwidget_manual() -> ( last_token_usage: TokenUsage::default(), stream: StreamController::new(cfg), running_commands: HashMap::new(), + running_mcp: HashMap::new(), pending_exec_completions: Vec::new(), task_complete_pending: false, interrupts: InterruptManager::new(), diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs index 0cbb4888df..c6008882a2 100644 --- a/codex-rs/tui/src/history_cell.rs +++ b/codex-rs/tui/src/history_cell.rs @@ -36,6 +36,9 @@ use shlex::try_join as shlex_try_join; use std::collections::HashMap; use std::io::Cursor; use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use tracing::error; @@ -375,16 +378,62 @@ fn new_exec_command_generic( } pub(crate) fn new_active_mcp_tool_call(invocation: McpInvocation) -> PlainHistoryCell { + // Backwards-compatible helper retained for call sites; create a non-animated cell. + // Prefer `new_running_mcp_tool_call` for spinner support. let title_line = Line::from(vec!["tool".magenta(), " running...".dim()]); let lines: Vec = vec![ - title_line, - format_mcp_invocation(invocation.clone()), Line::from(""), + title_line, + format_mcp_invocation(invocation), ]; - PlainHistoryCell { lines } } +#[derive(Debug)] +pub(crate) struct RunningMcpCell { + invocation: McpInvocation, + start_time: Instant, + is_running: Arc, +} + +impl RunningMcpCell { + pub(crate) fn new(invocation: McpInvocation, is_running: Arc) -> Self { + Self { + invocation, + start_time: Instant::now(), + is_running, + } + } +} + +impl HistoryCell for RunningMcpCell { + fn display_lines(&self) -> Vec> { + let mut lines: Vec> = Vec::new(); + lines.push(Line::from("")); + + // Spinner or final status marker + let running = self.is_running.load(Ordering::Relaxed); + let status = if running { + const FRAMES: &[char] = &['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']; + let idx = ((self.start_time.elapsed().as_millis() / 100) as usize) % FRAMES.len(); + Span::raw(format!("{}", FRAMES[idx])) + } else { + "✓".green() + }; + + let title = Line::from(vec![ + "tool".magenta(), + " ".into(), + status, + " ".into(), + "running...".dim(), + ]); + lines.push(title); + lines.push(format_mcp_invocation(self.invocation.clone())); + lines + } +} + /// If the first content is an image, return a new cell with the image. /// TODO(rgwood-dd): Handle images properly even if they're not the first result. fn try_new_completed_mcp_tool_call_with_image_output(