fix: race pending (#16561)

codex-rs/core/src/codex.rs

@@ -4023,6 +4023,7 @@ impl Session {
 
         let mut turn_state = active_turn.turn_state.lock().await;
         turn_state.push_pending_input(input.into());
+        turn_state.accept_mailbox_delivery_for_current_turn();
         Ok(active_turn_id.clone())
     }
 
@@ -4044,6 +4045,25 @@ impl Session {
         }
     }
 
+    pub(crate) async fn defer_mailbox_delivery_to_next_turn(&self, sub_id: &str) {
+        let turn_state = {
+            let active = self.active_turn.lock().await;
+            active.as_ref().and_then(|active_turn| {
+                active_turn
+                    .tasks
+                    .contains_key(sub_id)
+                    .then(|| Arc::clone(&active_turn.turn_state))
+            })
+        };
+        let Some(turn_state) = turn_state else {
+            return;
+        };
+        turn_state
+            .lock()
+            .await
+            .defer_mailbox_delivery_to_next_turn();
+    }
+
     pub(crate) fn subscribe_mailbox_seq(&self) -> watch::Receiver<u64> {
         self.mailbox.subscribe()
     }
@@ -4069,16 +4089,22 @@ impl Session {
     }
 
     pub async fn get_pending_input(&self) -> Vec<ResponseInputItem> {
-        let pending_input = {
+        let (pending_input, accepts_mailbox_delivery) = {
             let mut active = self.active_turn.lock().await;
             match active.as_mut() {
                 Some(at) => {
                     let mut ts = at.turn_state.lock().await;
-                    ts.take_pending_input()
+                    (
+                        ts.take_pending_input(),
+                        ts.accepts_mailbox_delivery_for_current_turn(),
+                    )
                 }
-                None => Vec::new(),
+                None => (Vec::new(), true),
             }
         };
+        if !accepts_mailbox_delivery {
+            return pending_input;
+        }
         let mailbox_items = {
             let mut mailbox_rx = self.mailbox_rx.lock().await;
             mailbox_rx
@@ -4118,17 +4144,26 @@ impl Session {
     }
 
     pub async fn has_pending_input(&self) -> bool {
-        if self.mailbox_rx.lock().await.has_pending() {
+        let (has_turn_pending_input, accepts_mailbox_delivery) = {
+            let active = self.active_turn.lock().await;
+            match active.as_ref() {
+                Some(at) => {
+                    let ts = at.turn_state.lock().await;
+                    (
+                        ts.has_pending_input(),
+                        ts.accepts_mailbox_delivery_for_current_turn(),
+                    )
+                }
+                None => (false, true),
+            }
+        };
+        if has_turn_pending_input {
             return true;
         }
-        let active = self.active_turn.lock().await;
-        match active.as_ref() {
-            Some(at) => {
-                let ts = at.turn_state.lock().await;
-                ts.has_pending_input()
-            }
-            None => false,
+        if !accepts_mailbox_delivery {
+            return false;
         }
+        self.mailbox_rx.lock().await.has_pending()
     }
 
     pub async fn list_resources(

codex-rs/core/src/codex_tests.rs

@@ -18,6 +18,7 @@ use crate::tools::format_exec_output_str;
 use codex_features::Features;
 use codex_login::CodexAuth;
 use codex_mcp::mcp_connection_manager::ToolInfo;
+use codex_protocol::AgentPath;
 use codex_protocol::ThreadId;
 use codex_protocol::models::FunctionCallOutputBody;
 use codex_protocol::models::FunctionCallOutputPayload;
@@ -69,6 +70,7 @@ use codex_protocol::protocol::ConversationAudioParams;
 use codex_protocol::protocol::CreditsSnapshot;
 use codex_protocol::protocol::GranularApprovalConfig;
 use codex_protocol::protocol::InitialHistory;
+use codex_protocol::protocol::InterAgentCommunication;
 use codex_protocol::protocol::NetworkApprovalProtocol;
 use codex_protocol::protocol::RateLimitSnapshot;
 use codex_protocol::protocol::RateLimitWindow;
@@ -4781,6 +4783,119 @@ async fn queued_response_items_for_next_turn_move_into_next_active_turn() {
     assert_eq!(sess.get_pending_input().await, vec![queued_item]);
 }
 
+#[tokio::test]
+async fn queue_only_mailbox_mail_waits_for_next_turn_after_answer_boundary() {
+    let (sess, tc, _rx) = make_session_and_context_with_rx().await;
+    let communication = InterAgentCommunication::new(
+        AgentPath::try_from("/root/worker").expect("worker path should parse"),
+        AgentPath::root(),
+        Vec::new(),
+        "late queue-only update".to_string(),
+        /*trigger_turn*/ false,
+    );
+    sess.spawn_task(
+        Arc::clone(&tc),
+        Vec::new(),
+        NeverEndingTask {
+            kind: TaskKind::Regular,
+            listen_to_cancellation_token: true,
+        },
+    )
+    .await;
+
+    sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
+    sess.enqueue_mailbox_communication(communication.clone());
+
+    assert!(
+        !sess.has_pending_input().await,
+        "queue-only mailbox mail should stay buffered once the current turn emitted its answer"
+    );
+    assert_eq!(sess.get_pending_input().await, Vec::new());
+
+    sess.abort_all_tasks(TurnAbortReason::Replaced).await;
+
+    assert_eq!(
+        sess.get_pending_input().await,
+        vec![communication.to_response_input_item()],
+    );
+}
+
+#[tokio::test]
+async fn trigger_turn_mailbox_mail_waits_for_next_turn_after_answer_boundary() {
+    let (sess, tc, _rx) = make_session_and_context_with_rx().await;
+    sess.spawn_task(
+        Arc::clone(&tc),
+        Vec::new(),
+        NeverEndingTask {
+            kind: TaskKind::Regular,
+            listen_to_cancellation_token: true,
+        },
+    )
+    .await;
+
+    sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
+    sess.enqueue_mailbox_communication(InterAgentCommunication::new(
+        AgentPath::try_from("/root/worker").expect("worker path should parse"),
+        AgentPath::root(),
+        Vec::new(),
+        "late trigger update".to_string(),
+        /*trigger_turn*/ true,
+    ));
+
+    assert!(
+        !sess.has_pending_input().await,
+        "trigger-turn mailbox mail should not extend the current turn after its answer boundary"
+    );
+
+    sess.abort_all_tasks(TurnAbortReason::Replaced).await;
+
+    assert!(sess.has_trigger_turn_mailbox_items().await);
+}
+
+#[tokio::test]
+async fn steered_input_reopens_mailbox_delivery_for_current_turn() {
+    let (sess, tc, _rx) = make_session_and_context_with_rx().await;
+    let communication = InterAgentCommunication::new(
+        AgentPath::try_from("/root/worker").expect("worker path should parse"),
+        AgentPath::root(),
+        Vec::new(),
+        "queued child update".to_string(),
+        /*trigger_turn*/ false,
+    );
+    sess.spawn_task(
+        Arc::clone(&tc),
+        Vec::new(),
+        NeverEndingTask {
+            kind: TaskKind::Regular,
+            listen_to_cancellation_token: true,
+        },
+    )
+    .await;
+
+    sess.defer_mailbox_delivery_to_next_turn(&tc.sub_id).await;
+    sess.enqueue_mailbox_communication(communication.clone());
+    sess.steer_input(
+        vec![UserInput::Text {
+            text: "follow up".to_string(),
+            text_elements: Vec::new(),
+        }],
+        Some(&tc.sub_id),
+    )
+    .await
+    .expect("steered input should be accepted");
+
+    assert_eq!(
+        sess.get_pending_input().await,
+        vec![
+            ResponseInputItem::from(vec![UserInput::Text {
+                text: "follow up".to_string(),
+                text_elements: Vec::new(),
+            }]),
+            communication.to_response_input_item(),
+        ],
+    );
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn abort_review_task_emits_exited_then_aborted_and_records_history() {
     let (sess, tc, rx) = make_session_and_context_with_rx().await;

codex-rs/core/src/state/turn.rs

@@ -29,6 +29,17 @@ pub(crate) struct ActiveTurn {
     pub(crate) turn_state: Arc<Mutex<TurnState>>,
 }
 
+/// Whether mailbox deliveries should still be folded into the current turn.
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+pub(crate) enum MailboxDeliveryPhase {
+    /// Incoming mailbox messages can still be consumed by the current turn.
+    #[default]
+    CurrentTurn,
+    /// The current turn already emitted visible final answer text; mailbox
+    /// messages should remain queued for a later turn.
+    NextTurn,
+}
+
 impl Default for ActiveTurn {
     fn default() -> Self {
         Self {
@@ -81,6 +92,7 @@ pub(crate) struct TurnState {
     pending_elicitations: HashMap<(String, RequestId), oneshot::Sender<ElicitationResponse>>,
     pending_dynamic_tools: HashMap<String, oneshot::Sender<DynamicToolResponse>>,
     pending_input: Vec<ResponseInputItem>,
+    mailbox_delivery_phase: MailboxDeliveryPhase,
     granted_permissions: Option<PermissionProfile>,
     pub(crate) tool_calls: u64,
     pub(crate) token_usage_at_turn_start: TokenUsage,
@@ -202,6 +214,18 @@ impl TurnState {
         !self.pending_input.is_empty()
     }
 
+    pub(crate) fn defer_mailbox_delivery_to_next_turn(&mut self) {
+        self.mailbox_delivery_phase = MailboxDeliveryPhase::NextTurn;
+    }
+
+    pub(crate) fn accept_mailbox_delivery_for_current_turn(&mut self) {
+        self.mailbox_delivery_phase = MailboxDeliveryPhase::CurrentTurn;
+    }
+
+    pub(crate) fn accepts_mailbox_delivery_for_current_turn(&self) -> bool {
+        self.mailbox_delivery_phase == MailboxDeliveryPhase::CurrentTurn
+    }
+
     pub(crate) fn record_granted_permissions(&mut self, permissions: PermissionProfile) {
         self.granted_permissions =
             merge_permission_profiles(self.granted_permissions.as_ref(), Some(&permissions));

codex-rs/core/src/stream_events_utils.rs

@@ -23,6 +23,7 @@ use crate::tools::router::ToolRouter;
 use codex_protocol::models::DeveloperInstructions;
 use codex_protocol::models::FunctionCallOutputBody;
 use codex_protocol::models::FunctionCallOutputPayload;
+use codex_protocol::models::MessagePhase;
 use codex_protocol::models::ResponseInputItem;
 use codex_protocol::models::ResponseItem;
 use codex_rollout::state_db;
@@ -129,6 +130,13 @@ pub(crate) async fn record_completed_response_item(
 ) {
     sess.record_conversation_items(turn_context, std::slice::from_ref(item))
         .await;
+    if completed_item_defers_mailbox_delivery_to_next_turn(
+        item,
+        turn_context.collaboration_mode.mode == ModeKind::Plan,
+    ) {
+        sess.defer_mailbox_delivery_to_next_turn(&turn_context.sub_id)
+            .await;
+    }
     maybe_mark_thread_memory_mode_polluted_from_web_search(sess, turn_context, item).await;
     record_stage1_output_usage_for_completed_item(turn_context, item).await;
 }
@@ -426,6 +434,24 @@ pub(crate) fn last_assistant_message_from_item(
     None
 }
 
+fn completed_item_defers_mailbox_delivery_to_next_turn(
+    item: &ResponseItem,
+    plan_mode: bool,
+) -> bool {
+    match item {
+        ResponseItem::Message { role, phase, .. } => {
+            if role != "assistant" || matches!(phase, Some(MessagePhase::Commentary)) {
+                return false;
+            }
+            // Treat `None` like final-answer text so untagged providers default
+            // to the safer "defer mailbox mail" behavior.
+            last_assistant_message_from_item(item, plan_mode).is_some()
+        }
+        ResponseItem::ImageGenerationCall { .. } => true,
+        _ => false,
+    }
+}
+
 pub(crate) fn response_input_to_response_item(input: &ResponseInputItem) -> Option<ResponseItem> {
     match input {
         ResponseInputItem::FunctionCallOutput { call_id, output } => {

codex-rs/core/src/stream_events_utils_tests.rs

@@ -1,3 +1,4 @@
+use super::completed_item_defers_mailbox_delivery_to_next_turn;
 use super::handle_non_tool_response_item;
 use super::image_generation_artifact_path;
 use super::last_assistant_message_from_item;
@@ -6,10 +7,15 @@ use crate::codex::make_session_and_context;
 use crate::error::CodexErr;
 use codex_protocol::items::TurnItem;
 use codex_protocol::models::ContentItem;
+use codex_protocol::models::MessagePhase;
 use codex_protocol::models::ResponseItem;
 use pretty_assertions::assert_eq;
 
 fn assistant_output_text(text: &str) -> ResponseItem {
+    assistant_output_text_with_phase(text, None)
+}
+
+fn assistant_output_text_with_phase(text: &str, phase: Option<MessagePhase>) -> ResponseItem {
     ResponseItem::Message {
         id: Some("msg-1".to_string()),
         role: "assistant".to_string(),
@@ -17,7 +23,7 @@ fn assistant_output_text(text: &str) -> ResponseItem {
             text: text.to_string(),
         }],
         end_turn: Some(true),
-        phase: None,
+        phase,
     }
 }
 
@@ -87,6 +93,38 @@ fn last_assistant_message_from_item_returns_none_for_plan_only_hidden_message()
     );
 }
 
+#[test]
+fn completed_item_defers_mailbox_delivery_for_unknown_phase_messages() {
+    let item = assistant_output_text("final answer");
+
+    assert!(completed_item_defers_mailbox_delivery_to_next_turn(
+        &item, /*plan_mode*/ false,
+    ));
+}
+
+#[test]
+fn completed_item_keeps_mailbox_delivery_open_for_commentary_messages() {
+    let item = assistant_output_text_with_phase("still working", Some(MessagePhase::Commentary));
+
+    assert!(!completed_item_defers_mailbox_delivery_to_next_turn(
+        &item, /*plan_mode*/ false,
+    ));
+}
+
+#[test]
+fn completed_item_defers_mailbox_delivery_for_image_generation_calls() {
+    let item = ResponseItem::ImageGenerationCall {
+        id: "ig-1".to_string(),
+        status: "completed".to_string(),
+        revised_prompt: None,
+        result: "Zm9v".to_string(),
+    };
+
+    assert!(completed_item_defers_mailbox_delivery_to_next_turn(
+        &item, /*plan_mode*/ false,
+    ));
+}
+
 #[tokio::test]
 async fn save_image_generation_result_saves_base64_to_png_in_codex_home() {
     let codex_home = tempfile::tempdir().expect("create codex home");