diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 92d75390d..86105c67f 100644 --- a/codex-rs/core/src/codex.rs +++ b/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 { self.mailbox.subscribe() } @@ -4069,16 +4089,22 @@ impl Session { } pub async fn get_pending_input(&self) -> Vec { - 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( diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 8f1bba647..441095a5a 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/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; diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index b573123ad..3cc6200a3 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -29,6 +29,17 @@ pub(crate) struct ActiveTurn { pub(crate) turn_state: Arc>, } +/// 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>, pending_dynamic_tools: HashMap>, pending_input: Vec, + mailbox_delivery_phase: MailboxDeliveryPhase, granted_permissions: Option, 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)); diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 5cc9574db..90a340958 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/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 { match input { ResponseInputItem::FunctionCallOutput { call_id, output } => { diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index 23f52a068..5d61870b3 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/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) -> 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");