Separate thread preparation from execution persistence (#47199)

## Why

Injecting context before the first turn currently forces a full flush. Stores need a preparation checkpoint that can leave disposable context in memory, while local storage must keep prepared context durable before the first user input.

## What changed

- Add `PersistContext::ThreadPreparation` and `CodexThread::checkpoint_preparation()`, and use the checkpoint in `inject_response_items()`.
- Allow stores to defer rollout-derived metadata writes through `record_thread_metadata()`, including during preparation and execution checkpoints. Keep standard metadata updates synchronous.
- Flush local storage at preparation checkpoints so injected context survives restart.

## Testing

Add tests showing that preparation and first sampling can proceed without waiting for durable metadata. Cover prepared context in first and follow-up requests, plus durability before first input and restoration after restart, in both legacy and paginated history modes.

GitOrigin-RevId: 74f6bb1a5f8d47192364b8bba6afd703a3b7d688

codex-rs/core/src/codex_thread.rs

@@ -300,6 +300,13 @@ impl CodexThread {
         self.session.emit_thread_idle_lifecycle_if_idle(cause).await;
     }
 
+    /// Checkpoint initialization without activating speculative persistence.
+    pub async fn checkpoint_preparation(&self) -> std::io::Result<()> {
+        self.session
+            .try_ensure_rollout_materialized(PersistContext::ThreadPreparation)
+            .await
+    }
+
     #[doc(hidden)]
     pub async fn ensure_rollout_materialized(&self) {
         self.session
@@ -692,7 +699,7 @@ impl CodexThread {
     /// Record raw Responses API items without starting a new turn.
     pub async fn inject_response_items(&self, items: Vec<ResponseItem>) -> CodexResult<()> {
         self.inject_response_items_for_turn(items).await?;
-        self.session.flush_rollout().await?;
+        self.checkpoint_preparation().await?;
         Ok(())
     }
 

codex-rs/core/tests/suite/pending_input_persistence.rs

@@ -1,4 +1,4 @@
-//! Exercises the checkpoint between a steered input and its next model request.
+//! Exercises preparation and execution checkpoints before model requests.
 
 use std::any::Any;
 use std::sync::Arc;
@@ -9,10 +9,16 @@ use std::time::Duration;
 use codex_core::TurnInput;
 use codex_core::TurnInputRequest;
 use codex_core::TurnInputSubmission;
+use codex_history::InitialHistory;
+use codex_history::ResumedHistory;
 use codex_protocol::ThreadId;
+use codex_protocol::mcp::ClientMcpExtensions;
+use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::ThreadHistoryMode;
 use codex_protocol::user_input::UserInput;
+use codex_rollout::RolloutItem;
+use codex_rollout::RolloutRecorder;
 use codex_thread_store::AppendThreadItemsParams;
 use codex_thread_store::ArchiveThreadParams;
 use codex_thread_store::CreateThreadParams;
@@ -44,6 +50,7 @@ use tokio::time::timeout;
 
 #[derive(Clone, Copy, PartialEq, Eq)]
 enum CheckpointPolicy {
+    Preparation,
     Background,
     Synchronous,
 }
@@ -90,7 +97,6 @@ impl ThreadStore for GatedCheckpointStore {
         fn read_thread(params: ReadThreadParams) -> StoredThread;
         fn read_thread_by_rollout_path(params: ReadThreadByRolloutPathParams) -> StoredThread;
         fn list_threads(params: ListThreadsParams) -> ThreadPage;
-        fn update_thread_metadata(params: UpdateThreadMetadataParams) -> Option<StoredThread>;
         fn archive_thread(params: ArchiveThreadParams) -> ();
         fn unarchive_thread(params: ArchiveThreadParams) -> StoredThread;
         fn delete_thread(params: DeleteThreadParams) -> ();
@@ -98,12 +104,36 @@ impl ThreadStore for GatedCheckpointStore {
         fn shutdown_thread(thread_id: ThreadId) -> ();
     }
 
+    fn update_thread_metadata(
+        &self,
+        params: UpdateThreadMetadataParams,
+    ) -> ThreadStoreFuture<'_, Option<StoredThread>> {
+        Box::pin(async move {
+            if self.policy == CheckpointPolicy::Preparation {
+                std::future::pending::<()>().await;
+            }
+            self.inner.update_thread_metadata(params).await
+        })
+    }
+
+    fn record_thread_metadata(
+        &self,
+        params: UpdateThreadMetadataParams,
+    ) -> ThreadStoreFuture<'_, ()> {
+        Box::pin(async move { self.inner.update_thread_metadata(params).await.map(|_| ()) })
+    }
+
     fn persist_thread(
         &self,
         thread_id: ThreadId,
         context: PersistContext,
     ) -> ThreadStoreFuture<'_, ()> {
         Box::pin(async move {
+            if self.policy == CheckpointPolicy::Preparation
+                && context == PersistContext::ThreadPreparation
+            {
+                return Ok(());
+            }
             if self.armed.swap(false, Ordering::SeqCst) {
                 let (complete, completed) = oneshot::channel();
                 self.checkpoints
@@ -250,3 +280,192 @@ async fn steered_input_checkpoint_controls_next_request(
     server.shutdown().await;
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata() -> anyhow::Result<()> {
+    let (complete, completion) = oneshot::channel();
+    let (server, _completions) = start_streaming_sse_server(vec![vec![StreamingSseChunk {
+        gate: Some(completion),
+        body: responses::sse(vec![
+            responses::ev_response_created("first"),
+            responses::ev_completed("first"),
+        ]),
+    }]])
+    .await;
+    let (checkpoints, mut checkpoint_requests) = mpsc::unbounded_channel();
+    let store = Arc::new(GatedCheckpointStore {
+        inner: InMemoryThreadStore::default(),
+        policy: CheckpointPolicy::Preparation,
+        armed: AtomicBool::new(false),
+        checkpoints,
+    });
+    let base_url = format!("{}/v1", server.uri());
+    let config_server = responses::start_mock_server().await;
+    let test = test_codex()
+        .with_thread_store(store.clone())
+        .with_history_mode(ThreadHistoryMode::Legacy)
+        .with_config(move |config| config.model_provider.base_url = Some(base_url))
+        .build_with_auto_env(&config_server)
+        .await?;
+    store.armed.store(true, Ordering::SeqCst);
+    timeout(Duration::from_secs(10), test.codex.inject_response_items(vec![serde_json::from_value(json!({
+        "type": "message", "role": "developer", "content": [{"type": "input_text", "text": "prepared context"}]
+    }))?])).await??;
+    assert!(checkpoint_requests.try_recv().is_err());
+    test.codex
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "first prompt".to_string(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    let checkpoint = timeout(Duration::from_secs(10), checkpoint_requests.recv())
+        .await?
+        .expect("first execution checkpoint");
+    assert_eq!(checkpoint.context, PersistContext::TurnStart);
+    timeout(
+        Duration::from_secs(10),
+        server.wait_for_request_count(/*count*/ 1),
+    )
+    .await?;
+    let requests = server.requests().await;
+    assert_eq!(requests.len(), 1);
+    let request: serde_json::Value = serde_json::from_slice(&requests[0])?;
+    let prepared_input = request["input"]
+        .as_array()
+        .expect("request input")
+        .iter()
+        .flat_map(|item| {
+            item["content"]
+                .as_array()
+                .into_iter()
+                .flatten()
+                .filter_map(move |content| {
+                    let text = content["text"].as_str()?;
+                    matches!(text, "prepared context" | "first prompt")
+                        .then(|| (item["role"].as_str().expect("message role"), text))
+                })
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(
+        prepared_input,
+        vec![("developer", "prepared context"), ("user", "first prompt")]
+    );
+    drop(checkpoint.complete);
+    complete.send(()).expect("complete inference");
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    Ok(())
+}
+
+#[test_case(ThreadHistoryMode::Legacy; "legacy")]
+#[test_case(ThreadHistoryMode::Paginated; "paginated")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn local_preparation_is_durable_before_first_input_and_survives_restart(
+    history_mode: ThreadHistoryMode,
+) -> anyhow::Result<()> {
+    const PREPARED_CONTEXT: &str = "prepared context persisted before first input";
+    const FIRST_PROMPT: &str = "first prompt after restart";
+
+    let server = responses::start_mock_server().await;
+    let mock = responses::mount_sse_once(
+        &server,
+        responses::sse(vec![
+            responses::ev_response_created("first-response"),
+            responses::ev_assistant_message("first-message", "prepared context restored"),
+            responses::ev_completed("first-response"),
+        ]),
+    )
+    .await;
+    let mut builder = test_codex().with_history_mode(history_mode);
+    let initial = builder.build_with_auto_env(&server).await?;
+    let prepared: ResponseItem = serde_json::from_value(json!({
+        "type": "message",
+        "id": "msg_preparation_before_input",
+        "role": "developer",
+        "content": [{"type": "input_text", "text": PREPARED_CONTEXT}],
+    }))?;
+    initial
+        .codex
+        .inject_response_items(vec![prepared.clone()])
+        .await?;
+
+    // Read disk before shutdown or a first turn can add another durability barrier.
+    let rollout_path = initial.codex.rollout_path().expect("local rollout path");
+    let (items, _, parse_errors) = RolloutRecorder::load_rollout_items(&rollout_path).await?;
+    let persisted_preparation = items
+        .into_iter()
+        .filter_map(|item| match item {
+            RolloutItem::ResponseItem(envelope) if envelope.item.id() == prepared.id() => {
+                // Core preserves this ID and adds runtime turn/creation-time metadata.
+                Some(responses::strip_metadata(envelope.item))
+            }
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+    assert_eq!((parse_errors, persisted_preparation), (0, vec![prepared]));
+    assert!(mock.requests().is_empty());
+
+    initial.codex.shutdown_and_wait().await?;
+    let thread_id = initial.session_configured.thread_id;
+    initial.thread_manager.remove_thread(&thread_id).await;
+    // Reopen persisted context by ID, which supports both local history modes.
+    let store = codex_core::thread_store_from_config(
+        &initial.config,
+        codex_core::init_state_db(&initial.config).await,
+    );
+    let context = store
+        .load_latest_model_context(LoadThreadHistoryParams {
+            thread_id,
+            include_archived: false,
+        })
+        .await?;
+    let resumed = initial
+        .thread_manager
+        .resume_thread_with_history(
+            initial.config.clone(),
+            InitialHistory::Resumed(ResumedHistory {
+                conversation_id: context.thread_id,
+                history: Arc::new(context.items),
+                rollout_path: Some(rollout_path),
+            }),
+            initial.thread_manager.auth_manager(),
+            /*parent_trace*/ None,
+            ClientMcpExtensions::default(),
+        )
+        .await?;
+    assert!(!Arc::ptr_eq(&initial.codex, &resumed.thread));
+    resumed
+        .thread
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: FIRST_PROMPT.to_string(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    wait_for_event(&resumed.thread, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    let input = mock.single_request().input();
+    let prepared_input = input
+        .iter()
+        .flat_map(|item| {
+            item["content"]
+                .as_array()
+                .into_iter()
+                .flatten()
+                .filter_map(move |content| {
+                    let text = content["text"].as_str()?;
+                    matches!(text, PREPARED_CONTEXT | FIRST_PROMPT)
+                        .then(|| (item["role"].as_str().expect("message role"), text))
+                })
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(
+        prepared_input,
+        vec![("developer", PREPARED_CONTEXT), ("user", FIRST_PROMPT)]
+    );
+    resumed.thread.shutdown_and_wait().await?;
+    Ok(())
+}

codex-rs/core/tests/suite/scenarios.rs

@@ -78,6 +78,9 @@ mod agent_message_board;
 #[path = "scenarios_guardian_extra_policy.rs"]
 mod guardian_extra_policy;
 
+#[path = "scenarios_preparation.rs"]
+mod preparation;
+
 #[path = "scenarios_shared_instructions.rs"]
 mod shared_instructions;
 

codex-rs/core/tests/suite/scenarios_preparation.rs

@@ -0,0 +1,67 @@
+//! Verifies prepared context reaches both the first model request and a follow-up.
+
+use super::*;
+use codex_protocol::protocol::ThreadHistoryMode;
+use pretty_assertions::assert_eq;
+use test_case::test_case;
+
+#[test_case(ThreadHistoryMode::Legacy; "legacy")]
+#[test_case(ThreadHistoryMode::Paginated; "paginated")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn prepared_context_reaches_first_request_and_survives_follow_up(
+    history_mode: ThreadHistoryMode,
+) -> Result<()> {
+    const PREPARED_CONTEXT: &str = "Prepared project context: the launch owner is Mira.";
+
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let test = test_codex()
+        .with_history_mode(history_mode)
+        .build_with_auto_env(&server)
+        .await?;
+    let mock = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("first-response"),
+                ev_assistant_message("first-message", "The launch owner is Mira."),
+                ev_completed("first-response"),
+            ]),
+            sse(vec![
+                ev_response_created("follow-up-response"),
+                ev_assistant_message("follow-up-message", "Mira will send the launch update."),
+                ev_completed("follow-up-response"),
+            ]),
+        ],
+    )
+    .await;
+
+    // A hosted caller can inject preparation before the user starts the first turn.
+    test.codex
+        .inject_response_items(vec![serde_json::from_value(json!({
+            "type": "message",
+            "role": "developer",
+            "content": [{
+                "type": "input_text",
+                "text": PREPARED_CONTEXT
+            }],
+        }))?])
+        .await?;
+    assert!(mock.requests().is_empty());
+
+    test.submit_turn("Who owns the launch?").await?;
+    test.submit_turn("Who will send the launch update?").await?;
+    let requests = mock.requests();
+    assert_eq!(requests.len(), 2);
+    for request in requests {
+        let prepared_context = request
+            .message_input_texts("developer")
+            .into_iter()
+            .filter(|text| text == PREPARED_CONTEXT)
+            .collect::<Vec<_>>();
+        assert_eq!(prepared_context, vec![PREPARED_CONTEXT]);
+    }
+    test.codex.shutdown_and_wait().await?;
+    Ok(())
+}

codex-rs/thread-store/src/live_thread.rs

@@ -247,7 +247,7 @@ impl LiveThread {
             .observe_appended_items(items.as_slice());
         if let Some(update) = update {
             self.thread_store
-                .update_thread_metadata(UpdateThreadMetadataParams {
+                .record_thread_metadata(UpdateThreadMetadataParams {
                     thread_id: self.thread_id,
                     patch: update.patch.clone(),
                     include_archived: true,
@@ -291,13 +291,18 @@ impl LiveThread {
 
     pub async fn persist(&self, context: PersistContext) -> ThreadStoreResult<()> {
         if context.allows_background_persistence() {
-            self.flush_pending_metadata_update_for_existing_history()
-                .await?;
+            let update = self
+                .metadata_sync
+                .lock()
+                .await
+                .take_pending_update_for_existing_history();
+            self.apply_pending_metadata_update(update, context).await?;
         }
         self.thread_store
             .persist_thread(self.thread_id, context)
             .await?;
-        self.flush_pending_metadata_update().await
+        let update = self.metadata_sync.lock().await.take_pending_update();
+        self.apply_pending_metadata_update(update, context).await
     }
 
     pub async fn flush(&self) -> ThreadStoreResult<()> {
@@ -417,7 +422,8 @@ impl LiveThread {
 
     async fn flush_pending_metadata_update(&self) -> ThreadStoreResult<()> {
         let update = self.metadata_sync.lock().await.take_pending_update();
-        self.apply_pending_metadata_update(update).await
+        self.apply_pending_metadata_update(update, PersistContext::Standard)
+            .await
     }
 
     async fn flush_pending_metadata_update_for_existing_history(&self) -> ThreadStoreResult<()> {
@@ -426,23 +432,28 @@ impl LiveThread {
             .lock()
             .await
             .take_pending_update_for_existing_history();
-        self.apply_pending_metadata_update(update).await
+        self.apply_pending_metadata_update(update, PersistContext::Standard)
+            .await
     }
 
     async fn apply_pending_metadata_update(
         &self,
         update: Option<crate::thread_metadata_sync::PendingThreadMetadataPatch>,
+        context: PersistContext,
     ) -> ThreadStoreResult<()> {
         let Some(update) = update else {
             return Ok(());
         };
-        self.thread_store
-            .update_thread_metadata(UpdateThreadMetadataParams {
-                thread_id: self.thread_id,
-                patch: update.patch.clone(),
-                include_archived: true,
-            })
-            .await?;
+        let params = UpdateThreadMetadataParams {
+            thread_id: self.thread_id,
+            patch: update.patch.clone(),
+            include_archived: true,
+        };
+        if context == PersistContext::Standard {
+            self.thread_store.update_thread_metadata(params).await?;
+        } else {
+            self.thread_store.record_thread_metadata(params).await?;
+        }
         self.metadata_sync
             .lock()
             .await

codex-rs/thread-store/src/local/mod.rs

@@ -527,9 +527,15 @@ impl ThreadStore for LocalThreadStore {
     fn persist_thread(
         &self,
         thread_id: ThreadId,
-        _context: PersistContext,
+        context: PersistContext,
     ) -> ThreadStoreFuture<'_, ()> {
-        Box::pin(async move { live_writer::persist_thread(self, thread_id).await })
+        Box::pin(async move {
+            if context == PersistContext::ThreadPreparation {
+                live_writer::flush_thread(self, thread_id).await
+            } else {
+                live_writer::persist_thread(self, thread_id).await
+            }
+        })
     }
 
     fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> {

codex-rs/thread-store/src/store.rs

@@ -64,6 +64,8 @@ pub type ThreadStoreFuture<'a, T> = Pin<Box<dyn Future<Output = ThreadStoreResul
 /// Why thread persistence is being requested.
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub enum PersistContext {
+    /// Thread initialization or context injection is complete without starting a turn.
+    ThreadPreparation,
     /// Standard persistence makes the thread and all queued items durable and readable.
     Standard,
     /// A turn is about to begin sampling after its input has been recorded.
@@ -78,7 +80,7 @@ impl PersistContext {
     /// durability barrier. Stores may still choose to persist synchronously.
     pub fn allows_background_persistence(self) -> bool {
         match self {
-            Self::Standard => false,
+            Self::ThreadPreparation | Self::Standard => false,
             Self::TurnStart | Self::SteeredUserInput => true,
         }
     }
@@ -142,9 +144,22 @@ pub trait ThreadStore: Any + Send + Sync {
     /// replay history and before updating any implementation-owned projections.
     fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreFuture<'_, ()>;
 
+    /// Records metadata derived from rollout items.
+    /// Unlike [`Self::update_thread_metadata`], this allows the write to be deferred.
+    ///
+    /// The default implementation awaits the update. Deferred writes must respect the
+    /// flush and shutdown guarantees documented by [`Self::persist_thread`].
+    fn record_thread_metadata(
+        &self,
+        params: UpdateThreadMetadataParams,
+    ) -> ThreadStoreFuture<'_, ()> {
+        Box::pin(async move { self.update_thread_metadata(params).await.map(|_| ()) })
+    }
+
     /// Materializes the thread if persistence is lazy, then persists all queued items.
     ///
-    /// Standard persistence must complete before returning. Contexts that allow background
+    /// Preparation checkpoints may leave disposable preparation data in memory without
+    /// activating storage. Standard persistence must complete before returning. Contexts that allow background
     /// persistence may complete asynchronously when the implementation enqueues the checkpoint
     /// before returning, fences it with subsequent flush or shutdown operations, and surfaces
     /// failures through those operations.
@@ -158,6 +173,7 @@ pub trait ThreadStore: Any + Send + Sync {
     fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()>;
 
     /// Flushes pending items and closes the live thread writer.
+    /// Stores with opt-in disposable preparations may discard an unactivated preparation.
     fn shutdown_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()>;
 
     /// Discards the live thread writer without forcing pending in-memory items to become durable.