From ce2c2759ebee2d64565922f6f7365082284f9570 Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 9 Sep 2026 15:43:24 +0000 Subject: [PATCH] Release persistent writers when session startup is cancelled (#44183) ## Why Cancelling a session after persistence opens but while MCP startup is still pending can leave its writer held, blocking a subsequent resume. ## What changed Create `LiveThreadInitGuard` inside the persistence startup future so it protects the live thread while `tokio::join!` waits for other startup work. Dropping the startup future then schedules writer cleanup. ## Testing Add a regression test that cancels a resume while MCP startup is blocked, verifies the local writer is released, and resumes the same thread successfully. GitOrigin-RevId: 57d8624e783a7b993579a6386aea35ba845da0e3 --- codex-rs/core/src/session/session.rs | 14 +-- codex-rs/core/tests/suite/mod.rs | 1 + .../core/tests/suite/startup_cancellation.rs | 115 ++++++++++++++++++ 3 files changed, 123 insertions(+), 7 deletions(-) create mode 100644 codex-rs/core/tests/suite/startup_cancellation.rs diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 261e180167..2c3ea5fc9e 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -869,7 +869,7 @@ impl Session { // - load history metadata (skipped for subagents) let thread_persistence_fut = async { if config.ephemeral { - Ok::<_, anyhow::Error>(None) + Ok::<_, anyhow::Error>(LiveThreadInitGuard::new(/*live_thread*/ None)) } else { let live_thread = match &initial_history { InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { @@ -947,7 +947,8 @@ impl Session { .await? } }; - Ok(Some(live_thread)) + // The completed result can wait in join! while the other startup work is pending. + Ok(LiveThreadInitGuard::new(Some(live_thread))) } } .instrument(info_span!( @@ -1027,11 +1028,10 @@ 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 mut live_thread_init = - LiveThreadInitGuard::new(thread_persistence_result.map_err(|e| { - error!("failed to initialize thread persistence: {e:#}"); - e - })?); + let mut live_thread_init = thread_persistence_result.map_err(|e| { + error!("failed to initialize thread persistence: {e:#}"); + e + })?; let session_result: anyhow::Result> = async { let rollout_path = if let Some(live_thread) = live_thread_init.as_ref() { live_thread.local_rollout_path().await? diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 95eb4a4e0b..73c08c5452 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -165,6 +165,7 @@ mod skills; mod skills_extension; mod spawn_agent_description; mod sqlite_state; +mod startup_cancellation; mod step_settings; mod step_settings_snapshots; mod stream_error_allows_next_turn; diff --git a/codex-rs/core/tests/suite/startup_cancellation.rs b/codex-rs/core/tests/suite/startup_cancellation.rs new file mode 100644 index 0000000000..5901e1feb5 --- /dev/null +++ b/codex-rs/core/tests/suite/startup_cancellation.rs @@ -0,0 +1,115 @@ +//! Cancelling partially initialized sessions must release their persistent writers. + +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_core::StartThreadOptions; +use codex_core::config::Config; +use codex_core::config::ThreadStoreConfig; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::McpServerContribution; +use codex_extension_api::McpServerContributionContext; +use codex_extension_api::McpServerContributor; +use codex_rollout::RolloutRecorder; +use codex_thread_store::LocalThreadStore; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use tokio::sync::Notify; +use tokio::time::timeout; + +struct StartupMcpBarrier { + block_next: AtomicBool, + entered: Notify, +} + +impl McpServerContributor for StartupMcpBarrier { + fn id(&self) -> &'static str { + "cancelled_startup_barrier" + } + + fn contribute<'a>( + &'a self, + _context: McpServerContributionContext<'a, Config>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + if self.block_next.swap(false, Ordering::SeqCst) { + self.entered.notify_one(); + std::future::pending::<()>().await; + } + Vec::new() + }) + } +} + +#[tokio::test] +async fn cancelled_resume_releases_writer_while_mcp_startup_is_pending() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let barrier = Arc::new(StartupMcpBarrier { + block_next: AtomicBool::new(false), + entered: Notify::new(), + }); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.mcp_server_contributor(barrier.clone()); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| config.experimental_thread_store = ThreadStoreConfig::Local) + .build_with_auto_env(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let environments = test.codex.environment_selections().await; + test.codex.ensure_rollout_materialized().await; + let rollout_path = test.codex.rollout_path().context("thread rollout")?; + test.codex.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&thread_id).await; + let history = RolloutRecorder::get_rollout_history(&rollout_path).await?; + let store = test + .thread_store + .as_any() + .downcast_ref::() + .context("local thread store")?; + let resume_options = || StartThreadOptions { + initial_history: history.clone(), + environments: Some(environments.clone()), + ..StartThreadOptions::new(test.config.clone()) + }; + + barrier.block_next.store(true, Ordering::SeqCst); + let mut resume = Box::pin(test.thread_manager.start_thread(resume_options())); + timeout(Duration::from_secs(10), async { + tokio::select! { + biased; + _ = &mut resume => panic!("resume must wait for MCP startup"), + _ = async { + barrier.entered.notified().await; + while store.live_rollout_path(thread_id).await.is_err() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + } => {} + } + }) + .await + .context("resume should open persistence while MCP startup waits")?; + drop(resume); + + // The guard schedules asynchronous cleanup. Resuming then also waits for discard's writer lock. + let resumed = timeout(Duration::from_secs(10), async { + while store.live_rollout_path(thread_id).await.is_ok() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + test.thread_manager.start_thread(resume_options()).await + }) + .await + .context("cancelled startup should release its writer")??; + assert_eq!(resumed.thread_id, thread_id); + resumed.thread.shutdown_and_wait().await?; + Ok(()) +}