diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 88ced2da01..da0e02801f 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/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) -> CodexResult<()> { self.inject_response_items_for_turn(items).await?; - self.session.flush_rollout().await?; + self.checkpoint_preparation().await?; Ok(()) } diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs index 24515c4be1..84f71d94a4 100644 --- a/codex-rs/core/tests/suite/pending_input_persistence.rs +++ b/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; 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> { + 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::>(); + 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::>(); + 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::>(); + assert_eq!( + prepared_input, + vec![("developer", PREPARED_CONTEXT), ("user", FIRST_PROMPT)] + ); + resumed.thread.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index fda39a149e..fa86826cea 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/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; diff --git a/codex-rs/core/tests/suite/scenarios_preparation.rs b/codex-rs/core/tests/suite/scenarios_preparation.rs new file mode 100644 index 0000000000..d8d5998f24 --- /dev/null +++ b/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::>(); + assert_eq!(prepared_context, vec![PREPARED_CONTEXT]); + } + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index 5f75a30404..f6f57e6dd0 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/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, + 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 diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 2261ec8c00..66af0429ed 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/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<'_, ()> { diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 58c54db156..38d52ee72c 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -64,6 +64,8 @@ pub type ThreadStoreFuture<'a, T> = Pin 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.