diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index f8d1b260fd..44c200a852 100644 --- 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) { diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index b2c1adade0..414ceba186 100644 --- 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: [] }, ); }); diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 58e4586887..22746b1d48 100644 --- 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( diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts index b14b492e4b..fcf9194c5e 100644 --- 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 { @@ -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, diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts index d5f92806f5..e1a2b091fa 100644 --- 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; + 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; + expect(response['output']).toBe('TruncatedContent...'); + expect(result.response.outputFile).toBe('/tmp/truncated_output.txt'); + } + }); }); diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index e358c53c8b..6edea96742 100644 --- 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 { - 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 { + 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 { + 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, diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 6afded3faa..f6a71eef0f 100644 --- 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( diff --git a/packages/core/src/utils/tool-utils.ts b/packages/core/src/utils/tool-utils.ts index 17ccbda8d6..44c72c7105 100644 --- 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. */