Files
codex/codex-rs/exec/src/lib.rs
Michael Bolin b323d10ea7 feat: add ZDR support to Rust implementation (#642)
This adds support for the `--disable-response-storage` flag across our
multiple Rust CLIs to support customers who have opted into Zero-Data
Retention (ZDR). The analogous changes to the TypeScript CLI were:

* https://github.com/openai/codex/pull/481
* https://github.com/openai/codex/pull/543

For a client using ZDR, `previous_response_id` will never be available,
so the `input` field of an API request must include the full transcript
of the conversation thus far. As such, this PR changes the type of
`Prompt.input` from `Vec<ResponseInputItem>` to `Vec<ResponseItem>`.

Practically speaking, `ResponseItem` was effectively a "superset" of
`ResponseInputItem` already. The main difference for us is that
`ResponseItem` includes the `FunctionCall` variant that we have to
include as part of the conversation history in the ZDR case.

Another key change in this PR is modifying `try_run_turn()` so that it
returns the `Vec<ResponseItem>` for the turn in addition to the
`Vec<ResponseInputItem>` produced by `try_run_turn()`. This is because
the caller of `run_turn()` needs to record the `Vec<ResponseItem>` when
ZDR is enabled.

To that end, this PR introduces `ZdrTranscript` (and adds
`zdr_transcript: Option<ZdrTranscript>` to `struct State` in `codex.rs`)
to take responsibility for maintaining the conversation transcript in
the ZDR case.
2025-04-25 12:08:18 -07:00

215 lines
7.0 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
mod cli;
use std::sync::Arc;
pub use cli::Cli;
use codex_core::codex_wrapper;
use codex_core::protocol::AskForApproval;
use codex_core::protocol::Event;
use codex_core::protocol::EventMsg;
use codex_core::protocol::FileChange;
use codex_core::protocol::InputItem;
use codex_core::protocol::Op;
use codex_core::protocol::SandboxPolicy;
use codex_core::util::is_inside_git_repo;
use tracing::debug;
use tracing::error;
use tracing::info;
use tracing_subscriber::EnvFilter;
pub async fn run_main(cli: Cli) -> anyhow::Result<()> {
// TODO(mbolin): Take a more thoughtful approach to logging.
let default_level = "error";
let allow_ansi = true;
let _ = tracing_subscriber::fmt()
.with_env_filter(
EnvFilter::try_from_default_env()
.or_else(|_| EnvFilter::try_new(default_level))
.unwrap(),
)
.with_ansi(allow_ansi)
.with_writer(std::io::stderr)
.try_init();
let Cli {
images,
model,
skip_git_repo_check,
disable_response_storage,
prompt,
..
} = cli;
if !skip_git_repo_check && !is_inside_git_repo() {
eprintln!("Not inside a Git repo and --skip-git-repo-check was not specified.");
std::process::exit(1);
} else if images.is_empty() && prompt.is_none() {
eprintln!("No images or prompt specified.");
std::process::exit(1);
}
// TODO(mbolin): We are reworking the CLI args right now, so this will
// likely come from a new --execution-policy arg.
let approval_policy = AskForApproval::Never;
let sandbox_policy = SandboxPolicy::NetworkAndFileWriteRestricted;
let (codex_wrapper, event, ctrl_c) = codex_wrapper::init_codex(
approval_policy,
sandbox_policy,
disable_response_storage,
model,
)
.await?;
let codex = Arc::new(codex_wrapper);
info!("Codex initialized with event: {event:?}");
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
{
let codex = codex.clone();
tokio::spawn(async move {
loop {
let interrupted = ctrl_c.notified();
tokio::select! {
_ = interrupted => {
// Forward an interrupt to the codex so it can abort any inflight task.
let _ = codex
.submit(
Op::Interrupt,
)
.await;
// Exit the inner loop and return to the main input prompt. The codex
// will emit a `TurnInterrupted` (Error) event which is drained later.
break;
}
res = codex.next_event() => match res {
Ok(event) => {
debug!("Received event: {event:?}");
process_event(&event);
if let Err(e) = tx.send(event) {
error!("Error sending event: {e:?}");
break;
}
},
Err(e) => {
error!("Error receiving event: {e:?}");
break;
}
}
}
}
});
}
if !images.is_empty() {
// Send images first.
let items: Vec<InputItem> = images
.into_iter()
.map(|path| InputItem::LocalImage { path })
.collect();
let initial_images_event_id = codex.submit(Op::UserInput { items }).await?;
info!("Sent images with event ID: {initial_images_event_id}");
while let Ok(event) = codex.next_event().await {
if event.id == initial_images_event_id && matches!(event.msg, EventMsg::TaskComplete) {
break;
}
}
}
if let Some(prompt) = prompt {
// Send the prompt.
let items: Vec<InputItem> = vec![InputItem::Text { text: prompt }];
let initial_prompt_task_id = codex.submit(Op::UserInput { items }).await?;
info!("Sent prompt with event ID: {initial_prompt_task_id}");
while let Some(event) = rx.recv().await {
if event.id == initial_prompt_task_id && matches!(event.msg, EventMsg::TaskComplete) {
break;
}
}
}
Ok(())
}
fn process_event(event: &Event) {
let Event { id, msg } = event;
match msg {
EventMsg::Error { message } => {
println!("Error: {message}");
}
EventMsg::BackgroundEvent { .. } => {
// Ignore these for now.
}
EventMsg::TaskStarted => {
println!("Task started: {id}");
}
EventMsg::TaskComplete => {
println!("Task complete: {id}");
}
EventMsg::AgentMessage { message } => {
println!("Agent message: {message}");
}
EventMsg::ExecCommandBegin {
call_id,
command,
cwd,
} => {
println!("exec('{call_id}'): {:?} in {cwd}", command);
}
EventMsg::ExecCommandEnd {
call_id,
stdout,
stderr,
exit_code,
} => {
let output = if *exit_code == 0 { stdout } else { stderr };
let truncated_output = output.lines().take(5).collect::<Vec<_>>().join("\n");
println!("exec('{call_id}') exited {exit_code}:\n{truncated_output}");
}
EventMsg::PatchApplyBegin {
call_id,
auto_approved,
changes,
} => {
let changes = changes
.iter()
.map(|(path, change)| {
format!("{} {}", format_file_change(change), path.to_string_lossy())
})
.collect::<Vec<_>>()
.join("\n");
println!("apply_patch('{call_id}') auto_approved={auto_approved}:\n{changes}");
}
EventMsg::PatchApplyEnd {
call_id,
stdout,
stderr,
success,
} => {
let (exit_code, output) = if *success { (0, stdout) } else { (1, stderr) };
let truncated_output = output.lines().take(5).collect::<Vec<_>>().join("\n");
println!("apply_patch('{call_id}') exited {exit_code}:\n{truncated_output}");
}
EventMsg::ExecApprovalRequest { .. } => {
// Should we exit?
}
EventMsg::ApplyPatchApprovalRequest { .. } => {
// Should we exit?
}
_ => {
// Ignore event.
}
}
}
fn format_file_change(change: &FileChange) -> &'static str {
match change {
FileChange::Add { .. } => "A",
FileChange::Delete => "D",
FileChange::Update {
move_path: Some(_), ..
} => "R",
FileChange::Update {
move_path: None, ..
} => "M",
}
}