diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index f52b464936..0dff26d89a 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/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 { if self.session_source.is_non_root_agent() { return Err(CodexErr::UnsupportedOperation( diff --git a/codex-rs/core/src/session/turn_suspension.rs b/codex-rs/core/src/session/turn_suspension.rs index 64e40d6480..324099bd9e 100644 --- a/codex-rs/core/src/session/turn_suspension.rs +++ b/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. diff --git a/codex-rs/core/tests/suite/abort_tasks.rs b/codex-rs/core/tests/suite/abort_tasks.rs index 19b39e5c9e..ae4a17ad0f 100644 --- a/codex-rs/core/tests/suite/abort_tasks.rs +++ b/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 diff --git a/codex-rs/protocol/src/turn_input.rs b/codex-rs/protocol/src/turn_input.rs index 687e0c73d9..c1998070ea 100644 --- a/codex-rs/protocol/src/turn_input.rs +++ b/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, }