From 6021e4c3baed2e1d9ff2c40d6e8ac8a13c1d9beb Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Wed, 14 Jan 2026 17:22:44 -0500 Subject: [PATCH] feat(scheduler): add types needed for event driven scheduler (#16641) --- packages/a2a-server/src/agent/task.ts | 6 +- packages/cli/src/ui/hooks/useGeminiStream.ts | 11 ++-- .../cli/src/ui/hooks/useReactToolScheduler.ts | 4 +- packages/core/src/confirmation-bus/types.ts | 57 ++++++++++++++++++- .../core/src/core/coreToolScheduler.test.ts | 28 ++++----- packages/core/src/scheduler/types.ts | 14 ++++- 6 files changed, 98 insertions(+), 22 deletions(-) diff --git a/packages/a2a-server/src/agent/task.ts b/packages/a2a-server/src/agent/task.ts index 7bf5e5ad4c..c82e9e6992 100644 --- a/packages/a2a-server/src/agent/task.ts +++ b/packages/a2a-server/src/agent/task.ts @@ -378,7 +378,7 @@ export class Task { if (tc.status === 'awaiting_approval' && tc.confirmationDetails) { this.pendingToolConfirmationDetails.set( tc.request.callId, - tc.confirmationDetails, + tc.confirmationDetails as ToolCallConfirmationDetails, ); } @@ -412,7 +412,9 @@ export class Task { toolCalls.forEach((tc: ToolCall) => { if (tc.status === 'awaiting_approval' && tc.confirmationDetails) { // eslint-disable-next-line @typescript-eslint/no-floating-promises - tc.confirmationDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce); + (tc.confirmationDetails as ToolCallConfirmationDetails).onConfirm( + ToolConfirmationOutcome.ProceedOnce, + ); this.pendingToolConfirmationDetails.delete(tc.request.callId); } }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index ab88962047..21add85556 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -43,6 +43,7 @@ import type { ToolCallRequestInfo, GeminiErrorEventValue, RetryAttemptPayload, + ToolCallConfirmationDetails, } from '@google/gemini-cli-core'; import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import type { @@ -1132,11 +1133,13 @@ export const useGeminiStream = ( // Process pending tool calls sequentially to reduce UI chaos for (const call of awaitingApprovalCalls) { - if (call.confirmationDetails?.onConfirm) { + if ( + (call.confirmationDetails as ToolCallConfirmationDetails)?.onConfirm + ) { try { - await call.confirmationDetails.onConfirm( - ToolConfirmationOutcome.ProceedOnce, - ); + await ( + call.confirmationDetails as ToolCallConfirmationDetails + ).onConfirm(ToolConfirmationOutcome.ProceedOnce); } catch (error) { debugLogger.warn( `Failed to auto-approve tool call ${call.request.callId}:`, diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index af32b18288..8939f5aa80 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -17,6 +17,7 @@ import type { AllToolCallsCompleteHandler, ToolCallsUpdateHandler, ToolCall, + ToolCallConfirmationDetails, Status as CoreStatus, EditorType, } from '@google/gemini-cli-core'; @@ -306,7 +307,8 @@ export function mapToDisplay( ...baseDisplayProperties, status: mapCoreStatusToDisplayStatus(trackedCall.status), resultDisplay: undefined, - confirmationDetails: trackedCall.confirmationDetails, + confirmationDetails: + trackedCall.confirmationDetails as ToolCallConfirmationDetails, }; case 'executing': return { diff --git a/packages/core/src/confirmation-bus/types.ts b/packages/core/src/confirmation-bus/types.ts index a26b786701..d9caa9c5c2 100644 --- a/packages/core/src/confirmation-bus/types.ts +++ b/packages/core/src/confirmation-bus/types.ts @@ -5,6 +5,11 @@ */ import { type FunctionCall } from '@google/genai'; +import type { + ToolConfirmationOutcome, + ToolConfirmationPayload, +} from '../tools/tools.js'; +import type { ToolCall } from '../scheduler/types.js'; export enum MessageBusType { TOOL_CONFIRMATION_REQUEST = 'tool-confirmation-request', @@ -16,6 +21,12 @@ export enum MessageBusType { HOOK_EXECUTION_REQUEST = 'hook-execution-request', HOOK_EXECUTION_RESPONSE = 'hook-execution-response', HOOK_POLICY_DECISION = 'hook-policy-decision', + TOOL_CALLS_UPDATE = 'tool-calls-update', +} + +export interface ToolCallsUpdateMessage { + type: MessageBusType.TOOL_CALLS_UPDATE; + toolCalls: ToolCall[]; } export interface ToolConfirmationRequest { @@ -23,12 +34,26 @@ export interface ToolConfirmationRequest { toolCall: FunctionCall; correlationId: string; serverName?: string; + /** + * Optional rich details for the confirmation UI (diffs, counts, etc.) + */ + details?: SerializableConfirmationDetails; } export interface ToolConfirmationResponse { type: MessageBusType.TOOL_CONFIRMATION_RESPONSE; correlationId: string; confirmed: boolean; + /** + * The specific outcome selected by the user. + * + * TODO: Make required after migration. + */ + outcome?: ToolConfirmationOutcome; + /** + * Optional payload (e.g., modified content for 'modify_with_editor'). + */ + payload?: ToolConfirmationPayload; /** * When true, indicates that policy decision was ASK_USER and the tool should * show its legacy confirmation UI instead of auto-proceeding. @@ -36,6 +61,35 @@ export interface ToolConfirmationResponse { requiresUserConfirmation?: boolean; } +/** + * Data-only versions of ToolCallConfirmationDetails for bus transmission. + */ +export type SerializableConfirmationDetails = + | { type: 'info'; title: string; prompt: string; urls?: string[] } + | { + type: 'edit'; + title: string; + fileName: string; + filePath: string; + fileDiff: string; + originalContent: string | null; + newContent: string; + } + | { + type: 'exec'; + title: string; + command: string; + rootCommand: string; + rootCommands: string[]; + } + | { + type: 'mcp'; + title: string; + serverName: string; + toolName: string; + toolDisplayName: string; + }; + export interface UpdatePolicy { type: MessageBusType.UPDATE_POLICY; toolName: string; @@ -94,4 +148,5 @@ export type Message = | UpdatePolicy | HookExecutionRequest | HookExecutionResponse - | HookPolicyDecision; + | HookPolicyDecision + | ToolCallsUpdateMessage; diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index 90b8ea7938..1497f7ac02 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -547,9 +547,9 @@ describe('CoreToolScheduler', () => { )) as WaitingToolCall; // Cancel the first tool via its confirmation handler - await awaitingCall.confirmationDetails.onConfirm( - ToolConfirmationOutcome.Cancel, - ); + const confirmationDetails = + awaitingCall.confirmationDetails as ToolCallConfirmationDetails; + await confirmationDetails.onConfirm(ToolConfirmationOutcome.Cancel); abortController.abort(); // User cancelling often involves an abort signal await vi.waitFor(() => { @@ -749,7 +749,7 @@ describe('CoreToolScheduler with payload', () => { if (confirmationDetails) { const payload: ToolConfirmationPayload = { newContent: 'final version' }; - await confirmationDetails.onConfirm( + await (confirmationDetails as ToolCallConfirmationDetails).onConfirm( ToolConfirmationOutcome.ProceedOnce, payload, ); @@ -762,9 +762,9 @@ describe('CoreToolScheduler with payload', () => { )) as WaitingToolCall; // Now confirm for real to execute. - await updatedAwaitingCall.confirmationDetails.onConfirm( - ToolConfirmationOutcome.ProceedOnce, - ); + await ( + updatedAwaitingCall.confirmationDetails as ToolCallConfirmationDetails + ).onConfirm(ToolConfirmationOutcome.ProceedOnce); // Wait for the tool execution to complete await vi.waitFor(() => { @@ -897,7 +897,9 @@ describe('CoreToolScheduler edit cancellation', () => { // Cancel the edit const confirmationDetails = awaitingCall.confirmationDetails; if (confirmationDetails) { - await confirmationDetails.onConfirm(ToolConfirmationOutcome.Cancel); + await (confirmationDetails as ToolCallConfirmationDetails).onConfirm( + ToolConfirmationOutcome.Cancel, + ); } expect(onAllToolCallsComplete).toHaveBeenCalled(); @@ -1447,14 +1449,14 @@ describe('CoreToolScheduler request queueing', () => { toolCalls.forEach((call) => { if (call.status === 'awaiting_approval') { const waitingCall = call; - if (waitingCall.confirmationDetails?.onConfirm) { + const details = + waitingCall.confirmationDetails as ToolCallConfirmationDetails; + if (details?.onConfirm) { const originalHandler = pendingConfirmations.find( - (h) => h === waitingCall.confirmationDetails.onConfirm, + (h) => h === details.onConfirm, ); if (!originalHandler) { - pendingConfirmations.push( - waitingCall.confirmationDetails.onConfirm, - ); + pendingConfirmations.push(details.onConfirm); } } } diff --git a/packages/core/src/scheduler/types.ts b/packages/core/src/scheduler/types.ts index 3a43a47704..2f2baf77e3 100644 --- a/packages/core/src/scheduler/types.ts +++ b/packages/core/src/scheduler/types.ts @@ -14,6 +14,7 @@ import type { } from '../tools/tools.js'; import type { AnsiOutput } from '../utils/terminalSerializer.js'; import type { ToolErrorType } from '../tools/tool-error.js'; +import type { SerializableConfirmationDetails } from '../confirmation-bus/types.js'; export interface ToolCallRequestInfo { callId: string; @@ -98,7 +99,18 @@ export type WaitingToolCall = { request: ToolCallRequestInfo; tool: AnyDeclarativeTool; invocation: AnyToolInvocation; - confirmationDetails: ToolCallConfirmationDetails; + /** + * Supports both legacy (with callbacks) and new (serializable) details. + * New code should treat this as SerializableConfirmationDetails. + * + * TODO: Remove ToolCallConfirmationDetails and collapse to just + * SerializableConfirmationDetails after migration. + */ + confirmationDetails: + | ToolCallConfirmationDetails + | SerializableConfirmationDetails; + // TODO: Make required after migration. + correlationId?: string; startTime?: number; outcome?: ToolConfirmationOutcome; };