openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-02 19:43:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retain thread persistence acquisition through session cancellation (#45487)

Browse Source 
## 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
This commit is contained in:
jif
2026-09-14 18:40:02 +00:00
committed by copyberry
parent 5fb3b7e401
commit 21b1ef18c6
4 changed files with 126 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/session/session.rs
+18 -12
View File
@@ -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<Arc<Self>> = 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,
codex-rs/thread-store/src/live_thread.rs
+52 -17
View File
@@ -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<LiveThread>,
acquiring: Option<ThreadStoreFuture<'static, LiveThread>>,
}
impl LiveThreadInitGuard {
pub fn new(live_thread: Option<LiveThread>) -> 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<Output = ThreadStoreResult<LiveThread>> + Send + 'static,
) -> ThreadStoreResult<LiveThread> {
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<dyn ThreadStore>,
mut params: CreateThreadParams,
inherited_model_context: &[RolloutItem],
guard: &mut LiveThreadInitGuard,
) -> ThreadStoreResult<Self> {
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)
}
codex-rs/thread-store/src/local/mod.rs
+5
View File
@@ -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
codex-rs/thread-store/src/local/tests/acquisition_tests.rs
+51
View File
@@ -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 { .. })
));
}