feat(scheduler): add types needed for event driven scheduler (#16641)

This commit is contained in:
Abhi
2026-01-14 17:22:44 -05:00
committed by GitHub
parent 7e6817da5b
commit 6021e4c3ba
6 changed files with 98 additions and 22 deletions

View File

@@ -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);
}
});

View File

@@ -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}:`,

View File

@@ -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 {

View File

@@ -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;

View File

@@ -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);
}
}
}

View File

@@ -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;
};