Expose non-consuming input activity waits to hosts (#58411)

## Why

Host-owned tools need to wait for agent messages or steering input with the same behavior as native `wait_agent`, including when native multi-agent support is disabled.

## What changed

- Add `CodexThread::wait_for_input_activity` and export `InputWaitOutcome` through core and core-api, sharing the wait implementation with `wait_agent`.
- Distinguish notifying agent messages from steering input, with steering taking precedence when both are pending. Passive response items do not wake waits.
- Observe input without consuming it. Waits remain active after the original turn ends, support independent observers, and can be cancelled by dropping the future.

## Testing

Add coverage for pending and incoming input, agent-message classification, cancellation, independent observers, timeouts, and input arriving after the original turn completes or is interrupted.

GitOrigin-RevId: d0c0bd776c3b6fd4ba3021603d87eaaf1d93976f

codex-rs/core-api/src/lib.rs

@@ -39,6 +39,7 @@ pub use codex_core::EnvironmentConfig;
 pub use codex_core::EnvironmentMcpPolicy;
 pub use codex_core::EnvironmentNetworkPolicy;
 pub use codex_core::ForkSnapshot;
+pub use codex_core::InputWaitOutcome;
 pub use codex_core::LoadedAgentsMd;
 pub use codex_core::McpManager;
 pub use codex_core::NetworkDomainPermission;

codex-rs/core/src/codex_thread.rs

@@ -460,6 +460,20 @@ impl CodexThread {
             .await
     }
 
+    /// Waits for input activity without consuming input, following native `wait_agent` behavior.
+    ///
+    /// The named turn supplies any already-buffered steering input. Subsequent activity is
+    /// thread-wide: ending or replacing that turn does not end this wait. Mailbox input
+    /// already pending is also observed. Callers own cancellation; dropping this future
+    /// does not consume input or affect execution. Each call observes independently.
+    pub async fn wait_for_input_activity(
+        &self,
+        turn_id: &str,
+        timeout: std::time::Duration,
+    ) -> crate::InputWaitOutcome {
+        crate::thread_input::wait_for_input_activity(&self.session, turn_id, timeout).await
+    }
+
     /// Submits turn input without requiring the caller to inspect thread state.
     ///
     /// The result describes whether Core started a turn, steered an active

codex-rs/core/src/lib.rs

@@ -31,6 +31,8 @@ pub use codex_protocol::turn_input::TurnStartOptions;
 pub use responses_metadata::CodexResponsesMetadata;
 pub use turn_metadata::detached_memory_responses_metadata;
 mod codex_thread;
+mod thread_input;
+pub use thread_input::InputWaitOutcome;
 mod turn_extension_data;
 pub use turn_extension_data::WithTurnExtensionData;
 mod compact_model_fallback;

codex-rs/core/src/session/input_queue.rs

@@ -82,18 +82,34 @@ mod turn_input_response_item {
     }
 }
 
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+/// Steering takes precedence when several kinds of input are pending together.
+#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
 pub(crate) enum InputQueueActivity {
     Mailbox,
     Steer,
 }
 
+impl TurnInput {
+    fn wait_activity(&self) -> InputQueueActivity {
+        match self {
+            Self::InterAgentCommunication(_)
+            | Self::ResponseItem(ResponseItemEnvelope {
+                item: ResponseItem::AgentMessage { .. },
+                ..
+            }) => InputQueueActivity::Mailbox,
+            Self::UserInput { .. } | Self::FunctionCallOutput(_) | Self::ResponseItem(_) => {
+                InputQueueActivity::Steer
+            }
+        }
+    }
+}
+
 /// Turn-local pending input storage owned by the input queue flow.
 #[derive(Default)]
 pub(crate) struct TurnInputQueue {
     items: Vec<TurnInput>,
-    /// Response items queued through a notifying injection, unlike passive context.
-    has_pending_injection: bool,
+    /// Activity from notifying writes. Passive response items do not wake a wait.
+    notified_activity: Option<InputQueueActivity>,
 }
 
 /// Session-scoped pending input storage and active-turn mailbox delivery coordination.
