chore: fix truncation logic and test duplications

packages/core/src/scheduler/tool-executor.test.ts

@@ -5,6 +5,7 @@
  */
 
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import fsPromises from 'node:fs/promises';
 import { ToolExecutor } from './tool-executor.js';
 import {
   type Config,
@@ -32,6 +33,7 @@ import {
 vi.mock('../utils/fileUtils.js', () => ({
   saveTruncatedToolOutput: vi.fn(),
   formatTruncatedToolOutput: vi.fn(),
+  moveToolOutputToFile: vi.fn(),
 }));
 
 // Mock executeToolWithHooks
@@ -431,68 +433,6 @@ describe('ToolExecutor', () => {
       10, // threshold (maxChars)
     );
 
-    expect(result.status).toBe(CoreToolCallStatus.Success);
-    if (result.status === CoreToolCallStatus.Success) {
-      const response = result.response.responseParts[0]?.functionResponse
-        ?.response as Record<string, unknown>;
-      // The content should be the *truncated* version returned by the mock formatTruncatedToolOutput
-      expect(response).toEqual({ output: 'TruncatedContent...' });
-      expect(result.response.outputFile).toBe('/tmp/truncated_output.txt');
-    }
-  });
-
-  it('should truncate large output and move file when fullOutputFilePath is provided', async () => {
-    // 1. Setup Config for Truncation
-    vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10);
-    vi.spyOn(config.storage, 'getProjectTempDir').mockReturnValue('/tmp');
-
-    const mockTool = new MockTool({ name: SHELL_TOOL_NAME });
-    const invocation = mockTool.build({});
-    const longOutput = 'This is a very long output that should be truncated.';
-
-    // 2. Mock execution returning long content AND fullOutputFilePath
-    vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
-      llmContent: longOutput,
-      returnDisplay: longOutput,
-      fullOutputFilePath: '/tmp/temp_full_output.txt',
-    });
-
-    const scheduledCall: ScheduledToolCall = {
-      status: CoreToolCallStatus.Scheduled,
-      request: {
-        callId: 'call-trunc-full',
-        name: SHELL_TOOL_NAME,
-        args: { command: 'echo long' },
-        isClientInitiated: false,
-        prompt_id: 'prompt-trunc-full',
-      },
-      tool: mockTool,
-      invocation: invocation as unknown as AnyToolInvocation,
-      startTime: Date.now(),
-    };
-
-    // 3. Execute
-    const result = await executor.execute({
-      call: scheduledCall,
-      signal: new AbortController().signal,
-      onUpdateToolCall: vi.fn(),
-    });
-
-    // 4. Verify Truncation Logic
-    expect(fileUtils.moveToolOutputToFile).toHaveBeenCalledWith(
-      '/tmp/temp_full_output.txt',
-      SHELL_TOOL_NAME,
-      'call-trunc-full',
-      expect.any(String), // temp dir
-      'test-session-id', // session id from makeFakeConfig
-    );
-
-    expect(fileUtils.formatTruncatedToolOutput).toHaveBeenCalledWith(
-      longOutput,
-      '/tmp/moved_output.txt',
-      10, // threshold (maxChars)
-    );
-
     expect(result.status).toBe(CoreToolCallStatus.Success);
     if (result.status === CoreToolCallStatus.Success) {
       const response = result.response.responseParts[0]?.functionResponse
@@ -500,76 +440,8 @@ describe('ToolExecutor', () => {
       // The content should be the *truncated* version returned by the mock formatTruncatedToolOutput
       expect(response).toEqual({
         output: 'TruncatedContent...',
-        outputFile: '/tmp/moved_output.txt',
+        outputFile: '/tmp/truncated_output.txt',
       });
-      expect(result.response.outputFile).toBe('/tmp/moved_output.txt');
-    }
-  });
-
-  it('should truncate large MCP tool output with single text Part', async () => {
-    // 1. Setup Config for Truncation
-    vi.spyOn(config, 'getTruncateToolOutputThreshold').mockReturnValue(10);
-    vi.spyOn(config.storage, 'getProjectTempDir').mockReturnValue('/tmp');
-
-    const mcpToolName = 'get_big_text';
-    const messageBus = createMockMessageBus();
-    const mcpTool = new DiscoveredMCPTool(
-      {} as CallableTool,
-      'my-server',
-      'get_big_text',
-      'A test MCP tool',
-      {},
-      messageBus,
-    );
-    const invocation = mcpTool.build({});
-    const longText = 'This is a very long MCP output that should be truncated.';
-
-    // 2. Mock execution returning Part[] with single text Part
-    // We do NOT provide fullOutputFilePath here because we want to test the path
-    // that uses saveTruncatedToolOutput for MCP tools.
-    vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
-      llmContent: [{ text: longText }],
-      returnDisplay: longText,
-    });
-
-    const scheduledCall: ScheduledToolCall = {
-      status: CoreToolCallStatus.Scheduled,
-      request: {
-        callId: 'call-mcp-trunc',
-        name: mcpToolName,
-        args: { query: 'test' },
-        isClientInitiated: false,
-        prompt_id: 'prompt-mcp-trunc',
-      },
-      tool: mcpTool,
-      invocation: invocation as unknown as AnyToolInvocation,
-      startTime: Date.now(),
-    };
-
-    // 3. Execute
-    const result = await executor.execute({
-      call: scheduledCall,
-      signal: new AbortController().signal,
-      onUpdateToolCall: vi.fn(),
-    });
-
-    // 4. Verify Truncation Logic
-    expect(fileUtils.saveTruncatedToolOutput).toHaveBeenCalledWith(
-      longText,
-      mcpToolName,
-      'call-mcp-trunc',
-      expect.any(String),
-      'test-session-id',
-    );
-
-    expect(fileUtils.formatTruncatedToolOutput).toHaveBeenCalledWith(
-      longText,
-      '/tmp/truncated_output.txt',
-      10,
-    );
-
-    expect(result.status).toBe(CoreToolCallStatus.Success);
-    if (result.status === CoreToolCallStatus.Success) {
       expect(result.response.outputFile).toBe('/tmp/truncated_output.txt');
     }
   });

