From ee40dddbf6a50a6f0641180ca299be7a3a03fd22 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 15 Jun 2026 19:08:15 +0100 Subject: [PATCH] core: let steer interrupt wait_agent (#28341) ## Why `wait_agent` can block for a long timeout while waiting for sub-agent mailbox activity. Although same-turn user steer is accepted during that tool call, the input remains pending until the wait returns, so an explicit request to change direction can appear unresponsive. ## What changed - Notify active `wait_agent` calls when user input is steered into the current turn. - Check for already-pending steer input when subscribing so input that races with tool startup is not missed. - Distinguish mailbox activity, steered input, and timeout outcomes, returning `Wait interrupted by new input.` for the steer path. - Update the `wait_agent` tool description to document the early-return behavior. ## Testing - `just test -p codex-core input_queue_` - `just test -p codex-core wait_agent` The coverage includes steer notification before and after subscription, plus an end-to-end test that verifies the interrupted wait result and steered user input are both included exactly once in the follow-up model request. --- codex-rs/core/src/session/input_queue.rs | 116 +++++++++++++++--- codex-rs/core/src/session/mod.rs | 1 + .../src/tools/handlers/multi_agents_spec.rs | 2 +- .../tools/handlers/multi_agents_v2/wait.rs | 55 ++++++--- codex-rs/core/tests/suite/pending_input.rs | 73 +++++++++++ 5 files changed, 216 insertions(+), 31 deletions(-) diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 199ffd004..e46a71892 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -19,6 +19,12 @@ pub(crate) enum TurnInput { 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 { @@ -27,25 +33,40 @@ pub(crate) struct TurnInputQueue { /// Session-scoped pending input storage and active-turn mailbox delivery coordination. pub(crate) struct InputQueue { - mailbox_tx: watch::Sender<()>, + activity_tx: watch::Sender, mailbox_pending_mails: Mutex>, } impl InputQueue { pub(crate) fn new() -> Self { - let (mailbox_tx, _) = watch::channel(()); + let (activity_tx, _) = watch::channel(InputQueueActivity::Mailbox); Self { - mailbox_tx, + activity_tx, mailbox_pending_mails: Mutex::new(VecDeque::new()), } } - pub(crate) async fn subscribe_mailbox(&self) -> watch::Receiver<()> { - let mut mailbox_rx = self.mailbox_tx.subscribe(); - if self.has_pending_mailbox_items().await { - mailbox_rx.mark_changed(); - } - mailbox_rx + pub(crate) async fn subscribe_activity( + &self, + turn_state: Option<&Mutex>, + ) -> ( + watch::Receiver, + Option, + ) { + 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( @@ -56,7 +77,7 @@ impl InputQueue { .lock() .await .push_back(communication); - self.mailbox_tx.send_replace(()); + self.activity_tx.send_replace(InputQueueActivity::Mailbox); } pub(crate) async fn has_pending_mailbox_items(&self) -> bool { @@ -146,9 +167,12 @@ impl InputQueue { turn_state: &Mutex, input: Vec, ) { - let mut turn_state = turn_state.lock().await; - turn_state.pending_input.items.extend(input); - turn_state.accept_mailbox_delivery_for_current_turn(); + { + 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( @@ -228,6 +252,14 @@ impl InputQueue { } } +impl TurnInputQueue { + fn has_user_input(&self) -> bool { + self.items + .iter() + .any(|input| matches!(input, TurnInput::UserInput { .. })) + } +} + #[cfg(test)] mod tests { use super::*; @@ -252,7 +284,9 @@ mod tests { #[tokio::test] async fn input_queue_notifies_mailbox_subscribers() { let input_queue = InputQueue::new(); - let mut mailbox_rx = input_queue.subscribe_mailbox().await; + 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( @@ -271,7 +305,59 @@ mod tests { )) .await; - mailbox_rx.changed().await.expect("mailbox update"); + 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] diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8e6702088..ed9e58113 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -220,6 +220,7 @@ use self::config_lock::validate_config_lock_if_configured; #[cfg(test)] use self::handlers::submission_dispatch_span; use self::handlers::submission_loop; +pub(crate) use self::input_queue::InputQueueActivity; pub(crate) use self::input_queue::TurnInput; pub(crate) use self::input_queue::TurnInputQueue; use self::review::spawn_review_thread; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs index 98f68b145..70f71f311 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs @@ -255,7 +255,7 @@ pub fn create_wait_agent_tool_v1(options: WaitAgentTimeoutOptions) -> ToolSpec { pub fn create_wait_agent_tool_v2(options: WaitAgentTimeoutOptions) -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "wait_agent".to_string(), - description: "Wait for a mailbox update from any live agent, including queued messages and final-status notifications. Does not return the content; returns either a summary of which agents have updates (if any), or a timeout summary if no mailbox update arrives before the deadline." + description: "Wait for a mailbox update from any live agent, including queued messages and final-status notifications. The wait also ends early when new user input is steered into the active turn. Does not return the content; returns either a summary of which agents have updates (if any), an interruption summary for steered input, or a timeout summary if no activity arrives before the deadline." .to_string(), strict: false, defer_loading: None, diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs index 36e5aa24f..e97f47238 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs @@ -1,4 +1,5 @@ use super::*; +use crate::session::InputQueueActivity; use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions; use crate::tools::handlers::multi_agents_spec::create_wait_agent_tool_v2; use crate::turn_timing::now_unix_timestamp_ms; @@ -65,7 +66,14 @@ impl Handler { None => default_timeout_ms, }; - let mut mailbox_rx = session.input_queue.subscribe_mailbox().await; + let turn_state = session + .input_queue + .turn_state_for_sub_id(&session.active_turn, &turn.sub_id) + .await; + let (mut activity_rx, pending_activity) = session + .input_queue + .subscribe_activity(turn_state.as_deref()) + .await; session .send_event( @@ -82,8 +90,8 @@ impl Handler { .await; let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64); - let timed_out = !wait_for_mailbox_change(&mut mailbox_rx, deadline).await; - let result = WaitAgentResult::from_timed_out(timed_out); + let outcome = wait_for_activity(&mut activity_rx, pending_activity, deadline).await; + let result = WaitAgentResult::from_outcome(outcome); session .send_event( @@ -122,15 +130,15 @@ pub(crate) struct WaitAgentResult { } impl WaitAgentResult { - fn from_timed_out(timed_out: bool) -> Self { - let message = if timed_out { - "Wait timed out." - } else { - "Wait completed." + fn from_outcome(outcome: WaitOutcome) -> Self { + let message = match outcome { + WaitOutcome::MailboxActivity => "Wait completed.", + WaitOutcome::Steered => "Wait interrupted by new input.", + WaitOutcome::TimedOut => "Wait timed out.", }; Self { message: message.to_string(), - timed_out, + timed_out: outcome == WaitOutcome::TimedOut, } } } @@ -153,12 +161,29 @@ impl ToolOutput for WaitAgentResult { } } -async fn wait_for_mailbox_change( - mailbox_rx: &mut tokio::sync::watch::Receiver<()>, +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum WaitOutcome { + MailboxActivity, + Steered, + TimedOut, +} + +async fn wait_for_activity( + activity_rx: &mut tokio::sync::watch::Receiver, + pending_activity: Option, deadline: Instant, -) -> bool { - match timeout_at(deadline, mailbox_rx.changed()).await { - Ok(Ok(())) => true, - Ok(Err(_)) | Err(_) => false, +) -> WaitOutcome { + if let Some(activity) = pending_activity { + return match activity { + InputQueueActivity::Mailbox => WaitOutcome::MailboxActivity, + InputQueueActivity::Steer => WaitOutcome::Steered, + }; + } + match timeout_at(deadline, activity_rx.changed()).await { + Ok(Ok(())) => match *activity_rx.borrow_and_update() { + InputQueueActivity::Mailbox => WaitOutcome::MailboxActivity, + InputQueueActivity::Steer => WaitOutcome::Steered, + }, + Ok(Err(_)) | Err(_) => WaitOutcome::TimedOut, } } diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 55a00a211..60775dbe9 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -2,6 +2,7 @@ use core_test_support::test_codex::local_selections; use std::sync::Arc; use codex_core::CodexThread; +use codex_features::Feature; use codex_protocol::AgentPath; use codex_protocol::items::TurnItem; use codex_protocol::models::PermissionProfile; @@ -206,6 +207,78 @@ async fn wait_for_turn_complete(codex: &CodexThread) { wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn steer_interrupts_wait_agent_and_is_sent_in_follow_up_request() { + const WAIT_CALL_ID: &str = "wait-call"; + const INITIAL_PROMPT: &str = "wait for an agent"; + const STEER_PROMPT: &str = "stop waiting and continue"; + + let first_chunks = vec![ + chunk(ev_response_created("resp-1")), + chunk(ev_function_call( + WAIT_CALL_ID, + "wait_agent", + r#"{"timeout_ms":10000}"#, + )), + chunk(ev_completed("resp-1")), + ]; + let (server, _completions) = + start_streaming_sse_server(vec![first_chunks, response_completed_chunks("resp-2")]).await; + let codex = test_codex() + .with_model("gpt-5.4") + .with_config(|config| { + config + .features + .enable(Feature::MultiAgentV2) + .expect("test config should allow feature update"); + }) + .build_with_streaming_server(&server) + .await + .expect("build Codex test session") + .codex; + + submit_user_input(&codex, INITIAL_PROMPT).await; + wait_for_event(&codex, |event| { + matches!(event, EventMsg::CollabWaitingBegin(_)) + }) + .await; + + steer_user_input(&codex, STEER_PROMPT).await; + wait_for_turn_complete(&codex).await; + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let second: Value = from_slice(&requests[1]).expect("parse second request"); + let relevant_user_input = message_input_texts(&second, "user") + .into_iter() + .filter(|text| text == INITIAL_PROMPT || text == STEER_PROMPT) + .collect::>(); + assert_eq!( + relevant_user_input, + vec![INITIAL_PROMPT.to_string(), STEER_PROMPT.to_string()] + ); + let wait_output = second["input"] + .as_array() + .expect("second request input") + .iter() + .find(|item| { + item.get("type").and_then(Value::as_str) == Some("function_call_output") + && item.get("call_id").and_then(Value::as_str) == Some(WAIT_CALL_ID) + }) + .and_then(|item| item.get("output")) + .and_then(Value::as_str) + .expect("wait_agent output"); + assert_eq!( + serde_json::from_str::(wait_output).expect("parse wait_agent output"), + json!({ + "message": "Wait interrupted by new input.", + "timed_out": false, + }) + ); + + server.shutdown().await; +} + fn assert_two_responses_input_snapshot(snapshot_name: &str, requests: &[Vec]) { assert_eq!(requests.len(), 2); let options = ContextSnapshotOptions::default().strip_capability_instructions();