From 9717ac87c1b8d8d6ba663264d3e62815e957c651 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 24 Apr 2025 23:54:20 -0700 Subject: [PATCH] feat: add ZDR support to Rust implementation --- codex-rs/core/src/client.rs | 12 ++++++- codex-rs/core/src/codex.rs | 38 +++++++++++++++++++-- codex-rs/core/src/codex_wrapper.rs | 2 ++ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/protocol.rs | 3 ++ codex-rs/core/src/zdr_transcript.rs | 26 ++++++++++++++ codex-rs/core/tests/live_agent.rs | 1 + codex-rs/core/tests/previous_response_id.rs | 1 + codex-rs/core/tests/stream_no_completed.rs | 1 + codex-rs/exec/src/cli.rs | 6 +++- codex-rs/exec/src/lib.rs | 3 +- codex-rs/repl/src/lib.rs | 2 ++ codex-rs/tui/src/chatwidget.rs | 3 +- 13 files changed, 93 insertions(+), 6 deletions(-) create mode 100644 codex-rs/core/src/zdr_transcript.rs diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 57f593a884..c4460be9aa 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -33,11 +33,17 @@ use crate::models::ResponseInputItem; use crate::models::ResponseItem; use crate::util::backoff; +/// API request payload for a single model turn #[derive(Default, Debug, Clone)] pub struct Prompt { + /// Conversation context input items pub input: Vec, + /// Optional previous response ID (when storage is enabled) pub prev_id: Option, + /// Optional initial instructions (only sent on first turn) pub instructions: Option, + /// Whether to store response on server side (disable_response_storage = !store) + pub store: bool, } #[derive(Debug)] @@ -58,7 +64,9 @@ struct Payload<'a> { reasoning: Option, #[serde(skip_serializing_if = "Option::is_none")] previous_response_id: Option, - stream: bool, + store: bool, + /// Stream responses via SSE + pub stream: bool, } #[derive(Debug, Serialize)] @@ -152,6 +160,8 @@ impl ModelClient { generate_summary: None, }), previous_response_id: prompt.prev_id.clone(), + // store response on server side when enabled + store: prompt.store, stream: true, }; diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index e57d3bbf07..08b6df733c 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -55,6 +55,7 @@ use crate::safety::assess_command_safety; use crate::safety::assess_patch_safety; use crate::safety::SafetyCheck; use crate::util::backoff; +use crate::zdr_transcript::ZdrTranscript; /// The high-level interface to the Codex system. /// It operates as a queue pair where you send submissions and receive events. @@ -190,6 +191,9 @@ impl Recorder { } } +/// Context for an initialized model agent +/// +/// A session has at most 1 running task at a time, and can be interrupted by user input. /// Context for an initialized model agent /// /// A session has at most 1 running task at a time, and can be interrupted by user input. @@ -201,6 +205,9 @@ struct Session { instructions: Option, approval_policy: AskForApproval, sandbox_policy: SandboxPolicy, + /// When true, omit `previous_response_id` in requests and send full context + disable_response_storage: bool, + /// Additional writable roots for sandbox execution writable_roots: Mutex>, state: Mutex, @@ -214,6 +221,7 @@ struct State { previous_response_id: Option, pending_approvals: HashMap>, pending_input: Vec, + zdr_transcript: Option, } impl Session { @@ -399,6 +407,7 @@ impl State { Self { approved_commands: self.approved_commands.clone(), previous_response_id: self.previous_response_id.clone(), + zdr_transcript: self.zdr_transcript.clone(), ..Default::default() } } @@ -489,6 +498,7 @@ async fn submission_loop( instructions, approval_policy, sandbox_policy, + disable_response_storage, } => { let model = model.unwrap_or_else(|| OPENAI_DEFAULT_MODEL.to_string()); info!(model, "Configuring session"); @@ -500,7 +510,14 @@ async fn submission_loop( sess.abort(); sess.state.lock().unwrap().partial_clone() } - None => State::default(), + None => State { + zdr_transcript: if disable_response_storage { + Some(ZdrTranscript::new()) + } else { + None + }, + ..Default::default() + }, }; // update session @@ -511,6 +528,7 @@ async fn submission_loop( instructions, approval_policy, sandbox_policy, + disable_response_storage, writable_roots: Mutex::new(get_writable_roots()), state: Mutex::new(state), })); @@ -592,12 +610,21 @@ async fn run_task(sess: Arc, sub_id: String, input: Vec) { let pending_input = sess.get_pending_input(); turn_input.splice(0..0, pending_input); + // If sess.state.transcript.is_some(), it should be written back into + // the prompt. match run_turn(&sess, sub_id.clone(), turn_input).await { Ok(turn_output) => { if turn_output.is_empty() { debug!("Turn completed"); break; } + + if let Some(transcript) = sess.state.lock().unwrap().zdr_transcript.as_mut() { + for item in &turn_output { + transcript.add_item(item.clone()); + } + } + turn_input = turn_output; } Err(e) => { @@ -626,19 +653,26 @@ async fn run_turn( sub_id: String, input: Vec, ) -> CodexResult> { + // Decide whether to use server-side storage (previous_response_id) or disable it let prev_id = { let state = sess.state.lock().unwrap(); - state.previous_response_id.clone() + if sess.disable_response_storage { + None + } else { + state.previous_response_id.clone() + } }; let instructions = match prev_id { Some(_) => None, None => sess.instructions.clone(), }; + // Build prompt payload, including store flag let prompt = Prompt { input, prev_id, instructions, + store: !sess.disable_response_storage, }; let mut retries = 0; diff --git a/codex-rs/core/src/codex_wrapper.rs b/codex-rs/core/src/codex_wrapper.rs index 426b5373c5..8d19683ffa 100644 --- a/codex-rs/core/src/codex_wrapper.rs +++ b/codex-rs/core/src/codex_wrapper.rs @@ -21,6 +21,7 @@ use tracing::debug; pub async fn init_codex( approval_policy: AskForApproval, sandbox_policy: SandboxPolicy, + disable_response_storage: bool, model_override: Option, ) -> anyhow::Result<(CodexWrapper, Event, Arc)> { let ctrl_c = notify_on_sigint(); @@ -33,6 +34,7 @@ pub async fn init_codex( instructions: config.instructions, approval_policy, sandbox_policy, + disable_response_storage, }) .await?; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 7d3309152c..d517e68824 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -19,6 +19,7 @@ mod models; pub mod protocol; mod safety; pub mod util; +mod zdr_transcript; pub use codex::Codex; diff --git a/codex-rs/core/src/protocol.rs b/codex-rs/core/src/protocol.rs index 42c8478e6b..96c4ea4832 100644 --- a/codex-rs/core/src/protocol.rs +++ b/codex-rs/core/src/protocol.rs @@ -33,6 +33,9 @@ pub enum Op { approval_policy: AskForApproval, /// How to sandbox commands executed in the system sandbox_policy: SandboxPolicy, + /// Disable server-side response storage (send full context each request) + #[serde(default)] + disable_response_storage: bool, }, /// Abort current task. diff --git a/codex-rs/core/src/zdr_transcript.rs b/codex-rs/core/src/zdr_transcript.rs new file mode 100644 index 0000000000..43c2d2620f --- /dev/null +++ b/codex-rs/core/src/zdr_transcript.rs @@ -0,0 +1,26 @@ +use crate::models::ResponseInputItem; + +/// Transcript that needs to be maintained for ZDR clients for which +/// previous_response_id is not available, so we must include the transcript +/// with every API call. +#[derive(Debug, Clone)] +pub(crate) struct ZdrTranscript { + items: Vec, +} + +impl ZdrTranscript { + pub(crate) fn new() -> Self { + Self { items: Vec::new() } + } + + pub(crate) fn add_item(&mut self, item: ResponseInputItem) { + if is_api_message(&item) { + // Note agent-loop.ts also does filtering on some of the fields. + self.items.push(item); + } + } +} + +fn is_api_message(message: &ResponseInputItem) -> bool { + !matches!(message, ResponseInputItem::Message { role, .. } if role.as_str() == "system") +} diff --git a/codex-rs/core/tests/live_agent.rs b/codex-rs/core/tests/live_agent.rs index 6562654c23..823cd73a01 100644 --- a/codex-rs/core/tests/live_agent.rs +++ b/codex-rs/core/tests/live_agent.rs @@ -55,6 +55,7 @@ async fn spawn_codex() -> Codex { instructions: None, approval_policy: AskForApproval::OnFailure, sandbox_policy: SandboxPolicy::NetworkAndFileWriteRestricted, + disable_response_storage: false, }, }) .await diff --git a/codex-rs/core/tests/previous_response_id.rs b/codex-rs/core/tests/previous_response_id.rs index 56fa9a6c0b..de1309e856 100644 --- a/codex-rs/core/tests/previous_response_id.rs +++ b/codex-rs/core/tests/previous_response_id.rs @@ -95,6 +95,7 @@ async fn keeps_previous_response_id_between_tasks() { instructions: None, approval_policy: AskForApproval::OnFailure, sandbox_policy: SandboxPolicy::NetworkAndFileWriteRestricted, + disable_response_storage: false, }, }) .await diff --git a/codex-rs/core/tests/stream_no_completed.rs b/codex-rs/core/tests/stream_no_completed.rs index da0cfb276b..c732a5fdbb 100644 --- a/codex-rs/core/tests/stream_no_completed.rs +++ b/codex-rs/core/tests/stream_no_completed.rs @@ -78,6 +78,7 @@ async fn retries_on_early_close() { instructions: None, approval_policy: AskForApproval::OnFailure, sandbox_policy: SandboxPolicy::NetworkAndFileWriteRestricted, + disable_response_storage: false, }, }) .await diff --git a/codex-rs/exec/src/cli.rs b/codex-rs/exec/src/cli.rs index a934aba003..938de29d23 100644 --- a/codex-rs/exec/src/cli.rs +++ b/codex-rs/exec/src/cli.rs @@ -1,4 +1,4 @@ -use clap::Parser; +use clap::{Parser, ArgAction}; use std::path::PathBuf; #[derive(Parser, Debug)] @@ -16,6 +16,10 @@ pub struct Cli { #[arg(long = "skip-git-repo-check", default_value_t = false)] pub skip_git_repo_check: bool, + /// Disable server-side response storage (omits previous_response_id and controls store flag) + #[arg(long = "disable-response-storage", action = ArgAction::SetTrue, default_value_t = false)] + pub disable_response_storage: bool, + /// Initial instructions for the agent. pub prompt: Option, } diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index c22b6bd694..e9a8518d7c 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -32,6 +32,7 @@ pub async fn run_main(cli: Cli) -> anyhow::Result<()> { let Cli { skip_git_repo_check, + disable_response_storage, model, images, prompt, @@ -51,7 +52,7 @@ pub async fn run_main(cli: Cli) -> anyhow::Result<()> { let approval_policy = AskForApproval::Never; let sandbox_policy = SandboxPolicy::NetworkAndFileWriteRestricted; let (codex_wrapper, event, ctrl_c) = - codex_wrapper::init_codex(approval_policy, sandbox_policy, model).await?; + 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:?}"); diff --git a/codex-rs/repl/src/lib.rs b/codex-rs/repl/src/lib.rs index 2266718ed9..7f3cd4a414 100644 --- a/codex-rs/repl/src/lib.rs +++ b/codex-rs/repl/src/lib.rs @@ -97,6 +97,8 @@ async fn codex_main(mut cli: Cli, cfg: Config, ctrl_c: Arc) -> anyhow::R instructions: cfg.instructions, approval_policy: cli.approval_policy.into(), sandbox_policy: cli.sandbox_policy.into(), + // by default, use server-side storage + disable_response_storage: false, }, }; diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 149cea42c4..d8228a131a 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -63,8 +63,9 @@ impl ChatWidget<'_> { let app_event_tx_clone = app_event_tx.clone(); // Create the Codex asynchronously so the UI loads as quickly as possible. tokio::spawn(async move { + // Initialize session; storage enabled by default let (codex, session_event, _ctrl_c) = - match init_codex(approval_policy, sandbox_policy, model).await { + match init_codex(approval_policy, sandbox_policy, false, model).await { Ok(vals) => vals, Err(e) => { // TODO(mbolin): This error needs to be surfaced to the user.