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
This commit is contained in:
jif
2026-09-14 19:30:11 +00:00
committed by copyberry
parent ad8a5e3a1b
commit 08d3748cf0
10 changed files with 434 additions and 10 deletions
+1
View File
@@ -87,6 +87,7 @@ pub(crate) async fn run_codex_thread_interactive(
thread_extension_init.insert(isolation);
thread_extension_init.insert(crate::guardian::GuardianReviewSessionHost::default());
let (session, io) = Session::spawn(SessionSpawnArgs {
startup: None,
config,
allow_provider_model_fallback: false,
instructions,
+9
View File
@@ -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))
}
}
+15 -8
View File
@@ -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);
+37
View File
@@ -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;
}
}
}
+3
View File
@@ -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),
@@ -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(),
+13 -2
View File
@@ -1,3 +1,5 @@
mod managed;
use crate::CodexAppsToolsCache;
use crate::agent::AgentControl;
use crate::agents_md_manager::SessionInstructions;
@@ -291,6 +293,7 @@ impl StartThreadOptions {
}
struct ThreadSpawnRequest {
startup: Option<Arc<crate::session::startup::SessionStartup>>,
options: StartThreadOptions,
auth_manager: Arc<AuthManager>,
agent_control: AgentControl,
@@ -310,6 +313,7 @@ impl ThreadSpawnRequest {
agent_control: AgentControl,
) -> Self {
Self {
startup: None,
options,
auth_manager,
agent_control,
@@ -987,7 +991,10 @@ impl ThreadManager {
}
pub async fn start_thread(&self, options: StartThreadOptions) -> CodexResult<NewThread> {
Box::pin(self.start_thread_inner(options, /*forked_from_thread_id*/ None)).await
Box::pin(self.start_thread_inner(
options, /*forked_from_thread_id*/ None, /*startup*/ None,
))
.await
}
/// Starts a fresh internal session associated with an existing parent thread.
@@ -1059,6 +1066,7 @@ impl ThreadManager {
&self,
mut options: StartThreadOptions,
forked_from_thread_id: Option<ThreadId>,
startup: Option<Arc<crate::session::startup::SessionStartup>>,
) -> CodexResult<NewThread> {
let agent_control = self.agent_control_for_config(&options.config);
let (resumed_session_source, resumed_thread_source) = options
@@ -1075,6 +1083,7 @@ impl ThreadManager {
let mut request =
ThreadSpawnRequest::new(options, Arc::clone(&self.state.auth_manager), agent_control);
request.forked_from_thread_id = forked_from_thread_id;
request.startup = startup;
Box::pin(self.state.spawn_thread(request)).await
}
@@ -1111,7 +1120,7 @@ impl ThreadManager {
inherited_multi_agent_version,
),
);
self.start_thread_inner(options, Some(forked_from_thread_id))
self.start_thread_inner(options, Some(forked_from_thread_id), /*startup*/ None)
.await
}
@@ -1924,6 +1933,7 @@ impl ThreadManagerState {
/// Spawn a new thread with optional history and register it with the manager.
async fn spawn_thread(&self, request: ThreadSpawnRequest) -> CodexResult<NewThread> {
let ThreadSpawnRequest {
startup,
options,
auth_manager,
agent_control,
@@ -2094,6 +2104,7 @@ impl ThreadManagerState {
codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile
};
let (session, io) = Session::spawn(SessionSpawnArgs {
startup,
config,
allow_provider_model_fallback,
instructions,
+109
View File
@@ -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)
}
}
@@ -0,0 +1,244 @@
//! Exercises owned thread startup and shutdown through the real runtime and persistence store.
use std::future::Future;
use std::sync::Arc;
use std::task::Context as TaskContext;
use std::task::Wake;
use std::time::Duration;
use anyhow::Context;
use codex_core::StartThreadOptions;
use codex_core::TurnInputRequest;
use codex_extension_api::SessionIsolation;
use codex_protocol::protocol::EventMsg;
use codex_protocol::user_input::UserInput;
use codex_thread_store::ThreadStoreError;
use core_test_support::apps_test_server::AppsTestServer;
use core_test_support::responses;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn dropping_startup_cleans_up_while_required_mcp_is_stalled() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let fixture = test_codex().build_with_auto_env(&server).await?;
let mcp_server = responses::start_mock_server().await;
let (http_server, control) = AppsTestServer::mount_with_startup_control(&mcp_server).await?;
let release = control.hold_next_successful_initialize();
let mut config = fixture.config.clone();
let mut servers = config.mcp_servers.get().clone();
servers.insert(
"stalled".to_owned(),
serde_json::from_value(json!({
"url": format!("{}/api/codex/ps/mcp", http_server.chatgpt_base_url),
"http_headers": { "Authorization": "Bearer synthetic-test-token" },
"required": true,
"startup_timeout_sec": 120,
}))?,
);
config.mcp_servers.set(servers)?;
let thread_id = fixture.thread_manager.reserve_thread_id();
let mut options = StartThreadOptions::new(config);
options.reserved_thread_id = Some(thread_id);
options.environments = Some(fixture.codex.environment_selections().await);
options
.thread_extension_init
.insert(SessionIsolation::Isolated);
let tasks = TaskTracker::new();
let mut start = Box::pin(fixture.thread_manager.start_thread_until(
options,
std::future::pending(),
&tasks,
));
tokio::select! {
_ = &mut start => panic!("startup must wait for the required MCP server"),
result = tokio::time::timeout(Duration::from_secs(10), async {
while control.initialize_attempts() == 0 {
tokio::task::yield_now().await;
}
}) => result.context("required MCP initialization did not begin")?,
}
// Persistence is already open while initialization is blocked on the server.
fixture.thread_store.flush_thread(thread_id).await?;
drop(start);
tasks.close();
tokio::time::timeout(Duration::from_secs(10), tasks.wait())
.await
.context("dropping startup must finish cleanup without releasing MCP initialization")?;
assert_eq!(
fixture.thread_manager.list_thread_ids().await,
vec![fixture.session_configured.thread_id],
);
assert!(matches!(
fixture.thread_store.flush_thread(thread_id).await,
Err(ThreadStoreError::ThreadNotFound { .. })
));
drop(release);
fixture.codex.shutdown_and_wait().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn dropping_startup_before_receiving_the_agent_finishes_cleanup() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
// Startup first suspends on its result channel. Observe that channel's wake without
// polling again, so the result stays sent but unconsumed until we drop the caller.
struct ResultReady(tokio::sync::Notify);
impl Wake for ResultReady {
fn wake(self: Arc<Self>) {
self.0.notify_one();
}
}
let server = responses::start_mock_server().await;
let fixture = test_codex().build_with_auto_env(&server).await?;
let mut options = StartThreadOptions::new(fixture.config.clone());
let thread_id = fixture.thread_manager.reserve_thread_id();
options.reserved_thread_id = Some(thread_id);
options.environments = Some(fixture.codex.environment_selections().await);
options
.thread_extension_init
.insert(SessionIsolation::Isolated);
let tasks = TaskTracker::new();
let mut start = Box::pin(fixture.thread_manager.start_thread_until(
options,
std::future::pending(),
&tasks,
));
let ready = Arc::new(ResultReady(tokio::sync::Notify::new()));
let waker = Arc::clone(&ready).into();
assert!(
start
.as_mut()
.poll(&mut TaskContext::from_waker(&waker))
.is_pending()
);
tokio::time::timeout(Duration::from_secs(10), ready.0.notified()).await?;
fixture.thread_manager.get_thread(thread_id).await?;
drop(start);
tasks.close();
tokio::time::timeout(Duration::from_secs(10), tasks.wait())
.await
.context("dropping an unconsumed result must also finish cleanup")?;
assert_eq!(
fixture.thread_manager.list_thread_ids().await,
vec![fixture.session_configured.thread_id]
);
assert!(matches!(
fixture.thread_store.flush_thread(thread_id).await,
Err(ThreadStoreError::ThreadNotFound { .. })
));
fixture.codex.shutdown_and_wait().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn owner_cancellation_closes_agent_and_preserves_history_and_parent() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let fixture = test_codex().build_with_auto_env(&server).await?;
let mut options = StartThreadOptions::new(fixture.config.clone());
options.environments = Some(fixture.codex.environment_selections().await);
options
.thread_extension_init
.insert(SessionIsolation::Isolated);
let cancelled = CancellationToken::new();
let tasks = TaskTracker::new();
let agent = fixture
.thread_manager
.start_thread_until(options, cancelled.clone().cancelled_owned(), &tasks)
.await?;
// A failed startup with the same ID must leave the original agent's writer intact.
let mut duplicate = StartThreadOptions::new(fixture.config.clone());
duplicate.reserved_thread_id = Some(agent.thread_id);
duplicate.environments = Some(fixture.codex.environment_selections().await);
duplicate
.thread_extension_init
.insert(SessionIsolation::Isolated);
let failed_tasks = TaskTracker::new();
assert!(
fixture
.thread_manager
.start_thread_until(duplicate, std::future::pending(), &failed_tasks,)
.await
.is_err()
);
failed_tasks.close();
tokio::time::timeout(Duration::from_secs(10), failed_tasks.wait()).await?;
fixture.thread_store.flush_thread(agent.thread_id).await?;
let response = responses::mount_sse_once(
&server,
responses::sse(vec![
responses::ev_response_created("agent-response"),
responses::ev_assistant_message("agent-message", "agent finished"),
responses::ev_completed("agent-response"),
]),
)
.await;
agent
.thread
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: "run the owned agent".to_owned(),
text_elements: Vec::new(),
}]))
.await?;
wait_for_event(&agent.thread, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
assert!(
response
.single_request()
.body_contains_text("run the owned agent")
);
cancelled.cancel();
tasks.close();
tokio::time::timeout(Duration::from_secs(10), tasks.wait())
.await
.context("the owner must be able to join shutdown and deregistration")?;
tokio::time::timeout(Duration::from_secs(1), agent.thread.wait_until_terminated()).await?;
assert_eq!(
fixture.thread_manager.list_thread_ids().await,
vec![fixture.session_configured.thread_id],
);
assert!(matches!(
fixture.thread_store.flush_thread(agent.thread_id).await,
Err(ThreadStoreError::ThreadNotFound { .. })
));
let history = agent
.thread
.load_history(/*include_archived*/ false)
.await?;
assert!(history.items.iter().any(|item| matches!(
item,
codex_history::RolloutItem::EventMsg(EventMsg::TurnComplete(_))
)));
let parent_response = responses::mount_sse_once(
&server,
responses::sse(vec![
responses::ev_response_created("parent-response"),
responses::ev_completed("parent-response"),
]),
)
.await;
fixture.submit_turn("parent still works").await?;
assert!(
parent_response
.single_request()
.body_contains_text("parent still works")
);
fixture.codex.shutdown_and_wait().await?;
Ok(())
}
+2
View File
@@ -99,6 +99,8 @@ mod interrupt_hooks;
mod items;
mod json_result;
mod live_cli;
#[path = "managed_threads_tests.rs"]
mod managed_threads;
mod mcp_auth_elicitation;
mod mcp_auth_refresh;
mod mcp_ema_config;