Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 32 additions & 1 deletion codex-rs/core/src/tools/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use crate::tools::TELEMETRY_PREVIEW_MAX_BYTES;
use crate::tools::TELEMETRY_PREVIEW_MAX_LINES;
use crate::tools::TELEMETRY_PREVIEW_TRUNCATION_NOTICE;
use crate::turn_diff_tracker::TurnDiffTracker;
use crate::unified_exec::format_output_omission_marker;
use crate::unified_exec::resolve_max_tokens;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::models::FunctionCallOutputBody;
Expand All @@ -17,10 +18,13 @@ use codex_protocol::models::function_call_output_content_items_to_text;
use codex_tools::LoadableToolSpec;
use codex_tools::ToolName;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::approx_token_count;
use codex_utils_output_truncation::formatted_truncate_text;
use codex_utils_output_truncation::truncate_text;
use codex_utils_string::take_bytes_at_char_boundary;
use serde::Serialize;
use serde_json::Value as JsonValue;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
Expand Down Expand Up @@ -319,6 +323,8 @@ pub struct ExecCommandToolOutput {
pub process_id: Option<i32>,
pub exit_code: Option<i32>,
pub original_token_count: Option<usize>,
/// Bytes omitted by the output collection cap before model-facing truncation.
pub output_omitted_bytes: Option<NonZeroUsize>,
pub hook_command: Option<String>,
}

Expand Down Expand Up @@ -406,7 +412,32 @@ impl ExecCommandToolOutput {

pub(crate) fn truncated_output(&self, max_tokens: usize) -> String {
let text = String::from_utf8_lossy(&self.raw_output).to_string();
formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens))
let policy = TruncationPolicy::Tokens(max_tokens);
let Some(omitted_bytes) = self.output_omitted_bytes else {
return formatted_truncate_text(&text, policy);
};

let marker = format_output_omission_marker(omitted_bytes.get());
if text.len() <= policy.byte_budget() {
return if text.contains(&marker) {
text
} else {
format!("{marker}\n{text}")
};
}

let original_token_count = self
.original_token_count
.unwrap_or_else(|| approx_token_count(&text));
let truncated = truncate_text(&text, policy);
let omission_notice = if truncated.contains(&marker) {
String::new()
} else {
format!("{marker}\n")
};
format!(
"Warning: truncated output (original token count: {original_token_count})\n{omission_notice}\n{truncated}"
)
}

fn response_text(&self) -> String {
Expand Down
40 changes: 40 additions & 0 deletions codex-rs/core/src/tools/context_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,7 @@ fn exec_command_tool_output_formats_truncated_response() {
process_id: None,
exit_code: Some(0),
original_token_count: Some(10),
output_omitted_bytes: None,
hook_command: None,
}
.to_response_item("call-42", &payload);
Expand Down Expand Up @@ -461,3 +462,42 @@ fn exec_command_tool_output_formats_truncated_response() {
other => panic!("expected FunctionCallOutput, got {other:?}"),
}
}

