From 6df44f08dd556ec2f2554561d53bc75f6f1f4313 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Fri, 30 Oct 2026 18:23:26 +0000 Subject: [PATCH] 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 | 1 + codex-rs/core/src/codex_thread.rs | 14 + codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/session/input_queue.rs | 80 +++-- codex-rs/core/src/thread_input.rs | 64 ++++ .../tools/handlers/multi_agents_v2/wait.rs | 50 +--- codex-rs/core/tests/suite/pending_input.rs | 278 ++++++++++++++++++ 7 files changed, 425 insertions(+), 64 deletions(-) create mode 100644 codex-rs/core/src/thread_input.rs diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 97cb3f90bd..5934637065 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/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; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index bdce74a843..e694937ba9 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/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 diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 45becc3d82..51573406cb 100644 --- a/codex-rs/core/src/lib.rs +++ b/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; diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 923f9f5628..a22fc2684f 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/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, - /// 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, } /// Session-scoped pending input storage and active-turn mailbox delivery coordination. @@ -378,15 +394,16 @@ impl InputQueue { turn_state: &Mutex, input: Vec, ) { + 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 { - 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] diff --git a/codex-rs/core/src/thread_input.rs b/codex-rs/core/src/thread_input.rs new file mode 100644 index 0000000000..d68f03649d --- /dev/null +++ b/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, + Option, +) { + 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, + pending: Option, +) -> 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, + } +} 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 83458d1891..703da4f8f2 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,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, - pending_activity: Option, - 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, - } -} diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 30fbb97704..5464800cd2 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/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::>(); + 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![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!["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")]