diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index fef0d7a1ea..566c2bc62f 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -46,6 +46,7 @@ pub use codex_core::StartIfIdleSubmission; pub use codex_core::StartThreadOptions; pub use codex_core::StateDbHandle; pub use codex_core::SteerSubmission; +pub use codex_core::SuspendTurnOutcome; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; pub use codex_core::TurnInput; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 74bc996113..59723d818b 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -3,6 +3,7 @@ use crate::config::ConstraintResult; use crate::elicitation::ElicitationRegistration; use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; +use crate::session::new_submission_id; use crate::session::session::Session; use codex_diagnostics::Gauge; use codex_diagnostics::GaugeGuard; @@ -12,6 +13,7 @@ use codex_extension_api::ThreadIdleCause; use codex_features::Feature; use codex_history::RolloutItem; use codex_otel::SessionTelemetry; +use codex_otel::current_span_w3c_trace_context; use codex_protocol::ThreadId; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::CollaborationMode; @@ -36,6 +38,7 @@ use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionConfiguredEvent; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::ThreadSettingsSnapshot; @@ -47,6 +50,7 @@ use codex_protocol::protocol::W3cTraceContext; use codex_protocol::turn_input::RecoverTurnRequest; use codex_protocol::turn_input::StartIfIdleSubmission; use codex_protocol::turn_input::SteerSubmission; +use codex_protocol::turn_input::SuspendTurnOutcome; use codex_protocol::turn_input::TurnInputMode; use codex_protocol::turn_input::TurnInputRequest; use codex_protocol::turn_input::TurnInputSubmission; @@ -63,6 +67,7 @@ use rmcp::model::ReadResourceRequestParams; use std::path::PathBuf; use std::sync::Arc; use tokio::sync::Mutex; +use tokio::sync::oneshot; use tokio::sync::watch; use tokio_util::sync::CancellationToken; @@ -397,6 +402,47 @@ 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. + /// + /// 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. + pub async fn suspend_turn_and_shutdown(&self) -> CodexResult { + if self.session_source.is_non_root_agent() { + return Err(CodexErr::UnsupportedOperation( + "turn suspension requires the owning root thread".to_string(), + )); + } + + // The session owns accepted suspension, so dropping this caller cannot interrupt + // cancellation, persistence, or writer shutdown halfway through a handoff. + let (reply, result) = oneshot::channel(); + self.io + .tx_sub + .send(Submission { + id: new_submission_id(), + op: Op::SuspendTurnAndShutdown { reply }, + trace: current_span_w3c_trace_context(), + parent_turn_id: None, + root_turn_id: None, + }) + .await + .map_err(|_| CodexErr::Fatal("thread session has stopped".to_string()))?; + let outcome = result + .await + .map_err(|_| CodexErr::Fatal("thread suspension reply was lost".to_string()))??; + if matches!(&outcome, SuspendTurnOutcome::Suspended { .. }) { + self.io.session_loop_termination.clone().await; + } + Ok(outcome) + } + /// Steers only if `expected_turn_id` is still the active regular turn. pub async fn steer_turn( &self, diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 1970907324..f7c636a5b8 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -19,6 +19,7 @@ pub use codex_protocol::turn_input::NotSubmittedReason; pub use codex_protocol::turn_input::RecoverTurnRequest; pub use codex_protocol::turn_input::StartIfIdleSubmission; pub use codex_protocol::turn_input::SteerSubmission; +pub use codex_protocol::turn_input::SuspendTurnOutcome; pub use codex_protocol::turn_input::TurnInput; pub use codex_protocol::turn_input::TurnInputRequest; pub use codex_protocol::turn_input::TurnInputSubmission; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index a928090cc8..644bc84fad 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -394,7 +394,7 @@ pub async fn set_thread_memory_mode(sess: &Arc, sub_id: String, mode: T } } -async fn shutdown_session_runtime(sess: &Arc) { +pub(super) async fn shutdown_session_runtime(sess: &Arc) { if let Some(startup_prewarm) = sess.take_session_startup_prewarm().await { startup_prewarm.abort().await; } @@ -421,7 +421,7 @@ async fn shutdown_session_runtime(sess: &Arc) { crate::hook_runtime::run_session_end_hooks(sess).await; } -async fn emit_thread_stop_lifecycle(sess: &Session) { +pub(super) async fn emit_thread_stop_lifecycle(sess: &Session) { for contributor in sess.services.extensions.thread_lifecycle_contributors() { contributor .on_thread_stop(codex_extension_api::ThreadStopInput { @@ -585,6 +585,19 @@ pub(super) async fn submission_loop( let _ = reply.send(result); false } + Op::SuspendTurnAndShutdown { reply } => { + let result = + super::turn_suspension::suspend_turn_and_shutdown(&sess, sub.id.clone()) + .await; + // Exit only after history is durable and its writer has closed; an error + // must leave responsibility for the thread with the current worker. + let should_exit = matches!( + &result, + Ok(codex_protocol::turn_input::SuspendTurnOutcome::Suspended { .. }) + ); + let _ = reply.send(result); + should_exit + } Op::ThreadSettings { thread_settings } => { thread_settings::update(&sess, sub.id.clone(), thread_settings).await; false diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index fe9d7377df..50fe4b9df6 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -229,6 +229,7 @@ mod token_budget; pub(crate) mod turn; pub(crate) mod turn_context; mod turn_input; +mod turn_suspension; mod world_state; use self::code_mode_warning::unsupported_code_mode_warning; #[cfg(test)] diff --git a/codex-rs/core/src/session/turn_suspension.rs b/codex-rs/core/src/session/turn_suspension.rs new file mode 100644 index 0000000000..64a573b813 --- /dev/null +++ b/codex-rs/core/src/session/turn_suspension.rs @@ -0,0 +1,120 @@ +use super::handlers; +use super::session::Session; +use crate::state::TaskKind; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::turn_input::SuspendTurnOutcome; +use std::sync::Arc; +use std::time::Duration; +use tracing::warn; + +pub(super) async fn suspend_turn_and_shutdown( + session: &Arc, + submission_id: String, +) -> CodexResult { + { + let active = session.active_turn.lock().await; + let Some(task) = active.as_ref().and_then(|turn| turn.task.as_ref()) else { + return Ok(SuspendTurnOutcome::NotActive); + }; + if task.kind != TaskKind::Regular { + return Ok(SuspendTurnOutcome::UnsupportedTask); + } + } + + // 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 + .agent_control + .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()))?; + // Flush before canceling execution so a persistence failure leaves the original turn running. + live_thread.flush().await.map_err(|error| { + CodexErr::Fatal(format!("flush before root turn suspension failed: {error}")) + })?; + + // The flush can yield while the active turn completes or changes. Recheck its + // kind under the same lock used to remove it. + let mut turn = { + let mut active = session.active_turn.lock().await; + let Some(active_turn) = active.as_ref() else { + return Ok(SuspendTurnOutcome::NotActive); + }; + let Some(task) = active_turn.task.as_ref() else { + return Ok(SuspendTurnOutcome::NotActive); + }; + if task.kind != TaskKind::Regular { + return Ok(SuspendTurnOutcome::UnsupportedTask); + } + active.take().ok_or_else(|| { + CodexErr::Fatal("accepted root turn suspension had no running turn".to_string()) + })? + }; + + let task = turn.task.take().ok_or_else(|| { + CodexErr::Fatal("accepted root turn suspension had no running task".to_string()) + })?; + let turn_id = task.turn_context.sub_id.clone(); + // Normal shutdown records a terminal turn event, preventing another worker from + // recovering this turn under its original ID. Cancel the task without that event. + task.cancellation_token.cancel(); + task.turn_context + .turn_metadata_state + .cancel_git_enrichment_task(); + let mut task_handle = task.handle.detach(); + match tokio::time::timeout( + Duration::from_millis(crate::tasks::GRACEFULL_INTERRUPTION_TIMEOUT_MS), + &mut task_handle, + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(error)) => { + warn!(thread_id = %session.thread_id, %error, "suspended turn task exited abnormally"); + } + Err(_) => { + warn!( + thread_id = %session.thread_id, + "suspended turn task did not stop gracefully; aborting it" + ); + task_handle.abort(); + let _ = task_handle.await; + } + } + // Pending accepted input and interactive waiters live only in this process. Handoff + // intentionally drops that state; persisting or replaying it needs a separate protocol. + session.input_queue.clear_pending(&turn).await; + + // 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. + handlers::shutdown_session_runtime(session).await; + live_thread.flush().await.map_err(|error| { + CodexErr::Fatal(format!("flush after root turn suspension failed: {error}")) + })?; + live_thread.shutdown().await.map_err(|error| { + CodexErr::Fatal(format!("close suspended root turn writer failed: {error}")) + })?; + // Announce thread shutdown only after its writer closes so a replacement worker + // cannot write the same thread concurrently. + handlers::emit_thread_stop_lifecycle(session.as_ref()).await; + session + .deliver_event_raw(Event { + id: submission_id, + msg: EventMsg::ShutdownComplete, + }) + .await; + Ok(SuspendTurnOutcome::Suspended { turn_id }) +} diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index cc71acb3b5..08448dc7d6 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -63,7 +63,7 @@ pub(crate) use user_shell::UserShellCommandMode; pub(crate) use user_shell::UserShellCommandTask; pub(crate) use user_shell::execute_user_shell_command; -const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100; +pub(crate) const GRACEFULL_INTERRUPTION_TIMEOUT_MS: u64 = 100; const TASK_COMPACT_METRIC: &str = "codex.task.compact"; static ACTIVE_TURNS: Gauge = Gauge::new("core.turns.active"); diff --git a/codex-rs/core/tests/suite/abort_tasks.rs b/codex-rs/core/tests/suite/abort_tasks.rs index b648235f85..8d8bf98ca4 100644 --- a/codex-rs/core/tests/suite/abort_tasks.rs +++ b/codex-rs/core/tests/suite/abort_tasks.rs @@ -1,17 +1,25 @@ use assert_matches::assert_matches; +use codex_core::StartThreadOptions; +use codex_core::SuspendTurnOutcome; use codex_core::TurnInputRequest; +use codex_history::RolloutItem; +use codex_history::RolloutLine; 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_response_created; +use core_test_support::responses::mount_response_once; use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; +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; @@ -63,6 +71,139 @@ async fn interrupt_long_running_tool_emits_turn_aborted() { codex.submit(Op::CleanBackgroundTerminals).await.unwrap(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn root_turn_suspension_preserves_unfinished_turn_history() { + 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( + &server, + sse_response(sse(vec![ + ev_response_created("suspended_response"), + ev_completed("suspended_response"), + ])) + .set_delay(Duration::from_secs(60)), + ) + .await; + let test = test_codex() + .with_model("gpt-5.4") + .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 submitted = codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "preserve this exact unfinished turn".into(), + text_elements: Vec::new(), + }])) + .await + .expect("start root turn"); + 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() + .await + .expect("stop the descendant before root handoff"); + test.thread_manager + .remove_thread(&descendant.thread_id) + .await + .expect("remove the stopped descendant from the live thread inventory"); + + // 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() + .await + .expect("stop and close the old writer"), + SuspendTurnOutcome::Suspended { + turn_id: turn_id.clone(), + }, + ); + let rollout_path = codex.rollout_path().expect("rollout path"); + let rollout = tokio::fs::read_to_string(&rollout_path) + .await + .expect("read durable rollout"); + let items = rollout + .lines() + .map(|line| { + serde_json::from_str::(line) + .expect("parse durable rollout") + .item + }) + .collect::>(); + assert!(items.iter().all(|item| !matches!( + item, + RolloutItem::EventMsg(EventMsg::TurnAborted(_) | EventMsg::TurnComplete(_)) + ))); + test.thread_manager + .remove_thread(&test.session_configured.thread_id) + .await + .expect("unload the suspended root"); + let recovery_server = start_mock_server().await; + mount_sse_once( + &recovery_server, + sse(vec![ + ev_response_created("recovered_response"), + ev_completed("recovered_response"), + ]), + ) + .await; + let resumed = test_codex() + .with_model("gpt-5.4") + .resume(&recovery_server, Arc::clone(&test.home), rollout_path) + .await + .expect("resume the suspended root on a replacement runtime"); + + assert_eq!( + resumed + .codex + .recover_turn_if_idle(codex_core::RecoverTurnRequest { + turn_id: turn_id.clone(), + thread_settings: Default::default(), + trace: None, + }) + .await + .expect("recover the unfinished turn"), + codex_core::StartIfIdleSubmission::Started { + turn_id: turn_id.clone(), + }, + ); + let completed = wait_for_event(&resumed.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let EventMsg::TurnComplete(completed) = completed else { + unreachable!("wait_for_event returned unexpected event"); + }; + assert_eq!(completed.turn_id, turn_id); +} + /// After an interrupt we expect the next request to the model to include both /// the original tool call and an `"aborted"` `function_call_output`. This test /// exercises the follow-up flow: it sends another user turn, inspects the mock diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 9b4ad5650f..e88eca7d41 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -56,6 +56,7 @@ use crate::plan_tool::UpdatePlanArgs; use crate::request_permissions::RequestPermissionsEvent; use crate::request_permissions::RequestPermissionsResponse; use crate::request_user_input::RequestUserInputResponse; +use crate::turn_input::SuspendTurnOutcome; use crate::turn_input::TurnInputMode; use crate::turn_input::TurnInputRequest; use crate::turn_input::TurnInputSubmission; @@ -581,6 +582,11 @@ pub enum Op { reply: oneshot::Sender>, }, + /// Stop the active root turn without recording a terminal turn event. + SuspendTurnAndShutdown { + reply: oneshot::Sender>, + }, + /// Apply persistent thread-settings overrides without starting a turn. /// /// This uses the same submission queue as turn starts so app-server can @@ -875,6 +881,7 @@ impl Op { Self::RealtimeConversationListVoices => "realtime_conversation_list_voices", Self::TurnInput { .. } => "turn_input", Self::RecoverTurn { .. } => "recover_turn", + Self::SuspendTurnAndShutdown { .. } => "suspend_turn_and_shutdown", Self::ThreadSettings { .. } => "thread_settings", Self::InterAgentCommunication { .. } => "inter_agent_communication", Self::ExecApproval { .. } => "exec_approval", diff --git a/codex-rs/protocol/src/turn_input.rs b/codex-rs/protocol/src/turn_input.rs index bc5adbe8b5..e7b58a957a 100644 --- a/codex-rs/protocol/src/turn_input.rs +++ b/codex-rs/protocol/src/turn_input.rs @@ -13,6 +13,18 @@ use serde_json::Value; use std::collections::BTreeMap; use std::collections::HashMap; +/// 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, + }, + NotActive, + /// A currently loaded descendant would remain running after root handoff. + HasLiveDescendants, + UnsupportedTask, +} + /// Input consumed by a regular turn. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub enum TurnInput {