Add coordinated shutdown for local agent trees (#49814)

## What changed

Expose `ThreadManager::request_agent_tree_shutdown` and an `AgentTreeShutdown::wait` handle. Requests reject new starts in the shared runtime, cancel pending startup, and signal every session in that tree, including delegates outside the thread registry. Waiting covers admitted operations and session teardown, and reports cleanup or persistence writer failures. Other root trees remain unaffected; registry removal and persistence handoff remain caller responsibilities.

Retain tree membership through startup, cancelled spawn cleanup, residency eviction, and completion watchers. Release persistence resources before returning startup failures, and stop completion watchers during tree shutdown without reporting shutdown as a child result.

## Testing

Add coverage for tree isolation, children removed from the registry, delegate teardown, cancellation during required MCP initialization, aborted teardown, persistence writer failures, and successful retry after a failed resume releases its writer.

GitOrigin-RevId: a109e25adf23d20baf7595891099947050c6b6b5
This commit is contained in:
Owen Lin
2026-10-01 01:57:56 +00:00
committed by copyberry
parent d4a475adda
commit 8e44ad94b0
22 changed files with 795 additions and 51 deletions
+25 -5
View File
@@ -27,6 +27,8 @@ use crate::session::SessionIo;
use crate::session::SessionSpawnArgs;
use crate::session::emit_subagent_session_started;
use crate::session::session::Session;
use crate::session::startup::SessionStartup;
use crate::session::startup::SessionStartupGuard;
use crate::session::turn_context::TurnContext;
use codex_history::InitialHistory;
use codex_login::AuthManager;
@@ -73,6 +75,10 @@ pub(crate) async fn run_codex_thread_interactive(
let conversation_history = initial_history.unwrap_or(InitialHistory::New);
let forked_from_thread_id = conversation_history.forked_from_id();
let runtime = parent_session.services.local_agent_runtime.clone();
let startup = Arc::new(SessionStartup::default());
startup.hold_membership(runtime.admit_start()?);
let startup_guard = SessionStartupGuard::new(Arc::clone(&startup));
let instructions = parent_session.inherited_instructions().await;
let session_source = SessionSource::SubAgent(subagent_source.clone());
let is_guardian_reviewer = crate::guardian::is_basic_session_source(&session_source);
@@ -83,8 +89,8 @@ pub(crate) async fn run_codex_thread_interactive(
};
let mut thread_extension_init = codex_extension_api::ExtensionDataInit::default();
thread_extension_init.insert(isolation);
let (session, io) = Session::spawn(SessionSpawnArgs {
startup: None,
let spawn_result = Session::spawn(SessionSpawnArgs {
startup: Some(Arc::clone(&startup)),
config,
allow_provider_model_fallback: false,
instructions,
@@ -116,7 +122,7 @@ pub(crate) async fn run_codex_thread_interactive(
originator: parent_ctx.originator.clone(),
agent_control: crate::agent::control::AgentControlInit::Provided {
control: Arc::clone(&parent_session.services.agent_control),
runtime: parent_session.services.local_agent_runtime.clone(),
runtime,
},
dynamic_tools: Vec::new(),
metrics_service_name: None,
@@ -139,7 +145,18 @@ pub(crate) async fn run_codex_thread_interactive(
windows_sandbox_proxy_settings_mode,
})
.or_cancel(&cancel_token)
.await??;
.await;
let (session, io) = match spawn_result {
Ok(Ok(spawned)) => spawned,
Ok(Err(error)) => {
startup_guard.cleanup().await;
return Err(error);
}
Err(error) => {
startup_guard.cleanup().await;
return Err(error.into());
}
};
let thread_config = session.thread_config_snapshot().await;
let client_metadata = parent_session.app_server_client_metadata().await;
emit_subagent_session_started(
@@ -151,7 +168,10 @@ pub(crate) async fn run_codex_thread_interactive(
thread_config,
subagent_source,
);
Ok((session, forward_session_io(Arc::new(io), cancel_token)))
let caller_io = forward_session_io(Arc::new(io), cancel_token);
startup.release_membership();
startup_guard.disarm();
Ok((session, caller_io))
}
/// Keeps delegate IO cancellation identical for standalone and manager-owned reviewers.