Merge branch 'main' into memory_usage3

2026-05-02 16:04:38 -07:00 · 2026-03-03 13:22:09 -05:00
parent c61bef6789 50af050623
commit 1f357fcd89
510 changed files with 25967 additions and 7445 deletions
@@ -11,7 +11,8 @@ import type {
  ToolCallResponseInfo,
  ToolResult,
  Config,
-  AnsiOutput,
+  ToolResultDisplay,
+  ToolLiveOutput,
 } from '../index.js';
 import {
  ToolErrorType,
@@ -20,6 +21,7 @@ import {
  runInDevTraceSpan,
 } 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 { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
 import {
@@ -37,11 +39,17 @@ import type {
  CancelledToolCall,
 } from './types.js';
 import { CoreToolCallStatus } from './types.js';
+import {
+  GeminiCliOperation,
+  GEN_AI_TOOL_CALL_ID,
+  GEN_AI_TOOL_DESCRIPTION,
+  GEN_AI_TOOL_NAME,
+} from '../telemetry/constants.js';

 export interface ToolExecutionContext {
  call: ToolCall;
  signal: AbortSignal;
-  outputUpdateHandler?: (callId: string, output: string | AnsiOutput) => void;
+  outputUpdateHandler?: (callId: string, output: ToolLiveOutput) => void;
  onUpdateToolCall: (updatedCall: ToolCall) => void;
 }

@@ -64,7 +72,7 @@ export class ToolExecutor {
    // Setup live output handling
    const liveOutputCallback =
      tool.canUpdateOutput && outputUpdateHandler
-        ? (outputChunk: string | AnsiOutput) => {
+        ? (outputChunk: ToolLiveOutput) => {
            outputUpdateHandler(callId, outputChunk);
          }
        : undefined;
@@ -73,11 +81,17 @@ export class ToolExecutor {

    return runInDevTraceSpan(
      {
-        name: tool.name,
-        attributes: { type: 'tool-call' },
+        operation: GeminiCliOperation.ToolCall,
+        attributes: {
+          [GEN_AI_TOOL_NAME]: toolName,
+          [GEN_AI_TOOL_CALL_ID]: callId,
+          [GEN_AI_TOOL_DESCRIPTION]: tool.description,
+        },
      },
      async ({ metadata: spanMetadata }) => {
-        spanMetadata.input = { request };
+        spanMetadata.input = request;
+
+        let completedToolCall: CompletedToolCall;

        try {
          let promise: Promise<ToolResult>;
@@ -119,7 +133,6 @@ export class ToolExecutor {
          }

          const toolResult: ToolResult = await promise;
-          spanMetadata.output = toolResult;

          if (signal.aborted) {
            if (toolResult.fullOutputFilePath) {
@@ -132,12 +145,16 @@ export class ToolExecutor {
                  );
                });
            }
-            return this.createCancelledResult(
+            completedToolCall = this.createCancelledResult(
              call,
              'User cancelled tool execution.',
+              toolResult.returnDisplay,
            );
          } else if (toolResult.error === undefined) {
-            return await this.createSuccessResult(call, toolResult);
+            completedToolCall = await this.createSuccessResult(
+              call,
+              toolResult,
+            );
          } else {
            if (toolResult.fullOutputFilePath) {
              await fsPromises
@@ -153,7 +170,7 @@ export class ToolExecutor {
              typeof toolResult.returnDisplay === 'string'
                ? toolResult.returnDisplay
                : undefined;
-            return this.createErrorResult(
+            completedToolCall = this.createErrorResult(
              call,
              new Error(toolResult.error.message),
              toolResult.error.type,
@@ -163,22 +180,31 @@ export class ToolExecutor {
          }
        } catch (executionError: unknown) {
          spanMetadata.error = executionError;
-          if (signal.aborted) {
-            return this.createCancelledResult(
+          const isAbortError =
+            executionError instanceof Error &&
+            (executionError.name === 'AbortError' ||
+              executionError.message.includes('Operation cancelled by user'));
+
+          if (signal.aborted || isAbortError) {
+            completedToolCall = this.createCancelledResult(
              call,
              'User cancelled tool execution.',
            );
+          } else {
+            const error =
+              executionError instanceof Error
+                ? executionError
+                : new Error(String(executionError));
+            completedToolCall = this.createErrorResult(
+              call,
+              error,
+              ToolErrorType.UNHANDLED_EXCEPTION,
+            );
          }
-          const error =
-            executionError instanceof Error
-              ? executionError
-              : new Error(String(executionError));
-          return this.createErrorResult(
-            call,
-            error,
-            ToolErrorType.UNHANDLED_EXCEPTION,
-          );
        }
+
+        spanMetadata.output = completedToolCall;
+        return completedToolCall;
      },
    );
  }
@@ -186,6 +212,7 @@ export class ToolExecutor {
  private createCancelledResult(
    call: ToolCall,
    reason: string,
+    resultDisplay?: ToolResultDisplay,
  ): CancelledToolCall {
    const errorMessage = `[Operation Cancelled] ${reason}`;
    const startTime = 'startTime' in call ? call.startTime : undefined;
@@ -210,7 +237,7 @@ export class ToolExecutor {
            },
          },
        ],
-        resultDisplay: undefined,
+        resultDisplay,
        error: undefined,
        errorType: undefined,
        contentLength: errorMessage.length,
@@ -285,6 +312,45 @@ export class ToolExecutor {
          }),
        );
      }
+    } else if (
+      Array.isArray(content) &&
+      content.length === 1 &&
+      'tool' in call &&
+      call.tool instanceof DiscoveredMCPTool
+    ) {
+      const firstPart = content[0];
+      if (typeof firstPart === 'object' && typeof firstPart.text === 'string') {
+        const textContent = firstPart.text;
+        const threshold = this.config.getTruncateToolOutputThreshold();
+
+        if (threshold > 0 && textContent.length > threshold) {
+          const originalContentLength = textContent.length;
+          const { outputFile: savedPath } = await saveTruncatedToolOutput(
+            textContent,
+            toolName,
+            callId,
+            this.config.storage.getProjectTempDir(),
+            this.config.getSessionId(),
+          );
+          outputFile = savedPath;
+          const truncatedText = formatTruncatedToolOutput(
+            textContent,
+            outputFile,
+            threshold,
+          );
+          content[0] = { ...firstPart, text: truncatedText };
+
+          logToolOutputTruncated(
+            this.config,
+            new ToolOutputTruncatedEvent(call.request.prompt_id, {
+              toolName,
+              originalContentLength,
+              truncatedContentLength: truncatedText.length,
+              threshold,
+            }),
+          );
+        }
+      }
    }

    const response = convertToFunctionResponse(