feat: add ZDR support to Rust implementation

codex-rs/core/src/client.rs

@@ -29,15 +29,20 @@ use crate::flags::OPENAI_API_BASE;
 use crate::flags::OPENAI_REQUEST_MAX_RETRIES;
 use crate::flags::OPENAI_STREAM_IDLE_TIMEOUT_MS;
 use crate::flags::OPENAI_TIMEOUT_MS;
-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 {
-    pub input: Vec<ResponseInputItem>,
+    /// Conversation context input items.
+    pub input: Vec<ResponseItem>,
+    /// Optional previous response ID (when storage is enabled).
     pub prev_id: Option<String>,
+    /// Optional initial instructions (only sent on first turn).
     pub instructions: Option<String>,
+    /// Whether to store response on server side (disable_response_storage = !store).
+    pub store: bool,
 }
 
 #[derive(Debug)]
@@ -51,13 +56,15 @@ struct Payload<'a> {
     model: &'a str,
     #[serde(skip_serializing_if = "Option::is_none")]
     instructions: Option<&'a String>,
-    input: &'a Vec<ResponseInputItem>,
+    // TODO(mbolin): ResponseItem::Other should not be serialized.
+    input: &'a Vec<ResponseItem>,
     tools: &'a [Tool],
     tool_choice: &'static str,
     parallel_tool_calls: bool,
     reasoning: Option<Reasoning>,
     #[serde(skip_serializing_if = "Option::is_none")]
     previous_response_id: Option<String>,
+    store: bool,
     stream: bool,
 }
 
@@ -152,6 +159,7 @@ impl ModelClient {
                 generate_summary: None,
             }),
             previous_response_id: prompt.prev_id.clone(),
