Add managed thread lifetimes with cancellation-safe startup (#45502)

## Why

Callers need to tie isolated threads to an explicit lifetime and wait for cleanup even when startup is cancelled or its result is never received.

## What changed

- Add `ThreadManager::start_thread_until` to run an isolated thread until a caller-provided future completes or the thread exits, with a `TaskTracker` covering cleanup and deregistration.
- Retain persistence and session resources across interrupted startup so cleanup can release partially initialized resources or shut down a running session using normal history rules.
- Reject resumed history and startup without explicit session isolation.

## Testing

Add integration tests for cancellation during stalled required MCP initialization, dropping an unconsumed startup result, and owner cancellation that preserves history and parent usability. Also verify that failed duplicate startup leaves the existing thread's writer intact.

GitOrigin-RevId: 47251c1b0c43821ab8f950816c6d341d39cdc823

codex-rs/core/src/session/mod.rs

@@ -1,3 +1,5 @@
+pub(crate) mod startup;
+
 use std::borrow::Cow;
 use std::collections::HashMap;
 use std::collections::HashSet;
@@ -397,6 +399,7 @@ use codex_utils_stream_parser::ProposedPlanSegment;
 /// Runtime state lives on `Session`; keeping these endpoints separate lets all
 /// submission senders be dropped to terminate the session loop. The shared
 /// completion future observes that shutdown.
+#[derive(Clone)]
 pub(crate) struct SessionIo {
     pub(crate) tx_sub: Sender<Submission>,
     pub(crate) rx_event: Receiver<Event>,
@@ -425,6 +428,7 @@ pub(crate) enum ForkPersistence {
 }
 
 pub(crate) struct SessionSpawnArgs {
+    pub(crate) startup: Option<Arc<startup::SessionStartup>>,
     pub(crate) config: Config,
     pub(crate) allow_provider_model_fallback: bool,
     pub(crate) instructions: SessionInstructions,
@@ -530,6 +534,7 @@ impl Session {
 
     async fn spawn_internal(args: SessionSpawnArgs) -> CodexResult<(Arc<Self>, SessionIo)> {
         let SessionSpawnArgs {
+            startup,
             config,
             allow_provider_model_fallback,
             instructions,
@@ -822,6 +827,7 @@ impl Session {
         let (agent_status_tx, agent_status_rx) = watch::channel(AgentStatus::PendingInit);
 
         let session = Box::pin(Session::new(
+            startup.clone(),
             session_configuration,
             &environment_selections,
             config.clone(),
@@ -889,6 +895,9 @@ impl Session {
             session_loop_termination: session_loop_termination_from_handle(session_loop_handle),
         };
 
+        if let Some(startup) = startup {
+            let _ = startup.io.set(io.clone());
+        }
         Ok((session, io))
     }
 }

codex-rs/core/src/session/session.rs

@@ -725,6 +725,7 @@ impl Session {
     #[instrument(name = "session_init", level = "info", skip_all)]
     #[allow(clippy::too_many_arguments)]
     pub(crate) async fn new(
+        startup: Option<Arc<super::startup::SessionStartup>>,
         mut session_configuration: SessionConfiguration,
         environment_selections: &[TurnEnvironmentSelection],
         config: Arc<Config>,
@@ -960,10 +961,15 @@ impl Session {
         // - load history metadata (skipped for subagents)
         let thread_persistence_fut = async {
             if config.ephemeral {
-                Ok::<_, anyhow::Error>(LiveThreadInitGuard::new(/*live_thread*/ None))
+                Ok::<_, anyhow::Error>((None, LiveThreadInitGuard::new(/*live_thread*/ None)))
             } else {
-                let mut guard = LiveThreadInitGuard::default();
-                match &initial_history {
+                let mut local_guard = LiveThreadInitGuard::default();
+                let mut managed_guard = match &startup {
+                    Some(startup) => Some(startup.persistence.lock().await),
+                    None => None,
+                };
+                let guard = managed_guard.as_deref_mut().unwrap_or(&mut local_guard);
+                let live_thread = match &initial_history {
                     InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => {
                         let params = CreateThreadParams {
                             session_id,
@@ -1009,7 +1015,7 @@ impl Session {
                                 Arc::clone(&thread_store),
                                 params,
                                 items,
-                                &mut guard,
+                                guard,
                             )
                             .await?
                         } else {
@@ -1043,8 +1049,7 @@ impl Session {
                             .await?
                     }
                 };
-                // The completed result can wait in join! while the other startup work is pending.
-                Ok(guard)
+                Ok((Some(live_thread), local_guard))
             }
         }
         .instrument(info_span!(
@@ -1126,11 +1131,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 = thread_persistence_result.map_err(|e| {
+        let (live_thread, mut live_thread_init) = thread_persistence_result.map_err(|e| {
             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.as_ref() {
                 live_thread.local_rollout_path().await?
@@ -1710,6 +1714,9 @@ impl Session {
                 forked_from_ordinal_exclusive,
                 next_internal_sub_id: AtomicU64::new(0),
             });
+            if let Some(startup) = &startup {
+                let _ = startup.session.set(Arc::clone(&sess));
+            }
             if let Some(network_policy_decider_session) = network_policy_decider_session {
                 let mut guard = network_policy_decider_session.write().await;
                 *guard = Arc::downgrade(&sess);

codex-rs/core/src/session/startup.rs

@@ -0,0 +1,37 @@
+//! Retains persistence acquisition and resources throughout managed startup.
+//! The thread manager drops initialization, then joins acquisition and resource cleanup here.
+
+use std::sync::Arc;
+use std::sync::OnceLock;
+
+use codex_protocol::protocol::Op;
+use codex_thread_store::LiveThreadInitGuard;
+use tokio::sync::Mutex;
+
+use super::SessionIo;
+use super::session::Session;
+
+#[derive(Default)]
+pub(crate) struct SessionStartup {
+    pub(crate) persistence: Mutex<LiveThreadInitGuard>,
+    pub(crate) session: OnceLock<Arc<Session>>,
+    pub(crate) io: OnceLock<SessionIo>,
+}
+
+impl SessionStartup {
+    pub(crate) async fn cleanup(&self) {
+        if let Some(io) = self.io.get() {
+            // The session loop owns persistence now. Preserve its shutdown semantics even
+            // if registration or the caller's handoff was interrupted after the loop started.
+            self.persistence.lock().await.commit();
+            let _ = io.submit(Op::Interrupt).await;
+            let _ = io.shutdown_and_wait().await;
+        } else {
+            if let Some(session) = self.session.get() {
+                super::handlers::shutdown_session_runtime(session).await;
+            }
+            let mut persistence = std::mem::take(&mut *self.persistence.lock().await);
+            persistence.discard().await;
+        }
+    }
+}

codex-rs/core/src/session/tests.rs

@@ -5787,6 +5787,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() {
     ));
     let environment_manager = Arc::new(EnvironmentManager::default_for_tests());
     let result = Session::new(
+        /*startup*/ None,
         session_configuration,
         /*environment_selections*/ &[],
         Arc::clone(&config),
@@ -6270,6 +6271,7 @@ async fn make_session_with_config_and_rx(
     let environment_manager = Arc::new(EnvironmentManager::default_for_tests());
 
     let session = Session::new(
+        /*startup*/ None,
         session_configuration,
         &default_environments,
         Arc::clone(&config),
@@ -6399,6 +6401,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx(
     let environment_manager = Arc::new(EnvironmentManager::default_for_tests());
 
     let session = Session::new(
+        /*startup*/ None,
         session_configuration,
         &default_environments,
         Arc::clone(&config),

codex-rs/core/src/session/tests/guardian_tests.rs

@@ -1260,6 +1260,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
     let mut thread_extension_init = codex_extension_api::ExtensionDataInit::default();
     thread_extension_init.insert(codex_extension_api::SessionIsolation::Isolated);
     let (session, io) = Session::spawn(SessionSpawnArgs {
+        startup: None,
         config,
         allow_provider_model_fallback: false,
         instructions: Default::default(),