Inline code-mode exec output assertions

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

@@ -5,7 +5,6 @@ use base64::Engine;
 use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
 use codex_config::types::McpServerConfig;
 use codex_config::types::McpServerTransportConfig;
-use codex_core::config::Config;
 use codex_features::Feature;
 use codex_login::CodexAuth;
 use codex_models_manager::bundled_models_response;
@@ -44,14 +43,6 @@ use std::time::Duration;
 use std::time::Instant;
 use wiremock::MockServer;
 
-const CODE_MODE_EXEC_LONG_OUTPUT_COMMAND: &str =
-    "printf '0123456789012345678901234567890123456789'";
-const CODE_MODE_EXEC_RAW_OUTPUT_LEN: usize = 50_000;
-const CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS: usize = 20_000;
-const CODE_MODE_EXEC_SMALL_TOOL_OUTPUT_LIMIT: usize = 50;
-const CODE_MODE_EXEC_TRUNCATED_OUTPUT: &str =
-    "Total output lines: 1\n\n0123456789…5 tokens truncated…0123456789";
-
 fn custom_tool_output_items(req: &ResponsesRequest, call_id: &str) -> Vec<Value> {
     match req.custom_tool_call_output(call_id).get("output") {
         Some(Value::Array(items)) => items.clone(),
@@ -153,21 +144,11 @@ async fn run_code_mode_turn(
     server: &MockServer,
     prompt: &str,
     code: &str,
-) -> Result<(TestCodex, ResponseMock)> {
-    run_code_mode_turn_with_config(server, prompt, code, |_| {}).await
-}
-
-async fn run_code_mode_turn_with_config(
-    server: &MockServer,
-    prompt: &str,
-    code: &str,
-    configure: impl FnOnce(&mut Config) + Send + 'static,
 ) -> Result<(TestCodex, ResponseMock)> {
     let mut builder = test_codex()
         .with_model("test-gpt-5.1-codex")
         .with_config(move |config| {
             let _ = config.features.enable(Feature::CodeMode);
-            configure(config);
         });
     let test = builder.build(server).await?;
 
@@ -194,57 +175,6 @@ async fn run_code_mode_turn_with_config(
     Ok((test, second_mock))
 }
 
-async fn code_mode_exec_output_items(code: &str) -> Result<Vec<Value>> {
-    code_mode_exec_output_items_with_config(code, |_| {}).await
-}
-
-async fn code_mode_exec_output_items_with_config(
-    code: &str,
-    configure: impl FnOnce(&mut Config) + Send + 'static,
-) -> Result<Vec<Value>> {
-    let server = responses::start_mock_server().await;
-    let (_test, second_mock) =
-        run_code_mode_turn_with_config(&server, "use exec_command from code mode", code, configure)
-            .await?;
-
-    let req = second_mock.single_request();
-    Ok(custom_tool_output_items(&req, "call-1"))
-}
-
-fn assert_code_mode_exec_output_items(items: &[Value], expected_output: &str) {
-    assert_eq!(items.len(), 2);
-    assert_regex_match(
-        concat!(
-            r"(?s)\A",
-            r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
-        ),
-        text_item(items, /*index*/ 0),
-    );
-    assert_eq!(text_item(items, /*index*/ 1), expected_output);
-}
-
-async fn assert_code_mode_exec_output_text(code: &str, expected_output: &str) -> Result<()> {
-    let items = code_mode_exec_output_items(code).await?;
-    assert_code_mode_exec_output_items(&items, expected_output);
-
-    Ok(())
-}
-
-async fn assert_code_mode_exec_output_text_with_config(
-    code: &str,
-    configure: impl FnOnce(&mut Config) + Send + 'static,
-    expected_output: &str,
-) -> Result<()> {
-    let items = code_mode_exec_output_items_with_config(code, configure).await?;
-    assert_code_mode_exec_output_items(&items, expected_output);
-
-    Ok(())
-}
-
-fn code_mode_exec_repeated_output_command(len: usize) -> String {
-    format!("python3 -c \"import sys; sys.stdout.write('x' * {len})\"")
-}
-
 async fn run_code_mode_turn_with_rmcp(
     server: &MockServer,
     prompt: &str,
@@ -362,8 +292,7 @@ text(JSON.stringify(await tools.exec_command({ cmd: "printf code_mode_exec_marke
     )
     .await?;
 
-    let req = second_mock.single_request();
-    let items = custom_tool_output_items(&req, "call-1");
+    let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
     assert_eq!(items.len(), 2);
     assert_regex_match(
         concat!(
@@ -718,17 +647,33 @@ text(JSON.stringify(results));
 async fn code_mode_exec_command_explicit_max_output_tokens_truncates() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
-    let code = format!(
+    let server = responses::start_mock_server().await;
+    let (_test, second_mock) = run_code_mode_turn(
+        &server,
+        "use exec_command from code mode",
         r#"
-const result = await tools.exec_command({{
-  cmd: {CODE_MODE_EXEC_LONG_OUTPUT_COMMAND:?},
+const result = await tools.exec_command({
+  cmd: "printf '0123456789012345678901234567890123456789'",
   max_output_tokens: 5
-}});
+});
 text(result.output);
-"#
-    );
+"#,
+    )
+    .await?;
 
-    assert_code_mode_exec_output_text(&code, CODE_MODE_EXEC_TRUNCATED_OUTPUT).await?;
+    let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
+    assert_eq!(items.len(), 2);
+    assert_regex_match(
+        concat!(
+            r"(?s)\A",
+            r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
+        ),
+        text_item(&items, /*index*/ 0),
+    );
+    assert_eq!(
+        text_item(&items, /*index*/ 1),
+        "Total output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
+    );
 
     Ok(())
 }
@@ -738,19 +683,30 @@ text(result.output);
 async fn code_mode_exec_explicit_max_above_default_preserves_output() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
-    let command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    let code = format!(
-        r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
-const result = await tools.exec_command({{
-  cmd: {command:?},
-  max_output_tokens: {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}
-}});
+    let server = responses::start_mock_server().await;
+    let (_test, second_mock) = run_code_mode_turn(
+        &server,
+        "use exec_command from code mode",
+        r#"// @exec: {"max_output_tokens": 20000}
+const result = await tools.exec_command({
+  cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\"",
+  max_output_tokens: 20000
+});
 text(result.output);
-"#
-    );
+"#,
+    )
+    .await?;
 
-    let expected_output = "x".repeat(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    assert_code_mode_exec_output_text(&code, &expected_output).await?;
+    let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
+    assert_eq!(items.len(), 2);
+    assert_regex_match(
+        concat!(
+            r"(?s)\A",
+            r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
+        ),
+        text_item(&items, /*index*/ 0),
+    );
+    assert_eq!(text_item(&items, /*index*/ 1), "x".repeat(50_000));
 
     Ok(())
 }
@@ -760,26 +716,56 @@ text(result.output);
 async fn code_mode_exec_explicit_max_above_truncation_policy_preserves_output() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
-    let command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    let code = format!(
-        r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
-const result = await tools.exec_command({{
-  cmd: {command:?},
-  max_output_tokens: {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}
-}});
-text(result.output);
-"#
-    );
+    let server = responses::start_mock_server().await;
+    let mut builder = test_codex()
+        .with_model("test-gpt-5.1-codex")
+        .with_config(|config| {
+            let _ = config.features.enable(Feature::CodeMode);
+            config.tool_output_token_limit = Some(50);
+        });
+    let test = builder.build(&server).await?;
 
-    let expected_output = "x".repeat(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    assert_code_mode_exec_output_text_with_config(
-        &code,
-        |config| {
-            config.tool_output_token_limit = Some(CODE_MODE_EXEC_SMALL_TOOL_OUTPUT_LIMIT);
-        },
-        &expected_output,
+    responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_response_created("resp-1"),
+            ev_custom_tool_call(
+                "call-1",
+                "exec",
+                r#"// @exec: {"max_output_tokens": 20000}
+const result = await tools.exec_command({
+  cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\"",
+  max_output_tokens: 20000
+});
+text(result.output);
+"#,
+            ),
+            ev_completed("resp-1"),
+        ]),
     )
-    .await?;
+    .await;
+
+    let second_mock = responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_assistant_message("msg-1", "done"),
+            ev_completed("resp-2"),
+        ]),
+    )
+    .await;
+
+    test.submit_turn("use exec_command from code mode").await?;
+
+    let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
+    assert_eq!(items.len(), 2);
+    assert_regex_match(
+        concat!(
+            r"(?s)\A",
+            r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
+        ),
+        text_item(&items, /*index*/ 0),
+    );
+    assert_eq!(text_item(&items, /*index*/ 1), "x".repeat(50_000));
 
     Ok(())
 }
@@ -789,18 +775,29 @@ text(result.output);
 async fn code_mode_exec_without_max_preserves_output_beyond_default() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
-    let command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    let code = format!(
-        r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
-const result = await tools.exec_command({{
-  cmd: {command:?}
-}});
+    let server = responses::start_mock_server().await;
+    let (_test, second_mock) = run_code_mode_turn(
+        &server,
+        "use exec_command from code mode",
+        r#"// @exec: {"max_output_tokens": 20000}
+const result = await tools.exec_command({
+  cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\""
+});
 text(result.output);
-"#
-    );
+"#,
+    )
+    .await?;
 
-    let expected_output = "x".repeat(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    assert_code_mode_exec_output_text(&code, &expected_output).await?;
+    let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
+    assert_eq!(items.len(), 2);
+    assert_regex_match(
+        concat!(
+            r"(?s)\A",
+            r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
+        ),
+        text_item(&items, /*index*/ 0),
+    );
+    assert_eq!(text_item(&items, /*index*/ 1), "x".repeat(50_000));
 
     Ok(())
 }
@@ -810,25 +807,55 @@ text(result.output);
 async fn code_mode_exec_without_max_preserves_output_beyond_truncation_policy() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
-    let command = code_mode_exec_repeated_output_command(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    let code = format!(
-        r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}}
-const result = await tools.exec_command({{
-  cmd: {command:?}
-}});
-text(result.output);
-"#
-    );
+    let server = responses::start_mock_server().await;
+    let mut builder = test_codex()
+        .with_model("test-gpt-5.1-codex")
+        .with_config(|config| {
+            let _ = config.features.enable(Feature::CodeMode);
+            config.tool_output_token_limit = Some(50);
+        });
+    let test = builder.build(&server).await?;
 
-    let expected_output = "x".repeat(CODE_MODE_EXEC_RAW_OUTPUT_LEN);
-    assert_code_mode_exec_output_text_with_config(
-        &code,
-        |config| {
-            config.tool_output_token_limit = Some(CODE_MODE_EXEC_SMALL_TOOL_OUTPUT_LIMIT);
-        },
-        &expected_output,
+    responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_response_created("resp-1"),
+            ev_custom_tool_call(
+                "call-1",
+                "exec",
+                r#"// @exec: {"max_output_tokens": 20000}
+const result = await tools.exec_command({
+  cmd: "python3 -c \"import sys; sys.stdout.write('x' * 50000)\""
+});
+text(result.output);
+"#,
+            ),
+            ev_completed("resp-1"),
+        ]),
     )
