From d04071d3ff517193d405901c46ac2839041f2a4e Mon Sep 17 00:00:00 2001 From: aakash-oai Date: Thu, 29 Oct 2026 20:29:13 +0000 Subject: [PATCH] Stabilize cancellation and tracing tests (#58171) ## Why The SIGTERM cleanup test depends on the OS scheduling the shell within the production grace period. Tracing tests also need fixture creation and cleanup to use the appropriate subscriber. ## What changed - Pause Tokio time while waiting for the shell's cleanup marker, keeping the cancellation future runnable to prevent automatic clock advancement, then resume time after execution. - Adjust subscriber setup and fixture lifetimes in session and tool dispatch tests so instrumented resources are created and dropped under the appropriate subscriber. GitOrigin-RevId: 12c654a031a8a4e61fc66488a5f687852f5ca099 --- codex-rs/core/src/exec_tests.rs | 17 ++++++++++++++++ codex-rs/core/src/session/tests.rs | 9 +++++---- codex-rs/core/src/tools/parallel.rs | 20 ++++++++++--------- .../core/src/tools/parallel_trace_tests.rs | 2 ++ 4 files changed, 35 insertions(+), 13 deletions(-) diff --git a/codex-rs/core/src/exec_tests.rs b/codex-rs/core/src/exec_tests.rs index 87d730a56a..98a542daaf 100644 --- a/codex-rs/core/src/exec_tests.rs +++ b/codex-rs/core/src/exec_tests.rs @@ -1526,11 +1526,25 @@ while :; do sleep 1; done"# ); let cancel_token = CancellationToken::new(); let cancel_tx = cancel_token.clone(); + let clock_paused = std::cell::Cell::new(false); let cancel_when_ready = async { while !ready_marker.exists() { tokio::time::sleep(Duration::from_millis(20)).await; } + // Exercise the TERM trap without requiring the OS to schedule the shell + // within the production grace period. Keep this future runnable so the + // paused clock cannot auto-advance while the real process handles TERM. + tokio::time::pause(); + clock_paused.set(true); cancel_tx.cancel(); + let cleanup_deadline = std::time::Instant::now() + Duration::from_secs(5); + while !matches!( + std::fs::read_to_string(&cleanup_marker).as_deref(), + Ok("cleaned") + ) && std::time::Instant::now() < cleanup_deadline + { + tokio::task::yield_now().await; + } }; let params = ExecParams { command, @@ -1567,6 +1581,9 @@ while :; do sleep 1; done"# } }) .await; + if clock_paused.get() { + tokio::time::resume(); + } if result.is_err() { cancel_tx.cancel(); let _ = execution.await; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 9cf370f6e6..d2d62152a0 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -628,6 +628,8 @@ fn skill_message(text: &str) -> ResponseItem { #[tokio::test] async fn regular_turn_emits_turn_started_with_trace_id_without_waiting_for_startup_prewarm() { + // Prepare filesystem-backed fixtures before replacing the ambient subscriber. + let (sess, _initial_turn, rx) = make_session_and_context_with_rx().await; let _trace_test_context = install_test_tracing("codex-core-tests"); let request_parent = W3cTraceContext { traceparent: Some("00-00000000000000000000000000000011-0000000000000022-01".into()), @@ -638,9 +640,7 @@ async fn regular_turn_emits_turn_started_with_trace_id_without_waiting_for_start &request_span, &request_parent )); - let (sess, tc, rx) = make_session_and_context_with_rx() - .instrument(request_span) - .await; + let tc = sess.new_default_turn().instrument(request_span).await; assert_eq!( tc.trace_id.as_deref(), Some("00000000000000000000000000000011") @@ -8754,6 +8754,8 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() { } } + // Keep fixture creation and cleanup under the ambient subscriber. + let (sess, tc, rx) = make_session_and_context_with_rx().await; let _trace_test_context = install_test_tracing("codex-core-tests"); let request_parent = W3cTraceContext { @@ -8783,7 +8785,6 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() { }); let dispatch_span_id = dispatch_span.context().span().span_context().span_id(); - let (sess, tc, rx) = make_session_and_context_with_rx().await; let captured_trace = Arc::new(std::sync::Mutex::new(None)); async { diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index a20ca34f02..4361a91aa7 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -767,6 +767,17 @@ mod tests { dynamic: bool, critical_path_timing: bool, ) -> anyhow::Result<()> { + // Create and drop instrumented session resources under the same subscriber. + // A prior traced test may have installed a different global subscriber. + let buffer: &'static std::sync::Mutex> = + Box::leak(Box::new(std::sync::Mutex::new(Vec::new()))); + let subscriber = tracing_subscriber::fmt() + .with_ansi(false) + .with_max_level(tracing::Level::INFO) + .with_writer(MockWriter::new(buffer)) + .finish(); + let _subscriber_guard = tracing::subscriber::set_default(subscriber); + let (session, mut turn_context) = crate::session::tests::make_session_and_context().await; Arc::make_mut(&mut turn_context.config) .features @@ -832,15 +843,6 @@ mod tests { .expect("test should release the execution gate"); }); - let buffer: &'static std::sync::Mutex> = - Box::leak(Box::new(std::sync::Mutex::new(Vec::new()))); - let subscriber = tracing_subscriber::fmt() - .with_ansi(false) - .with_max_level(tracing::Level::INFO) - .with_writer(MockWriter::new(buffer)) - .finish(); - let _subscriber_guard = tracing::subscriber::set_default(subscriber); - let cancellation_token = CancellationToken::new(); let call = ToolCall { tool_name, diff --git a/codex-rs/core/src/tools/parallel_trace_tests.rs b/codex-rs/core/src/tools/parallel_trace_tests.rs index 1c4f6b11b3..9863fc4f0a 100644 --- a/codex-rs/core/src/tools/parallel_trace_tests.rs +++ b/codex-rs/core/src/tools/parallel_trace_tests.rs @@ -61,8 +61,10 @@ async fn spawned_dispatch_preserves_caller_context(level: LevelFilter) -> anyhow tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default()); let _subscriber_guard = tracing::subscriber::set_default(subscriber); async { + // Retain the fixture until its ambient subscriber is restored. // The unknown-tool result still crosses the real spawned dispatch boundary. runtime + .clone() .handle_tool_call( ToolCall { tool_name: codex_tools::ToolName::plain("unknown_tool"),