Centralize unified exec output rendering (#57185)

## What changed

Introduce `CapturedStream` to keep retained output text and omitted-byte counts together. Use its shared renderer for unified exec logs and model-facing output, preserving omission markers and truncation notices.

GitOrigin-RevId: 045d925df1e31945755196dafed5f78ed507b33b

codex-rs/core/src/unified_exec/mod.rs

@@ -49,6 +49,7 @@ use codex_core_plugins::PluginMetricsSidecar;
 mod async_watcher;
 mod errors;
 mod oneshot;
+pub(crate) mod output;
 mod process;
 mod process_manager;
 pub(crate) use process_manager::exec_env_policy_from_shell_policy;

codex-rs/core/src/unified_exec/mod_tests.rs

@@ -9,6 +9,7 @@ use crate::session::tests::make_session_and_context;
 use crate::session::turn_context::TurnContext;
 use crate::tools::context::ExecCommandToolOutput;
 use crate::unified_exec::WriteStdinRequest;
+use crate::unified_exec::output::CapturedStream;
 use codex_exec_server::ExecProcess;
 use codex_exec_server::ExecProcessEventReceiver;
 use codex_exec_server::ExecProcessFuture;
@@ -27,7 +28,6 @@ use core_test_support::skip_if_sandbox;
 use core_test_support::test_codex::test_env as remote_test_env;
 use pretty_assertions::assert_eq;
 use std::collections::HashMap;
-use std::num::NonZeroUsize;
 use std::path::PathBuf;
 use std::sync::Arc;
 use tokio::sync::Notify;
@@ -179,8 +179,7 @@ async fn exec_command_with_tty(
         collected_output.total_bytes(),
     ))
     .unwrap_or(usize::MAX);
-    let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes());
-    let collected = collected_output.to_bytes_with_omission_marker();
+    let collected = CapturedStream::from(&collected_output);
     let has_exited = process.has_exited();
     let exit_code = process.exit_code();
     let response_process_id = if process_started_alive && !has_exited {
@@ -212,7 +211,6 @@ async fn exec_command_with_tty(
         process_id: response_process_id,
         exit_code,
         original_token_count: Some(original_token_count),
-        output_omitted_bytes,
         hook_command: Some(cmd.to_string()),
     })
 }

codex-rs/core/src/unified_exec/output.rs

@@ -0,0 +1,63 @@
+//! Text retained for one unified exec output stream and the bytes omitted by its collection cap.
+
+use codex_utils_head_tail_buffer::HeadTailBuffer;
+use codex_utils_head_tail_buffer::format_output_omission_marker;
+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;
+
+#[derive(Clone, Debug, Default, PartialEq)]
+pub(crate) struct CapturedStream {
+    pub(crate) text: String,
+    pub(crate) omitted_bytes: usize,
+}
+
+impl From<&str> for CapturedStream {
+    fn from(text: &str) -> Self {
+        Self {
+            text: text.to_string(),
+            omitted_bytes: 0,
+        }
+    }
+}
+
+impl<const N: usize> From<&HeadTailBuffer<N>> for CapturedStream {
+    fn from(buffer: &HeadTailBuffer<N>) -> Self {
+        Self {
+            text: String::from_utf8_lossy(&buffer.to_bytes_with_omission_marker()).into_owned(),
+            omitted_bytes: buffer.omitted_bytes(),
+        }
+    }
+}
+
+impl CapturedStream {
+    pub(crate) fn render(
+        &self,
+        policy: Option<TruncationPolicy>,
+        original_tokens: Option<usize>,
+    ) -> String {
+        let truncated = match policy {
+            Some(policy) if self.omitted_bytes == 0 => {
+                return formatted_truncate_text(&self.text, policy);
+            }
+            Some(policy) => truncate_text(&self.text, policy),
+            None if self.omitted_bytes == 0 => return self.text.clone(),
+            None => self.text.clone(),
+        };
+        let marker = format_output_omission_marker(self.omitted_bytes);
+        let notice = if truncated.contains(&marker) {
+            String::new()
+        } else {
+            format!("{marker}\n")
+        };
+        if policy.is_none_or(|policy| self.text.len() <= policy.byte_budget()) {
+            if notice.is_empty() {
+                return truncated;
+            }
+            return format!("{notice}{truncated}");
+        }
+        let tokens = original_tokens.unwrap_or_else(|| approx_token_count(&self.text));
+        format!("Warning: truncated output (original token count: {tokens})\n{notice}\n{truncated}")
+    }
+}

codex-rs/core/src/unified_exec/process_manager.rs

@@ -1,3 +1,4 @@
+use crate::unified_exec::output::CapturedStream;
 use rand::Rng;
 use std::cmp::Reverse;
 use std::collections::HashMap;
@@ -686,8 +687,8 @@ impl UnifiedExecProcessManager {
         ))
         .unwrap_or(usize::MAX);
         let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes());
-        let collected = collected_output.to_bytes_with_omission_marker();
-        let text = String::from_utf8_lossy(&collected).to_string();
+        let collected = CapturedStream::from(&collected_output);
+        let text = collected.text.clone();
         let chunk_id = generate_chunk_id();
         if deferred_network_approval
             .as_ref()
@@ -872,7 +873,6 @@ impl UnifiedExecProcessManager {
             process_id: response_process_id,
             exit_code,
             original_token_count: Some(original_token_count),
-            output_omitted_bytes,
             hook_command: Some(request.hook_command.clone()),
         };
 
@@ -1061,8 +1061,7 @@ impl UnifiedExecProcessManager {
             collected_output.total_bytes(),
         ))
         .unwrap_or(usize::MAX);
-        let output_omitted_bytes = NonZeroUsize::new(collected_output.omitted_bytes());
-        let collected = collected_output.to_bytes_with_omission_marker();
+        let collected = CapturedStream::from(&collected_output);
         let chunk_id = generate_chunk_id();
         if network_approval
             .as_ref()
@@ -1126,7 +1125,6 @@ impl UnifiedExecProcessManager {
             process_id,
             exit_code,
             original_token_count: Some(original_token_count),
-            output_omitted_bytes,
             hook_command: Some(hook_command),
         };