fix: pending input race

codex-rs/core/src/agent/mailbox.rs

@@ -19,6 +19,11 @@ pub(crate) struct MailboxReceiver {
     pending_mails: VecDeque<InterAgentCommunication>,
 }
 
+pub(crate) struct MailboxDrain {
+    pub(crate) mails: Vec<InterAgentCommunication>,
+    pub(crate) saw_trigger_turn: bool,
+}
+
 impl Mailbox {
     pub(crate) fn new() -> (Self, MailboxReceiver) {
         let (tx, rx) = mpsc::unbounded_channel();
@@ -65,9 +70,14 @@ impl MailboxReceiver {
         self.pending_mails.iter().any(|mail| mail.trigger_turn)
     }
 
-    pub(crate) fn drain(&mut self) -> Vec<InterAgentCommunication> {
+    pub(crate) fn drain(&mut self) -> MailboxDrain {
         self.sync_pending_mails();
-        self.pending_mails.drain(..).collect()
+        let mails = self.pending_mails.drain(..).collect::<Vec<_>>();
+        let saw_trigger_turn = mails.iter().any(|mail| mail.trigger_turn);
+        MailboxDrain {
+            mails,
+            saw_trigger_turn,
+        }
     }
 }
 
@@ -134,28 +144,38 @@ mod tests {
         mailbox.send(mail_one.clone());
         mailbox.send(mail_two.clone());
 
-        assert_eq!(receiver.drain(), vec![mail_one, mail_two]);
+        let drain = receiver.drain();
+        assert_eq!(drain.mails, vec![mail_one, mail_two]);
+        assert!(!drain.saw_trigger_turn);
         assert!(!receiver.has_pending());
     }
 
     #[tokio::test]
-    async fn mailbox_tracks_pending_trigger_turn_mail() {
+    async fn mailbox_drain_tracks_pending_trigger_turn_mail() {
         let (mailbox, mut receiver) = Mailbox::new();
 
-        mailbox.send(make_mail(
+        let queue_only_mail = make_mail(
             AgentPath::root(),
             AgentPath::try_from("/root/worker").expect("agent path"),
             "queued",
             /*trigger_turn*/ false,
-        ));
-        assert!(!receiver.has_pending_trigger_turn());
+        );
+        mailbox.send(queue_only_mail.clone());
 
-        mailbox.send(make_mail(
+        let drain = receiver.drain();
+        assert_eq!(drain.mails, vec![queue_only_mail]);
+        assert!(!drain.saw_trigger_turn);
+
+        let trigger_turn_mail = make_mail(
             AgentPath::root(),
             AgentPath::try_from("/root/worker").expect("agent path"),
             "wake",
             /*trigger_turn*/ true,
-        ));
-        assert!(receiver.has_pending_trigger_turn());
+        );
+        mailbox.send(trigger_turn_mail.clone());
+
+        let drain = receiver.drain();
+        assert_eq!(drain.mails, vec![trigger_turn_mail]);
+        assert!(drain.saw_trigger_turn);
     }
 }

codex-rs/core/src/agent/mod.rs

@@ -8,6 +8,7 @@ pub(crate) mod status;
 pub(crate) use codex_protocol::protocol::AgentStatus;
 pub(crate) use control::AgentControl;
 pub(crate) use mailbox::Mailbox;
+pub(crate) use mailbox::MailboxDrain;
 pub(crate) use mailbox::MailboxReceiver;
 pub(crate) use registry::exceeds_thread_spawn_depth_limit;
 pub(crate) use registry::next_thread_spawn_depth;

codex-rs/core/src/codex.rs

@@ -9,6 +9,7 @@ use std::sync::atomic::AtomicU64;
 use crate::agent::AgentControl;
 use crate::agent::AgentStatus;
 use crate::agent::Mailbox;
+use crate::agent::MailboxDrain;
 use crate::agent::MailboxReceiver;
 use crate::agent::agent_status_from_event;
 use crate::agent::status::is_final;
@@ -4052,6 +4053,10 @@ impl Session {
         self.mailbox.send(communication);
     }
 
+    pub(crate) async fn drain_mailbox(&self) -> MailboxDrain {
+        self.mailbox_rx.lock().await.drain()
+    }
+
     pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool {
         self.mailbox_rx.lock().await.has_pending_trigger_turn()
     }
@@ -4080,9 +4085,9 @@ impl Session {
             }
         };
         let mailbox_items = {
-            let mut mailbox_rx = self.mailbox_rx.lock().await;
-            mailbox_rx
-                .drain()
+            self.drain_mailbox()
+                .await
+                .mails
                 .into_iter()
                 .map(|mail| mail.to_response_input_item())
                 .collect::<Vec<_>>()
@@ -4099,7 +4104,6 @@ impl Session {
     }
 
     /// Queue response items to be injected into the next active turn created for this session.
-    #[cfg(test)]
     pub(crate) async fn queue_response_items_for_next_turn(&self, items: Vec<ResponseInputItem>) {
         if items.is_empty() {
             return;

codex-rs/core/src/codex_tests.rs

@@ -54,6 +54,7 @@ use codex_execpolicy::NetworkRuleProtocol;
 use codex_execpolicy::Policy;
 use codex_network_proxy::NetworkProxyConfig;
 use codex_otel::TelemetryAuthMode;
+use codex_protocol::AgentPath;
 use codex_protocol::config_types::CollaborationMode;
 use codex_protocol::config_types::ModeKind;
 use codex_protocol::config_types::Settings;
@@ -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,40 @@ 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 on_task_finished_records_queue_only_mailbox_mail_without_restart() {
+    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: false,
+        },
+    )
+    .await;
+
+    let communication = InterAgentCommunication::new(
+        AgentPath::try_from("/root/worker").expect("worker path should parse"),
+        AgentPath::root(),
+        Vec::new(),
+        "queued completion".to_string(),
+        /*trigger_turn*/ false,
+    );
+    sess.enqueue_mailbox_communication(communication.clone());
+
+    sess.on_task_finished(Arc::clone(&tc), /*last_agent_message*/ None)
+        .await;
+    sleep(Duration::from_millis(50)).await;
+
+    assert!(!sess.has_pending_input().await);
+    assert!(sess.active_turn.lock().await.is_none());
+    assert_eq!(
+        sess.clone_history().await.raw_items(),
+        &[ResponseItem::from(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/tasks/mod.rs

@@ -356,6 +356,20 @@ impl Session {
             turn_tool_calls = ts.tool_calls;
             token_usage_at_turn_start = Some(ts.token_usage_at_turn_start.clone());
         }
+
+        let mailbox_drain = self.drain_mailbox().await;
+        let mut next_turn_mailbox_input = Vec::new();
+        let mut queue_mail_for_next_turn = false;
+        for mail in mailbox_drain.mails {
+            let mailbox_input = mail.to_response_input_item();
+            queue_mail_for_next_turn |= mail.trigger_turn;
+            if queue_mail_for_next_turn {
+                next_turn_mailbox_input.push(mailbox_input);
+            } else {
+                pending_input.push(mailbox_input);
+            }
+        }
+
         if !pending_input.is_empty() {
             for pending_input_item in pending_input {
                 match inspect_pending_input(self, &turn_context, pending_input_item).await {
@@ -370,6 +384,11 @@ impl Session {
                 }
             }
         }
+        if mailbox_drain.saw_trigger_turn {
+            self.queue_response_items_for_next_turn(next_turn_mailbox_input)
+                .await;
+        }
+
         // Emit token usage metrics.
         if let Some(token_usage_at_turn_start) = token_usage_at_turn_start {
             // TODO(jif): drop this