Codex Handbook
core/src/session/input_queue.rs 420 lines
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);    }}