Allow background persistence for steered user input (#45506)

## Why

Persisting user input received during an active turn currently blocks the next model request. Stores that support background persistence can overlap this checkpoint with inference.

## What changed

- Add `PersistContext::SteeredUserInput` and `allows_background_persistence()` so stores may enqueue these checkpoints, with durability and error reporting enforced by later flush or shutdown operations.
- Use the new context for accepted steered user input and apply the same metadata handling as turn-start persistence.
- Keep tool outputs synchronous, including in mixed input batches, and allow stores to retain synchronous persistence for all contexts.

## Testing

Add gated-store integration tests covering background user-input persistence, synchronous stores, and synchronous tool-output checkpoints. Verify that the next request includes the steered input and waits for persistence when required.

GitOrigin-RevId: c60b7b6c9b483245fd3169306bcf0de248ccdf35

codex-rs/core/src/session/turn.rs

@@ -438,7 +438,7 @@ pub(crate) async fn run_turn(
             &turn_context,
             &turn_context.capture_current_model_info(),
             &pending_input,
-            PersistContext::Standard,
+            PersistContext::SteeredUserInput,
         )
         .await
         {
@@ -834,13 +834,21 @@ pub(crate) async fn run_hooks_and_record_inputs(
             if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) {
                 accepted_user_input = true;
             }
+            // Tool outputs retain their durability barrier, including in mixed input batches.
+            let input_persist_context = if persist_context == PersistContext::SteeredUserInput
+                && matches!(input_item, TurnInput::FunctionCallOutput(_))
+            {
+                PersistContext::Standard
+            } else {
+                persist_context
+            };
             record_pending_input(
                 sess,
                 turn_context,
                 model_info,
                 input_item.clone(),
                 hook_outcome.additional_contexts,
-                persist_context,
+                input_persist_context,
             )
             .await;
         }

codex-rs/core/tests/common/test_codex.rs

@@ -341,9 +341,15 @@ pub struct TestCodexBuilder {
     code_mode_host_program: Option<PathBuf>,
     history_mode: Option<ThreadHistoryMode>,
     models_manager: Option<SharedModelsManager>,
+    thread_store: Option<Arc<dyn ThreadStore>>,
 }
 
 impl TestCodexBuilder {
+    pub fn with_thread_store(mut self, thread_store: Arc<dyn ThreadStore>) -> Self {
+        self.thread_store = Some(thread_store);
+        self
+    }
+
     pub fn with_config<T>(mut self, mutator: T) -> Self
     where
         T: FnOnce(&mut Config) + Send + 'static,
@@ -694,7 +700,10 @@ impl TestCodexBuilder {
     ) -> anyhow::Result<TestCodex> {
         let auth = self.auth.clone();
         let state_db = codex_core::init_state_db(&config).await;
-        let thread_store = thread_store_from_config(&config, state_db.clone());
+        let thread_store = self
+            .thread_store
+            .clone()
+            .unwrap_or_else(|| thread_store_from_config(&config, state_db.clone()));
         let installation_id = resolve_installation_id(&config.codex_home).await?;
         let user_instructions_provider =
             self.user_instructions_provider.clone().unwrap_or_else(|| {
@@ -1387,6 +1396,7 @@ pub fn test_codex() -> TestCodexBuilder {
         code_mode_host_program: None,
         history_mode: None,
         models_manager: None,
+        thread_store: None,
     }
 }
 

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

@@ -131,6 +131,7 @@ mod openai_file_mcp;
 mod otel;
 mod override_updates;
 mod pending_input;
+mod pending_input_persistence;
 mod permissions_messages;
 mod personality;
 mod plugins;

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

@@ -0,0 +1,252 @@
+//! Exercises the checkpoint between a steered input and its next model request.
+
+use std::any::Any;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use codex_core::TurnInput;
+use codex_core::TurnInputRequest;
+use codex_core::TurnInputSubmission;
+use codex_protocol::ThreadId;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::ThreadHistoryMode;
+use codex_protocol::user_input::UserInput;
+use codex_thread_store::AppendThreadItemsParams;
+use codex_thread_store::ArchiveThreadParams;
+use codex_thread_store::CreateThreadParams;
+use codex_thread_store::DeleteThreadParams;
+use codex_thread_store::InMemoryThreadStore;
+use codex_thread_store::ListThreadsParams;
+use codex_thread_store::LoadThreadHistoryParams;
+use codex_thread_store::PersistContext;
+use codex_thread_store::ReadThreadByRolloutPathParams;
+use codex_thread_store::ReadThreadParams;
+use codex_thread_store::ResumeThreadParams;
+use codex_thread_store::StoredThread;
+use codex_thread_store::StoredThreadHistory;
+use codex_thread_store::ThreadPage;
+use codex_thread_store::ThreadStore;
+use codex_thread_store::ThreadStoreFuture;
+use codex_thread_store::UpdateThreadMetadataParams;
+use core_test_support::responses;
+use core_test_support::streaming_sse::StreamingSseChunk;
+use core_test_support::streaming_sse::start_streaming_sse_server;
+use core_test_support::test_codex::test_codex;
+use core_test_support::wait_for_event;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+use test_case::test_case;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio::time::timeout;
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum CheckpointPolicy {
+    Background,
+    Synchronous,
+}
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum InputKind {
+    User,
+    ToolOutput,
+}
+
+#[derive(Debug)]
+struct PendingCheckpoint {
+    context: PersistContext,
+    complete: oneshot::Sender<()>,
+}
+
+/// Records one checkpoint and gates it only when its persistence must be synchronous.
+struct GatedCheckpointStore {
+    inner: InMemoryThreadStore,
+    policy: CheckpointPolicy,
+    armed: AtomicBool,
+    checkpoints: mpsc::UnboundedSender<PendingCheckpoint>,
+}
+
+macro_rules! delegate_store_methods {
+    ($(fn $name:ident($param:ident: $params:ty) -> $result:ty;)*) => {
+        $(fn $name(&self, $param: $params) -> ThreadStoreFuture<'_, $result> {
+            ThreadStore::$name(&self.inner, $param)
+        })*
+    };
+}
+
+impl ThreadStore for GatedCheckpointStore {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    delegate_store_methods! {
+        fn create_thread(params: CreateThreadParams) -> ();
+        fn resume_thread(params: ResumeThreadParams) -> ();
+        fn append_items(params: AppendThreadItemsParams) -> ();
+        fn discard_thread(thread_id: ThreadId) -> ();
+        fn load_history(params: LoadThreadHistoryParams) -> StoredThreadHistory;
+        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) -> ();
+        fn flush_thread(thread_id: ThreadId) -> ();
+        fn shutdown_thread(thread_id: ThreadId) -> ();
+    }
+
+    fn persist_thread(
+        &self,
+        thread_id: ThreadId,
+        context: PersistContext,
+    ) -> ThreadStoreFuture<'_, ()> {
+        Box::pin(async move {
+            if self.armed.swap(false, Ordering::SeqCst) {
+                let (complete, completed) = oneshot::channel();
+                self.checkpoints
+                    .send(PendingCheckpoint { context, complete })
+                    .expect("checkpoint receiver should stay alive");
+                if self.policy == CheckpointPolicy::Synchronous
+                    || !context.allows_background_persistence()
+                {
+                    completed.await.expect("test should complete checkpoint");
+                }
+            }
+            self.inner.persist_thread(thread_id, context).await
+        })
+    }
+}
+
+#[test_case(CheckpointPolicy::Background, InputKind::User; "background_user_input")]
+#[test_case(CheckpointPolicy::Synchronous, InputKind::User; "synchronous_store")]
+#[test_case(CheckpointPolicy::Background, InputKind::ToolOutput; "tool_output_stays_synchronous")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn steered_input_checkpoint_controls_next_request(
+    policy: CheckpointPolicy,
+    input_kind: InputKind,
+) -> anyhow::Result<()> {
+    let (first_completed, first_completion) = oneshot::channel();
+    let (second_completed, second_completion) = oneshot::channel();
+    let (server, _completions) = start_streaming_sse_server(vec![
+        vec![
+            StreamingSseChunk {
+                gate: None,
+                body: responses::sse(vec![
+                    responses::ev_response_created("first"),
+                    responses::ev_message_item_added("first-message", ""),
+                    responses::ev_output_text_delta("original answer"),
+                ]),
+            },
+            StreamingSseChunk {
+                gate: Some(first_completion),
+                body: responses::sse(vec![
+                    responses::ev_assistant_message("first-message", "original answer"),
+                    responses::ev_completed("first"),
+                ]),
+            },
+        ],
+        vec![StreamingSseChunk {
+            gate: Some(second_completion),
+            body: responses::sse(vec![
+                responses::ev_response_created("second"),
+                responses::ev_completed("second"),
+            ]),
+        }],
+    ])
+    .await;
+    let (checkpoints, mut checkpoint_requests) = mpsc::unbounded_channel();
+    let store = Arc::new(GatedCheckpointStore {
+        inner: InMemoryThreadStore::default(),
+        policy,
+        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?;
+    let first = test
+        .codex
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "first prompt".to_string(),
+            text_elements: Vec::new(),
+        }]))
+        .await?;
+    let TurnInputSubmission::Started { turn_id } = first else {
+        panic!("first input should start a turn");
+    };
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::AgentMessageContentDelta(_))
+    })
+    .await;
+    store.armed.store(true, Ordering::SeqCst);
+    let input = match input_kind {
+        InputKind::User => TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "steered input".to_string(),
+            text_elements: Vec::new(),
+        }]),
+        InputKind::ToolOutput => {
+            TurnInputRequest::new(TurnInput::ResponseItem(serde_json::from_value(json!({
+                "type": "function_call_output",
+                "name": "send_message_to_thread",
+                "namespace": "codex_app",
+                "output": "steered input",
+            }))?))
+        }
+    };
+    assert_eq!(
+        test.codex.start_or_steer_turn(input).await?,
+        TurnInputSubmission::Steered { turn_id }
+    );
+    // Steering still waits for the existing inference stream to finish.
+    assert!(checkpoint_requests.try_recv().is_err());
+    first_completed.send(()).expect("finish original inference");
+    let checkpoint = timeout(Duration::from_secs(10), checkpoint_requests.recv())
+        .await?
+        .expect("Core should checkpoint the accepted input");
+    assert_eq!(
+        checkpoint.context,
+        match input_kind {
+            InputKind::User => PersistContext::SteeredUserInput,
+            InputKind::ToolOutput => PersistContext::Standard,
+        }
+    );
+    let should_overlap = policy == CheckpointPolicy::Background && input_kind == InputKind::User;
+    if !should_overlap {
+        assert!(
+            timeout(
+                Duration::from_millis(50),
+                server.wait_for_request_count(/*count*/ 2)
+            )
+            .await
+            .is_err()
+        );
+        checkpoint.complete.send(()).expect("complete checkpoint");
+    }
+    timeout(
+        Duration::from_secs(10),
+        server.wait_for_request_count(/*count*/ 2),
+    )
+    .await?;
+    let requests = server.requests().await;
+    assert_eq!(requests.len(), 2);
+    assert!(!String::from_utf8_lossy(&requests[0]).contains("steered input"));
+    assert!(String::from_utf8_lossy(&requests[1]).contains("steered input"));
+    second_completed
+        .send(())
+        .expect("finish follow-up inference");
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnComplete(_))
+    })
+    .await;
+    test.codex.shutdown_and_wait().await?;
+    server.shutdown().await;
+    Ok(())
+}

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

