From d7b0d4aa663172876871b605d53283bc0120882b Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Wed, 30 Sep 2026 11:44:52 +0000 Subject: [PATCH] Return authoritative replay history when resuming a thread (#49599) ## Why A stored history snapshot can become stale before a resumed thread acquires writer ownership. Resume must include writes committed before that acquisition. ## What changed Make `ThreadStore::resume_thread` return the replay context and use it for session restoration and metadata initialization. The local store refreshes stored snapshots while holding writer ownership, reads through the installed writer to honor explicit rollout paths, and releases the writer if loading context fails. Explicit histories without a canonical session header retain their override semantics. ## Testing Extend resume tests to cover edits made after the initial snapshot in both legacy and paginated history modes, stale SQLite rollout paths, and preservation of explicitly supplied history. GitOrigin-RevId: 2361f524d324c859f5dfcf2863320e2045736c2e --- codex-rs/core/src/session/session.rs | 63 +++++-------- .../tests/suite/guardian_persistence_tests.rs | 2 +- .../tests/suite/multi_agent_restore_tests.rs | 3 +- .../tests/suite/pending_input_persistence.rs | 2 +- codex-rs/thread-store/src/in_memory.rs | 28 ++++-- codex-rs/thread-store/src/live_thread.rs | 48 ++++------ .../thread-store/src/local/live_writer.rs | 90 ++++++++++++------- codex-rs/thread-store/src/local/mod.rs | 76 +++++++++++++--- codex-rs/thread-store/src/store.rs | 15 +++- codex-rs/thread-store/src/types.rs | 3 +- 10 files changed, 203 insertions(+), 127 deletions(-) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 966ce1f92d..65cabca09d 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1011,9 +1011,10 @@ impl Session { let mcp_auth = persistence_auth.clone(); let thread_persistence_fut = async { if config.ephemeral { - Ok::<_, anyhow::Error>((None, LiveThreadInitGuard::new(/*live_thread*/ None))) + Ok::<_, anyhow::Error>((None, LiveThreadInitGuard::new(/*live_thread*/ None), None)) } else { let mut local_guard = LiveThreadInitGuard::default(); + let mut resume_context = None; let mut managed_guard = match &startup { Some(startup) => Some(startup.persistence.lock().await), None => None, @@ -1096,16 +1097,23 @@ impl Session { }, }, }; - guard - .acquire(LiveThread::resume( - Arc::clone(&thread_store), - session_configuration.history_mode, - params, - )) - .await? + let store = Arc::clone(&thread_store); + let history_mode = session_configuration.history_mode; + let (context_tx, context_rx) = tokio::sync::oneshot::channel(); + let live_thread = guard + .acquire(async move { + let (live_thread, history) = + LiveThread::resume(store, history_mode, params).await?; + // Cancellation still leaves the writer with the acquisition guard. + let _ = context_tx.send(history); + Ok(live_thread) + }) + .await?; + resume_context = Some(context_rx.await?); + live_thread } }; - Ok((Some(live_thread), local_guard)) + Ok((Some(live_thread), local_guard, resume_context)) } } .instrument(info_span!( @@ -1187,39 +1195,16 @@ impl Session { let (thread_persistence_result, state_db_ctx, (auth, mcp_projection)) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); - let (live_thread, mut live_thread_init) = thread_persistence_result.map_err(|e| { - error!("failed to initialize thread persistence: {e:#}"); - e - })?; + let (live_thread, mut live_thread_init, resume_context) = thread_persistence_result + .map_err(|e| { + error!("failed to initialize thread persistence: {e:#}"); + e + })?; let session_result: anyhow::Result> = async { if let InitialHistory::Resumed(resumed) = &mut initial_history - // A header identifies a stored snapshot. Headerless histories are explicitly - // supplied by callers (for example, imports) and must not be replaced from disk. - && matches!( - resumed.history.first(), - Some(RolloutItem::SessionMeta(meta)) if meta.meta.id == resumed.conversation_id - ) - && let Some(live_thread) = live_thread.as_ref() - && thread_store.as_any().is::() + && let Some(history) = resume_context { - // The initial snapshot predates writer ownership. A cold goal edit may have - // completed in between, including while V1 resume_agent was loading history. - // Reload under the live writer's ownership before publishing replayed context. - let items = match session_configuration.history_mode { - ThreadHistoryMode::Legacy => { - live_thread.load_history(/*include_archived*/ true).await?.items - } - ThreadHistoryMode::Paginated => { - thread_store - .load_latest_model_context(codex_thread_store::LoadThreadHistoryParams { - thread_id, - include_archived: true, - }) - .await? - .items - } - }; - resumed.history = Arc::new(items); + resumed.history = history; } let rollout_path = if let Some(live_thread) = live_thread.as_ref() { live_thread.local_rollout_path().await? diff --git a/codex-rs/core/tests/suite/guardian_persistence_tests.rs b/codex-rs/core/tests/suite/guardian_persistence_tests.rs index f306967d3a..60a072856a 100644 --- a/codex-rs/core/tests/suite/guardian_persistence_tests.rs +++ b/codex-rs/core/tests/suite/guardian_persistence_tests.rs @@ -73,7 +73,7 @@ impl ThreadStore for GatedReviewerStore { } delegate_store_methods! { - fn resume_thread(params: ResumeThreadParams) -> (); + fn resume_thread(params: ResumeThreadParams) -> Arc>; fn append_items(params: AppendThreadItemsParams) -> (); fn discard_thread(thread_id: ThreadId) -> (); fn load_history(params: LoadThreadHistoryParams) -> StoredThreadHistory; diff --git a/codex-rs/core/tests/suite/multi_agent_restore_tests.rs b/codex-rs/core/tests/suite/multi_agent_restore_tests.rs index 6e7abcc142..906816fad7 100644 --- a/codex-rs/core/tests/suite/multi_agent_restore_tests.rs +++ b/codex-rs/core/tests/suite/multi_agent_restore_tests.rs @@ -10,6 +10,7 @@ use super::request_has_model; use anyhow::Context; use anyhow::Result; use codex_features::Feature; +use codex_history::RolloutItem; use codex_protocol::ThreadId; use codex_protocol::error::CodexErrorDetails; use codex_protocol::protocol::EventMsg; @@ -76,7 +77,7 @@ impl ThreadStore for GatedChildMetadataStore { delegate_store_methods! { fn create_thread(params: CreateThreadParams) -> (); - fn resume_thread(params: ResumeThreadParams) -> (); + fn resume_thread(params: ResumeThreadParams) -> Arc>; fn append_items(params: AppendThreadItemsParams) -> (); fn flush_thread(thread_id: ThreadId) -> (); fn shutdown_thread(thread_id: ThreadId) -> (); diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs index e94db81946..9d84ca5324 100644 --- a/codex-rs/core/tests/suite/pending_input_persistence.rs +++ b/codex-rs/core/tests/suite/pending_input_persistence.rs @@ -104,7 +104,7 @@ impl ThreadStore for GatedCheckpointStore { delegate_store_methods! { fn create_thread(params: CreateThreadParams) -> (); - fn resume_thread(params: ResumeThreadParams) -> (); + fn resume_thread(params: ResumeThreadParams) -> Arc>; fn append_items(params: AppendThreadItemsParams) -> (); fn discard_thread(thread_id: ThreadId) -> (); fn load_history(params: LoadThreadHistoryParams) -> StoredThreadHistory; diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index b7a3899260..7b11fab096 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -616,20 +616,33 @@ impl InMemoryThreadStore { Ok(()) } - async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()> { + async fn resume_thread( + &self, + params: ResumeThreadParams, + ) -> ThreadStoreResult>> { let mut state = self.state.lock().await; state.calls.resume_thread += 1; if let Some(history) = params.history { - state - .histories - .insert(params.thread_id, Arc::unwrap_or_clone(history)); + if matches!( + history.first(), + Some(RolloutItem::SessionMeta(meta)) if meta.meta.id == params.thread_id + ) { + state + .histories + .entry(params.thread_id) + .or_insert_with(|| Arc::unwrap_or_clone(history)); + } else { + state + .histories + .insert(params.thread_id, Arc::unwrap_or_clone(history)); + } } else { state.histories.entry(params.thread_id).or_default(); } if let Some(rollout_path) = params.rollout_path { state.rollout_paths.insert(rollout_path, params.thread_id); } - Ok(()) + Ok(Arc::new(state.histories[¶ms.thread_id].clone())) } async fn append_items(&self, params: AppendThreadItemsParams) -> ThreadStoreResult<()> { @@ -900,7 +913,10 @@ impl ThreadStore for InMemoryThreadStore { Box::pin(InMemoryThreadStore::create_thread(self, params)) } - fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()> { + fn resume_thread( + &self, + params: ResumeThreadParams, + ) -> ThreadStoreFuture<'_, Arc>> { Box::pin(InMemoryThreadStore::resume_thread(self, params)) } diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index f6f57e6dd0..f9bb507fbc 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -179,14 +179,13 @@ impl LiveThread { Ok(live_thread) } + /// Reopens persistence and returns the replay history selected under writer ownership. pub async fn resume( thread_store: Arc, history_mode: ThreadHistoryMode, - params: ResumeThreadParams, - ) -> ThreadStoreResult { + mut params: ResumeThreadParams, + ) -> ThreadStoreResult<(Self, Arc>)> { let thread_id = params.thread_id; - let should_load_history = params.history.is_none(); - let include_archived = params.include_archived; let metadata = if history_mode == ThreadHistoryMode::Paginated && let Some(local_store) = thread_store.as_any().downcast_ref::() && let Some(state_db) = local_store.state_db().await @@ -200,34 +199,19 @@ impl LiveThread { } else { None }; - let mut metadata_sync = ThreadMetadataSync::for_resume(¶ms, metadata.as_ref()); - thread_store.resume_thread(params).await?; - if should_load_history { - match thread_store - .load_history(LoadThreadHistoryParams { - thread_id, - include_archived, - }) - .await - { - Ok(history) => metadata_sync.record_resume_history(&history.items), - Err(err) => { - if let Err(discard_err) = thread_store.discard_thread(thread_id).await { - warn!( - "failed to discard thread persistence after resume history load failed: {discard_err}" - ); - } - return Err(err); - } - } - } - Ok(Self { - thread_id, - history_mode, - thread_store, - metadata_sync: Arc::new(Mutex::new(metadata_sync)), - persistence_telemetry: RolloutPersistenceTelemetry::new(thread_id), - }) + let history = thread_store.resume_thread(params.clone()).await?; + params.history = Some(Arc::clone(&history)); + let metadata_sync = ThreadMetadataSync::for_resume(¶ms, metadata.as_ref()); + Ok(( + Self { + thread_id, + history_mode, + thread_store, + metadata_sync: Arc::new(Mutex::new(metadata_sync)), + persistence_telemetry: RolloutPersistenceTelemetry::new(thread_id), + }, + history, + )) } #[tracing::instrument( diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 6a2d95f2b0..450f0304da 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -1,4 +1,5 @@ use std::path::PathBuf; +use std::sync::Arc; use codex_protocol::ThreadId; use codex_protocol::protocol::ThreadHistoryMode; @@ -14,7 +15,7 @@ use super::LocalThreadStore; use super::create_thread; use crate::AppendThreadItemsParams; use crate::CreateThreadParams; -use crate::ReadThreadParams; +use crate::LoadThreadHistoryParams; use crate::ResumeThreadParams; use crate::ThreadStoreError; use crate::ThreadStoreResult; @@ -40,13 +41,42 @@ pub(super) async fn create_thread( pub(super) async fn resume_thread( store: &LocalThreadStore, params: ResumeThreadParams, -) -> ThreadStoreResult<()> { +) -> ThreadStoreResult>> { let _live_writer_guard = store.live_writer_locks.lock(params.thread_id).await; store.ensure_live_recorder_absent(params.thread_id).await?; let writer_lock = store.acquire_writer_lock(params.thread_id)?; + let supplied_history = params + .history + .as_ref() + .filter(|history| { + !matches!( + history.first(), + Some(RolloutItem::SessionMeta(meta)) if meta.meta.id == params.thread_id + ) + }) + .cloned(); + let rollout_path = match params.rollout_path { + Some(rollout_path) => rollout_path, + None => { + let resolved = if params.include_archived { + super::thread_rollout_resolver::resolve_current_including_archived( + store, + params.thread_id, + ) + .await? + } else { + super::thread_rollout_resolver::resolve_current(store, params.thread_id).await? + }; + resolved + .ok_or(ThreadStoreError::ThreadNotFound { + thread_id: params.thread_id, + })? + .path + } + }; let history_mode = if let Some(history) = params.history.as_deref() { canonical_history_mode_from_rollout_items(history) - } else if let Some(rollout_path) = params.rollout_path.as_ref() { + } else { super::read_thread::read_thread_by_rollout_path( store, rollout_path.clone(), @@ -55,36 +85,6 @@ pub(super) async fn resume_thread( ) .await? .history_mode - } else { - super::read_thread::read_thread( - store, - ReadThreadParams { - thread_id: params.thread_id, - include_archived: params.include_archived, - include_history: false, - }, - ) - .await? - .history_mode - }; - let rollout_path = match (params.rollout_path, params.history) { - (Some(rollout_path), _history) => rollout_path, - (None, history) => { - let thread = super::read_thread::read_thread( - store, - ReadThreadParams { - thread_id: params.thread_id, - include_archived: params.include_archived, - include_history: history.is_none(), - }, - ) - .await?; - thread - .rollout_path - .ok_or_else(|| ThreadStoreError::Internal { - message: format!("thread {} does not have a rollout path", params.thread_id), - })? - } }; let cwd = params .metadata @@ -122,7 +122,29 @@ pub(super) async fn resume_thread( history_mode, writer_lock, ) - .await + .await?; + if let Some(history) = supplied_history { + return Ok(history); + } + // Resolve through the installed writer so explicit paths and materialized compressed rollouts + // take precedence over SQLite. Both ownership guards remain held throughout this read. + let history = super::model_context::load_latest_model_context( + store, + LoadThreadHistoryParams { + thread_id: params.thread_id, + include_archived: params.include_archived, + }, + ) + .await; + if history.is_err() { + let entry = store.live_recorders.lock().await.remove(¶ms.thread_id); + if let Some(entry) = entry + && let Err(err) = entry.recorder.discard().await + { + warn!("failed to discard writer after resume history load failed: {err}"); + } + } + history.map(|history| Arc::new(history.items)) } #[tracing::instrument( diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 05f2d9061d..55822428a1 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -44,6 +44,7 @@ mod timestamp_metadata_tests; use codex_protocol::ThreadId; use codex_protocol::protocol::ThreadHistoryMode; +use codex_rollout::RolloutItem; use codex_rollout::RolloutRecorder; use codex_rollout::StateDbHandle; use codex_rollout::WriterLockCoordinator; @@ -526,7 +527,10 @@ impl ThreadStore for LocalThreadStore { }) } - fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()> { + fn resume_thread( + &self, + params: ResumeThreadParams, + ) -> ThreadStoreFuture<'_, Arc>> { Box::pin(async move { live_writer::resume_thread(self, params).await }) } @@ -804,6 +808,7 @@ mod tests { use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; use codex_protocol::models::BaseInstructions; + use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; @@ -992,13 +997,14 @@ mod tests { .await .expect("update stale sqlite rollout path"); - let resumed = LiveThread::resume( + let history = Arc::new(vec![user_message_item("bounded suffix")]); + let (resumed, supplied_history) = LiveThread::resume( store, ThreadHistoryMode::Paginated, ResumeThreadParams { thread_id, rollout_path: Some(rollout_path.clone()), - history: Some(Arc::new(vec![user_message_item("bounded suffix")])), + history: Some(Arc::clone(&history)), include_archived: false, metadata: ThreadPersistenceMetadata { cwd: Some(home.path().to_path_buf()), @@ -1009,6 +1015,7 @@ mod tests { ) .await .expect("resume paginated thread from its requested rollout"); + assert!(Arc::ptr_eq(&supplied_history, &history)); assert_eq!( resumed.local_rollout_path().await.expect("live rollout"), Some(rollout_path) @@ -1344,7 +1351,7 @@ mod tests { let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(home.path(), "2025-01-03T17-00-00", uuid).expect("session file"); - let live_thread = LiveThread::resume( + let (live_thread, _history) = LiveThread::resume( store, ThreadHistoryMode::Legacy, ResumeThreadParams { @@ -1399,7 +1406,7 @@ mod tests { let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); let rollout_path = write_session_file(external_home.path(), "2025-01-03T17-30-00", uuid) .expect("external session file"); - let live_thread = LiveThread::resume( + let (live_thread, _history) = LiveThread::resume( store, ThreadHistoryMode::Legacy, ResumeThreadParams { @@ -1529,12 +1536,18 @@ mod tests { } #[tokio::test] - async fn resume_thread_reopens_live_writer_and_appends() { + async fn resume_thread_reopens_live_writer_and_appends_with_stale_sqlite_path() { let home = TempDir::new().expect("temp dir"); let config = test_config(home.path()); + let runtime = codex_state::StateRuntime::init( + config.sqlite.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); let thread_id = ThreadId::default(); - let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None); + let first_store = LocalThreadStore::new(config.clone(), Some(runtime.clone())); first_store .create_thread(create_thread_params(thread_id)) .await @@ -1563,7 +1576,21 @@ mod tests { .await .expect("shutdown initial writer"); - let resumed_store = LocalThreadStore::new(config, /*state_db*/ None); + let mut builder = codex_state::ThreadMetadataBuilder::new( + thread_id, + home.path().join("missing-rollout.jsonl"), + chrono::Utc::now(), + SessionSource::Exec, + ); + builder.history_mode = ThreadHistoryMode::Legacy; + builder.cwd = home.path().to_path_buf(); + let metadata = builder.build("test-provider"); + runtime + .upsert_thread(&metadata) + .await + .expect("store stale sqlite path"); + + let resumed_store = LocalThreadStore::new(config, Some(runtime)); resumed_store .resume_thread(ResumeThreadParams { thread_id, @@ -1614,10 +1641,18 @@ mod tests { .live_rollout_path(thread_id) .await .expect("load rollout path"); + let mut expected = primary + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: true, + }) + .await + .expect("load initial snapshot") + .items; let resume_params = ResumeThreadParams { thread_id, rollout_path: Some(rollout_path), - history: None, + history: Some(Arc::new(expected.clone())), include_archived: true, metadata: thread_metadata(), }; @@ -1634,14 +1669,35 @@ mod tests { .expect_err("competing resume should fail"); assert!(matches!(error, ThreadStoreError::Conflict { .. })); + let late_edit = RolloutItem::ResponseItem( + ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "edit after the initial snapshot".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + } + .into(), + ); + primary + .append_items(AppendThreadItemsParams { + thread_id, + items: vec![late_edit.clone()], + }) + .await + .expect("append before handing off writer ownership"); + expected.push(late_edit); primary .shutdown_thread(thread_id) .await .expect("shutdown should release writer ownership"); - secondary + let resumed = secondary .resume_thread(resume_params) .await .expect("resume after shutdown should acquire writer ownership"); + assert_eq!(serde_json::json!(*resumed), serde_json::json!(expected)); } } diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index d67c3b766a..fb7b93803d 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -1,8 +1,10 @@ use codex_protocol::ThreadId; use codex_protocol::protocol::ThreadHistoryMode; +use codex_rollout::RolloutItem; use std::any::Any; use std::future::Future; use std::pin::Pin; +use std::sync::Arc; use crate::AddThreadAttachmentOutcome; use crate::AddThreadAttachmentParams; @@ -138,8 +140,17 @@ pub trait ThreadStore: Any + Send + Sync { }) } - /// Reopens an existing thread for live appends. - fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreFuture<'_, ()>; + /// Reopens an existing thread for live appends and returns its authoritative replay context. + /// + /// Stored snapshots supplied by callers may predate writer acquisition. Implementations must + /// include writes committed before acquisition in the returned context. Explicitly supplied + /// histories without a canonical session header retain their import/override semantics. + /// Returned histories can share the caller's allocation when no refresh is needed. + /// On failure, implementations must release any writer acquired by this operation. + fn resume_thread( + &self, + params: ResumeThreadParams, + ) -> ThreadStoreFuture<'_, Arc>>; /// Appends raw rollout items to a live thread. /// diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 5030f8864d..4e842556d3 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -125,7 +125,8 @@ pub struct ResumeThreadParams { pub thread_id: ThreadId, /// Known local rollout path when the caller resumed from a specific file. pub rollout_path: Option, - /// Known replay history for the resumed thread, if already loaded by the caller. + /// Previously loaded replay history. A canonical session header identifies a stored snapshot + /// that may need refreshing under writer ownership; other histories are explicit overrides. pub history: Option>>, /// Whether archived threads may be reopened. pub include_archived: bool,