#[test]
fn exec_command_tool_output_preserves_omission_metadata_when_truncated() {
let payload = ToolPayload::Function {
arguments: "{}".to_string(),
};
let marker = format_output_omission_marker(/*omitted_bytes*/ 123_456);
let raw_output = format!(
"HEAD-{}\n{marker}\nTAIL-{}",
"a".repeat(/*n*/ 100),
"z".repeat(/*n*/ 100)
)
.into_bytes();
let response = ExecCommandToolOutput {
event_call_id: "call-omitted".to_string(),
chunk_id: "abc123".to_string(),
wall_time: std::time::Duration::from_millis(/*millis*/ 1250),
raw_output,
truncation_policy: TruncationPolicy::Tokens(10_000),
max_output_tokens: Some(4),
process_id: None,
exit_code: Some(0),
original_token_count: Some(42_000),
output_omitted_bytes: NonZeroUsize::new(/*n*/ 123_456),
hook_command: None,
}
.to_response_item("call-omitted", &payload);

let ResponseInputItem::FunctionCallOutput { output, .. } = response else {
panic!("expected FunctionCallOutput");
};
let text = output
.body
.to_text()
.expect("exec output should serialize as text");
assert!(text.contains("Original token count: 42000"));
assert!(text.contains("Warning: truncated output (original token count: 42000)"));
assert_eq!(text.matches(&marker).count(), 1);
}
12 changes: 10 additions & 2 deletions codex-rs/core/src/tools/handlers/unified_exec/exec_command.rs
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,7 @@ impl ExecCommandHandler {
process_id: None,
exit_code: None,
original_token_count: None,
output_omitted_bytes: None,
hook_command: None,
}));
}
Expand Down Expand Up @@ -367,9 +368,15 @@ impl ExecCommandHandler {
.await
{
Ok(response) => Ok(boxed_tool_output(response)),
Err(UnifiedExecError::SandboxDenied { output, .. }) => {
Err(UnifiedExecError::SandboxDenied {
output,
original_token_count,
output_omitted_bytes,
..
}) => {
let output_text = output.aggregated_output.text;
let original_token_count = approx_token_count(&output_text);
let original_token_count =
original_token_count.unwrap_or_else(|| approx_token_count(&output_text));
Ok(boxed_tool_output(ExecCommandToolOutput {
event_call_id: context.call_id.clone(),
chunk_id: generate_chunk_id(),
Expand All @@ -382,6 +389,7 @@ impl ExecCommandHandler {
process_id: None,
exit_code: Some(output.exit_code),
original_token_count: Some(original_token_count),
output_omitted_bytes,
hook_command: Some(hook_command),
}))
}
Expand Down
6 changes: 6 additions & 0 deletions codex-rs/core/src/tools/handlers/unified_exec_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_noninteractive_one_s
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("echo three".to_string()),
};
let invocation = invocation_for_payload("exec_command", "call-43", payload).await;
Expand Down Expand Up @@ -331,6 +332,7 @@ async fn exec_command_post_tool_use_payload_uses_output_for_interactive_completi
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("echo three".to_string()),
};
let invocation = invocation_for_payload("exec_command", "call-44", payload).await;
Expand Down Expand Up @@ -362,6 +364,7 @@ async fn exec_command_post_tool_use_payload_skips_running_sessions() {
process_id: Some(45),
exit_code: None,
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("echo three".to_string()),
};
let invocation = invocation_for_payload("exec_command", "call-45", payload).await;
Expand All @@ -388,6 +391,7 @@ async fn write_stdin_post_tool_use_payload_uses_original_exec_call_id_and_comman
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("sleep 1; echo finished".to_string()),
};
let invocation = invocation_for_payload("write_stdin", "write-stdin-call", payload).await;
Expand Down Expand Up @@ -419,6 +423,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("sleep 2; echo alpha".to_string()),
};
let output_b = ExecCommandToolOutput {
Expand All @@ -431,6 +436,7 @@ async fn write_stdin_post_tool_use_payload_keeps_parallel_session_metadata_separ
process_id: None,
exit_code: Some(0),
original_token_count: None,
output_omitted_bytes: None,
hook_command: Some("sleep 1; echo beta".to_string()),
};
let invocation_b = invocation_for_payload("write_stdin", "write-call-b", payload.clone()).await;
Expand Down
2 changes: 1 addition & 1 deletion codex-rs/core/src/unified_exec/async_watcher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -326,7 +326,7 @@ async fn resolve_aggregated_output(
return fallback;
}

String::from_utf8_lossy(&guard.to_bytes()).to_string()
String::from_utf8_lossy(&guard.to_bytes_with_omission_marker()).to_string()
}

#[cfg(test)]
Expand Down
28 changes: 27 additions & 1 deletion codex-rs/core/src/unified_exec/errors.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use codex_protocol::exec_output::ExecToolCallOutput;
use codex_utils_path_uri::PathUri;
use std::num::NonZeroUsize;
use thiserror::Error;

#[derive(Debug, Error)]
Expand All @@ -23,6 +24,8 @@ pub(crate) enum UnifiedExecError {
SandboxDenied {
message: String,
output: ExecToolCallOutput,
original_token_count: Option<usize>,
output_omitted_bytes: Option<NonZeroUsize>,
},
#[error("{path} is not valid on {}", std::env::consts::OS)]
ForeignPath { path: PathUri },
Expand All @@ -38,6 +41,29 @@ impl UnifiedExecError {
}

pub(crate) fn sandbox_denied(message: String, output: ExecToolCallOutput) -> Self {
Self::SandboxDenied { message, output }
Self::SandboxDenied {
message,
output,
original_token_count: None,
output_omitted_bytes: None,
}
}

pub(crate) fn with_output_collection_metadata(
self,
original_token_count: usize,
output_omitted_bytes: Option<NonZeroUsize>,
) -> Self {
match self {
Self::SandboxDenied {
message, output, ..
} => Self::SandboxDenied {
message,
output,
original_token_count: Some(original_token_count),
output_omitted_bytes,
},
other => other,
}
}
}
Loading
Loading