packages/core/src/scheduler/tool-executor.ts

@@ -26,6 +26,7 @@ import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
 import {
   saveTruncatedToolOutput,
   formatTruncatedToolOutput,
+  moveToolOutputToFile,
 } from '../utils/fileUtils.js';
 import { convertToFunctionResponse } from '../utils/generateContentResponseUtilities.js';
 import {
@@ -161,6 +162,16 @@ export class ToolExecutor {
               toolResult,
             );
           } else {
+            if (toolResult.fullOutputFilePath) {
+              await fsPromises
+                .unlink(toolResult.fullOutputFilePath)
+                .catch((error) => {
+                  debugLogger.warn(
+                    `Failed to delete temporary tool output file on error: ${toolResult.fullOutputFilePath}`,
+                    error,
+                  );
+                });
+            }
             const displayText =
               typeof toolResult.returnDisplay === 'string'
                 ? toolResult.returnDisplay
@@ -209,6 +220,7 @@ export class ToolExecutor {
   private async truncateOutputIfNeeded(
     call: ToolCall,
     content: PartListUnion,
+    fullOutputFilePath?: string,
   ): Promise<{ truncatedContent: PartListUnion; outputFile?: string }> {
     if (this.config.isContextManagementEnabled()) {
       const distiller = new ToolOutputDistillationService(
@@ -216,7 +228,22 @@ export class ToolExecutor {
         this.context.geminiClient,
         this.context.promptId,
       );
-      return distiller.distill(call.request.name, call.request.callId, content);
+      const result = await distiller.distill(
+        call.request.name,
+        call.request.callId,
+        content,
+      );
+      if (fullOutputFilePath && !result.outputFile) {
+        try {
+          await fsPromises.unlink(fullOutputFilePath);
+        } catch (error) {
+          debugLogger.warn(
+            `Failed to delete temporary tool output file: ${fullOutputFilePath}`,
+            error,
+          );
+        }
+      }
+      return result;
     }
 
     const toolName = call.request.name;
@@ -228,13 +255,28 @@ export class ToolExecutor {
 
       if (threshold > 0 && content.length > threshold) {
         const originalContentLength = content.length;
-        const { outputFile: savedPath } = await saveTruncatedToolOutput(
+
-          content,
+        let savedPath: string;
-          toolName,
+        if (fullOutputFilePath) {
-          callId,
+          const { outputFile: movedPath } = await moveToolOutputToFile(
-          this.config.storage.getProjectTempDir(),
+            fullOutputFilePath,
-          this.context.promptId,
+            toolName,
-        );
+            callId,
+            this.config.storage.getProjectTempDir(),
+            this.config.getSessionId(),
+          );
+          savedPath = movedPath;
+        } else {
+          const { outputFile: writtenPath } = await saveTruncatedToolOutput(
+            content,
+            toolName,
+            callId,
+            this.config.storage.getProjectTempDir(),
+            this.context.promptId,
+          );
+          savedPath = writtenPath;
+        }
+
         outputFile = savedPath;
         const truncatedContent = formatTruncatedToolOutput(
           content,
@@ -297,11 +339,33 @@ export class ToolExecutor {
             }),
           );
 
+          if (fullOutputFilePath) {
+            try {
+              await fsPromises.unlink(fullOutputFilePath);
+            } catch (error) {
+              debugLogger.warn(
+                `Failed to delete temporary tool output file: ${fullOutputFilePath}`,
+                error,
+              );
+            }
+          }
+
           return { truncatedContent, outputFile };
         }
       }
     }
 
+    if (fullOutputFilePath && !outputFile) {
+      try {
+        await fsPromises.unlink(fullOutputFilePath);
+      } catch (error) {
+        debugLogger.warn(
+          `Failed to delete temporary tool output file: ${fullOutputFilePath}`,
+          error,
+        );
+      }
+    }
+
     return { truncatedContent: content, outputFile };
   }
 
@@ -326,7 +390,11 @@ export class ToolExecutor {
       // Attempt to truncate and save output if we have content, even in cancellation case
       // This is to handle cases where the tool may have produced output before cancellation
       const { truncatedContent: output, outputFile: truncatedOutputFile } =
-        await this.truncateOutputIfNeeded(call, toolResult?.llmContent);
+        await this.truncateOutputIfNeeded(
+          call,
+          toolResult.llmContent,
+          toolResult.fullOutputFilePath,
+        );
 
       outputFile = truncatedOutputFile;
       responseParts = convertToFunctionResponse(
@@ -335,6 +403,7 @@ export class ToolExecutor {
         output,
         this.config.getActiveModel(),
         this.config,
+        outputFile,
       );
 
       // Inject the cancellation error into the response object
@@ -344,6 +413,16 @@ export class ToolExecutor {
         respObj['error'] = errorMessage;
       }
     } else {
+      if (toolResult?.fullOutputFilePath) {
+        try {
+          await fsPromises.unlink(toolResult.fullOutputFilePath);
+        } catch (error) {
+          debugLogger.warn(
+            `Failed to delete temporary tool output file: ${toolResult.fullOutputFilePath}`,
+            error,
+          );
+        }
+      }
       responseParts = [
         {
           functionResponse: {
@@ -380,56 +459,22 @@ export class ToolExecutor {
     call: ToolCall,
     toolResult: ToolResult,
   ): Promise<SuccessfulToolCall> {
-    let { truncatedContent: content, outputFile } =
+    const { truncatedContent: content, outputFile } =
-      await this.truncateOutputIfNeeded(call, toolResult.llmContent);
+      await this.truncateOutputIfNeeded(
+        call,
+        toolResult.llmContent,
+        toolResult.fullOutputFilePath,
+      );
     const toolName = call.request.originalRequestName || call.request.name;
     const callId = call.request.callId;
 
-    if (toolResult.fullOutputFilePath) {
-      const threshold = this.config.getTruncateToolOutputThreshold();
-      if (
-        threshold > 0 &&
-        typeof content === 'string' &&
-        content.length > threshold
-      ) {
-        const { outputFile: savedPath } = await moveToolOutputToFile(
-          toolResult.fullOutputFilePath,
-          toolName,
-          callId,
-          this.config.storage.getProjectTempDir(),
-          this.config.getSessionId(),
-        );
-        outputFile = savedPath;
-        content = formatTruncatedToolOutput(content, outputFile, threshold);
-
-        logToolOutputTruncated(
-          this.config,
-          new ToolOutputTruncatedEvent(call.request.prompt_id, {
-            toolName,
-            originalContentLength: content.length, // approximation
-            truncatedContentLength: content.length,
-            threshold,
-          }),
-        );
-      } else {
-        // If the content is not truncated, we don't need the temporary file.
-        try {
-          await fsPromises.unlink(toolResult.fullOutputFilePath);
-        } catch (error) {
-          debugLogger.warn(
-            `Failed to delete temporary tool output file: ${toolResult.fullOutputFilePath}`,
-            error,
-          );
-        }
-      }
-    }
-
     const response = convertToFunctionResponse(
       toolName,
       callId,
       content,
       this.config.getActiveModel(),
       this.config,
+      outputFile,
     );
 
     const successResponse: ToolCallResponseInfo = {

packages/core/src/services/shellExecutionService.ts

@@ -1182,7 +1182,8 @@ export class ShellExecutionService {
 
       let lastYdisp = 0;
       let hasReachedMax = false;
-      const scrollbackLimit = shellExecutionConfig.scrollback ?? SCROLLBACK_LIMIT;
+      const scrollbackLimit =
+        shellExecutionConfig.scrollback ?? SCROLLBACK_LIMIT;
 
       headlessTerminal.onScroll((ydisp) => {
         if (!isWriting) {

packages/core/src/tools/shell.ts

@@ -9,7 +9,6 @@ import fs from 'node:fs';
 import path from 'node:path';
 import os from 'node:os';
 import crypto from 'node:crypto';
-import type { Config } from '../config/config.js';
 import { debugLogger } from '../utils/debugLogger.js';
 import {
   type SandboxPermissions,
@@ -937,7 +936,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
           this.context.geminiClient,
           signal,
         );
-        const threshold = this.config.getTruncateToolOutputThreshold();
+        const threshold = this.context.config.getTruncateToolOutputThreshold();
         const fullOutputFilePath =
           threshold > 0 && totalBytesWritten >= threshold
             ? outputFilePath
@@ -945,7 +944,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
 
         const toolResult: ToolResult = {
           llmContent: summary,
-          returnDisplay: typeof returnDisplayMessage !== 'undefined' ? returnDisplayMessage : returnDisplay,
+          returnDisplay,
           fullOutputFilePath,
           ...executionError,
         };
@@ -955,7 +954,7 @@ export class ShellToolInvocation extends BaseToolInvocation<
         return toolResult;
       }
 
-      const threshold = this.config.getTruncateToolOutputThreshold();
+      const threshold = this.context.config.getTruncateToolOutputThreshold();
       const fullOutputFilePath =
         threshold > 0 && totalBytesWritten >= threshold
           ? outputFilePath