diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 4cba4bddc5..9670577711 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -49,6 +49,7 @@ use crate::session::Codex; use crate::session::CodexSpawnArgs; use crate::session::CodexSpawnOk; use crate::session::SUBMISSION_CHANNEL_CAPACITY; +use crate::session::SubmissionTransport; use crate::session::emit_subagent_session_started; use crate::session::session::Session; use crate::session::turn_context::TurnContext; @@ -186,6 +187,7 @@ pub(crate) async fn run_codex_thread_interactive( agent_status: codex.agent_status.clone(), session: Arc::clone(&codex.session), session_loop_termination: codex.session_loop_termination.clone(), + submission_transport: SubmissionTransport::ForwardingProxy, }) } @@ -232,7 +234,6 @@ pub(crate) async fn run_codex_thread_one_shot( // Bridge events so we can observe completion and shut down automatically. let (tx_bridge, rx_bridge) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); - let ops_tx = io.tx_sub.clone(); let agent_status = io.agent_status.clone(); let session = Arc::clone(&io.session); let session_loop_termination = io.session_loop_termination.clone(); @@ -245,14 +246,7 @@ pub(crate) async fn run_codex_thread_one_shot( ); let _ = tx_bridge.send(event).await; if should_shutdown { - let _ = ops_tx - .send(Submission { - id: "shutdown".to_string(), - op: Op::Shutdown {}, - client_user_message_id: None, - trace: None, - }) - .await; + let _ = io_for_bridge.submit(Op::Shutdown {}).await; child_cancel.cancel(); break; } @@ -271,6 +265,7 @@ pub(crate) async fn run_codex_thread_one_shot( agent_status, session, session_loop_termination, + submission_transport: SubmissionTransport::ForwardingProxy, }) } diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 5cb2251a4a..3c746debd3 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -48,6 +48,7 @@ async fn forward_events_filters_private_events_before_blocked_send_is_cancelled( agent_status, session: Arc::clone(&session), session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, }); let (tx_out, rx_out) = bounded(1); @@ -136,7 +137,7 @@ async fn forward_events_filters_private_events_before_blocked_send_is_cancelled( } #[tokio::test] -async fn forward_ops_preserves_submission_trace_context() { +async fn forwarding_proxy_reserves_once_and_delivers_shutdown() { let (tx_sub, rx_sub) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_tx_events, rx_events) = bounded(SUBMISSION_CHANNEL_CAPACITY); let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); @@ -145,10 +146,20 @@ async fn forward_ops_preserves_submission_trace_context() { tx_sub, rx_event: rx_events, agent_status, - session, + session: Arc::clone(&session), session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, }); let (tx_ops, rx_ops) = bounded(1); + let (_tx_proxy_events, rx_proxy_events) = bounded(SUBMISSION_CHANNEL_CAPACITY); + let proxy = Codex { + tx_sub: tx_ops, + rx_event: rx_proxy_events, + agent_status: codex.agent_status.clone(), + session: Arc::clone(&session), + session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::ForwardingProxy, + }; let cancel = CancellationToken::new(); let forward = tokio::spawn(forward_ops(Arc::clone(&codex), rx_ops, cancel)); @@ -163,8 +174,7 @@ async fn forward_ops_preserves_submission_trace_context() { tracestate: Some("vendor=state".to_string()), }), }; - tx_ops.send(submission.clone()).await.unwrap(); - drop(tx_ops); + proxy.submit_with_id(submission.clone()).await.unwrap(); let forwarded = timeout(Duration::from_secs(1), rx_sub.recv()) .await @@ -173,6 +183,23 @@ async fn forward_ops_preserves_submission_trace_context() { assert_eq!(submission.id, forwarded.id); assert_eq!(submission.op, forwarded.op); assert_eq!(submission.trace, forwarded.trace); + session.finish_submission(); + + let idle_reservation = session + .try_claim_idle_shutdown() + .await + .expect("the forwarded submission should own exactly one reservation"); + drop(idle_reservation); + + proxy.submit(Op::Shutdown).await.unwrap(); + let forwarded_shutdown = timeout(Duration::from_secs(1), rx_sub.recv()) + .await + .expect("forward_ops hung while delivering shutdown") + .expect("forwarded shutdown missing"); + assert!(matches!(forwarded_shutdown.op, Op::Shutdown)); + session.finish_submission(); + + drop(proxy); timeout(Duration::from_secs(1), forward) .await @@ -223,6 +250,7 @@ async fn handle_request_permissions_uses_tool_call_id_for_round_trip() { agent_status, session: Arc::clone(&parent_session), session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, }); let call_id = "tool-call-1".to_string(); @@ -328,6 +356,7 @@ async fn handle_exec_approval_uses_call_id_for_guardian_review_and_approval_id_f agent_status, session: Arc::clone(&parent_session), session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, }); let cancel_token = CancellationToken::new(); diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index c80268e2f9..fc7f74f78c 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -211,6 +211,10 @@ impl CodexThread { self.codex.shutdown_and_wait().await } + pub async fn try_shutdown_if_idle(&self) -> CodexResult { + self.codex.try_shutdown_if_idle().await + } + /// Wait until the underlying session loop has terminated. pub async fn wait_until_terminated(&self) { self.codex.session_loop_termination.clone().await; @@ -452,6 +456,9 @@ impl CodexThread { /// Records a user-role session-prefix message without creating a new user turn boundary. pub(crate) async fn inject_user_message_without_turn(&self, message: String) { + let Some(_reservation) = self.codex.session.try_reserve_activity() else { + return; + }; let item = ResponseItem::Message { id: None, role: "user".to_string(), @@ -472,6 +479,11 @@ impl CodexThread { "items must not be empty".to_string(), )); } + let Some(_reservation) = self.codex.session.try_reserve_activity() else { + return Err(CodexErr::InvalidRequest( + "thread is shutting down".to_string(), + )); + }; let turn_context = self.codex.session.new_default_turn().await; if self.codex.session.reference_context_item().await.is_none() { @@ -634,6 +646,11 @@ impl CodexThread { server: &str, uri: &str, ) -> anyhow::Result { + let _reservation = self + .codex + .session + .try_reserve_activity() + .ok_or_else(|| anyhow::anyhow!("thread is shutting down"))?; let result = self .current_mcp_runtime() .await @@ -651,6 +668,11 @@ impl CodexThread { arguments: Option, meta: Option, ) -> anyhow::Result { + let _reservation = self + .codex + .session + .try_reserve_activity() + .ok_or_else(|| anyhow::anyhow!("thread is shutting down"))?; self.current_mcp_runtime() .await .manager_arc() @@ -663,6 +685,11 @@ impl CodexThread { } pub async fn increment_out_of_band_elicitation_count(&self) -> CodexResult { + let Some(_reservation) = self.codex.session.try_reserve_activity() else { + return Err(CodexErr::InvalidRequest( + "thread is shutting down".to_string(), + )); + }; let mut elicitations = self.out_of_band_elicitations.lock().await; let incremented = elicitations.count.checked_add(1).ok_or_else(|| { CodexErr::Fatal("out-of-band elicitation count overflowed".to_string()) @@ -675,6 +702,11 @@ impl CodexThread { } pub async fn decrement_out_of_band_elicitation_count(&self) -> CodexResult { + let Some(_reservation) = self.codex.session.try_reserve_activity() else { + return Err(CodexErr::InvalidRequest( + "thread is shutting down".to_string(), + )); + }; let mut elicitations = self.out_of_band_elicitations.lock().await; if elicitations.count == 0 { return Err(CodexErr::InvalidRequest( diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index de56de41a6..5e20a25e23 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -1154,6 +1154,7 @@ mod tests { agent_status, session, session_loop_termination: crate::session::completed_session_loop_termination(), + submission_transport: crate::session::SubmissionTransport::Canonical, }, cancel_token: CancellationToken::new(), reuse_key, diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 924750bc34..62ddcea799 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -841,6 +841,7 @@ pub(super) async fn submission_loop( } .instrument(dispatch_span) .await; + sess.finish_submission(); if should_exit { shutdown_received = true; break; diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 5c5b7d1feb..2ad7285757 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -62,6 +62,12 @@ impl Session { )); } + let Some(reservation) = self.try_reserve_activity() else { + return Err(TryStartTurnIfIdleError::new( + TryStartTurnIfIdleRejectionReason::Busy, + input, + )); + }; let turn_state = { let mut active_turn = self.active_turn.lock().await; if active_turn.is_some() { @@ -73,6 +79,7 @@ impl Session { let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); Arc::clone(&active_turn.turn_state) }; + drop(reservation); if self.input_queue.has_trigger_turn_mailbox_items().await { self.clear_reserved_idle_turn(&turn_state).await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index ba8eacb395..70fea7bc58 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -393,10 +393,17 @@ pub struct Codex { // Shared future for the background submission loop completion so multiple // callers can wait for shutdown. pub(crate) session_loop_termination: SessionLoopTermination, + pub(crate) submission_transport: SubmissionTransport, } pub(crate) type SessionLoopTermination = Shared>; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum SubmissionTransport { + Canonical, + ForwardingProxy, +} + /// Wrapper returned by [`Codex::spawn`] containing the spawned [`Codex`] and /// the unique session id. pub struct CodexSpawnOk { @@ -721,6 +728,7 @@ impl Codex { agent_status: agent_status_rx, session, session_loop_termination: session_loop_termination_from_handle(session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }; Ok(CodexSpawnOk { codex, thread_id }) @@ -767,17 +775,73 @@ impl Codex { /// Use sparingly: prefer `submit()` so Codex is responsible for generating /// unique IDs for each submission. + #[expect( + clippy::await_holding_invalid_type, + reason = "submission ordering must remain serialized while a bounded-channel send waits" + )] pub async fn submit_with_id(&self, mut sub: Submission) -> CodexResult<()> { if sub.trace.is_none() { sub.trace = current_span_w3c_trace_context(); } - self.tx_sub - .send(sub) - .await - .map_err(|_| CodexErr::InternalAgentDied)?; + if self.submission_transport == SubmissionTransport::ForwardingProxy { + self.tx_sub + .send(sub) + .await + .map_err(|_| CodexErr::InternalAgentDied)?; + return Ok(()); + } + let is_shutdown = matches!(&sub.op, Op::Shutdown); + let _send_guard = self.session.submission_send_lock.lock().await; + let Some(reservation) = self.session.submission_lifecycle.try_reserve(is_shutdown) else { + if is_shutdown { + return Ok(()); + } + return Err(CodexErr::InvalidRequest( + "thread is shutting down".to_string(), + )); + }; + if self.tx_sub.send(sub).await.is_err() { + reservation.release_after_failed_delivery(); + return Err(CodexErr::InternalAgentDied); + } + reservation.commit(); Ok(()) } + /// Atomically submits shutdown only when no work is active, queued, or + /// being dispatched. Once claimed, new submissions and automatic idle work + /// are rejected so shutdown cannot race a newly starting turn. + #[expect( + clippy::await_holding_invalid_type, + reason = "the send lock serializes the idle claim with submission delivery" + )] + pub async fn try_shutdown_if_idle(&self) -> CodexResult { + if self.submission_transport == SubmissionTransport::ForwardingProxy { + return Err(CodexErr::InvalidRequest( + "idle shutdown is unavailable through a forwarding proxy".to_string(), + )); + } + let _send_guard = self.session.submission_send_lock.lock().await; + if self.session.submission_lifecycle.is_closing() { + return Ok(true); + } + let Some(reservation) = self.session.try_claim_idle_shutdown().await else { + return Ok(false); + }; + let sub = Submission { + id: new_submission_id(), + op: Op::Shutdown, + client_user_message_id: None, + trace: current_span_w3c_trace_context(), + }; + if self.tx_sub.send(sub).await.is_err() { + reservation.release_after_failed_delivery(); + return Err(CodexErr::InternalAgentDied); + } + reservation.commit(); + Ok(true) + } + /// Persist a thread-level memory mode update for the active session. /// /// This is a local-only operation that updates rollout metadata directly @@ -978,6 +1042,28 @@ fn push_prompt_fragment( } impl Session { + async fn has_idle_shutdown_blocker(&self) -> bool { + matches!(self.agent_status.borrow().clone(), AgentStatus::Running) + || self.conversation.running_state().await.is_some() + || *self.services.elicitations.subscribe().borrow() + || self.input_queue.has_pending_mailbox_items().await + || self.active_turn.lock().await.is_some() + || !self.list_background_terminals().await.is_empty() + } + + pub(crate) async fn try_claim_idle_shutdown( + &self, + ) -> Option> { + if self.has_idle_shutdown_blocker().await { + return None; + } + let reservation = self.submission_lifecycle.try_reserve_idle_shutdown()?; + if self.has_idle_shutdown_blocker().await { + return None; + } + Some(reservation) + } + pub(crate) async fn app_server_client_metadata(&self) -> AppServerClientMetadata { let state = self.state.lock().await; AppServerClientMetadata { @@ -1194,6 +1280,9 @@ impl Session { } pub(crate) async fn route_realtime_text_input(self: &Arc, text: String) { + let Some(_reservation) = self.try_reserve_activity() else { + return; + }; handlers::user_input_or_turn_inner( self, Uuid::now_v7().to_string(), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 2261eaa495..3a05b9354d 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -20,11 +20,114 @@ use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnEnvironmentSelections; use std::sync::OnceLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use tokio::sync::Semaphore; /// Context for an initialized model agent /// /// A session has at most 1 running task at a time, and can be interrupted by user input. +pub(crate) struct SubmissionLifecycle { + state: AtomicUsize, +} + +const SUBMISSION_CLOSING: usize = 1 << (usize::BITS - 1); +const SUBMISSION_COUNT_MASK: usize = SUBMISSION_CLOSING - 1; + +impl Default for SubmissionLifecycle { + fn default() -> Self { + Self { + state: AtomicUsize::new(0), + } + } +} + +impl SubmissionLifecycle { + pub(crate) fn is_closing(&self) -> bool { + self.state.load(Ordering::SeqCst) & SUBMISSION_CLOSING != 0 + } + + pub(crate) fn try_reserve(&self, close: bool) -> Option> { + self.state + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |state| { + if state & SUBMISSION_CLOSING != 0 { + return None; + } + let count = state & SUBMISSION_COUNT_MASK; + assert!(count < SUBMISSION_COUNT_MASK, "submission count overflow"); + Some((state + 1) | if close { SUBMISSION_CLOSING } else { 0 }) + }) + .ok()?; + Some(SubmissionReservation { + lifecycle: self, + clear_closing_on_drop: close, + active: true, + }) + } + + pub(crate) fn try_reserve_idle_shutdown(&self) -> Option> { + self.state + .compare_exchange( + 0, + SUBMISSION_CLOSING | 1, + Ordering::SeqCst, + Ordering::SeqCst, + ) + .ok() + .map(|_| SubmissionReservation { + lifecycle: self, + clear_closing_on_drop: true, + active: true, + }) + } + + pub(crate) fn finish_submission(&self) { + self.release(/*clear_closing*/ false); + } + + fn release(&self, clear_closing: bool) { + let _ = self + .state + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |state| { + if state & SUBMISSION_COUNT_MASK == 0 { + return None; + } + let mut next = state - 1; + if clear_closing { + next &= !SUBMISSION_CLOSING; + } + Some(next) + }); + } +} + +pub(crate) struct SubmissionReservation<'a> { + lifecycle: &'a SubmissionLifecycle, + clear_closing_on_drop: bool, + active: bool, +} + +impl SubmissionReservation<'_> { + pub(crate) fn commit(mut self) { + self.active = false; + } + + pub(crate) fn release_after_failed_delivery(mut self) { + if self.active { + self.lifecycle.release(/*clear_closing*/ false); + self.active = false; + } + } +} + +impl Drop for SubmissionReservation<'_> { + fn drop(&mut self) { + if self.active { + self.lifecycle.release(self.clear_closing_on_drop); + } + } +} + pub(crate) struct Session { pub(crate) thread_id: ThreadId, pub(crate) installation_id: String, @@ -41,6 +144,8 @@ pub(crate) struct Session { pub(super) pending_mcp_server_refresh_config: Mutex>, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, + pub(crate) submission_lifecycle: SubmissionLifecycle, + pub(crate) submission_send_lock: Mutex<()>, pub(crate) input_queue: InputQueue, pub(crate) guardian_review_session: GuardianReviewSessionManager, pub(crate) services: SessionServices, @@ -465,6 +570,14 @@ impl Session { self.thread_id } + pub(crate) fn try_reserve_activity(&self) -> Option> { + self.submission_lifecycle.try_reserve(/*close*/ false) + } + + pub(crate) fn finish_submission(&self) { + self.submission_lifecycle.finish_submission(); + } + /// Returns the identity shared by the root thread and all descendant threads. pub(crate) fn session_id(&self) -> SessionId { self.services.agent_control.session_id() @@ -1137,6 +1250,8 @@ impl Session { pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + submission_lifecycle: SubmissionLifecycle::default(), + submission_send_lock: Mutex::new(()), input_queue: InputQueue::new(), guardian_review_session: GuardianReviewSessionManager::default(), services, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6cc223cd7d..ad19e4093d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1,3 +1,4 @@ +use super::session::SubmissionLifecycle; use super::turn_context::TurnEnvironment; use super::*; use crate::agents_md_manager::AgentsMdManager; @@ -5518,6 +5519,8 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + submission_lifecycle: SubmissionLifecycle::default(), + submission_send_lock: Mutex::new(()), input_queue: super::input_queue::InputQueue::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, @@ -6448,6 +6451,7 @@ async fn submit_with_id_captures_current_span_trace_context() { agent_status, session: Arc::new(session), session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, }; let _trace_test_context = install_test_tracing("codex-core-tests"); @@ -6483,6 +6487,146 @@ async fn submit_with_id_captures_current_span_trace_context() { assert_eq!(submitted.trace, Some(expected_trace)); } +async fn codex_with_disconnected_submission_channel() -> (Codex, Arc) { + let (session, _turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let (tx_sub, rx_sub) = async_channel::bounded(1); + drop(rx_sub); + let (_tx_event, rx_event) = async_channel::unbounded(); + let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); + let codex = Codex { + tx_sub, + rx_event, + agent_status, + session: Arc::clone(&session), + session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, + }; + (codex, session) +} + +fn assert_failed_shutdown_stays_closed(session: &Session) { + assert!(session.submission_lifecycle.is_closing()); +} + +#[tokio::test] +async fn failed_idle_shutdown_delivery_stays_closed() { + let (codex, session) = codex_with_disconnected_submission_channel().await; + + assert!(matches!( + codex.try_shutdown_if_idle().await, + Err(CodexErr::InternalAgentDied) + )); + assert_failed_shutdown_stays_closed(&session); +} + +#[tokio::test] +async fn failed_explicit_shutdown_delivery_stays_closed() { + let (codex, session) = codex_with_disconnected_submission_channel().await; + + assert!(matches!( + codex + .submit_with_id(Submission { + id: "shutdown".to_string(), + op: Op::Shutdown, + client_user_message_id: None, + trace: None, + }) + .await, + Err(CodexErr::InternalAgentDied) + )); + assert_failed_shutdown_stays_closed(&session); +} + +#[tokio::test] +async fn cancelled_submission_and_idle_shutdown_release_lifecycle_reservations() { + let (session, _turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let (tx_sub, rx_sub) = async_channel::bounded(1); + let (_tx_event, rx_event) = async_channel::unbounded(); + let (_agent_status_tx, agent_status) = watch::channel(AgentStatus::PendingInit); + + tx_sub + .send(Submission { + id: "filler".into(), + op: Op::Interrupt, + client_user_message_id: None, + trace: None, + }) + .await + .expect("fill submission channel"); + + let codex = Arc::new(Codex { + tx_sub, + rx_event, + agent_status, + session: Arc::clone(&session), + session_loop_termination: completed_session_loop_termination(), + submission_transport: SubmissionTransport::Canonical, + }); + + let blocked_submit = { + let codex = Arc::clone(&codex); + tokio::spawn(async move { codex.submit(Op::Interrupt).await }) + }; + tokio::time::timeout(StdDuration::from_secs(1), async { + while session.submission_send_lock.try_lock().is_ok() { + tokio::task::yield_now().await; + } + }) + .await + .expect("submission should hold the send lock while blocking"); + blocked_submit.abort(); + assert!( + blocked_submit + .await + .expect_err("blocked submission should be cancelled") + .is_cancelled() + ); + let released_submission = session + .try_claim_idle_shutdown() + .await + .expect("cancelled submission should release its lifecycle reservation"); + drop(released_submission); + assert!(!session.submission_lifecycle.is_closing()); + + let blocked_shutdown = { + let codex = Arc::clone(&codex); + tokio::spawn(async move { codex.try_shutdown_if_idle().await }) + }; + tokio::time::timeout(StdDuration::from_secs(1), async { + while !session.submission_lifecycle.is_closing() { + tokio::task::yield_now().await; + } + }) + .await + .expect("idle shutdown should claim lifecycle state before blocking"); + blocked_shutdown.abort(); + assert!( + blocked_shutdown + .await + .expect_err("blocked idle shutdown should be cancelled") + .is_cancelled() + ); + let released_shutdown = session + .try_claim_idle_shutdown() + .await + .expect("cancelled idle shutdown should release its lifecycle reservation"); + drop(released_shutdown); + assert!(!session.submission_lifecycle.is_closing()); + + let filler = rx_sub.recv().await.expect("filler submission"); + assert_eq!(filler.id, "filler"); + assert!( + codex + .try_shutdown_if_idle() + .await + .expect("idle shutdown should submit after cancellation") + ); + let shutdown = rx_sub.recv().await.expect("shutdown submission"); + assert_eq!(shutdown.op, Op::Shutdown); +} + #[tokio::test] async fn new_default_turn_captures_current_span_trace_id() { let (session, _turn_context) = make_session_and_context().await; @@ -7180,6 +7324,7 @@ async fn shutdown_and_wait_allows_multiple_waiters() { agent_status, session: Arc::new(session), session_loop_termination: session_loop_termination_from_handle(session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }); let waiter_1 = { @@ -7218,6 +7363,7 @@ async fn shutdown_and_wait_waits_when_shutdown_is_already_in_progress() { agent_status, session: Arc::new(session), session_loop_termination: session_loop_termination_from_handle(session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }); let waiter = { @@ -7256,6 +7402,7 @@ async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() { agent_status: parent_agent_status, session: Arc::clone(&parent_session), session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }; let (child_session, _child_turn_context) = make_session_and_context().await; @@ -7279,6 +7426,7 @@ async fn shutdown_and_wait_shuts_down_cached_guardian_subagent() { agent_status: child_agent_status, session: Arc::new(child_session), session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }; parent_session .guardian_review_session @@ -7312,6 +7460,7 @@ async fn cached_guardian_subagent_exposes_its_rollout_path() { agent_status: child_agent_status, session: Arc::new(child_session), session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }; parent_session .guardian_review_session @@ -7345,6 +7494,7 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() { agent_status: parent_agent_status, session: Arc::clone(&parent_session), session_loop_termination: session_loop_termination_from_handle(parent_session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }; let (child_session, _child_turn_context) = make_session_and_context().await; @@ -7368,6 +7518,7 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() { agent_status: child_agent_status, session: Arc::new(child_session), session_loop_termination: session_loop_termination_from_handle(child_session_loop_handle), + submission_transport: SubmissionTransport::Canonical, }; parent_session .guardian_review_session @@ -7644,6 +7795,8 @@ where pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + submission_lifecycle: SubmissionLifecycle::default(), + submission_send_lock: Mutex::new(()), input_queue: super::input_queue::InputQueue::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, @@ -9951,6 +10104,43 @@ async fn try_start_turn_if_idle_rejects_active_turn_without_injecting() { sess.abort_all_tasks(TurnAbortReason::Interrupted).await; } +#[tokio::test] +async fn idle_shutdown_claim_waits_for_submissions_and_blocks_automatic_turns() { + let (sess, _tc, _rx) = make_session_and_context_with_rx().await; + let pending = sess + .try_reserve_activity() + .expect("open session should reserve a submission"); + assert!(sess.try_claim_idle_shutdown().await.is_none()); + drop(pending); + + sess.input_queue + .enqueue_mailbox_communication(InterAgentCommunication::new( + AgentPath::root(), + AgentPath::root(), + Vec::new(), + "pending context".to_string(), + /*trigger_turn*/ false, + )) + .await; + assert!(sess.try_claim_idle_shutdown().await.is_none()); + sess.input_queue.drain_mailbox_input_items().await; + + let shutdown = sess + .try_claim_idle_shutdown() + .await + .expect("idle session should claim shutdown"); + let item = user_message("synthetic idle input"); + let err = sess + .try_start_turn_if_idle(vec![item.clone()]) + .await + .expect_err("claimed shutdown should reject automatic idle work"); + + assert_eq!(TryStartTurnIfIdleRejectionReason::Busy, err.reason()); + assert_eq!(vec![item], err.into_input()); + assert!(sess.active_turn.lock().await.is_none()); + drop(shutdown); +} + #[tokio::test] async fn try_start_turn_if_idle_rejects_plan_mode_without_injecting() { let (sess, _tc, _rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index c6d9d2990c..a7145f8db8 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -474,6 +474,9 @@ impl Session { return; } + let Some(reservation) = self.try_reserve_activity() else { + return; + }; { let mut active_turn = self.active_turn.lock().await; if active_turn.is_some() { @@ -481,6 +484,7 @@ impl Session { } *active_turn = Some(ActiveTurn::default()); } + drop(reservation); let turn_context = self.new_default_turn_with_sub_id(sub_id).await; self.maybe_emit_model_warnings_for_turn(turn_context.as_ref())