https://github.com/openai/codex/pull/642 introduced support for the `--disable-response-storage` flag, but if you are a ZDR customer, it is tedious to set this every time, so this PR makes it possible to set this once in `config.toml` and be done with it. Incidentally, this tidies things up such that now `init_codex()` takes only one parameter: `Config`.
220 lines
7.2 KiB
Rust
220 lines
7.2 KiB
Rust
mod cli;
|
||
use std::sync::Arc;
|
||
|
||
pub use cli::Cli;
|
||
use codex_core::codex_wrapper;
|
||
use codex_core::config::Config;
|
||
use codex_core::config::ConfigOverrides;
|
||
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::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,
|
||
sandbox_policy,
|
||
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);
|
||
}
|
||
|
||
// Load configuration and determine approval policy
|
||
let overrides = ConfigOverrides {
|
||
model: model.clone(),
|
||
// This CLI is intended to be headless and has no affordances for asking
|
||
// the user for approval.
|
||
approval_policy: Some(AskForApproval::Never),
|
||
sandbox_policy: sandbox_policy.map(Into::into),
|
||
disable_response_storage: if disable_response_storage {
|
||
Some(true)
|
||
} else {
|
||
None
|
||
},
|
||
};
|
||
let config = Config::load_with_overrides(overrides)?;
|
||
let (codex_wrapper, event, ctrl_c) = codex_wrapper::init_codex(config).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 in‑flight 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",
|
||
}
|
||
}
|