Add unfinished root turn suspension (#40038)

## Why

An active root turn needs to be stopped without marking it complete or aborted
before another runtime can recover the same turn ID.

## What changed

- Add `CodexThread::suspend_turn_and_shutdown` and `SuspendTurnOutcome`.
- Flush history, stop the active regular task, close the history writer, and
  shut down the session without recording a terminal turn event.
- Reject suspension when no supported turn is active or the loaded agent
  subtree still contains a live descendant.

## Testing

- Verify that suspension preserves unfinished history and allows the turn to be
  recovered under its original ID after the descendant guard is cleared.

GitOrigin-RevId: 1f9b019d07c51474ec2d991d263bc15cdd4f89ad
This commit is contained in:
richardopenai
2026-08-22 00:42:19 +00:00
committed by copyberry
parent 970b7f2ff4
commit 4f39251a01
10 changed files with 345 additions and 3 deletions
+1
View File
@@ -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;
+46
View File
@@ -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<SuspendTurnOutcome> {
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,
+1
View File
@@ -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;
+15 -2
View File
@@ -394,7 +394,7 @@ pub async fn set_thread_memory_mode(sess: &Arc<Session>, sub_id: String, mode: T
}
}
async fn shutdown_session_runtime(sess: &Arc<Session>) {
pub(super) async fn shutdown_session_runtime(sess: &Arc<Session>) {
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<Session>) {
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
+1
View File
@@ -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)]
@@ -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<Session>,
submission_id: String,
) -> CodexResult<SuspendTurnOutcome> {
{
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 })
}
+1 -1
View File
@@ -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");
+141
View File
@@ -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::<RolloutLine>(line)
.expect("parse durable rollout")
.item
})
.collect::<Vec<_>>();
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
+7
View File
@@ -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<CodexResult<TurnInputSubmission>>,
},
/// Stop the active root turn without recording a terminal turn event.
SuspendTurnAndShutdown {
reply: oneshot::Sender<CodexResult<SuspendTurnOutcome>>,
},
/// 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",
+12
View File
@@ -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 {