From 5020d8fa5769151f5aeed50865b3442800177d40 Mon Sep 17 00:00:00 2001 From: Adam Weidman <65992621+adamfweidman@users.noreply.github.com> Date: Tue, 10 Mar 2026 22:38:38 -0400 Subject: [PATCH] refactor(a2a): remove legacy CoreToolScheduler (#21955) --- packages/a2a-server/src/agent/task.test.ts | 221 +----------------- packages/a2a-server/src/agent/task.ts | 173 +------------- packages/a2a-server/src/config/config.ts | 2 - packages/a2a-server/src/config/settings.ts | 3 - packages/a2a-server/src/http/app.test.ts | 7 +- .../a2a-server/src/utils/testing_utils.ts | 6 +- 6 files changed, 16 insertions(+), 396 deletions(-) diff --git a/packages/a2a-server/src/agent/task.test.ts b/packages/a2a-server/src/agent/task.test.ts index bf15d7fc49..26039ae3aa 100644 --- a/packages/a2a-server/src/agent/task.test.ts +++ b/packages/a2a-server/src/agent/task.test.ts @@ -4,25 +4,14 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { - describe, - it, - expect, - vi, - beforeEach, - afterEach, - type Mock, -} from 'vitest'; +import { describe, it, expect, vi, beforeEach, type Mock } from 'vitest'; import { Task } from './task.js'; import { GeminiEventType, - ApprovalMode, - ToolConfirmationOutcome, type Config, type ToolCallRequestInfo, type GitService, type CompletedToolCall, - type ToolCall, } from '@google/gemini-cli-core'; import { createMockConfig } from '../utils/testing_utils.js'; import type { ExecutionEventBus, RequestContext } from '@a2a-js/sdk/server'; @@ -389,214 +378,6 @@ describe('Task', () => { ); }); - describe('_schedulerToolCallsUpdate', () => { - let task: Task; - type SpyInstance = ReturnType; - let setTaskStateAndPublishUpdateSpy: SpyInstance; - let mockConfig: Config; - let mockEventBus: ExecutionEventBus; - - beforeEach(() => { - mockConfig = createMockConfig() as Config; - mockEventBus = { - publish: vi.fn(), - on: vi.fn(), - off: vi.fn(), - once: vi.fn(), - removeAllListeners: vi.fn(), - finished: vi.fn(), - }; - - // @ts-expect-error - Calling private constructor - task = new Task('task-id', 'context-id', mockConfig, mockEventBus); - - // Spy on the method we want to check calls for - setTaskStateAndPublishUpdateSpy = vi.spyOn( - task, - 'setTaskStateAndPublishUpdate', - ); - }); - - afterEach(() => { - vi.restoreAllMocks(); - }); - - it('should set state to input-required when a tool is awaiting approval and none are executing', () => { - const toolCalls = [ - { request: { callId: '1' }, status: 'awaiting_approval' }, - ] as ToolCall[]; - - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(toolCalls); - - // The last call should be the final state update - expect(setTaskStateAndPublishUpdateSpy).toHaveBeenLastCalledWith( - 'input-required', - { kind: 'state-change' }, - undefined, - undefined, - true, // final: true - ); - }); - - it('should NOT set state to input-required if a tool is awaiting approval but another is executing', () => { - const toolCalls = [ - { request: { callId: '1' }, status: 'awaiting_approval' }, - { request: { callId: '2' }, status: 'executing' }, - ] as ToolCall[]; - - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(toolCalls); - - // It will be called for status updates, but not with final: true - const finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find( - (call) => call[4] === true, - ); - expect(finalCall).toBeUndefined(); - }); - - it('should set state to input-required once an executing tool finishes, leaving one awaiting approval', () => { - const initialToolCalls = [ - { request: { callId: '1' }, status: 'awaiting_approval' }, - { request: { callId: '2' }, status: 'executing' }, - ] as ToolCall[]; - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(initialToolCalls); - - // No final call yet - let finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find( - (call) => call[4] === true, - ); - expect(finalCall).toBeUndefined(); - - // Now, the executing tool finishes. The scheduler would call _resolveToolCall for it. - // @ts-expect-error - Calling private method - task._resolveToolCall('2'); - - // Then another update comes in for the awaiting tool (e.g., a re-check) - const subsequentToolCalls = [ - { request: { callId: '1' }, status: 'awaiting_approval' }, - ] as ToolCall[]; - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(subsequentToolCalls); - - // NOW we should get the final call - finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find( - (call) => call[4] === true, - ); - expect(finalCall).toBeDefined(); - expect(finalCall?.[0]).toBe('input-required'); - }); - - it('should NOT set state to input-required if skipFinalTrueAfterInlineEdit is true', () => { - task.skipFinalTrueAfterInlineEdit = true; - const toolCalls = [ - { request: { callId: '1' }, status: 'awaiting_approval' }, - ] as ToolCall[]; - - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(toolCalls); - - const finalCall = setTaskStateAndPublishUpdateSpy.mock.calls.find( - (call) => call[4] === true, - ); - expect(finalCall).toBeUndefined(); - }); - - describe('auto-approval', () => { - it('should NOT publish ToolCallConfirmationEvent when autoExecute is true', () => { - task.autoExecute = true; - const onConfirmSpy = vi.fn(); - const toolCalls = [ - { - request: { callId: '1' }, - status: 'awaiting_approval', - correlationId: 'test-corr-id', - confirmationDetails: { - type: 'edit', - onConfirm: onConfirmSpy, - }, - }, - ] as unknown as ToolCall[]; - - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(toolCalls); - - expect(onConfirmSpy).toHaveBeenCalledWith( - ToolConfirmationOutcome.ProceedOnce, - ); - const calls = (mockEventBus.publish as Mock).mock.calls; - // Search if ToolCallConfirmationEvent was published - const confEvent = calls.find( - (call) => - call[0].metadata?.coderAgent?.kind === - CoderAgentEvent.ToolCallConfirmationEvent, - ); - expect(confEvent).toBeUndefined(); - }); - - it('should NOT publish ToolCallConfirmationEvent when approval mode is YOLO', () => { - (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.YOLO); - task.autoExecute = false; - const onConfirmSpy = vi.fn(); - const toolCalls = [ - { - request: { callId: '1' }, - status: 'awaiting_approval', - correlationId: 'test-corr-id', - confirmationDetails: { - type: 'edit', - onConfirm: onConfirmSpy, - }, - }, - ] as unknown as ToolCall[]; - - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(toolCalls); - - expect(onConfirmSpy).toHaveBeenCalledWith( - ToolConfirmationOutcome.ProceedOnce, - ); - const calls = (mockEventBus.publish as Mock).mock.calls; - // Search if ToolCallConfirmationEvent was published - const confEvent = calls.find( - (call) => - call[0].metadata?.coderAgent?.kind === - CoderAgentEvent.ToolCallConfirmationEvent, - ); - expect(confEvent).toBeUndefined(); - }); - - it('should NOT auto-approve when autoExecute is false and mode is not YOLO', () => { - task.autoExecute = false; - (mockConfig.getApprovalMode as Mock).mockReturnValue( - ApprovalMode.DEFAULT, - ); - const onConfirmSpy = vi.fn(); - const toolCalls = [ - { - request: { callId: '1' }, - status: 'awaiting_approval', - confirmationDetails: { onConfirm: onConfirmSpy }, - }, - ] as unknown as ToolCall[]; - - // @ts-expect-error - Calling private method - task._schedulerToolCallsUpdate(toolCalls); - - expect(onConfirmSpy).not.toHaveBeenCalled(); - const calls = (mockEventBus.publish as Mock).mock.calls; - // Search if ToolCallConfirmationEvent was published - const confEvent = calls.find( - (call) => - call[0].metadata?.coderAgent?.kind === - CoderAgentEvent.ToolCallConfirmationEvent, - ); - expect(confEvent).toBeDefined(); - }); - }); - }); - describe('currentPromptId and promptCount', () => { it('should correctly initialize and update promptId and promptCount', async () => { const mockConfig = createMockConfig(); diff --git a/packages/a2a-server/src/agent/task.ts b/packages/a2a-server/src/agent/task.ts index e0124f4585..94a03171d7 100644 --- a/packages/a2a-server/src/agent/task.ts +++ b/packages/a2a-server/src/agent/task.ts @@ -6,7 +6,6 @@ import { Scheduler, - CoreToolScheduler, type GeminiClient, GeminiEventType, ToolConfirmationOutcome, @@ -69,37 +68,10 @@ import type { PartUnion, Part as genAiPart } from '@google/genai'; type UnionKeys = T extends T ? keyof T : never; -type ConfirmationType = ToolCallConfirmationDetails['type']; - -const VALID_CONFIRMATION_TYPES: readonly ConfirmationType[] = [ - 'edit', - 'exec', - 'mcp', - 'info', - 'ask_user', - 'exit_plan_mode', -] as const; - -function isToolCallConfirmationDetails( - value: unknown, -): value is ToolCallConfirmationDetails { - if ( - typeof value !== 'object' || - value === null || - !('onConfirm' in value) || - typeof value.onConfirm !== 'function' || - !('type' in value) || - typeof value.type !== 'string' - ) { - return false; - } - return (VALID_CONFIRMATION_TYPES as readonly string[]).includes(value.type); -} - export class Task { id: string; contextId: string; - scheduler: Scheduler | CoreToolScheduler; + scheduler: Scheduler; config: Config; geminiClient: GeminiClient; pendingToolConfirmationDetails: Map; @@ -140,11 +112,7 @@ export class Task { this.contextId = contextId; this.config = config; - if (this.config.isEventDrivenSchedulerEnabled()) { - this.scheduler = this.setupEventDrivenScheduler(); - } else { - this.scheduler = this.createLegacyScheduler(); - } + this.scheduler = this.setupEventDrivenScheduler(); this.geminiClient = this.config.getGeminiClient(); this.pendingToolConfirmationDetails = new Map(); @@ -260,11 +228,7 @@ export class Task { this.pendingToolCalls.clear(); this.pendingCorrelationIds.clear(); - if (this.scheduler instanceof Scheduler) { - this.scheduler.cancelAll(); - } else { - this.scheduler.cancelAll(new AbortController().signal); - } + this.scheduler.cancelAll(); // Reset the promise for any future operations, ensuring it's in a clean state. this._resetToolCompletionPromise(); } @@ -409,133 +373,6 @@ export class Task { this.eventBus?.publish(artifactEvent); } - private async _schedulerAllToolCallsComplete( - completedToolCalls: CompletedToolCall[], - ): Promise { - logger.info( - '[Task] All tool calls completed by scheduler (batch):', - completedToolCalls.map((tc) => tc.request.callId), - ); - this.completedToolCalls.push(...completedToolCalls); - completedToolCalls.forEach((tc) => { - this._resolveToolCall(tc.request.callId); - }); - } - - private _schedulerToolCallsUpdate(toolCalls: ToolCall[]): void { - logger.info( - '[Task] Scheduler tool calls updated:', - toolCalls.map((tc) => `${tc.request.callId} (${tc.status})`), - ); - - // Update state and send continuous, non-final updates - toolCalls.forEach((tc) => { - const previousStatus = this.pendingToolCalls.get(tc.request.callId); - const hasChanged = previousStatus !== tc.status; - - // Resolve tool call if it has reached a terminal state - if (['success', 'error', 'cancelled'].includes(tc.status)) { - this._resolveToolCall(tc.request.callId); - } else { - // This will update the map - this._registerToolCall(tc.request.callId, tc.status); - } - - if (tc.status === 'awaiting_approval' && tc.confirmationDetails) { - const details = tc.confirmationDetails; - if (isToolCallConfirmationDetails(details)) { - this.pendingToolConfirmationDetails.set(tc.request.callId, details); - } - } - - // Only send an update if the status has actually changed. - if (hasChanged) { - // Skip sending confirmation event if we are going to auto-approve it anyway - if ( - tc.status === 'awaiting_approval' && - tc.confirmationDetails && - this.isYoloMatch - ) { - logger.info( - `[Task] Skipping ToolCallConfirmationEvent for ${tc.request.callId} due to YOLO mode.`, - ); - } else { - const coderAgentMessage: CoderAgentMessage = - tc.status === 'awaiting_approval' - ? { kind: CoderAgentEvent.ToolCallConfirmationEvent } - : { kind: CoderAgentEvent.ToolCallUpdateEvent }; - const message = this.toolStatusMessage(tc, this.id, this.contextId); - - const event = this._createStatusUpdateEvent( - this.taskState, - coderAgentMessage, - message, - false, // Always false for these continuous updates - ); - this.eventBus?.publish(event); - } - } - }); - - if (this.isYoloMatch) { - logger.info( - '[Task] ' + - (this.autoExecute ? '' : 'YOLO mode enabled. ') + - 'Auto-approving all tool calls.', - ); - toolCalls.forEach((tc: ToolCall) => { - if (tc.status === 'awaiting_approval' && tc.confirmationDetails) { - const details = tc.confirmationDetails; - if (isToolCallConfirmationDetails(details)) { - // eslint-disable-next-line @typescript-eslint/no-floating-promises - details.onConfirm(ToolConfirmationOutcome.ProceedOnce); - this.pendingToolConfirmationDetails.delete(tc.request.callId); - } - } - }); - return; - } - - const allPendingStatuses = Array.from(this.pendingToolCalls.values()); - const isAwaitingApproval = allPendingStatuses.some( - (status) => status === 'awaiting_approval', - ); - const isExecuting = allPendingStatuses.some( - (status) => status === 'executing', - ); - - // The turn is complete and requires user input if at least one tool - // is waiting for the user's decision, and no other tool is actively - // running in the background. - if ( - isAwaitingApproval && - !isExecuting && - !this.skipFinalTrueAfterInlineEdit - ) { - this.skipFinalTrueAfterInlineEdit = false; - - // We don't need to send another message, just a final status update. - this.setTaskStateAndPublishUpdate( - 'input-required', - { kind: CoderAgentEvent.StateChangeEvent }, - undefined, - undefined, - /*final*/ true, - ); - } - } - - private createLegacyScheduler(): CoreToolScheduler { - const scheduler = new CoreToolScheduler({ - outputUpdateHandler: this._schedulerOutputUpdate.bind(this), - onAllToolCallsComplete: this._schedulerAllToolCallsComplete.bind(this), - onToolCallsUpdate: this._schedulerToolCallsUpdate.bind(this), - getPreferredEditor: () => DEFAULT_GUI_EDITOR, - config: this.config, - }); - return scheduler; - } - private messageBusListener?: (message: ToolCallsUpdateMessage) => void; private setupEventDrivenScheduler(): Scheduler { @@ -564,9 +401,7 @@ export class Task { this.messageBusListener = undefined; } - if (this.scheduler instanceof Scheduler) { - this.scheduler.dispose(); - } + this.scheduler.dispose(); } private handleEventDrivenToolCallsUpdate( diff --git a/packages/a2a-server/src/config/config.ts b/packages/a2a-server/src/config/config.ts index 229abc65c9..5b6757701d 100644 --- a/packages/a2a-server/src/config/config.ts +++ b/packages/a2a-server/src/config/config.ts @@ -106,8 +106,6 @@ export async function loadConfig( trustedFolder: true, extensionLoader, checkpointing, - enableEventDrivenScheduler: - settings.experimental?.enableEventDrivenScheduler ?? true, interactive: !isHeadlessMode(), enableInteractiveShell: !isHeadlessMode(), ptyInfo: 'auto', diff --git a/packages/a2a-server/src/config/settings.ts b/packages/a2a-server/src/config/settings.ts index 0c353b46aa..da9db4e069 100644 --- a/packages/a2a-server/src/config/settings.ts +++ b/packages/a2a-server/src/config/settings.ts @@ -40,9 +40,6 @@ export interface Settings { general?: { previewFeatures?: boolean; }; - experimental?: { - enableEventDrivenScheduler?: boolean; - }; // Git-aware file filtering settings fileFiltering?: { diff --git a/packages/a2a-server/src/http/app.test.ts b/packages/a2a-server/src/http/app.test.ts index 7262be42a8..4a883992b5 100644 --- a/packages/a2a-server/src/http/app.test.ts +++ b/packages/a2a-server/src/http/app.test.ts @@ -65,7 +65,12 @@ vi.mock('../utils/logger.js', () => ({ })); let config: Config; -const getToolRegistrySpy = vi.fn().mockReturnValue(ApprovalMode.DEFAULT); +const getToolRegistrySpy = vi.fn().mockReturnValue({ + getTool: vi.fn(), + getAllToolNames: vi.fn().mockReturnValue([]), + getAllTools: vi.fn().mockReturnValue([]), + getToolsByServer: vi.fn().mockReturnValue([]), +}); const getApprovalModeSpy = vi.fn(); const getShellExecutionConfigSpy = vi.fn(); const getExtensionsSpy = vi.fn(); diff --git a/packages/a2a-server/src/utils/testing_utils.ts b/packages/a2a-server/src/utils/testing_utils.ts index 4981dbbd67..f63e66e85e 100644 --- a/packages/a2a-server/src/utils/testing_utils.ts +++ b/packages/a2a-server/src/utils/testing_utils.ts @@ -20,6 +20,7 @@ import { tmpdir, type Config, type Storage, + type ToolRegistry, } from '@google/gemini-cli-core'; import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js'; import { expect, vi } from 'vitest'; @@ -30,6 +31,10 @@ export function createMockConfig( const tmpDir = tmpdir(); // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion const mockConfig = { + get toolRegistry(): ToolRegistry { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return (this as unknown as Config).getToolRegistry(); + }, getToolRegistry: vi.fn().mockReturnValue({ getTool: vi.fn(), getAllToolNames: vi.fn().mockReturnValue([]), @@ -64,7 +69,6 @@ export function createMockConfig( getEmbeddingModel: vi.fn().mockReturnValue('text-embedding-004'), getSessionId: vi.fn().mockReturnValue('test-session-id'), getUserTier: vi.fn(), - isEventDrivenSchedulerEnabled: vi.fn().mockReturnValue(false), getMessageBus: vi.fn(), getPolicyEngine: vi.fn(), getEnableExtensionReloading: vi.fn().mockReturnValue(false),