+            store: prompt.store,
             stream: true,
         };
 

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.
@@ -214,6 +215,7 @@ struct State {
     previous_response_id: Option<String>,
     pending_approvals: HashMap<String, oneshot::Sender<ReviewDecision>>,
     pending_input: Vec<ResponseInputItem>,
+    zdr_transcript: Option<ZdrTranscript>,
 }
 
 impl Session {
@@ -399,6 +401,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 +492,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 +504,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
@@ -587,10 +598,23 @@ async fn run_task(sess: Arc<Session>, sub_id: String, input: Vec<InputItem>) {
         return;
     }
 
-    let mut turn_input = vec![ResponseInputItem::from(input)];
+    let mut pending_response_input: Vec<ResponseInputItem> = vec![ResponseInputItem::from(input)];
     loop {
-        let pending_input = sess.get_pending_input();
-        turn_input.splice(0..0, pending_input);
+        let mut turn_input: Vec<ResponseItem> =
+            if let Some(transcript) = &sess.state.lock().unwrap().zdr_transcript {
+                // If we are using ZDR, we need to send the transcript with every turn.
+                transcript.contents()
+            } else {
+                Vec::new()
+            };
+
+        turn_input.extend(pending_response_input.drain(..).map(ResponseItem::from));
+
+        // Note that pending_input would be something like a message the user
+        // submitted through the UI while the model was running. Though the UI
+        // may support this, the model might not.
+        let pending_input = sess.get_pending_input().into_iter().map(ResponseItem::from);
+        turn_input.extend(pending_input);
 
         match run_turn(&sess, sub_id.clone(), turn_input).await {
             Ok(turn_output) => {
@@ -598,7 +622,17 @@ async fn run_task(sess: Arc<Session>, sub_id: String, input: Vec<InputItem>) {
                     debug!("Turn completed");
                     break;
                 }
-                turn_input = turn_output;
+
+                if let Some(transcript) = sess.state.lock().unwrap().zdr_transcript.as_mut() {
+                    let num_added = transcript.record_items(turn_output.iter().map(|i| &i.item));
+                    if num_added == 0 {
+                        debug!("Turn completed");
+                        break;
+                    }
+                }
+
+                pending_response_input =
+                    turn_output.into_iter().filter_map(|i| i.response).collect();
             }
             Err(e) => {
                 info!("Turn error: {e:#}");
@@ -624,21 +658,27 @@ async fn run_task(sess: Arc<Session>, sub_id: String, input: Vec<InputItem>) {
 async fn run_turn(
     sess: &Session,
     sub_id: String,
-    input: Vec<ResponseInputItem>,
-) -> CodexResult<Vec<ResponseInputItem>> {
-    let prev_id = {
+    input: Vec<ResponseItem>,
+) -> CodexResult<Vec<ProcessedResponseItem>> {
+    // Decide whether to use server-side storage (previous_response_id) or disable it
+    let (prev_id, store) = {
         let state = sess.state.lock().unwrap();
-        state.previous_response_id.clone()
+        (
+            state.previous_response_id.clone(),
+            state.zdr_transcript.is_none(),
+        )
     };
 
     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,
     };
 
     let mut retries = 0;
@@ -676,11 +716,20 @@ async fn run_turn(
     }
 }
 
+/// When the model is prompted, it returns a stream of events. Some of these
+/// events map to a `ResponseItem`. A `ResponseItem` may need to be
+/// "handled" such that it produces a `ResponseInputItem` that needs to be
+/// sent back to the model on the next turn.
+struct ProcessedResponseItem {
+    item: ResponseItem,
+    response: Option<ResponseInputItem>,
+}
+
 async fn try_run_turn(
     sess: &Session,
     sub_id: &str,
     prompt: &Prompt,
-) -> CodexResult<Vec<ResponseInputItem>> {
+) -> CodexResult<Vec<ProcessedResponseItem>> {
     let mut stream = sess.client.clone().stream(prompt).await?;
 
     // Buffer all the incoming messages from the stream first, then execute them.
@@ -694,9 +743,8 @@ async fn try_run_turn(
     for event in input {
         match event {
             ResponseEvent::OutputItemDone(item) => {
-                if let Some(item) = handle_response_item(sess, sub_id, item).await? {
-                    output.push(item);
-                }
+                let response = handle_response_item(sess, sub_id, item.clone()).await?;
+                output.push(ProcessedResponseItem { item, response });
             }
             ResponseEvent::Completed { response_id } => {
                 let mut state = sess.state.lock().unwrap();

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<String>,
 ) -> anyhow::Result<(CodexWrapper, Event, Arc<Notify>)> {
     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?;
 

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;
 

codex-rs/core/src/models.rs

@@ -56,6 +56,17 @@ pub enum ResponseItem {
     Other,
 }
 
+impl From<ResponseInputItem> for ResponseItem {
+    fn from(item: ResponseInputItem) -> Self {
+        match item {
+            ResponseInputItem::Message { role, content } => Self::Message { role, content },
+            ResponseInputItem::FunctionCallOutput { call_id, output } => {
+                Self::FunctionCallOutput { call_id, output }
+            }
+        }
+    }
+}
+
 impl From<Vec<InputItem>> for ResponseInputItem {
     fn from(items: Vec<InputItem>) -> Self {
         Self::Message {

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.

codex-rs/core/src/zdr_transcript.rs

@@ -0,0 +1,51 @@
+use crate::models::ResponseItem;
+
+/// 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. This must include each `function_call` and its
+/// corresponding `function_call_output`.
+#[derive(Debug, Clone)]
+pub(crate) struct ZdrTranscript {
+    /// The oldest items are at the beginning of the vector.
+    items: Vec<ResponseItem>,
+}
+
+impl ZdrTranscript {
+    pub(crate) fn new() -> Self {
+        Self { items: Vec::new() }
+    }
+
+    /// Returns a clone of the contents in the transcript.
+    pub(crate) fn contents(&self) -> Vec<ResponseItem> {
+        self.items.clone()
+    }
+
+    /// `items` is ordered from oldest to newest.
+    pub(crate) fn record_items<'a, I>(&mut self, items: I) -> usize
+    where
+        I: IntoIterator<Item = &'a ResponseItem>,
+    {
+        let mut count = 0;
+        for item in items {
+            if is_api_message(item) {
+                // Note agent-loop.ts also does filtering on some of the fields.
+                self.items.push(item.clone());
+                count += 1;
+            }
+        }
+        count
+    }
+}
+
+/// Anything that is not a system message or "reasoning" message is considered
+/// an API message.
+fn is_api_message(message: &ResponseItem) -> bool {
+    match message {
+        ResponseItem::Message { role, .. } => {
+            role.as_str() != "system" && role.as_str() != "assistant"
+        }
+        ResponseItem::FunctionCall { .. } => true,
+        ResponseItem::FunctionCallOutput { .. } => true,
+        _ => false,
+    }
+}

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

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

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

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<String>,
 }

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:?}");
 

codex-rs/repl/src/lib.rs

@@ -97,6 +97,8 @@ async fn codex_main(mut cli: Cli, cfg: Config, ctrl_c: Arc<Notify>) -> 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,
         },
     };
 

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.