From 875de1ec111beed474f590f0d329f7b9cf1835e3 Mon Sep 17 00:00:00 2001 From: Friel Date: Sat, 27 Jun 2026 06:56:16 +0000 Subject: [PATCH] Add late steer integration coverage --- codex-rs/core/tests/suite/pending_input.rs | 156 +++++++++++++++++++++ 1 file changed, 156 insertions(+) diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 9c9852a276..7f3fbdcb1a 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -1,5 +1,9 @@ use core_test_support::test_codex::local_selections; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::sync::mpsc; +use std::time::Duration; use codex_core::CodexThread; use codex_core::config::CurrentTimeReminderConfig; @@ -39,6 +43,50 @@ use serde_json::Value; use serde_json::from_slice; use serde_json::json; use tokio::sync::oneshot; +use tracing::Subscriber; +use tracing::span::Id; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context; +use tracing_subscriber::prelude::*; +use tracing_subscriber::registry::LookupSpan; + +/// Pauses task finalization after `RegularTask::run` has made its last pending-input decision. +/// +/// The test thread uses this span boundary to submit a steer through `CodexThread`, exercising the +/// production gap between the regular task returning and `Session::on_task_finished` atomically +/// deciding whether to continue or finish the turn. +struct RegularTaskRunCloseBarrier { + armed: Arc, + target_thread: std::thread::ThreadId, + run_closed: mpsc::SyncSender<()>, + steer_finished: std::sync::Mutex>, +} + +impl Layer for RegularTaskRunCloseBarrier +where + S: Subscriber + for<'lookup> LookupSpan<'lookup>, +{ + fn on_close(&self, id: Id, ctx: Context<'_, S>) { + let is_regular_task_run = ctx + .metadata(&id) + .is_some_and(|metadata| metadata.name() == "session_task.run"); + if std::thread::current().id() != self.target_thread + || !is_regular_task_run + || !self.armed.swap(false, Ordering::SeqCst) + { + return; + } + + self.run_closed + .send(()) + .expect("late-steer worker should wait for regular task completion"); + self.steer_finished + .lock() + .expect("late-steer barrier mutex should not be poisoned") + .recv_timeout(Duration::from_secs(15)) + .expect("late-steer worker should finish before task finalization resumes"); + } +} fn ev_message_item_done(id: &str, text: &str) -> Value { serde_json::json!({ @@ -287,6 +335,114 @@ async fn wait_for_sleep_item_completed(codex: &CodexThread, call_id: &str, durat ); } +#[tokio::test(flavor = "current_thread")] +async fn steer_accepted_after_regular_task_returns_continues_same_turn() { + const INITIAL_PROMPT: &str = "first prompt"; + const STEER_PROMPT: &str = "late steer"; + const CLIENT_ID: &str = "late-steer-client-id"; + const FINAL_MESSAGE: &str = "processed late steer"; + + let armed = Arc::new(AtomicBool::new(false)); + let (run_closed_tx, run_closed_rx) = mpsc::sync_channel(0); + let (steer_finished_tx, steer_finished_rx) = mpsc::sync_channel(0); + let subscriber = tracing_subscriber::registry().with(RegularTaskRunCloseBarrier { + armed: Arc::clone(&armed), + target_thread: std::thread::current().id(), + run_closed: run_closed_tx, + steer_finished: std::sync::Mutex::new(steer_finished_rx), + }); + // A global subscriber keeps callsite interest stable while other integration tests run in + // parallel. The layer only acts on this current-thread runtime and disarms after one span. + tracing::subscriber::set_global_default(subscriber) + .expect("pending-input integration test should install the global tracing subscriber"); + + let final_chunks = vec![ + chunk(ev_response_created("resp-2")), + chunk(ev_message_item_done("msg-2", FINAL_MESSAGE)), + chunk(ev_completed("resp-2")), + ]; + let (server, _completions) = + start_streaming_sse_server(vec![response_completed_chunks("resp-1"), final_chunks]).await; + let codex = build_codex(&server).await; + + let runtime_handle = tokio::runtime::Handle::current(); + let codex_for_steer = Arc::clone(&codex); + let steer_thread = std::thread::spawn(move || { + run_closed_rx + .recv_timeout(Duration::from_secs(15)) + .expect("regular task should reach its post-run finalization gap"); + runtime_handle.block_on(async { + codex_for_steer + .steer_input( + vec![UserInput::Text { + text: STEER_PROMPT.to_string(), + text_elements: Vec::new(), + }], + /*additional_context*/ Default::default(), + /*expected_turn_id*/ None, + /*client_user_message_id*/ + Some(CLIENT_ID.to_string()), + /*responsesapi_client_metadata*/ None, + ) + .await + .expect("steer should be accepted before task finalization"); + }); + steer_finished_tx + .send(()) + .expect("task finalization should wait for the steer"); + }); + + armed.store(true, Ordering::SeqCst); + submit_user_input(&codex, INITIAL_PROMPT).await; + + let mut turn_started_count = 0; + let mut late_user_message_client_ids = Vec::new(); + let mut received_final_message = false; + tokio::time::timeout(Duration::from_secs(30), async { + loop { + let event = codex + .next_event() + .await + .expect("event stream should remain open"); + match event.msg { + EventMsg::TurnStarted(_) => turn_started_count += 1, + EventMsg::UserMessage(message) if message.message == STEER_PROMPT => { + late_user_message_client_ids.push(message.client_id); + } + EventMsg::AgentMessage(message) if message.message == FINAL_MESSAGE => { + received_final_message = true; + } + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + }) + .await + .expect("continued regular task should complete"); + steer_thread + .join() + .expect("late-steer worker should not panic"); + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let continued_body: Value = from_slice(&requests[1]).expect("parse continued request"); + assert_eq!( + message_input_texts(&continued_body, "user") + .into_iter() + .filter(|text| text == STEER_PROMPT) + .count(), + 1 + ); + assert_eq!(turn_started_count, 1); + assert_eq!( + late_user_message_client_ids, + vec![Some(CLIENT_ID.to_string())] + ); + assert!(received_final_message); + + server.shutdown().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";