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
2026-11-03 19:53:25 +00:00 · 2026-09-14 19:30:11 +00:00
parent ad8a5e3a1b
commit 08d3748cf0
10 changed files with 434 additions and 10 deletions
@@ -0,0 +1,109 @@
+//! Runs a thread for its owner's chosen lifetime, including cancellation during startup.
+//! Cleanup remains tracked when the caller drops its startup future or stops waiting.
+
+use std::future::Future;
+use std::sync::Arc;
+
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::Result as CodexResult;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+use tokio_util::task::TaskTracker;
+
+use super::NewThread;
+use super::StartThreadOptions;
+use super::ThreadManager;
+use crate::session::startup::SessionStartup;
+
+impl ThreadManager {
+    /// Starts a new isolated thread and owns it until `until` completes or the thread exits.
+    /// Dropping the startup future also cancels initialization. `tasks` tracks the
+    /// entire lifetime, including partial startup cleanup and deregistration.
+    /// Cleanup waits for any in-flight persistence acquisition before releasing its writer.
+    /// Once the session runtime starts, cleanup follows normal shutdown history rules,
+    /// even if the caller has not received the thread yet.
+    /// Resumed threads are excluded because they can refer to another owner's live runtime.
+    pub async fn start_thread_until(
+        self: &Arc<Self>,
+        options: StartThreadOptions,
+        until: impl Future<Output = ()> + Send + 'static,
+        tasks: &TaskTracker,
+    ) -> CodexResult<NewThread> {
+        if matches!(
+            options.initial_history,
+            codex_history::InitialHistory::Resumed(_)
+        ) {
+            return Err(CodexErr::InvalidRequest(
+                "owned thread startup requires new or forked history".to_owned(),
+            ));
+        }
+        if options
+            .thread_extension_init
+            .get::<codex_extension_api::SessionIsolation>()
+            .as_deref()
+            != Some(&codex_extension_api::SessionIsolation::Isolated)
+        {
+            return Err(CodexErr::InvalidRequest(
+                "owned thread startup requires explicit session isolation".to_owned(),
+            ));
+        }
+        let task = tasks.token();
+        let manager = Arc::clone(self);
+        let state = Arc::downgrade(&self.state);
+        let (sender, receiver) = oneshot::channel();
+        let abandoned = CancellationToken::new();
+        let handoff = abandoned.clone().drop_guard();
+        tokio::spawn(async move {
+            let _task = task;
+            let startup = Arc::new(SessionStartup::default());
+            tokio::pin!(until);
+            // Drop initialization before cleanup; startup still owns any unfinished acquisition.
+            let result = {
+                let start = manager.start_thread_inner(
+                    options,
+                    /*forked_from_thread_id*/ None,
+                    Some(Arc::clone(&startup)),
+                );
+                tokio::select! {
+                    biased;
+                    _ = &mut until => Err(CodexErr::TurnAborted),
+                    _ = abandoned.cancelled() => Err(CodexErr::TurnAborted),
+                    result = start => result,
+                }
+            };
+            // The lifetime task must not keep the manager and its parent threads alive.
+            drop(manager);
+            match result {
+                Ok(thread) => {
+                    let running = Arc::clone(&thread.thread);
+                    if sender.send(Ok(thread)).is_ok() {
+                        tokio::select! {
+                            _ = abandoned.cancelled() => {}
+                            _ = &mut until => {}
+                            _ = running.wait_until_terminated() => {}
+                        }
+                    }
+                }
+                Err(error) => {
+                    let _ = sender.send(Err(error));
+                }
+            }
+            startup.cleanup().await;
+            if let Some(state) = state.upgrade()
+                && let Some(session) = startup.session.get()
+            {
+                let mut threads = state.threads.write().await;
+                let thread_id = session.thread_id();
+                if threads
+                    .get(&thread_id)
+                    .is_some_and(|thread| Arc::ptr_eq(&thread.session, session))
+                {
+                    threads.remove(&thread_id);
+                }
+            }
+        });
+        let thread = receiver.await.map_err(|_| CodexErr::InternalAgentDied)??;
+        handoff.disarm();
+        Ok(thread)
+    }
+}