fix(core): send shell output to model on cancel (#20501)

2026-03-10 22:21:22 -07:00 · 2026-03-03 14:10:16 -08:00
parent 28e79831ac
commit f3bbe6e77a
8 changed files with 315 additions and 103 deletions
--- a/packages/core/src/core/coreToolScheduler.ts
+++ b/packages/core/src/core/coreToolScheduler.ts
@@ -20,7 +20,10 @@ import { ToolErrorType } from '../tools/tool-error.js';
 import { ToolCallEvent } from '../telemetry/types.js';
 import { runInDevTraceSpan } from '../telemetry/trace.js';
 import { ToolModificationHandler } from '../scheduler/tool-modifier.js';
-import { getToolSuggestion } from '../utils/tool-utils.js';
+import {
+  getToolSuggestion,
+  isToolCallResponseInfo,
+} from '../utils/tool-utils.js';
 import type { ToolConfirmationRequest } from '../confirmation-bus/types.js';
 import { MessageBusType } from '../confirmation-bus/types.js';
 import type { MessageBus } from '../confirmation-bus/message-bus.js';
@@ -225,32 +228,36 @@ export class CoreToolScheduler {
          const durationMs = existingStartTime
            ? Date.now() - existingStartTime
            : undefined;
-          return {
-            request: currentCall.request,
-            tool: toolInstance,
-            invocation,
-            status: CoreToolCallStatus.Success,
-            // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
-            response: auxiliaryData as ToolCallResponseInfo,
-            durationMs,
-            outcome,
-            approvalMode,
-          } as SuccessfulToolCall;
+          if (isToolCallResponseInfo(auxiliaryData)) {
+            return {
+              request: currentCall.request,
+              tool: toolInstance,
+              invocation,
+              status: CoreToolCallStatus.Success,
+              response: auxiliaryData,
+              durationMs,
+              outcome,
+              approvalMode,
+            } as SuccessfulToolCall;
+          }
+          throw new Error('Invalid response data for tool success');
        }
        case CoreToolCallStatus.Error: {
          const durationMs = existingStartTime
            ? Date.now() - existingStartTime
            : undefined;
-          return {
-            request: currentCall.request,
-            status: CoreToolCallStatus.Error,
-            tool: toolInstance,
-            // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
-            response: auxiliaryData as ToolCallResponseInfo,
-            durationMs,
-            outcome,
-            approvalMode,
-          } as ErroredToolCall;
+          if (isToolCallResponseInfo(auxiliaryData)) {
+            return {
+              request: currentCall.request,
+              status: CoreToolCallStatus.Error,
+              tool: toolInstance,
+              response: auxiliaryData,
+              durationMs,
+              outcome,
+              approvalMode,
+            } as ErroredToolCall;
+          }
+          throw new Error('Invalid response data for tool error');
        }
        case CoreToolCallStatus.AwaitingApproval:
          return {
@@ -280,6 +287,19 @@ export class CoreToolScheduler {
            ? Date.now() - existingStartTime
            : undefined;

+          if (isToolCallResponseInfo(auxiliaryData)) {
+            return {
+              request: currentCall.request,
+              tool: toolInstance,
+              invocation,
+              status: CoreToolCallStatus.Cancelled,
+              response: auxiliaryData,
+              durationMs,
+              outcome,
+              approvalMode,
+            } as CancelledToolCall;
+          }
+
          // Preserve diff for cancelled edit operations
          let resultDisplay: ToolResultDisplay | undefined = undefined;
          if (currentCall.status === CoreToolCallStatus.AwaitingApproval) {
--- a/packages/core/src/scheduler/scheduler.test.ts
+++ b/packages/core/src/scheduler/scheduler.test.ts
@@ -946,7 +946,7 @@ describe('Scheduler (Orchestrator)', () => {
      expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
        'call-1',
        CoreToolCallStatus.Cancelled,
-        'Operation cancelled',
+        { callId: 'call-1', responseParts: [] },
      );
    });

--- a/packages/core/src/scheduler/scheduler.ts
+++ b/packages/core/src/scheduler/scheduler.ts
@@ -741,7 +741,7 @@ export class Scheduler {
      this.state.updateStatus(
        callId,
        CoreToolCallStatus.Cancelled,
-        'Operation cancelled',
+        result.response,
      );
    } else {
      this.state.updateStatus(
--- a/packages/core/src/scheduler/state-manager.ts
+++ b/packages/core/src/scheduler/state-manager.ts
@@ -30,6 +30,7 @@ import {
  MessageBusType,
  type SerializableConfirmationDetails,
 } from '../confirmation-bus/types.js';
+import { isToolCallResponseInfo } from '../utils/tool-utils.js';

 /**
 * Handler for terminal tool calls.
@@ -127,7 +128,7 @@ export class SchedulerStateManager {
  updateStatus(
    callId: string,
    status: CoreToolCallStatus.Cancelled,
-    data: string,
+    data: string | ToolCallResponseInfo,
  ): void;
  updateStatus(
    callId: string,
@@ -264,7 +265,7 @@ export class SchedulerStateManager {
  ): ToolCall {
    switch (newStatus) {
      case CoreToolCallStatus.Success: {
-        if (!this.isToolCallResponseInfo(auxiliaryData)) {
+        if (!isToolCallResponseInfo(auxiliaryData)) {
          throw new Error(
            `Invalid data for 'success' transition (callId: ${call.request.callId})`,
          );
@@ -272,7 +273,7 @@ export class SchedulerStateManager {
        return this.toSuccess(call, auxiliaryData);
      }
      case CoreToolCallStatus.Error: {
-        if (!this.isToolCallResponseInfo(auxiliaryData)) {
+        if (!isToolCallResponseInfo(auxiliaryData)) {
          throw new Error(
            `Invalid data for 'error' transition (callId: ${call.request.callId})`,
          );
@@ -290,9 +291,12 @@ export class SchedulerStateManager {
      case CoreToolCallStatus.Scheduled:
        return this.toScheduled(call);
      case CoreToolCallStatus.Cancelled: {
-        if (typeof auxiliaryData !== 'string') {
+        if (
+          typeof auxiliaryData !== 'string' &&
+          !isToolCallResponseInfo(auxiliaryData)
+        ) {
          throw new Error(
-            `Invalid reason (string) for 'cancelled' transition (callId: ${call.request.callId})`,
+            `Invalid reason (string) or response for 'cancelled' transition (callId: ${call.request.callId})`,
          );
        }
        return this.toCancelled(call, auxiliaryData);
@@ -317,15 +321,6 @@ export class SchedulerStateManager {
    }
  }

-  private isToolCallResponseInfo(data: unknown): data is ToolCallResponseInfo {
-    return (
-      typeof data === 'object' &&
-      data !== null &&
-      'callId' in data &&
-      'responseParts' in data
-    );
-  }
-
  private isExecutingToolCallPatch(
    data: unknown,
  ): data is Partial<ExecutingToolCall> {
@@ -451,7 +446,10 @@ export class SchedulerStateManager {
    };
  }

-  private toCancelled(call: ToolCall, reason: string): CancelledToolCall {
+  private toCancelled(
+    call: ToolCall,
+    reason: string | ToolCallResponseInfo,
+  ): CancelledToolCall {
    this.validateHasToolAndInvocation(call, CoreToolCallStatus.Cancelled);
    const startTime = 'startTime' in call ? call.startTime : undefined;

@@ -478,6 +476,20 @@ export class SchedulerStateManager {
      }
    }

+    if (isToolCallResponseInfo(reason)) {
+      return {
+        request: call.request,
+        tool: call.tool,
+        invocation: call.invocation,
+        status: CoreToolCallStatus.Cancelled,
+        response: reason,
+        durationMs: startTime ? Date.now() - startTime : undefined,
+        outcome: call.outcome,
+        schedulerId: call.schedulerId,
+        approvalMode: call.approvalMode,
+      };
+    }
+
    const errorMessage = `[Operation Cancelled] Reason: ${reason}`;
    return {
      request: call.request,
--- a/packages/core/src/scheduler/tool-executor.test.ts
+++ b/packages/core/src/scheduler/tool-executor.test.ts
@@ -534,4 +534,113 @@ describe('ToolExecutor', () => {
      }),
    );
  });
+
+  it('should return cancelled result with partial output when signal is aborted', async () => {
+    const mockTool = new MockTool({
+      name: 'slowTool',
+    });
+    const invocation = mockTool.build({});
+
+    const partialOutput = 'Some partial output before cancellation';
+    vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockImplementation(
+      async () => ({
+        llmContent: partialOutput,
+        returnDisplay: `[Cancelled] ${partialOutput}`,
+      }),
+    );
+
+    const scheduledCall: ScheduledToolCall = {
+      status: CoreToolCallStatus.Scheduled,
+      request: {
+        callId: 'call-cancel-partial',
+        name: 'slowTool',
+        args: {},
+        isClientInitiated: false,
+        prompt_id: 'prompt-cancel',
+      },
+      tool: mockTool,
+      invocation: invocation as unknown as AnyToolInvocation,
+      startTime: Date.now(),
+    };
+
+    const controller = new AbortController();
+    controller.abort();
+
+    const result = await executor.execute({
+      call: scheduledCall,
+      signal: controller.signal,
+      onUpdateToolCall: vi.fn(),
+    });
+
+    expect(result.status).toBe(CoreToolCallStatus.Cancelled);
+    if (result.status === CoreToolCallStatus.Cancelled) {
+      const response = result.response.responseParts[0]?.functionResponse
+        ?.response as Record<string, unknown>;
+      expect(response).toEqual({
+        error: '[Operation Cancelled] User cancelled tool execution.',
+        output: partialOutput,
+      });
+      expect(result.response.resultDisplay).toBe(
+        `[Cancelled] ${partialOutput}`,
+      );
+    }
+  });
+
+  it('should truncate large shell output even on cancellation', 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
+    vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockResolvedValue({
+      llmContent: longOutput,
+      returnDisplay: longOutput,
+    });
+
+    const scheduledCall: ScheduledToolCall = {
+      status: CoreToolCallStatus.Scheduled,
+      request: {
+        callId: 'call-trunc-cancel',
+        name: SHELL_TOOL_NAME,
+        args: { command: 'echo long' },
+        isClientInitiated: false,
+        prompt_id: 'prompt-trunc-cancel',
+      },
+      tool: mockTool,
+      invocation: invocation as unknown as AnyToolInvocation,
+      startTime: Date.now(),
+    };
+
+    // 3. Abort immediately
+    const controller = new AbortController();
+    controller.abort();
+
+    // 4. Execute
+    const result = await executor.execute({
+      call: scheduledCall,
+      signal: controller.signal,
+      onUpdateToolCall: vi.fn(),
+    });
+
+    // 5. Verify Truncation Logic was applied in cancelled path
+    expect(fileUtils.saveTruncatedToolOutput).toHaveBeenCalledWith(
+      longOutput,
+      SHELL_TOOL_NAME,
+      'call-trunc-cancel',
+      expect.any(String),
+      'test-session-id',
+    );
+
+    expect(result.status).toBe(CoreToolCallStatus.Cancelled);
+    if (result.status === CoreToolCallStatus.Cancelled) {
+      const response = result.response.responseParts[0]?.functionResponse
+        ?.response as Record<string, unknown>;
+      expect(response['output']).toBe('TruncatedContent...');
+      expect(result.response.outputFile).toBe('/tmp/truncated_output.txt');
+    }
+  });
 });
--- a/packages/core/src/scheduler/tool-executor.ts
+++ b/packages/core/src/scheduler/tool-executor.ts
@@ -9,7 +9,6 @@ import type {
  ToolCallResponseInfo,
  ToolResult,
  Config,
-  ToolResultDisplay,
  ToolLiveOutput,
 } from '../index.js';
 import {
@@ -19,8 +18,8 @@ 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 { DiscoveredMCPTool } from '../tools/mcp-tool.js';
 import { executeToolWithHooks } from '../core/coreToolHookTriggers.js';
 import {
  saveTruncatedToolOutput,
@@ -36,6 +35,7 @@ import type {
  CancelledToolCall,
 } from './types.js';
 import { CoreToolCallStatus } from './types.js';
+import type { PartListUnion, Part } from '@google/genai';
 import {
  GeminiCliOperation,
  GEN_AI_TOOL_CALL_ID,
@@ -132,10 +132,10 @@ export class ToolExecutor {
          const toolResult: ToolResult = await promise;

          if (signal.aborted) {
-            completedToolCall = this.createCancelledResult(
+            completedToolCall = await this.createCancelledResult(
              call,
              'User cancelled tool execution.',
-              toolResult.returnDisplay,
+              toolResult,
            );
          } else if (toolResult.error === undefined) {
            completedToolCall = await this.createSuccessResult(
@@ -163,7 +163,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.',
            );
@@ -186,56 +186,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 (typeof content === 'string' && toolName === SHELL_TOOL_NAME) {
      const threshold = this.config.getTruncateToolOutputThreshold();
@@ -250,17 +207,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) &&
@@ -288,7 +251,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,
@@ -299,10 +267,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,
--- a/packages/core/src/tools/shell.ts
+++ b/packages/core/src/tools/shell.ts
@@ -388,16 +388,17 @@ export class ShellToolInvocation extends BaseToolInvocation<
      } else {
        if (this.params.is_background || result.backgrounded) {
          returnDisplayMessage = `Command moved to background (PID: ${result.pid}). Output hidden. Press Ctrl+B to view.`;
+        } else if (result.aborted) {
+          const cancelMsg = timeoutMessage || 'Command cancelled by user.';
+          if (result.output.trim()) {
+            returnDisplayMessage = `${cancelMsg}\n\nOutput before cancellation:\n${result.output}`;
+          } else {
+            returnDisplayMessage = cancelMsg;
+          }
        } else if (result.output.trim()) {
          returnDisplayMessage = result.output;
        } else {
-          if (result.aborted) {
-            if (timeoutMessage) {
-              returnDisplayMessage = timeoutMessage;
-            } else {
-              returnDisplayMessage = 'Command cancelled by user.';
-            }
-          } else if (result.signal) {
+          if (result.signal) {
            returnDisplayMessage = `Command terminated by signal: ${result.signal}`;
          } else if (result.error) {
            returnDisplayMessage = `Command failed: ${getErrorMessage(
--- a/packages/core/src/utils/tool-utils.ts
+++ b/packages/core/src/utils/tool-utils.ts
@@ -9,13 +9,30 @@ import { isTool } from '../index.js';
 import { SHELL_TOOL_NAMES } from './shell-utils.js';
 import levenshtein from 'fast-levenshtein';
 import { ApprovalMode } from '../policy/types.js';
-import { CoreToolCallStatus } from '../scheduler/types.js';
+import {
+  CoreToolCallStatus,
+  type ToolCallResponseInfo,
+} from '../scheduler/types.js';
 import {
  ASK_USER_DISPLAY_NAME,
  WRITE_FILE_DISPLAY_NAME,
  EDIT_DISPLAY_NAME,
 } from '../tools/tool-names.js';

+/**
+ * Validates if an object is a ToolCallResponseInfo.
+ */
+export function isToolCallResponseInfo(
+  data: unknown,
+): data is ToolCallResponseInfo {
+  return (
+    typeof data === 'object' &&
+    data !== null &&
+    'callId' in data &&
+    'responseParts' in data
+  );
+}
+
 /**
 * Options for determining if a tool call should be hidden in the CLI history.
 */