use crate::state::ActiveTurn;use crate::state::MailboxDeliveryPhase;use crate::state::TurnState;use codex_protocol::models::ResponseItem;use codex_protocol::protocol::InterAgentCommunication;use codex_protocol::user_input::UserInput;use std::collections::VecDeque;use std::sync::Arc;use tokio::sync::Mutex;use tokio::sync::watch;#[derive(Clone, Debug, PartialEq)]pub(crate) enum TurnInput { UserInput { content: Vec<UserInput>, client_id: Option<String>, }, ResponseItem(ResponseItem), InterAgentCommunication(InterAgentCommunication),}#[derive(Clone, Copy, Debug, Eq, PartialEq)]pub(crate) enum InputQueueActivity { Mailbox, Steer,}/// Turn-local pending input storage owned by the input queue flow.#[derive(Default)]pub(crate) struct TurnInputQueue { items: Vec<TurnInput>,}/// Session-scoped pending input storage and active-turn mailbox delivery coordination.pub(crate) struct InputQueue { activity_tx: watch::Sender<InputQueueActivity>, mailbox_pending_mails: Mutex<VecDeque<InterAgentCommunication>>,}impl InputQueue { pub(crate) fn new() -> Self { let (activity_tx, _) = watch::channel(InputQueueActivity::Mailbox); Self { activity_tx, mailbox_pending_mails: Mutex::new(VecDeque::new()), } } pub(crate) async fn subscribe_activity( &self, turn_state: Option<&Mutex<TurnState>>, ) -> ( watch::Receiver<InputQueueActivity>, Option<InputQueueActivity>, ) { let activity_rx = self.activity_tx.subscribe(); let has_pending_steer = if let Some(turn_state) = turn_state { turn_state.lock().await.pending_input.has_user_input() } else { false }; let pending_activity = if has_pending_steer { Some(InputQueueActivity::Steer) } else if self.has_pending_mailbox_items().await { Some(InputQueueActivity::Mailbox) } else { None }; (activity_rx, pending_activity) } pub(crate) async fn enqueue_mailbox_communication( &self, communication: InterAgentCommunication, ) { self.mailbox_pending_mails .lock() .await .push_back(communication); self.activity_tx.send_replace(InputQueueActivity::Mailbox); } pub(crate) async fn has_pending_mailbox_items(&self) -> bool { !self.mailbox_pending_mails.lock().await.is_empty() } pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool { self.mailbox_pending_mails .lock() .await .iter() .any(|mail| mail.trigger_turn) } pub(crate) async fn drain_mailbox_input_items(&self) -> Vec<TurnInput> { self.mailbox_pending_mails .lock() .await .drain(..) .map(TurnInput::InterAgentCommunication) .collect() } pub(crate) async fn turn_state_for_sub_id( &self, active_turn: &Mutex<Option<ActiveTurn>>, sub_id: &str, ) -> Option<Arc<Mutex<TurnState>>> { let active = active_turn.lock().await; active.as_ref().and_then(|active_turn| { active_turn .task .as_ref() .is_some_and(|task| task.turn_context.sub_id == sub_id) .then(|| Arc::clone(&active_turn.turn_state)) }) } /// Clear any pending waiters and input buffered for the current turn. pub(crate) async fn clear_pending(&self, active_turn: &ActiveTurn) { let mut turn_state = active_turn.turn_state.lock().await; turn_state.clear_pending_waiters(); turn_state.pending_input.items.clear(); } pub(crate) async fn defer_mailbox_delivery_to_next_turn( &self, active_turn: &Mutex<Option<ActiveTurn>>, sub_id: &str, ) { let turn_state = self.turn_state_for_sub_id(active_turn, sub_id).await; let Some(turn_state) = turn_state else { return; }; let mut turn_state = turn_state.lock().await; if !turn_state.pending_input.items.is_empty() { return; } turn_state.set_mailbox_delivery_phase(MailboxDeliveryPhase::NextTurn); } pub(crate) async fn accept_mailbox_delivery_for_current_turn( &self, active_turn: &Mutex<Option<ActiveTurn>>, sub_id: &str, ) { let turn_state = self.turn_state_for_sub_id(active_turn, sub_id).await; let Some(turn_state) = turn_state else { return; }; self.accept_mailbox_delivery_for_turn_state(turn_state.as_ref()) .await; } pub(super) async fn accept_mailbox_delivery_for_turn_state( &self, turn_state: &Mutex<TurnState>, ) { turn_state .lock() .await .accept_mailbox_delivery_for_current_turn(); } pub(super) async fn extend_pending_input_and_accept_mailbox_delivery_for_turn_state( &self, turn_state: &Mutex<TurnState>, input: Vec<TurnInput>, ) { { let mut turn_state = turn_state.lock().await; turn_state.pending_input.items.extend(input); turn_state.accept_mailbox_delivery_for_current_turn(); } self.activity_tx.send_replace(InputQueueActivity::Steer); } pub(crate) async fn extend_pending_input_for_turn_state( &self, turn_state: &Mutex<TurnState>, input: Vec<TurnInput>, ) { turn_state.lock().await.pending_input.items.extend(input); } pub(crate) async fn take_pending_input_for_turn_state( &self, turn_state: &Mutex<TurnState>, ) -> Vec<TurnInput> { turn_state.lock().await.pending_input.items.split_off(0) } #[expect( clippy::await_holding_invalid_type, reason = "active turn checks and turn state updates must remain atomic" )] pub(crate) async fn get_pending_input( &self, active_turn: &Mutex<Option<ActiveTurn>>, ) -> Vec<TurnInput> { let (pending_input, accepts_mailbox_delivery) = { let mut active = active_turn.lock().await; match active.as_mut() { Some(active_turn) => { let mut turn_state = active_turn.turn_state.lock().await; ( turn_state.pending_input.items.split_off(0), turn_state.accepts_mailbox_delivery_for_current_turn(), ) } None => (Vec::new(), true), } }; if !accepts_mailbox_delivery { return pending_input; } let mailbox_items = self.drain_mailbox_input_items().await.into_iter(); if pending_input.is_empty() { mailbox_items.collect() } else { let mut pending_input = pending_input; pending_input.extend(mailbox_items); pending_input } } #[expect( clippy::await_holding_invalid_type, reason = "active turn checks and turn state reads must remain atomic" )] pub(crate) async fn has_pending_input(&self, active_turn: &Mutex<Option<ActiveTurn>>) -> bool { let (has_turn_pending_input, accepts_mailbox_delivery) = { let active = active_turn.lock().await; match active.as_ref() { Some(active_turn) => { let turn_state = active_turn.turn_state.lock().await; ( !turn_state.pending_input.items.is_empty(), turn_state.accepts_mailbox_delivery_for_current_turn(), ) } None => (false, true), } }; if has_turn_pending_input { return true; } if !accepts_mailbox_delivery { return false; } self.has_pending_mailbox_items().await }}impl TurnInputQueue { fn has_user_input(&self) -> bool { self.items .iter() .any(|input| matches!(input, TurnInput::UserInput { .. })) }}#[cfg(test)]mod tests { use super::*; use codex_protocol::AgentPath; use pretty_assertions::assert_eq; fn make_mail( author: AgentPath, recipient: AgentPath, content: &str, trigger_turn: bool, ) -> InterAgentCommunication { InterAgentCommunication::new( author, recipient, Vec::new(), content.to_string(), trigger_turn, ) } #[tokio::test] async fn input_queue_notifies_mailbox_subscribers() { let input_queue = InputQueue::new(); let (mut activity_rx, pending_activity) = input_queue.subscribe_activity(/*turn_state*/ None).await; assert_eq!(pending_activity, None); input_queue .enqueue_mailbox_communication(make_mail( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), "one", /*trigger_turn*/ false, )) .await; input_queue .enqueue_mailbox_communication(make_mail( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), "two", /*trigger_turn*/ false, )) .await; activity_rx.changed().await.expect("mailbox update"); assert_eq!( *activity_rx.borrow_and_update(), InputQueueActivity::Mailbox ); } #[tokio::test] async fn input_queue_notifies_steer_subscribers() { let input_queue = InputQueue::new(); let turn_state = Mutex::new(TurnState::default()); let (mut activity_rx, pending_activity) = input_queue.subscribe_activity(Some(&turn_state)).await; assert_eq!(pending_activity, None); input_queue .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( &turn_state, vec![TurnInput::UserInput { content: vec![UserInput::Text { text: "steer".to_string(), text_elements: Vec::new(), }], client_id: None, }], ) .await; activity_rx.changed().await.expect("steer update"); assert_eq!(*activity_rx.borrow_and_update(), InputQueueActivity::Steer); } #[tokio::test] async fn input_queue_reports_already_pending_steer() { let input_queue = InputQueue::new(); let turn_state = Mutex::new(TurnState::default()); input_queue .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( &turn_state, vec![TurnInput::UserInput { content: vec![UserInput::Text { text: "already pending".to_string(), text_elements: Vec::new(), }], client_id: None, }], ) .await; let (_activity_rx, pending_activity) = input_queue.subscribe_activity(Some(&turn_state)).await; assert_eq!(pending_activity, Some(InputQueueActivity::Steer)); } #[tokio::test] async fn input_queue_drains_mailbox_in_delivery_order() { let input_queue = InputQueue::new(); let mail_one = make_mail( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), "one", /*trigger_turn*/ false, ); let mail_two = make_mail( AgentPath::try_from("/root/worker").expect("agent path"), AgentPath::root(), "two", /*trigger_turn*/ true, ); input_queue .enqueue_mailbox_communication(mail_one.clone()) .await; input_queue .enqueue_mailbox_communication(mail_two.clone()) .await; assert_eq!( input_queue.drain_mailbox_input_items().await, vec![ TurnInput::InterAgentCommunication(mail_one), TurnInput::InterAgentCommunication(mail_two) ] ); assert!(!input_queue.has_pending_mailbox_items().await); } #[tokio::test] async fn input_queue_tracks_pending_trigger_turn_mail() { let input_queue = InputQueue::new(); input_queue .enqueue_mailbox_communication(make_mail( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), "queued", /*trigger_turn*/ false, )) .await; assert!(!input_queue.has_trigger_turn_mailbox_items().await); input_queue .enqueue_mailbox_communication(make_mail( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), "wake", /*trigger_turn*/ true, )) .await; assert!(input_queue.has_trigger_turn_mailbox_items().await); }}