@@ -290,7 +290,7 @@ impl LiveThread {
     }
 
     pub async fn persist(&self, context: PersistContext) -> ThreadStoreResult<()> {
-        if context == PersistContext::TurnStart {
+        if context.allows_background_persistence() {
             self.flush_pending_metadata_update_for_existing_history()
                 .await?;
         }

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

@@ -68,6 +68,20 @@ pub enum PersistContext {
     Standard,
     /// A turn is about to begin sampling after its input has been recorded.
     TurnStart,
+    /// Accepted user input is being recorded before an active turn's next sampling request.
+    /// This does not apply to tool outputs, cancellation, or task cleanup.
+    SteeredUserInput,
+}
+
+impl PersistContext {
+    /// Whether a store may enqueue this checkpoint before returning and fence it at a later
+    /// durability barrier. Stores may still choose to persist synchronously.
+    pub fn allows_background_persistence(self) -> bool {
+        match self {
+            Self::Standard => false,
+            Self::TurnStart | Self::SteeredUserInput => true,
+        }
+    }
 }
 
 /// Storage-neutral thread persistence boundary.
@@ -122,9 +136,10 @@ pub trait ThreadStore: Any + Send + Sync {
 
     /// Materializes the thread if persistence is lazy, then persists all queued items.
     ///
-    /// Standard persistence must complete before returning. Turn-start persistence may complete
-    /// in the background when the implementation enqueues it before returning, fences it with
-    /// subsequent flush or shutdown operations, and surfaces failures through those operations.
+    /// 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.
     fn persist_thread(
         &self,
         thread_id: ThreadId,