agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: better UI for unified_exec (#6515)

Browse Source 
<img width="376" height="132" alt="Screenshot 2025-11-12 at 17 36 22"
src="https://github.com/user-attachments/assets/ce693f0d-5ca0-462e-b170-c20811dcc8d5"
/>
This commit is contained in:
jif-oai
2025-11-14 16:31:12 +01:00
committed by GitHub
parent 4788fb179a
commit 63c8c01f40
17 changed files with 563 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tasks/user_shell.rs
+3 -1
View File
@@ -20,6 +20,7 @@ use crate::parse_command::parse_command;
use crate::protocol::EventMsg;
use crate::protocol::ExecCommandBeginEvent;
use crate::protocol::ExecCommandEndEvent;
use crate::protocol::ExecCommandSource;
use crate::protocol::SandboxPolicy;
use crate::protocol::TaskStartedEvent;
use crate::sandboxing::ExecEnv;
@@ -80,7 +81,8 @@ impl SessionTask for UserShellCommandTask {
command: shell_invocation.clone(),
cwd: turn_context.cwd.clone(),
parsed_cmd,
is_user_shell_command: true,
source: ExecCommandSource::UserShell,
interaction_input: None,
}),
)
.await;
codex-rs/core/src/tools/context.rs
-26
View File
@@ -4,17 +4,13 @@ use crate::tools::TELEMETRY_PREVIEW_MAX_BYTES;
use crate::tools::TELEMETRY_PREVIEW_MAX_LINES;
use crate::tools::TELEMETRY_PREVIEW_TRUNCATION_NOTICE;
use crate::turn_diff_tracker::TurnDiffTracker;
use codex_otel::otel_event_manager::OtelEventManager;
use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::FunctionCallOutputPayload;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ShellToolCallParams;
use codex_protocol::protocol::FileChange;
use codex_utils_string::take_bytes_at_char_boundary;
use mcp_types::CallToolResult;
use std::borrow::Cow;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Mutex;
@@ -244,25 +240,3 @@ mod tests {
assert_eq!(lines.last(), Some(&TELEMETRY_PREVIEW_TRUNCATION_NOTICE));
}
}
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub(crate) struct ExecCommandContext {
pub(crate) turn: Arc<TurnContext>,
pub(crate) call_id: String,
pub(crate) command_for_display: Vec<String>,
pub(crate) cwd: PathBuf,
pub(crate) apply_patch: Option<ApplyPatchCommandContext>,
pub(crate) tool_name: String,
pub(crate) otel_event_manager: OtelEventManager,
// TODO(abhisek-oai): Find a better way to track this.
// https://github.com/openai/codex/pull/2471/files#r2470352242
pub(crate) is_user_shell_command: bool,
}
#[derive(Clone, Debug)]
#[allow(dead_code)]
pub(crate) struct ApplyPatchCommandContext {
pub(crate) user_explicitly_approved_this_action: bool,
pub(crate) changes: HashMap<PathBuf, FileChange>,
}
codex-rs/core/src/tools/events.rs
+37 -14
View File
@@ -8,6 +8,7 @@ use crate::parse_command::parse_command;
use crate::protocol::EventMsg;
use crate::protocol::ExecCommandBeginEvent;
use crate::protocol::ExecCommandEndEvent;
use crate::protocol::ExecCommandSource;
use crate::protocol::FileChange;
use crate::protocol::PatchApplyBeginEvent;
use crate::protocol::PatchApplyEndEvent;
@@ -60,7 +61,8 @@ pub(crate) async fn emit_exec_command_begin(
ctx: ToolEventCtx<'_>,
command: &[String],
cwd: &Path,
is_user_shell_command: bool,
source: ExecCommandSource,
interaction_input: Option<String>,
) {
ctx.session
.send_event(
@@ -70,7 +72,8 @@ pub(crate) async fn emit_exec_command_begin(
command: command.to_vec(),
cwd: cwd.to_path_buf(),
parsed_cmd: parse_command(command),
is_user_shell_command,
source,
interaction_input,
}),
)
.await;
@@ -80,7 +83,7 @@ pub(crate) enum ToolEmitter {
Shell {
command: Vec<String>,
cwd: PathBuf,
is_user_shell_command: bool,
source: ExecCommandSource,
},
ApplyPatch {
changes: HashMap<PathBuf, FileChange>,
@@ -89,18 +92,17 @@ pub(crate) enum ToolEmitter {
UnifiedExec {
command: Vec<String>,
cwd: PathBuf,
// True for `exec_command` and false for `write_stdin`.
#[allow(dead_code)]
is_startup_command: bool,
source: ExecCommandSource,
interaction_input: Option<String>,
},
}
impl ToolEmitter {
pub fn shell(command: Vec<String>, cwd: PathBuf, is_user_shell_command: bool) -> Self {
pub fn shell(command: Vec<String>, cwd: PathBuf, source: ExecCommandSource) -> Self {
Self::Shell {
command,
cwd,
is_user_shell_command,
source,
}
}
@@ -111,11 +113,17 @@ impl ToolEmitter {
}
}
pub fn unified_exec(command: &[String], cwd: PathBuf, is_startup_command: bool) -> Self {
pub fn unified_exec(
command: &[String],
cwd: PathBuf,
source: ExecCommandSource,
interaction_input: Option<String>,
) -> Self {
Self::UnifiedExec {
command: command.to_vec(),
cwd,
is_startup_command,
source,
interaction_input,
}
}
@@ -125,11 +133,11 @@ impl ToolEmitter {
Self::Shell {
command,
cwd,
is_user_shell_command,
source,
},
ToolEventStage::Begin,
) => {
emit_exec_command_begin(ctx, command, cwd.as_path(), *is_user_shell_command).await;
emit_exec_command_begin(ctx, command, cwd.as_path(), *source, None).await;
}
(Self::Shell { .. }, ToolEventStage::Success(output)) => {
emit_exec_end(
@@ -217,8 +225,23 @@ impl ToolEmitter {
) => {
emit_patch_end(ctx, String::new(), (*message).to_string(), false).await;
}
(Self::UnifiedExec { command, cwd, .. }, ToolEventStage::Begin) => {
emit_exec_command_begin(ctx, command, cwd.as_path(), false).await;
(
Self::UnifiedExec {
command,
cwd,
source,
interaction_input,
},
ToolEventStage::Begin,
) => {
emit_exec_command_begin(
ctx,
command,
cwd.as_path(),
*source,
interaction_input.clone(),
)
.await;
}
(Self::UnifiedExec { .. }, ToolEventStage::Success(output)) => {
emit_exec_end(
codex-rs/core/src/tools/handlers/shell.rs
+8 -5
View File
@@ -11,6 +11,7 @@ use crate::exec::ExecParams;
use crate::exec_env::create_env;
use crate::function_tool::FunctionCallError;
use crate::is_safe_command::is_known_safe_command;
use crate::protocol::ExecCommandSource;
use crate::tools::context::ToolInvocation;
use crate::tools::context::ToolOutput;
use crate::tools::context::ToolPayload;
@@ -285,11 +286,13 @@ impl ShellHandler {
}
// Regular shell execution path.
let emitter = ToolEmitter::shell(
exec_params.command.clone(),
exec_params.cwd.clone(),
is_user_shell_command,
);
let source = if is_user_shell_command {
ExecCommandSource::UserShell
} else {
ExecCommandSource::Agent
};
let emitter =
ToolEmitter::shell(exec_params.command.clone(), exec_params.cwd.clone(), source);
let event_ctx = ToolEventCtx::new(session.as_ref(), turn.as_ref(), &call_id, None);
emitter.begin(event_ctx).await;
codex-rs/core/src/tools/handlers/unified_exec.rs
+8 -2
View File
@@ -4,6 +4,7 @@ use crate::function_tool::FunctionCallError;
use crate::is_safe_command::is_known_safe_command;
use crate::protocol::EventMsg;
use crate::protocol::ExecCommandOutputDeltaEvent;
use crate::protocol::ExecCommandSource;
use crate::protocol::ExecOutputStream;
use crate::shell::get_shell_by_model_provided_path;
use crate::tools::context::ToolInvocation;
@@ -162,8 +163,12 @@ impl ToolHandler for UnifiedExecHandler {
&context.call_id,
None,
);
let emitter = ToolEmitter::unified_exec(&command, cwd.clone(), true);
let emitter = ToolEmitter::unified_exec(
&command,
cwd.clone(),
ExecCommandSource::UnifiedExecStartup,
None,
);
emitter.emit(event_ctx, ToolEventStage::Begin).await;
manager
@@ -191,6 +196,7 @@ impl ToolHandler for UnifiedExecHandler {
})?;
manager
.write_stdin(WriteStdinRequest {
call_id: &call_id,
session_id: args.session_id,
input: &args.chars,
yield_time_ms: args.yield_time_ms,
codex-rs/core/src/unified_exec/mod.rs
+3
View File
@@ -74,6 +74,7 @@ pub(crate) struct ExecCommandRequest {
#[derive(Debug)]
pub(crate) struct WriteStdinRequest<'a> {
pub call_id: &'a str,
pub session_id: i32,
pub input: &'a str,
pub yield_time_ms: u64,
@@ -89,6 +90,7 @@ pub(crate) struct UnifiedExecResponse {
pub session_id: Option<i32>,
pub exit_code: Option<i32>,
pub original_token_count: Option<usize>,
pub session_command: Option<Vec<String>>,
}
#[derive(Default)]
@@ -213,6 +215,7 @@ mod tests {
.services
.unified_exec_manager
.write_stdin(WriteStdinRequest {
call_id: "write-stdin",
session_id,
input,
yield_time_ms,
codex-rs/core/src/unified_exec/session_manager.rs
+136 -19
View File
@@ -6,12 +6,18 @@ use tokio::sync::mpsc;
use tokio::time::Duration;
use tokio::time::Instant;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::exec::ExecToolCallOutput;
use crate::exec::StreamOutput;
use crate::exec_env::create_env;
use crate::protocol::BackgroundEventEvent;
use crate::protocol::EventMsg;
use crate::protocol::ExecCommandSource;
use crate::sandboxing::ExecEnv;
use crate::tools::events::ToolEmitter;
use crate::tools::events::ToolEventCtx;
use crate::tools::events::ToolEventFailure;
use crate::tools::events::ToolEventStage;
use crate::tools::orchestrator::ToolOrchestrator;
use crate::tools::runtimes::unified_exec::UnifiedExecRequest as UnifiedExecToolRequest;
@@ -66,15 +72,18 @@ impl UnifiedExecSessionManager {
let text = String::from_utf8_lossy(&collected).to_string();
let (output, original_token_count) = truncate_output_to_tokens(&text, max_tokens);
let chunk_id = generate_chunk_id();
let exit_code = session.exit_code();
let session_id = if session.has_exited() {
None
} else {
Some(
self.store_session(session, context, &request.command, cwd.clone(), start)
.await,
)
};
let has_exited = session.has_exited();
let stored_id = self
.store_session(session, context, &request.command, cwd.clone(), start)
.await;
let exit_code = self
.sessions
.lock()
.await
.get(&stored_id)
.map(|entry| entry.session.exit_code());
// Only include a session_id in the response if the process is still alive.
let session_id = if has_exited { None } else { Some(stored_id) };
let response = UnifiedExecResponse {
event_call_id: context.call_id.clone(),
@@ -82,10 +91,15 @@ impl UnifiedExecSessionManager {
wall_time,
output,
session_id,
exit_code,
exit_code: exit_code.flatten(),
original_token_count,
session_command: Some(request.command.clone()),
};
if response.session_id.is_some() {
Self::emit_waiting_status(&context.session, &context.turn, &request.command).await;
}
// If the command completed during this call, emit an ExecCommandEnd via the emitter.
if response.session_id.is_none() {
let exit = response.exit_code.unwrap_or(-1);
@@ -109,11 +123,46 @@ impl UnifiedExecSessionManager {
) -> Result<UnifiedExecResponse, UnifiedExecError> {
let session_id = request.session_id;
let (writer_tx, output_buffer, output_notify) =
self.prepare_session_handles(session_id).await?;
let (
writer_tx,
output_buffer,
output_notify,
session_ref,
turn_ref,
session_command,
session_cwd,
) = self.prepare_session_handles(session_id).await?;
let interaction_emitter = ToolEmitter::unified_exec(
&session_command,
session_cwd.clone(),
ExecCommandSource::UnifiedExecInteraction,
(!request.input.is_empty()).then(|| request.input.to_string()),
);
let make_event_ctx = || {
ToolEventCtx::new(
session_ref.as_ref(),
turn_ref.as_ref(),
request.call_id,
None,
)
};
interaction_emitter
.emit(make_event_ctx(), ToolEventStage::Begin)
.await;
if !request.input.is_empty() {
Self::send_input(&writer_tx, request.input.as_bytes()).await?;
if let Err(err) = Self::send_input(&writer_tx, request.input.as_bytes()).await {
interaction_emitter
.emit(
make_event_ctx(),
ToolEventStage::Failure(ToolEventFailure::Message(format!(
"write_stdin failed: {err:?}"
))),
)
.await;
return Err(err);
}
tokio::time::sleep(Duration::from_millis(100)).await;
}
@@ -151,8 +200,28 @@ impl UnifiedExecSessionManager {
session_id,
exit_code,
original_token_count,
session_command: Some(session_command.clone()),
};
let interaction_output = ExecToolCallOutput {
exit_code: response.exit_code.unwrap_or(0),
stdout: StreamOutput::new(response.output.clone()),
stderr: StreamOutput::new(String::new()),
aggregated_output: StreamOutput::new(response.output.clone()),
duration: response.wall_time,
timed_out: false,
};
interaction_emitter
.emit(
make_event_ctx(),
ToolEventStage::Success(interaction_output),
)
.await;
if response.session_id.is_some() {
Self::emit_waiting_status(&session_ref, &turn_ref, &session_command).await;
}
if let (Some(exit), Some(entry)) = (response.exit_code, completion_entry) {
let total_duration = Instant::now().saturating_duration_since(entry.started_at);
Self::emit_exec_end_from_entry(entry, response.output.clone(), exit, total_duration)
@@ -189,17 +258,44 @@ impl UnifiedExecSessionManager {
async fn prepare_session_handles(
&self,
session_id: i32,
) -> Result<(mpsc::Sender<Vec<u8>>, OutputBuffer, Arc<Notify>), UnifiedExecError> {
) -> Result<
(
mpsc::Sender<Vec<u8>>,
OutputBuffer,
Arc<Notify>,
Arc<Session>,
Arc<TurnContext>,
Vec<String>,
PathBuf,
),
UnifiedExecError,
> {
let sessions = self.sessions.lock().await;
let (output_buffer, output_notify, writer_tx) =
let (output_buffer, output_notify, writer_tx, session, turn, command, cwd) =
if let Some(entry) = sessions.get(&session_id) {
let (buffer, notify) = entry.session.output_handles();
(buffer, notify, entry.session.writer_sender())
(
buffer,
notify,
entry.session.writer_sender(),
Arc::clone(&entry.session_ref),
Arc::clone(&entry.turn_ref),
entry.command.clone(),
entry.cwd.clone(),
)
} else {
return Err(UnifiedExecError::UnknownSessionId { session_id });
};
Ok((writer_tx, output_buffer, output_notify))
Ok((
writer_tx,
output_buffer,
output_notify,
session,
turn,
command,
cwd,
))
}
async fn send_input(
@@ -256,7 +352,12 @@ impl UnifiedExecSessionManager {
&entry.call_id,
None,
);
let emitter = ToolEmitter::unified_exec(&entry.command, entry.cwd, true);
let emitter = ToolEmitter::unified_exec(
&entry.command,
entry.cwd,
ExecCommandSource::UnifiedExecStartup,
None,
);
emitter
.emit(event_ctx, ToolEventStage::Success(output))
.await;
@@ -284,12 +385,28 @@ impl UnifiedExecSessionManager {
&context.call_id,
None,
);
let emitter = ToolEmitter::unified_exec(command, cwd, true);
let emitter =
ToolEmitter::unified_exec(command, cwd, ExecCommandSource::UnifiedExecStartup, None);
emitter
.emit(event_ctx, ToolEventStage::Success(output))
.await;
}
async fn emit_waiting_status(
session: &Arc<Session>,
turn: &Arc<TurnContext>,
command: &[String],
) {
let command_display = command.join(" ");
let message = format!("Waiting for `{command_display}`");
session
.send_event(
turn.as_ref(),
EventMsg::BackgroundEvent(BackgroundEventEvent { message }),
)
.await;
}
pub(crate) async fn open_session_with_exec_env(
&self,
env: &ExecEnv,