Add user command event types (#6246)
adding new user command event, logic in TUI to render user command events
This commit is contained in:
zhao-oai
GitHub
@@ -14,6 +14,7 @@ use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::user_instructions::UserInstructions;
|
||||
use crate::user_shell_command::is_user_shell_command_text;
|
||||
|
||||
fn is_session_prefix(text: &str) -> bool {
|
||||
let trimmed = text.trim_start();
|
||||
@@ -31,7 +32,7 @@ fn parse_user_message(message: &[ContentItem]) -> Option<UserMessageItem> {
|
||||
for content_item in message.iter() {
|
||||
match content_item {
|
||||
ContentItem::InputText { text } => {
|
||||
if is_session_prefix(text) {
|
||||
if is_session_prefix(text) || is_user_shell_command_text(text) {
|
||||
return None;
|
||||
}
|
||||
content.push(UserInput::Text { text: text.clone() });
|
||||
@@ -197,7 +198,14 @@ mod tests {
|
||||
text: "# AGENTS.md instructions for test_directory\n\n<INSTRUCTIONS>\ntest_text\n</INSTRUCTIONS>".to_string(),
|
||||
}],
|
||||
},
|
||||
];
|
||||
ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text: "<user_shell_command>echo 42</user_shell_command>".to_string(),
|
||||
}],
|
||||
},
|
||||
];
|
||||
|
||||
for item in items {
|
||||
let turn_item = parse_turn_item(&item);
|
||||
|
||||
@@ -81,6 +81,7 @@ mod function_tool;
|
||||
mod state;
|
||||
mod tasks;
|
||||
mod user_notification;
|
||||
mod user_shell_command;
|
||||
pub mod util;
|
||||
|
||||
pub use apply_patch::CODEX_APPLY_PATCH_ARG1;
|
||||
|
||||
@@ -1,28 +1,35 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use codex_protocol::models::ShellToolCallParams;
|
||||
use codex_async_utils::CancelErr;
|
||||
use codex_async_utils::OrCancelExt;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::error;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::codex::TurnContext;
|
||||
use crate::exec::ExecToolCallOutput;
|
||||
use crate::exec::SandboxType;
|
||||
use crate::exec::StdoutStream;
|
||||
use crate::exec::StreamOutput;
|
||||
use crate::exec::execute_exec_env;
|
||||
use crate::exec_env::create_env;
|
||||
use crate::parse_command::parse_command;
|
||||
use crate::protocol::EventMsg;
|
||||
use crate::protocol::ExecCommandBeginEvent;
|
||||
use crate::protocol::ExecCommandEndEvent;
|
||||
use crate::protocol::SandboxPolicy;
|
||||
use crate::protocol::TaskStartedEvent;
|
||||
use crate::sandboxing::ExecEnv;
|
||||
use crate::state::TaskKind;
|
||||
use crate::tools::context::ToolPayload;
|
||||
use crate::tools::parallel::ToolCallRuntime;
|
||||
use crate::tools::router::ToolCall;
|
||||
use crate::tools::router::ToolRouter;
|
||||
use crate::turn_diff_tracker::TurnDiffTracker;
|
||||
use crate::tools::format_exec_output_str;
|
||||
use crate::user_shell_command::user_shell_command_record_item;
|
||||
|
||||
use super::SessionTask;
|
||||
use super::SessionTaskContext;
|
||||
|
||||
const USER_SHELL_TOOL_NAME: &str = "local_shell";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct UserShellCommandTask {
|
||||
command: String,
|
||||
@@ -78,34 +85,126 @@ impl SessionTask for UserShellCommandTask {
|
||||
}
|
||||
};
|
||||
|
||||
let params = ShellToolCallParams {
|
||||
let call_id = Uuid::new_v4().to_string();
|
||||
let raw_command = self.command.clone();
|
||||
|
||||
let parsed_cmd = parse_command(&shell_invocation);
|
||||
session
|
||||
.send_event(
|
||||
turn_context.as_ref(),
|
||||
EventMsg::ExecCommandBegin(ExecCommandBeginEvent {
|
||||
call_id: call_id.clone(),
|
||||
command: shell_invocation.clone(),
|
||||
cwd: turn_context.cwd.clone(),
|
||||
parsed_cmd,
|
||||
is_user_shell_command: true,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let exec_env = ExecEnv {
|
||||
command: shell_invocation,
|
||||
workdir: None,
|
||||
cwd: turn_context.cwd.clone(),
|
||||
env: create_env(&turn_context.shell_environment_policy),
|
||||
timeout_ms: None,
|
||||
sandbox: SandboxType::None,
|
||||
with_escalated_permissions: None,
|
||||
justification: None,
|
||||
arg0: None,
|
||||
};
|
||||
|
||||
let tool_call = ToolCall {
|
||||
tool_name: USER_SHELL_TOOL_NAME.to_string(),
|
||||
call_id: Uuid::new_v4().to_string(),
|
||||
payload: ToolPayload::LocalShell { params },
|
||||
};
|
||||
let stdout_stream = Some(StdoutStream {
|
||||
sub_id: turn_context.sub_id.clone(),
|
||||
call_id: call_id.clone(),
|
||||
tx_event: session.get_tx_event(),
|
||||
});
|
||||
|
||||
let router = Arc::new(ToolRouter::from_config(&turn_context.tools_config, None));
|
||||
let tracker = Arc::new(Mutex::new(TurnDiffTracker::new()));
|
||||
let runtime = ToolCallRuntime::new(
|
||||
Arc::clone(&router),
|
||||
Arc::clone(&session),
|
||||
Arc::clone(&turn_context),
|
||||
Arc::clone(&tracker),
|
||||
);
|
||||
let sandbox_policy = SandboxPolicy::DangerFullAccess;
|
||||
let exec_result = execute_exec_env(exec_env, &sandbox_policy, stdout_stream)
|
||||
.or_cancel(&cancellation_token)
|
||||
.await;
|
||||
|
||||
if let Err(err) = runtime
|
||||
.handle_tool_call(tool_call, cancellation_token)
|
||||
.await
|
||||
{
|
||||
error!("user shell command failed: {err:?}");
|
||||
match exec_result {
|
||||
Err(CancelErr::Cancelled) => {
|
||||
let aborted_message = "command aborted by user".to_string();
|
||||
let exec_output = ExecToolCallOutput {
|
||||
exit_code: -1,
|
||||
stdout: StreamOutput::new(String::new()),
|
||||
stderr: StreamOutput::new(aborted_message.clone()),
|
||||
aggregated_output: StreamOutput::new(aborted_message.clone()),
|
||||
duration: Duration::ZERO,
|
||||
timed_out: false,
|
||||
};
|
||||
let output_items = [user_shell_command_record_item(&raw_command, &exec_output)];
|
||||
session
|
||||
.record_conversation_items(turn_context.as_ref(), &output_items)
|
||||
.await;
|
||||
session
|
||||
.send_event(
|
||||
turn_context.as_ref(),
|
||||
EventMsg::ExecCommandEnd(ExecCommandEndEvent {
|
||||
call_id,
|
||||
stdout: String::new(),
|
||||
stderr: aborted_message.clone(),
|
||||
aggregated_output: aborted_message.clone(),
|
||||
exit_code: -1,
|
||||
duration: Duration::ZERO,
|
||||
formatted_output: aborted_message,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Ok(Ok(output)) => {
|
||||
session
|
||||
.send_event(
|
||||
turn_context.as_ref(),
|
||||
EventMsg::ExecCommandEnd(ExecCommandEndEvent {
|
||||
call_id: call_id.clone(),
|
||||
stdout: output.stdout.text.clone(),
|
||||
stderr: output.stderr.text.clone(),
|
||||
aggregated_output: output.aggregated_output.text.clone(),
|
||||
exit_code: output.exit_code,
|
||||
duration: output.duration,
|
||||
formatted_output: format_exec_output_str(&output),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
let output_items = [user_shell_command_record_item(&raw_command, &output)];
|
||||
session
|
||||
.record_conversation_items(turn_context.as_ref(), &output_items)
|
||||
.await;
|
||||
}
|
||||
Ok(Err(err)) => {
|
||||
error!("user shell command failed: {err:?}");
|
||||
let message = format!("execution error: {err:?}");
|
||||
let exec_output = ExecToolCallOutput {
|
||||
exit_code: -1,
|
||||
stdout: StreamOutput::new(String::new()),
|
||||
stderr: StreamOutput::new(message.clone()),
|
||||
aggregated_output: StreamOutput::new(message.clone()),
|
||||
duration: Duration::ZERO,
|
||||
timed_out: false,
|
||||
};
|
||||
session
|
||||
.send_event(
|
||||
turn_context.as_ref(),
|
||||
EventMsg::ExecCommandEnd(ExecCommandEndEvent {
|
||||
call_id,
|
||||
stdout: exec_output.stdout.text.clone(),
|
||||
stderr: exec_output.stderr.text.clone(),
|
||||
aggregated_output: exec_output.aggregated_output.text.clone(),
|
||||
exit_code: exec_output.exit_code,
|
||||
duration: exec_output.duration,
|
||||
formatted_output: format_exec_output_str(&exec_output),
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
let output_items = [user_shell_command_record_item(&raw_command, &exec_output)];
|
||||
session
|
||||
.record_conversation_items(turn_context.as_ref(), &output_items)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
108
codex-rs/core/src/user_shell_command.rs
Normal file
108
codex-rs/core/src/user_shell_command.rs
Normal file
@@ -0,0 +1,108 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
|
||||
use crate::exec::ExecToolCallOutput;
|
||||
use crate::tools::format_exec_output_str;
|
||||
|
||||
pub const USER_SHELL_COMMAND_OPEN: &str = "<user_shell_command>";
|
||||
pub const USER_SHELL_COMMAND_CLOSE: &str = "</user_shell_command>";
|
||||
|
||||
pub fn is_user_shell_command_text(text: &str) -> bool {
|
||||
let trimmed = text.trim_start();
|
||||
let lowered = trimmed.to_ascii_lowercase();
|
||||
lowered.starts_with(USER_SHELL_COMMAND_OPEN)
|
||||
}
|
||||
|
||||
fn format_duration_line(duration: Duration) -> String {
|
||||
let duration_seconds = duration.as_secs_f64();
|
||||
format!("Duration: {duration_seconds:.4} seconds")
|
||||
}
|
||||
|
||||
fn format_user_shell_command_body(command: &str, exec_output: &ExecToolCallOutput) -> String {
|
||||
let mut sections = Vec::new();
|
||||
sections.push("<command>".to_string());
|
||||
sections.push(command.to_string());
|
||||
sections.push("</command>".to_string());
|
||||
sections.push("<result>".to_string());
|
||||
sections.push(format!("Exit code: {}", exec_output.exit_code));
|
||||
sections.push(format_duration_line(exec_output.duration));
|
||||
sections.push("Output:".to_string());
|
||||
sections.push(format_exec_output_str(exec_output));
|
||||
sections.push("</result>".to_string());
|
||||
sections.join("\n")
|
||||
}
|
||||
|
||||
pub fn format_user_shell_command_record(command: &str, exec_output: &ExecToolCallOutput) -> String {
|
||||
let body = format_user_shell_command_body(command, exec_output);
|
||||
format!("{USER_SHELL_COMMAND_OPEN}\n{body}\n{USER_SHELL_COMMAND_CLOSE}")
|
||||
}
|
||||
|
||||
pub fn user_shell_command_record_item(
|
||||
command: &str,
|
||||
exec_output: &ExecToolCallOutput,
|
||||
) -> ResponseItem {
|
||||
ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text: format_user_shell_command_record(command, exec_output),
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::exec::StreamOutput;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn detects_user_shell_command_text_variants() {
|
||||
assert!(is_user_shell_command_text(
|
||||
"<user_shell_command>\necho hi\n</user_shell_command>"
|
||||
));
|
||||
assert!(!is_user_shell_command_text("echo hi"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formats_basic_record() {
|
||||
let exec_output = ExecToolCallOutput {
|
||||
exit_code: 0,
|
||||
stdout: StreamOutput::new("hi".to_string()),
|
||||
stderr: StreamOutput::new(String::new()),
|
||||
aggregated_output: StreamOutput::new("hi".to_string()),
|
||||
duration: Duration::from_secs(1),
|
||||
timed_out: false,
|
||||
};
|
||||
let item = user_shell_command_record_item("echo hi", &exec_output);
|
||||
let ResponseItem::Message { content, .. } = item else {
|
||||
panic!("expected message");
|
||||
};
|
||||
let [ContentItem::InputText { text }] = content.as_slice() else {
|
||||
panic!("expected input text");
|
||||
};
|
||||
assert_eq!(
|
||||
text,
|
||||
"<user_shell_command>\n<command>\necho hi\n</command>\n<result>\nExit code: 0\nDuration: 1.0000 seconds\nOutput:\nhi\n</result>\n</user_shell_command>"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uses_aggregated_output_over_streams() {
|
||||
let exec_output = ExecToolCallOutput {
|
||||
exit_code: 42,
|
||||
stdout: StreamOutput::new("stdout-only".to_string()),
|
||||
stderr: StreamOutput::new("stderr-only".to_string()),
|
||||
aggregated_output: StreamOutput::new("combined output wins".to_string()),
|
||||
duration: Duration::from_millis(120),
|
||||
timed_out: false,
|
||||
};
|
||||
let record = format_user_shell_command_record("false", &exec_output);
|
||||
assert_eq!(
|
||||
record,
|
||||
"<user_shell_command>\n<command>\nfalse\n</command>\n<result>\nExit code: 42\nDuration: 0.1200 seconds\nOutput:\ncombined output wins\n</result>\n</user_shell_command>"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user