From 919e5da58187530c2247c105f64575fda79e3d67 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Mon, 23 Mar 2026 17:21:14 -0400 Subject: [PATCH] refactor(core): delete obsolete coreToolScheduler (#23502) --- .../core/src/code_assist/telemetry.test.ts | 6 +- .../core/src/core/coreToolScheduler.test.ts | 2451 ----------------- packages/core/src/core/coreToolScheduler.ts | 1164 -------- .../src/core/coreToolSchedulerHooks.test.ts | 313 --- packages/core/src/scheduler/policy.test.ts | 79 +- .../clearcut-logger/clearcut-logger.test.ts | 2 +- .../core/src/telemetry/uiTelemetry.test.ts | 2 +- 7 files changed, 30 insertions(+), 3987 deletions(-) delete mode 100644 packages/core/src/core/coreToolScheduler.test.ts delete mode 100644 packages/core/src/core/coreToolScheduler.ts delete mode 100644 packages/core/src/core/coreToolSchedulerHooks.test.ts diff --git a/packages/core/src/code_assist/telemetry.test.ts b/packages/core/src/code_assist/telemetry.test.ts index 66f1e631eb..f1404ecfb0 100644 --- a/packages/core/src/code_assist/telemetry.test.ts +++ b/packages/core/src/code_assist/telemetry.test.ts @@ -24,14 +24,16 @@ import { } from '@google/genai'; import * as codeAssist from './codeAssist.js'; import type { CodeAssistServer } from './server.js'; -import type { CompletedToolCall } from '../core/coreToolScheduler.js'; +import type { + CompletedToolCall, + ToolCallResponseInfo, +} from '../scheduler/types.js'; import { ToolConfirmationOutcome, type AnyDeclarativeTool, type AnyToolInvocation, } from '../tools/tools.js'; import type { Config } from '../config/config.js'; -import type { ToolCallResponseInfo } from '../scheduler/types.js'; function createMockResponse( candidates: GenerateContentResponse['candidates'] = [], diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts deleted file mode 100644 index 28350fef10..0000000000 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ /dev/null @@ -1,2451 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, it, expect, vi, type Mock } from 'vitest'; -import type { CallableTool } from '@google/genai'; -import { CoreToolScheduler } from './coreToolScheduler.js'; -import { - type ToolCall, - type WaitingToolCall, - type ErroredToolCall, - CoreToolCallStatus, -} from '../scheduler/types.js'; -import { - type ToolCallConfirmationDetails, - type ToolConfirmationPayload, - type ToolInvocation, - type ToolResult, - type Config, - type ToolRegistry, - type MessageBus, - DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, - BaseDeclarativeTool, - BaseToolInvocation, - ToolConfirmationOutcome, - Kind, - ApprovalMode, - HookSystem, - PolicyDecision, - ToolErrorType, - DiscoveredMCPTool, - GeminiCliOperation, -} from '../index.js'; -import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; -import { NoopSandboxManager } from '../services/sandboxManager.js'; -import { - MockModifiableTool, - MockTool, - MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, -} from '../test-utils/mock-tool.js'; -import * as modifiableToolModule from '../tools/modifiable-tool.js'; -import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; -import type { PolicyEngine } from '../policy/policy-engine.js'; -import { runInDevTraceSpan, type SpanMetadata } from '../telemetry/trace.js'; - -vi.mock('fs/promises', () => ({ - writeFile: vi.fn(), -})); - -vi.mock('../telemetry/trace.js', () => ({ - runInDevTraceSpan: vi.fn(async (opts, fn) => { - const metadata = { attributes: opts.attributes || {} }; - return fn({ - metadata, - endSpan: vi.fn(), - }); - }), -})); - -class TestApprovalTool extends BaseDeclarativeTool<{ id: string }, ToolResult> { - static readonly Name = 'testApprovalTool'; - - constructor( - private config: Config, - messageBus: MessageBus, - ) { - super( - TestApprovalTool.Name, - 'TestApprovalTool', - 'A tool for testing approval logic', - Kind.Edit, - { - properties: { id: { type: 'string' } }, - required: ['id'], - type: 'object', - }, - messageBus, - ); - } - - protected createInvocation( - params: { id: string }, - messageBus: MessageBus, - _toolName?: string, - _toolDisplayName?: string, - ): ToolInvocation<{ id: string }, ToolResult> { - return new TestApprovalInvocation(this.config, params, messageBus); - } -} - -class TestApprovalInvocation extends BaseToolInvocation< - { id: string }, - ToolResult -> { - constructor( - private config: Config, - params: { id: string }, - messageBus: MessageBus, - ) { - super(params, messageBus); - } - - getDescription(): string { - return `Test tool ${this.params.id}`; - } - - override async shouldConfirmExecute(): Promise< - ToolCallConfirmationDetails | false - > { - // Need confirmation unless approval mode is AUTO_EDIT - if (this.config.getApprovalMode() === ApprovalMode.AUTO_EDIT) { - return false; - } - - return { - type: 'edit', - title: `Confirm Test Tool ${this.params.id}`, - fileName: `test-${this.params.id}.txt`, - filePath: `/test-${this.params.id}.txt`, - fileDiff: 'Test diff content', - originalContent: '', - newContent: 'Test content', - onConfirm: async (outcome: ToolConfirmationOutcome) => { - if (outcome === ToolConfirmationOutcome.ProceedAlways) { - this.config.setApprovalMode(ApprovalMode.AUTO_EDIT); - } - }, - }; - } - - async execute(): Promise { - return { - llmContent: `Executed test tool ${this.params.id}`, - returnDisplay: `Executed test tool ${this.params.id}`, - }; - } -} - -class AbortDuringConfirmationInvocation extends BaseToolInvocation< - Record, - ToolResult -> { - constructor( - private readonly abortController: AbortController, - private readonly abortError: Error, - params: Record, - messageBus: MessageBus, - ) { - super(params, messageBus); - } - - override async shouldConfirmExecute( - _signal: AbortSignal, - ): Promise { - this.abortController.abort(); - throw this.abortError; - } - - async execute(_abortSignal: AbortSignal): Promise { - throw new Error('execute should not be called when confirmation fails'); - } - - getDescription(): string { - return 'Abort during confirmation invocation'; - } -} - -class AbortDuringConfirmationTool extends BaseDeclarativeTool< - Record, - ToolResult -> { - constructor( - private readonly abortController: AbortController, - private readonly abortError: Error, - messageBus: MessageBus, - ) { - super( - 'abortDuringConfirmationTool', - 'Abort During Confirmation Tool', - 'A tool that aborts while confirming execution.', - Kind.Other, - { - type: 'object', - properties: {}, - }, - messageBus, - ); - } - - protected createInvocation( - params: Record, - messageBus: MessageBus, - _toolName?: string, - _toolDisplayName?: string, - ): ToolInvocation, ToolResult> { - return new AbortDuringConfirmationInvocation( - this.abortController, - this.abortError, - params, - messageBus, - ); - } -} - -async function waitForStatus( - onToolCallsUpdate: Mock, - status: CoreToolCallStatus, - timeout = 5000, -): Promise { - return new Promise((resolve, reject) => { - const startTime = Date.now(); - const check = () => { - if (Date.now() - startTime > timeout) { - const seenStatuses = onToolCallsUpdate.mock.calls - .flatMap((call) => call[0]) - .map((toolCall: ToolCall) => toolCall.status); - reject( - new Error( - `Timed out waiting for status "${status}". Seen statuses: ${seenStatuses.join( - ', ', - )}`, - ), - ); - return; - } - - const foundCall = onToolCallsUpdate.mock.calls - .flatMap((call) => call[0]) - .find((toolCall: ToolCall) => toolCall.status === status); - if (foundCall) { - resolve(foundCall); - } else { - setTimeout(check, 10); // Check again in 10ms - } - }; - check(); - }); -} - -function createMockConfig(overrides: Partial = {}): Config { - const defaultToolRegistry = { - getTool: () => undefined, - getToolByName: () => undefined, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => undefined, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - getExperiments: () => {}, - } as unknown as ToolRegistry; - - const baseConfig = { - getSessionId: () => 'test-session-id', - getUsageStatisticsEnabled: () => true, - getDebugMode: () => false, - isInteractive: () => true, - getApprovalMode: () => ApprovalMode.DEFAULT, - setApprovalMode: () => {}, - getAllowedTools: () => [], - getContentGeneratorConfig: () => ({ - model: 'test-model', - authType: 'oauth-personal', - }), - getShellExecutionConfig: () => ({ - terminalWidth: 90, - terminalHeight: 30, - sanitizationConfig: { - enableEnvironmentVariableRedaction: true, - allowedEnvironmentVariables: [], - blockedEnvironmentVariables: [], - }, - sandboxManager: new NoopSandboxManager(), - }), - storage: { - getProjectTempDir: () => '/tmp', - }, - getTruncateToolOutputThreshold: () => - DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, - getToolRegistry: () => defaultToolRegistry, - getActiveModel: () => DEFAULT_GEMINI_MODEL, - getGeminiClient: () => null, - getMessageBus: () => createMockMessageBus(), - getEnableHooks: () => false, - getHookSystem: () => undefined, - getExperiments: () => {}, - } as unknown as Config; - - // eslint-disable-next-line @typescript-eslint/no-misused-spread - const finalConfig = { ...baseConfig, ...overrides } as Config; - - (finalConfig as unknown as { config: Config }).config = finalConfig; - - // Patch the policy engine to use the final config if not overridden - if (!overrides.getPolicyEngine) { - finalConfig.getPolicyEngine = () => - ({ - check: async ( - toolCall: { name: string; args: object }, - _serverName?: string, - ) => { - // Mock simple policy logic for tests - const mode = finalConfig.getApprovalMode(); - if (mode === ApprovalMode.YOLO) { - return { decision: PolicyDecision.ALLOW }; - } - const allowed = finalConfig.getAllowedTools(); - if ( - allowed && - (allowed.includes(toolCall.name) || - allowed.some((p) => toolCall.name.startsWith(p))) - ) { - return { decision: PolicyDecision.ALLOW }; - } - return { decision: PolicyDecision.ASK_USER }; - }, - }) as unknown as PolicyEngine; - } - - Object.defineProperty(finalConfig, 'toolRegistry', { - get: () => finalConfig.getToolRegistry?.() || defaultToolRegistry, - }); - Object.defineProperty(finalConfig, 'messageBus', { - get: () => finalConfig.getMessageBus?.(), - }); - Object.defineProperty(finalConfig, 'geminiClient', { - get: () => finalConfig.getGeminiClient?.(), - }); - - return finalConfig; -} - -describe('CoreToolScheduler', () => { - it('should cancel a tool call if the signal is aborted before confirmation', async () => { - const mockTool = new MockTool({ - name: 'mockTool', - shouldConfirmExecute: MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, - }); - const declarativeTool = mockTool; - const mockToolRegistry = { - getTool: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: () => declarativeTool, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - isInteractive: () => false, - }); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }; - - abortController.abort(); - await scheduler.schedule([request], abortController.signal); - - expect(onAllToolCallsComplete).toHaveBeenCalled(); - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe(CoreToolCallStatus.Cancelled); - - expect(runInDevTraceSpan).toHaveBeenCalledWith( - expect.objectContaining({ - operation: GeminiCliOperation.ScheduleToolCalls, - }), - expect.any(Function), - ); - - const spanArgs = vi.mocked(runInDevTraceSpan).mock.calls[0]; - const fn = spanArgs[1]; - const metadata: SpanMetadata = { name: '', attributes: {} }; - await fn({ metadata, endSpan: vi.fn() }); - expect(metadata).toMatchObject({ - input: [request], - }); - }); - - it('should cancel all tools when cancelAll is called', async () => { - const mockTool1 = new MockTool({ - name: 'mockTool1', - shouldConfirmExecute: MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, - }); - const mockTool2 = new MockTool({ name: 'mockTool2' }); - const mockTool3 = new MockTool({ name: 'mockTool3' }); - - const mockToolRegistry = { - getTool: (name: string) => { - if (name === 'mockTool1') return mockTool1; - if (name === 'mockTool2') return mockTool2; - if (name === 'mockTool3') return mockTool3; - return undefined; - }, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: (name: string) => { - if (name === 'mockTool1') return mockTool1; - if (name === 'mockTool2') return mockTool2; - if (name === 'mockTool3') return mockTool3; - return undefined; - }, - getToolByDisplayName: () => undefined, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getHookSystem: () => undefined, - }); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const requests = [ - { - callId: '1', - name: 'mockTool1', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }, - { - callId: '2', - name: 'mockTool2', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }, - { - callId: '3', - name: 'mockTool3', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }, - ]; - - // Don't await, let it run in the background - void scheduler.schedule(requests, abortController.signal); - - // Wait for the first tool to be awaiting approval - await waitForStatus(onToolCallsUpdate, CoreToolCallStatus.AwaitingApproval); - - // Cancel all operations - scheduler.cancelAll(abortController.signal); - abortController.abort(); // Also fire the signal - - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - - expect(completedCalls).toHaveLength(3); - expect(completedCalls.find((c) => c.request.callId === '1')?.status).toBe( - CoreToolCallStatus.Cancelled, - ); - expect(completedCalls.find((c) => c.request.callId === '2')?.status).toBe( - CoreToolCallStatus.Cancelled, - ); - expect(completedCalls.find((c) => c.request.callId === '3')?.status).toBe( - CoreToolCallStatus.Cancelled, - ); - }); - - it('should cancel all tools in a batch when one is cancelled via confirmation', async () => { - const mockTool1 = new MockTool({ - name: 'mockTool1', - shouldConfirmExecute: MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, - }); - const mockTool2 = new MockTool({ name: 'mockTool2' }); - const mockTool3 = new MockTool({ name: 'mockTool3' }); - - const mockToolRegistry = { - getTool: (name: string) => { - if (name === 'mockTool1') return mockTool1; - if (name === 'mockTool2') return mockTool2; - if (name === 'mockTool3') return mockTool3; - return undefined; - }, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: (name: string) => { - if (name === 'mockTool1') return mockTool1; - if (name === 'mockTool2') return mockTool2; - if (name === 'mockTool3') return mockTool3; - return undefined; - }, - getToolByDisplayName: () => undefined, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getHookSystem: () => undefined, - }); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const requests = [ - { - callId: '1', - name: 'mockTool1', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }, - { - callId: '2', - name: 'mockTool2', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }, - { - callId: '3', - name: 'mockTool3', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }, - ]; - - // Don't await, let it run in the background - void scheduler.schedule(requests, abortController.signal); - - // Wait for the first tool to be awaiting approval - const awaitingCall = (await waitForStatus( - onToolCallsUpdate, - CoreToolCallStatus.AwaitingApproval, - )) as WaitingToolCall; - - // Cancel the first tool via its confirmation handler - const confirmationDetails = - awaitingCall.confirmationDetails as ToolCallConfirmationDetails; - await confirmationDetails.onConfirm(ToolConfirmationOutcome.Cancel); - abortController.abort(); // User cancelling often involves an abort signal - - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - - expect(completedCalls).toHaveLength(3); - expect(completedCalls.find((c) => c.request.callId === '1')?.status).toBe( - CoreToolCallStatus.Cancelled, - ); - expect(completedCalls.find((c) => c.request.callId === '2')?.status).toBe( - CoreToolCallStatus.Cancelled, - ); - expect(completedCalls.find((c) => c.request.callId === '3')?.status).toBe( - CoreToolCallStatus.Cancelled, - ); - }); - - it('should mark tool call as cancelled when abort happens during confirmation error', async () => { - const abortController = new AbortController(); - const abortError = new Error('Abort requested during confirmation'); - const declarativeTool = new AbortDuringConfirmationTool( - abortController, - abortError, - createMockMessageBus(), - ); - - const mockToolRegistry = { - getTool: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: () => declarativeTool, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - isInteractive: () => true, - }); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const request = { - callId: 'abort-1', - name: 'abortDuringConfirmationTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-abort', - }; - - await scheduler.schedule([request], abortController.signal); - - expect(onAllToolCallsComplete).toHaveBeenCalled(); - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - 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(CoreToolCallStatus.Error); - }); - - it('should error when tool requires confirmation in non-interactive mode', async () => { - const mockTool = new MockTool({ - name: 'mockTool', - shouldConfirmExecute: MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, - }); - const declarativeTool = mockTool; - const mockToolRegistry = { - getTool: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: () => declarativeTool, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - isInteractive: () => false, - }); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }; - - await scheduler.schedule([request], abortController.signal); - - expect(onAllToolCallsComplete).toHaveBeenCalled(); - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe(CoreToolCallStatus.Error); - - const erroredCall = completedCalls[0] as ErroredToolCall; - const errorResponse = erroredCall.response; - const errorParts = errorResponse.responseParts; - // @ts-expect-error - accessing internal structure of FunctionResponsePart - const errorMessage = errorParts[0].functionResponse.response.error; - expect(errorMessage).toContain( - 'Tool execution for "mockTool" requires user confirmation, which is not supported in non-interactive mode.', - ); - }); -}); - -describe('CoreToolScheduler with payload', () => { - it('should update args and diff and execute tool when payload is provided', async () => { - const mockTool = new MockModifiableTool(); - mockTool.executeFn = vi.fn(); - const declarativeTool = mockTool; - const mockToolRegistry = { - getTool: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: () => declarativeTool, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockModifiableTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-2', - }; - - await scheduler.schedule([request], abortController.signal); - - const awaitingCall = (await waitForStatus( - onToolCallsUpdate, - CoreToolCallStatus.AwaitingApproval, - )) as WaitingToolCall; - const confirmationDetails = awaitingCall.confirmationDetails; - - if (confirmationDetails) { - const payload: ToolConfirmationPayload = { newContent: 'final version' }; - await (confirmationDetails as ToolCallConfirmationDetails).onConfirm( - ToolConfirmationOutcome.ProceedOnce, - payload, - ); - } - - // After internal update, the tool should be awaiting approval again with the NEW content. - const updatedAwaitingCall = (await waitForStatus( - onToolCallsUpdate, - CoreToolCallStatus.AwaitingApproval, - )) as WaitingToolCall; - - // Now confirm for real to execute. - await ( - updatedAwaitingCall.confirmationDetails as ToolCallConfirmationDetails - ).onConfirm(ToolConfirmationOutcome.ProceedOnce); - - // Wait for the tool execution to complete - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe(CoreToolCallStatus.Success); - expect(mockTool.executeFn).toHaveBeenCalledWith({ - newContent: 'final version', - }); - }); -}); - -class MockEditToolInvocation extends BaseToolInvocation< - Record, - ToolResult -> { - constructor(params: Record, messageBus: MessageBus) { - super(params, messageBus); - } - - getDescription(): string { - return 'A mock edit tool invocation'; - } - - override async shouldConfirmExecute( - _abortSignal: AbortSignal, - ): Promise { - return { - type: 'edit', - title: 'Confirm Edit', - fileName: 'test.txt', - filePath: 'test.txt', - fileDiff: - '--- test.txt\n+++ test.txt\n@@ -1,1 +1,1 @@\n-old content\n+new content', - originalContent: 'old content', - newContent: 'new content', - onConfirm: async () => {}, - }; - } - - async execute(_abortSignal: AbortSignal): Promise { - return { - llmContent: 'Edited successfully', - returnDisplay: 'Edited successfully', - }; - } -} - -class MockEditTool extends BaseDeclarativeTool< - Record, - ToolResult -> { - constructor(messageBus: MessageBus) { - super( - 'mockEditTool', - 'mockEditTool', - 'A mock edit tool', - Kind.Edit, - {}, - messageBus, - ); - } - - protected createInvocation( - params: Record, - messageBus: MessageBus, - _toolName?: string, - _toolDisplayName?: string, - ): ToolInvocation, ToolResult> { - return new MockEditToolInvocation(params, messageBus); - } -} - -describe('CoreToolScheduler edit cancellation', () => { - it('should preserve diff when an edit is cancelled', async () => { - const mockEditTool = new MockEditTool(createMockMessageBus()); - const mockToolRegistry = { - getTool: () => mockEditTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: () => mockEditTool, - getToolByDisplayName: () => mockEditTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockEditTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }; - - await scheduler.schedule([request], abortController.signal); - - const awaitingCall = (await waitForStatus( - onToolCallsUpdate, - CoreToolCallStatus.AwaitingApproval, - )) as WaitingToolCall; - - // Cancel the edit - const confirmationDetails = awaitingCall.confirmationDetails; - if (confirmationDetails) { - await (confirmationDetails as ToolCallConfirmationDetails).onConfirm( - ToolConfirmationOutcome.Cancel, - ); - } - - expect(onAllToolCallsComplete).toHaveBeenCalled(); - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - - expect(completedCalls[0].status).toBe(CoreToolCallStatus.Cancelled); - - // Check that the diff is preserved - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const cancelledCall = completedCalls[0] as any; - expect(cancelledCall.response.resultDisplay).toBeDefined(); - expect(cancelledCall.response.resultDisplay.fileDiff).toBe( - '--- test.txt\n+++ test.txt\n@@ -1,1 +1,1 @@\n-old content\n+new content', - ); - expect(cancelledCall.response.resultDisplay.fileName).toBe('test.txt'); - }); -}); - -describe('CoreToolScheduler YOLO mode', () => { - it('should execute tool requiring confirmation directly without waiting', async () => { - // Arrange - const executeFn = vi.fn().mockResolvedValue({ - llmContent: 'Tool executed', - returnDisplay: 'Tool executed', - }); - const mockTool = new MockTool({ - name: 'mockTool', - execute: executeFn, - shouldConfirmExecute: MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, - }); - const declarativeTool = mockTool; - - const mockToolRegistry = { - getTool: () => declarativeTool, - getToolByName: () => declarativeTool, - // Other properties are not needed for this test but are included for type consistency. - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - // Configure the scheduler for YOLO mode. - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.YOLO, - isInteractive: () => false, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockTool', - args: { param: 'value' }, - isClientInitiated: false, - prompt_id: 'prompt-id-yolo', - }; - - // Act - await scheduler.schedule([request], abortController.signal); - - // Wait for the tool execution to complete - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - // Assert - // 1. The tool's execute method was called directly. - expect(executeFn).toHaveBeenCalledWith( - { param: 'value' }, - expect.anything(), - undefined, - expect.anything(), - ); - - // 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(CoreToolCallStatus.AwaitingApproval); - expect(statusUpdates).toEqual([ - CoreToolCallStatus.Validating, - CoreToolCallStatus.Scheduled, - CoreToolCallStatus.Executing, - CoreToolCallStatus.Success, - ]); - - // 3. The final callback indicates the tool call was successful. - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls).toHaveLength(1); - const completedCall = completedCalls[0]; - expect(completedCall.status).toBe(CoreToolCallStatus.Success); - if (completedCall.status === CoreToolCallStatus.Success) { - expect(completedCall.response.resultDisplay).toBe('Tool executed'); - } - }); -}); - -describe('CoreToolScheduler request queueing', () => { - it('should queue a request if another is running', async () => { - let resolveFirstCall: (result: ToolResult) => void; - const firstCallPromise = new Promise((resolve) => { - resolveFirstCall = resolve; - }); - - const executeFn = vi.fn().mockImplementation(() => firstCallPromise); - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - const declarativeTool = mockTool; - - const mockToolRegistry = { - getTool: () => declarativeTool, - getToolByName: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.YOLO, // Use YOLO to avoid confirmation prompts - isInteractive: () => false, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request1 = { - callId: '1', - name: 'mockTool', - args: { a: 1 }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - const request2 = { - callId: '2', - name: 'mockTool', - args: { b: 2 }, - isClientInitiated: false, - prompt_id: 'prompt-2', - }; - - // Schedule the first call, which will pause execution. - // eslint-disable-next-line @typescript-eslint/no-floating-promises - scheduler.schedule([request1], abortController.signal); - - // 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( - [request2], - abortController.signal, - ); - - // Ensure the second tool call hasn't been executed yet. - expect(executeFn).toHaveBeenCalledWith( - { a: 1 }, - expect.anything(), - undefined, - expect.anything(), - ); - - // Complete the first tool call. - resolveFirstCall!({ - llmContent: 'First call complete', - returnDisplay: 'First call complete', - }); - - // Wait for the second schedule promise to resolve. - await schedulePromise2; - - // Let the second call finish. - const secondCallResult = { - llmContent: 'Second call complete', - returnDisplay: 'Second call complete', - }; - // Since the mock is shared, we need to resolve the current promise. - // In a real scenario, a new promise would be created for the second call. - resolveFirstCall!(secondCallResult); - - await vi.waitFor(() => { - // Now the second tool call should have been executed. - expect(executeFn).toHaveBeenCalledTimes(2); - }); - expect(executeFn).toHaveBeenCalledWith( - { b: 2 }, - expect.anything(), - undefined, - expect.anything(), - ); - - // Wait for the second completion. - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalledTimes(2); - }); - - // Verify the completion callbacks were called correctly. - 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 () => { - // Arrange - const executeFn = vi.fn().mockResolvedValue({ - llmContent: 'Tool executed', - returnDisplay: 'Tool executed', - }); - const mockTool = new MockTool({ - name: 'mockTool', - execute: executeFn, - shouldConfirmExecute: MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, - }); - const declarativeTool = mockTool; - - const toolRegistry = { - getTool: () => declarativeTool, - getToolByName: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - // Configure the scheduler to auto-approve the specific tool call. - const mockConfig = createMockConfig({ - getAllowedTools: () => ['mockTool'], // Auto-approve this tool - getToolRegistry: () => toolRegistry, - getShellExecutionConfig: () => ({ - terminalWidth: 80, - terminalHeight: 24, - sanitizationConfig: { - enableEnvironmentVariableRedaction: true, - allowedEnvironmentVariables: [], - blockedEnvironmentVariables: [], - }, - sandboxManager: new NoopSandboxManager(), - }), - isInteractive: () => false, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockTool', - args: { param: 'value' }, - isClientInitiated: false, - prompt_id: 'prompt-auto-approved', - }; - - // Act - await scheduler.schedule([request], abortController.signal); - - // Wait for the tool execution to complete - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - // Assert - // 1. The tool's execute method was called directly. - expect(executeFn).toHaveBeenCalledWith( - { param: 'value' }, - expect.anything(), - undefined, - expect.anything(), - ); - - // 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(CoreToolCallStatus.AwaitingApproval); - expect(statusUpdates).toEqual([ - CoreToolCallStatus.Validating, - CoreToolCallStatus.Scheduled, - CoreToolCallStatus.Executing, - CoreToolCallStatus.Success, - ]); - - // 3. The final callback indicates the tool call was successful. - expect(onAllToolCallsComplete).toHaveBeenCalled(); - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls).toHaveLength(1); - const completedCall = completedCalls[0]; - expect(completedCall.status).toBe(CoreToolCallStatus.Success); - if (completedCall.status === CoreToolCallStatus.Success) { - expect(completedCall.response.resultDisplay).toBe('Tool executed'); - } - }); - - it('should require approval for a chained shell command even when prefix is allowlisted', async () => { - const executeFn = vi.fn().mockResolvedValue({ - llmContent: 'Shell command executed', - returnDisplay: 'Shell command executed', - }); - - const mockShellTool = new MockTool({ - name: 'run_shell_command', - shouldConfirmExecute: (params) => - Promise.resolve({ - type: 'exec', - title: 'Confirm Shell Command', - command: String(params['command'] ?? ''), - rootCommand: 'git', - rootCommands: ['git'], - onConfirm: async () => {}, - }), - execute: () => executeFn({}), - }); - - const toolRegistry = { - getTool: () => mockShellTool, - getToolByName: () => mockShellTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => mockShellTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getAllowedTools: () => ['run_shell_command(git)'], - getShellExecutionConfig: () => ({ - terminalWidth: 80, - terminalHeight: 24, - sanitizationConfig: { - enableEnvironmentVariableRedaction: true, - allowedEnvironmentVariables: [], - blockedEnvironmentVariables: [], - }, - sandboxManager: new NoopSandboxManager(), - }), - getToolRegistry: () => toolRegistry, - getHookSystem: () => undefined, - getPolicyEngine: () => - ({ - check: async () => ({ decision: PolicyDecision.ASK_USER }), - }) as unknown as PolicyEngine, - }); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: 'shell-1', - name: 'run_shell_command', - args: { command: 'git status && rm -rf /tmp/should-not-run' }, - isClientInitiated: false, - prompt_id: 'prompt-shell-auto-approved', - }; - - await scheduler.schedule([request], abortController.signal); - - const statusUpdates = onToolCallsUpdate.mock.calls - .map((call) => (call[0][0] as ToolCall)?.status) - .filter(Boolean); - - expect(statusUpdates).toContain(CoreToolCallStatus.AwaitingApproval); - expect(executeFn).not.toHaveBeenCalled(); - expect(onAllToolCallsComplete).not.toHaveBeenCalled(); - }, 20000); - - it('should handle two synchronous calls to schedule', async () => { - const executeFn = vi.fn().mockResolvedValue({ - llmContent: 'Tool executed', - returnDisplay: 'Tool executed', - }); - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - const declarativeTool = mockTool; - const mockToolRegistry = { - getTool: () => declarativeTool, - getToolByName: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.YOLO, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request1 = { - callId: '1', - name: 'mockTool', - args: { a: 1 }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - const request2 = { - callId: '2', - name: 'mockTool', - args: { b: 2 }, - isClientInitiated: false, - prompt_id: 'prompt-2', - }; - - // Schedule two calls synchronously. - const schedulePromise1 = scheduler.schedule( - [request1], - abortController.signal, - ); - const schedulePromise2 = scheduler.schedule( - [request2], - abortController.signal, - ); - - // Wait for both promises to resolve. - await Promise.all([schedulePromise1, schedulePromise2]); - - // Ensure the tool was called twice with the correct arguments. - expect(executeFn).toHaveBeenCalledTimes(2); - expect(executeFn).toHaveBeenCalledWith( - { a: 1 }, - expect.anything(), - undefined, - expect.anything(), - ); - expect(executeFn).toHaveBeenCalledWith( - { b: 2 }, - expect.anything(), - undefined, - expect.anything(), - ); - - // Ensure completion callbacks were called twice. - expect(onAllToolCallsComplete).toHaveBeenCalledTimes(2); - }); - - it('should auto-approve remaining tool calls when first tool call is approved with ProceedAlways', async () => { - let approvalMode = ApprovalMode.DEFAULT; - const mockConfig = createMockConfig({ - getApprovalMode: () => approvalMode, - setApprovalMode: (mode: ApprovalMode) => { - approvalMode = mode; - }, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const testTool = new TestApprovalTool(mockConfig, mockMessageBus); - const toolRegistry = { - getTool: () => testTool, - getFunctionDeclarations: () => [], - getFunctionDeclarationsFiltered: () => [], - registerTool: () => {}, - discoverAllTools: async () => {}, - discoverMcpTools: async () => {}, - discoverToolsForServer: async () => {}, - removeMcpToolsByServer: () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - tools: new Map(), - context: mockConfig, - mcpClientManager: undefined, - getToolByName: () => testTool, - getToolByDisplayName: () => testTool, - getTools: () => [], - discoverTools: async () => {}, - discovery: {}, - } as unknown as ToolRegistry; - - mockConfig.getToolRegistry = () => toolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - const pendingConfirmations: Array< - (outcome: ToolConfirmationOutcome) => void - > = []; - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate: (toolCalls) => { - onToolCallsUpdate(toolCalls); - // Capture confirmation handlers for awaiting_approval tools - toolCalls.forEach((call) => { - if (call.status === CoreToolCallStatus.AwaitingApproval) { - const waitingCall = call; - const details = - waitingCall.confirmationDetails as ToolCallConfirmationDetails; - if (details?.onConfirm) { - const originalHandler = pendingConfirmations.find( - (h) => h === details.onConfirm, - ); - if (!originalHandler) { - pendingConfirmations.push(details.onConfirm); - } - } - } - }); - }, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - - // Schedule multiple tools that need confirmation - const requests = [ - { - callId: '1', - name: 'testApprovalTool', - args: { id: 'first' }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - { - callId: '2', - name: 'testApprovalTool', - args: { id: 'second' }, - isClientInitiated: false, - prompt_id: 'prompt-2', - }, - { - callId: '3', - name: 'testApprovalTool', - args: { id: 'third' }, - isClientInitiated: false, - prompt_id: 'prompt-3', - }, - ]; - - await scheduler.schedule(requests, abortController.signal); - - // Wait for the FIRST tool to be awaiting approval - await vi.waitFor(() => { - 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(CoreToolCallStatus.AwaitingApproval); - expect(calls?.[0].request.callId).toBe('1'); - // 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); - - // Approve the first tool with ProceedAlways - const firstConfirmation = pendingConfirmations[0]; - firstConfirmation(ToolConfirmationOutcome.ProceedAlways); - - // Wait for all tools to be completed - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - const completedCalls = onAllToolCallsComplete.mock.calls.at( - -1, - )?.[0] as ToolCall[]; - expect(completedCalls?.length).toBe(3); - expect( - completedCalls?.every( - (call) => call.status === CoreToolCallStatus.Success, - ), - ).toBe(true); - - // Verify approval mode was changed - expect(approvalMode).toBe(ApprovalMode.AUTO_EDIT); - }); -}); - -describe('CoreToolScheduler Sequential Execution', () => { - it('should execute tool calls in a batch sequentially', async () => { - // Arrange - let firstCallFinished = false; - const executeFn = vi - .fn() - .mockImplementation(async (args: { call: number }) => { - if (args.call === 1) { - // First call, wait for a bit to simulate work - await new Promise((resolve) => setTimeout(resolve, 50)); - firstCallFinished = true; - return { llmContent: 'First call done' }; - } - if (args.call === 2) { - // Second call, should only happen after the first is finished - if (!firstCallFinished) { - throw new Error( - 'Second tool call started before the first one finished!', - ); - } - return { llmContent: 'Second call done' }; - } - return { llmContent: 'default' }; - }); - - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - const declarativeTool = mockTool; - - const mockToolRegistry = { - getTool: () => declarativeTool, - getToolByName: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.YOLO, // Use YOLO to avoid confirmation prompts - isInteractive: () => false, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const requests = [ - { - callId: '1', - name: 'mockTool', - args: { call: 1 }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - { - callId: '2', - name: 'mockTool', - args: { call: 2 }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - ]; - - // Act - await scheduler.schedule(requests, abortController.signal); - - // Assert - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - // Check that execute was called twice - expect(executeFn).toHaveBeenCalledTimes(2); - - // Check the order of calls - const calls = executeFn.mock.calls; - expect(calls[0][0]).toEqual({ call: 1 }); - expect(calls[1][0]).toEqual({ call: 2 }); - - // The onAllToolCallsComplete should be called once with both results - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls).toHaveLength(2); - 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 () => { - // Arrange - const abortController = new AbortController(); - let secondCallStarted = false; - - const executeFn = vi - .fn() - .mockImplementation(async (args: { call: number }) => { - if (args.call === 1) { - return { llmContent: 'First call done' }; - } - if (args.call === 2) { - secondCallStarted = true; - // This call will be cancelled while it's "running". - await new Promise((resolve) => setTimeout(resolve, 100)); - // It should not return a value because it will be cancelled. - return { llmContent: 'Second call should not complete' }; - } - if (args.call === 3) { - return { llmContent: 'Third call done' }; - } - return { llmContent: 'default' }; - }); - - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - const declarativeTool = mockTool; - - const mockToolRegistry = { - getTool: () => declarativeTool, - getToolByName: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.YOLO, - isInteractive: () => false, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const requests = [ - { - callId: '1', - name: 'mockTool', - args: { call: 1 }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - { - callId: '2', - name: 'mockTool', - args: { call: 2 }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - { - callId: '3', - name: 'mockTool', - args: { call: 3 }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - ]; - - // Act - const schedulePromise = scheduler.schedule( - requests, - abortController.signal, - ); - - // Wait for the second call to start, then abort. - await vi.waitFor(() => { - expect(secondCallStarted).toBe(true); - }); - abortController.abort(); - - await schedulePromise; - - // Assert - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - // Check that execute was called for the first two tools only - expect(executeFn).toHaveBeenCalledTimes(2); - expect(executeFn).toHaveBeenCalledWith( - { call: 1 }, - expect.anything(), - undefined, - expect.anything(), - ); - expect(executeFn).toHaveBeenCalledWith( - { call: 2 }, - expect.anything(), - undefined, - expect.anything(), - ); - - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls).toHaveLength(3); - - const call1 = completedCalls.find((c) => c.request.callId === '1'); - const call2 = completedCalls.find((c) => c.request.callId === '2'); - const call3 = completedCalls.find((c) => c.request.callId === '3'); - - 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 () => { - const modifyWithEditorSpy = vi - .spyOn(modifiableToolModule, 'modifyWithEditor') - .mockResolvedValue({ - updatedParams: { param: 'updated' }, - updatedDiff: 'updated diff', - }); - - const mockModifiableTool = new MockModifiableTool('mockModifiableTool'); - const mockToolRegistry = { - getTool: () => mockModifiableTool, - getToolByName: () => mockModifiableTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => mockModifiableTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - - await scheduler.schedule( - [ - { - callId: '1', - name: 'mockModifiableTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - ], - abortController.signal, - ); - - const toolCall = (scheduler as unknown as { toolCalls: ToolCall[] }) - .toolCalls[0] as WaitingToolCall; - expect(toolCall.status).toBe(CoreToolCallStatus.AwaitingApproval); - - const confirmationSignal = new AbortController().signal; - await scheduler.handleConfirmationResponse( - toolCall.request.callId, - async () => {}, - ToolConfirmationOutcome.ModifyWithEditor, - confirmationSignal, - ); - - expect(modifyWithEditorSpy).toHaveBeenCalled(); - const overrides = - modifyWithEditorSpy.mock.calls[ - modifyWithEditorSpy.mock.calls.length - 1 - ][4]; - expect(overrides).toEqual({ - currentContent: 'originalContent', - proposedContent: 'newContent', - }); - - modifyWithEditorSpy.mockRestore(); - }); - - it('should handle inline modify with empty new content', async () => { - // Mock the modifiable check to return true for this test - const isModifiableSpy = vi - .spyOn(modifiableToolModule, 'isModifiableDeclarativeTool') - .mockReturnValue(true); - - const mockTool = new MockModifiableTool(); - const mockToolRegistry = { - getTool: () => mockTool, - getAllToolNames: () => [], - } as unknown as ToolRegistry; - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - isInteractive: () => true, - }); - mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - getPreferredEditor: () => 'vscode', - }); - - // Manually inject a waiting tool call - const callId = 'call-1'; - const toolCall: WaitingToolCall = { - status: CoreToolCallStatus.AwaitingApproval, - request: { - callId, - name: 'mockModifiableTool', - args: {}, - isClientInitiated: false, - prompt_id: 'p1', - }, - tool: mockTool, - invocation: {} as unknown as ToolInvocation< - Record, - ToolResult - >, - confirmationDetails: { - type: 'edit', - title: 'Confirm', - fileName: 'test.txt', - filePath: 'test.txt', - fileDiff: 'diff', - originalContent: 'old', - newContent: 'new', - onConfirm: async () => {}, - }, - startTime: Date.now(), - }; - - const schedulerInternals = scheduler as unknown as { - toolCalls: ToolCall[]; - toolModifier: { applyInlineModify: Mock }; - }; - schedulerInternals.toolCalls = [toolCall]; - - const applyInlineModifySpy = vi - .spyOn(schedulerInternals.toolModifier, 'applyInlineModify') - .mockResolvedValue({ - updatedParams: { content: '' }, - updatedDiff: 'diff-empty', - }); - - await scheduler.handleConfirmationResponse( - callId, - async () => {}, - ToolConfirmationOutcome.ProceedOnce, - new AbortController().signal, - { newContent: '' } as ToolConfirmationPayload, - ); - - expect(applyInlineModifySpy).toHaveBeenCalled(); - isModifiableSpy.mockRestore(); - }); - - it('should pass serverName and toolAnnotations to policy engine for DiscoveredMCPTool', async () => { - const mockMcpTool = { - tool: async () => ({ functionDeclarations: [] }), - callTool: async () => [], - }; - const serverName = 'test-server'; - const toolName = 'test-tool'; - const annotations = { readOnlyHint: true }; - const mcpTool = new DiscoveredMCPTool( - mockMcpTool as unknown as CallableTool, - serverName, - toolName, - 'description', - { type: 'object', properties: {} }, - createMockMessageBus() as unknown as MessageBus, - undefined, // trust - true, // isReadOnly - undefined, // nameOverride - undefined, // cliConfig - undefined, // extensionName - undefined, // extensionId - annotations, // toolAnnotations - ); - - const mockToolRegistry = { - getTool: () => mcpTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByName: () => mcpTool, - getToolByDisplayName: () => mcpTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const mockPolicyEngineCheck = vi.fn().mockResolvedValue({ - decision: PolicyDecision.ALLOW, - }); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getPolicyEngine: () => - ({ - check: mockPolicyEngineCheck, - }) as unknown as PolicyEngine, - isInteractive: () => false, - }); - mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: toolName, - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-id-1', - }; - - await scheduler.schedule(request, abortController.signal); - - expect(mockPolicyEngineCheck).toHaveBeenCalledWith( - expect.objectContaining({ name: toolName }), - serverName, - annotations, - ); - }); - - it('should not double-report completed tools when concurrent completions occur', async () => { - // Arrange - const executeFn = vi - .fn() - .mockResolvedValue({ llmContent: CoreToolCallStatus.Success }); - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - const declarativeTool = mockTool; - - const mockToolRegistry = { - getTool: () => declarativeTool, - getToolByName: () => declarativeTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => declarativeTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - let completionCallCount = 0; - const onAllToolCallsComplete = vi.fn().mockImplementation(async () => { - completionCallCount++; - // Simulate slow reporting (e.g. Gemini API call) - await new Promise((resolve) => setTimeout(resolve, 50)); - }); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.YOLO, - isInteractive: () => false, - }); - const mockMessageBus = createMockMessageBus(); - mockConfig.getMessageBus = vi.fn().mockReturnValue(mockMessageBus); - mockConfig.getEnableHooks = vi.fn().mockReturnValue(false); - mockConfig.getHookSystem = vi - .fn() - .mockReturnValue(new HookSystem(mockConfig)); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - - // Act - // 1. Start execution - const schedulePromise = scheduler.schedule( - [request], - abortController.signal, - ); - - // 2. Wait just enough for it to finish and enter checkAndNotifyCompletion - // (awaiting our slow mock) - await vi.waitFor(() => { - expect(completionCallCount).toBe(1); - }); - - // 3. Trigger a concurrent completion event (e.g. via cancelAll) - scheduler.cancelAll(abortController.signal); - - await schedulePromise; - - // Assert - // Even though cancelAll was called while the first completion was in progress, - // it should not have triggered a SECOND completion call because the first one - // was still 'finalizing' and will drain any new tools. - expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1); - }); - - it('should complete reporting all tools even mid-callback during abort', async () => { - // Arrange - const onAllToolCallsComplete = vi.fn().mockImplementation(async () => { - // Simulate slow reporting - await new Promise((resolve) => setTimeout(resolve, 50)); - }); - - const mockTool = new MockTool({ name: 'mockTool' }); - const mockToolRegistry = { - getTool: () => mockTool, - getToolByName: () => mockTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => mockTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.YOLO, - isInteractive: () => false, - }); - mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const signal = abortController.signal; - - // Act - // 1. Start execution of two tools - const schedulePromise = scheduler.schedule( - [ - { - callId: '1', - name: 'mockTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - { - callId: '2', - name: 'mockTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }, - ], - signal, - ); - - // 2. Wait for reporting to start - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - // 3. Abort the signal while reporting is in progress - abortController.abort(); - - await schedulePromise; - - // Assert - // Verify that onAllToolCallsComplete was called and processed the tools, - // and that the scheduler didn't just drop them because of the abort. - expect(onAllToolCallsComplete).toHaveBeenCalled(); - - const reportedTools = onAllToolCallsComplete.mock.calls.flatMap((call) => - // eslint-disable-next-line @typescript-eslint/no-explicit-any - call[0].map((t: any) => t.request.callId), - ); - - // Both tools should have been reported exactly once with success status - expect(reportedTools).toContain('1'); - expect(reportedTools).toContain('2'); - - const allStatuses = onAllToolCallsComplete.mock.calls.flatMap((call) => - // eslint-disable-next-line @typescript-eslint/no-explicit-any - call[0].map((t: any) => t.status), - ); - expect(allStatuses).toEqual([ - CoreToolCallStatus.Success, - CoreToolCallStatus.Success, - ]); - - expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1); - }); - - describe('Policy Decisions in Plan Mode', () => { - it('should return POLICY_VIOLATION error type and informative message when denied in Plan Mode', async () => { - const mockTool = new MockTool({ - name: 'dangerous_tool', - displayName: 'Dangerous Tool', - description: 'Does risky stuff', - }); - const mockToolRegistry = { - getTool: () => mockTool, - getAllToolNames: () => ['dangerous_tool'], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.PLAN, - getPolicyEngine: () => - ({ - check: async () => ({ decision: PolicyDecision.DENY }), - }) as unknown as PolicyEngine, - }); - mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - getPreferredEditor: () => 'vscode', - }); - - const request = { - callId: 'call-1', - name: 'dangerous_tool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - - await scheduler.schedule(request, new AbortController().signal); - - expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1); - const reportedTools = onAllToolCallsComplete.mock.calls[0][0]; - const result = reportedTools[0]; - - 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.', - ); - }); - - it('should return custom deny message when denied in Plan Mode with a specific rule message', async () => { - const mockTool = new MockTool({ - name: 'dangerous_tool', - displayName: 'Dangerous Tool', - description: 'Does risky stuff', - }); - const mockToolRegistry = { - getTool: () => mockTool, - getAllToolNames: () => ['dangerous_tool'], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const customDenyMessage = 'Custom denial message for testing'; - - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.PLAN, - getPolicyEngine: () => - ({ - check: async () => ({ - decision: PolicyDecision.DENY, - rule: { denyMessage: customDenyMessage }, - }), - }) as unknown as PolicyEngine, - }); - mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - getPreferredEditor: () => 'vscode', - }); - - const request = { - callId: 'call-1', - name: 'dangerous_tool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - - await scheduler.schedule(request, new AbortController().signal); - - expect(onAllToolCallsComplete).toHaveBeenCalledTimes(1); - const reportedTools = onAllToolCallsComplete.mock.calls[0][0]; - const result = reportedTools[0]; - - 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}`, - ); - }); - }); - - describe('ApprovalMode Preservation', () => { - it('should preserve approvalMode throughout tool lifecycle', async () => { - // Arrange - const executeFn = vi.fn().mockResolvedValue({ - llmContent: 'Tool executed', - returnDisplay: 'Tool executed', - }); - const mockTool = new MockTool({ - name: 'mockTool', - execute: executeFn, - shouldConfirmExecute: MOCK_TOOL_SHOULD_CONFIRM_EXECUTE, - }); - - const mockToolRegistry = { - getTool: () => mockTool, - getAllToolNames: () => ['mockTool'], - } as unknown as ToolRegistry; - - const onAllToolCallsComplete = vi.fn(); - const onToolCallsUpdate = vi.fn(); - - // Set approval mode to PLAN - const mockConfig = createMockConfig({ - getToolRegistry: () => mockToolRegistry, - getApprovalMode: () => ApprovalMode.PLAN, - // Ensure policy engine returns ASK_USER to trigger AwaitingApproval state - getPolicyEngine: () => - ({ - check: async () => ({ decision: PolicyDecision.ASK_USER }), - }) as unknown as PolicyEngine, - }); - mockConfig.getHookSystem = vi.fn().mockReturnValue(undefined); - - const scheduler = new CoreToolScheduler({ - context: mockConfig, - onAllToolCallsComplete, - onToolCallsUpdate, - getPreferredEditor: () => 'vscode', - }); - - const abortController = new AbortController(); - const request = { - callId: '1', - name: 'mockTool', - args: { param: 'value' }, - isClientInitiated: false, - prompt_id: 'test-prompt', - }; - - // Act - Schedule - const schedulePromise = scheduler.schedule( - request, - abortController.signal, - ); - - // Assert - Check AwaitingApproval state - const awaitingCall = (await waitForStatus( - onToolCallsUpdate, - CoreToolCallStatus.AwaitingApproval, - )) as WaitingToolCall; - - expect(awaitingCall).toBeDefined(); - expect(awaitingCall.approvalMode).toBe(ApprovalMode.PLAN); - - // Act - Confirm - - await ( - awaitingCall.confirmationDetails as ToolCallConfirmationDetails - ).onConfirm(ToolConfirmationOutcome.ProceedOnce); - - // Wait for completion - await schedulePromise; - - // Assert - Check Success state - expect(onAllToolCallsComplete).toHaveBeenCalled(); - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls).toHaveLength(1); - expect(completedCalls[0].status).toBe(CoreToolCallStatus.Success); - expect(completedCalls[0].approvalMode).toBe(ApprovalMode.PLAN); - }); - }); -}); diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts deleted file mode 100644 index 8aabd709c2..0000000000 --- a/packages/core/src/core/coreToolScheduler.ts +++ /dev/null @@ -1,1164 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - type ToolResultDisplay, - type AnyDeclarativeTool, - type AnyToolInvocation, - type ToolCallConfirmationDetails, - type ToolConfirmationPayload, - ToolConfirmationOutcome, -} from '../tools/tools.js'; -import type { EditorType } from '../utils/editor.js'; -import { PolicyDecision } from '../policy/types.js'; -import { logToolCall } from '../telemetry/loggers.js'; -import { ToolErrorType } from '../tools/tool-error.js'; -import { ToolCallEvent } from '../telemetry/types.js'; -import { runInDevTraceSpan } from '../telemetry/trace.js'; -import { ToolModificationHandler } from '../scheduler/tool-modifier.js'; -import { - getToolSuggestion, - isToolCallResponseInfo, -} from '../utils/tool-utils.js'; -import type { ToolConfirmationRequest } from '../confirmation-bus/types.js'; -import { MessageBusType } from '../confirmation-bus/types.js'; -import type { MessageBus } from '../confirmation-bus/message-bus.js'; -import { - CoreToolCallStatus, - type ToolCall, - type ValidatingToolCall, - type ScheduledToolCall, - type ErroredToolCall, - type SuccessfulToolCall, - type ExecutingToolCall, - type CancelledToolCall, - type WaitingToolCall, - type Status, - type CompletedToolCall, - type ConfirmHandler, - type OutputUpdateHandler, - type AllToolCallsCompleteHandler, - type ToolCallsUpdateHandler, - type ToolCallRequestInfo, - type ToolCallResponseInfo, -} from '../scheduler/types.js'; -import { ToolExecutor } from '../scheduler/tool-executor.js'; -import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; -import { getPolicyDenialError } from '../scheduler/policy.js'; -import { GeminiCliOperation } from '../telemetry/constants.js'; -import { evaluateBeforeToolHook } from '../scheduler/hook-utils.js'; -import type { AgentLoopContext } from '../config/agent-loop-context.js'; - -export type { - ToolCall, - ValidatingToolCall, - ScheduledToolCall, - ErroredToolCall, - SuccessfulToolCall, - ExecutingToolCall, - CancelledToolCall, - WaitingToolCall, - Status, - CompletedToolCall, - ConfirmHandler, - OutputUpdateHandler, - AllToolCallsCompleteHandler, - ToolCallsUpdateHandler, - ToolCallRequestInfo, - ToolCallResponseInfo, -}; - -const createErrorResponse = ( - request: ToolCallRequestInfo, - error: Error, - errorType: ToolErrorType | undefined, -): ToolCallResponseInfo => ({ - callId: request.callId, - error, - responseParts: [ - { - functionResponse: { - id: request.callId, - name: request.name, - response: { error: error.message }, - }, - }, - ], - resultDisplay: error.message, - errorType, - contentLength: error.message.length, -}); - -interface CoreToolSchedulerOptions { - context: AgentLoopContext; - outputUpdateHandler?: OutputUpdateHandler; - onAllToolCallsComplete?: AllToolCallsCompleteHandler; - onToolCallsUpdate?: ToolCallsUpdateHandler; - getPreferredEditor: () => EditorType | undefined; -} - -export class CoreToolScheduler { - // Static WeakMap to track which MessageBus instances already have a handler subscribed - // This prevents duplicate subscriptions when multiple CoreToolScheduler instances are created - private static subscribedMessageBuses = new WeakMap< - MessageBus, - (request: ToolConfirmationRequest) => void - >(); - - private toolCalls: ToolCall[] = []; - private outputUpdateHandler?: OutputUpdateHandler; - private onAllToolCallsComplete?: AllToolCallsCompleteHandler; - private onToolCallsUpdate?: ToolCallsUpdateHandler; - private getPreferredEditor: () => EditorType | undefined; - private context: AgentLoopContext; - private isFinalizingToolCalls = false; - private isScheduling = false; - private isCancelling = false; - private requestQueue: Array<{ - request: ToolCallRequestInfo | ToolCallRequestInfo[]; - signal: AbortSignal; - resolve: () => void; - reject: (reason?: Error) => void; - }> = []; - private toolCallQueue: ToolCall[] = []; - private completedToolCallsForBatch: CompletedToolCall[] = []; - private toolExecutor: ToolExecutor; - private toolModifier: ToolModificationHandler; - - constructor(options: CoreToolSchedulerOptions) { - this.context = options.context; - this.outputUpdateHandler = options.outputUpdateHandler; - this.onAllToolCallsComplete = options.onAllToolCallsComplete; - this.onToolCallsUpdate = options.onToolCallsUpdate; - this.getPreferredEditor = options.getPreferredEditor; - this.toolExecutor = new ToolExecutor(this.context); - this.toolModifier = new ToolModificationHandler(); - - // Subscribe to message bus for ASK_USER policy decisions - // Use a static WeakMap to ensure we only subscribe ONCE per MessageBus instance - // This prevents memory leaks when multiple CoreToolScheduler instances are created - // (e.g., on every React render, or for each non-interactive tool call) - const messageBus = this.context.messageBus; - - // Check if we've already subscribed a handler to this message bus - if (!CoreToolScheduler.subscribedMessageBuses.has(messageBus)) { - // Create a shared handler that will be used for this message bus - const sharedHandler = (request: ToolConfirmationRequest) => { - // When ASK_USER policy decision is made, respond with requiresUserConfirmation=true - // to tell tools to use their legacy confirmation flow - // eslint-disable-next-line @typescript-eslint/no-floating-promises - messageBus.publish({ - type: MessageBusType.TOOL_CONFIRMATION_RESPONSE, - correlationId: request.correlationId, - confirmed: false, - requiresUserConfirmation: true, - }); - }; - - messageBus.subscribe( - MessageBusType.TOOL_CONFIRMATION_REQUEST, - sharedHandler, - ); - - // Store the handler in the WeakMap so we don't subscribe again - CoreToolScheduler.subscribedMessageBuses.set(messageBus, sharedHandler); - } - } - - private setStatusInternal( - targetCallId: string, - status: CoreToolCallStatus.Success, - signal: AbortSignal, - response: ToolCallResponseInfo, - ): void; - private setStatusInternal( - targetCallId: string, - status: CoreToolCallStatus.AwaitingApproval, - signal: AbortSignal, - confirmationDetails: ToolCallConfirmationDetails, - ): void; - private setStatusInternal( - targetCallId: string, - status: CoreToolCallStatus.Error, - signal: AbortSignal, - response: ToolCallResponseInfo, - ): void; - private setStatusInternal( - targetCallId: string, - status: CoreToolCallStatus.Cancelled, - signal: AbortSignal, - reason: string, - ): void; - private setStatusInternal( - targetCallId: string, - status: - | CoreToolCallStatus.Executing - | CoreToolCallStatus.Scheduled - | CoreToolCallStatus.Validating, - signal: AbortSignal, - ): void; - private setStatusInternal( - targetCallId: string, - newStatus: Status, - signal: AbortSignal, - auxiliaryData?: unknown, - ): void { - this.toolCalls = this.toolCalls.map((currentCall) => { - if ( - currentCall.request.callId !== targetCallId || - currentCall.status === CoreToolCallStatus.Success || - currentCall.status === CoreToolCallStatus.Error || - currentCall.status === CoreToolCallStatus.Cancelled - ) { - return currentCall; - } - - // currentCall is a non-terminal state here and should have startTime and tool. - const existingStartTime = currentCall.startTime; - const toolInstance = currentCall.tool; - const invocation = currentCall.invocation; - - const outcome = currentCall.outcome; - const approvalMode = currentCall.approvalMode; - - switch (newStatus) { - case CoreToolCallStatus.Success: { - const durationMs = existingStartTime - ? Date.now() - existingStartTime - : undefined; - if (isToolCallResponseInfo(auxiliaryData)) { - return { - request: currentCall.request, - tool: toolInstance, - invocation, - status: CoreToolCallStatus.Success, - response: auxiliaryData, - durationMs, - outcome, - approvalMode, - } as SuccessfulToolCall; - } - throw new Error('Invalid response data for tool success'); - } - case CoreToolCallStatus.Error: { - const durationMs = existingStartTime - ? Date.now() - existingStartTime - : undefined; - if (isToolCallResponseInfo(auxiliaryData)) { - return { - request: currentCall.request, - status: CoreToolCallStatus.Error, - tool: toolInstance, - response: auxiliaryData, - durationMs, - outcome, - approvalMode, - } as ErroredToolCall; - } - throw new Error('Invalid response data for tool error'); - } - case CoreToolCallStatus.AwaitingApproval: - return { - request: currentCall.request, - tool: toolInstance, - status: CoreToolCallStatus.AwaitingApproval, - confirmationDetails: - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - auxiliaryData as ToolCallConfirmationDetails, - startTime: existingStartTime, - outcome, - invocation, - approvalMode, - } as WaitingToolCall; - case CoreToolCallStatus.Scheduled: - return { - request: currentCall.request, - tool: toolInstance, - status: CoreToolCallStatus.Scheduled, - startTime: existingStartTime, - outcome, - invocation, - approvalMode, - } as ScheduledToolCall; - case CoreToolCallStatus.Cancelled: { - const durationMs = existingStartTime - ? Date.now() - existingStartTime - : undefined; - - if (isToolCallResponseInfo(auxiliaryData)) { - return { - request: currentCall.request, - tool: toolInstance, - invocation, - status: CoreToolCallStatus.Cancelled, - response: auxiliaryData, - durationMs, - outcome, - approvalMode, - } as CancelledToolCall; - } - - // Preserve diff for cancelled edit operations - let resultDisplay: ToolResultDisplay | undefined = undefined; - if (currentCall.status === CoreToolCallStatus.AwaitingApproval) { - const waitingCall = currentCall; - if (waitingCall.confirmationDetails.type === 'edit') { - resultDisplay = { - fileDiff: waitingCall.confirmationDetails.fileDiff, - fileName: waitingCall.confirmationDetails.fileName, - originalContent: - waitingCall.confirmationDetails.originalContent, - newContent: waitingCall.confirmationDetails.newContent, - filePath: waitingCall.confirmationDetails.filePath, - }; - } - } - - const errorMessage = `[Operation Cancelled] Reason: ${auxiliaryData}`; - return { - request: currentCall.request, - tool: toolInstance, - invocation, - status: CoreToolCallStatus.Cancelled, - response: { - callId: currentCall.request.callId, - responseParts: [ - { - functionResponse: { - id: currentCall.request.callId, - name: currentCall.request.name, - response: { - error: errorMessage, - }, - }, - }, - ], - resultDisplay, - error: undefined, - errorType: undefined, - contentLength: errorMessage.length, - }, - durationMs, - outcome, - approvalMode, - } as CancelledToolCall; - } - case CoreToolCallStatus.Validating: - return { - request: currentCall.request, - tool: toolInstance, - status: CoreToolCallStatus.Validating, - startTime: existingStartTime, - outcome, - invocation, - approvalMode, - } as ValidatingToolCall; - case CoreToolCallStatus.Executing: - return { - request: currentCall.request, - tool: toolInstance, - status: CoreToolCallStatus.Executing, - startTime: existingStartTime, - outcome, - invocation, - approvalMode, - } as ExecutingToolCall; - default: { - const exhaustiveCheck: never = newStatus; - return exhaustiveCheck; - } - } - }); - this.notifyToolCallsUpdate(); - } - - private setArgsInternal(targetCallId: string, args: unknown): void { - 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 === CoreToolCallStatus.Error - ) { - return call; - } - - const invocationOrError = this.buildInvocation( - call.tool, - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - args as Record, - ); - if (invocationOrError instanceof Error) { - const response = createErrorResponse( - call.request, - invocationOrError, - ToolErrorType.INVALID_TOOL_PARAMS, - ); - return { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - request: { ...call.request, args: args as Record }, - status: CoreToolCallStatus.Error, - tool: call.tool, - response, - approvalMode: call.approvalMode, - } as ErroredToolCall; - } - - return { - ...call, - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - request: { ...call.request, args: args as Record }, - invocation: invocationOrError, - }; - }); - } - - private isRunning(): boolean { - return ( - this.isFinalizingToolCalls || - this.toolCalls.some( - (call) => - call.status === CoreToolCallStatus.Executing || - call.status === CoreToolCallStatus.AwaitingApproval, - ) - ); - } - - private buildInvocation( - tool: AnyDeclarativeTool, - args: object, - ): AnyToolInvocation | Error { - try { - return tool.build(args); - } catch (e) { - if (e instanceof Error) { - return e; - } - return new Error(String(e)); - } - } - - schedule( - request: ToolCallRequestInfo | ToolCallRequestInfo[], - signal: AbortSignal, - ): Promise { - return runInDevTraceSpan( - { operation: GeminiCliOperation.ScheduleToolCalls }, - async ({ metadata: spanMetadata }) => { - spanMetadata.input = request; - if (this.isRunning() || this.isScheduling) { - return new Promise((resolve, reject) => { - const abortHandler = () => { - // Find and remove the request from the queue - const index = this.requestQueue.findIndex( - (item) => item.request === request, - ); - if (index > -1) { - this.requestQueue.splice(index, 1); - reject(new Error('Tool call cancelled while in queue.')); - } - }; - - signal.addEventListener('abort', abortHandler, { once: true }); - - this.requestQueue.push({ - request, - signal, - resolve: () => { - signal.removeEventListener('abort', abortHandler); - resolve(); - }, - reject: (reason?: Error) => { - signal.removeEventListener('abort', abortHandler); - reject(reason); - }, - }); - }); - } - return this._schedule(request, signal); - }, - ); - } - - cancelAll(signal: AbortSignal): void { - if (this.isCancelling) { - return; - } - this.isCancelling = true; - // Cancel the currently active tool call, if there is one. - if (this.toolCalls.length > 0) { - const activeCall = this.toolCalls[0]; - // Only cancel if it's in a cancellable state. - if ( - activeCall.status === CoreToolCallStatus.AwaitingApproval || - activeCall.status === CoreToolCallStatus.Executing || - activeCall.status === CoreToolCallStatus.Scheduled || - activeCall.status === CoreToolCallStatus.Validating - ) { - this.setStatusInternal( - activeCall.request.callId, - CoreToolCallStatus.Cancelled, - signal, - 'User cancelled the operation.', - ); - } - } - - // Clear the queue and mark all queued items as cancelled for completion reporting. - this._cancelAllQueuedCalls(); - - // Finalize the batch immediately. - void this.checkAndNotifyCompletion(signal); - } - - private async _schedule( - request: ToolCallRequestInfo | ToolCallRequestInfo[], - signal: AbortSignal, - ): Promise { - this.isScheduling = true; - this.isCancelling = false; - try { - if (this.isRunning()) { - throw new Error( - 'Cannot schedule new tool calls while other tool calls are actively running (executing or awaiting approval).', - ); - } - const requestsToProcess = Array.isArray(request) ? request : [request]; - const currentApprovalMode = this.context.config.getApprovalMode(); - this.completedToolCallsForBatch = []; - - const newToolCalls: ToolCall[] = requestsToProcess.map( - (reqInfo): ToolCall => { - const toolInstance = this.context.toolRegistry.getTool(reqInfo.name); - if (!toolInstance) { - const suggestion = getToolSuggestion( - reqInfo.name, - this.context.toolRegistry.getAllToolNames(), - ); - const errorMessage = `Tool "${reqInfo.name}" not found in registry. Tools must use the exact names that are registered.${suggestion}`; - return { - status: CoreToolCallStatus.Error, - request: reqInfo, - response: createErrorResponse( - reqInfo, - new Error(errorMessage), - ToolErrorType.TOOL_NOT_REGISTERED, - ), - durationMs: 0, - approvalMode: currentApprovalMode, - }; - } - - const invocationOrError = this.buildInvocation( - toolInstance, - reqInfo.args, - ); - if (invocationOrError instanceof Error) { - return { - status: CoreToolCallStatus.Error, - request: reqInfo, - tool: toolInstance, - response: createErrorResponse( - reqInfo, - invocationOrError, - ToolErrorType.INVALID_TOOL_PARAMS, - ), - durationMs: 0, - approvalMode: currentApprovalMode, - }; - } - - return { - status: CoreToolCallStatus.Validating, - request: reqInfo, - tool: toolInstance, - invocation: invocationOrError, - startTime: Date.now(), - approvalMode: currentApprovalMode, - }; - }, - ); - - this.toolCallQueue.push(...newToolCalls); - await this._processNextInQueue(signal); - } finally { - this.isScheduling = false; - } - } - - private async _processNextInQueue(signal: AbortSignal): Promise { - // If there's already a tool being processed, or the queue is empty, stop. - if (this.toolCalls.length > 0 || this.toolCallQueue.length === 0) { - return; - } - - // If cancellation happened between steps, handle it. - if (signal.aborted) { - this._cancelAllQueuedCalls(); - // Finalize the batch. - await this.checkAndNotifyCompletion(signal); - return; - } - - let toolCall = this.toolCallQueue.shift()!; - - // This is now the single active tool call. - this.toolCalls = [toolCall]; - this.notifyToolCallsUpdate(); - - // Handle tools that were already errored during creation. - 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); - return; - } - - // This logic is moved from the old `for` loop in `_schedule`. - if (toolCall.status === CoreToolCallStatus.Validating) { - let { request: reqInfo } = toolCall; - - try { - if (signal.aborted) { - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.Cancelled, - signal, - 'Tool call cancelled by user.', - ); - // The completion check will handle the cascade. - await this.checkAndNotifyCompletion(signal); - return; - } - - // 1. Hook Check (BeforeTool) - const hookResult = await evaluateBeforeToolHook( - this.context.config, - toolCall.tool, - toolCall.request, - toolCall.invocation, - ); - - if (hookResult.status === 'error') { - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.Error, - signal, - createErrorResponse( - toolCall.request, - hookResult.error, - hookResult.errorType, - ), - ); - await this.checkAndNotifyCompletion(signal); - return; - } - - const { hookDecision, hookSystemMessage, modifiedArgs, newInvocation } = - hookResult; - - if (modifiedArgs && newInvocation) { - this.setArgsInternal(reqInfo.callId, modifiedArgs); - // Re-retrieve toolCall as it was updated in the array by setArgsInternal - const updatedCall = this.toolCalls.find( - (c) => c.request.callId === reqInfo.callId, - ); - if ( - updatedCall && - updatedCall.status === CoreToolCallStatus.Validating - ) { - toolCall = updatedCall; - } - toolCall.request.inputModifiedByHook = true; - reqInfo = toolCall.request; - } - - // 2. Policy Check using PolicyEngine - // We must reconstruct the FunctionCall format expected by PolicyEngine - const toolCallForPolicy = { - name: toolCall.request.name, - args: toolCall.request.args, - }; - const serverName = - toolCall.tool instanceof DiscoveredMCPTool - ? toolCall.tool.serverName - : undefined; - const toolAnnotations = toolCall.tool.toolAnnotations; - - const { decision: policyDecision, rule } = await this.context.config - .getPolicyEngine() - .check(toolCallForPolicy, serverName, toolAnnotations); - - let finalDecision = policyDecision; - if (hookDecision === 'ask') { - finalDecision = PolicyDecision.ASK_USER; - } - - if (finalDecision === PolicyDecision.DENY) { - const { errorMessage, errorType } = getPolicyDenialError( - this.context.config, - rule, - ); - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.Error, - signal, - createErrorResponse(reqInfo, new Error(errorMessage), errorType), - ); - await this.checkAndNotifyCompletion(signal); - return; - } - - if (finalDecision === PolicyDecision.ALLOW) { - this.setToolCallOutcome( - reqInfo.callId, - ToolConfirmationOutcome.ProceedAlways, - ); - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.Scheduled, - signal, - ); - } else { - // PolicyDecision.ASK_USER - - // We need confirmation details to show to the user - const confirmationDetails = - await toolCall.invocation.shouldConfirmExecute( - signal, - hookDecision === 'ask' ? 'ask_user' : undefined, - ); - - if (!confirmationDetails) { - this.setToolCallOutcome( - reqInfo.callId, - ToolConfirmationOutcome.ProceedAlways, - ); - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.Scheduled, - signal, - ); - } else { - if (!this.context.config.isInteractive()) { - throw new Error( - `Tool execution for "${ - toolCall.tool.displayName || toolCall.tool.name - }" requires user confirmation, which is not supported in non-interactive mode.`, - ); - } - - if (hookSystemMessage) { - confirmationDetails.systemMessage = hookSystemMessage; - } - - // Fire Notification hook before showing confirmation to user - const hookSystem = this.context.config.getHookSystem(); - if (hookSystem) { - await hookSystem.fireToolNotificationEvent(confirmationDetails); - } - - // Allow IDE to resolve confirmation - if ( - confirmationDetails.type === 'edit' && - confirmationDetails.ideConfirmation - ) { - // eslint-disable-next-line @typescript-eslint/no-floating-promises - confirmationDetails.ideConfirmation.then((resolution) => { - if (resolution.status === 'accepted') { - // eslint-disable-next-line @typescript-eslint/no-floating-promises - this.handleConfirmationResponse( - reqInfo.callId, - confirmationDetails.onConfirm, - ToolConfirmationOutcome.ProceedOnce, - signal, - ); - } else { - // eslint-disable-next-line @typescript-eslint/no-floating-promises - this.handleConfirmationResponse( - reqInfo.callId, - confirmationDetails.onConfirm, - ToolConfirmationOutcome.Cancel, - signal, - ); - } - }); - } - - const originalOnConfirm = confirmationDetails.onConfirm; - const wrappedConfirmationDetails: ToolCallConfirmationDetails = { - ...confirmationDetails, - onConfirm: ( - outcome: ToolConfirmationOutcome, - payload?: ToolConfirmationPayload, - ) => - this.handleConfirmationResponse( - reqInfo.callId, - originalOnConfirm, - outcome, - signal, - payload, - ), - }; - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.AwaitingApproval, - signal, - wrappedConfirmationDetails, - ); - } - } - } catch (error) { - if (signal.aborted) { - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.Cancelled, - signal, - 'Tool call cancelled by user.', - ); - await this.checkAndNotifyCompletion(signal); - } else { - this.setStatusInternal( - reqInfo.callId, - CoreToolCallStatus.Error, - signal, - createErrorResponse( - reqInfo, - error instanceof Error ? error : new Error(String(error)), - ToolErrorType.UNHANDLED_EXCEPTION, - ), - ); - await this.checkAndNotifyCompletion(signal); - } - } - } - await this.attemptExecutionOfScheduledCalls(signal); - } - - async handleConfirmationResponse( - callId: string, - originalOnConfirm: (outcome: ToolConfirmationOutcome) => Promise, - outcome: ToolConfirmationOutcome, - signal: AbortSignal, - payload?: ToolConfirmationPayload, - ): Promise { - const toolCall = this.toolCalls.find( - (c) => - c.request.callId === callId && - c.status === CoreToolCallStatus.AwaitingApproval, - ); - - if (toolCall && toolCall.status === CoreToolCallStatus.AwaitingApproval) { - await originalOnConfirm(outcome); - } - - this.setToolCallOutcome(callId, outcome); - - if (outcome === ToolConfirmationOutcome.Cancel || signal.aborted) { - // Instead of just cancelling one tool, trigger the full cancel cascade. - this.cancelAll(signal); - return; // `cancelAll` calls `checkAndNotifyCompletion`, so we can exit here. - } else if (outcome === ToolConfirmationOutcome.ModifyWithEditor) { - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - const waitingToolCall = toolCall as WaitingToolCall; - - const editorType = this.getPreferredEditor(); - if (!editorType) { - return; - } - - /* 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, - editorType, - signal, - ); - - // Restore status (isModifying: false) and update diff if result exists - if (result) { - this.setArgsInternal(callId, result.updatedParams); - /* 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 @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 - // re-confirmation. - if (payload && 'newContent' in payload && toolCall) { - const result = await this.toolModifier.applyInlineModify( - // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion - toolCall as WaitingToolCall, - payload, - signal, - ); - if (result) { - this.setArgsInternal(callId, result.updatedParams); - /* 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, CoreToolCallStatus.Scheduled, signal); - } - await this.attemptExecutionOfScheduledCalls(signal); - } - - private async attemptExecutionOfScheduledCalls( - signal: AbortSignal, - ): Promise { - const allCallsFinalOrScheduled = this.toolCalls.every( - (call) => - 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 === CoreToolCallStatus.Scheduled, - ); - - for (const toolCall of callsToExecute) { - if (toolCall.status !== CoreToolCallStatus.Scheduled) continue; - - this.setStatusInternal( - toolCall.request.callId, - CoreToolCallStatus.Executing, - signal, - ); - const executingCall = this.toolCalls.find( - (c) => c.request.callId === toolCall.request.callId, - ); - - if (!executingCall) { - // Should not happen, but safe guard - continue; - } - - const completedCall = await this.toolExecutor.execute({ - call: executingCall, - signal, - outputUpdateHandler: (callId, output) => { - if (this.outputUpdateHandler) { - this.outputUpdateHandler(callId, output); - } - this.toolCalls = this.toolCalls.map((tc) => - tc.request.callId === callId && - tc.status === CoreToolCallStatus.Executing - ? { ...tc, liveOutput: output } - : tc, - ); - this.notifyToolCallsUpdate(); - }, - onUpdateToolCall: (updatedCall) => { - this.toolCalls = this.toolCalls.map((tc) => - tc.request.callId === updatedCall.request.callId - ? updatedCall - : tc, - ); - this.notifyToolCallsUpdate(); - }, - }); - - this.toolCalls = this.toolCalls.map((tc) => - tc.request.callId === completedCall.request.callId - ? { ...completedCall, approvalMode: tc.approvalMode } - : tc, - ); - this.notifyToolCallsUpdate(); - - await this.checkAndNotifyCompletion(signal); - } - } - } - - private async checkAndNotifyCompletion(signal: AbortSignal): Promise { - // This method is now only concerned with the single active tool call. - if (this.toolCalls.length === 0) { - // It's possible to be called when a batch is cancelled before any tool has started. - if (signal.aborted && this.toolCallQueue.length > 0) { - this._cancelAllQueuedCalls(); - } - } else { - const activeCall = this.toolCalls[0]; - const isTerminal = - 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 CoreToolCallStatus.Executing or CoreToolCallStatus.AwaitingApproval), - // then the scheduler is still busy or paused. We should not proceed. - if (!isTerminal) { - return; - } - - // The active tool is finished. Move it to the completed batch. - const completedCall = activeCall as CompletedToolCall; - this.completedToolCallsForBatch.push(completedCall); - logToolCall(this.context.config, new ToolCallEvent(completedCall)); - - // Clear the active tool slot. This is crucial for the sequential processing. - this.toolCalls = []; - } - - // Now, check if the entire batch is complete. - // The batch is complete if the queue is empty or the operation was cancelled. - if (this.toolCallQueue.length === 0 || signal.aborted) { - if (signal.aborted) { - this._cancelAllQueuedCalls(); - } - - // If we are already finalizing, another concurrent call to - // checkAndNotifyCompletion will just return. The ongoing finalized loop - // will pick up any new tools added to completedToolCallsForBatch. - if (this.isFinalizingToolCalls) { - return; - } - - // If there's nothing to report and we weren't cancelled, we can stop. - // But if we were cancelled, we must proceed to potentially start the next queued request. - if (this.completedToolCallsForBatch.length === 0 && !signal.aborted) { - return; - } - - this.isFinalizingToolCalls = true; - try { - // We use a while loop here to ensure that if new tools are added to the - // batch (e.g., via cancellation) while we are awaiting - // onAllToolCallsComplete, they are also reported before we finish. - while (this.completedToolCallsForBatch.length > 0) { - const batchToReport = [...this.completedToolCallsForBatch]; - this.completedToolCallsForBatch = []; - if (this.onAllToolCallsComplete) { - await this.onAllToolCallsComplete(batchToReport); - } - } - } finally { - this.isFinalizingToolCalls = false; - this.isCancelling = false; - this.notifyToolCallsUpdate(); - } - - // After completion of the entire batch, process the next item in the main request queue. - if (this.requestQueue.length > 0) { - const next = this.requestQueue.shift()!; - this._schedule(next.request, next.signal) - .then(next.resolve) - .catch(next.reject); - } - } else { - // The batch is not yet complete, so continue processing the current batch sequence. - await this._processNextInQueue(signal); - } - } - - private _cancelAllQueuedCalls(): void { - while (this.toolCallQueue.length > 0) { - const queuedCall = this.toolCallQueue.shift()!; - // Don't cancel tools that already errored during validation. - if (queuedCall.status === CoreToolCallStatus.Error) { - this.completedToolCallsForBatch.push(queuedCall); - continue; - } - const durationMs = - 'startTime' in queuedCall && queuedCall.startTime - ? Date.now() - queuedCall.startTime - : undefined; - const errorMessage = - '[Operation Cancelled] User cancelled the operation.'; - this.completedToolCallsForBatch.push({ - request: queuedCall.request, - tool: queuedCall.tool, - invocation: queuedCall.invocation, - status: CoreToolCallStatus.Cancelled, - response: { - callId: queuedCall.request.callId, - responseParts: [ - { - functionResponse: { - id: queuedCall.request.callId, - name: queuedCall.request.name, - response: { - error: errorMessage, - }, - }, - }, - ], - resultDisplay: undefined, - error: undefined, - errorType: undefined, - contentLength: errorMessage.length, - }, - durationMs, - outcome: ToolConfirmationOutcome.Cancel, - approvalMode: queuedCall.approvalMode, - }); - } - } - - private notifyToolCallsUpdate(): void { - if (this.onToolCallsUpdate) { - this.onToolCallsUpdate([ - ...this.completedToolCallsForBatch, - ...this.toolCalls, - ...this.toolCallQueue, - ]); - } - } - - private setToolCallOutcome(callId: string, outcome: ToolConfirmationOutcome) { - this.toolCalls = this.toolCalls.map((call) => { - if (call.request.callId !== callId) return call; - return { - ...call, - outcome, - }; - }); - } -} diff --git a/packages/core/src/core/coreToolSchedulerHooks.test.ts b/packages/core/src/core/coreToolSchedulerHooks.test.ts deleted file mode 100644 index a6c2e470d0..0000000000 --- a/packages/core/src/core/coreToolSchedulerHooks.test.ts +++ /dev/null @@ -1,313 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, it, expect, vi } from 'vitest'; -import { CoreToolScheduler } from './coreToolScheduler.js'; -import type { ToolCall, ErroredToolCall } from '../scheduler/types.js'; -import type { Config, ToolRegistry, AgentLoopContext } from '../index.js'; -import { - ApprovalMode, - DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, -} from '../index.js'; -import { createMockMessageBus } from '../test-utils/mock-message-bus.js'; -import { MockTool } from '../test-utils/mock-tool.js'; -import { DEFAULT_GEMINI_MODEL } from '../config/models.js'; -import type { PolicyEngine } from '../policy/policy-engine.js'; -import type { HookSystem } from '../hooks/hookSystem.js'; -import { BeforeToolHookOutput } from '../hooks/types.js'; - -function createMockConfig(overrides: Partial = {}): Config { - const defaultToolRegistry = { - getTool: () => undefined, - getToolByName: () => undefined, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => undefined, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - getExperiments: () => {}, - } as unknown as ToolRegistry; - - const baseConfig = { - getSessionId: () => 'test-session-id', - getUsageStatisticsEnabled: () => true, - getDebugMode: () => false, - isInteractive: () => true, - getApprovalMode: () => ApprovalMode.DEFAULT, - setApprovalMode: () => {}, - getAllowedTools: () => [], - getContentGeneratorConfig: () => ({ - model: 'test-model', - authType: 'oauth-personal', - }), - getShellExecutionConfig: () => ({ - terminalWidth: 90, - terminalHeight: 30, - sanitizationConfig: { - enableEnvironmentVariableRedaction: true, - allowedEnvironmentVariables: [], - blockedEnvironmentVariables: [], - }, - }), - storage: { - getProjectTempDir: () => '/tmp', - }, - getTruncateToolOutputThreshold: () => - DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, - getTruncateToolOutputLines: () => 1000, - getToolRegistry: () => defaultToolRegistry, - getActiveModel: () => DEFAULT_GEMINI_MODEL, - getGeminiClient: () => null, - getMessageBus: () => createMockMessageBus(), - getEnableHooks: () => true, // Enabled for these tests - getExperiments: () => {}, - getPolicyEngine: () => - ({ - check: async () => ({ decision: 'allow' }), // Default allow for hook tests - }) as unknown as PolicyEngine, - } as unknown as Config; - - // eslint-disable-next-line @typescript-eslint/no-misused-spread - return { ...baseConfig, ...overrides } as Config; -} - -describe('CoreToolScheduler Hooks', () => { - it('should stop execution if BeforeTool hook requests stop', async () => { - const executeFn = vi.fn().mockResolvedValue({ - llmContent: 'Tool executed', - returnDisplay: 'Tool executed', - }); - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - - const toolRegistry = { - getTool: () => mockTool, - getToolByName: () => mockTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => mockTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const mockMessageBus = createMockMessageBus(); - const mockHookSystem = { - fireBeforeToolEvent: vi.fn().mockResolvedValue({ - shouldStopExecution: () => true, - getEffectiveReason: () => 'Hook stopped execution', - getBlockingError: () => ({ blocked: false }), - isAskDecision: () => false, - }), - } as unknown as HookSystem; - - const mockConfig = createMockConfig({ - getToolRegistry: () => toolRegistry, - getMessageBus: () => mockMessageBus, - getHookSystem: () => mockHookSystem, - getApprovalMode: () => ApprovalMode.YOLO, - }); - - const onAllToolCallsComplete = vi.fn(); - const scheduler = new CoreToolScheduler({ - context: { - config: mockConfig, - messageBus: mockMessageBus, - toolRegistry, - } as unknown as AgentLoopContext, - onAllToolCallsComplete, - getPreferredEditor: () => 'vscode', - }); - - const request = { - callId: '1', - name: 'mockTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - - await scheduler.schedule([request], new AbortController().signal); - - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe('error'); - const erroredCall = completedCalls[0] as ErroredToolCall; - - // Check error type/message - expect(erroredCall.response.error?.message).toContain( - 'Hook stopped execution', - ); - expect(executeFn).not.toHaveBeenCalled(); - }); - - it('should block tool execution if BeforeTool hook requests block', async () => { - const executeFn = vi.fn(); - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - - const toolRegistry = { - getTool: () => mockTool, - getToolByName: () => mockTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => mockTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const mockMessageBus = createMockMessageBus(); - const mockHookSystem = { - fireBeforeToolEvent: vi.fn().mockResolvedValue({ - shouldStopExecution: () => false, - getBlockingError: () => ({ - blocked: true, - reason: 'Hook blocked execution', - }), - isAskDecision: () => false, - }), - } as unknown as HookSystem; - - const mockConfig = createMockConfig({ - getToolRegistry: () => toolRegistry, - getMessageBus: () => mockMessageBus, - getHookSystem: () => mockHookSystem, - getApprovalMode: () => ApprovalMode.YOLO, - }); - - const onAllToolCallsComplete = vi.fn(); - const scheduler = new CoreToolScheduler({ - context: { - config: mockConfig, - messageBus: mockMessageBus, - toolRegistry, - } as unknown as AgentLoopContext, - onAllToolCallsComplete, - getPreferredEditor: () => 'vscode', - }); - - const request = { - callId: '1', - name: 'mockTool', - args: {}, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - - await scheduler.schedule([request], new AbortController().signal); - - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe('error'); - const erroredCall = completedCalls[0] as ErroredToolCall; - expect(erroredCall.response.error?.message).toContain( - 'Hook blocked execution', - ); - expect(executeFn).not.toHaveBeenCalled(); - }); - - it('should update tool input if BeforeTool hook provides modified input', async () => { - const executeFn = vi.fn().mockResolvedValue({ - llmContent: 'Tool executed', - returnDisplay: 'Tool executed', - }); - const mockTool = new MockTool({ name: 'mockTool', execute: executeFn }); - - const toolRegistry = { - getTool: () => mockTool, - getToolByName: () => mockTool, - getFunctionDeclarations: () => [], - tools: new Map(), - discovery: {}, - registerTool: () => {}, - getToolByDisplayName: () => mockTool, - getTools: () => [], - discoverTools: async () => {}, - getAllTools: () => [], - getToolsByServer: () => [], - } as unknown as ToolRegistry; - - const mockMessageBus = createMockMessageBus(); - const mockBeforeOutput = new BeforeToolHookOutput({ - continue: true, - hookSpecificOutput: { - hookEventName: 'BeforeTool', - tool_input: { newParam: 'modifiedValue' }, - }, - }); - - const mockHookSystem = { - fireBeforeToolEvent: vi.fn().mockResolvedValue(mockBeforeOutput), - fireAfterToolEvent: vi.fn(), - } as unknown as HookSystem; - - const mockConfig = createMockConfig({ - getToolRegistry: () => toolRegistry, - getMessageBus: () => mockMessageBus, - getHookSystem: () => mockHookSystem, - getApprovalMode: () => ApprovalMode.YOLO, - }); - - const onAllToolCallsComplete = vi.fn(); - const scheduler = new CoreToolScheduler({ - context: { - config: mockConfig, - messageBus: mockMessageBus, - toolRegistry, - } as unknown as AgentLoopContext, - onAllToolCallsComplete, - getPreferredEditor: () => 'vscode', - }); - - const request = { - callId: '1', - name: 'mockTool', - args: { originalParam: 'originalValue' }, - isClientInitiated: false, - prompt_id: 'prompt-1', - }; - - await scheduler.schedule([request], new AbortController().signal); - - await vi.waitFor(() => { - expect(onAllToolCallsComplete).toHaveBeenCalled(); - }); - - const completedCalls = onAllToolCallsComplete.mock - .calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe('success'); - - // Verify execute was called with modified args - expect(executeFn).toHaveBeenCalledWith( - { newParam: 'modifiedValue' }, - expect.anything(), - undefined, - expect.anything(), - ); - - // Verify call request args were updated in the completion report - expect(completedCalls[0].request.args).toEqual({ - newParam: 'modifiedValue', - }); - }); -}); diff --git a/packages/core/src/scheduler/policy.test.ts b/packages/core/src/scheduler/policy.test.ts index 435fe6524d..abcfc422cd 100644 --- a/packages/core/src/scheduler/policy.test.ts +++ b/packages/core/src/scheduler/policy.test.ts @@ -34,11 +34,9 @@ import { ROOT_SCHEDULER_ID, type ValidatingToolCall, type ToolCallRequestInfo, - type CompletedToolCall, } from './types.js'; import type { PolicyEngine } from '../policy/policy-engine.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; -import { CoreToolScheduler } from '../core/coreToolScheduler.js'; import { Scheduler } from './scheduler.js'; import { ToolErrorType } from '../tools/tool-error.js'; import type { ToolRegistry } from '../tools/tool-registry.js'; @@ -840,61 +838,32 @@ describe('Plan Mode Denial Consistency', () => { vi.clearAllMocks(); }); - describe.each([ - { enableEventDrivenScheduler: false, name: 'Legacy CoreToolScheduler' }, - { enableEventDrivenScheduler: true, name: 'Event-Driven Scheduler' }, - ])('$name', ({ enableEventDrivenScheduler }) => { - it('should return the correct Plan Mode denial message when policy denies execution', async () => { - let resultMessage: string | undefined; - let resultErrorType: ToolErrorType | undefined; + it('should return the correct Plan Mode denial message when policy denies execution', async () => { + let resultMessage: string | undefined; + let resultErrorType: ToolErrorType | undefined; - const signal = new AbortController().signal; + const signal = new AbortController().signal; - if (enableEventDrivenScheduler) { - const scheduler = new Scheduler({ - context: { - config: mockConfig, - messageBus: mockMessageBus, - toolRegistry: mockToolRegistry, - } as unknown as AgentLoopContext, - getPreferredEditor: () => undefined, - schedulerId: ROOT_SCHEDULER_ID, - }); - - const results = await scheduler.schedule(req, signal); - const result = results[0]; - - expect(result.status).toBe('error'); - if (result.status === 'error') { - resultMessage = result.response.error?.message; - resultErrorType = result.response.errorType; - } - } else { - let capturedCalls: CompletedToolCall[] = []; - const scheduler = new CoreToolScheduler({ - context: { - config: mockConfig, - messageBus: mockMessageBus, - toolRegistry: mockToolRegistry, - } as unknown as AgentLoopContext, - getPreferredEditor: () => undefined, - onAllToolCallsComplete: async (calls) => { - capturedCalls = calls; - }, - }); - - await scheduler.schedule(req, signal); - - expect(capturedCalls.length).toBeGreaterThan(0); - const call = capturedCalls[0]; - if (call.status === 'error') { - resultMessage = call.response.error?.message; - resultErrorType = call.response.errorType; - } - } - - expect(resultMessage).toBe('Tool execution denied by policy.'); - expect(resultErrorType).toBe(ToolErrorType.POLICY_VIOLATION); + const scheduler = new Scheduler({ + context: { + config: mockConfig, + messageBus: mockMessageBus, + toolRegistry: mockToolRegistry, + } as unknown as AgentLoopContext, + getPreferredEditor: () => undefined, + schedulerId: ROOT_SCHEDULER_ID, }); + + const results = await scheduler.schedule(req, signal); + const result = results[0]; + + expect(result.status).toBe('error'); + if (result.status === 'error') { + resultMessage = result.response.error?.message; + resultErrorType = result.response.errorType; + } + + expect(resultMessage).toBe('Tool execution denied by policy.'); + expect(resultErrorType).toBe(ToolErrorType.POLICY_VIOLATION); }); }); diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts index 0ea6c390d3..69ac326d7f 100644 --- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts +++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.test.ts @@ -25,7 +25,7 @@ import { AuthType, type ContentGeneratorConfig, } from '../../core/contentGenerator.js'; -import type { SuccessfulToolCall } from '../../core/coreToolScheduler.js'; +import type { SuccessfulToolCall } from '../../scheduler/types.js'; import type { ConfigParameters } from '../../config/config.js'; import { EventMetadataKey } from './event-metadata-key.js'; import { makeFakeConfig } from '../../test-utils/config.js'; diff --git a/packages/core/src/telemetry/uiTelemetry.test.ts b/packages/core/src/telemetry/uiTelemetry.test.ts index 9669a5ae59..263f904b5a 100644 --- a/packages/core/src/telemetry/uiTelemetry.test.ts +++ b/packages/core/src/telemetry/uiTelemetry.test.ts @@ -20,7 +20,7 @@ import type { CompletedToolCall, ErroredToolCall, SuccessfulToolCall, -} from '../core/coreToolScheduler.js'; +} from '../scheduler/types.js'; import { ToolErrorType } from '../tools/tool-error.js'; import { ToolConfirmationOutcome } from '../tools/tools.js'; import { MockTool } from '../test-utils/mock-tool.js';