From 60be42f0957ebf729037616fe461e09136c854aa Mon Sep 17 00:00:00 2001 From: Jerop Kipruto Date: Fri, 13 Feb 2026 11:27:20 -0500 Subject: [PATCH] refactor(core): adopt `CoreToolCallStatus` enum for type safety (#18998) --- packages/cli/src/config/extension-manager.ts | 11 +- packages/cli/src/ui/hooks/toolMapping.test.ts | 37 +-- packages/cli/src/ui/hooks/toolMapping.ts | 34 +-- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 67 +++--- .../cli/src/zed-integration/zedIntegration.ts | 20 +- .../core/src/core/coreToolScheduler.test.ts | 136 ++++++----- packages/core/src/core/coreToolScheduler.ts | 215 +++++++++++------- packages/core/src/scheduler/confirmation.ts | 8 +- packages/core/src/scheduler/scheduler.test.ts | 96 ++++---- packages/core/src/scheduler/scheduler.ts | 84 ++++--- .../core/src/scheduler/state-manager.test.ts | 134 +++++++---- packages/core/src/scheduler/state-manager.ts | 75 +++--- .../core/src/scheduler/tool-executor.test.ts | 27 +-- packages/core/src/scheduler/tool-executor.ts | 9 +- .../core/src/scheduler/tool-modifier.test.ts | 3 +- packages/core/src/scheduler/types.ts | 27 ++- .../src/services/chatRecordingService.test.ts | 11 +- .../src/telemetry/loggers.test.circular.ts | 5 +- packages/core/src/telemetry/loggers.test.ts | 27 +-- packages/core/src/telemetry/types.ts | 31 +-- packages/core/src/tools/edit.ts | 3 +- packages/core/src/utils/tool-utils.test.ts | 2 +- 22 files changed, 631 insertions(+), 431 deletions(-) diff --git a/packages/cli/src/config/extension-manager.ts b/packages/cli/src/config/extension-manager.ts index 7544231c98..4756b95b97 100644 --- a/packages/cli/src/config/extension-manager.ts +++ b/packages/cli/src/config/extension-manager.ts @@ -50,6 +50,7 @@ import { coreEvents, applyAdminAllowlist, getAdminBlockedMcpServersMessage, + CoreToolCallStatus, } from '@google/gemini-cli-core'; import { maybeRequestConsentOrFail } from './extensions/consent.js'; import { resolveEnvVarsInObject } from '../utils/envVarResolver.js'; @@ -383,7 +384,7 @@ Would you like to attempt to install via "git clone" instead?`, newExtensionConfig.version, previousExtensionConfig.version, installMetadata.type, - 'success', + CoreToolCallStatus.Success, ), ); } else { @@ -395,7 +396,7 @@ Would you like to attempt to install via "git clone" instead?`, getExtensionId(newExtensionConfig, installMetadata), newExtensionConfig.version, installMetadata.type, - 'success', + CoreToolCallStatus.Success, ), ); await this.enableExtension( @@ -433,7 +434,7 @@ Would you like to attempt to install via "git clone" instead?`, newExtensionConfig?.version ?? '', previousExtensionConfig.version, installMetadata.type, - 'error', + CoreToolCallStatus.Error, ), ); } else { @@ -445,7 +446,7 @@ Would you like to attempt to install via "git clone" instead?`, extensionId ?? '', newExtensionConfig?.version ?? '', installMetadata.type, - 'error', + CoreToolCallStatus.Error, ), ); } @@ -491,7 +492,7 @@ Would you like to attempt to install via "git clone" instead?`, extension.name, hashValue(extension.name), extension.id, - 'success', + CoreToolCallStatus.Success, ), ); } diff --git a/packages/cli/src/ui/hooks/toolMapping.test.ts b/packages/cli/src/ui/hooks/toolMapping.test.ts index 16900f3ad7..ad7b147b20 100644 --- a/packages/cli/src/ui/hooks/toolMapping.test.ts +++ b/packages/cli/src/ui/hooks/toolMapping.test.ts @@ -18,6 +18,7 @@ import { type ExecutingToolCall, type WaitingToolCall, type CancelledToolCall, + CoreToolCallStatus, } from '@google/gemini-cli-core'; import { ToolCallStatus } from '../types.js'; @@ -28,13 +29,13 @@ describe('toolMapping', () => { describe('mapCoreStatusToDisplayStatus', () => { it.each([ - ['validating', ToolCallStatus.Pending], - ['awaiting_approval', ToolCallStatus.Confirming], - ['executing', ToolCallStatus.Executing], - ['success', ToolCallStatus.Success], - ['cancelled', ToolCallStatus.Canceled], - ['error', ToolCallStatus.Error], - ['scheduled', ToolCallStatus.Pending], + [CoreToolCallStatus.Validating, ToolCallStatus.Pending], + [CoreToolCallStatus.AwaitingApproval, ToolCallStatus.Confirming], + [CoreToolCallStatus.Executing, ToolCallStatus.Executing], + [CoreToolCallStatus.Success, ToolCallStatus.Success], + [CoreToolCallStatus.Cancelled, ToolCallStatus.Canceled], + [CoreToolCallStatus.Error, ToolCallStatus.Error], + [CoreToolCallStatus.Scheduled, ToolCallStatus.Pending], ] as const)('maps %s to %s', (coreStatus, expectedDisplayStatus) => { expect(mapCoreStatusToDisplayStatus(coreStatus)).toBe( expectedDisplayStatus, @@ -77,7 +78,7 @@ describe('toolMapping', () => { it('handles a single tool call input', () => { const toolCall: ScheduledToolCall = { - status: 'scheduled', + status: CoreToolCallStatus.Scheduled, request: mockRequest, tool: mockTool, invocation: mockInvocation, @@ -91,13 +92,13 @@ describe('toolMapping', () => { it('handles an array of tool calls', () => { const toolCall1: ScheduledToolCall = { - status: 'scheduled', + status: CoreToolCallStatus.Scheduled, request: mockRequest, tool: mockTool, invocation: mockInvocation, }; const toolCall2: ScheduledToolCall = { - status: 'scheduled', + status: CoreToolCallStatus.Scheduled, request: { ...mockRequest, callId: 'call-2' }, tool: mockTool, invocation: mockInvocation, @@ -111,7 +112,7 @@ describe('toolMapping', () => { it('maps successful tool call properties correctly', () => { const toolCall: SuccessfulToolCall = { - status: 'success', + status: CoreToolCallStatus.Success, request: mockRequest, tool: mockTool, invocation: mockInvocation, @@ -139,7 +140,7 @@ describe('toolMapping', () => { it('maps executing tool call properties correctly with live output and ptyId', () => { const toolCall: ExecutingToolCall = { - status: 'executing', + status: CoreToolCallStatus.Executing, request: mockRequest, tool: mockTool, invocation: mockInvocation, @@ -166,7 +167,7 @@ describe('toolMapping', () => { }; const toolCall: WaitingToolCall = { - status: 'awaiting_approval', + status: CoreToolCallStatus.AwaitingApproval, request: mockRequest, tool: mockTool, invocation: mockInvocation, @@ -193,7 +194,7 @@ describe('toolMapping', () => { }; const toolCall: WaitingToolCall = { - status: 'awaiting_approval', + status: CoreToolCallStatus.AwaitingApproval, request: mockRequest, tool: mockTool, invocation: mockInvocation, @@ -211,7 +212,7 @@ describe('toolMapping', () => { it('maps error tool call missing tool definition', () => { // e.g. "TOOL_NOT_REGISTERED" errors const toolCall: ToolCall = { - status: 'error', + status: CoreToolCallStatus.Error, request: mockRequest, // name: 'test_tool' response: { ...mockResponse, resultDisplay: 'Tool not found' }, // notice: no `tool` or `invocation` defined here @@ -229,7 +230,7 @@ describe('toolMapping', () => { it('maps cancelled tool call properties correctly', () => { const toolCall: CancelledToolCall = { - status: 'cancelled', + status: CoreToolCallStatus.Cancelled, request: mockRequest, tool: mockTool, invocation: mockInvocation, @@ -248,7 +249,7 @@ describe('toolMapping', () => { it('propagates borderTop and borderBottom options correctly', () => { const toolCall: ScheduledToolCall = { - status: 'scheduled', + status: CoreToolCallStatus.Scheduled, request: mockRequest, tool: mockTool, invocation: mockInvocation, @@ -264,7 +265,7 @@ describe('toolMapping', () => { it('sets resultDisplay to undefined for pre-execution statuses', () => { const toolCall: ScheduledToolCall = { - status: 'scheduled', + status: CoreToolCallStatus.Scheduled, request: mockRequest, tool: mockTool, invocation: mockInvocation, diff --git a/packages/cli/src/ui/hooks/toolMapping.ts b/packages/cli/src/ui/hooks/toolMapping.ts index 00072b3d14..9de82754eb 100644 --- a/packages/cli/src/ui/hooks/toolMapping.ts +++ b/packages/cli/src/ui/hooks/toolMapping.ts @@ -10,6 +10,8 @@ import { type SerializableConfirmationDetails, type ToolResultDisplay, debugLogger, + CoreToolCallStatus, + checkExhaustive, } from '@google/gemini-cli-core'; import { ToolCallStatus, @@ -17,25 +19,23 @@ import { type IndividualToolCallDisplay, } from '../types.js'; -import { checkExhaustive } from '@google/gemini-cli-core'; - export function mapCoreStatusToDisplayStatus( coreStatus: CoreStatus, ): ToolCallStatus { switch (coreStatus) { - case 'validating': + case CoreToolCallStatus.Validating: return ToolCallStatus.Pending; - case 'awaiting_approval': + case CoreToolCallStatus.AwaitingApproval: return ToolCallStatus.Confirming; - case 'executing': + case CoreToolCallStatus.Executing: return ToolCallStatus.Executing; - case 'success': + case CoreToolCallStatus.Success: return ToolCallStatus.Success; - case 'cancelled': + case CoreToolCallStatus.Cancelled: return ToolCallStatus.Canceled; - case 'error': + case CoreToolCallStatus.Error: return ToolCallStatus.Error; - case 'scheduled': + case CoreToolCallStatus.Scheduled: return ToolCallStatus.Pending; default: return checkExhaustive(coreStatus); @@ -60,7 +60,7 @@ export function mapToDisplay( const displayName = call.tool?.displayName ?? call.request.name; - if (call.status === 'error') { + if (call.status === CoreToolCallStatus.Error) { description = JSON.stringify(call.request.args); } else { description = call.invocation.getDescription(); @@ -82,25 +82,25 @@ export function mapToDisplay( let correlationId: string | undefined = undefined; switch (call.status) { - case 'success': + case CoreToolCallStatus.Success: resultDisplay = call.response.resultDisplay; outputFile = call.response.outputFile; break; - case 'error': - case 'cancelled': + case CoreToolCallStatus.Error: + case CoreToolCallStatus.Cancelled: resultDisplay = call.response.resultDisplay; break; - case 'awaiting_approval': + case CoreToolCallStatus.AwaitingApproval: correlationId = call.correlationId; // Pass through details. Context handles dispatch (callback vs bus). confirmationDetails = call.confirmationDetails; break; - case 'executing': + case CoreToolCallStatus.Executing: resultDisplay = call.liveOutput; ptyId = call.pid; break; - case 'scheduled': - case 'validating': + case CoreToolCallStatus.Scheduled: + case CoreToolCallStatus.Validating: break; default: { const exhaustiveCheck: never = call; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 3130ee6365..5220b80fad 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -27,6 +27,7 @@ import type { AnyToolInvocation, } from '@google/gemini-cli-core'; import { + CoreToolCallStatus, ApprovalMode, AuthType, GeminiEventType as ServerGeminiEventType, @@ -343,14 +344,14 @@ describe('useGeminiStream', () => { mockCancelAllToolCalls(...args); lastToolCalls = lastToolCalls.map((tc) => { if ( - tc.status === 'awaiting_approval' || - tc.status === 'executing' || - tc.status === 'scheduled' || - tc.status === 'validating' + tc.status === CoreToolCallStatus.AwaitingApproval || + tc.status === CoreToolCallStatus.Executing || + tc.status === CoreToolCallStatus.Scheduled || + tc.status === CoreToolCallStatus.Validating ) { return { ...tc, - status: 'cancelled', + status: CoreToolCallStatus.Cancelled, response: { callId: tc.request.callId, responseParts: [], @@ -406,7 +407,8 @@ describe('useGeminiStream', () => { toolName: string, callId: string, confirmationType: 'edit' | 'info', - status: TrackedToolCall['status'] = 'awaiting_approval', + status: TrackedToolCall['status'] = CoreToolCallStatus.AwaitingApproval, + mockOnConfirm: Mock = vi.fn(), ): TrackedWaitingToolCall => ({ request: { callId, @@ -415,7 +417,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-1', }, - status: status as 'awaiting_approval', + status: status as CoreToolCallStatus.AwaitingApproval, responseSubmittedToGemini: false, confirmationDetails: confirmationType === 'edit' @@ -427,11 +429,13 @@ describe('useGeminiStream', () => { fileDiff: 'fake diff', originalContent: 'old', newContent: 'new', + onConfirm: mockOnConfirm, } : { type: 'info', title: `${toolName} confirmation`, prompt: `Execute ${toolName}?`, + onConfirm: mockOnConfirm, }, tool: { name: toolName, @@ -500,7 +504,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-1', }, - status: 'success', + status: CoreToolCallStatus.Success, responseSubmittedToGemini: false, response: { callId: 'call1', @@ -528,7 +532,7 @@ describe('useGeminiStream', () => { args: {}, prompt_id: 'prompt-id-1', }, - status: 'executing', + status: CoreToolCallStatus.Executing, responseSubmittedToGemini: false, tool: { name: 'tool2', @@ -566,7 +570,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-2', }, - status: 'success', + status: CoreToolCallStatus.Success, responseSubmittedToGemini: false, response: { callId: 'call1', @@ -588,7 +592,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-2', }, - status: 'error', + status: CoreToolCallStatus.Error, responseSubmittedToGemini: false, response: { callId: 'call2', @@ -675,10 +679,10 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-3', }, - status: 'cancelled', + status: CoreToolCallStatus.Cancelled, response: { callId: '1', - responseParts: [{ text: 'cancelled' }], + responseParts: [{ text: CoreToolCallStatus.Cancelled }], errorType: undefined, // FIX: Added missing property }, responseSubmittedToGemini: false, @@ -744,7 +748,7 @@ describe('useGeminiStream', () => { expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']); expect(client.addHistory).toHaveBeenCalledWith({ role: 'user', - parts: [{ text: 'cancelled' }], + parts: [{ text: CoreToolCallStatus.Cancelled }], }); // Ensure we do NOT call back to the API expect(mockSendMessageStream).not.toHaveBeenCalled(); @@ -761,7 +765,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-stop', }, - status: 'error', + status: CoreToolCallStatus.Error, response: { callId: 'stop-call', responseParts: [{ text: 'error occurred' }], @@ -825,7 +829,7 @@ describe('useGeminiStream', () => { invocation: { getDescription: () => `Mock description`, } as unknown as AnyToolInvocation, - status: 'cancelled', + status: CoreToolCallStatus.Cancelled, response: { callId: 'cancel-1', responseParts: [ @@ -854,7 +858,7 @@ describe('useGeminiStream', () => { invocation: { getDescription: () => `Mock description`, } as unknown as AnyToolInvocation, - status: 'cancelled', + status: CoreToolCallStatus.Cancelled, response: { callId: 'cancel-2', responseParts: [ @@ -954,7 +958,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-4', }, - status: 'executing', + status: CoreToolCallStatus.Executing, responseSubmittedToGemini: false, tool: { name: 'tool1', @@ -972,7 +976,7 @@ describe('useGeminiStream', () => { const completedToolCalls: TrackedToolCall[] = [ { ...(initialToolCalls[0] as TrackedExecutingToolCall), - status: 'success', + status: CoreToolCallStatus.Success, response: { callId: 'call1', responseParts: toolCallResponseParts, @@ -1278,7 +1282,7 @@ describe('useGeminiStream', () => { const toolCalls: TrackedToolCall[] = [ { request: { callId: 'call1', name: 'tool1', args: {} }, - status: 'executing', + status: CoreToolCallStatus.Executing, responseSubmittedToGemini: false, tool: { name: 'tool1', @@ -1318,7 +1322,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-1', }, - status: 'awaiting_approval', + status: CoreToolCallStatus.AwaitingApproval, responseSubmittedToGemini: false, tool: { name: 'some_tool', @@ -1630,7 +1634,7 @@ describe('useGeminiStream', () => { isClientInitiated: true, prompt_id: 'prompt-id-6', }, - status: 'success', + status: CoreToolCallStatus.Success, responseSubmittedToGemini: false, response: { callId: 'save-mem-call-1', @@ -1875,7 +1879,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-1', }, - status: 'awaiting_approval', + status: CoreToolCallStatus.AwaitingApproval, responseSubmittedToGemini: false, // No confirmationDetails tool: { @@ -1900,8 +1904,15 @@ describe('useGeminiStream', () => { }); it('should only process tool calls with awaiting_approval status', async () => { + const mockOnConfirmAwaiting = vi.fn().mockResolvedValue(undefined); const mixedStatusToolCalls: TrackedToolCall[] = [ - createMockToolCall('replace', 'call1', 'edit'), + createMockToolCall( + 'replace', + 'call1', + 'edit', + CoreToolCallStatus.AwaitingApproval, + mockOnConfirmAwaiting, + ), { request: { callId: 'call2', @@ -1910,7 +1921,7 @@ describe('useGeminiStream', () => { isClientInitiated: false, prompt_id: 'prompt-id-1', }, - status: 'executing', + status: CoreToolCallStatus.Executing, responseSubmittedToGemini: false, tool: { name: 'write_file', @@ -2206,7 +2217,7 @@ describe('useGeminiStream', () => { // were added to history during the await scheduleToolCalls(...) block. const tools = requests.map((r: any) => ({ request: r, - status: 'success', + status: CoreToolCallStatus.Success, tool: { displayName: r.name, name: r.name }, invocation: { getDescription: () => 'desc' }, response: { responseParts: [], resultDisplay: 'done' }, @@ -2681,7 +2692,7 @@ describe('useGeminiStream', () => { const newToolCalls: TrackedToolCall[] = [ { request: { callId: 'call1', name: 'tool1', args: {} }, - status: 'executing', + status: CoreToolCallStatus.Executing, tool: { name: 'tool1', displayName: 'tool1', @@ -2809,7 +2820,7 @@ describe('useGeminiStream', () => { await waitFor(() => { expect(mockAddItem).toHaveBeenCalledWith( expect.objectContaining({ - type: 'error', + type: CoreToolCallStatus.Error, }), expect.any(Number), ); diff --git a/packages/cli/src/zed-integration/zedIntegration.ts b/packages/cli/src/zed-integration/zedIntegration.ts index 57d8dec3a8..1d976e5de6 100644 --- a/packages/cli/src/zed-integration/zedIntegration.ts +++ b/packages/cli/src/zed-integration/zedIntegration.ts @@ -13,6 +13,7 @@ import type { ConversationRecord, } from '@google/gemini-cli-core'; import { + CoreToolCallStatus, AuthType, logToolCall, convertToFunctionResponse, @@ -451,7 +452,10 @@ export class Session { await this.sendUpdate({ sessionUpdate: 'tool_call', toolCallId: toolCall.id, - status: toolCall.status === 'success' ? 'completed' : 'failed', + status: + toolCall.status === CoreToolCallStatus.Success + ? 'completed' + : 'failed', title: toolCall.displayName || toolCall.name, content: toolCallContent, kind: tool ? toAcpToolKind(tool.kind) : 'other', @@ -477,7 +481,7 @@ export class Session { while (nextMessage !== null) { if (pendingSend.signal.aborted) { chat.addHistory(nextMessage); - return { stopReason: 'cancelled' }; + return { stopReason: CoreToolCallStatus.Cancelled }; } const functionCalls: FunctionCall[] = []; @@ -494,7 +498,7 @@ export class Session { for await (const resp of responseStream) { if (pendingSend.signal.aborted) { - return { stopReason: 'cancelled' }; + return { stopReason: CoreToolCallStatus.Cancelled }; } if ( @@ -529,7 +533,7 @@ export class Session { } if (pendingSend.signal.aborted) { - return { stopReason: 'cancelled' }; + return { stopReason: CoreToolCallStatus.Cancelled }; } } catch (error) { if (getErrorStatus(error) === 429) { @@ -543,7 +547,7 @@ export class Session { pendingSend.signal.aborted || (error instanceof Error && error.name === 'AbortError') ) { - return { stopReason: 'cancelled' }; + return { stopReason: CoreToolCallStatus.Cancelled }; } throw new acp.RequestError( @@ -663,7 +667,7 @@ export class Session { const output = await this.connection.requestPermission(params); const outcome = - output.outcome.outcome === 'cancelled' + output.outcome.outcome === CoreToolCallStatus.Cancelled ? ToolConfirmationOutcome.Cancel : z .nativeEnum(ToolConfirmationOutcome) @@ -728,7 +732,7 @@ export class Session { this.chat.recordCompletedToolCalls(this.config.getActiveModel(), [ { - status: 'success', + status: CoreToolCallStatus.Success, request: { callId, name: fc.name, @@ -773,7 +777,7 @@ export class Session { this.chat.recordCompletedToolCalls(this.config.getActiveModel(), [ { - status: 'error', + status: CoreToolCallStatus.Error, request: { callId, name: fc.name, diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 2755303c80..b8862b2b31 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -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 { 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}`, diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index d3346c9ffa..2c4b18b578 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -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 }, - 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 { 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 { 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: [ diff --git a/packages/core/src/scheduler/confirmation.ts b/packages/core/src/scheduler/confirmation.ts index 8840900bdd..5d98dd526e 100644 --- a/packages/core/src/scheduler/confirmation.ts +++ b/packages/core/src/scheduler/confirmation.ts @@ -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, }); diff --git a/packages/core/src/scheduler/scheduler.test.ts b/packages/core/src/scheduler/scheduler.test.ts index a3979f43a6..222b4993b9 100644 --- a/packages/core/src/scheduler/scheduler.test.ts +++ b/packages/core/src/scheduler/scheduler.test.ts @@ -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((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, diff --git a/packages/core/src/scheduler/scheduler.ts b/packages/core/src/scheduler/scheduler.ts index 1cd8dc3317..a8c68290f1 100644 --- a/packages/core/src/scheduler/scheduler.ts +++ b/packages/core/src/scheduler/scheduler.ts @@ -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 { - 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, + ); } } diff --git a/packages/core/src/scheduler/state-manager.test.ts b/packages/core/src/scheduler/state-manager.test.ts index d0369fdcb1..9314d70f41 100644 --- a/packages/core/src/scheduler/state-manager.test.ts +++ b/packages/core/src/scheduler/state-manager.test.ts @@ -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 diff --git a/packages/core/src/scheduler/state-manager.ts b/packages/core/src/scheduler/state-manager.ts index 21e931a18a..170873bc83 100644 --- a/packages/core/src/scheduler/state-manager.ts +++ b/packages/core/src/scheduler/state-manager.ts @@ -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, ): 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(call: T, patch: Partial): 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 | 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, diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts index c6fac5734f..53b244031d 100644 --- a/packages/core/src/scheduler/tool-executor.test.ts +++ b/packages/core/src/scheduler/tool-executor.test.ts @@ -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; 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; // 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, }), ); diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index 76b25f7c67..116598a2b9 100644 --- a/packages/core/src/scheduler/tool-executor.ts +++ b/packages/core/src/scheduler/tool-executor.ts @@ -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, diff --git a/packages/core/src/scheduler/tool-modifier.test.ts b/packages/core/src/scheduler/tool-modifier.test.ts index 0dc1a55a49..35ff2cd79c 100644 --- a/packages/core/src/scheduler/tool-modifier.test.ts +++ b/packages/core/src/scheduler/tool-modifier.test.ts @@ -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 { return { - status: 'awaiting_approval', + status: CoreToolCallStatus.AwaitingApproval, request: { callId: 'test-call-id', name: 'test-tool', diff --git a/packages/core/src/scheduler/types.ts b/packages/core/src/scheduler/types.ts index c0b6cae3d7..c63d6d3c8e 100644 --- a/packages/core/src/scheduler/types.ts +++ b/packages/core/src/scheduler/types.ts @@ -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; diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 28d458c14b..504999a671 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -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(), }, ]); diff --git a/packages/core/src/telemetry/loggers.test.circular.ts b/packages/core/src/telemetry/loggers.test.circular.ts index 6da8b31cd3..b8c365cd2c 100644 --- a/packages/core/src/telemetry/loggers.test.circular.ts +++ b/packages/core/src/telemetry/loggers.test.circular.ts @@ -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, diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts index fd2d1bc221..39b884148e 100644 --- a/packages/core/src/telemetry/loggers.test.ts +++ b/packages/core/src/telemetry/loggers.test.ts @@ -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, }, }); }); diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts index cf0e5f853f..c46255744a 100644 --- a/packages/core/src/telemetry/types.ts +++ b/packages/core/src/telemetry/types.ts @@ -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; diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts index 032e9bee74..2e79ebcb6b 100644 --- a/packages/core/src/tools/edit.ts +++ b/packages/core/src/tools/edit.ts @@ -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 { diff --git a/packages/core/src/utils/tool-utils.test.ts b/packages/core/src/utils/tool-utils.test.ts index 0a36615005..0ded23f2e3 100644 --- a/packages/core/src/utils/tool-utils.test.ts +++ b/packages/core/src/utils/tool-utils.test.ts @@ -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', () => {