From b9bfc0aff84de6af4c549f1ae544de569474f5f4 Mon Sep 17 00:00:00 2001 From: Bryan Ashley Date: Mon, 14 Sep 2026 19:24:17 +0000 Subject: [PATCH] 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 | 12 +- codex-rs/core/tests/common/test_codex.rs | 12 +- codex-rs/core/tests/suite/mod.rs | 1 + .../tests/suite/pending_input_persistence.rs | 252 ++++++++++++++++++ codex-rs/thread-store/src/live_thread.rs | 2 +- codex-rs/thread-store/src/store.rs | 21 +- 6 files changed, 293 insertions(+), 7 deletions(-) create mode 100644 codex-rs/core/tests/suite/pending_input_persistence.rs diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index e2ade90aaa..8ab44553be 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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; } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index f81000b450..3823c84730 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -341,9 +341,15 @@ pub struct TestCodexBuilder { code_mode_host_program: Option, history_mode: Option, models_manager: Option, + thread_store: Option>, } impl TestCodexBuilder { + pub fn with_thread_store(mut self, thread_store: Arc) -> Self { + self.thread_store = Some(thread_store); + self + } + pub fn with_config(mut self, mutator: T) -> Self where T: FnOnce(&mut Config) + Send + 'static, @@ -694,7 +700,10 @@ impl TestCodexBuilder { ) -> anyhow::Result { 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, } } diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 78a793a46e..67cfcb47eb 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/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; diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs new file mode 100644 index 0000000000..24515c4be1 --- /dev/null +++ b/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, +} + +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; + 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(()) +} diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index 6feabaebe6..5f75a30404 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/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?; } diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 267df60076..bf2ffb3169 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/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,