From 5a9a5b51b273dc11269a91a13c30838e8ce0f8c9 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 25 Feb 2026 13:19:21 +0000 Subject: [PATCH] feat: add large stack test macro (#12768) This PR adds the macro `#[large_stack_test]` This spawns the tests in a dedicated tokio runtime with a larger stack. It is useful for tests that needs the full recursion on the harness (which is now too deep for windows for example) --- codex-rs/Cargo.lock | 10 ++ codex-rs/Cargo.toml | 2 + codex-rs/core/Cargo.toml | 1 + codex-rs/core/tests/suite/apply_patch_cli.rs | 59 +++---- codex-rs/test-macros/BUILD.bazel | 7 + codex-rs/test-macros/Cargo.toml | 16 ++ codex-rs/test-macros/src/lib.rs | 155 +++++++++++++++++++ 7 files changed, 221 insertions(+), 29 deletions(-) create mode 100644 codex-rs/test-macros/BUILD.bazel create mode 100644 codex-rs/test-macros/Cargo.toml create mode 100644 codex-rs/test-macros/src/lib.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 36f4fec3b..06acc6c97 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1718,6 +1718,7 @@ dependencies = [ "codex-shell-escalation", "codex-skills", "codex-state", + "codex-test-macros", "codex-utils-absolute-path", "codex-utils-cargo-bin", "codex-utils-home-dir", @@ -2298,6 +2299,15 @@ dependencies = [ "uds_windows", ] +[[package]] +name = "codex-test-macros" +version = "0.0.0" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "codex-tui" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index b3bb85d5a..51f541f24 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -62,6 +62,7 @@ members = [ "codex-api", "state", "codex-experimental-api-macros", + "test-macros", ] resolver = "2" @@ -116,6 +117,7 @@ codex-shell-command = { path = "shell-command" } codex-shell-escalation = { path = "shell-escalation" } codex-skills = { path = "skills" } codex-state = { path = "state" } +codex-test-macros = { path = "test-macros" } codex-stdio-to-uds = { path = "stdio-to-uds" } codex-tui = { path = "tui" } codex-utils-absolute-path = { path = "utils/absolute-path" } diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 820f35cb1..9a0be02a9 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -148,6 +148,7 @@ codex-arg0 = { workspace = true } codex-otel = { workspace = true, features = [ "disable-default-metrics-exporter", ] } +codex-test-macros = { workspace = true } codex-utils-cargo-bin = { workspace = true } core_test_support = { workspace = true } ctor = { workspace = true } diff --git a/codex-rs/core/tests/suite/apply_patch_cli.rs b/codex-rs/core/tests/suite/apply_patch_cli.rs index 9ef2715d8..45494124e 100644 --- a/codex-rs/core/tests/suite/apply_patch_cli.rs +++ b/codex-rs/core/tests/suite/apply_patch_cli.rs @@ -1,6 +1,7 @@ #![allow(clippy::expect_used)] use anyhow::Result; +use codex_test_macros::large_stack_test; use core_test_support::responses::ev_apply_patch_call; use core_test_support::responses::ev_apply_patch_custom_tool_call; use core_test_support::responses::ev_shell_command_call; @@ -85,7 +86,7 @@ fn apply_patch_responses( ] } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -132,7 +133,7 @@ D delete.txt Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -159,7 +160,7 @@ async fn apply_patch_cli_multiple_chunks(model_output: ApplyPatchModelOutput) -> Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -188,7 +189,7 @@ async fn apply_patch_cli_moves_file_to_new_directory( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -216,7 +217,7 @@ async fn apply_patch_cli_updates_file_appends_trailing_newline( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -242,7 +243,7 @@ async fn apply_patch_cli_insert_only_hunk_modifies_file( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -273,7 +274,7 @@ async fn apply_patch_cli_move_overwrites_existing_destination( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -334,7 +335,7 @@ async fn apply_patch_cli_move_without_content_change_has_no_turn_diff( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -360,7 +361,7 @@ async fn apply_patch_cli_add_overwrites_existing_file( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -392,7 +393,7 @@ async fn apply_patch_cli_rejects_invalid_hunk_header( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -426,7 +427,7 @@ async fn apply_patch_cli_reports_missing_context( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -462,7 +463,7 @@ async fn apply_patch_cli_reports_missing_target_file( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -499,7 +500,7 @@ async fn apply_patch_cli_delete_missing_file_reports_error( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -524,7 +525,7 @@ async fn apply_patch_cli_rejects_empty_patch(model_output: ApplyPatchModelOutput Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -551,7 +552,7 @@ async fn apply_patch_cli_delete_directory_reports_verification_error( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -605,7 +606,7 @@ async fn apply_patch_cli_rejects_path_traversal_outside_workspace( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -660,7 +661,7 @@ async fn apply_patch_cli_rejects_move_path_traversal_outside_workspace( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -694,7 +695,7 @@ async fn apply_patch_cli_verification_failure_has_no_side_effects( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] async fn apply_patch_shell_command_heredoc_with_cd_updates_relative_workdir() -> Result<()> { skip_if_no_network!(Ok(())); @@ -732,7 +733,7 @@ async fn apply_patch_shell_command_heredoc_with_cd_updates_relative_workdir() -> Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] async fn apply_patch_cli_can_use_shell_command_output_as_patch_input() -> Result<()> { skip_if_no_network!(Ok(())); @@ -862,7 +863,7 @@ async fn apply_patch_cli_can_use_shell_command_output_as_patch_input() -> Result Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] async fn apply_patch_shell_command_heredoc_with_cd_emits_turn_diff() -> Result<()> { skip_if_no_network!(Ok(())); @@ -945,7 +946,7 @@ async fn apply_patch_shell_command_heredoc_with_cd_emits_turn_diff() -> Result<( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] async fn apply_patch_shell_command_failure_propagates_error_and_skips_diff() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1021,7 +1022,7 @@ async fn apply_patch_shell_command_failure_propagates_error_and_skips_diff() -> Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::ShellViaHeredoc)] #[test_case(ApplyPatchModelOutput::ShellCommandViaHeredoc)] async fn apply_patch_function_accepts_lenient_heredoc_wrapped_patch( @@ -1044,7 +1045,7 @@ async fn apply_patch_function_accepts_lenient_heredoc_wrapped_patch( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -1067,7 +1068,7 @@ async fn apply_patch_cli_end_of_file_anchor(model_output: ApplyPatchModelOutput) Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -1102,7 +1103,7 @@ async fn apply_patch_cli_missing_second_chunk_context_rejected( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -1161,7 +1162,7 @@ async fn apply_patch_emits_turn_diff_event_with_unified_diff( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] @@ -1227,7 +1228,7 @@ async fn apply_patch_turn_diff_for_rename_with_content_change( Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] async fn apply_patch_aggregates_diff_across_multiple_tool_calls() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1295,7 +1296,7 @@ async fn apply_patch_aggregates_diff_across_multiple_tool_calls() -> Result<()> Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] async fn apply_patch_aggregates_diff_preserves_success_after_failure() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1385,7 +1386,7 @@ async fn apply_patch_aggregates_diff_preserves_success_after_failure() -> Result Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[large_stack_test] #[test_case(ApplyPatchModelOutput::Freeform)] #[test_case(ApplyPatchModelOutput::Function)] #[test_case(ApplyPatchModelOutput::Shell)] diff --git a/codex-rs/test-macros/BUILD.bazel b/codex-rs/test-macros/BUILD.bazel new file mode 100644 index 000000000..0ff074f5e --- /dev/null +++ b/codex-rs/test-macros/BUILD.bazel @@ -0,0 +1,7 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "test-macros", + crate_name = "codex_test_macros", + proc_macro = True, +) diff --git a/codex-rs/test-macros/Cargo.toml b/codex-rs/test-macros/Cargo.toml new file mode 100644 index 000000000..275580322 --- /dev/null +++ b/codex-rs/test-macros/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "codex-test-macros" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lib] +proc-macro = true + +[dependencies] +proc-macro2 = "1" +quote = "1" +syn = { version = "2", features = ["full"] } + +[lints] +workspace = true diff --git a/codex-rs/test-macros/src/lib.rs b/codex-rs/test-macros/src/lib.rs new file mode 100644 index 000000000..16c87f858 --- /dev/null +++ b/codex-rs/test-macros/src/lib.rs @@ -0,0 +1,155 @@ +use proc_macro::TokenStream; +use proc_macro2::TokenStream as TokenStream2; +use quote::quote; +use syn::Attribute; +use syn::ItemFn; +use syn::parse::Nothing; +use syn::parse_macro_input; +use syn::parse_quote; + +const LARGE_STACK_TEST_STACK_SIZE_BYTES: usize = 16 * 1024 * 1024; + +/// Run a test body on a dedicated thread with a larger stack. +/// +/// For async tests, this macro creates a Tokio multi-thread runtime with two +/// worker threads and blocks on the original async body inside the large-stack +/// thread. +#[proc_macro_attribute] +pub fn large_stack_test(attr: TokenStream, item: TokenStream) -> TokenStream { + parse_macro_input!(attr as Nothing); + + let item = parse_macro_input!(item as ItemFn); + expand_large_stack_test(item).into() +} + +fn expand_large_stack_test(mut item: ItemFn) -> TokenStream2 { + let attrs = filtered_attributes(&item.attrs); + item.attrs = attrs; + + let is_async = item.sig.asyncness.take().is_some(); + let name = &item.sig.ident; + let body = &item.block; + + let thread_body = if is_async { + quote! { + { + let runtime = ::tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap_or_else(|error| { + panic!("failed to build tokio runtime for large-stack test: {error}") + }); + runtime.block_on(async move #body) + } + } + } else { + quote! { #body } + }; + + *item.block = parse_quote!({ + let handle = ::std::thread::Builder::new() + .name(::std::string::String::from(::std::stringify!(#name))) + .stack_size(#LARGE_STACK_TEST_STACK_SIZE_BYTES) + .spawn(move || #thread_body) + .unwrap_or_else(|error| { + panic!("failed to spawn large-stack test thread: {error}") + }); + + match handle.join() { + Ok(result) => result, + Err(payload) => ::std::panic::resume_unwind(payload), + } + }); + + quote! { #item } +} + +fn filtered_attributes(attrs: &[Attribute]) -> Vec { + let mut filtered = Vec::with_capacity(attrs.len() + 1); + let mut has_test_attr = false; + + for attr in attrs { + if is_tokio_test_attr(attr) { + continue; + } + if is_test_attr(attr) || is_test_case_attr(attr) { + has_test_attr = true; + } + filtered.push(attr.clone()); + } + + if !has_test_attr { + filtered.push(parse_quote!(#[test])); + } + + filtered +} + +fn is_test_attr(attr: &Attribute) -> bool { + attr.path().is_ident("test") +} + +fn is_test_case_attr(attr: &Attribute) -> bool { + attr.path().is_ident("test_case") +} + +fn is_tokio_test_attr(attr: &Attribute) -> bool { + let mut segments = attr.path().segments.iter(); + matches!( + (segments.next(), segments.next(), segments.next()), + (Some(first), Some(second), None) if first.ident == "tokio" && second.ident == "test" + ) +} + +#[cfg(test)] +mod tests { + use super::expand_large_stack_test; + use syn::ItemFn; + use syn::parse_quote; + + fn has_attr(item: &ItemFn, name: &str) -> bool { + item.attrs.iter().any(|attr| attr.path().is_ident(name)) + } + + #[test] + fn adds_test_attribute_when_missing() { + let item: ItemFn = parse_quote! { + fn sample() {} + }; + + let expanded_tokens = expand_large_stack_test(item); + let expanded: ItemFn = match syn::parse2(expanded_tokens) { + Ok(expanded) => expanded, + Err(error) => panic!("failed to parse expanded function: {error}"), + }; + + assert!(has_attr(&expanded, "test")); + let body = quote::quote!(#expanded).to_string(); + assert!(body.contains("stack_size")); + } + + #[test] + fn removes_tokio_test_and_keeps_test_case() { + let item: ItemFn = parse_quote! { + #[test_case(1)] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn sample(value: usize) -> anyhow::Result<()> { + let _ = value; + Ok(()) + } + }; + + let expanded_tokens = expand_large_stack_test(item); + let expanded: ItemFn = match syn::parse2(expanded_tokens) { + Ok(expanded) => expanded, + Err(error) => panic!("failed to parse expanded function: {error}"), + }; + + assert!(has_attr(&expanded, "test_case")); + assert!(!has_attr(&expanded, "test")); + let body = quote::quote!(#expanded).to_string(); + assert!(body.contains("tokio :: runtime :: Builder")); + assert!(!body.contains("tokio :: test")); + } +}