-    .await?;
+    .await;
+
+    let second_mock = responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_assistant_message("msg-1", "done"),
+            ev_completed("resp-2"),
+        ]),
+    )
+    .await;
+
+    test.submit_turn("use exec_command from code mode").await?;
+
+    let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
+    assert_eq!(items.len(), 2);
+    assert_regex_match(
+        concat!(
+            r"(?s)\A",
+            r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
+        ),
+        text_item(&items, /*index*/ 0),
+    );
+    assert_eq!(text_item(&items, /*index*/ 1), "x".repeat(50_000));
 
     Ok(())
 }
@@ -838,16 +865,32 @@ text(result.output);
 async fn code_mode_exec_explicit_max_output_tokens_truncates() -> Result<()> {
     skip_if_no_network!(Ok(()));
 
-    let code = format!(
-        r#"// @exec: {{"max_output_tokens": 5}}
-const result = await tools.exec_command({{
-  cmd: {CODE_MODE_EXEC_LONG_OUTPUT_COMMAND:?}
-}});
+    let server = responses::start_mock_server().await;
+    let (_test, second_mock) = run_code_mode_turn(
+        &server,
+        "use exec_command from code mode",
+        r#"// @exec: {"max_output_tokens": 5}
+const result = await tools.exec_command({
+  cmd: "printf '0123456789012345678901234567890123456789'"
+});
 text(result.output);
-"#
-    );
+"#,
+    )
+    .await?;
 
-    assert_code_mode_exec_output_text(&code, CODE_MODE_EXEC_TRUNCATED_OUTPUT).await?;
+    let items = custom_tool_output_items(&second_mock.single_request(), "call-1");
+    assert_eq!(items.len(), 2);
+    assert_regex_match(
+        concat!(
+            r"(?s)\A",
+            r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z"
+        ),
+        text_item(&items, /*index*/ 0),
+    );
+    assert_eq!(
+        text_item(&items, /*index*/ 1),
+        "Total output lines: 1\n\n0123456789…5 tokens truncated…0123456789"
+    );
 
     Ok(())
 }