diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index 46df387891..57cf0d80b3 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -358,7 +358,9 @@ impl ToolOutput for ExecCommandToolOutput { return None; } - Some(JsonValue::String(self.truncated_output())) + Some(JsonValue::String( + self.truncated_output(self.model_output_max_tokens()), + )) } fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue { @@ -383,7 +385,7 @@ impl ToolOutput for ExecCommandToolOutput { session_id: self.process_id, original_token_count: self.original_token_count, output: match self.max_output_tokens { - Some(_) => self.truncated_output(), + Some(max_tokens) => self.truncated_output(max_tokens), None => String::from_utf8_lossy(&self.raw_output).to_string(), }, }; @@ -395,10 +397,12 @@ impl ToolOutput for ExecCommandToolOutput { } impl ExecCommandToolOutput { - pub(crate) fn truncated_output(&self) -> String { + fn model_output_max_tokens(&self) -> usize { + resolve_max_tokens(self.max_output_tokens).min(self.truncation_policy.token_budget()) + } + + pub(crate) fn truncated_output(&self, max_tokens: usize) -> String { let text = String::from_utf8_lossy(&self.raw_output).to_string(); - let max_tokens = - resolve_max_tokens(self.max_output_tokens).min(self.truncation_policy.token_budget()); formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens)) } @@ -425,7 +429,7 @@ impl ExecCommandToolOutput { } sections.push("Output:".to_string()); - sections.push(self.truncated_output()); + sections.push(self.truncated_output(self.model_output_max_tokens())); sections.join("\n") } diff --git a/codex-rs/core/src/unified_exec/mod_tests.rs b/codex-rs/core/src/unified_exec/mod_tests.rs index e524d09722..d416be2e33 100644 --- a/codex-rs/core/src/unified_exec/mod_tests.rs +++ b/codex-rs/core/src/unified_exec/mod_tests.rs @@ -263,7 +263,9 @@ async fn unified_exec_persists_across_requests() -> anyhow::Result<()> { ) .await?; assert!( - out_2.truncated_output().contains("codex"), + out_2 + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains("codex"), "expected environment variable output" ); @@ -304,7 +306,9 @@ async fn multi_unified_exec_sessions() -> anyhow::Result<()> { "short command should not report a process id if it exits quickly" ); assert!( - !out_2.truncated_output().contains("codex"), + !out_2 + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains("codex"), "short command should run in a fresh shell" ); @@ -316,7 +320,9 @@ async fn multi_unified_exec_sessions() -> anyhow::Result<()> { ) .await?; assert!( - out_3.truncated_output().contains("codex"), + out_3 + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains("codex"), "session should preserve state" ); @@ -353,7 +359,9 @@ async fn unified_exec_timeouts() -> anyhow::Result<()> { ) .await?; assert!( - !out_2.truncated_output().contains(TEST_VAR_VALUE), + !out_2 + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains(TEST_VAR_VALUE), "timeout too short should yield incomplete output" ); @@ -362,7 +370,9 @@ async fn unified_exec_timeouts() -> anyhow::Result<()> { let out_3 = write_stdin(&session, process_id, "", /*yield_time_ms*/ 100).await?; assert!( - out_3.truncated_output().contains(TEST_VAR_VALUE), + out_3 + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains(TEST_VAR_VALUE), "subsequent poll should retrieve output" ); @@ -397,7 +407,9 @@ async fn unified_exec_pause_blocks_yield_timeout() -> anyhow::Result<()> { "pause should block the unified exec yield timeout" ); assert!( - response.truncated_output().contains("unified-exec-done"), + response + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains("unified-exec-done"), "exec_command should wait for output after the pause lifts" ); assert!( @@ -423,7 +435,11 @@ async fn requests_with_large_timeout_are_capped() -> anyhow::Result<()> { .await?; assert!(result.process_id.is_some()); - assert!(result.truncated_output().contains("codex")); + assert!( + result + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains("codex") + ); Ok(()) } @@ -445,7 +461,11 @@ async fn completed_commands_do_not_persist_sessions() -> anyhow::Result<()> { result.process_id.is_some(), "completed command should report a process id" ); - assert!(result.truncated_output().contains("codex")); + assert!( + result + .truncated_output(DEFAULT_MAX_OUTPUT_TOKENS) + .contains("codex") + ); assert!( session diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 29492f286e..52d7e34b8a 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -47,7 +47,7 @@ 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_OUTER_MAX_OUTPUT_TOKENS: usize = 20_000; +const CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS: usize = 20_000; const CODE_MODE_EXEC_SMALL_TOOL_OUTPUT_LIMIT: usize = 50; fn custom_tool_output_items(req: &ResponsesRequest, call_id: &str) -> Vec { @@ -714,6 +714,53 @@ text(result.output); Ok(()) } +#[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +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} +}}); +text(result.output); +"# + ); + + let output = code_mode_exec_output_text(&code).await?; + assert_eq!(output.len(), CODE_MODE_EXEC_RAW_OUTPUT_LEN); + + Ok(()) +} + +#[cfg_attr(windows, ignore = "no exec_command on Windows")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +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 output = code_mode_exec_output_text_with_config(&code, |config| { + config.tool_output_token_limit = Some(CODE_MODE_EXEC_SMALL_TOOL_OUTPUT_LIMIT); + }) + .await?; + assert_eq!(output.len(), CODE_MODE_EXEC_RAW_OUTPUT_LEN); + + Ok(()) +} + #[cfg_attr(windows, ignore = "no exec_command on Windows")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_exec_without_max_preserves_output_beyond_default() -> Result<()> { @@ -721,7 +768,7 @@ async fn code_mode_exec_without_max_preserves_output_beyond_default() -> Result< 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_OUTER_MAX_OUTPUT_TOKENS}}} + r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}} const result = await tools.exec_command({{ cmd: {command:?} }}); @@ -742,7 +789,7 @@ async fn code_mode_exec_without_max_preserves_output_beyond_truncation_policy() 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_OUTER_MAX_OUTPUT_TOKENS}}} + r#"// @exec: {{"max_output_tokens": {CODE_MODE_EXEC_LARGE_MAX_OUTPUT_TOKENS}}} const result = await tools.exec_command({{ cmd: {command:?} }});