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

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
@@ -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}`,
+131 -84
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: [
+6 -2
View File
@@ -17,7 +17,11 @@ import {
type ToolConfirmationPayload,
type ToolCallConfirmationDetails,
} from '../tools/tools.js';
import type { ValidatingToolCall, WaitingToolCall } from './types.js';
import {
type ValidatingToolCall,
type WaitingToolCall,
CoreToolCallStatus,
} from './types.js';
import type { Config } from '../config/config.js';
import type { SchedulerStateManager } from './state-manager.js';
import type { ToolModificationHandler } from './tool-modifier.js';
@@ -145,7 +149,7 @@ export async function resolveConfirmation(
const ideConfirmation =
'ideConfirmation' in details ? details.ideConfirmation : undefined;
state.updateStatus(callId, 'awaiting_approval', {
state.updateStatus(callId, CoreToolCallStatus.AwaitingApproval, {
confirmationDetails: serializableDetails,
correlationId,
});
+51 -45
View File
@@ -77,7 +77,7 @@ import type {
CompletedToolCall,
ToolCallResponseInfo,
} from './types.js';
import { ROOT_SCHEDULER_ID } from './types.js';
import { CoreToolCallStatus, ROOT_SCHEDULER_ID } from './types.js';
import { ToolErrorType } from '../tools/tool-error.js';
import * as ToolUtils from '../utils/tool-utils.js';
import type { EditorType } from '../utils/editor.js';
@@ -276,7 +276,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.enqueue).toHaveBeenCalledWith(
expect.arrayContaining([
expect.objectContaining({
status: 'error',
status: CoreToolCallStatus.Error,
response: expect.objectContaining({
errorType: ToolErrorType.TOOL_NOT_REGISTERED,
}),
@@ -295,7 +295,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.enqueue).toHaveBeenCalledWith(
expect.arrayContaining([
expect.objectContaining({
status: 'error',
status: CoreToolCallStatus.Error,
response: expect.objectContaining({
errorType: ToolErrorType.INVALID_TOOL_PARAMS,
}),
@@ -310,7 +310,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.enqueue).toHaveBeenCalledWith(
expect.arrayContaining([
expect.objectContaining({
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation,
@@ -325,7 +325,7 @@ describe('Scheduler (Orchestrator)', () => {
describe('Phase 2: Queue Management', () => {
it('should drain the queue if multiple calls are scheduled', async () => {
const validatingCall: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -355,7 +355,7 @@ describe('Scheduler (Orchestrator)', () => {
// Execute is the end of the loop, stub it
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall);
await scheduler.schedule(req1, signal);
@@ -382,14 +382,14 @@ describe('Scheduler (Orchestrator)', () => {
});
const validatingCall1: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
};
const validatingCall2: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req2,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -419,7 +419,9 @@ describe('Scheduler (Orchestrator)', () => {
// Yield to the event loop deterministically using queueMicrotask
await new Promise<void>((resolve) => queueMicrotask(resolve));
executionLog.push(`end-${id}`);
return { status: 'success' } as unknown as SuccessfulToolCall;
return {
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall;
});
// Action: Schedule batch of 2 tools
@@ -436,14 +438,14 @@ describe('Scheduler (Orchestrator)', () => {
it('should queue and process multiple schedule() calls made synchronously', async () => {
const validatingCall1: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
};
const validatingCall2: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req2, // Second request
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -483,7 +485,7 @@ describe('Scheduler (Orchestrator)', () => {
// Executor succeeds instantly
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall);
// ACT: Call schedule twice synchronously (without awaiting the first)
@@ -500,14 +502,14 @@ describe('Scheduler (Orchestrator)', () => {
it('should queue requests when scheduler is busy (overlapping batches)', async () => {
const validatingCall1: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
};
const validatingCall2: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req2, // Second request
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -554,13 +556,17 @@ describe('Scheduler (Orchestrator)', () => {
executionLog.push('start-batch-1');
await firstBatchPromise; // Simulating long-running tool execution
executionLog.push('end-batch-1');
return { status: 'success' } as unknown as SuccessfulToolCall;
return {
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall;
});
mockExecutor.execute.mockImplementationOnce(async () => {
executionLog.push('start-batch-2');
executionLog.push('end-batch-2');
return { status: 'success' } as unknown as SuccessfulToolCall;
return {
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall;
});
// 3. ACTIONS
@@ -608,7 +614,7 @@ describe('Scheduler (Orchestrator)', () => {
it('cancelAll() should cancel active call and clear queue', () => {
const activeCall: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -623,7 +629,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'cancelled',
CoreToolCallStatus.Cancelled,
'Operation cancelled by user',
);
// finalizeCall is handled by the processing loop, not synchronously by cancelAll
@@ -656,7 +662,7 @@ describe('Scheduler (Orchestrator)', () => {
describe('Phase 3: Policy & Confirmation Loop', () => {
const validatingCall: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -684,7 +690,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'error',
CoreToolCallStatus.Error,
expect.objectContaining({
errorType: ToolErrorType.POLICY_VIOLATION,
}),
@@ -706,7 +712,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'error',
CoreToolCallStatus.Error,
expect.objectContaining({
errorType: ToolErrorType.POLICY_VIOLATION,
responseParts: expect.arrayContaining([
@@ -731,7 +737,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'error',
CoreToolCallStatus.Error,
expect.objectContaining({
errorType: ToolErrorType.UNHANDLED_EXCEPTION,
responseParts: expect.arrayContaining([
@@ -757,7 +763,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'error',
CoreToolCallStatus.Error,
expect.objectContaining({
errorType: ToolErrorType.POLICY_VIOLATION,
responseParts: expect.arrayContaining([
@@ -786,7 +792,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'error',
CoreToolCallStatus.Error,
expect.objectContaining({
errorType: ToolErrorType.POLICY_VIOLATION,
responseParts: expect.arrayContaining([
@@ -810,7 +816,7 @@ describe('Scheduler (Orchestrator)', () => {
// Provide a mock execute to finish the loop
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall);
await scheduler.schedule(req1, signal);
@@ -827,7 +833,7 @@ describe('Scheduler (Orchestrator)', () => {
// Triggered execution
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'executing',
CoreToolCallStatus.Executing,
);
expect(mockExecutor.execute).toHaveBeenCalled();
});
@@ -835,13 +841,13 @@ describe('Scheduler (Orchestrator)', () => {
it('should auto-approve remaining identical tools in batch after ProceedAlways', async () => {
// Setup: two identical tools
const validatingCall1: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
};
const validatingCall2: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req2,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -874,7 +880,7 @@ describe('Scheduler (Orchestrator)', () => {
});
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall);
await scheduler.schedule([req1, req2], signal);
@@ -904,7 +910,7 @@ describe('Scheduler (Orchestrator)', () => {
vi.mocked(resolveConfirmation).mockResolvedValue(resolution);
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
} as unknown as SuccessfulToolCall);
await scheduler.schedule(req1, signal);
@@ -949,7 +955,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'cancelled',
CoreToolCallStatus.Cancelled,
'User denied execution.',
);
expect(mockStateManager.cancelAllQueued).toHaveBeenCalledWith(
@@ -979,7 +985,7 @@ describe('Scheduler (Orchestrator)', () => {
// Because the signal is aborted, the catch block should convert the error to a cancellation
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'cancelled',
CoreToolCallStatus.Cancelled,
'Operation cancelled',
);
});
@@ -1010,7 +1016,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'cancelled',
CoreToolCallStatus.Cancelled,
'User denied execution.',
);
// We assume the state manager stores these details.
@@ -1021,7 +1027,7 @@ describe('Scheduler (Orchestrator)', () => {
describe('Phase 4: Execution Outcomes', () => {
const validatingCall: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -1047,7 +1053,7 @@ describe('Scheduler (Orchestrator)', () => {
} as unknown as ToolCallResponseInfo;
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
response: mockResponse,
} as unknown as SuccessfulToolCall);
@@ -1055,14 +1061,14 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'success',
CoreToolCallStatus.Success,
mockResponse,
);
});
it('should update state to cancelled when executor returns cancelled status', async () => {
mockExecutor.execute.mockResolvedValue({
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
response: { callId: 'call-1', responseParts: [] },
} as unknown as CancelledToolCall);
@@ -1070,7 +1076,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'cancelled',
CoreToolCallStatus.Cancelled,
'Operation cancelled',
);
});
@@ -1082,7 +1088,7 @@ describe('Scheduler (Orchestrator)', () => {
} as unknown as ToolCallResponseInfo;
mockExecutor.execute.mockResolvedValue({
status: 'error',
status: CoreToolCallStatus.Error,
response: mockResponse,
} as unknown as ErroredToolCall);
@@ -1090,7 +1096,7 @@ describe('Scheduler (Orchestrator)', () => {
expect(mockStateManager.updateStatus).toHaveBeenCalledWith(
'call-1',
'error',
CoreToolCallStatus.Error,
mockResponse,
);
});
@@ -1103,14 +1109,14 @@ describe('Scheduler (Orchestrator)', () => {
// Mock the execution so the state advances
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
response: mockResponse,
} as unknown as SuccessfulToolCall);
// Mock the state manager to return a SUCCESS state when getToolCall is
// called
const successfulCall: SuccessfulToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: req1,
response: mockResponse,
tool: mockTool,
@@ -1145,7 +1151,7 @@ describe('Scheduler (Orchestrator)', () => {
};
mockExecutor.execute.mockResolvedValue({
status: 'success',
status: CoreToolCallStatus.Success,
response,
} as unknown as SuccessfulToolCall);
@@ -1172,7 +1178,7 @@ describe('Scheduler (Orchestrator)', () => {
});
const validatingCall: ValidatingToolCall = {
status: 'validating',
status: CoreToolCallStatus.Validating,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
@@ -1200,7 +1206,7 @@ describe('Scheduler (Orchestrator)', () => {
mockExecutor.execute.mockImplementation(async () => {
capturedContext = getToolCallContext();
return {
status: 'success',
status: CoreToolCallStatus.Success,
request: req1,
tool: mockTool,
invocation: mockInvocation as unknown as AnyToolInvocation,
+58 -26
View File
@@ -19,6 +19,7 @@ import {
type ExecutingToolCall,
type ValidatingToolCall,
type ErroredToolCall,
CoreToolCallStatus,
} from './types.js';
import { ToolErrorType } from '../tools/tool-error.js';
import { PolicyDecision } from '../policy/types.js';
@@ -212,7 +213,7 @@ export class Scheduler {
if (activeCall && !this.isTerminal(activeCall.status)) {
this.state.updateStatus(
activeCall.request.callId,
'cancelled',
CoreToolCallStatus.Cancelled,
'Operation cancelled by user',
);
}
@@ -226,7 +227,11 @@ export class Scheduler {
}
private isTerminal(status: string) {
return status === 'success' || status === 'error' || status === 'cancelled';
return (
status === CoreToolCallStatus.Success ||
status === CoreToolCallStatus.Error ||
status === CoreToolCallStatus.Cancelled
);
}
// --- Phase 1: Ingestion & Resolution ---
@@ -250,10 +255,12 @@ export class Scheduler {
const tool = toolRegistry.getTool(request.name);
if (!tool) {
return this._createToolNotFoundErroredToolCall(
enrichedRequest,
toolRegistry.getAllToolNames(),
);
return {
...this._createToolNotFoundErroredToolCall(
enrichedRequest,
toolRegistry.getAllToolNames(),
),
};
}
return this._validateAndCreateToolCall(enrichedRequest, tool);
@@ -275,7 +282,7 @@ export class Scheduler {
): ErroredToolCall {
const suggestion = getToolSuggestion(request.name, toolNames);
return {
status: 'error',
status: CoreToolCallStatus.Error,
request,
response: createErrorResponse(
request,
@@ -301,7 +308,7 @@ export class Scheduler {
try {
const invocation = tool.build(request.args);
return {
status: 'validating',
status: CoreToolCallStatus.Validating,
request,
tool,
invocation,
@@ -310,7 +317,7 @@ export class Scheduler {
};
} catch (e) {
return {
status: 'error',
status: CoreToolCallStatus.Error,
request,
tool,
response: createErrorResponse(
@@ -349,8 +356,12 @@ export class Scheduler {
const next = this.state.dequeue();
if (!next) return false;
if (next.status === 'error') {
this.state.updateStatus(next.request.callId, 'error', next.response);
if (next.status === CoreToolCallStatus.Error) {
this.state.updateStatus(
next.request.callId,
CoreToolCallStatus.Error,
next.response,
);
this.state.finalizeCall(next.request.callId);
return true;
}
@@ -359,7 +370,7 @@ export class Scheduler {
const active = this.state.firstActiveCall;
if (!active) return false;
if (active.status === 'validating') {
if (active.status === CoreToolCallStatus.Validating) {
await this._processValidatingCall(active, signal);
}
@@ -379,13 +390,13 @@ export class Scheduler {
if (signal.aborted || err.name === 'AbortError') {
this.state.updateStatus(
active.request.callId,
'cancelled',
CoreToolCallStatus.Cancelled,
'Operation cancelled',
);
} else {
this.state.updateStatus(
active.request.callId,
'error',
CoreToolCallStatus.Error,
createErrorResponse(
active.request,
err,
@@ -417,7 +428,7 @@ export class Scheduler {
this.state.updateStatus(
callId,
'error',
CoreToolCallStatus.Error,
createErrorResponse(
toolCall.request,
new Error(errorMessage),
@@ -456,7 +467,11 @@ export class Scheduler {
// Handle cancellation (cascades to entire batch)
if (outcome === ToolConfirmationOutcome.Cancel) {
this.state.updateStatus(callId, 'cancelled', 'User denied execution.');
this.state.updateStatus(
callId,
CoreToolCallStatus.Cancelled,
'User denied execution.',
);
this.state.finalizeCall(callId);
this.state.cancelAllQueued('User cancelled operation');
return; // Skip execution
@@ -472,9 +487,9 @@ export class Scheduler {
* Executes the tool and records the result.
*/
private async _execute(callId: string, signal: AbortSignal): Promise<void> {
this.state.updateStatus(callId, 'scheduled');
this.state.updateStatus(callId, CoreToolCallStatus.Scheduled);
if (signal.aborted) throw new Error('Operation cancelled');
this.state.updateStatus(callId, 'executing');
this.state.updateStatus(callId, CoreToolCallStatus.Executing);
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const activeCall = this.state.firstActiveCall as ExecutingToolCall;
@@ -490,10 +505,15 @@ export class Scheduler {
call: activeCall,
signal,
outputUpdateHandler: (id, out) =>
this.state.updateStatus(id, 'executing', { liveOutput: out }),
this.state.updateStatus(id, CoreToolCallStatus.Executing, {
liveOutput: out,
}),
onUpdateToolCall: (updated) => {
if (updated.status === 'executing' && updated.pid) {
this.state.updateStatus(callId, 'executing', {
if (
updated.status === CoreToolCallStatus.Executing &&
updated.pid
) {
this.state.updateStatus(callId, CoreToolCallStatus.Executing, {
pid: updated.pid,
});
}
@@ -501,12 +521,24 @@ export class Scheduler {
}),
);
if (result.status === 'success') {
this.state.updateStatus(callId, 'success', result.response);
} else if (result.status === 'cancelled') {
this.state.updateStatus(callId, 'cancelled', 'Operation cancelled');
if (result.status === CoreToolCallStatus.Success) {
this.state.updateStatus(
callId,
CoreToolCallStatus.Success,
result.response,
);
} else if (result.status === CoreToolCallStatus.Cancelled) {
this.state.updateStatus(
callId,
CoreToolCallStatus.Cancelled,
'Operation cancelled',
);
} else {
this.state.updateStatus(callId, 'error', result.response);
this.state.updateStatus(
callId,
CoreToolCallStatus.Error,
result.response,
);
}
}
@@ -16,6 +16,7 @@ import type {
ToolCallRequestInfo,
ToolCallResponseInfo,
} from './types.js';
import { CoreToolCallStatus, ROOT_SCHEDULER_ID } from './types.js';
import {
ToolConfirmationOutcome,
type AnyDeclarativeTool,
@@ -23,7 +24,6 @@ import {
} from '../tools/tools.js';
import { MessageBusType } from '../confirmation-bus/types.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { ROOT_SCHEDULER_ID } from './types.js';
describe('SchedulerStateManager', () => {
const mockRequest: ToolCallRequestInfo = {
@@ -44,7 +44,7 @@ describe('SchedulerStateManager', () => {
} as unknown as AnyToolInvocation;
const createValidatingCall = (id = 'call-1'): ValidatingToolCall => ({
status: 'validating',
status: CoreToolCallStatus.Validating,
request: { ...mockRequest, callId: id },
tool: mockTool,
invocation: mockInvocation,
@@ -97,7 +97,7 @@ describe('SchedulerStateManager', () => {
manager.dequeue();
manager.updateStatus(
call.request.callId,
'success',
CoreToolCallStatus.Success,
createMockResponse(call.request.callId),
);
manager.finalizeCall(call.request.callId);
@@ -105,7 +105,7 @@ describe('SchedulerStateManager', () => {
expect(onTerminalCall).toHaveBeenCalledTimes(1);
expect(onTerminalCall).toHaveBeenCalledWith(
expect.objectContaining({
status: 'success',
status: CoreToolCallStatus.Success,
request: expect.objectContaining({ callId: call.request.callId }),
}),
);
@@ -125,13 +125,13 @@ describe('SchedulerStateManager', () => {
expect(onTerminalCall).toHaveBeenCalledTimes(2);
expect(onTerminalCall).toHaveBeenCalledWith(
expect.objectContaining({
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
request: expect.objectContaining({ callId: '1' }),
}),
);
expect(onTerminalCall).toHaveBeenCalledWith(
expect.objectContaining({
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
request: expect.objectContaining({ callId: '2' }),
}),
);
@@ -167,7 +167,7 @@ describe('SchedulerStateManager', () => {
stateManager.dequeue();
stateManager.updateStatus(
'completed-1',
'success',
CoreToolCallStatus.Success,
createMockResponse('completed-1'),
);
stateManager.finalizeCall('completed-1');
@@ -212,10 +212,13 @@ describe('SchedulerStateManager', () => {
stateManager.enqueue([call]);
stateManager.dequeue();
stateManager.updateStatus(call.request.callId, 'scheduled');
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Scheduled,
);
const snapshot = stateManager.getSnapshot();
expect(snapshot[0].status).toBe('scheduled');
expect(snapshot[0].status).toBe(CoreToolCallStatus.Scheduled);
expect(snapshot[0].request.callId).toBe(call.request.callId);
});
@@ -223,11 +226,19 @@ describe('SchedulerStateManager', () => {
const call = createValidatingCall();
stateManager.enqueue([call]);
stateManager.dequeue();
stateManager.updateStatus(call.request.callId, 'scheduled');
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Scheduled,
);
stateManager.updateStatus(call.request.callId, 'executing');
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Executing,
);
expect(stateManager.firstActiveCall?.status).toBe('executing');
expect(stateManager.firstActiveCall?.status).toBe(
CoreToolCallStatus.Executing,
);
});
it('should transition to success and move to completed batch', () => {
@@ -244,7 +255,11 @@ describe('SchedulerStateManager', () => {
};
vi.mocked(onUpdate).mockClear();
stateManager.updateStatus(call.request.callId, 'success', response);
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Success,
response,
);
expect(onUpdate).toHaveBeenCalledTimes(1);
vi.mocked(onUpdate).mockClear();
@@ -254,7 +269,7 @@ describe('SchedulerStateManager', () => {
expect(stateManager.isActive).toBe(false);
expect(stateManager.completedBatch).toHaveLength(1);
const completed = stateManager.completedBatch[0] as SuccessfulToolCall;
expect(completed.status).toBe('success');
expect(completed.status).toBe(CoreToolCallStatus.Success);
expect(completed.response).toEqual(response);
expect(completed.durationMs).toBeDefined();
});
@@ -272,13 +287,17 @@ describe('SchedulerStateManager', () => {
errorType: undefined,
};
stateManager.updateStatus(call.request.callId, 'error', response);
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Error,
response,
);
stateManager.finalizeCall(call.request.callId);
expect(stateManager.isActive).toBe(false);
expect(stateManager.completedBatch).toHaveLength(1);
const completed = stateManager.completedBatch[0] as ErroredToolCall;
expect(completed.status).toBe('error');
expect(completed.status).toBe(CoreToolCallStatus.Error);
expect(completed.response).toEqual(response);
});
@@ -296,12 +315,12 @@ describe('SchedulerStateManager', () => {
stateManager.updateStatus(
call.request.callId,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
details,
);
const active = stateManager.firstActiveCall as WaitingToolCall;
expect(active.status).toBe('awaiting_approval');
expect(active.status).toBe(CoreToolCallStatus.AwaitingApproval);
expect(active.confirmationDetails).toEqual(details);
});
@@ -322,12 +341,12 @@ describe('SchedulerStateManager', () => {
stateManager.updateStatus(
call.request.callId,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
eventDrivenData,
);
const active = stateManager.firstActiveCall as WaitingToolCall;
expect(active.status).toBe('awaiting_approval');
expect(active.status).toBe(CoreToolCallStatus.AwaitingApproval);
expect(active.correlationId).toBe('corr-123');
expect(active.confirmationDetails).toEqual(details);
});
@@ -350,18 +369,18 @@ describe('SchedulerStateManager', () => {
stateManager.updateStatus(
call.request.callId,
'awaiting_approval',
CoreToolCallStatus.AwaitingApproval,
details,
);
stateManager.updateStatus(
call.request.callId,
'cancelled',
CoreToolCallStatus.Cancelled,
'User said no',
);
stateManager.finalizeCall(call.request.callId);
const completed = stateManager.completedBatch[0] as CancelledToolCall;
expect(completed.status).toBe('cancelled');
expect(completed.status).toBe(CoreToolCallStatus.Cancelled);
expect(completed.response.resultDisplay).toEqual({
fileDiff: 'diff',
fileName: 'test.txt',
@@ -372,7 +391,7 @@ describe('SchedulerStateManager', () => {
});
it('should ignore status updates for non-existent callIds', () => {
stateManager.updateStatus('unknown', 'scheduled');
stateManager.updateStatus('unknown', CoreToolCallStatus.Scheduled);
expect(onUpdate).not.toHaveBeenCalled();
});
@@ -382,13 +401,16 @@ describe('SchedulerStateManager', () => {
stateManager.dequeue();
stateManager.updateStatus(
call.request.callId,
'success',
CoreToolCallStatus.Success,
createMockResponse(call.request.callId),
);
stateManager.finalizeCall(call.request.callId);
vi.mocked(onUpdate).mockClear();
stateManager.updateStatus(call.request.callId, 'scheduled');
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Scheduled,
);
expect(onUpdate).not.toHaveBeenCalled();
});
@@ -397,7 +419,10 @@ describe('SchedulerStateManager', () => {
stateManager.enqueue([call]);
stateManager.dequeue();
stateManager.updateStatus(call.request.callId, 'executing');
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Executing,
);
stateManager.finalizeCall(call.request.callId);
expect(stateManager.isActive).toBe(true);
@@ -405,7 +430,7 @@ describe('SchedulerStateManager', () => {
stateManager.updateStatus(
call.request.callId,
'success',
CoreToolCallStatus.Success,
createMockResponse(call.request.callId),
);
stateManager.finalizeCall(call.request.callId);
@@ -420,30 +445,45 @@ describe('SchedulerStateManager', () => {
stateManager.dequeue();
// Start executing
stateManager.updateStatus(call.request.callId, 'executing');
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Executing,
);
let active = stateManager.firstActiveCall as ExecutingToolCall;
expect(active.status).toBe('executing');
expect(active.status).toBe(CoreToolCallStatus.Executing);
expect(active.liveOutput).toBeUndefined();
// Update with live output
stateManager.updateStatus(call.request.callId, 'executing', {
liveOutput: 'chunk 1',
});
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Executing,
{
liveOutput: 'chunk 1',
},
);
active = stateManager.firstActiveCall as ExecutingToolCall;
expect(active.liveOutput).toBe('chunk 1');
// Update with pid (should preserve liveOutput)
stateManager.updateStatus(call.request.callId, 'executing', {
pid: 1234,
});
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Executing,
{
pid: 1234,
},
);
active = stateManager.firstActiveCall as ExecutingToolCall;
expect(active.liveOutput).toBe('chunk 1');
expect(active.pid).toBe(1234);
// Update live output again (should preserve pid)
stateManager.updateStatus(call.request.callId, 'executing', {
liveOutput: 'chunk 2',
});
stateManager.updateStatus(
call.request.callId,
CoreToolCallStatus.Executing,
{
liveOutput: 'chunk 2',
},
);
active = stateManager.firstActiveCall as ExecutingToolCall;
expect(active.liveOutput).toBe('chunk 2');
expect(active.pid).toBe(1234);
@@ -475,7 +515,7 @@ describe('SchedulerStateManager', () => {
stateManager.dequeue();
stateManager.updateStatus(
call.request.callId,
'error',
CoreToolCallStatus.Error,
createMockResponse(call.request.callId),
);
stateManager.finalizeCall(call.request.callId);
@@ -521,7 +561,9 @@ describe('SchedulerStateManager', () => {
expect(stateManager.queueLength).toBe(0);
expect(stateManager.completedBatch).toHaveLength(2);
expect(
stateManager.completedBatch.every((c) => c.status === 'cancelled'),
stateManager.completedBatch.every(
(c) => c.status === CoreToolCallStatus.Cancelled,
),
).toBe(true);
expect(onUpdate).toHaveBeenCalledTimes(1);
});
@@ -538,7 +580,7 @@ describe('SchedulerStateManager', () => {
stateManager.dequeue();
stateManager.updateStatus(
call.request.callId,
'success',
CoreToolCallStatus.Success,
createMockResponse(call.request.callId),
);
stateManager.finalizeCall(call.request.callId);
@@ -555,7 +597,7 @@ describe('SchedulerStateManager', () => {
stateManager.dequeue();
stateManager.updateStatus(
call.request.callId,
'success',
CoreToolCallStatus.Success,
createMockResponse(call.request.callId),
);
stateManager.finalizeCall(call.request.callId);
@@ -578,7 +620,11 @@ describe('SchedulerStateManager', () => {
const call1 = createValidatingCall('1');
stateManager.enqueue([call1]);
stateManager.dequeue();
stateManager.updateStatus('1', 'success', createMockResponse('1'));
stateManager.updateStatus(
'1',
CoreToolCallStatus.Success,
createMockResponse('1'),
);
stateManager.finalizeCall('1');
// 2. Active
+45 -30
View File
@@ -17,6 +17,7 @@ import type {
ExecutingToolCall,
ToolCallResponseInfo,
} from './types.js';
import { CoreToolCallStatus } from './types.js';
import { ROOT_SCHEDULER_ID } from './types.js';
import type {
ToolConfirmationOutcome,
@@ -98,17 +99,17 @@ export class SchedulerStateManager {
*/
updateStatus(
callId: string,
status: 'success',
status: CoreToolCallStatus.Success,
data: ToolCallResponseInfo,
): void;
updateStatus(
callId: string,
status: 'error',
status: CoreToolCallStatus.Error,
data: ToolCallResponseInfo,
): void;
updateStatus(
callId: string,
status: 'awaiting_approval',
status: CoreToolCallStatus.AwaitingApproval,
data:
| ToolCallConfirmationDetails
| {
@@ -116,13 +117,20 @@ export class SchedulerStateManager {
confirmationDetails: SerializableConfirmationDetails;
},
): void;
updateStatus(callId: string, status: 'cancelled', data: string): void;
updateStatus(
callId: string,
status: 'executing',
status: CoreToolCallStatus.Cancelled,
data: string,
): void;
updateStatus(
callId: string,
status: CoreToolCallStatus.Executing,
data?: Partial<ExecutingToolCall>,
): void;
updateStatus(callId: string, status: 'scheduled' | 'validating'): void;
updateStatus(
callId: string,
status: CoreToolCallStatus.Scheduled | CoreToolCallStatus.Validating,
): void;
updateStatus(callId: string, status: Status, auxiliaryData?: unknown): void {
const call = this.activeCalls.get(callId);
if (!call) return;
@@ -152,7 +160,7 @@ export class SchedulerStateManager {
newInvocation: AnyToolInvocation,
): void {
const call = this.activeCalls.get(callId);
if (!call || call.status === 'error') return;
if (!call || call.status === CoreToolCallStatus.Error) return;
this.activeCalls.set(
callId,
@@ -179,7 +187,7 @@ export class SchedulerStateManager {
while (this.queue.length > 0) {
const queuedCall = this.queue.shift()!;
if (queuedCall.status === 'error') {
if (queuedCall.status === CoreToolCallStatus.Error) {
this._completedBatch.push(queuedCall);
this.onTerminalCall?.(queuedCall);
continue;
@@ -222,7 +230,11 @@ export class SchedulerStateManager {
private isTerminalCall(call: ToolCall): call is CompletedToolCall {
const { status } = call;
return status === 'success' || status === 'error' || status === 'cancelled';
return (
status === CoreToolCallStatus.Success ||
status === CoreToolCallStatus.Error ||
status === CoreToolCallStatus.Cancelled
);
}
private transitionCall(
@@ -231,7 +243,7 @@ export class SchedulerStateManager {
auxiliaryData?: unknown,
): ToolCall {
switch (newStatus) {
case 'success': {
case CoreToolCallStatus.Success: {
if (!this.isToolCallResponseInfo(auxiliaryData)) {
throw new Error(
`Invalid data for 'success' transition (callId: ${call.request.callId})`,
@@ -239,7 +251,7 @@ export class SchedulerStateManager {
}
return this.toSuccess(call, auxiliaryData);
}
case 'error': {
case CoreToolCallStatus.Error: {
if (!this.isToolCallResponseInfo(auxiliaryData)) {
throw new Error(
`Invalid data for 'error' transition (callId: ${call.request.callId})`,
@@ -247,7 +259,7 @@ export class SchedulerStateManager {
}
return this.toError(call, auxiliaryData);
}
case 'awaiting_approval': {
case CoreToolCallStatus.AwaitingApproval: {
if (!auxiliaryData) {
throw new Error(
`Missing data for 'awaiting_approval' transition (callId: ${call.request.callId})`,
@@ -255,9 +267,9 @@ export class SchedulerStateManager {
}
return this.toAwaitingApproval(call, auxiliaryData);
}
case 'scheduled':
case CoreToolCallStatus.Scheduled:
return this.toScheduled(call);
case 'cancelled': {
case CoreToolCallStatus.Cancelled: {
if (typeof auxiliaryData !== 'string') {
throw new Error(
`Invalid reason (string) for 'cancelled' transition (callId: ${call.request.callId})`,
@@ -265,9 +277,9 @@ export class SchedulerStateManager {
}
return this.toCancelled(call, auxiliaryData);
}
case 'validating':
case CoreToolCallStatus.Validating:
return this.toValidating(call);
case 'executing': {
case CoreToolCallStatus.Executing: {
if (
auxiliaryData !== undefined &&
!this.isExecutingToolCallPatch(auxiliaryData)
@@ -327,13 +339,13 @@ export class SchedulerStateManager {
call: ToolCall,
response: ToolCallResponseInfo,
): SuccessfulToolCall {
this.validateHasToolAndInvocation(call, 'success');
this.validateHasToolAndInvocation(call, CoreToolCallStatus.Success);
const startTime = 'startTime' in call ? call.startTime : undefined;
return {
request: call.request,
tool: call.tool,
invocation: call.invocation,
status: 'success',
status: CoreToolCallStatus.Success,
response,
durationMs: startTime ? Date.now() - startTime : undefined,
outcome: call.outcome,
@@ -348,7 +360,7 @@ export class SchedulerStateManager {
const startTime = 'startTime' in call ? call.startTime : undefined;
return {
request: call.request,
status: 'error',
status: CoreToolCallStatus.Error,
tool: 'tool' in call ? call.tool : undefined,
response,
durationMs: startTime ? Date.now() - startTime : undefined,
@@ -358,7 +370,10 @@ export class SchedulerStateManager {
}
private toAwaitingApproval(call: ToolCall, data: unknown): WaitingToolCall {
this.validateHasToolAndInvocation(call, 'awaiting_approval');
this.validateHasToolAndInvocation(
call,
CoreToolCallStatus.AwaitingApproval,
);
let confirmationDetails:
| ToolCallConfirmationDetails
@@ -377,7 +392,7 @@ export class SchedulerStateManager {
return {
request: call.request,
tool: call.tool,
status: 'awaiting_approval',
status: CoreToolCallStatus.AwaitingApproval,
correlationId,
confirmationDetails,
startTime: 'startTime' in call ? call.startTime : undefined,
@@ -400,11 +415,11 @@ export class SchedulerStateManager {
}
private toScheduled(call: ToolCall): ScheduledToolCall {
this.validateHasToolAndInvocation(call, 'scheduled');
this.validateHasToolAndInvocation(call, CoreToolCallStatus.Scheduled);
return {
request: call.request,
tool: call.tool,
status: 'scheduled',
status: CoreToolCallStatus.Scheduled,
startTime: 'startTime' in call ? call.startTime : undefined,
outcome: call.outcome,
invocation: call.invocation,
@@ -413,7 +428,7 @@ export class SchedulerStateManager {
}
private toCancelled(call: ToolCall, reason: string): CancelledToolCall {
this.validateHasToolAndInvocation(call, 'cancelled');
this.validateHasToolAndInvocation(call, CoreToolCallStatus.Cancelled);
const startTime = 'startTime' in call ? call.startTime : undefined;
// TODO: Refactor this tool-specific logic into the confirmation details payload.
@@ -444,7 +459,7 @@ export class SchedulerStateManager {
request: call.request,
tool: call.tool,
invocation: call.invocation,
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
response: {
callId: call.request.callId,
responseParts: [
@@ -468,7 +483,7 @@ export class SchedulerStateManager {
}
private isWaitingToolCall(call: ToolCall): call is WaitingToolCall {
return call.status === 'awaiting_approval';
return call.status === CoreToolCallStatus.AwaitingApproval;
}
private patchCall<T extends ToolCall>(call: T, patch: Partial<T>): T {
@@ -476,11 +491,11 @@ export class SchedulerStateManager {
}
private toValidating(call: ToolCall): ValidatingToolCall {
this.validateHasToolAndInvocation(call, 'validating');
this.validateHasToolAndInvocation(call, CoreToolCallStatus.Validating);
return {
request: call.request,
tool: call.tool,
status: 'validating',
status: CoreToolCallStatus.Validating,
startTime: 'startTime' in call ? call.startTime : undefined,
outcome: call.outcome,
invocation: call.invocation,
@@ -489,7 +504,7 @@ export class SchedulerStateManager {
}
private toExecuting(call: ToolCall, data?: unknown): ExecutingToolCall {
this.validateHasToolAndInvocation(call, 'executing');
this.validateHasToolAndInvocation(call, CoreToolCallStatus.Executing);
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const execData = data as Partial<ExecutingToolCall> | undefined;
const liveOutput =
@@ -500,7 +515,7 @@ export class SchedulerStateManager {
return {
request: call.request,
tool: call.tool,
status: 'executing',
status: CoreToolCallStatus.Executing,
startTime: 'startTime' in call ? call.startTime : undefined,
outcome: call.outcome,
invocation: call.invocation,
@@ -11,6 +11,7 @@ import type { ToolResult } from '../tools/tools.js';
import { makeFakeConfig } from '../test-utils/config.js';
import { MockTool } from '../test-utils/mock-tool.js';
import type { ScheduledToolCall } from './types.js';
import { CoreToolCallStatus } from './types.js';
import type { AnyToolInvocation } from '../index.js';
import { SHELL_TOOL_NAME } from '../tools/tool-names.js';
import * as fileUtils from '../utils/fileUtils.js';
@@ -71,7 +72,7 @@ describe('ToolExecutor', () => {
} as ToolResult);
const scheduledCall: ScheduledToolCall = {
status: 'scheduled',
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-1',
name: 'testTool',
@@ -91,8 +92,8 @@ describe('ToolExecutor', () => {
onUpdateToolCall,
});
expect(result.status).toBe('success');
if (result.status === 'success') {
expect(result.status).toBe(CoreToolCallStatus.Success);
if (result.status === CoreToolCallStatus.Success) {
const response = result.response.responseParts[0]?.functionResponse
?.response as Record<string, unknown>;
expect(response).toEqual({ output: 'Tool output' });
@@ -111,7 +112,7 @@ describe('ToolExecutor', () => {
);
const scheduledCall: ScheduledToolCall = {
status: 'scheduled',
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-2',
name: 'failTool',
@@ -130,8 +131,8 @@ describe('ToolExecutor', () => {
onUpdateToolCall: vi.fn(),
});
expect(result.status).toBe('error');
if (result.status === 'error') {
expect(result.status).toBe(CoreToolCallStatus.Error);
if (result.status === CoreToolCallStatus.Error) {
expect(result.response.error?.message).toBe('Tool Failed');
}
});
@@ -151,7 +152,7 @@ describe('ToolExecutor', () => {
);
const scheduledCall: ScheduledToolCall = {
status: 'scheduled',
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-3',
name: 'slowTool',
@@ -174,7 +175,7 @@ describe('ToolExecutor', () => {
controller.abort();
const result = await promise;
expect(result.status).toBe('cancelled');
expect(result.status).toBe(CoreToolCallStatus.Cancelled);
});
it('should truncate large shell output', async () => {
@@ -193,7 +194,7 @@ describe('ToolExecutor', () => {
});
const scheduledCall: ScheduledToolCall = {
status: 'scheduled',
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-trunc',
name: SHELL_TOOL_NAME,
@@ -228,8 +229,8 @@ describe('ToolExecutor', () => {
10, // threshold (maxChars)
);
expect(result.status).toBe('success');
if (result.status === 'success') {
expect(result.status).toBe(CoreToolCallStatus.Success);
if (result.status === CoreToolCallStatus.Success) {
const response = result.response.responseParts[0]?.functionResponse
?.response as Record<string, unknown>;
// The content should be the *truncated* version returned by the mock formatTruncatedToolOutput
@@ -262,7 +263,7 @@ describe('ToolExecutor', () => {
);
const scheduledCall: ScheduledToolCall = {
status: 'scheduled',
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-pid',
name: SHELL_TOOL_NAME,
@@ -287,7 +288,7 @@ describe('ToolExecutor', () => {
// 4. Verify PID was reported
expect(onUpdateToolCall).toHaveBeenCalledWith(
expect.objectContaining({
status: 'executing',
status: CoreToolCallStatus.Executing,
pid: testPid,
}),
);
+5 -4
View File
@@ -33,6 +33,7 @@ import type {
SuccessfulToolCall,
CancelledToolCall,
} from './types.js';
import { CoreToolCallStatus } from './types.js';
export interface ToolExecutionContext {
call: ToolCall;
@@ -81,7 +82,7 @@ export class ToolExecutor {
const setPidCallback = (pid: number) => {
const executingCall: ExecutingToolCall = {
...call,
status: 'executing',
status: CoreToolCallStatus.Executing,
tool,
invocation,
pid,
@@ -170,7 +171,7 @@ export class ToolExecutor {
}
return {
status: 'cancelled',
status: CoreToolCallStatus.Cancelled,
request: call.request,
response: {
callId: call.request.callId,
@@ -256,7 +257,7 @@ export class ToolExecutor {
}
return {
status: 'success',
status: CoreToolCallStatus.Success,
request: call.request,
tool: call.tool,
response: successResponse,
@@ -281,7 +282,7 @@ export class ToolExecutor {
const startTime = 'startTime' in call ? call.startTime : undefined;
return {
status: 'error',
status: CoreToolCallStatus.Error,
request: call.request,
response,
tool: call.tool,
@@ -7,6 +7,7 @@
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { ToolModificationHandler } from './tool-modifier.js';
import type { WaitingToolCall, ToolCallRequestInfo } from './types.js';
import { CoreToolCallStatus } from './types.js';
import * as modifiableToolModule from '../tools/modifiable-tool.js';
import * as Diff from 'diff';
import { MockModifiableTool, MockTool } from '../test-utils/mock-tool.js';
@@ -37,7 +38,7 @@ function createMockWaitingToolCall(
overrides: Partial<WaitingToolCall> = {},
): WaitingToolCall {
return {
status: 'awaiting_approval',
status: CoreToolCallStatus.AwaitingApproval,
request: {
callId: 'test-call-id',
name: 'test-tool',
+20 -7
View File
@@ -18,6 +18,19 @@ import type { SerializableConfirmationDetails } from '../confirmation-bus/types.
export const ROOT_SCHEDULER_ID = 'root';
/**
* Internal core statuses for the tool call state machine.
*/
export enum CoreToolCallStatus {
Validating = 'validating',
Scheduled = 'scheduled',
Error = 'error',
Success = 'success',
Executing = 'executing',
Cancelled = 'cancelled',
AwaitingApproval = 'awaiting_approval',
}
export interface ToolCallRequestInfo {
callId: string;
name: string;
@@ -45,7 +58,7 @@ export interface ToolCallResponseInfo {
}
export type ValidatingToolCall = {
status: 'validating';
status: CoreToolCallStatus.Validating;
request: ToolCallRequestInfo;
tool: AnyDeclarativeTool;
invocation: AnyToolInvocation;
@@ -55,7 +68,7 @@ export type ValidatingToolCall = {
};
export type ScheduledToolCall = {
status: 'scheduled';
status: CoreToolCallStatus.Scheduled;
request: ToolCallRequestInfo;
tool: AnyDeclarativeTool;
invocation: AnyToolInvocation;
@@ -65,7 +78,7 @@ export type ScheduledToolCall = {
};
export type ErroredToolCall = {
status: 'error';
status: CoreToolCallStatus.Error;
request: ToolCallRequestInfo;
response: ToolCallResponseInfo;
tool?: AnyDeclarativeTool;
@@ -75,7 +88,7 @@ export type ErroredToolCall = {
};
export type SuccessfulToolCall = {
status: 'success';
status: CoreToolCallStatus.Success;
request: ToolCallRequestInfo;
tool: AnyDeclarativeTool;
response: ToolCallResponseInfo;
@@ -86,7 +99,7 @@ export type SuccessfulToolCall = {
};
export type ExecutingToolCall = {
status: 'executing';
status: CoreToolCallStatus.Executing;
request: ToolCallRequestInfo;
tool: AnyDeclarativeTool;
invocation: AnyToolInvocation;
@@ -98,7 +111,7 @@ export type ExecutingToolCall = {
};
export type CancelledToolCall = {
status: 'cancelled';
status: CoreToolCallStatus.Cancelled;
request: ToolCallRequestInfo;
response: ToolCallResponseInfo;
tool: AnyDeclarativeTool;
@@ -109,7 +122,7 @@ export type CancelledToolCall = {
};
export type WaitingToolCall = {
status: 'awaiting_approval';
status: CoreToolCallStatus.AwaitingApproval;
request: ToolCallRequestInfo;
tool: AnyDeclarativeTool;
invocation: AnyToolInvocation;
@@ -13,6 +13,7 @@ import type {
ToolCallRecord,
MessageRecord,
} from './chatRecordingService.js';
import { CoreToolCallStatus } from '../scheduler/types.js';
import type { Content, Part } from '@google/genai';
import { ChatRecordingService } from './chatRecordingService.js';
import type { Config } from '../config/config.js';
@@ -247,7 +248,7 @@ describe('ChatRecordingService', () => {
id: 'tool-1',
name: 'testTool',
args: {},
status: 'awaiting_approval',
status: CoreToolCallStatus.AwaitingApproval,
timestamp: new Date().toISOString(),
};
chatRecordingService.recordToolCalls('gemini-pro', [toolCall]);
@@ -274,7 +275,7 @@ describe('ChatRecordingService', () => {
id: 'tool-1',
name: 'testTool',
args: {},
status: 'awaiting_approval',
status: CoreToolCallStatus.AwaitingApproval,
timestamp: new Date().toISOString(),
};
chatRecordingService.recordToolCalls('gemini-pro', [toolCall]);
@@ -571,7 +572,7 @@ describe('ChatRecordingService', () => {
name: 'list_files',
args: { path: '.' },
result: originalResult,
status: 'success',
status: CoreToolCallStatus.Success,
timestamp: new Date().toISOString(),
},
]);
@@ -650,7 +651,7 @@ describe('ChatRecordingService', () => {
name: 'read_file',
args: { path: 'image.png' },
result: originalResult,
status: 'success',
status: CoreToolCallStatus.Success,
timestamp: new Date().toISOString(),
},
]);
@@ -707,7 +708,7 @@ describe('ChatRecordingService', () => {
name: 'read_file',
args: { path: 'test.txt' },
result: [],
status: 'success',
status: CoreToolCallStatus.Success,
timestamp: new Date().toISOString(),
},
]);
@@ -17,6 +17,7 @@ import {
type ToolCallRequestInfo,
type ToolCallResponseInfo,
} from '../scheduler/types.js';
import { CoreToolCallStatus } from '../scheduler/types.js';
import { MockTool } from '../test-utils/mock-tool.js';
describe('Circular Reference Handling', () => {
@@ -62,7 +63,7 @@ describe('Circular Reference Handling', () => {
const tool = new MockTool({ name: 'mock-tool' });
const mockCompletedToolCall: CompletedToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: mockRequest,
response: mockResponse,
tool,
@@ -112,7 +113,7 @@ describe('Circular Reference Handling', () => {
const tool = new MockTool({ name: 'mock-tool' });
const mockCompletedToolCall: CompletedToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: mockRequest,
response: mockResponse,
tool,
+14 -13
View File
@@ -12,6 +12,7 @@ import type {
ErroredToolCall,
} from '../index.js';
import {
CoreToolCallStatus,
AuthType,
EditTool,
GeminiClient,
@@ -1070,7 +1071,7 @@ describe('loggers', () => {
it('should log a tool call with all fields', () => {
const tool = new EditTool(mockConfig, createMockMessageBus());
const call: CompletedToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: {
name: 'test-function',
args: {
@@ -1188,7 +1189,7 @@ describe('loggers', () => {
it('should merge data from response into metadata', () => {
const call: CompletedToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: {
name: 'ask_user',
args: { questions: [] },
@@ -1234,7 +1235,7 @@ describe('loggers', () => {
it('should log a tool call with a reject decision', () => {
const call: ErroredToolCall = {
status: 'error',
status: CoreToolCallStatus.Error,
request: {
name: 'test-function',
args: {
@@ -1312,7 +1313,7 @@ describe('loggers', () => {
it('should log a tool call with a modify decision', () => {
const call: CompletedToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: {
name: 'test-function',
args: {
@@ -1392,7 +1393,7 @@ describe('loggers', () => {
it('should log a tool call without a decision', () => {
const call: CompletedToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: {
name: 'test-function',
args: {
@@ -1472,7 +1473,7 @@ describe('loggers', () => {
it('should log a failed tool call with an error', () => {
const errorMessage = 'test-error';
const call: ErroredToolCall = {
status: 'error',
status: CoreToolCallStatus.Error,
request: {
name: 'test-function',
args: {
@@ -1573,7 +1574,7 @@ describe('loggers', () => {
);
const call: CompletedToolCall = {
status: 'success',
status: CoreToolCallStatus.Success,
request: {
name: 'mock_mcp_tool',
args: { arg1: 'value1', arg2: 2 },
@@ -1890,7 +1891,7 @@ describe('loggers', () => {
'testing-id',
'0.1.0',
'git',
'success',
CoreToolCallStatus.Success,
);
await logExtensionInstallEvent(mockConfig, event);
@@ -1911,7 +1912,7 @@ describe('loggers', () => {
extension_name: 'testing',
extension_version: '0.1.0',
extension_source: 'git',
status: 'success',
status: CoreToolCallStatus.Success,
},
});
});
@@ -1943,7 +1944,7 @@ describe('loggers', () => {
'0.1.0',
'0.1.1',
'git',
'success',
CoreToolCallStatus.Success,
);
await logExtensionUpdateEvent(mockConfig, event);
@@ -1965,7 +1966,7 @@ describe('loggers', () => {
extension_version: '0.1.0',
extension_previous_version: '0.1.1',
extension_source: 'git',
status: 'success',
status: CoreToolCallStatus.Success,
},
});
});
@@ -1993,7 +1994,7 @@ describe('loggers', () => {
'testing',
'testing-hash',
'testing-id',
'success',
CoreToolCallStatus.Success,
);
await logExtensionUninstall(mockConfig, event);
@@ -2012,7 +2013,7 @@ describe('loggers', () => {
'event.timestamp': '2025-01-01T00:00:00.000Z',
interactive: false,
extension_name: 'testing',
status: 'success',
status: CoreToolCallStatus.Success,
},
});
});
+16 -15
View File
@@ -14,6 +14,7 @@ import type { Config } from '../config/config.js';
import type { ApprovalMode } from '../policy/types.js';
import type { CompletedToolCall } from '../core/coreToolScheduler.js';
import { CoreToolCallStatus } from '../scheduler/types.js';
import { DiscoveredMCPTool } from '../tools/mcp-tool.js';
import { AuthType } from '../core/contentGenerator.js';
import type { LogAttributes, LogRecord } from '@opentelemetry/api-logs';
@@ -271,7 +272,7 @@ export class ToolCallEvent implements BaseTelemetryEvent {
this.function_name = call.request.name;
this.function_args = call.request.args;
this.duration_ms = call.durationMs ?? 0;
this.success = call.status === 'success';
this.success = call.status === CoreToolCallStatus.Success;
this.decision = call.outcome
? getDecisionFromOutcome(call.outcome)
: undefined;
@@ -296,7 +297,7 @@ export class ToolCallEvent implements BaseTelemetryEvent {
);
if (
call.status === 'success' &&
call.status === CoreToolCallStatus.Success &&
typeof call.response.resultDisplay === 'object' &&
call.response.resultDisplay !== null &&
fileDiff
@@ -317,7 +318,7 @@ export class ToolCallEvent implements BaseTelemetryEvent {
}
}
if (call.status === 'success' && call.response.data) {
if (call.status === CoreToolCallStatus.Success && call.response.data) {
this.metadata = { ...this.metadata, ...call.response.data };
}
} else {
@@ -352,7 +353,7 @@ export class ToolCallEvent implements BaseTelemetryEvent {
};
if (this.error) {
attributes['error'] = this.error;
attributes[CoreToolCallStatus.Error] = this.error;
attributes['error.message'] = this.error;
if (this.error_type) {
attributes['error_type'] = this.error_type;
@@ -891,8 +892,8 @@ export function makeSlashCommandEvent({
}
export enum SlashCommandStatus {
SUCCESS = 'success',
ERROR = 'error',
SUCCESS = CoreToolCallStatus.Success,
ERROR = CoreToolCallStatus.Error,
}
export const EVENT_REWIND = 'gemini_cli.rewind';
@@ -1294,7 +1295,7 @@ export class ExtensionInstallEvent implements BaseTelemetryEvent {
extension_id: string;
extension_version: string;
extension_source: string;
status: 'success' | 'error';
status: CoreToolCallStatus.Success | CoreToolCallStatus.Error;
constructor(
extension_name: string,
@@ -1302,7 +1303,7 @@ export class ExtensionInstallEvent implements BaseTelemetryEvent {
extension_id: string,
extension_version: string,
extension_source: string,
status: 'success' | 'error',
status: CoreToolCallStatus.Success | CoreToolCallStatus.Error,
) {
this['event.name'] = 'extension_install';
this['event.timestamp'] = new Date().toISOString();
@@ -1428,13 +1429,13 @@ export class ExtensionUninstallEvent implements BaseTelemetryEvent {
extension_name: string;
hashed_extension_name: string;
extension_id: string;
status: 'success' | 'error';
status: CoreToolCallStatus.Success | CoreToolCallStatus.Error;
constructor(
extension_name: string,
hashed_extension_name: string,
extension_id: string,
status: 'success' | 'error',
status: CoreToolCallStatus.Success | CoreToolCallStatus.Error,
) {
this['event.name'] = 'extension_uninstall';
this['event.timestamp'] = new Date().toISOString();
@@ -1469,7 +1470,7 @@ export class ExtensionUpdateEvent implements BaseTelemetryEvent {
extension_previous_version: string;
extension_version: string;
extension_source: string;
status: 'success' | 'error';
status: CoreToolCallStatus.Success | CoreToolCallStatus.Error;
constructor(
extension_name: string,
@@ -1478,7 +1479,7 @@ export class ExtensionUpdateEvent implements BaseTelemetryEvent {
extension_version: string,
extension_previous_version: string,
extension_source: string,
status: 'success' | 'error',
status: CoreToolCallStatus.Success | CoreToolCallStatus.Error,
) {
this['event.name'] = 'extension_update';
this['event.timestamp'] = new Date().toISOString();
@@ -1721,9 +1722,9 @@ export const EVENT_EDIT_CORRECTION = 'gemini_cli.edit_correction';
export class EditCorrectionEvent implements BaseTelemetryEvent {
'event.name': 'edit_correction';
'event.timestamp': string;
correction: 'success' | 'failure';
correction: CoreToolCallStatus.Success | 'failure';
constructor(correction: 'success' | 'failure') {
constructor(correction: CoreToolCallStatus.Success | 'failure') {
this['event.name'] = 'edit_correction';
this['event.timestamp'] = new Date().toISOString();
this.correction = correction;
@@ -2098,7 +2099,7 @@ export class HookCallEvent implements BaseTelemetryEvent {
if (this.error) {
// Always log errors
attributes['error'] = this.error;
attributes[CoreToolCallStatus.Error] = this.error;
}
return attributes;
+2 -1
View File
@@ -27,6 +27,7 @@ import { makeRelative, shortenPath } from '../utils/paths.js';
import { isNodeError } from '../utils/errors.js';
import type { Config } from '../config/config.js';
import { ApprovalMode } from '../policy/types.js';
import { CoreToolCallStatus } from '../scheduler/types.js';
import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js';
import {
@@ -505,7 +506,7 @@ class EditToolInvocation
};
}
const event = new EditCorrectionEvent('success');
const event = new EditCorrectionEvent(CoreToolCallStatus.Success);
logEditCorrectionEvent(this.config, event);
return {
+1 -1
View File
@@ -7,7 +7,7 @@
import { expect, describe, it } from 'vitest';
import { doesToolInvocationMatch, getToolSuggestion } from './tool-utils.js';
import type { AnyToolInvocation, Config } from '../index.js';
import { ReadFileTool } from '../tools/read-file.js';
import { ReadFileTool } from '../index.js';
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
describe('getToolSuggestion', () => {