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
@@ -0,0 +1,46 @@
//! Fences and signals one local agent tree, then exposes its teardown result.
use super::ThreadManager;
use crate::agent::control::AgentTreeShutdownState;
use codex_protocol::ThreadId;
use codex_protocol::error::Result as CodexResult;
use std::sync::Arc;
/// Completion handle for one exact local agent tree.
#[derive(Clone, Debug)]
#[must_use = "agent-tree shutdown is not complete until wait() finishes"]
pub struct AgentTreeShutdown {
state: Arc<AgentTreeShutdownState>,
}
impl AgentTreeShutdown {
/// Waits for every operation and session admitted before the shutdown fence to finish.
/// Returns an error if cleanup or persistence writer termination failed. Registry removal and
/// caller-owned persistence handoff remain with the caller. Dropping this future does not
/// cancel shutdown.
pub async fn wait(&self) -> CodexResult<()> {
self.state.wait().await
}
}
impl ThreadManager {
/// Requests shutdown of a loaded thread and every session sharing its local runtime.
///
/// Returns after fencing new starts and signalling admitted sessions. Callers can wait on the
/// returned handle before completing registry cleanup or a persistence handoff. This fences
/// only starts sharing the current runtime; callers must separately serialize top-level loads
/// or resumes of the same thread ID through that handoff.
pub async fn request_agent_tree_shutdown(
&self,
thread_id: ThreadId,
) -> CodexResult<AgentTreeShutdown> {
let thread = self.get_thread(thread_id).await?;
Ok(AgentTreeShutdown {
state: thread
.session
.services
.local_agent_runtime
.request_shutdown(),
})
}
}