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

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 {
+          if (isToolCallResponseInfo(auxiliaryData)) {
-            request: currentCall.request,
+            return {
-            tool: toolInstance,
+              request: currentCall.request,
-            invocation,
+              tool: toolInstance,
-            status: CoreToolCallStatus.Success,
+              invocation,
-            // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+              status: CoreToolCallStatus.Success,
-            response: auxiliaryData as ToolCallResponseInfo,
+              response: auxiliaryData,
-            durationMs,
+              durationMs,
-            outcome,
+              outcome,
-            approvalMode,
+              approvalMode,
-          } as SuccessfulToolCall;
+            } as SuccessfulToolCall;
+          }
+          throw new Error('Invalid response data for tool success');
         }
         case CoreToolCallStatus.Error: {
           const durationMs = existingStartTime
             ? Date.now() - existingStartTime
             : undefined;
-          return {
+          if (isToolCallResponseInfo(auxiliaryData)) {
-            request: currentCall.request,
+            return {
-            status: CoreToolCallStatus.Error,
+              request: currentCall.request,
-            tool: toolInstance,
+              status: CoreToolCallStatus.Error,
-            // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
+              tool: toolInstance,
-            response: auxiliaryData as ToolCallResponseInfo,
+              response: auxiliaryData,
-            durationMs,
+              durationMs,
-            outcome,
+              outcome,
-            approvalMode,
+              approvalMode,
-          } as ErroredToolCall;
+            } 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) {

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: [] },
       );
     });
 

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(

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,

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');
+    }
+  });
 });

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,
+    content: PartListUnion,
-    resultDisplay?: ToolResultDisplay,
+  ): Promise<{ truncatedContent: PartListUnion; outputFile?: string }> {
-  ): CancelledToolCall {
+    const toolName = call.request.name;
-    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;
     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,

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 (result.signal) {
-            if (timeoutMessage) {
-              returnDisplayMessage = timeoutMessage;
-            } else {
-              returnDisplayMessage = 'Command cancelled by user.';
-            }
-          } else if (result.signal) {
             returnDisplayMessage = `Command terminated by signal: ${result.signal}`;
           } else if (result.error) {
             returnDisplayMessage = `Command failed: ${getErrorMessage(

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.
  */