MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-09 21:00:56 -07:00
Code Issues Packages Projects Releases Wiki Activity

refactor(core): adopt CoreToolCallStatus enum for type safety (#18998)

Browse Source
This commit is contained in:
Jerop Kipruto
2026-02-13 11:27:20 -05:00
committed by GitHub
parent d0c6a56c65
commit 60be42f095
22 changed files with 631 additions and 431 deletions
Download Patch File Download Diff File Expand all files Collapse all files

136
packages/core/src/core/coreToolScheduler.test.ts
View File

@@ -8,10 +8,11 @@ import { describe, it, expect, vi } from 'vitest';
import type { Mock } from 'vitest';
import type { CallableTool } from '@google/genai';
import { CoreToolScheduler } from './coreToolScheduler.js';
import type {
ToolCall,
WaitingToolCall,
ErroredToolCall,
import {
type ToolCall,
type WaitingToolCall,
type ErroredToolCall,
CoreToolCallStatus,
} from '../scheduler/types.js';
import type {
ToolCallConfirmationDetails,
@@ -195,7 +196,7 @@ class AbortDuringConfirmationTool extends BaseDeclarativeTool<
async function waitForStatus(
onToolCallsUpdate: Mock,
status: 'awaiting_approval' | 'executing' | 'success' | 'error' | 'cancelled',
status: CoreToolCallStatus,
timeout = 5000,
): Promise<ToolCall> {
return new Promise((resolve, reject) => {
@@ -360,7 +361,7 @@ describe('CoreToolScheduler', () => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('cancelled');
expect(completedCalls[0].status).toBe(CoreToolCallStatus.Cancelled);
});
it('should cancel all tools when cancelAll is called', async () => {
@@ -439,7 +440,7 @@ describe('CoreToolScheduler', () => {
void scheduler.schedule(requests, abortController.signal);
// Wait for the first tool to be awaiting approval
await waitForStatus(onToolCallsUpdate, 'awaiting_approval');
await waitForStatus(onToolCallsUpdate, CoreToolCallStatus.AwaitingApproval);
// Cancel all operations
scheduler.cancelAll(abortController.signal);
@@ -454,13 +455,13 @@ describe('CoreToolScheduler', () => {
expect(completedCalls).toHaveLength(3);
expect(completedCalls.find((c) => c.request.callId === '1')?.status).toBe(
'cancelled',
CoreToolCallStatus.Cancelled,
);
expect(completedCalls.find((c) => c.request.callId === '2')?.status).toBe(
'cancelled',
CoreToolCallStatus.Cancelled,
);
expect(completedCalls.find((c) => c.request.callId === '3')?.status).toBe(
'cancelled',
CoreToolCallStatus.Cancelled,
);
});
@@ -542,7 +543,7 @@ describe('CoreToolScheduler', () => {
// Wait for the first tool to be awaiting approval
const awaitingCall = (await waitForStatus(
onToolCallsUpdate,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
)) as WaitingToolCall;
// Cancel the first tool via its confirmation handler
@@ -560,13 +561,13 @@ describe('CoreToolScheduler', () => {
expect(completedCalls).toHaveLength(3);
expect(completedCalls.find((c) => c.request.callId === '1')?.status).toBe(
'cancelled',
CoreToolCallStatus.Cancelled,
);
expect(completedCalls.find((c) => c.request.callId === '2')?.status).toBe(
'cancelled',
CoreToolCallStatus.Cancelled,
);
expect(completedCalls.find((c) => c.request.callId === '3')?.status).toBe(
'cancelled',
CoreToolCallStatus.Cancelled,
);
});
@@ -621,11 +622,11 @@ describe('CoreToolScheduler', () => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('cancelled');
expect(completedCalls[0].status).toBe(CoreToolCallStatus.Cancelled);
const statuses = onToolCallsUpdate.mock.calls.flatMap((call) =>
(call[0] as ToolCall[]).map((toolCall) => toolCall.status),
);
expect(statuses).not.toContain('error');
expect(statuses).not.toContain(CoreToolCallStatus.Error);
});
it('should error when tool requires confirmation in non-interactive mode', async () => {
@@ -677,7 +678,7 @@ describe('CoreToolScheduler', () => {
expect(onAllToolCallsComplete).toHaveBeenCalled();
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('error');
expect(completedCalls[0].status).toBe(CoreToolCallStatus.Error);
const erroredCall = completedCalls[0] as ErroredToolCall;
const errorResponse = erroredCall.response;
@@ -742,7 +743,7 @@ describe('CoreToolScheduler with payload', () => {
const awaitingCall = (await waitForStatus(
onToolCallsUpdate,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
)) as WaitingToolCall;
const confirmationDetails = awaitingCall.confirmationDetails;
@@ -757,7 +758,7 @@ describe('CoreToolScheduler with payload', () => {
// After internal update, the tool should be awaiting approval again with the NEW content.
const updatedAwaitingCall = (await waitForStatus(
onToolCallsUpdate,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
)) as WaitingToolCall;
// Now confirm for real to execute.
@@ -772,7 +773,7 @@ describe('CoreToolScheduler with payload', () => {
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('success');
expect(completedCalls[0].status).toBe(CoreToolCallStatus.Success);
expect(mockTool.executeFn).toHaveBeenCalledWith({
newContent: 'final version',
});
@@ -890,7 +891,7 @@ describe('CoreToolScheduler edit cancellation', () => {
const awaitingCall = (await waitForStatus(
onToolCallsUpdate,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
)) as WaitingToolCall;
// Cancel the edit
@@ -905,7 +906,7 @@ describe('CoreToolScheduler edit cancellation', () => {
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls[0].status).toBe('cancelled');
expect(completedCalls[0].status).toBe(CoreToolCallStatus.Cancelled);
// Check that the diff is preserved
// eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -991,16 +992,16 @@ describe('CoreToolScheduler YOLO mode', () => {
// 1. The tool's execute method was called directly.
expect(executeFn).toHaveBeenCalledWith({ param: 'value' });
// 2. The tool call status never entered 'awaiting_approval'.
// 2. The tool call status never entered CoreToolCallStatus.AwaitingApproval.
const statusUpdates = onToolCallsUpdate.mock.calls
.map((call) => (call[0][0] as ToolCall)?.status)
.filter(Boolean);
expect(statusUpdates).not.toContain('awaiting_approval');
expect(statusUpdates).not.toContain(CoreToolCallStatus.AwaitingApproval);
expect(statusUpdates).toEqual([
'validating',
'scheduled',
'executing',
'success',
CoreToolCallStatus.Validating,
CoreToolCallStatus.Scheduled,
CoreToolCallStatus.Executing,
CoreToolCallStatus.Success,
]);
// 3. The final callback indicates the tool call was successful.
@@ -1008,8 +1009,8 @@ describe('CoreToolScheduler YOLO mode', () => {
.calls[0][0] as ToolCall[];
expect(completedCalls).toHaveLength(1);
const completedCall = completedCalls[0];
expect(completedCall.status).toBe('success');
if (completedCall.status === 'success') {
expect(completedCall.status).toBe(CoreToolCallStatus.Success);
if (completedCall.status === CoreToolCallStatus.Success) {
expect(completedCall.response.resultDisplay).toBe('Tool executed');
}
});
@@ -1082,8 +1083,8 @@ describe('CoreToolScheduler request queueing', () => {
// eslint-disable-next-line @typescript-eslint/no-floating-promises
scheduler.schedule([request1], abortController.signal);
// Wait for the first call to be in the 'executing' state.
await waitForStatus(onToolCallsUpdate, 'executing');
// Wait for the first call to be in the CoreToolCallStatus.Executing state.
await waitForStatus(onToolCallsUpdate, CoreToolCallStatus.Executing);
// Schedule the second call while the first is "running".
const schedulePromise2 = scheduler.schedule(
@@ -1124,8 +1125,12 @@ describe('CoreToolScheduler request queueing', () => {
});
// Verify the completion callbacks were called correctly.
expect(onAllToolCallsComplete.mock.calls[0][0][0].status).toBe('success');
expect(onAllToolCallsComplete.mock.calls[1][0][0].status).toBe('success');
expect(onAllToolCallsComplete.mock.calls[0][0][0].status).toBe(
CoreToolCallStatus.Success,
);
expect(onAllToolCallsComplete.mock.calls[1][0][0].status).toBe(
CoreToolCallStatus.Success,
);
});
it('should auto-approve a tool call if it is on the allowedTools list', async () => {
@@ -1208,16 +1213,16 @@ describe('CoreToolScheduler request queueing', () => {
// 1. The tool's execute method was called directly.
expect(executeFn).toHaveBeenCalledWith({ param: 'value' });
// 2. The tool call status never entered 'awaiting_approval'.
// 2. The tool call status never entered CoreToolCallStatus.AwaitingApproval.
const statusUpdates = onToolCallsUpdate.mock.calls
.map((call) => (call[0][0] as ToolCall)?.status)
.filter(Boolean);
expect(statusUpdates).not.toContain('awaiting_approval');
expect(statusUpdates).not.toContain(CoreToolCallStatus.AwaitingApproval);
expect(statusUpdates).toEqual([
'validating',
'scheduled',
'executing',
'success',
CoreToolCallStatus.Validating,
CoreToolCallStatus.Scheduled,
CoreToolCallStatus.Executing,
CoreToolCallStatus.Success,
]);
// 3. The final callback indicates the tool call was successful.
@@ -1226,8 +1231,8 @@ describe('CoreToolScheduler request queueing', () => {
.calls[0][0] as ToolCall[];
expect(completedCalls).toHaveLength(1);
const completedCall = completedCalls[0];
expect(completedCall.status).toBe('success');
if (completedCall.status === 'success') {
expect(completedCall.status).toBe(CoreToolCallStatus.Success);
if (completedCall.status === CoreToolCallStatus.Success) {
expect(completedCall.response.resultDisplay).toBe('Tool executed');
}
});
@@ -1310,7 +1315,7 @@ describe('CoreToolScheduler request queueing', () => {
.map((call) => (call[0][0] as ToolCall)?.status)
.filter(Boolean);
expect(statusUpdates).toContain('awaiting_approval');
expect(statusUpdates).toContain(CoreToolCallStatus.AwaitingApproval);
expect(executeFn).not.toHaveBeenCalled();
expect(onAllToolCallsComplete).not.toHaveBeenCalled();
}, 20000);
@@ -1446,7 +1451,7 @@ describe('CoreToolScheduler request queueing', () => {
onToolCallsUpdate(toolCalls);
// Capture confirmation handlers for awaiting_approval tools
toolCalls.forEach((call) => {
if (call.status === 'awaiting_approval') {
if (call.status === CoreToolCallStatus.AwaitingApproval) {
const waitingCall = call;
const details =
waitingCall.confirmationDetails as ToolCallConfirmationDetails;
@@ -1498,11 +1503,11 @@ describe('CoreToolScheduler request queueing', () => {
const calls = onToolCallsUpdate.mock.calls.at(-1)?.[0] as ToolCall[];
// With the sequential scheduler, the update includes the active call and the queue.
expect(calls?.length).toBe(3);
expect(calls?.[0].status).toBe('awaiting_approval');
expect(calls?.[0].status).toBe(CoreToolCallStatus.AwaitingApproval);
expect(calls?.[0].request.callId).toBe('1');
// Check that the other two are in the queue (still in 'validating' state)
expect(calls?.[1].status).toBe('validating');
expect(calls?.[2].status).toBe('validating');
// Check that the other two are in the queue (still in CoreToolCallStatus.Validating state)
expect(calls?.[1].status).toBe(CoreToolCallStatus.Validating);
expect(calls?.[2].status).toBe(CoreToolCallStatus.Validating);
});
expect(pendingConfirmations.length).toBe(1);
@@ -1520,9 +1525,11 @@ describe('CoreToolScheduler request queueing', () => {
-1,
)?.[0] as ToolCall[];
expect(completedCalls?.length).toBe(3);
expect(completedCalls?.every((call) => call.status === 'success')).toBe(
true,
);
expect(
completedCalls?.every(
(call) => call.status === CoreToolCallStatus.Success,
),
).toBe(true);
// Verify approval mode was changed
expect(approvalMode).toBe(ApprovalMode.AUTO_EDIT);
@@ -1631,8 +1638,8 @@ describe('CoreToolScheduler Sequential Execution', () => {
const completedCalls = onAllToolCallsComplete.mock
.calls[0][0] as ToolCall[];
expect(completedCalls).toHaveLength(2);
expect(completedCalls[0].status).toBe('success');
expect(completedCalls[1].status).toBe('success');
expect(completedCalls[0].status).toBe(CoreToolCallStatus.Success);
expect(completedCalls[1].status).toBe(CoreToolCallStatus.Success);
});
it('should cancel subsequent tools when the signal is aborted.', async () => {
@@ -1754,9 +1761,9 @@ describe('CoreToolScheduler Sequential Execution', () => {
const call2 = completedCalls.find((c) => c.request.callId === '2');
const call3 = completedCalls.find((c) => c.request.callId === '3');
expect(call1?.status).toBe('success');
expect(call2?.status).toBe('cancelled');
expect(call3?.status).toBe('cancelled');
expect(call1?.status).toBe(CoreToolCallStatus.Success);
expect(call2?.status).toBe(CoreToolCallStatus.Cancelled);
expect(call3?.status).toBe(CoreToolCallStatus.Cancelled);
});
it('should pass confirmation diff data into modifyWithEditor overrides', async () => {
@@ -1819,7 +1826,7 @@ describe('CoreToolScheduler Sequential Execution', () => {
const toolCall = (scheduler as unknown as { toolCalls: ToolCall[] })
.toolCalls[0] as WaitingToolCall;
expect(toolCall.status).toBe('awaiting_approval');
expect(toolCall.status).toBe(CoreToolCallStatus.AwaitingApproval);
const confirmationSignal = new AbortController().signal;
await scheduler.handleConfirmationResponse(
@@ -1868,7 +1875,7 @@ describe('CoreToolScheduler Sequential Execution', () => {
// Manually inject a waiting tool call
const callId = 'call-1';
const toolCall: WaitingToolCall = {
status: 'awaiting_approval',
status: CoreToolCallStatus.AwaitingApproval,
request: {
callId,
name: 'mockModifiableTool',
@@ -1987,7 +1994,9 @@ describe('CoreToolScheduler Sequential Execution', () => {
it('should not double-report completed tools when concurrent completions occur', async () => {
// Arrange
const executeFn = vi.fn().mockResolvedValue({ llmContent: 'success' });
const executeFn = vi
.fn()
.mockResolvedValue({ llmContent: CoreToolCallStatus.Success });
const mockTool = new MockTool({ name: 'mockTool', execute: executeFn });
const declarativeTool = mockTool;
@@ -2152,7 +2161,10 @@ describe('CoreToolScheduler Sequential Execution', () => {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
call[0].map((t: any) => t.status),
);
expect(allStatuses).toEqual(['success', 'success']);
expect(allStatuses).toEqual([
CoreToolCallStatus.Success,
CoreToolCallStatus.Success,
]);
expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1);
});
@@ -2201,7 +2213,7 @@ describe('CoreToolScheduler Sequential Execution', () => {
const reportedTools = onAllToolCallsComplete.mock.calls[0][0];
const result = reportedTools[0];
expect(result.status).toBe('error');
expect(result.status).toBe(CoreToolCallStatus.Error);
expect(result.response.errorType).toBe(ToolErrorType.POLICY_VIOLATION);
expect(result.response.error.message).toBe(
'Tool execution denied by policy.',
@@ -2255,7 +2267,7 @@ describe('CoreToolScheduler Sequential Execution', () => {
const reportedTools = onAllToolCallsComplete.mock.calls[0][0];
const result = reportedTools[0];
expect(result.status).toBe('error');
expect(result.status).toBe(CoreToolCallStatus.Error);
expect(result.response.errorType).toBe(ToolErrorType.POLICY_VIOLATION);
expect(result.response.error.message).toBe(
`Tool execution denied by policy. ${customDenyMessage}`,

215
packages/core/src/core/coreToolScheduler.ts
View File

@@ -42,6 +42,7 @@ import {
type ToolCallRequestInfo,
type ToolCallResponseInfo,
} from '../scheduler/types.js';
import { CoreToolCallStatus } from '../scheduler/types.js';
import { ToolExecutor } from '../scheduler/tool-executor.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { getPolicyDenialError } from '../scheduler/policy.js';
@@ -164,31 +165,34 @@ export class CoreToolScheduler {
private setStatusInternal(
targetCallId: string,
status: 'success',
status: CoreToolCallStatus.Success,
signal: AbortSignal,
response: ToolCallResponseInfo,
): void;
private setStatusInternal(
targetCallId: string,
status: 'awaiting_approval',
status: CoreToolCallStatus.AwaitingApproval,
signal: AbortSignal,
confirmationDetails: ToolCallConfirmationDetails,
): void;
private setStatusInternal(
targetCallId: string,
status: 'error',
status: CoreToolCallStatus.Error,
signal: AbortSignal,
response: ToolCallResponseInfo,
): void;
private setStatusInternal(
targetCallId: string,
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
signal: AbortSignal,
reason: string,
): void;
private setStatusInternal(
targetCallId: string,
status: 'executing' | 'scheduled' | 'validating',
status:
| CoreToolCallStatus.Executing
| CoreToolCallStatus.Scheduled
| CoreToolCallStatus.Validating,
signal: AbortSignal,
): void;
private setStatusInternal(
@@ -200,9 +204,9 @@ export class CoreToolScheduler {
this.toolCalls = this.toolCalls.map((currentCall) => {
if (
currentCall.request.callId !== targetCallId ||
currentCall.status === 'success' ||
currentCall.status === 'error' ||
currentCall.status === 'cancelled'
currentCall.status === CoreToolCallStatus.Success ||
currentCall.status === CoreToolCallStatus.Error ||
currentCall.status === CoreToolCallStatus.Cancelled
) {
return currentCall;
}
@@ -215,7 +219,7 @@ export class CoreToolScheduler {
const outcome = currentCall.outcome;
switch (newStatus) {
case 'success': {
case CoreToolCallStatus.Success: {
const durationMs = existingStartTime
? Date.now() - existingStartTime
: undefined;
@@ -223,20 +227,20 @@ export class CoreToolScheduler {
request: currentCall.request,
tool: toolInstance,
invocation,
status: 'success',
status: CoreToolCallStatus.Success,
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
response: auxiliaryData as ToolCallResponseInfo,
durationMs,
outcome,
} as SuccessfulToolCall;
}
case 'error': {
case CoreToolCallStatus.Error: {
const durationMs = existingStartTime
? Date.now() - existingStartTime
: undefined;
return {
request: currentCall.request,
status: 'error',
status: CoreToolCallStatus.Error,
tool: toolInstance,
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
response: auxiliaryData as ToolCallResponseInfo,
@@ -244,34 +248,35 @@ export class CoreToolScheduler {
outcome,
} as ErroredToolCall;
}
case 'awaiting_approval':
case CoreToolCallStatus.AwaitingApproval:
return {
request: currentCall.request,
tool: toolInstance,
status: 'awaiting_approval',
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
confirmationDetails: auxiliaryData as ToolCallConfirmationDetails,
status: CoreToolCallStatus.AwaitingApproval,
confirmationDetails:
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
auxiliaryData as ToolCallConfirmationDetails,
startTime: existingStartTime,
outcome,
invocation,
} as WaitingToolCall;
case 'scheduled':
case CoreToolCallStatus.Scheduled:
return {
request: currentCall.request,
tool: toolInstance,
status: 'scheduled',
status: CoreToolCallStatus.Scheduled,
startTime: existingStartTime,
outcome,
invocation,
} as ScheduledToolCall;
case 'cancelled': {
case CoreToolCallStatus.Cancelled: {
const durationMs = existingStartTime
? Date.now() - existingStartTime
: undefined;
// Preserve diff for cancelled edit operations
let resultDisplay: ToolResultDisplay | undefined = undefined;
if (currentCall.status === 'awaiting_approval') {
if (currentCall.status === CoreToolCallStatus.AwaitingApproval) {
const waitingCall = currentCall;
if (waitingCall.confirmationDetails.type === 'edit') {
resultDisplay = {
@@ -290,7 +295,7 @@ export class CoreToolScheduler {
request: currentCall.request,
tool: toolInstance,
invocation,
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
response: {
callId: currentCall.request.callId,
responseParts: [
@@ -313,20 +318,20 @@ export class CoreToolScheduler {
outcome,
} as CancelledToolCall;
}
case 'validating':
case CoreToolCallStatus.Validating:
return {
request: currentCall.request,
tool: toolInstance,
status: 'validating',
status: CoreToolCallStatus.Validating,
startTime: existingStartTime,
outcome,
invocation,
} as ValidatingToolCall;
case 'executing':
case CoreToolCallStatus.Executing:
return {
request: currentCall.request,
tool: toolInstance,
status: 'executing',
status: CoreToolCallStatus.Executing,
startTime: existingStartTime,
outcome,
invocation,
@@ -344,7 +349,10 @@ export class CoreToolScheduler {
this.toolCalls = this.toolCalls.map((call) => {
// We should never be asked to set args on an ErroredToolCall, but
// we guard for the case anyways.
if (call.request.callId !== targetCallId || call.status === 'error') {
if (
call.request.callId !== targetCallId ||
call.status === CoreToolCallStatus.Error
) {
return call;
}
@@ -362,7 +370,7 @@ export class CoreToolScheduler {
return {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
request: { ...call.request, args: args as Record<string, unknown> },
status: 'error',
status: CoreToolCallStatus.Error,
tool: call.tool,
response,
} as ErroredToolCall;
@@ -382,7 +390,8 @@ export class CoreToolScheduler {
this.isFinalizingToolCalls ||
this.toolCalls.some(
(call) =>
call.status === 'executing' || call.status === 'awaiting_approval',
call.status === CoreToolCallStatus.Executing ||
call.status === CoreToolCallStatus.AwaitingApproval,
)
);
}
@@ -453,14 +462,14 @@ export class CoreToolScheduler {
const activeCall = this.toolCalls[0];
// Only cancel if it's in a cancellable state.
if (
activeCall.status === 'awaiting_approval' ||
activeCall.status === 'executing' ||
activeCall.status === 'scheduled' ||
activeCall.status === 'validating'
activeCall.status === CoreToolCallStatus.AwaitingApproval ||
activeCall.status === CoreToolCallStatus.Executing ||
activeCall.status === CoreToolCallStatus.Scheduled ||
activeCall.status === CoreToolCallStatus.Validating
) {
this.setStatusInternal(
activeCall.request.callId,
'cancelled',
CoreToolCallStatus.Cancelled,
signal,
'User cancelled the operation.',
);
@@ -501,7 +510,7 @@ export class CoreToolScheduler {
);
const errorMessage = `Tool "${reqInfo.name}" not found in registry. Tools must use the exact names that are registered.${suggestion}`;
return {
status: 'error',
status: CoreToolCallStatus.Error,
request: reqInfo,
response: createErrorResponse(
reqInfo,
@@ -518,7 +527,7 @@ export class CoreToolScheduler {
);
if (invocationOrError instanceof Error) {
return {
status: 'error',
status: CoreToolCallStatus.Error,
request: reqInfo,
tool: toolInstance,
response: createErrorResponse(
@@ -531,7 +540,7 @@ export class CoreToolScheduler {
}
return {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: reqInfo,
tool: toolInstance,
invocation: invocationOrError,
@@ -568,7 +577,7 @@ export class CoreToolScheduler {
this.notifyToolCallsUpdate();
// Handle tools that were already errored during creation.
if (toolCall.status === 'error') {
if (toolCall.status === CoreToolCallStatus.Error) {
// An error during validation means this "active" tool is already complete.
// We need to check for batch completion to either finish or process the next in queue.
await this.checkAndNotifyCompletion(signal);
@@ -576,14 +585,14 @@ export class CoreToolScheduler {
}
// This logic is moved from the old `for` loop in `_schedule`.
if (toolCall.status === 'validating') {
if (toolCall.status === CoreToolCallStatus.Validating) {
const { request: reqInfo, invocation } = toolCall;
try {
if (signal.aborted) {
this.setStatusInternal(
reqInfo.callId,
'cancelled',
CoreToolCallStatus.Cancelled,
signal,
'Tool call cancelled by user.',
);
@@ -614,7 +623,7 @@ export class CoreToolScheduler {
);
this.setStatusInternal(
reqInfo.callId,
'error',
CoreToolCallStatus.Error,
signal,
createErrorResponse(reqInfo, new Error(errorMessage), errorType),
);
@@ -627,7 +636,11 @@ export class CoreToolScheduler {
reqInfo.callId,
ToolConfirmationOutcome.ProceedAlways,
);
this.setStatusInternal(reqInfo.callId, 'scheduled', signal);
this.setStatusInternal(
reqInfo.callId,
CoreToolCallStatus.Scheduled,
signal,
);
} else {
// PolicyDecision.ASK_USER
@@ -640,7 +653,11 @@ export class CoreToolScheduler {
reqInfo.callId,
ToolConfirmationOutcome.ProceedAlways,
);
this.setStatusInternal(reqInfo.callId, 'scheduled', signal);
this.setStatusInternal(
reqInfo.callId,
CoreToolCallStatus.Scheduled,
signal,
);
} else {
if (!this.config.isInteractive()) {
throw new Error(
@@ -700,7 +717,7 @@ export class CoreToolScheduler {
};
this.setStatusInternal(
reqInfo.callId,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
signal,
wrappedConfirmationDetails,
);
@@ -710,7 +727,7 @@ export class CoreToolScheduler {
if (signal.aborted) {
this.setStatusInternal(
reqInfo.callId,
'cancelled',
CoreToolCallStatus.Cancelled,
signal,
'Tool call cancelled by user.',
);
@@ -718,7 +735,7 @@ export class CoreToolScheduler {
} else {
this.setStatusInternal(
reqInfo.callId,
'error',
CoreToolCallStatus.Error,
signal,
createErrorResponse(
reqInfo,
@@ -741,10 +758,12 @@ export class CoreToolScheduler {
payload?: ToolConfirmationPayload,
): Promise<void> {
const toolCall = this.toolCalls.find(
(c) => c.request.callId === callId && c.status === 'awaiting_approval',
(c) =>
c.request.callId === callId &&
c.status === CoreToolCallStatus.AwaitingApproval,
);
if (toolCall && toolCall.status === 'awaiting_approval') {
if (toolCall && toolCall.status === CoreToolCallStatus.AwaitingApproval) {
await originalOnConfirm(outcome);
}
@@ -763,11 +782,17 @@ export class CoreToolScheduler {
return;
}
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
this.setStatusInternal(callId, 'awaiting_approval', signal, {
...waitingToolCall.confirmationDetails,
isModifying: true,
} as ToolCallConfirmationDetails);
/* eslint-disable @typescript-eslint/no-unsafe-type-assertion */
this.setStatusInternal(
callId,
CoreToolCallStatus.AwaitingApproval,
signal,
{
...waitingToolCall.confirmationDetails,
isModifying: true,
} as ToolCallConfirmationDetails,
);
/* eslint-enable @typescript-eslint/no-unsafe-type-assertion */
const result = await this.toolModifier.handleModifyWithEditor(
waitingToolCall,
@@ -778,18 +803,30 @@ export class CoreToolScheduler {
// Restore status (isModifying: false) and update diff if result exists
if (result) {
this.setArgsInternal(callId, result.updatedParams);
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
this.setStatusInternal(callId, 'awaiting_approval', signal, {
...waitingToolCall.confirmationDetails,
fileDiff: result.updatedDiff,
isModifying: false,
} as ToolCallConfirmationDetails);
/* eslint-disable @typescript-eslint/no-unsafe-type-assertion */
this.setStatusInternal(
callId,
CoreToolCallStatus.AwaitingApproval,
signal,
{
...waitingToolCall.confirmationDetails,
fileDiff: result.updatedDiff,
isModifying: false,
} as ToolCallConfirmationDetails,
);
/* eslint-enable @typescript-eslint/no-unsafe-type-assertion */
} else {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
this.setStatusInternal(callId, 'awaiting_approval', signal, {
...waitingToolCall.confirmationDetails,
isModifying: false,
} as ToolCallConfirmationDetails);
/* eslint-disable @typescript-eslint/no-unsafe-type-assertion */
this.setStatusInternal(
callId,
CoreToolCallStatus.AwaitingApproval,
signal,
{
...waitingToolCall.confirmationDetails,
isModifying: false,
} as ToolCallConfirmationDetails,
);
/* eslint-enable @typescript-eslint/no-unsafe-type-assertion */
}
} else {
// If the client provided new content, apply it and wait for
@@ -803,17 +840,22 @@ export class CoreToolScheduler {
);
if (result) {
this.setArgsInternal(callId, result.updatedParams);
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
this.setStatusInternal(callId, 'awaiting_approval', signal, {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
...(toolCall as WaitingToolCall).confirmationDetails,
fileDiff: result.updatedDiff,
} as ToolCallConfirmationDetails);
/* eslint-disable @typescript-eslint/no-unsafe-type-assertion */
this.setStatusInternal(
callId,
CoreToolCallStatus.AwaitingApproval,
signal,
{
...(toolCall as WaitingToolCall).confirmationDetails,
fileDiff: result.updatedDiff,
} as ToolCallConfirmationDetails,
);
/* eslint-enable @typescript-eslint/no-unsafe-type-assertion */
// After an inline modification, wait for another user confirmation.
return;
}
}
this.setStatusInternal(callId, 'scheduled', signal);
this.setStatusInternal(callId, CoreToolCallStatus.Scheduled, signal);
}
await this.attemptExecutionOfScheduledCalls(signal);
}
@@ -823,21 +865,25 @@ export class CoreToolScheduler {
): Promise<void> {
const allCallsFinalOrScheduled = this.toolCalls.every(
(call) =>
call.status === 'scheduled' ||
call.status === 'cancelled' ||
call.status === 'success' ||
call.status === 'error',
call.status === CoreToolCallStatus.Scheduled ||
call.status === CoreToolCallStatus.Cancelled ||
call.status === CoreToolCallStatus.Success ||
call.status === CoreToolCallStatus.Error,
);
if (allCallsFinalOrScheduled) {
const callsToExecute = this.toolCalls.filter(
(call) => call.status === 'scheduled',
(call) => call.status === CoreToolCallStatus.Scheduled,
);
for (const toolCall of callsToExecute) {
if (toolCall.status !== 'scheduled') continue;
if (toolCall.status !== CoreToolCallStatus.Scheduled) continue;
this.setStatusInternal(toolCall.request.callId, 'executing', signal);
this.setStatusInternal(
toolCall.request.callId,
CoreToolCallStatus.Executing,
signal,
);
const executingCall = this.toolCalls.find(
(c) => c.request.callId === toolCall.request.callId,
);
@@ -855,7 +901,8 @@ export class CoreToolScheduler {
this.outputUpdateHandler(callId, output);
}
this.toolCalls = this.toolCalls.map((tc) =>
tc.request.callId === callId && tc.status === 'executing'
tc.request.callId === callId &&
tc.status === CoreToolCallStatus.Executing
? { ...tc, liveOutput: output }
: tc,
);
@@ -893,11 +940,11 @@ export class CoreToolScheduler {
} else {
const activeCall = this.toolCalls[0];
const isTerminal =
activeCall.status === 'success' ||
activeCall.status === 'error' ||
activeCall.status === 'cancelled';
activeCall.status === CoreToolCallStatus.Success ||
activeCall.status === CoreToolCallStatus.Error ||
activeCall.status === CoreToolCallStatus.Cancelled;
// If the active tool is not in a terminal state (e.g., it's 'executing' or 'awaiting_approval'),
// If the active tool is not in a terminal state (e.g., it's CoreToolCallStatus.Executing or CoreToolCallStatus.AwaitingApproval),
// then the scheduler is still busy or paused. We should not proceed.
if (!isTerminal) {
return;
@@ -967,7 +1014,7 @@ export class CoreToolScheduler {
while (this.toolCallQueue.length > 0) {
const queuedCall = this.toolCallQueue.shift()!;
// Don't cancel tools that already errored during validation.
if (queuedCall.status === 'error') {
if (queuedCall.status === CoreToolCallStatus.Error) {
this.completedToolCallsForBatch.push(queuedCall);
continue;
}
@@ -981,7 +1028,7 @@ export class CoreToolScheduler {
request: queuedCall.request,
tool: queuedCall.tool,
invocation: queuedCall.invocation,
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
response: {
callId: queuedCall.request.callId,
responseParts: [