From 722784e9366a0343a005f91fadf6d807c8a22c76 Mon Sep 17 00:00:00 2001 From: Bryan Ashley Date: Tue, 11 Aug 2026 03:25:25 +0000 Subject: [PATCH] Distinguish turn-start thread persistence (#37926) ## What changed - Add `PersistContext` to the thread-store persistence contract so stores can identify persistence requested immediately before model sampling. A `TurnStart` request may be enqueued in the background when later flush and shutdown operations fence it and report any failure; all other persistence remains synchronous. - Use the turn-start context only for the initial input of a regular turn, while admission acknowledgments, steered input, and other persistence paths keep the standard durability behavior. - Flush pending metadata before turn-start persistence and ensure shutdown is still attempted if that metadata update fails. ## Testing - Verify that developer instructions and user input are persisted before the model request completes for a default paginated-history thread. GitOrigin-RevId: 162f0ec796a61aebe66ca0b909fbcccec3047b85 --- .../session_importer.rs | 3 +- .../request_processors/thread_processor.rs | 7 ++- .../app-server/tests/suite/v2/thread_read.rs | 9 ++- codex-rs/core/src/agent/control/legacy.rs | 6 +- codex-rs/core/src/agent/control_tests.rs | 46 ++++++++++++--- codex-rs/core/src/codex_thread.rs | 5 +- codex-rs/core/src/guardian/review_session.rs | 10 +++- codex-rs/core/src/hook_runtime.rs | 5 ++ codex-rs/core/src/session/handlers.rs | 3 +- codex-rs/core/src/session/mod.rs | 24 +++++--- codex-rs/core/src/session/session.rs | 3 +- codex-rs/core/src/session/tests.rs | 24 ++++++-- codex-rs/core/src/session/turn.rs | 23 ++++++-- codex-rs/core/src/tasks/mod.rs | 9 ++- codex-rs/core/src/tasks/regular.rs | 3 +- codex-rs/core/src/tasks/review.rs | 5 +- codex-rs/core/src/tasks/user_shell.rs | 5 +- .../tests/suite/user_message_admission.rs | 59 +++++++++++++++++++ codex-rs/thread-store/src/in_memory.rs | 7 ++- codex-rs/thread-store/src/lib.rs | 1 + codex-rs/thread-store/src/live_thread.rs | 27 +++++++-- codex-rs/thread-store/src/local/mod.rs | 18 ++++-- .../thread_history_materialization_tests.rs | 43 +++++++------- codex-rs/thread-store/src/store.rs | 19 +++++- 24 files changed, 289 insertions(+), 75 deletions(-) diff --git a/codex-rs/app-server/src/external_agent_migration/session_importer.rs b/codex-rs/app-server/src/external_agent_migration/session_importer.rs index 4cd0d1f166..5f08aa502a 100644 --- a/codex-rs/app-server/src/external_agent_migration/session_importer.rs +++ b/codex-rs/app-server/src/external_agent_migration/session_importer.rs @@ -36,6 +36,7 @@ use codex_rollout::RolloutItem; use codex_rollout::is_persisted_rollout_item; use codex_thread_store::AppendThreadItemsParams; use codex_thread_store::CreateThreadParams; +use codex_thread_store::PersistContext; use codex_thread_store::ThreadMetadataPatch; use codex_thread_store::ThreadPersistenceMetadata; use codex_thread_store::ThreadStore; @@ -576,7 +577,7 @@ impl ExternalAgentSessionImporter { ) })?; self.thread_store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .map_err(|err| { SessionImportStepFailure::new( diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 7163358b2f..b1567c748c 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -16,6 +16,7 @@ use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; use codex_protocol::protocol::ThreadHistoryMode; +use codex_thread_store::PersistContext; pub(super) const THREAD_LIST_DEFAULT_LIMIT: usize = 25; pub(super) const THREAD_LIST_MAX_LIMIT: usize = 100; @@ -2526,7 +2527,7 @@ impl ThreadRequestProcessor { codex_app_server_protocol::ThreadHistoryMode::Paginated ) { self.thread_store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .map_err(|err| thread_read_history_load_error(thread_id, err))?; thread.turns = self @@ -3322,7 +3323,7 @@ impl ThreadRequestProcessor { if paginated_resume && let Err(error) = self .thread_store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .map_err(thread_store_resume_read_error) { @@ -3671,7 +3672,7 @@ impl ThreadRequestProcessor { } if paginated_resume && (include_turns || params.initial_turns_page.is_some()) { self.thread_store - .persist_thread(existing_thread_id) + .persist_thread(existing_thread_id, PersistContext::Standard) .await .map_err(thread_store_resume_read_error)?; } diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 0c888581f0..05e46dd2be 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -76,6 +76,7 @@ use codex_thread_store::CreateThreadParams; use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; +use codex_thread_store::PersistContext; use codex_thread_store::ThreadMetadataPatch; use codex_thread_store::ThreadPersistenceMetadata; use codex_thread_store::ThreadStore; @@ -487,7 +488,9 @@ async fn thread_search_occurrences_reads_paginated_projection() -> Result<()> { }, }) .await?; - store.persist_thread(thread_id).await?; + store + .persist_thread(thread_id, PersistContext::Standard) + .await?; store .append_items(AppendThreadItemsParams { thread_id, @@ -1535,7 +1538,9 @@ async fn paginated_history_lists_and_legacy_reads_use_projected_turns_and_items( }, }) .await?; - store.persist_thread(thread_id).await?; + store + .persist_thread(thread_id, PersistContext::Standard) + .await?; store .append_items(AppendThreadItemsParams { thread_id, diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs index f5f16f6851..72f757ab2b 100644 --- a/codex-rs/core/src/agent/control/legacy.rs +++ b/codex-rs/core/src/agent/control/legacy.rs @@ -1,5 +1,6 @@ use super::*; use codex_protocol::error::CodexErrorDetails; +use codex_thread_store::PersistContext; impl AgentControl { /// Submit a shutdown request for a live agent without marking it explicitly closed in @@ -7,7 +8,10 @@ impl AgentControl { pub(crate) async fn shutdown_live_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; let result = if let Ok(thread) = state.get_thread(agent_id).await { - thread.session.ensure_rollout_materialized().await; + thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; thread.session.flush_rollout().await?; let result = if matches!(thread.agent_status().await, AgentStatus::Shutdown) { Ok(String::new()) diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 0cbd809f0d..0f59b0d86f 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -54,6 +54,7 @@ use codex_thread_store::ArchiveThreadParams; use codex_thread_store::InMemoryThreadStore; use codex_thread_store::LocalThreadStore; use codex_thread_store::LocalThreadStoreConfig; +use codex_thread_store::PersistContext; use codex_thread_store::ThreadStore; use codex_utils_path_uri::PathUri; use core_test_support::responses::strip_response_item_ids; @@ -353,7 +354,10 @@ async fn persist_thread_for_tree_resume(thread: &Arc, message: &str thread .inject_user_message_without_turn(message.to_string()) .await; - thread.session.ensure_rollout_materialized().await; + thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; thread .session .flush_rollout() @@ -1419,7 +1423,10 @@ async fn spawn_agent_can_fork_parent_thread_history_with_sanitized_items() { parent_reference_context_item.clone(), )]) .await; - parent_thread.session.ensure_rollout_materialized().await; + parent_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; parent_thread .session .flush_rollout() @@ -1683,7 +1690,10 @@ async fn spawn_agent_fork_strips_parent_usage_hints_from_compacted_history() { RolloutItem::ResponseItem(spawn_agent_call(&parent_spawn_call_id)), ]) .await; - parent_thread.session.ensure_rollout_materialized().await; + parent_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; parent_thread .session .flush_rollout() @@ -1843,7 +1853,10 @@ async fn spawn_agent_full_fork_restores_instructions_after_compaction_discards_p RolloutItem::ResponseItem(spawn_agent_call(&parent_spawn_call_id)), ]) .await; - parent_thread.session.ensure_rollout_materialized().await; + parent_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; parent_thread .session .flush_rollout() @@ -1987,7 +2000,10 @@ async fn spawn_agent_full_fork_legacy_compaction_rebuilds_child_instructions_onc .session .persist_rollout_items(&rollout_items) .await; - parent_thread.session.ensure_rollout_materialized().await; + parent_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; parent_thread .session .flush_rollout() @@ -2180,7 +2196,10 @@ async fn spawn_agent_fork_last_n_turns_keeps_only_recent_turns() { spawn_turn_context.to_turn_context_item(), )]) .await; - parent_thread.session.ensure_rollout_materialized().await; + parent_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; parent_thread .session .flush_rollout() @@ -2303,7 +2322,10 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li &[spawn_agent_call(&parent_spawn_call_id)], ) .await; - parent_thread.session.ensure_rollout_materialized().await; + parent_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; parent_thread .session .flush_rollout() @@ -2429,7 +2451,10 @@ async fn spawn_agent_fork_last_n_turns_strips_parent_usage_hints() { ], ) .await; - parent_thread.session.ensure_rollout_materialized().await; + parent_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; parent_thread .session .flush_rollout() @@ -3220,7 +3245,10 @@ async fn resume_thread_subagent_restores_stored_metadata() { .get_thread(child_thread_id) .await .expect("child thread should exist"); - child_thread.session.ensure_rollout_materialized().await; + child_thread + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; child_thread .session .flush_rollout() diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 469afeed0f..7b57bd74e5 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -48,6 +48,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::user_input::UserInput; +use codex_thread_store::PersistContext; use codex_thread_store::StoredThread; use codex_thread_store::StoredThreadHistory; use codex_thread_store::ThreadMetadataPatch; @@ -284,7 +285,9 @@ impl CodexThread { #[doc(hidden)] pub async fn ensure_rollout_materialized(&self) { - self.session.ensure_rollout_materialized().await; + self.session + .ensure_rollout_materialized(PersistContext::Standard) + .await; } #[doc(hidden)] diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 7c6bff792d..bf28182c5a 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -51,6 +51,7 @@ use crate::session::turn_context::TurnContext; use codex_config::types::McpServerConfig; use codex_features::Feature; use codex_model_provider_info::ModelProviderInfo; +use codex_thread_store::PersistContext; use codex_utils_path_uri::PathUri; use super::ApprovalRequestReasons; @@ -380,7 +381,10 @@ impl GuardianReviewSessionManager { pub(crate) async fn trunk_rollout_path(&self) -> Option { let trunk = self.state.lock().await.trunk.clone()?; - trunk.session.ensure_rollout_materialized().await; + trunk + .session + .ensure_rollout_materialized(PersistContext::Standard) + .await; match trunk.session.current_rollout_path().await { Ok(path) => path, Err(err) => { @@ -966,7 +970,9 @@ async fn append_guardian_followup_reminder(review_session: &GuardianReviewSessio async fn load_rollout_items_for_fork( session: &Session, ) -> anyhow::Result>> { - session.try_ensure_rollout_materialized().await?; + session + .try_ensure_rollout_materialized(PersistContext::Standard) + .await?; session.flush_rollout().await?; let live_thread = session.live_thread_for_persistence("guardian review fork")?; let history = live_thread.load_history(/*include_archived*/ true).await?; diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index ba49e994e0..da4a8f2b83 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -38,6 +38,7 @@ use codex_protocol::protocol::HookStartedEvent; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::WarningEvent; +use codex_thread_store::PersistContext; use codex_thread_store::ReadThreadParams; use serde_json::Value; use tracing::instrument; @@ -582,6 +583,7 @@ pub(crate) async fn record_pending_input( turn_context: &Arc, pending_input: TurnInput, additional_contexts: Vec, + persist_context: PersistContext, ) -> Result<(), TryStartTurnIfIdleRejectionReason> { match pending_input { TurnInput::UserInput { content, client_id } => { @@ -596,10 +598,12 @@ pub(crate) async fn record_pending_input( .filter(|task| Arc::ptr_eq(&task.turn_context, turn_context)) .and_then(|task| task.input_persisted.take()); if awaiting_admission || input_persisted.is_some() { + // An admission acknowledgment always requires synchronous persistence. sess.record_user_prompt_and_emit_turn_item( turn_context.as_ref(), content.as_slice(), client_id.clone(), + PersistContext::Standard, ) .await; record_additional_contexts(sess, turn_context, additional_contexts).await; @@ -640,6 +644,7 @@ pub(crate) async fn record_pending_input( turn_context.as_ref(), content.as_slice(), client_id, + persist_context, ) .await; } diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 4f6b0f9c22..3a648bd848 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -49,6 +49,7 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::WarningEvent; use codex_protocol::request_permissions::RequestPermissionsResponse; use codex_protocol::request_user_input::RequestUserInputResponse; +use codex_thread_store::PersistContext; use crate::context_manager::is_user_turn_boundary; use codex_protocol::dynamic_tools::DynamicToolResponse; @@ -557,7 +558,7 @@ pub(super) async fn persist_thread_memory_mode_update( mode: ThreadMemoryMode, ) -> anyhow::Result<()> { let live_thread = sess.live_thread_for_persistence("update thread memory mode")?; - live_thread.persist().await?; + live_thread.persist(PersistContext::Standard).await?; live_thread.flush().await?; live_thread .update_memory_mode(mode, /*include_archived*/ false) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 402179accf..31e60e2d2b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -151,6 +151,7 @@ use codex_thread_store::CreateThreadParams; use codex_thread_store::LiveThread; use codex_thread_store::LiveThreadInitGuard; use codex_thread_store::LocalThreadStore; +use codex_thread_store::PersistContext; use codex_thread_store::ReadThreadParams; use codex_thread_store::ResumeThreadParams; use codex_thread_store::ThreadPersistenceMetadata; @@ -1223,15 +1224,21 @@ impl Session { } } - pub(crate) async fn try_ensure_rollout_materialized(&self) -> std::io::Result<()> { + pub(crate) async fn try_ensure_rollout_materialized( + &self, + context: PersistContext, + ) -> std::io::Result<()> { if let Some(live_thread) = self.live_thread() { - live_thread.persist().await.map_err(std::io::Error::other)?; + live_thread + .persist(context) + .await + .map_err(std::io::Error::other)?; } Ok(()) } - pub(crate) async fn ensure_rollout_materialized(&self) { - if let Err(e) = self.try_ensure_rollout_materialized().await { + pub(crate) async fn ensure_rollout_materialized(&self, context: PersistContext) { + if let Err(e) = self.try_ensure_rollout_materialized(context).await { warn!("failed to materialize thread persistence: {e}"); } } @@ -1454,7 +1461,8 @@ impl Session { self.persist_rollout_items(&rollout_items).await; // Forked threads should remain file-backed immediately after startup. - self.ensure_rollout_materialized().await; + self.ensure_rollout_materialized(PersistContext::Standard) + .await; // Flush after seeding history and any persisted rollout copy. if !is_subagent { @@ -3989,6 +3997,7 @@ impl Session { turn_context: &TurnContext, input: &[UserInput], client_id: Option, + persist_context: PersistContext, ) { // Persist the user message to history, but emit the turn item from `UserInput` so // UI-only `text_elements` are preserved. `ResponseItem::Message` does not carry @@ -4001,7 +4010,7 @@ impl Session { let turn_item = TurnItem::UserMessage(user_message_item); self.emit_turn_item_started(turn_context, &turn_item).await; self.emit_turn_item_completed(turn_context, turn_item).await; - self.ensure_rollout_materialized().await; + self.ensure_rollout_materialized(persist_context).await; } pub(crate) async fn notify_stream_error( @@ -4152,7 +4161,8 @@ impl Session { return None; } }; - self.ensure_rollout_materialized().await; + self.ensure_rollout_materialized(PersistContext::Standard) + .await; Some(rollout_path) } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 9a95ea47ed..e440be7336 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1431,7 +1431,8 @@ impl Session { } if matches!(&sess.fork_persistence, ForkPersistence::Referenced { .. }) { // Keep the source reserved until the child's history reference is durable. - sess.try_ensure_rollout_materialized().await?; + sess.try_ensure_rollout_materialized(PersistContext::Standard) + .await?; } { let mut state = sess.state.lock().await; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index fc305616f1..a3b9fb9023 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4512,7 +4512,9 @@ async fn open_thread_persistence(session: &mut Session) -> PathBuf { async fn attach_thread_persistence(session: &mut Session) -> PathBuf { let rollout_path = open_thread_persistence(session).await; - session.ensure_rollout_materialized().await; + session + .ensure_rollout_materialized(PersistContext::Standard) + .await; session .flush_rollout() .await @@ -9655,7 +9657,9 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline serde_json::to_value(Some(turn_context.to_turn_context_item())) .expect("serialize expected context item") ); - session.ensure_rollout_materialized().await; + session + .ensure_rollout_materialized(PersistContext::Standard) + .await; session.flush_rollout().await.expect("rollout should flush"); let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path) @@ -9698,7 +9702,9 @@ async fn record_context_updates_and_set_reference_context_item_persists_split_fi .record_context_updates_and_set_reference_context_item(&step_context) .await .expect("world state should build"); - session.ensure_rollout_materialized().await; + session + .ensure_rollout_materialized(PersistContext::Standard) + .await; session.flush_rollout().await.expect("rollout should flush"); let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path) @@ -9786,7 +9792,9 @@ async fn record_context_updates_and_set_reference_context_item_persists_full_rei .record_context_updates_and_set_reference_context_item(&step_context) .await .expect("world state should build"); - session.ensure_rollout_materialized().await; + session + .ensure_rollout_materialized(PersistContext::Standard) + .await; session.flush_rollout().await.expect("rollout should flush"); let InitialHistory::Resumed(resumed) = RolloutRecorder::get_rollout_history(&rollout_path) @@ -10007,7 +10015,13 @@ async fn failed_user_message_persistence_stops_turn_processing() { ]; assert!( - super::turn::run_hooks_and_record_inputs(&session, &turn_context, &input).await, + super::turn::run_hooks_and_record_inputs( + &session, + &turn_context, + &input, + PersistContext::Standard, + ) + .await, "a rollout flush failure should stop turn processing" ); assert!(matches!( diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 869d6ab312..ef9dad66f2 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -112,6 +112,7 @@ use codex_skills::app_id_from_path; use codex_skills::tool_kind_for_path; use codex_skills_extension::HostSkillPrompts; use codex_skills_extension::InjectedHostSkillPrompts; +use codex_thread_store::PersistContext; use codex_tools::DiscoverableTool; use codex_tools::ToolName; use codex_tools::filter_request_plugin_install_discoverable_tools_for_client; @@ -175,7 +176,8 @@ pub(crate) async fn run_turn( .await { if matches!(err.details(), CodexErrorDetails::TurnAborted) { - run_hooks_and_record_inputs(&sess, &turn_context, &input).await; + run_hooks_and_record_inputs(&sess, &turn_context, &input, PersistContext::Standard) + .await; return Err(err); } if matches!(err.details(), CodexErrorDetails::ToolCollision(_)) { @@ -196,7 +198,8 @@ pub(crate) async fn run_turn( { Ok(requirements) => requirements, Err(err) => { - run_hooks_and_record_inputs(&sess, &turn_context, &input).await; + run_hooks_and_record_inputs(&sess, &turn_context, &input, PersistContext::Standard) + .await; return Err(err.into()); } }; @@ -212,7 +215,8 @@ pub(crate) async fn run_turn( { Ok(step_context) => step_context, Err(err) if matches!(err.details(), CodexErrorDetails::TurnAborted) => { - run_hooks_and_record_inputs(&sess, &turn_context, &input).await; + run_hooks_and_record_inputs(&sess, &turn_context, &input, PersistContext::Standard) + .await; return Err(err); } Err(err) => return Err(err), @@ -240,7 +244,7 @@ pub(crate) async fn run_turn( return Ok(None); } let mut can_drain_pending_input = input.is_empty(); - if run_hooks_and_record_inputs(&sess, &turn_context, &input).await { + if run_hooks_and_record_inputs(&sess, &turn_context, &input, PersistContext::TurnStart).await { return Ok(None); } @@ -288,7 +292,14 @@ pub(crate) async fn run_turn( Vec::new() }; - if run_hooks_and_record_inputs(&sess, &turn_context, &pending_input).await { + if run_hooks_and_record_inputs( + &sess, + &turn_context, + &pending_input, + PersistContext::Standard, + ) + .await + { break; } @@ -588,6 +599,7 @@ pub(crate) async fn run_hooks_and_record_inputs( sess: &Arc, turn_context: &Arc, input: &[TurnInput], + persist_context: PersistContext, ) -> bool { let mut blocked_input = false; let mut accepted_user_input = false; @@ -607,6 +619,7 @@ pub(crate) async fn run_hooks_and_record_inputs( turn_context, input_item.clone(), hook_outcome.additional_contexts, + persist_context, ) .await .is_err() diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 3c5cae4c4f..d4aa3a7256 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -51,6 +51,7 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::WarningEvent; +use codex_thread_store::PersistContext; use codex_features::Feature; use codex_protocol::error::CodexErrorDetails; @@ -627,7 +628,13 @@ impl Session { ts.token_usage_at_turn_start.clone(), ) }; - run_hooks_and_record_inputs(self, &turn_context, &pending_input).await; + run_hooks_and_record_inputs( + self, + &turn_context, + &pending_input, + PersistContext::Standard, + ) + .await; task_ended_before_persistence |= self .pending_user_message_admissions .complete_task_end(&turn_context.sub_id); diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index b88bd51b03..c7079e4071 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -11,6 +11,7 @@ use crate::session_startup_prewarm::SessionStartupPrewarmResolution; use crate::state::TaskKind; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::TurnStartedEvent; +use codex_thread_store::PersistContext; use tracing::Instrument; use tracing::trace_span; @@ -62,7 +63,7 @@ impl SessionTask for RegularTask { .await; let prewarmed_client_session = match prewarmed_client_session { SessionStartupPrewarmResolution::Cancelled => { - run_hooks_and_record_inputs(&sess, &ctx, &input).await; + run_hooks_and_record_inputs(&sess, &ctx, &input, PersistContext::Standard).await; return Ok(None); } SessionStartupPrewarmResolution::Unavailable { .. } => None, diff --git a/codex-rs/core/src/tasks/review.rs b/codex-rs/core/src/tasks/review.rs index d244cd0c44..94859375dd 100644 --- a/codex-rs/core/src/tasks/review.rs +++ b/codex-rs/core/src/tasks/review.rs @@ -28,6 +28,7 @@ use crate::session::turn_context::TurnContext; use crate::state::TaskKind; use codex_features::Feature; use codex_protocol::user_input::UserInput; +use codex_thread_store::PersistContext; use super::SessionTask; use super::SessionTaskResult; @@ -269,5 +270,7 @@ pub(crate) async fn exit_review_mode( // Review turns can run before any regular user turn, so explicitly // materialize rollout persistence. Do this after emitting review output so // file creation + git metadata collection cannot delay client-facing items. - session.ensure_rollout_materialized().await; + session + .ensure_rollout_materialized(PersistContext::Standard) + .await; } diff --git a/codex-rs/core/src/tasks/user_shell.rs b/codex-rs/core/src/tasks/user_shell.rs index d3312c8b46..0a88656ac8 100644 --- a/codex-rs/core/src/tasks/user_shell.rs +++ b/codex-rs/core/src/tasks/user_shell.rs @@ -39,6 +39,7 @@ use codex_protocol::protocol::ExecCommandSource; use codex_protocol::protocol::TurnStartedEvent; use codex_sandboxing::SandboxType; use codex_shell_command::parse_command::parse_command; +use codex_thread_store::PersistContext; use super::SessionTask; use super::SessionTaskResult; @@ -463,7 +464,9 @@ async fn persist_user_shell_output( .await; // Standalone shell turns can run before any regular user turn, so // explicitly materialize rollout persistence after recording output. - session.ensure_rollout_materialized().await; + session + .ensure_rollout_materialized(PersistContext::Standard) + .await; return; } diff --git a/codex-rs/core/tests/suite/user_message_admission.rs b/codex-rs/core/tests/suite/user_message_admission.rs index 2d1a4f7fe2..b3f3c7d8b7 100644 --- a/codex-rs/core/tests/suite/user_message_admission.rs +++ b/codex-rs/core/tests/suite/user_message_admission.rs @@ -4,6 +4,8 @@ use codex_history::RolloutItem; use codex_history::RolloutLine; use codex_protocol::error::CodexErrorDetails; use codex_protocol::items::TurnItem; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; @@ -37,6 +39,63 @@ fn user_input(text: &str) -> Op { } } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn default_turn_start_persists_developer_and_user_input_before_model_request() { + let (release_response, response_gate) = oneshot::channel(); + let (server, _completions) = start_streaming_sse_server(vec![vec![StreamingSseChunk { + gate: Some(response_gate), + body: responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + }]]) + .await; + let test = test_codex() + .with_model("gpt-5.4") + .with_history_mode(ThreadHistoryMode::Paginated) + .build_with_streaming_server(&server) + .await + .expect("build default thread-store session"); + test.codex + .inject_response_items(vec![ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "turn-start developer instructions".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }]) + .await + .expect("inject developer instructions"); + test.codex + .submit(user_input("turn-start user input")) + .await + .expect("submit user input"); + + timeout( + Duration::from_secs(5), + server.wait_for_request_count(/*count*/ 1), + ) + .await + .expect("turn should reach the model request"); + let rollout_path = test.codex.rollout_path().expect("local rollout path"); + let rollout = tokio::fs::read_to_string(rollout_path) + .await + .expect("read rollout while the model response is blocked"); + assert!(rollout.contains("turn-start developer instructions")); + assert!(rollout.contains("turn-start user input")); + + release_response.send(()).expect("release model response"); + loop { + let event = timeout(Duration::from_secs(5), test.codex.next_event()) + .await + .expect("turn should finish") + .expect("event stream should remain available"); + if matches!(event.msg, EventMsg::TurnComplete(_)) { + break; + } + } + server.shutdown().await; +} + /// Concurrent submissions must start exactly one turn, steer the other message, and persist both client IDs. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn user_message_admission_reports_started_and_steered_for_concurrent_submissions() { diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index f2e18c7d51..bd4c678484 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -26,6 +26,7 @@ use crate::DeleteThreadParams; use crate::ListThreadsParams; use crate::LoadThreadHistoryParams; use crate::MoveThreadToSectionParams; +use crate::PersistContext; use crate::ReadThreadByRolloutPathParams; use crate::ReadThreadParams; use crate::ResumeThreadParams; @@ -816,7 +817,11 @@ impl ThreadStore for InMemoryThreadStore { Box::pin(InMemoryThreadStore::append_items(self, params)) } - fn persist_thread(&self, _thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + fn persist_thread( + &self, + _thread_id: ThreadId, + _context: PersistContext, + ) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { self.state.lock().await.calls.persist_thread += 1; Ok(()) diff --git a/codex-rs/thread-store/src/lib.rs b/codex-rs/thread-store/src/lib.rs index 812a1b1bdf..f7bf9baefb 100644 --- a/codex-rs/thread-store/src/lib.rs +++ b/codex-rs/thread-store/src/lib.rs @@ -32,6 +32,7 @@ pub use local::RolloutMigrationReport; pub use local::RolloutMigrationStatus; pub use queue_store::LocalQueueStore; pub use queue_store::QueueStore; +pub use store::PersistContext; pub use store::ThreadStore; pub use store::ThreadStoreFuture; pub use thread_sections::CreateThreadSectionParams; diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index 0052c7fd0b..dc9afcf9e7 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -15,6 +15,7 @@ use crate::AppendThreadItemsParams; use crate::CreateThreadParams; use crate::LoadThreadHistoryParams; use crate::LocalThreadStore; +use crate::PersistContext; use crate::ReadThreadParams; use crate::ResumeThreadParams; use crate::StoredThread; @@ -253,8 +254,14 @@ impl LiveThread { Ok(items) } - pub async fn persist(&self) -> ThreadStoreResult<()> { - self.thread_store.persist_thread(self.thread_id).await?; + pub async fn persist(&self, context: PersistContext) -> ThreadStoreResult<()> { + if context == PersistContext::TurnStart { + self.flush_pending_metadata_update_for_existing_history() + .await?; + } + self.thread_store + .persist_thread(self.thread_id, context) + .await?; self.flush_pending_metadata_update().await } @@ -265,9 +272,19 @@ impl LiveThread { } pub async fn shutdown(&self) -> ThreadStoreResult<()> { - self.flush_pending_metadata_update_for_existing_history() - .await?; - self.thread_store.shutdown_thread(self.thread_id).await + let metadata_result = self + .flush_pending_metadata_update_for_existing_history() + .await; + let shutdown_result = self.thread_store.shutdown_thread(self.thread_id).await; + match (metadata_result, shutdown_result) { + (Err(metadata_error), Err(shutdown_error)) => Err(ThreadStoreError::Internal { + message: format!( + "thread metadata update failed: {metadata_error}; thread shutdown failed: {shutdown_error}" + ), + }), + (Err(metadata_error), Ok(())) => Err(metadata_error), + (Ok(()), result) => result, + } } pub async fn discard(&self) -> ThreadStoreResult<()> { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index c58c16beee..af6bbb3423 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -54,6 +54,7 @@ use crate::ListThreadsParams; use crate::ListTurnsParams; use crate::LoadThreadHistoryParams; use crate::MoveThreadToSectionParams; +use crate::PersistContext; use crate::PrepareForkParams; use crate::PreparedFork; use crate::ReadThreadByRolloutPathParams; @@ -400,7 +401,11 @@ impl ThreadStore for LocalThreadStore { Box::pin(async move { live_writer::append_items(self, params).await }) } - fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()> { + fn persist_thread( + &self, + thread_id: ThreadId, + _context: PersistContext, + ) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { live_writer::persist_thread(self, thread_id).await }) } @@ -615,7 +620,7 @@ mod tests { .await .expect("append live item"); store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist live thread"); store @@ -843,7 +848,10 @@ mod tests { ) .await .expect("create live thread with inherited context"); - live_thread.persist().await.expect("persist thread"); + live_thread + .persist(PersistContext::Standard) + .await + .expect("persist thread"); let inherited_metadata = runtime .get_thread(thread_id) .await @@ -1277,7 +1285,7 @@ mod tests { .await .expect("append initial item"); first_store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist initial thread"); first_store @@ -1337,7 +1345,7 @@ mod tests { .await .expect("create live thread"); primary - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist thread for resume"); let rollout_path = primary diff --git a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs index 516c80ecb3..b3edce228b 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs @@ -44,6 +44,7 @@ use crate::DeleteThreadParams; use crate::ForkBoundary; use crate::ListThreadsParams; use crate::ListTurnsParams; +use crate::PersistContext; use crate::PrepareForkParams; use crate::PreparedFork; use crate::ResumeThreadParams; @@ -72,7 +73,7 @@ async fn paginated_history_without_state_db_does_not_initialize_sqlite() { .await .expect("append paginated rollout"); store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist paginated rollout"); store @@ -288,7 +289,7 @@ async fn paginated_live_append_materializes_turn_items_and_state() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -459,7 +460,7 @@ async fn referenced_paginated_rollout_projects_inherited_ordinal_range() { ) .await; store - .persist_thread(child_id) + .persist_thread(child_id, PersistContext::Standard) .await .expect("persist child metadata"); assert_eq!( @@ -563,7 +564,7 @@ async fn named_fork_boundaries_reject_invisible_and_noncanonical_turns() { ) .await; store - .persist_thread(child_id) + .persist_thread(child_id, PersistContext::Standard) .await .expect("persist child metadata"); store @@ -730,7 +731,7 @@ async fn paginated_fork_materializes_compressed_source_and_ancestor() { let ancestor_thread_id = ThreadId::default(); create_paginated_thread(&store, ancestor_thread_id).await; store - .persist_thread(ancestor_thread_id) + .persist_thread(ancestor_thread_id, PersistContext::Standard) .await .expect("persist ancestor meta"); store @@ -766,7 +767,7 @@ async fn paginated_fork_materializes_compressed_source_and_ancestor() { ) .await; store - .persist_thread(source_thread_id) + .persist_thread(source_thread_id, PersistContext::Standard) .await .expect("persist source meta"); store @@ -863,7 +864,7 @@ async fn cancelled_fork_keeps_source_reserved_until_lineage_materialization_fini ) .await; store - .persist_thread(source_thread_id) + .persist_thread(source_thread_id, PersistContext::Standard) .await .expect("persist source metadata"); let source_path = store @@ -924,7 +925,7 @@ async fn prepared_fork_reserves_source_until_child_reference_is_durable() { let source_thread_id = ThreadId::default(); create_paginated_thread(&store, source_thread_id).await; store - .persist_thread(source_thread_id) + .persist_thread(source_thread_id, PersistContext::Standard) .await .expect("persist source metadata"); store @@ -982,7 +983,7 @@ async fn prepared_fork_reserves_source_until_child_reference_is_durable() { ) .await; store - .persist_thread(child_thread_id) + .persist_thread(child_thread_id, PersistContext::Standard) .await .expect("persist child history reference"); drop(prepared); @@ -1010,7 +1011,7 @@ async fn subagent_prefix_advances_projection_without_materializing_history() { ) .await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -1269,7 +1270,7 @@ async fn summary_items_use_final_answers_and_ignore_commentary() { let store = LocalThreadStore::new(config, Some(runtime)); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -1371,7 +1372,7 @@ async fn paginated_projection_accepts_float_rate_limits_and_later_final_answers( let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); store @@ -1512,7 +1513,7 @@ async fn next_write_catches_up_unprojected_durable_suffix() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -1639,7 +1640,7 @@ async fn catch_up_preserves_trailing_partial_line_boundaries() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -1725,7 +1726,7 @@ async fn catch_up_rejects_invalid_complete_suffixes_without_advancing_state() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -1796,7 +1797,7 @@ async fn catch_up_rejects_missing_rollout_after_projection() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); let rollout_path = store @@ -1854,7 +1855,7 @@ async fn blank_and_rejected_rollout_lines_do_not_poison_projection() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -1942,7 +1943,7 @@ async fn unprojectable_rollout_lines_wait_for_later_ordinals() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -2007,7 +2008,7 @@ async fn event_timestamps_allow_invalid_rollout_timestamps() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -2066,7 +2067,7 @@ async fn malformed_rollout_lines_skip_inferred_ordinal_gaps() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); @@ -2153,7 +2154,7 @@ async fn delete_waits_for_in_flight_projection_before_removing_rows() { let thread_id = ThreadId::default(); create_paginated_thread(&store, thread_id).await; store - .persist_thread(thread_id) + .persist_thread(thread_id, PersistContext::Standard) .await .expect("persist session metadata"); let write_permit = store.live_writer_locks.lock(thread_id).await; diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 357282a472..9d88724489 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -43,6 +43,15 @@ use crate::UpdateThreadMetadataParams; /// Future returned by [`ThreadStore`] operations. pub type ThreadStoreFuture<'a, T> = Pin> + Send + 'a>>; +/// Why thread persistence is being requested. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum PersistContext { + /// 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. + TurnStart, +} + /// Storage-neutral thread persistence boundary. pub trait ThreadStore: Any + Send + Sync { /// Return this store as [`Any`] for implementation-owned escape hatches. @@ -69,7 +78,15 @@ pub trait ThreadStore: Any + Send + Sync { fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreFuture<'_, ()>; /// Materializes the thread if persistence is lazy, then persists all queued items. - fn persist_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()>; + /// + /// 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. + fn persist_thread( + &self, + thread_id: ThreadId, + context: PersistContext, + ) -> ThreadStoreFuture<'_, ()>; /// Flushes all queued items and returns once they are durable/readable. fn flush_thread(&self, thread_id: ThreadId) -> ThreadStoreFuture<'_, ()>;