openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

wip

Browse Source
This commit is contained in:
Jeremy Rose
2025-08-22 16:23:16 -07:00
parent 097782c775
commit cacf8c91bb
5 changed files with 304 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
+1 -1
View File
@@ -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,
codex-rs/core/src/codex.rs
+237 -9
View File
@@ -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<String>,
call_id: Option<String>,
action: crate::models::LocalShellAction,
},
Mcp {
item: ResponseItem,
server: String,
tool: String,
args: String,
call_id: String,
},
}
let mut staged: Vec<StagedCall> = 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<ProcessedResponseItem> = 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) => {
codex-rs/tui/src/chatwidget.rs
+13 -1
View File
@@ -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<String, RunningCommand>,
running_mcp: HashMap<String, Arc<AtomicBool>>,
pending_exec_completions: Vec<(Vec<String>, Vec<ParsedCommand>, 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(),
codex-rs/tui/src/chatwidget/tests.rs
+1
View File
@@ -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(),
codex-rs/tui/src/history_cell.rs
+52 -3
View File
@@ -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<Line> = 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<AtomicBool>,
}
impl RunningMcpCell {
pub(crate) fn new(invocation: McpInvocation, is_running: Arc<AtomicBool>) -> Self {
Self {
invocation,
start_time: Instant::now(),
is_running,
}
}
}
impl HistoryCell for RunningMcpCell {
fn display_lines(&self) -> Vec<Line<'static>> {
let mut lines: Vec<Line<'static>> = 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(