Merge branch 'main' into memory_usage3

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

@@ -6,23 +6,21 @@
 
 import fsPromises from 'node:fs/promises';
 import { debugLogger } from '../utils/debugLogger.js';
-import type {
-  ToolCallRequestInfo,
-  ToolCallResponseInfo,
-  ToolResult,
-  Config,
-  ToolResultDisplay,
-  ToolLiveOutput,
-} from '../index.js';
 import {
   ToolErrorType,
   ToolOutputTruncatedEvent,
   logToolOutputTruncated,
   runInDevTraceSpan,
+  type ToolCallRequestInfo,
+  type ToolCallResponseInfo,
+  type ToolResult,
+  type ToolResultDisplay,
+  type Config,
+  type ToolLiveOutput,
 } from '../index.js';
 import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
-import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
 import { ShellToolInvocation } from '../tools/shell.js';
+import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
 import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
 import {
   saveTruncatedToolOutput,
@@ -30,15 +28,16 @@ import {
   moveToolOutputToFile,
 } from '../utils/fileUtils.js';
 import { convertToFunctionResponse } from '../utils/generateContentResponseUtilities.js';
-import type {
-  CompletedToolCall,
-  ToolCall,
-  ExecutingToolCall,
-  ErroredToolCall,
-  SuccessfulToolCall,
-  CancelledToolCall,
+import {
+  CoreToolCallStatus,
+  type CompletedToolCall,
+  type ToolCall,
+  type ExecutingToolCall,
+  type ErroredToolCall,
+  type SuccessfulToolCall,
+  type CancelledToolCall,
 } from './types.js';
-import { CoreToolCallStatus } from './types.js';
+import type { PartListUnion, Part } from '@google/genai';
 import {
   GeminiCliOperation,
   GEN_AI_TOOL_CALL_ID,
@@ -145,10 +144,10 @@ export class ToolExecutor {
                   );
                 });
             }
