Move exec command truncation into ExecCommandToolOutput (#14169)

Summary - relocate truncation logic for exec command output into the new `ExecCommandToolOutput` response helper instead of centralized handler code - update all affected tools and unified exec handling to use the new response item structure and eliminate `Function(FunctionToolOutput)` responses - adjust context, registry, and handler interfaces to align with the new response semantics and error fields Testing - Not run (not requested)
2026-07-01 00:31:56 +08:00 · 2026-03-09 23:13:48 -06:00
parent 0c33af7746
commit a9ae43621b
4 changed files with 154 additions and 75 deletions
@@ -26,7 +26,6 @@ use std::collections::HashSet;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::Weak;
-use std::time::Duration;

 use codex_network_proxy::NetworkProxy;
 use codex_protocol::models::PermissionProfile;
@@ -108,20 +107,6 @@ pub(crate) struct WriteStdinRequest<'a> {
    pub max_output_tokens: Option<usize>,
 }

-#[derive(Debug, Clone, PartialEq)]
-pub(crate) struct UnifiedExecResponse {
-    pub event_call_id: String,
-    pub chunk_id: String,
-    pub wall_time: Duration,
-    pub output: String,
-    /// Raw bytes returned for this unified exec call before any truncation.
-    pub raw_output: Vec<u8>,
-    pub process_id: Option<String>,
-    pub exit_code: Option<i32>,
-    pub original_token_count: Option<usize>,
-    pub session_command: Option<Vec<String>>,
-}
-
 #[derive(Default)]
 pub(crate) struct ProcessStore {
    processes: HashMap<String, ProcessEntry>,
@@ -192,6 +177,7 @@ mod tests {
    use crate::codex::make_session_and_context;
    use crate::protocol::AskForApproval;
    use crate::protocol::SandboxPolicy;
+    use crate::tools::context::ExecCommandToolOutput;
    use crate::unified_exec::ExecCommandRequest;
    use crate::unified_exec::WriteStdinRequest;
    use core_test_support::skip_if_sandbox;
@@ -218,7 +204,7 @@ mod tests {
        turn: &Arc<TurnContext>,
        cmd: &str,
        yield_time_ms: u64,
-    ) -> Result<UnifiedExecResponse, UnifiedExecError> {
+    ) -> Result<ExecCommandToolOutput, UnifiedExecError> {
        let context =
            UnifiedExecContext::new(Arc::clone(session), Arc::clone(turn), "call".to_string());
        let process_id = session
@@ -255,7 +241,7 @@ mod tests {
        process_id: &str,
        input: &str,
        yield_time_ms: u64,
-    ) -> Result<UnifiedExecResponse, UnifiedExecError> {
+    ) -> Result<ExecCommandToolOutput, UnifiedExecError> {
        session
            .services
            .unified_exec_manager
@@ -330,7 +316,7 @@ mod tests {
        )
        .await?;
        assert!(
-            out_2.output.contains("codex"),
+            out_2.truncated_output().contains("codex"),
            "expected environment variable output"
        );

@@ -366,7 +352,7 @@ mod tests {
            "short command should not report a process id if it exits quickly"
        );
        assert!(
-            !out_2.output.contains("codex"),
+            !out_2.truncated_output().contains("codex"),
            "short command should run in a fresh shell"
        );

@@ -382,7 +368,7 @@ mod tests {
        )
        .await?;
        assert!(
-            out_3.output.contains("codex"),
+            out_3.truncated_output().contains("codex"),
            "session should preserve state"
        );

@@ -420,7 +406,7 @@ mod tests {
        )
        .await?;
        assert!(
-            !out_2.output.contains(TEST_VAR_VALUE),
+            !out_2.truncated_output().contains(TEST_VAR_VALUE),
            "timeout too short should yield incomplete output"
        );

@@ -429,7 +415,7 @@ mod tests {
        let out_3 = write_stdin(&session, process_id, "", 100).await?;

        assert!(
-            out_3.output.contains(TEST_VAR_VALUE),
+            out_3.truncated_output().contains(TEST_VAR_VALUE),
            "subsequent poll should retrieve output"
        );

@@ -444,7 +430,7 @@ mod tests {
        let result = exec_command(&session, &turn, "echo codex", 120_000).await?;

        assert!(result.process_id.is_some());
-        assert!(result.output.contains("codex"));
+        assert!(result.truncated_output().contains("codex"));

        Ok(())
    }
@@ -459,7 +445,7 @@ mod tests {
            result.process_id.is_some(),
            "completed command should report a process id"
        );
-        assert!(result.output.contains("codex"));
+        assert!(result.truncated_output().contains("codex"));

        assert!(
            session
@@ -16,6 +16,7 @@ use crate::exec_env::create_env;
 use crate::exec_policy::ExecApprovalRequest;
 use crate::protocol::ExecCommandSource;
 use crate::sandboxing::ExecRequest;
+use crate::tools::context::ExecCommandToolOutput;
 use crate::tools::events::ToolEmitter;
 use crate::tools::events::ToolEventCtx;
 use crate::tools::events::ToolEventStage;
@@ -25,9 +26,7 @@ use crate::tools::orchestrator::ToolOrchestrator;
 use crate::tools::runtimes::unified_exec::UnifiedExecRequest as UnifiedExecToolRequest;
 use crate::tools::runtimes::unified_exec::UnifiedExecRuntime;
 use crate::tools::sandboxing::ToolCtx;
-use crate::truncate::TruncationPolicy;
 use crate::truncate::approx_token_count;
-use crate::truncate::formatted_truncate_text;
 use crate::unified_exec::ExecCommandRequest;
 use crate::unified_exec::MAX_UNIFIED_EXEC_PROCESSES;
 use crate::unified_exec::MAX_YIELD_TIME_MS;
@@ -38,7 +37,6 @@ use crate::unified_exec::ProcessStore;
 use crate::unified_exec::UnifiedExecContext;
 use crate::unified_exec::UnifiedExecError;
 use crate::unified_exec::UnifiedExecProcessManager;
-use crate::unified_exec::UnifiedExecResponse;
 use crate::unified_exec::WARNING_UNIFIED_EXEC_PROCESSES;
 use crate::unified_exec::WriteStdinRequest;
 use crate::unified_exec::async_watcher::emit_exec_end_for_unified_exec;
@@ -51,7 +49,6 @@ use crate::unified_exec::process::OutputBuffer;
 use crate::unified_exec::process::OutputHandles;
 use crate::unified_exec::process::SpawnLifecycleHandle;
 use crate::unified_exec::process::UnifiedExecProcess;
-use crate::unified_exec::resolve_max_tokens;

 const UNIFIED_EXEC_ENV: [(&str, &str); 10] = [
    ("NO_COLOR", "1"),
@@ -159,7 +156,7 @@ impl UnifiedExecProcessManager {
        &self,
        request: ExecCommandRequest,
        context: &UnifiedExecContext,
-    ) -> Result<UnifiedExecResponse, UnifiedExecError> {
+    ) -> Result<ExecCommandToolOutput, UnifiedExecError> {
        let cwd = request
            .workdir
            .clone()
@@ -194,7 +191,6 @@ impl UnifiedExecProcessManager {
        emitter.emit(event_ctx, ToolEventStage::Begin).await;

        start_streaming_output(&process, context, Arc::clone(&transcript));
-        let max_tokens = resolve_max_tokens(request.max_output_tokens);
        let yield_time_ms = clamp_yield_time(request.yield_time_ms);

        let start = Instant::now();
@@ -221,7 +217,6 @@ impl UnifiedExecProcessManager {
        let wall_time = Instant::now().saturating_duration_since(start);

        let text = String::from_utf8_lossy(&collected).to_string();
-        let output = formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens));
        let exit_code = process.exit_code();
        let has_exited = process.has_exited() || exit_code.is_some();
        let chunk_id = generate_chunk_id();
@@ -240,7 +235,7 @@ impl UnifiedExecProcessManager {
                cwd.clone(),
                Some(process_id),
                Arc::clone(&transcript),
-                output.clone(),
+                text.clone(),
                exit,
                wall_time,
            )
@@ -276,12 +271,12 @@ impl UnifiedExecProcessManager {
        };

        let original_token_count = approx_token_count(&text);
-        let response = UnifiedExecResponse {
+        let response = ExecCommandToolOutput {
            event_call_id: context.call_id.clone(),
            chunk_id,
            wall_time,
-            output,
            raw_output: collected,
+            max_output_tokens: request.max_output_tokens,
            process_id: if has_exited {
                None
            } else {
@@ -298,7 +293,7 @@ impl UnifiedExecProcessManager {
    pub(crate) async fn write_stdin(
        &self,
        request: WriteStdinRequest<'_>,
-    ) -> Result<UnifiedExecResponse, UnifiedExecError> {
+    ) -> Result<ExecCommandToolOutput, UnifiedExecError> {
        let process_id = request.process_id.to_string();

        let PreparedProcessHandles {
@@ -324,7 +319,6 @@ impl UnifiedExecProcessManager {
            tokio::time::sleep(Duration::from_millis(100)).await;
        }

-        let max_tokens = resolve_max_tokens(request.max_output_tokens);
        let yield_time_ms = {
            // Empty polls use configurable background timeout bounds. Non-empty
            // writes keep a fixed max cap so interactive stdin remains responsive.
@@ -349,7 +343,6 @@ impl UnifiedExecProcessManager {
        let wall_time = Instant::now().saturating_duration_since(start);

        let text = String::from_utf8_lossy(&collected).to_string();
-        let output = formatted_truncate_text(&text, TruncationPolicy::Tokens(max_tokens));
        let original_token_count = approx_token_count(&text);
        let chunk_id = generate_chunk_id();

@@ -375,12 +368,12 @@ impl UnifiedExecProcessManager {
            }
        };

-        let response = UnifiedExecResponse {
+        let response = ExecCommandToolOutput {
            event_call_id,
            chunk_id,
            wall_time,
-            output,
            raw_output: collected,
+            max_output_tokens: request.max_output_tokens,
            process_id,
            exit_code,
            original_token_count: Some(original_token_count),