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,