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,