From 7c098d87418a357bafb241a579a18938801989e0 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Tue, 8 Sep 2026 23:29:47 +0000 Subject: [PATCH] Gate new turn submissions on host shutdown admission (#43943) ## Why Hosts need to stop new turn-input work during shutdown without consuming pending input or preventing already-running delegated work from finishing. ## What changed - Add an optional `TurnStartAdmission` extension gate, checked before reserving or starting a new turn. Hosts without a gate retain existing behavior. - Return `NotSubmittedReason::ServerDraining` for refused starts and surface an app-server error instructing clients to reconnect and retry. - Keep steering, parent-delegated subagent input, and memory-only mailbox wakeups available during drain, while gating automatic starts. - Close realtime conversations with an ordered handoff, error, and close event sequence when a handoff is refused during drain. ## Testing Add regression coverage for rejected input staying out of subsequent requests, persisted queue items remaining available for later starts, delegated agent and review work completing during drain, mailbox wakeups, and realtime handoff error ordering. GitOrigin-RevId: 03dbdcd71eab200e597c0649e6eb39bd92dbc82f --- codex-rs/app-server/src/error_code.rs | 4 + .../thread_queue_processor.rs | 3 + .../src/request_processors/turn_processor.rs | 9 +- codex-rs/core/src/realtime_conversation.rs | 29 +- codex-rs/core/src/session/turn_input.rs | 59 +++- .../core/tests/suite/turn_input_submission.rs | 321 ++++++++++++++++++ codex-rs/ext/extension-api/src/lib.rs | 2 + codex-rs/ext/extension-api/src/registry.rs | 17 + .../ext/extension-api/src/turn_admission.rs | 12 + codex-rs/ext/queue/tests/queue_service.rs | 59 ++++ codex-rs/protocol/src/turn_input.rs | 3 + 11 files changed, 510 insertions(+), 8 deletions(-) create mode 100644 codex-rs/ext/extension-api/src/turn_admission.rs diff --git a/codex-rs/app-server/src/error_code.rs b/codex-rs/app-server/src/error_code.rs index 48e401f7bc..b5b44c0274 100644 --- a/codex-rs/app-server/src/error_code.rs +++ b/codex-rs/app-server/src/error_code.rs @@ -7,6 +7,10 @@ pub(crate) const INTERNAL_ERROR_CODE: i64 = -32603; pub(crate) const OVERLOADED_ERROR_CODE: i64 = -32001; pub const INPUT_TOO_LARGE_ERROR_CODE: &str = "input_too_large"; +pub(crate) fn server_draining_error() -> JSONRPCErrorError { + invalid_request("Server is draining; retry after reconnecting") +} + pub(crate) fn invalid_request(message: impl Into) -> JSONRPCErrorError { error(INVALID_REQUEST_ERROR_CODE, message) } diff --git a/codex-rs/app-server/src/request_processors/thread_queue_processor.rs b/codex-rs/app-server/src/request_processors/thread_queue_processor.rs index 6fe5336dc8..a561272e3f 100644 --- a/codex-rs/app-server/src/request_processors/thread_queue_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_queue_processor.rs @@ -198,6 +198,9 @@ impl ThreadQueueRequestProcessor { .map_err(queue_error)?; let turn_id = match submission { StartIfIdleSubmission::Started { turn_id } => turn_id, + StartIfIdleSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining, + } => return Err(crate::error_code::server_draining_error()), StartIfIdleSubmission::NotSubmitted { reason: NotSubmittedReason::NotIdle | NotSubmittedReason::PendingTriggerTurn, } => { diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index 2beaa87033..b939a0f3f4 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -667,7 +667,11 @@ impl TurnRequestProcessor { TurnInputSubmission::Started { turn_id } => (turn_id, true), TurnInputSubmission::Steered { turn_id } => (turn_id, false), TurnInputSubmission::NotSubmitted { reason } => { - let error = internal_error(format!("failed to submit turn input: {reason:?}")); + let error = if reason == NotSubmittedReason::ServerDraining { + crate::error_code::server_draining_error() + } else { + internal_error(format!("failed to submit turn input: {reason:?}")) + }; self.track_error_response(&request_id, &error, /*error_type*/ None); return Err(error); } @@ -1067,6 +1071,9 @@ impl TurnRequestProcessor { SteerSubmission::Steered { turn_id } => turn_id, SteerSubmission::NotSubmitted { reason } => { let (message, data, error_type) = match reason { + NotSubmittedReason::ServerDraining => { + return Err(crate::error_code::server_draining_error()); + } NotSubmittedReason::NoActiveTurn | NotSubmittedReason::NotIdle => ( "no active turn to steer".to_string(), None, diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 63f59d5408..5ef019d75e 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -144,13 +144,14 @@ impl RealtimeHandoffAdmission { } } - async fn route(&self, session: &Arc, text: String) { + async fn route(&self, session: &Arc, text: String) -> Result<(), &'static str> { let Ok(_permit) = self.gate.acquire().await else { - return; + return Ok(()); }; if !self.retired.load(Ordering::Acquire) { - session.route_realtime_text_input(text).await; + session.route_realtime_text_input(text).await?; } + Ok(()) } async fn retire(&self) { @@ -1645,6 +1646,7 @@ async fn handle_start_inner( msg, }; let mut end = RealtimeConversationEnd::TransportClosed; + let mut handoff_error = None; // Drain already-parsed events so a queued handoff is routed before the final tail. while let Ok(event) = events_rx.recv().await { match &event { @@ -1667,7 +1669,7 @@ async fn handle_start_inner( if let Some(text) = maybe_routed_text { // The routed text can contain spoken prompts or workspace secrets. debug!("[realtime-text] realtime conversation text output"); - route_handoffs.route(&sess_clone, text).await; + handoff_error = route_handoffs.route(&sess_clone, text).await.err(); } sess_clone .send_event_raw(ev(EventMsg::RealtimeConversationRealtime( @@ -1676,9 +1678,24 @@ async fn handle_start_inner( }, ))) .await; + if handoff_error.is_some() { + break; + } } - if let Ok(text) = transcript_tail_rx.recv().await { - route_handoffs.route(&sess_clone, text).await; + if handoff_error.is_none() + && let Ok(text) = transcript_tail_rx.recv().await + { + handoff_error = route_handoffs.route(&sess_clone, text).await.err(); + } + if let Some(error) = handoff_error { + end = RealtimeConversationEnd::Error; + sess_clone + .send_event_raw(ev(EventMsg::RealtimeConversationRealtime( + RealtimeConversationRealtimeEvent { + payload: RealtimeEvent::Error(error.to_string()), + }, + ))) + .await; } if fanout_realtime_active.swap(false, Ordering::Relaxed) { match end { diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index b064e24566..d08dcd62a9 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -7,6 +7,9 @@ //! //! Persistent thread settings apply on Started and Steered. Turn start //! options only apply on Started. +//! Host shutdown admission is checked before reserving or starting a new turn. +//! Parent-delegated subagent input bypasses drain; automatic starts remain gated. +//! Realtime drain refusals are returned to the fanout for ordered session teardown. use super::TurnInput; use super::session::Session; @@ -28,6 +31,8 @@ use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::NonSteerableTurnKind; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::turn_input::NotSubmittedReason; use codex_protocol::turn_input::TurnInput as SubmittedTurnInput; @@ -276,6 +281,28 @@ async fn start_or_steer( Ok(TurnInputSubmission::Steered { turn_id }) } Err(NotSubmittedReason::NoActiveTurn) => { + // MAv1 sends explicit input to spawned agents as part of an existing + // parent's work. Client RPCs are gated separately by the host. + let is_delegated_input = settings.start_options.parent_turn_id.is_some() + && matches!( + session + .state + .lock() + .await + .session_configuration + .session_source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ); + let _admission = if is_delegated_input { + None + } else { + let Some(admission) = session.services.extensions.admit_turn_start() else { + return Ok(TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining, + }); + }; + Some(admission) + }; let Some(turn_context) = settings .apply_started(session, submission_id.clone(), TurnStartKind::User) .await? @@ -337,6 +364,27 @@ async fn start_if_idle( }); } + let _admission = session.services.extensions.admit_turn_start(); + // A one-shot review delegate completes its already-running parent's work. + // Its explicit input carries parent lineage; automatic starts do not qualify. + if _admission.is_none() + && !(kind == TurnStartKind::User + && start.parent_turn_id.is_some() + && matches!( + session + .state + .lock() + .await + .session_configuration + .session_source, + SessionSource::SubAgent(SubAgentSource::Review) + )) + { + return Ok(TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining, + }); + } + let turn_state = { let mut active_turn = session.active_turn.lock().await; if active_turn.is_some() { @@ -460,7 +508,10 @@ async fn steer( } impl Session { - pub(crate) async fn route_realtime_text_input(self: &Arc, text: String) { + pub(crate) async fn route_realtime_text_input( + self: &Arc, + text: String, + ) -> Result<(), &'static str> { let submission_id = Uuid::now_v7().to_string(); let submission = handle( self, @@ -478,6 +529,11 @@ impl Session { .await; match submission { Ok(TurnInputSubmission::Started { .. } | TurnInputSubmission::Steered { .. }) => {} + Ok(TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining, + }) => { + return Err("Server is draining; retry the turn after reconnecting"); + } Ok(TurnInputSubmission::NotSubmitted { reason }) => { self.send_event_raw(Event { id: submission_id, @@ -497,6 +553,7 @@ impl Session { .await; } } + Ok(()) } async fn clear_reserved_idle_turn(&self, turn_state: &Arc>) { diff --git a/codex-rs/core/tests/suite/turn_input_submission.rs b/codex-rs/core/tests/suite/turn_input_submission.rs index 837d7b2add..5a4c2ce586 100644 --- a/codex-rs/core/tests/suite/turn_input_submission.rs +++ b/codex-rs/core/tests/suite/turn_input_submission.rs @@ -1,17 +1,25 @@ use codex_core::NotSubmittedReason; use codex_core::RecoverTurnRequest; use codex_core::StartIfIdleSubmission; +use codex_core::StartThreadOptions; use codex_core::SteerSubmission; use codex_core::TurnInput; use codex_core::TurnInputRequest; use codex_core::TurnInputSubmission; use codex_core::TurnStartOptions; use codex_core::config::Constrained; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::TurnStartAdmission; +use codex_protocol::AgentPath; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::user_input::UserInput; @@ -27,12 +35,325 @@ use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::Value; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use test_case::test_case; use tokio::sync::Barrier; use tokio::sync::oneshot; use tokio::time::timeout; +#[derive(Debug)] +struct TestAdmission(AtomicBool); + +impl TurnStartAdmission for TestAdmission { + fn admit_turn_start(&self) -> Option> { + if self.0.load(Ordering::SeqCst) { + None + } else { + Some(Box::new(())) + } + } +} + +#[tokio::test] +async fn host_drain_rejects_turn_start_paths_without_recording_input() -> anyhow::Result<()> { + let server = responses::start_mock_server().await; + let response = responses::mount_sse_once(&server, responses::sse_completed("allowed")).await; + let admission = Arc::new(TestAdmission(AtomicBool::new(true))); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_start_admission(admission.clone()); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&server) + .await?; + assert_eq!( + test.codex + .start_or_steer_turn(user_message_request("rejected direct input")) + .await?, + TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining + }, + ); + for request in [ + user_message_request("rejected queued input"), + TurnInputRequest::user_input(Vec::new()), + TurnInputRequest::new(TurnInput::ResponseItem(responses::user_message_item( + "rejected continuation", + ))), + ] { + assert_eq!( + test.codex.start_turn_if_idle(request).await?, + StartIfIdleSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining + }, + ); + } + assert!(response.requests().is_empty()); + admission.0.store(false, Ordering::SeqCst); + test.codex + .start_turn_if_idle(user_message_request("allowed input")) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let request = response.single_request(); + assert!(request.body_contains_text("allowed input")); + assert!(!request.body_contains_text("rejected direct input")); + assert!(!request.body_contains_text("rejected queued input")); + assert!(!request.body_contains_text("rejected continuation")); + Ok(()) +} + +#[tokio::test] +async fn host_drain_allows_spawned_agent_input_but_not_automatic_work() -> anyhow::Result<()> { + let server = responses::start_mock_server().await; + let response = responses::mount_sse_once(&server, responses::sse_completed("child")).await; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_start_admission(Arc::new(TestAdmission(AtomicBool::new(true)))); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&server) + .await?; + let child = 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, + })), + environments: Some(test.codex.environment_selections().await), + ..StartThreadOptions::new(test.config.clone()) + }) + .await? + .thread; + for request in [ + user_message_request("queued child input"), + TurnInputRequest::user_input(Vec::new()), + ] { + assert_eq!( + child.start_turn_if_idle(request).await?, + StartIfIdleSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining + }, + ); + } + assert_eq!( + child + .start_or_steer_turn(user_message_request("external child input")) + .await?, + TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining + }, + ); + assert!(matches!( + child + .start_or_steer_turn(user_message_request("delegated input").on_start( + TurnStartOptions { + parent_turn_id: Some("parent-turn".to_string()), + ..Default::default() + } + )) + .await?, + TurnInputSubmission::Started { .. } + )); + wait_for_event(&child, |event| matches!(event, EventMsg::TurnComplete(_))).await; + assert!( + response + .single_request() + .body_contains_text("delegated input") + ); + Ok(()) +} + +#[tokio::test] +async fn host_drain_allows_running_review_to_finish_its_delegate() -> anyhow::Result<()> { + use codex_protocol::protocol::ReviewOutputEvent; + use codex_protocol::protocol::ReviewRequest; + use codex_protocol::protocol::ReviewTarget; + + let server = responses::start_mock_server().await; + let expected = ReviewOutputEvent { + overall_explanation: "review completed during drain".to_string(), + ..Default::default() + }; + let response = responses::mount_sse_once( + &server, + responses::sse(vec![ + ev_response_created("review"), + responses::ev_assistant_message("result", &serde_json::to_string(&expected)?), + ev_completed("review"), + ]), + ) + .await; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_start_admission(Arc::new(TestAdmission(AtomicBool::new(true)))); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&server) + .await?; + // The host already admitted the parent review; only its child start hits Core admission. + test.codex + .submit(Op::Review { + review_request: ReviewRequest { + target: ReviewTarget::Custom { + instructions: "review these changes".to_string(), + }, + user_facing_hint: None, + }, + }) + .await?; + let event = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::ExitedReviewMode(_)) + }) + .await; + let EventMsg::ExitedReviewMode(event) = event else { + unreachable!() + }; + assert_eq!(event.review_output, Some(expected)); + response.single_request(); + Ok(()) +} + +#[tokio::test] +async fn host_drain_closes_realtime_after_handoff_error() -> anyhow::Result<()> { + use codex_protocol::protocol::ConversationStartParams; + use codex_protocol::protocol::RealtimeConversationClosedEvent; + use codex_protocol::protocol::RealtimeConversationRealtimeEvent; + use codex_protocol::protocol::RealtimeEvent; + use codex_protocol::protocol::RealtimeHandoffRequested; + use codex_protocol::protocol::RealtimeTranscriptEntry; + + let server = responses::start_mock_server().await; + let response = responses::mount_sse_once(&server, responses::sse_completed("unexpected")).await; + let realtime = responses::start_websocket_server(vec![vec![vec![ + serde_json::json!({ + "type": "session.updated", + "session": { "id": "draining", "instructions": "backend prompt" } + }), + serde_json::json!({ + "type": "conversation.handoff.requested", + "handoff_id": "rejected", + "item_id": "rejected", + "input_transcript": "must not start" + }), + ]]]) + .await; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_start_admission(Arc::new(TestAdmission(AtomicBool::new(true)))); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config({ + let realtime_url = realtime.uri().to_string(); + move |config| { + config.experimental_realtime_ws_base_url = Some(realtime_url); + config.realtime.version = codex_config::config_toml::RealtimeWsVersion::V1; + } + }) + .build_with_auto_env(&server) + .await?; + test.codex + .submit(Op::RealtimeConversationStart(ConversationStartParams { + client_managed_handoffs: false, + delegation_ack_filler: None, + flush_transcript_tail_on_session_end: false, + codex_responses_as_items: false, + codex_response_item_prefix: None, + codex_response_handoff_mode: + codex_protocol::protocol::CodexResponseHandoffMode::Thinking, + codex_response_handoff_channel_prefixes: None, + model: None, + output_modality: codex_protocol::protocol::RealtimeOutputModality::Audio, + include_startup_context: false, + initial_items: Vec::new(), + realtime_start_instructions: None, + realtime_end_instructions: None, + prompt: Some(Some("backend prompt".to_string())), + realtime_session_id: None, + transport: None, + version: None, + voice: None, + })) + .await?; + for expected in [ + EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent { + payload: RealtimeEvent::HandoffRequested(RealtimeHandoffRequested { + handoff_id: "rejected".to_string(), + item_id: "rejected".to_string(), + input_transcript: "must not start".to_string(), + active_transcript: vec![RealtimeTranscriptEntry { + role: "user".to_string(), + text: "must not start".to_string(), + }], + }), + }), + EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent { + payload: RealtimeEvent::Error( + "Server is draining; retry the turn after reconnecting".to_string(), + ), + }), + EventMsg::RealtimeConversationClosed(RealtimeConversationClosedEvent { + reason: Some("error".to_string()), + }), + ] { + let event = wait_for_event(&test.codex, |event| match event { + EventMsg::RealtimeConversationRealtime(RealtimeConversationRealtimeEvent { + payload: RealtimeEvent::HandoffRequested(_) | RealtimeEvent::Error(_), + }) + | EventMsg::RealtimeConversationClosed(_) => true, + EventMsg::TurnStarted(_) | EventMsg::Error(_) => panic!("unexpected event: {event:?}"), + _ => false, + }) + .await; + assert_eq!( + serde_json::to_value(event)?, + serde_json::to_value(expected)? + ); + } + assert!(response.requests().is_empty()); + realtime.shutdown().await; + Ok(()) +} + +#[tokio::test] +async fn host_drain_allows_mailbox_work_to_start_a_turn() -> anyhow::Result<()> { + let server = responses::start_mock_server().await; + let response = responses::mount_sse_once(&server, responses::sse_completed("mailbox")).await; + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_start_admission(Arc::new(TestAdmission(AtomicBool::new(true)))); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&server) + .await?; + // Mailbox input is memory-only and must be processed before the host exits. + test.codex + .submit(Op::InterAgentCommunication { + communication: InterAgentCommunication::new( + AgentPath::try_from("/root/worker").expect("valid agent path"), + AgentPath::root(), + Vec::new(), + "mail while draining".to_string(), + /*trigger_turn*/ true, + ), + start_options: Default::default(), + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert!( + response + .single_request() + .body_contains_text("mail while draining") + ); + Ok(()) +} + fn user_message_request(text: &str) -> TurnInputRequest { TurnInputRequest::user_input(vec![UserInput::Text { text: text.to_string(), diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 84f2ac0b8e..f92d36bc00 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -2,6 +2,7 @@ mod capabilities; mod contributors; mod registry; mod state; +mod turn_admission; mod user_instructions; pub use capabilities::AgentSpawnFuture; @@ -97,6 +98,7 @@ pub use registry::ExtensionRegistryBuilder; pub use registry::empty_extension_registry; pub use state::ExtensionData; pub use state::ExtensionDataInit; +pub use turn_admission::TurnStartAdmission; pub use user_instructions::Instructions; pub use user_instructions::LoadUserInstructionsFuture; pub use user_instructions::LoadedUserInstructions; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index e4aeba9b5e..dddd8b8ab0 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -14,6 +14,7 @@ use crate::ToolLifecycleContributor; use crate::TurnInputContributor; use crate::TurnItemContributor; use crate::TurnLifecycleContributor; +use crate::TurnStartAdmission; /// Mutable registry used while hosts register typed runtime contributions. pub struct ExtensionRegistryBuilder { @@ -25,6 +26,7 @@ impl Default for ExtensionRegistryBuilder { Self { registry: ExtensionRegistry { event_sink: Arc::new(NoopExtensionEventSink), + turn_start_admission: None, thread_lifecycle_contributors: Vec::new(), turn_lifecycle_contributors: Vec::new(), config_contributors: Vec::new(), @@ -60,6 +62,11 @@ impl ExtensionRegistryBuilder { Arc::clone(&self.registry.event_sink) } + /// Installs the host gate for turn-input submissions that start a new turn. + pub fn turn_start_admission(&mut self, admission: Arc) { + self.registry.turn_start_admission = Some(admission); + } + /// Registers one approval-review contributor. pub fn approval_review_contributor(&mut self, contributor: Arc) { self.registry.approval_review_contributors.push(contributor); @@ -139,6 +146,7 @@ impl ExtensionRegistryBuilder { /// Immutable typed registry produced after extensions are installed. pub struct ExtensionRegistry { event_sink: Arc, + turn_start_admission: Option>, thread_lifecycle_contributors: Vec>>, turn_lifecycle_contributors: Vec>, config_contributors: Vec>>, @@ -154,6 +162,15 @@ pub struct ExtensionRegistry { } impl ExtensionRegistry { + /// Acquires the host's turn-start permit, or an empty permit for ungated hosts. + /// A missing permit rejects the start before Core consumes pending input. + pub fn admit_turn_start(&self) -> Option> { + match &self.turn_start_admission { + Some(admission) => admission.admit_turn_start(), + None => Some(Box::new(())), + } + } + /// Returns the host event sink retained by this registry. pub fn event_sink(&self) -> Arc { Arc::clone(&self.event_sink) diff --git a/codex-rs/ext/extension-api/src/turn_admission.rs b/codex-rs/ext/extension-api/src/turn_admission.rs new file mode 100644 index 0000000000..9d5de89fc7 --- /dev/null +++ b/codex-rs/ext/extension-api/src/turn_admission.rs @@ -0,0 +1,12 @@ +//! Lets Core turn-input submissions participate in a host's shutdown drain. + +/// A host-provided gate checked before Core starts a turn-input submission. +/// +/// Implementations return a permit for work admitted before shutdown and +/// Core retains it through submission. `None` skips the start without consuming +/// pending input. Steering an existing turn does not acquire a new permit. +/// Memory-only mailbox wakeups and parent-delegated subagent input bypass this +/// gate so delegated work can finish before exit. Automatic starts remain gated. +pub trait TurnStartAdmission: std::fmt::Debug + Send + Sync { + fn admit_turn_start(&self) -> Option>; +} diff --git a/codex-rs/ext/queue/tests/queue_service.rs b/codex-rs/ext/queue/tests/queue_service.rs index 875dae3292..a597583d39 100644 --- a/codex-rs/ext/queue/tests/queue_service.rs +++ b/codex-rs/ext/queue/tests/queue_service.rs @@ -3,6 +3,8 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::OnceLock; use std::sync::Weak; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Context; @@ -23,6 +25,7 @@ use codex_extension_api::ThreadIdleCause; use codex_extension_api::ThreadIdleInput; use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadResumeInput; +use codex_extension_api::TurnStartAdmission; use codex_protocol::ThreadId; use codex_protocol::items::TurnItem; use codex_protocol::models::ImageDetail; @@ -59,6 +62,19 @@ const TINY_PNG_BYTES: &[u8] = &[ ]; const TINY_PNG_DATA_URL: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVR4nGNgAAIAAAUAAXpeqz8AAAAASUVORK5CYII="; +#[derive(Debug)] +struct TestAdmission(AtomicBool); + +impl TurnStartAdmission for TestAdmission { + fn admit_turn_start(&self) -> Option> { + if self.0.load(Ordering::SeqCst) { + None + } else { + Some(Box::new(())) + } + } +} + #[derive(Default)] struct RecordingEventSink { events: Mutex>, @@ -228,6 +244,49 @@ async fn emit_idle_with_cause( .await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn drain_leaves_persisted_queued_message_for_a_later_start() -> anyhow::Result<()> { + let server = start_mock_server().await; + let response = + responses::mount_sse_once(&server, responses::sse_completed("queued-turn")).await; + let admission = Arc::new(TestAdmission(AtomicBool::new(true))); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_start_admission(admission.clone()); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let queue = loaded_thread_queue(&test)?; + let staged = QueuedItemService::new( + Arc::clone(&queue), + Weak::new(), + Arc::new(NoopExtensionEventSink), + ); + let item = staged + .enqueue(thread_id, user_input("queued after drain")) + .await?; + let service = QueuedItemService::new( + Arc::clone(&queue), + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + ); + + emit_idle(&service, thread_id).await; + assert_eq!(vec![item.clone()], service.list(thread_id).await?); + assert!(response.requests().is_empty()); + + admission.0.store(false, Ordering::SeqCst); + emit_idle(&service, thread_id).await; + wait_for_event_match(test.codex.as_ref(), |event| { + matches!(event, EventMsg::TurnComplete(_)).then_some(()) + }) + .await; + assert!(service.list(thread_id).await?.is_empty()); + assert_eq!(1, response.requests().len()); + Ok(()) +} + #[tokio::test] async fn queued_input_and_unique_event_ids_round_trip() -> anyhow::Result<()> { let (queue, _home) = test_queue().await?; diff --git a/codex-rs/protocol/src/turn_input.rs b/codex-rs/protocol/src/turn_input.rs index 4819d0fb9b..61b1ffe893 100644 --- a/codex-rs/protocol/src/turn_input.rs +++ b/codex-rs/protocol/src/turn_input.rs @@ -215,6 +215,9 @@ pub enum SteerSubmission { /// Why Core did not accept submitted turn input for turn processing. #[derive(Clone, Debug, Eq, PartialEq)] pub enum NotSubmittedReason { + /// The host is draining and no longer permits new regular turns. + ServerDraining, + /// `start_turn_if_idle` found an active turn. NotIdle,