-            completedToolCall = this.createCancelledResult(
+            completedToolCall = await this.createCancelledResult(
               call,
               'User cancelled tool execution.',
-              toolResult.returnDisplay,
+              toolResult,
             );
           } else if (toolResult.error === undefined) {
             completedToolCall = await this.createSuccessResult(
@@ -186,7 +185,7 @@ export class ToolExecutor {
               executionError.message.includes('Operation cancelled by user'));
 
           if (signal.aborted || isAbortError) {
-            completedToolCall = this.createCancelledResult(
+            completedToolCall = await this.createCancelledResult(
               call,
               'User cancelled tool execution.',
             );
@@ -209,56 +208,13 @@ export class ToolExecutor {
     );
   }
 
-  private createCancelledResult(
+  private async truncateOutputIfNeeded(
     call: ToolCall,
-    reason: string,
-    resultDisplay?: ToolResultDisplay,
-  ): CancelledToolCall {
-    const errorMessage = `[Operation Cancelled] ${reason}`;
-    const startTime = 'startTime' in call ? call.startTime : undefined;
-
-    if (!('tool' in call) || !('invocation' in call)) {
-      // This should effectively never happen in execution phase, but we handle
-      // it safely
-      throw new Error('Cancelled tool call missing tool/invocation references');
-    }
-
-    return {
-      status: CoreToolCallStatus.Cancelled,
-      request: call.request,
-      response: {
-        callId: call.request.callId,
-        responseParts: [
-          {
-            functionResponse: {
-              id: call.request.callId,
-              name: call.request.name,
-              response: { error: errorMessage },
-            },
-          },
-        ],
-        resultDisplay,
-        error: undefined,
-        errorType: undefined,
-        contentLength: errorMessage.length,
-      },
-      tool: call.tool,
-      invocation: call.invocation,
-      durationMs: startTime ? Date.now() - startTime : undefined,
-      startTime,
-      endTime: Date.now(),
-      outcome: call.outcome,
-    };
-  }
-
-  private async createSuccessResult(
-    call: ToolCall,
-    toolResult: ToolResult,
-  ): Promise<SuccessfulToolCall> {
-    let content = toolResult.llmContent;
-    let outputFile: string | undefined;
-    const toolName = call.request.originalRequestName || call.request.name;
+    content: PartListUnion,
+  ): Promise<{ truncatedContent: PartListUnion; outputFile?: string }> {
+    const toolName = call.request.name;
     const callId = call.request.callId;
+    let outputFile: string | undefined;
 
     if (toolResult.fullOutputFilePath) {
       const threshold = this.config.getTruncateToolOutputThreshold();
@@ -300,17 +256,23 @@ export class ToolExecutor {
           this.config.getSessionId(),
         );
         outputFile = savedPath;
-        content = formatTruncatedToolOutput(content, outputFile, threshold);
+        const truncatedContent = formatTruncatedToolOutput(
+          content,
+          outputFile,
+          threshold,
+        );
 
         logToolOutputTruncated(
           this.config,
           new ToolOutputTruncatedEvent(call.request.prompt_id, {
             toolName,
             originalContentLength,
-            truncatedContentLength: content.length,
+            truncatedContentLength: truncatedContent.length,
             threshold,
           }),
         );
+
+        return { truncatedContent, outputFile };
       }
     } else if (
       Array.isArray(content) &&
@@ -338,7 +300,12 @@ export class ToolExecutor {
             outputFile,
             threshold,
           );
-          content[0] = { ...firstPart, text: truncatedText };
+
+          // We need to return a NEW array to avoid mutating the original toolResult if it matters,
+          // though here we are creating the response so it's probably fine to mutate or return new.
+          const truncatedContent: Part[] = [
+            { ...firstPart, text: truncatedText },
+          ];
 
           logToolOutputTruncated(
             this.config,
@@ -349,10 +316,95 @@ export class ToolExecutor {
               threshold,
             }),
           );
+
+          return { truncatedContent, outputFile };
         }
       }
     }
 
+    return { truncatedContent: content, outputFile };
+  }
+
+  private async createCancelledResult(
+    call: ToolCall,
+    reason: string,
+    toolResult?: ToolResult,
+  ): Promise<CancelledToolCall> {
+    const errorMessage = `[Operation Cancelled] ${reason}`;
+    const startTime = 'startTime' in call ? call.startTime : undefined;
+
+    if (!('tool' in call) || !('invocation' in call)) {
+      // This should effectively never happen in execution phase, but we handle
+      // it safely
+      throw new Error('Cancelled tool call missing tool/invocation references');
+    }
+
+    let responseParts: Part[] = [];
+    let outputFile: string | undefined;
+
+    if (toolResult?.llmContent) {
+      // 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);
+
+      outputFile = truncatedOutputFile;
+      responseParts = convertToFunctionResponse(
+        call.request.name,
+        call.request.callId,
+        output,
+        this.config.getActiveModel(),
+      );
+
+      // Inject the cancellation error into the response object
+      const mainPart = responseParts[0];
+      if (mainPart?.functionResponse?.response) {
+        const respObj = mainPart.functionResponse.response;
+        respObj['error'] = errorMessage;
+      }
+    } else {
+      responseParts = [
+        {
+          functionResponse: {
+            id: call.request.callId,
+            name: call.request.name,
+            response: { error: errorMessage },
+          },
+        },
+      ];
+    }
+
+    return {
+      status: CoreToolCallStatus.Cancelled,
+      request: call.request,
+      response: {
+        callId: call.request.callId,
+        responseParts,
+        resultDisplay: toolResult?.returnDisplay,
+        error: undefined,
+        errorType: undefined,
+        outputFile,
+        contentLength: JSON.stringify(responseParts).length,
+      },
+      tool: call.tool,
+      invocation: call.invocation,
+      durationMs: startTime ? Date.now() - startTime : undefined,
+      startTime,
+      endTime: Date.now(),
+      outcome: call.outcome,
+    };
+  }
+
+  private async createSuccessResult(
+    call: ToolCall,
+    toolResult: ToolResult,
+  ): Promise<SuccessfulToolCall> {
+    const { truncatedContent: content, outputFile } =
+      await this.truncateOutputIfNeeded(call, toolResult.llmContent);
+
+    const toolName = call.request.originalRequestName || call.request.name;
+    const callId = call.request.callId;
+
     const response = convertToFunctionResponse(
       toolName,
       callId,