Files
codex/codex-rs/repl/src/lib.rs
T
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

425 lines
15 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.
use std::io::stdin;
use std::io::stdout;
use std::io::Write;
use std::sync::Arc;
use codex_core::config::Config;
use codex_core::protocol;
use codex_core::protocol::FileChange;
use codex_core::util::is_inside_git_repo;
use codex_core::util::notify_on_sigint;
use codex_core::Codex;
use owo_colors::OwoColorize;
use owo_colors::Style;
use tokio::io::AsyncBufReadExt;
use tokio::io::BufReader;
use tokio::io::Lines;
use tokio::io::Stdin;
use tokio::sync::Notify;
use tracing::debug;
use tracing_subscriber::EnvFilter;
mod cli;
pub use cli::Cli;
/// Initialize the global logger once at startup based on the `--verbose` flag.
fn init_logger(verbose: u8, allow_ansi: bool) {
// Map -v occurrences to explicit log levels:
// 0 → warn (default)
// 1 → info
// 2 → debug
// ≥3 → trace
let default_level = match verbose {
0 => "warn",
1 => "info",
2 => "codex=debug",
_ => "codex=trace",
};
// Only initialize the logger once – repeated calls are ignored. `try_init` will return an
// error if another crate (like tests) initialized it first, which we can safely ignore.
// By default `tracing_subscriber::fmt()` writes formatted logs to stderr. That is fine when
// running the CLI manually but in our smoke tests we capture **stdout** (via `assert_cmd`) and
// ignore stderr. As a result none of the `tracing::info!` banners or warnings show up in the
// recorded output making it much harder to debug live runs.
// Switch the logger's writer to stdout so both human runs and the integration tests see the
// same stream. Disable ANSI colors because the binary already prints plain text and color
// escape codes make predicate matching brittle.
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::stdout)
.try_init();
}
pub async fn run_main(cli: Cli) -> anyhow::Result<()> {
let ctrl_c = notify_on_sigint();
// Abort early when the user runs Codex outside a Git repository unless
// they explicitly acknowledged the risks with `--allow-no-git-exec`.
if !cli.allow_no_git_exec && !is_inside_git_repo() {
eprintln!(
"We recommend running codex inside a git repository. \
If you understand the risks, you can proceed with \
`--allow-no-git-exec`."
);
std::process::exit(1);
}
// Initialize logging before any other work so early errors are captured.
init_logger(cli.verbose, !cli.no_ansi);
let config = Config::load().unwrap_or_default();
codex_main(cli, config, ctrl_c).await
}
async fn codex_main(mut cli: Cli, cfg: Config, ctrl_c: Arc<Notify>) -> anyhow::Result<()> {
let mut builder = Codex::builder();
if let Some(path) = cli.record_submissions {
builder = builder.record_submissions(path);
}
if let Some(path) = cli.record_events {
builder = builder.record_events(path);
}
let codex = builder.spawn(Arc::clone(&ctrl_c))?;
let init_id = random_id();
let init = protocol::Submission {
id: init_id.clone(),
op: protocol::Op::ConfigureSession {
model: cli.model.or(cfg.model),
instructions: cfg.instructions,
approval_policy: cli.approval_policy.into(),
sandbox_policy: cli.sandbox_policy.into(),
disable_response_storage: cli.disable_response_storage,
},
};
out(
"initializing session",
MessagePriority::BackgroundEvent,
MessageActor::User,
);
codex.submit(init).await?;
// init
loop {
out(
"waiting for session initialization",
MessagePriority::BackgroundEvent,
MessageActor::User,
);
let event = codex.next_event().await?;
if event.id == init_id {
if let protocol::EventMsg::Error { message } = event.msg {
anyhow::bail!("Error during initialization: {message}");
} else {
out(
"session initialized",
MessagePriority::BackgroundEvent,
MessageActor::User,
);
break;
}
}
}
// run loop
let mut reader = InputReader::new(ctrl_c.clone());
loop {
let text = match cli.prompt.take() {
Some(input) => input,
None => match reader.request_input().await? {
Some(input) => input,
None => {
// ctrl + d
println!();
return Ok(());
}
},
};
if text.trim().is_empty() {
continue;
}
// Interpret certain single‑word commands as immediate termination requests.
let trimmed = text.trim();
if trimmed == "q" {
// Exit gracefully.
println!("Exiting…");
return Ok(());
}
let sub = protocol::Submission {
id: random_id(),
op: protocol::Op::UserInput {
items: vec![protocol::InputItem::Text { text }],
},
};
out(
"sending request to model",
MessagePriority::TaskProgress,
MessageActor::User,
);
codex.submit(sub).await?;
// Wait for agent events **or** user interrupts (Ctrl+C).
'inner: loop {
// Listen for either the next agent event **or** a SIGINT notification. Using
// `tokio::select!` allows the user to cancel a long‑running request that would
// otherwise leave the CLI stuck waiting for a server response.
let event = {
let interrupted = ctrl_c.notified();
tokio::select! {
_ = interrupted => {
// Forward an interrupt to the agent so it can abort any in‑flight task.
let _ = codex
.submit(protocol::Submission {
id: random_id(),
op: protocol::Op::Interrupt,
})
.await;
// Exit the inner loop and return to the main input prompt. The agent
// will emit a `TurnInterrupted` (Error) event which is drained later.
break 'inner;
}
res = codex.next_event() => res?
}
};
debug!(?event, "Got event");
let id = event.id;
match event.msg {
protocol::EventMsg::Error { message } => {
println!("Error: {message}");
break 'inner;
}
protocol::EventMsg::TaskComplete => break 'inner,
protocol::EventMsg::AgentMessage { message } => {
out(&message, MessagePriority::UserMessage, MessageActor::Agent)
}
protocol::EventMsg::SessionConfigured { model } => {
debug!(model, "Session initialized");
}
protocol::EventMsg::ExecApprovalRequest {
command,
cwd,
reason,
} => {
let reason_str = reason
.as_deref()
.map(|r| format!(" [{r}]"))
.unwrap_or_default();
let prompt = format!(
"approve command in {} {}{} (y/N): ",
cwd.display(),
command.join(" "),
reason_str
);
let decision = request_user_approval2(prompt)?;
let sub = protocol::Submission {
id: random_id(),
op: protocol::Op::ExecApproval { id, decision },
};
out(
"submitting command approval",
MessagePriority::TaskProgress,
MessageActor::User,
);
codex.submit(sub).await?;
}
protocol::EventMsg::ApplyPatchApprovalRequest {
changes,
reason: _,
grant_root: _,
} => {
let file_list = changes
.keys()
.map(|path| path.to_string_lossy().to_string())
.collect::<Vec<_>>()
.join(", ");
let request =
format!("approve apply_patch that will touch? {file_list} (y/N): ");
let decision = request_user_approval2(request)?;
let sub = protocol::Submission {
id: random_id(),
op: protocol::Op::PatchApproval { id, decision },
};
out(
"submitting patch approval",
MessagePriority::UserMessage,
MessageActor::Agent,
);
codex.submit(sub).await?;
}
protocol::EventMsg::ExecCommandBegin {
command,
cwd,
call_id: _,
} => {
out(
&format!("running command: '{}' in '{}'", command.join(" "), cwd),
MessagePriority::BackgroundEvent,
MessageActor::Agent,
);
}
protocol::EventMsg::ExecCommandEnd {
stdout,
stderr,
exit_code,
call_id: _,
} => {
let msg = if exit_code == 0 {
"command completed (exit 0)".to_string()
} else {
// Prefer stderr but fall back to stdout if empty.
let err_snippet = if !stderr.trim().is_empty() {
stderr.trim()
} else {
stdout.trim()
};
format!("command failed (exit {exit_code}): {err_snippet}")
};
out(&msg, MessagePriority::BackgroundEvent, MessageActor::Agent);
out(
"sending results to model",
MessagePriority::TaskProgress,
MessageActor::Agent,
);
}
protocol::EventMsg::PatchApplyBegin { changes, .. } => {
// Emit PatchApplyBegin so the front‑end can show progress.
let summary = changes
.iter()
.map(|(path, change)| match change {
FileChange::Add { .. } => format!("A {}", path.display()),
FileChange::Delete => format!("D {}", path.display()),
FileChange::Update { .. } => format!("M {}", path.display()),
})
.collect::<Vec<_>>()
.join(", ");
out(
&format!("applying patch: {summary}"),
MessagePriority::BackgroundEvent,
MessageActor::Agent,
);
}
protocol::EventMsg::PatchApplyEnd { success, .. } => {
let status = if success { "success" } else { "failed" };
out(
&format!("patch application {status}"),
MessagePriority::BackgroundEvent,
MessageActor::Agent,
);
out(
"sending results to model",
MessagePriority::TaskProgress,
MessageActor::Agent,
);
}
// Broad fallback; if the CLI is unaware of an event type, it will just
// print it as a generic BackgroundEvent.
e => {
out(
&format!("event: {e:?}"),
MessagePriority::BackgroundEvent,
MessageActor::Agent,
);
}
}
}
}
}
fn random_id() -> String {
let id: u64 = rand::random();
id.to_string()
}
fn request_user_approval2(request: String) -> anyhow::Result<protocol::ReviewDecision> {
println!("{}", request);
let mut line = String::new();
stdin().read_line(&mut line)?;
let answer = line.trim().to_ascii_lowercase();
let is_accepted = answer == "y" || answer == "yes";
let decision = if is_accepted {
protocol::ReviewDecision::Approved
} else {
protocol::ReviewDecision::Denied
};
Ok(decision)
}
#[derive(Debug, Clone, Copy)]
enum MessagePriority {
BackgroundEvent,
TaskProgress,
UserMessage,
}
enum MessageActor {
Agent,
User,
}
impl From<MessageActor> for String {
fn from(actor: MessageActor) -> Self {
match actor {
MessageActor::Agent => "codex".to_string(),
MessageActor::User => "user".to_string(),
}
}
}
fn out(msg: &str, priority: MessagePriority, actor: MessageActor) {
let actor: String = actor.into();
let style = match priority {
MessagePriority::BackgroundEvent => Style::new().fg_rgb::<127, 127, 127>(),
MessagePriority::TaskProgress => Style::new().fg_rgb::<200, 200, 200>(),
MessagePriority::UserMessage => Style::new().white(),
};
println!("{}> {}", actor.bold(), msg.style(style));
}
struct InputReader {
reader: Lines<BufReader<Stdin>>,
ctrl_c: Arc<Notify>,
}
impl InputReader {
pub fn new(ctrl_c: Arc<Notify>) -> Self {
Self {
reader: BufReader::new(tokio::io::stdin()).lines(),
ctrl_c,
}
}
pub async fn request_input(&mut self) -> std::io::Result<Option<String>> {
print!("user> ");
stdout().flush()?;
let interrupted = self.ctrl_c.notified();
tokio::select! {
line = self.reader.next_line() => {
match line? {
Some(input) => Ok(Some(input.trim().to_string())),
None => Ok(None),
}
}
_ = interrupted => {
println!();
Ok(Some(String::new()))
}
}
}
}