Allow root turn suspension with live descendants (#54414)

## Why

Live descendants previously blocked suspension of an unfinished root turn, requiring callers to stop and unload them before recovering the root elsewhere.

## What changed

Allow `suspend_turn_and_shutdown` to suspend the root and initiate shutdown of all sessions sharing its local agent runtime, fencing new starts. Remove `SuspendTurnOutcome::HasLiveDescendants`.

Success confirms that the root has stopped and its history writer is flushed and closed; descendant cleanup completes separately. Callers must finish tree shutdown or fence and terminate the tree's persistence writers before transferring ownership. Interrupted descendants are not automatically recovered.

## Testing

Extend integration coverage to verify child shutdown, recovery under the original root turn ID without restarting children, and separate reporting of child writer failures.

GitOrigin-RevId: 5adc240ae6e34654209154ee9324621a2bd604d5

codex-rs/core/src/codex_thread.rs

@@ -521,14 +521,15 @@ impl CodexThread {
     /// Stops the active unfinished root turn without recording TurnAborted or
     /// TurnComplete, so another worker can recover its original turn ID.
     ///
-    /// Suspension is refused while a currently loaded descendant exists. Past
-    /// descendants do not prevent recovery, and concurrent descendant admission
-    /// is not sealed. Queued user input and outstanding approval, elicitation,
-    /// or server-request waiters remain best effort and may be discarded.
+    /// After stopping the root turn, fences new starts and initiates shutdown of all sessions
+    /// sharing its local agent runtime. Descendant turns are interrupted and are not automatically
+    /// recovered. Queued user input and outstanding approval, elicitation, or server-request
+    /// waiters remain best effort and may be discarded.
     ///
     /// The session processes an accepted request even if its caller disconnects.
-    /// Callers must not transfer ownership until suspension succeeds, which
-    /// requires stopping execution, flushing history, and closing its writer.
+    /// Success confirms the root's execution has stopped, its history is flushed, and its writer
+    /// is closed. It does not wait for descendant cleanup. Before transferring ownership,
+    /// callers must also finish tree shutdown or fence and terminate the tree's persistence writers.
     pub async fn suspend_turn_and_shutdown(&self) -> CodexResult<SuspendTurnOutcome> {
         if self.session_source.is_non_root_agent() {
             return Err(CodexErr::UnsupportedOperation(

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

@@ -24,19 +24,6 @@ pub(super) async fn suspend_turn_and_shutdown(
         }
     }
 
-    // This is a snapshot of currently loaded descendants, not a spawn-admission seal.
-    // Previously closed descendants and concurrent future spawns remain best effort.
-    if session
-        .services
-        .local_agent_runtime
-        .list_live_agent_subtree_thread_ids(session.thread_id)
-        .await?
-        .len()
-        > 1
-    {
-        return Ok(SuspendTurnOutcome::HasLiveDescendants);
-    }
-
     let live_thread = session
         .live_thread_for_persistence("suspend an unfinished root turn")
         .map_err(|error| CodexErr::Fatal(error.to_string()))?;
@@ -97,6 +84,11 @@ pub(super) async fn suspend_turn_and_shutdown(
     // intentionally drops that state; persisting or replaying it needs a separate protocol.
     session.input_queue.clear_pending(&turn).await;
 
+    // The root task has stopped without a terminal event. Fence new starts and stop every
+    // other local session, even if the suspension caller disconnects during cleanup.
+    // The host owns the separate tree-cleanup wait and any persistence-writer fencing.
+    session.services.local_agent_runtime.request_shutdown();
+
     // Stop all producers before flushing their final history and closing its writer.
     // If either persistence step fails, do not report success: the current worker
     // retains ownership until worker-failure recovery can take responsibility.

codex-rs/core/tests/suite/abort_tasks.rs

@@ -1,18 +1,17 @@
 use assert_matches::assert_matches;
-use codex_core::StartThreadOptions;
 use codex_core::SuspendTurnOutcome;
 use codex_core::TurnInputRequest;
+use codex_features::Feature;
 use codex_history::RolloutItem;
 use std::sync::Arc;
 use std::time::Duration;
 
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::Op;
-use codex_protocol::protocol::SessionSource;
-use codex_protocol::protocol::SubAgentSource;
 use codex_protocol::user_input::UserInput;
 use core_test_support::responses::ev_completed;
 use core_test_support::responses::ev_function_call;
+use core_test_support::responses::ev_function_call_with_namespace;
 use core_test_support::responses::ev_response_created;
 use core_test_support::responses::mount_response_once;
 use core_test_support::responses::mount_sse_once;
@@ -22,8 +21,10 @@ use core_test_support::responses::sse_response;
 use core_test_support::responses::start_mock_server;
 use core_test_support::test_codex::test_codex;
 use core_test_support::wait_for_event;
+use pretty_assertions::assert_eq;
 use regex_lite::Regex;
 use serde_json::json;
+use test_case::test_case;
 
 /// Integration test: spawn a long‑running exec_command tool via a mocked Responses SSE
 /// function call, then interrupt the session and expect TurnAborted.
@@ -70,40 +71,62 @@ async fn interrupt_long_running_tool_emits_turn_aborted() {
     codex.submit(Op::CleanBackgroundTerminals).await.unwrap();
 }
 
+/// Suspending a root preserves its unfinished turn and initiates child shutdown. The host
+/// separately waits for tree cleanup, including any child writer failure, before recovery.
+#[test_case(false; "recovers root after stopping child")]
+#[test_case(true; "reports child writer failure separately")]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn root_turn_suspension_preserves_unfinished_turn_history() {
+async fn root_turn_suspension_preserves_unfinished_turn_history(fail_child_writer: bool) {
     let server = start_mock_server().await;
-    // Waiting on the mocked response keeps the turn active on local and remote executors
-    // without requiring an OS-specific command or a working sandboxed child process.
-    mount_response_once(
+    mount_sse_once(
         &server,
-        sse_response(sse(vec![
-            ev_response_created("suspended_response"),
-            ev_completed("suspended_response"),
-        ]))
-        .set_delay(Duration::from_secs(60)),
+        sse(vec![
+            ev_response_created("spawn_response"),
+            ev_function_call_with_namespace(
+                "spawn_worker",
+                "collaboration",
+                "spawn_agent",
+                &json!({
+                    "message": "keep working until shutdown",
+                    "task_name": "worker",
+                    "fork_turns": "none",
+                })
+                .to_string(),
+            ),
+            ev_completed("spawn_response"),
+        ]),
     )
     .await;
+    // Waiting on the mocked response keeps the turn active on local and remote executors
+    // without requiring an OS-specific command or a working sandboxed child process.
+    for response_id in ["pending_child", "pending_root"] {
+        mount_response_once(
+            &server,
+            sse_response(sse(vec![
+                ev_response_created(response_id),
+                ev_completed(response_id),
+            ]))
+            .set_delay(Duration::from_secs(60)),
+        )
+        .await;
+    }
     let test = test_codex()
-        .with_model("gpt-5.4")
+        .with_model("gpt-5.6-sol")
+        .with_config(|config| {
+            config
+                .features
+                .enable(Feature::Collab)
+                .expect("enable collaboration for the suspension test");
+            config
+                .features
+                .enable(Feature::MultiAgentV2)
+                .expect("enable multi-agent v2 for the suspension test");
+        })
         .build_with_auto_env(&server)
         .await
         .expect("start persistent root thread");
     let codex = Arc::clone(&test.codex);
-    let descendant = test
-        .thread_manager
-        .start_thread(StartThreadOptions {
-            session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
-                parent_thread_id: test.session_configured.thread_id,
-                depth: 1,
-                agent_path: None,
-                agent_nickname: None,
-                agent_role: None,
-            })),
-            ..StartThreadOptions::new(test.config.clone())
-        })
-        .await
-        .expect("start a currently loaded descendant");
+    let mut created_threads = test.thread_manager.subscribe_thread_created();
     let submitted = codex
         .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
             text: "preserve this exact unfinished turn".into(),
@@ -114,36 +137,71 @@ async fn root_turn_suspension_preserves_unfinished_turn_history() {
     let codex_core::TurnInputSubmission::Started { turn_id, .. } = submitted else {
         panic!("expected a started root turn");
     };
-    wait_for_event(&codex, |event| matches!(event, EventMsg::TurnStarted(_))).await;
-
-    assert_eq!(
-        codex
-            .suspend_turn_and_shutdown()
-            .await
-            .expect("reject handoff while a descendant remains loaded"),
-        SuspendTurnOutcome::HasLiveDescendants,
-    );
-    descendant
-        .thread
-        .shutdown_and_wait()
+    let child_id = tokio::time::timeout(Duration::from_secs(10), created_threads.recv())
         .await
-        .expect("stop the descendant before root handoff");
-    test.thread_manager
-        .remove_thread(&descendant.thread_id)
+        .expect("child creation notification must arrive")
+        .expect("spawned child");
+    let descendant = test
+        .thread_manager
+        .get_thread(child_id)
         .await
-        .expect("remove the stopped descendant from the live thread inventory");
+        .expect("loaded child");
+    wait_for_event(&descendant, |event| {
+        matches!(event, EventMsg::TurnStarted(_))
+    })
+    .await;
 
-    // A previously admitted descendant no longer blocks handoff once it is stopped
-    // and removed; suspension only consults the current live subtree.
-    assert_eq!(
-        codex
-            .suspend_turn_and_shutdown()
+    if fail_child_writer {
+        test.thread_store
+            .shutdown_thread(child_id)
             .await
-            .expect("stop and close the old writer"),
+            .expect("close the child writer before shutdown to inject a persistence failure");
+    }
+    let suspension =
+        tokio::time::timeout(Duration::from_secs(10), codex.suspend_turn_and_shutdown())
+            .await
+            .expect("suspension must not wait for the child turn to finish naturally");
+    assert_eq!(
+        suspension.expect("suspend the root independently of child cleanup"),
         SuspendTurnOutcome::Suspended {
             turn_id: turn_id.clone(),
         },
     );
+    // Observe child termination before requesting the completion handle, so the test proves
+    // that root suspension itself initiated child shutdown.
+    tokio::time::timeout(Duration::from_secs(10), descendant.wait_until_terminated())
+        .await
+        .expect("root suspension must initiate child shutdown");
+    let shutdown = test
+        .thread_manager
+        .request_agent_tree_shutdown(test.session_configured.thread_id)
+        .await
+        .expect("get the tree cleanup completion handle");
+    let shutdown_result = tokio::time::timeout(Duration::from_secs(10), shutdown.wait_detailed())
+        .await
+        .expect("tree cleanup must finish");
+    if fail_child_writer {
+        let report =
+            shutdown_result.expect_err("tree cleanup must report the child writer failure");
+        assert!(
+            report
+                .failures
+                .iter()
+                .any(|failure| failure.thread_id == Some(child_id))
+        );
+        return;
+    }
+    shutdown_result.expect("finish tree cleanup before handing off the root");
+    let child_rollout =
+        tokio::fs::read_to_string(descendant.rollout_path().expect("child rollout"))
+            .await
+            .expect("read closed child writer");
+    assert!(child_rollout.lines().any(|line| matches!(
+        codex_rollout::parse_rollout_line(line)
+            .expect("parse child rollout")
+            .item,
+        RolloutItem::EventMsg(EventMsg::TurnAborted(_))
+    )));
     let rollout_path = codex.rollout_path().expect("rollout path");
     let rollout = tokio::fs::read_to_string(&rollout_path)
         .await
@@ -203,6 +261,11 @@ async fn root_turn_suspension_preserves_unfinished_turn_history() {
         unreachable!("wait_for_event returned unexpected event");
     };
     assert_eq!(completed.turn_id, turn_id);
+    assert_eq!(
+        resumed.thread_manager.list_thread_ids().await,
+        vec![resumed.session_configured.thread_id],
+        "root recovery must not automatically restart the interrupted child"
+    );
 }
 
 /// After an interrupt we expect the next request to the model to include both

codex-rs/protocol/src/turn_input.rs

@@ -19,12 +19,8 @@ use ts_rs::TS;
 /// Result of stopping an unfinished root turn so another worker can recover it.
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum SuspendTurnOutcome {
-    Suspended {
-        turn_id: String,
-    },
+    Suspended { turn_id: String },
     NotActive,
-    /// A currently loaded descendant would remain running after root handoff.
-    HasLiveDescendants,
     UnsupportedTask,
 }