@@ -378,15 +394,16 @@ impl InputQueue {
         turn_state: &Mutex<TurnState>,
         input: Vec<TurnInput>,
     ) {
+        let activity = input.iter().map(TurnInput::wait_activity).max();
         {
             let mut turn_state = turn_state.lock().await;
-            turn_state.pending_input.has_pending_injection |= input
-                .iter()
-                .any(|input| matches!(input, TurnInput::ResponseItem(_)));
+            turn_state.pending_input.notified_activity =
+                turn_state.pending_input.notified_activity.max(activity);
             turn_state.pending_input.items.extend(input);
             turn_state.accept_mailbox_delivery_for_current_turn();
         }
-        self.activity_tx.send_replace(InputQueueActivity::Steer);
+        self.activity_tx
+            .send_replace(activity.unwrap_or(InputQueueActivity::Steer));
     }
 
     pub(crate) async fn extend_pending_input_for_turn_state(
@@ -482,24 +499,15 @@ impl TurnInputQueue {
     }
 
     fn pending_activity(&self) -> Option<InputQueueActivity> {
-        if self.has_pending_injection
-            || self.items.iter().any(|input| {
-                matches!(
-                    input,
-                    TurnInput::UserInput { .. } | TurnInput::FunctionCallOutput(_)
-                )
-            })
-        {
-            Some(InputQueueActivity::Steer)
-        } else if self
-            .items
-            .iter()
-            .any(|input| matches!(input, TurnInput::InterAgentCommunication(_)))
-        {
-            Some(InputQueueActivity::Mailbox)
-        } else {
-            None
+        if self.notified_activity == Some(InputQueueActivity::Steer) {
+            return self.notified_activity;
         }
+        self.items
+            .iter()
+            .filter(|input| !matches!(input, TurnInput::ResponseItem(_)))
+            .map(TurnInput::wait_activity)
+            .chain(self.notified_activity)
+            .max()
     }
 }
 
@@ -612,7 +620,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn input_queue_notifies_steer_subscribers() {
+    async fn input_queue_distinguishes_new_agent_input_from_pending_steering() {
         let input_queue = InputQueue::new();
         let turn_state = Mutex::new(TurnState::default());
         let (mut activity_rx, pending_activity) =
@@ -635,6 +643,28 @@ mod tests {
 
         activity_rx.changed().await.expect("steer update");
         assert_eq!(*activity_rx.borrow_and_update(), InputQueueActivity::Steer);
+
+        let mail = make_mail(
+            AgentPath::root(),
+            AgentPath::try_from("/root/worker").expect("agent path"),
+            "follow-up",
+            /*trigger_turn*/ true,
+        );
+        input_queue
+            .extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
+                &turn_state,
+                vec![TurnInput::InterAgentCommunication(mail.into())],
+            )
+            .await;
+        activity_rx.changed().await.expect("mailbox update");
+        assert_eq!(
+            *activity_rx.borrow_and_update(),
+            InputQueueActivity::Mailbox
+        );
+        assert_eq!(
+            input_queue.subscribe_activity(Some(&turn_state)).await.1,
+            Some(InputQueueActivity::Steer)
+        );
     }
 
     #[tokio::test]

codex-rs/core/src/thread_input.rs

@@ -0,0 +1,64 @@
+//! Shared, non-consuming input-activity waits for native tools and hosts.
+
+use crate::session::InputQueueActivity;
+use crate::session::session::Session;
+use std::time::Duration;
+use tokio::sync::watch;
+
+/// Why a native or host input-activity wait finished.
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum InputWaitOutcome {
+    /// Agent mail was observed, including a typed agent message from a notifying injection.
+    MailboxActivity,
+    /// User steering or another kind of notifying input was observed.
+    Steered,
+    /// The deadline elapsed or the input-activity source closed.
+    TimedOut,
+}
+
+pub(crate) async fn wait_for_input_activity(
+    session: &Session,
+    turn_id: &str,
+    timeout: Duration,
+) -> InputWaitOutcome {
+    tokio::time::timeout(timeout, async {
+        let (activity, pending) = subscribe_input_activity(session, turn_id).await;
+        wait_for_activity(activity, pending).await
+    })
+    .await
+    .unwrap_or(InputWaitOutcome::TimedOut)
+}
+
+pub(crate) async fn subscribe_input_activity(
+    session: &Session,
+    turn_id: &str,
+) -> (
+    watch::Receiver<InputQueueActivity>,
+    Option<InputQueueActivity>,
+) {
+    let turn_state = session
+        .input_queue
+        .turn_state_for_sub_id(&session.active_turn, turn_id)
+        .await;
+    session
+        .input_queue
+        .subscribe_activity(turn_state.as_deref())
+        .await
+}
+
+pub(crate) async fn wait_for_activity(
+    mut activity: watch::Receiver<InputQueueActivity>,
+    pending: Option<InputQueueActivity>,
+) -> InputWaitOutcome {
+    let activity = match pending {
+        Some(activity) => activity,
+        None => match activity.changed().await {
+            Ok(()) => *activity.borrow_and_update(),
+            Err(_) => return InputWaitOutcome::TimedOut,
+        },
+    };
+    match activity {
+        InputQueueActivity::Mailbox => InputWaitOutcome::MailboxActivity,
+        InputQueueActivity::Steer => InputWaitOutcome::Steered,
+    }
+}

codex-rs/core/src/tools/handlers/multi_agents_v2/wait.rs

@@ -1,12 +1,13 @@
 use super::*;
-use crate::session::InputQueueActivity;
+use crate::thread_input::InputWaitOutcome as WaitOutcome;
+use crate::thread_input::subscribe_input_activity;
+use crate::thread_input::wait_for_activity;
 use crate::tools::handlers::multi_agents_spec::WaitAgentTimeoutOptions;
 use crate::tools::handlers::multi_agents_spec::create_wait_agent_tool_v2;
 use codex_tools::ToolSpec;
 use std::collections::HashMap;
 use std::time::Duration;
 use tokio::time::Instant;
-use tokio::time::timeout_at;
 
 #[derive(Default)]
 pub(crate) struct Handler {
@@ -64,14 +65,8 @@ impl Handler {
             None => default_timeout_ms,
         };
 
-        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;
+        let (activity_rx, pending_activity) =
+            subscribe_input_activity(&session, &turn.sub_id).await;
 
         session
             .emit_turn_item_started(
@@ -92,8 +87,12 @@ impl Handler {
             .await;
 
         let wait_started = Instant::now();
-        let deadline = wait_started + Duration::from_millis(timeout_ms as u64);
-        let outcome = wait_for_activity(&mut activity_rx, pending_activity, deadline).await;
+        let outcome = tokio::time::timeout(
+            Duration::from_millis(timeout_ms as u64),
+            wait_for_activity(activity_rx, pending_activity),
+        )
+        .await
+        .unwrap_or(WaitOutcome::TimedOut);
         // A completed wait may wake for a message, user input, or its timeout.
         // Dropped waits do not have an observed outcome and are not included.
         turn.session_telemetry.record_duration(
@@ -191,30 +190,3 @@ impl ToolOutput for WaitAgentResult {
         tool_output_code_mode_result(self, "wait_agent")
     }
 }
-
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-enum WaitOutcome {
-    MailboxActivity,
-    Steered,
-    TimedOut,
-}
-
-async fn wait_for_activity(
-    activity_rx: &mut tokio::sync::watch::Receiver<InputQueueActivity>,
-    pending_activity: Option<InputQueueActivity>,
-    deadline: Instant,
-) -> 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,
-    }
-}

codex-rs/core/tests/suite/pending_input.rs

@@ -1,7 +1,9 @@
 use core_test_support::test_codex::local_requests;
 use std::sync::Arc;
+use std::time::Duration;
 
 use codex_core::CodexThread;
+use codex_core::InputWaitOutcome;
 use codex_core::StartIfIdleSubmission;
 use codex_core::StartThreadOptions;
 use codex_core::TurnInput;
@@ -32,6 +34,7 @@ use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::CodexErrorInfo;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::InterAgentCommunication;
+use codex_protocol::protocol::MultiAgentVersion;
 use codex_protocol::protocol::Op;
 use codex_protocol::protocol::ThreadSettingsOverrides;
 use codex_protocol::turn_input::CyberAccessProgram;
@@ -786,6 +789,281 @@ async fn input_interrupts_wait_agent_and_is_sent_in_follow_up_request(input: Wai
     server.shutdown().await;
 }
 
+#[test_case(true, InputWaitOutcome::Steered; "already_pending")]
+#[test_case(false, InputWaitOutcome::Steered; "arrives_while_waiting")]
+#[test_case(false, InputWaitOutcome::MailboxActivity; "agent_message_while_waiting")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn plugin_wait_wakes_on_turn_input(
+    queue_before_wait: bool,
+    injection_outcome: InputWaitOutcome,
+) -> anyhow::Result<()> {
+    let server = responses::start_mock_server().await;
+    let test = test_codex()
+        .with_config(|config| {
+            config.agents_enabled = false;
+            config
+                .features
+                .disable(Feature::MultiAgentV2)
+                .expect("disable native MA");
+        })
+        .build_with_auto_env(&server)
+        .await?;
+    // The host owns this plugin tool; the thread has no native MA tools or runtime.
+    let thread = test
+        .thread_manager
+        .start_thread(StartThreadOptions {
+            dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec {
+                name: "plugin_wait".to_string(),
+                description: "Wait for worker results or new input.".to_string(),
+                input_schema: json!({"type": "object", "properties": {}}),
+                defer_loading: false,
+            })],
+            ..StartThreadOptions::new(test.config.clone())
+        })
+        .await?
+        .thread;
+    let first = responses::mount_sse_once(
+        &server,
+        responses::sse(vec![
+            ev_response_created("plugin-wait"),
+            ev_function_call("wait-call", "plugin_wait", "{}"),
+            ev_completed("plugin-wait"),
+        ]),
+    )
+    .await;
+    let TurnInputSubmission::Started { turn_id, .. } = thread
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "wait for the workers".to_string(),
+            text_elements: Vec::new(),
+        }]))
+        .await?
+    else {
+        panic!("initial input should start a turn");
+    };
+    let (notification, wait_message) = if injection_outcome == InputWaitOutcome::MailboxActivity {
+        (
+            json!({
+                "type": "agent_message", "author": "/root/worker", "recipient": "/root",
+                "content": [{"type": "input_text", "text": "worker is done"}],
+            }),
+            "Wait completed.",
+        )
+    } else {
+        (
+            json!({
+                "type": "function_call_output", "namespace": "local_ma",
+                "name": "worker_result", "output": "worker is done",
+            }),
+            "Wait interrupted by new input.",
+        )
+    };
+    let mut followups = Vec::new();
+    // Exercise injection and user steering in the same turn. The second wait
+    // must start clean after the model consumes the first delivery.
+    for (call_id, injected_input) in [
+        (
+            "wait-call",
+            Some(serde_json::from_value(notification.clone())?),
+        ),
+        ("wait-after-injection", None),
+    ] {
+        wait_for_event(&thread, |event| {
+            matches!(event, EventMsg::DynamicToolCallRequest(request) if request.call_id == call_id)
+        })
+        .await;
+        assert_eq!(
+            thread.multi_agent_version(),
+            Some(MultiAgentVersion::Disabled)
+        );
+        let response = if injected_input.is_some() {
+            responses::sse(vec![
+                ev_response_created("after-injection"),
+                ev_function_call("wait-after-injection", "plugin_wait", "{}"),
+                ev_completed("after-injection"),
+            ])
+        } else {
+            responses::sse_completed("followup")
+        };
+        followups.push(responses::mount_sse_once(&server, response).await);
+
+        // Dropping a wait must leave both the turn and its input untouched.
+        let mut cancelled_wait = Box::pin(thread.wait_for_input_activity(&turn_id, Duration::MAX));
+        assert!(futures::poll!(cancelled_wait.as_mut()).is_pending());
+        drop(cancelled_wait);
+        let wait = thread.wait_for_input_activity(&turn_id, Duration::from_secs(10));
+        let deliver = async {
+            if let Some(item) = injected_input {
+                thread.inject_response_items(vec![item]).await?;
+                anyhow::Ok((injection_outcome, wait_message))
+            } else {
+                steer_user_input(&thread, "stop waiting and work on X").await;
+                Ok((InputWaitOutcome::Steered, "Wait interrupted by new input."))
+            }
+        };
+        let (observed, (expected, message)) = if queue_before_wait {
+            let delivered = deliver.await?;
+            (wait.await, delivered)
+        } else {
+            tokio::try_join!(biased; async { anyhow::Ok(wait.await) }, deliver)?
+        };
+        assert_eq!(observed, expected);
+        // Another observer sees the same input: waiting never consumes it.
+        assert_eq!(
+            tokio::time::timeout(
+                Duration::from_secs(10),
+                thread.wait_for_input_activity(&turn_id, Duration::MAX),
+            )
+            .await?,
+            expected
+        );
+        thread
+            .submit(Op::DynamicToolResponse {
+                id: call_id.to_string(),
+                response: DynamicToolResponse {
+                    content_items: vec![DynamicToolCallOutputContentItem::InputText {
+                        text: message.to_string(),
+                    }],
+                    success: true,
+                },
+            })
+            .await?;
+    }
+    wait_for_turn_complete(&thread).await;
+    first.single_request();
+    let requests = followups
+        .iter()
+        .map(responses::ResponseMock::single_request)
+        .collect::<Vec<_>>();
+    for request in &requests {
+        assert_eq!(
+            request
+                .inputs_of_type(notification["type"].as_str().expect("item type"))
+                .into_iter()
+                .filter(|item| item["name"] == notification["name"])
+                .map(responses::strip_metadata_from_json)
+                .map(responses::strip_response_item_ids_from_json)
+                .collect::<Vec<_>>(),
+            vec![notification.clone()],
+        );
+    }
+    let request = requests.last().expect("final model request");
+    assert_eq!(
+        request
+            .message_input_texts("user")
+            .into_iter()
+            .filter(|text| matches!(
+                text.as_str(),
+                "wait for the workers" | "stop waiting and work on X"
+            ))
+            .collect::<Vec<_>>(),
+        vec!["wait for the workers", "stop waiting and work on X"],
+    );
+    assert!(
+        request.function_call_output("wait-call")["output"]
+            .as_str()
+            .expect("plugin output is text")
+            .contains(wait_message)
+    );
+    assert_eq!(
+        thread
+            .wait_for_input_activity(&turn_id, Duration::ZERO)
+            .await,
+        InputWaitOutcome::TimedOut
+    );
+    Ok(())
+}
+
+#[test_case(true; "interrupted")]
+#[test_case(false; "completed")]
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn host_input_wait_observes_activity_after_original_turn_ends(
+    interrupt: bool,
+) -> anyhow::Result<()> {
+    let (release_first, first_gate) = oneshot::channel();
+    let (release_second, second_gate) = oneshot::channel();
+    let (server, _) = start_streaming_sse_server(vec![
+        vec![gated_chunk(
+            first_gate,
+            vec![ev_response_created("first"), ev_completed("first")],
+        )],
+        vec![gated_chunk(
+            second_gate,
+            vec![ev_response_created("second"), ev_completed("second")],
+        )],
+    ])
+    .await;
+    let test = test_codex()
+        .with_config(|config| {
+            config.agents_enabled = false;
+            config
+                .features
+                .disable(Feature::MultiAgentV2)
+                .expect("disable native MA");
+            config
+                .features
+                .disable(Feature::InstantInterrupt)
+                .expect("keep the held response active");
+        })
+        .build_with_streaming_server(&server)
+        .await?;
+    let TurnInputSubmission::Started { turn_id, .. } = test
+        .codex
+        .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
+            text: "first task".to_string(),
+            text_elements: Vec::new(),
+        }]))
+        .await?
+    else {
+        panic!("initial input should start a turn");
+    };
+    server.wait_for_request_count(/*count*/ 1).await;
+    let wait = test
+        .codex
+        .wait_for_input_activity(&turn_id, Duration::from_secs(10));
+    tokio::pin!(wait);
+    let end_original_turn = async {
+        if interrupt {
+            test.codex.submit(Op::Interrupt).await?;
+            wait_for_event(&test.codex, |event| {
+                matches!(event, EventMsg::TurnAborted(_))
+            })
+            .await;
+        }
+        let _ = release_first.send(());
+        if !interrupt {
+            wait_for_turn_complete(&test.codex).await;
+        }
+        anyhow::Ok(())
+    };
+    tokio::select! {
+        biased;
+        outcome = &mut wait => panic!("wait ended before replacement input: {outcome:?}"),
+        result = end_original_turn => result?,
+    }
+
+    // Poll both observers before driving the replacement and while it acquires the turn lock.
+    let (original, late, ()) = tokio::join!(
+        biased;
+        wait,
+        test.codex.wait_for_input_activity(&turn_id, Duration::from_secs(10)),
+        async {
+            submit_user_input(&test.codex, "replacement task").await;
+            server.wait_for_request_count(/*count*/ 2).await;
+            steer_user_input(&test.codex, "input for the replacement").await;
+        },
+    );
+    assert_eq!(original, InputWaitOutcome::Steered);
+    assert_eq!(late, InputWaitOutcome::Steered);
+    test.codex.submit(Op::Interrupt).await?;
+    wait_for_event(&test.codex, |event| {
+        matches!(event, EventMsg::TurnAborted(_))
+    })
+    .await;
+    let _ = release_second.send(());
+    server.shutdown().await;
+    Ok(())
+}
+
 #[test_case(WaitInput::UserSteer; "user_steer")]
 #[test_case(WaitInput::PendingInjection; "already_pending_injection")]
 #[test_case(WaitInput::LiveInjection; "injection_while_waiting")]