From 21b1ef18c65b5fcdcddc1797ed6e450e94cf9194 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 14 Sep 2026 18:14:03 +0000 Subject: [PATCH] Retain thread persistence acquisition through session cancellation (#45487) ## Why Session initialization can be cancelled after a persistence writer is installed but before acquisition returns. Cleanup must wait for acquisition to finish so it can discard the writer. ## What changed Extend `LiveThreadInitGuard` to own in-flight acquisition and finish it before discarding persistence, including when the guard is dropped. Use the guard for thread creation, resume, and inherited model context initialization. ## Testing Add a regression test that cancels acquisition after writer installation, verifies cleanup waits for handoff, and confirms the writer is removed afterward. GitOrigin-RevId: 09041fa5cd675d082a86a4cc8a719c572afe11a8 --- codex-rs/core/src/session/session.rs | 30 ++++---- codex-rs/thread-store/src/live_thread.rs | 69 ++++++++++++++----- codex-rs/thread-store/src/local/mod.rs | 5 ++ .../src/local/tests/acquisition_tests.rs | 51 ++++++++++++++ 4 files changed, 126 insertions(+), 29 deletions(-) create mode 100644 codex-rs/thread-store/src/local/tests/acquisition_tests.rs diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 3f9a02f0d7..c2b62185ae 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -962,7 +962,8 @@ impl Session { if config.ephemeral { Ok::<_, anyhow::Error>(LiveThreadInitGuard::new(/*live_thread*/ None)) } else { - let live_thread = match &initial_history { + let mut guard = LiveThreadInitGuard::default(); + match &initial_history { InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { let params = CreateThreadParams { session_id, @@ -1008,10 +1009,13 @@ impl Session { Arc::clone(&thread_store), params, items, + &mut guard, ) .await? } else { - LiveThread::create(Arc::clone(&thread_store), params).await? + guard + .acquire(LiveThread::create(Arc::clone(&thread_store), params)) + .await? } } InitialHistory::Resumed(resumed_history) => { @@ -1030,16 +1034,17 @@ impl Session { }, }, }; - LiveThread::resume( - Arc::clone(&thread_store), - session_configuration.history_mode, - params, - ) - .await? + guard + .acquire(LiveThread::resume( + Arc::clone(&thread_store), + session_configuration.history_mode, + params, + )) + .await? } }; // The completed result can wait in join! while the other startup work is pending. - Ok(LiveThreadInitGuard::new(Some(live_thread))) + Ok(guard) } } .instrument(info_span!( @@ -1125,8 +1130,9 @@ impl Session { error!("failed to initialize thread persistence: {e:#}"); e })?; + let live_thread = live_thread_init.as_ref().cloned(); let session_result: anyhow::Result> = async { - let rollout_path = if let Some(live_thread) = live_thread_init.as_ref() { + let rollout_path = if let Some(live_thread) = live_thread.as_ref() { live_thread.local_rollout_path().await? } else { None @@ -1390,7 +1396,7 @@ impl Session { otel.name = "session_init.plugin_skill_warmup", )); let thread_name_lookup = - thread_title_from_thread_store(live_thread_init.as_ref(), &thread_store, thread_id) + thread_title_from_thread_store(live_thread.as_ref(), &thread_store, thread_id) .instrument(info_span!( "session_init.thread_name_lookup", otel.name = "session_init.thread_name_lookup", @@ -1625,7 +1631,7 @@ impl Session { managed_network_requirements_configured, network_approval: Arc::clone(&network_approval), state_db: state_db_ctx.clone(), - live_thread: live_thread_init.as_ref().cloned(), + live_thread: live_thread.clone(), thread_store: Arc::clone(&thread_store), attestation_provider: attestation_provider.clone(), time_provider, diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index 2f043800f2..6feabaebe6 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -1,3 +1,4 @@ +use std::future::Future; use std::path::PathBuf; use std::sync::Arc; @@ -23,6 +24,7 @@ use crate::StoredThreadHistory; use crate::ThreadMetadataPatch; use crate::ThreadStore; use crate::ThreadStoreError; +use crate::ThreadStoreFuture; use crate::ThreadStoreResult; use crate::UpdateThreadMetadataParams; use crate::thread_metadata_sync::ThreadMetadataSync; @@ -41,18 +43,49 @@ pub struct LiveThread { persistence_telemetry: RolloutPersistenceTelemetry, } -/// Owns a live thread while session initialization is still fallible. +/// Owns persistence acquisition and its live thread while initialization is still fallible. /// /// If initialization returns early after persistence has been opened, dropping this guard discards /// the live writer without forcing lazy in-memory state to become durable. Call [`commit`] once the -/// session owns the live thread for normal operation. +/// session owns the live thread for normal operation. Cancellation leaves an in-flight acquisition +/// owned by the guard: cleanup waits for it to finish before discarding any acquired writer. +#[derive(Default)] pub struct LiveThreadInitGuard { live_thread: Option, + acquiring: Option>, } impl LiveThreadInitGuard { pub fn new(live_thread: Option) -> Self { - Self { live_thread } + Self { + live_thread, + acquiring: None, + } + } + + /// Retains the operation even if the caller stops waiting. Store operations may install a + /// writer before returning or delegate to another runtime, so they cannot simply be dropped. + pub async fn acquire( + &mut self, + acquisition: impl Future> + Send + 'static, + ) -> ThreadStoreResult { + assert!(self.live_thread.is_none() && self.acquiring.is_none()); + self.acquiring = Some(Box::pin(acquisition)); + self.finish_acquisition().await?; + self.live_thread + .clone() + .ok_or_else(|| ThreadStoreError::Internal { + message: "persistence acquisition returned no live thread".to_owned(), + }) + } + + async fn finish_acquisition(&mut self) -> ThreadStoreResult<()> { + if let Some(acquiring) = self.acquiring.as_mut() { + let result = acquiring.await; + self.acquiring = None; + self.live_thread = Some(result?); + } + Ok(()) } pub fn as_ref(&self) -> Option<&LiveThread> { @@ -64,6 +97,7 @@ impl LiveThreadInitGuard { } pub async fn discard(&mut self) { + let _ = self.finish_acquisition().await; let Some(live_thread) = self.live_thread.take() else { return; }; @@ -75,15 +109,23 @@ impl LiveThreadInitGuard { impl Drop for LiveThreadInitGuard { fn drop(&mut self) { - let Some(live_thread) = self.live_thread.take() else { + if self.live_thread.is_none() && self.acquiring.is_none() { return; - }; + } let Ok(handle) = tokio::runtime::Handle::try_current() else { warn!("failed to discard thread persistence for failed session init: no Tokio runtime"); return; }; + let acquiring = self.acquiring.take(); + let live_thread = self.live_thread.take(); handle.spawn(async move { - if let Err(err) = live_thread.discard().await { + let live_thread = match acquiring { + Some(acquiring) => acquiring.await.ok(), + None => live_thread, + }; + if let Some(live_thread) = live_thread + && let Err(err) = live_thread.discard().await + { warn!("failed to discard thread persistence for failed session init: {err}"); } }); @@ -116,6 +158,7 @@ impl LiveThread { thread_store: Arc, mut params: CreateThreadParams, inherited_model_context: &[RolloutItem], + guard: &mut LiveThreadInitGuard, ) -> ThreadStoreResult { let persisted_prefix_item_count = persisted_rollout_items(inherited_model_context, params.history_mode).len(); @@ -129,18 +172,10 @@ impl LiveThread { message: "inherited model context is too large".to_string(), })?, ); - let live_thread = Self::create(thread_store, params).await?; - if let Err(err) = live_thread + let live_thread = guard.acquire(Self::create(thread_store, params)).await?; + live_thread .persist_appended_items(inherited_model_context) - .await - { - if let Err(discard_err) = live_thread.discard().await { - warn!( - "failed to discard thread persistence after inherited context append failed: {discard_err}" - ); - } - return Err(err); - } + .await?; Ok(live_thread) } diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 4e20031a50..4d8dfb24c1 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -733,6 +733,8 @@ impl ThreadStore for LocalThreadStore { #[cfg(test)] mod tests { + #[path = "acquisition_tests.rs"] + mod acquisition_tests; use std::sync::Arc; use codex_protocol::ThreadId; @@ -1014,13 +1016,16 @@ mod tests { }) }; + let mut guard = crate::LiveThreadInitGuard::default(); let live_thread = LiveThread::create_with_inherited_model_context( store, params, &[turn_context("parent-model", AskForApproval::Never)], + &mut guard, ) .await .expect("create live thread with inherited context"); + guard.commit(); live_thread .persist(PersistContext::Standard) .await diff --git a/codex-rs/thread-store/src/local/tests/acquisition_tests.rs b/codex-rs/thread-store/src/local/tests/acquisition_tests.rs new file mode 100644 index 0000000000..db02660571 --- /dev/null +++ b/codex-rs/thread-store/src/local/tests/acquisition_tests.rs @@ -0,0 +1,51 @@ +//! Persistence acquisition stays owned when its caller is cancelled after writer installation. + +use super::*; +use crate::LiveThreadInitGuard; +use futures::poll; +use tokio::sync::oneshot; + +#[tokio::test] +async fn cancelled_acquisition_waits_for_writer_handoff_before_discarding() { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(LocalThreadStore::new( + test_config(home.path()), + /*state_db*/ None, + )); + let thread_id = ThreadId::default(); + let params = create_thread_params(thread_id); + let (installed, installation) = oneshot::channel(); + let (release, released) = oneshot::channel(); + let creating_store = Arc::clone(&store); + // Model the cloud wrapper: creation runs independently, installs its writer, then awaits + // initial metadata before returning. Dropping the JoinHandle would detach that work. + let creating = tokio::spawn(async move { + let live_thread = LiveThread::create(creating_store, params).await?; + installed.send(()).expect("installation observed"); + released.await.expect("release acquisition"); + Ok(live_thread) + }); + let mut guard = LiveThreadInitGuard::default(); + let mut start = Box::pin(guard.acquire(async move { creating.await.expect("creation task") })); + assert!(poll!(&mut start).is_pending()); + installation + .await + .expect("writer installed before cancellation"); + drop(start); + store + .flush_thread(thread_id) + .await + .expect("writer installed"); + + let mut cleanup = Box::pin(guard.discard()); + assert!( + poll!(&mut cleanup).is_pending(), + "cleanup must retain the unfinished acquisition" + ); + release.send(()).expect("creation still owned"); + cleanup.await; + assert!(matches!( + store.flush_thread(thread_id).await, + Err(ThreadStoreError::ThreadNotFound { .. }) + )); +}