mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-21 02:24:09 -07:00
refactor(cli): consolidate useToolScheduler and delete legacy implementation (#18567)
This commit is contained in:
Abhi
@@ -18,6 +18,7 @@ export const createMockConfig = (overrides: Partial<Config> = {}): Config =>
|
||||
getSandbox: vi.fn(() => undefined),
|
||||
getQuestion: vi.fn(() => ''),
|
||||
isInteractive: vi.fn(() => false),
|
||||
isInitialized: vi.fn(() => true),
|
||||
setTerminalBackground: vi.fn(),
|
||||
storage: {
|
||||
getProjectTempDir: vi.fn().mockReturnValue('/tmp/gemini-test'),
|
||||
|
||||
@@ -220,10 +220,6 @@ describe('App', () => {
|
||||
} as UIState;
|
||||
|
||||
const configWithExperiment = makeFakeConfig();
|
||||
vi.spyOn(
|
||||
configWithExperiment,
|
||||
'isEventDrivenSchedulerEnabled',
|
||||
).mockReturnValue(true);
|
||||
vi.spyOn(configWithExperiment, 'isTrustedFolder').mockReturnValue(true);
|
||||
vi.spyOn(configWithExperiment, 'getIdeMode').mockReturnValue(false);
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ import { cleanup } from 'ink-testing-library';
|
||||
import { act, useContext, type ReactElement } from 'react';
|
||||
import { AppContainer } from './AppContainer.js';
|
||||
import { SettingsContext } from './contexts/SettingsContext.js';
|
||||
import { type TrackedToolCall } from './hooks/useReactToolScheduler.js';
|
||||
import { type TrackedToolCall } from './hooks/useToolScheduler.js';
|
||||
import {
|
||||
type Config,
|
||||
makeFakeConfig,
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
|
||||
|
||||
exports[`useReactToolScheduler > should handle live output updates 1`] = `
|
||||
{
|
||||
"callId": "liveCall",
|
||||
"contentLength": 12,
|
||||
"data": undefined,
|
||||
"error": undefined,
|
||||
"errorType": undefined,
|
||||
"outputFile": undefined,
|
||||
"responseParts": [
|
||||
{
|
||||
"functionResponse": {
|
||||
"id": "liveCall",
|
||||
"name": "mockToolWithLiveOutput",
|
||||
"response": {
|
||||
"output": "Final output",
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
"resultDisplay": "Final display",
|
||||
}
|
||||
`;
|
||||
|
||||
exports[`useReactToolScheduler > should handle tool requiring confirmation - approved 1`] = `
|
||||
{
|
||||
"callId": "callConfirm",
|
||||
"contentLength": 16,
|
||||
"data": undefined,
|
||||
"error": undefined,
|
||||
"errorType": undefined,
|
||||
"outputFile": undefined,
|
||||
"responseParts": [
|
||||
{
|
||||
"functionResponse": {
|
||||
"id": "callConfirm",
|
||||
"name": "mockToolRequiresConfirmation",
|
||||
"response": {
|
||||
"output": "Confirmed output",
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
"resultDisplay": "Confirmed display",
|
||||
}
|
||||
`;
|
||||
|
||||
exports[`useReactToolScheduler > should handle tool requiring confirmation - cancelled by user 1`] = `
|
||||
{
|
||||
"callId": "callConfirmCancel",
|
||||
"contentLength": 59,
|
||||
"error": undefined,
|
||||
"errorType": undefined,
|
||||
"responseParts": [
|
||||
{
|
||||
"functionResponse": {
|
||||
"id": "callConfirmCancel",
|
||||
"name": "mockToolRequiresConfirmation",
|
||||
"response": {
|
||||
"error": "[Operation Cancelled] Reason: User cancelled the operation.",
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
"resultDisplay": {
|
||||
"fileDiff": "Mock tool requires confirmation",
|
||||
"fileName": "mockToolRequiresConfirmation.ts",
|
||||
"filePath": undefined,
|
||||
"newContent": undefined,
|
||||
"originalContent": undefined,
|
||||
},
|
||||
}
|
||||
`;
|
||||
|
||||
exports[`useReactToolScheduler > should schedule and execute a tool call successfully 1`] = `
|
||||
{
|
||||
"callId": "call1",
|
||||
"contentLength": 11,
|
||||
"data": undefined,
|
||||
"error": undefined,
|
||||
"errorType": undefined,
|
||||
"outputFile": undefined,
|
||||
"responseParts": [
|
||||
{
|
||||
"functionResponse": {
|
||||
"id": "call1",
|
||||
"name": "mockTool",
|
||||
"response": {
|
||||
"output": "Tool output",
|
||||
},
|
||||
},
|
||||
},
|
||||
],
|
||||
"resultDisplay": "Formatted tool output",
|
||||
}
|
||||
`;
|
||||
@@ -246,7 +246,6 @@ describe('useGeminiStream', () => {
|
||||
getContentGenerator: vi.fn(),
|
||||
isInteractive: () => false,
|
||||
getExperiments: () => {},
|
||||
isEventDrivenSchedulerEnabled: vi.fn(() => false),
|
||||
getMaxSessionTurns: vi.fn(() => 100),
|
||||
isJitContextEnabled: vi.fn(() => false),
|
||||
getGlobalMemory: vi.fn(() => ''),
|
||||
|
||||
@@ -1,77 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { CoreToolScheduler } from '@google/gemini-cli-core';
|
||||
import type { Config } from '@google/gemini-cli-core';
|
||||
import { renderHook } from '../../test-utils/render.js';
|
||||
import { vi, describe, it, expect, beforeEach } from 'vitest';
|
||||
import { useReactToolScheduler } from './useReactToolScheduler.js';
|
||||
|
||||
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
|
||||
const actual =
|
||||
await importOriginal<typeof import('@google/gemini-cli-core')>();
|
||||
return {
|
||||
...actual,
|
||||
CoreToolScheduler: vi.fn(),
|
||||
};
|
||||
});
|
||||
|
||||
const mockCoreToolScheduler = vi.mocked(CoreToolScheduler);
|
||||
|
||||
describe('useReactToolScheduler', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('only creates one instance of CoreToolScheduler even if props change', () => {
|
||||
const onComplete = vi.fn();
|
||||
const getPreferredEditor = vi.fn();
|
||||
const config = {} as Config;
|
||||
|
||||
const { rerender } = renderHook(
|
||||
(props) =>
|
||||
useReactToolScheduler(
|
||||
props.onComplete,
|
||||
props.config,
|
||||
props.getPreferredEditor,
|
||||
),
|
||||
{
|
||||
initialProps: {
|
||||
onComplete,
|
||||
config,
|
||||
getPreferredEditor,
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Rerender with a new onComplete function
|
||||
const newOnComplete = vi.fn();
|
||||
rerender({
|
||||
onComplete: newOnComplete,
|
||||
config,
|
||||
getPreferredEditor,
|
||||
});
|
||||
expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1);
|
||||
|
||||
// Rerender with a new getPreferredEditor function
|
||||
const newGetPreferredEditor = vi.fn();
|
||||
rerender({
|
||||
onComplete: newOnComplete,
|
||||
config,
|
||||
getPreferredEditor: newGetPreferredEditor,
|
||||
});
|
||||
expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1);
|
||||
|
||||
rerender({
|
||||
onComplete: newOnComplete,
|
||||
config,
|
||||
getPreferredEditor: newGetPreferredEditor,
|
||||
});
|
||||
expect(mockCoreToolScheduler).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
@@ -1,221 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type {
|
||||
Config,
|
||||
ToolCallRequestInfo,
|
||||
OutputUpdateHandler,
|
||||
AllToolCallsCompleteHandler,
|
||||
ToolCallsUpdateHandler,
|
||||
ToolCall,
|
||||
EditorType,
|
||||
CompletedToolCall,
|
||||
ExecutingToolCall,
|
||||
ScheduledToolCall,
|
||||
ValidatingToolCall,
|
||||
WaitingToolCall,
|
||||
CancelledToolCall,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { CoreToolScheduler } from '@google/gemini-cli-core';
|
||||
import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
|
||||
|
||||
export type ScheduleFn = (
|
||||
request: ToolCallRequestInfo | ToolCallRequestInfo[],
|
||||
signal: AbortSignal,
|
||||
) => Promise<void>;
|
||||
export type MarkToolsAsSubmittedFn = (callIds: string[]) => void;
|
||||
export type CancelAllFn = (signal: AbortSignal) => void;
|
||||
|
||||
export type TrackedScheduledToolCall = ScheduledToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
export type TrackedValidatingToolCall = ValidatingToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
export type TrackedWaitingToolCall = WaitingToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
export type TrackedExecutingToolCall = ExecutingToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
export type TrackedCompletedToolCall = CompletedToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
export type TrackedCancelledToolCall = CancelledToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
|
||||
export type TrackedToolCall =
|
||||
| TrackedScheduledToolCall
|
||||
| TrackedValidatingToolCall
|
||||
| TrackedWaitingToolCall
|
||||
| TrackedExecutingToolCall
|
||||
| TrackedCompletedToolCall
|
||||
| TrackedCancelledToolCall;
|
||||
|
||||
/**
|
||||
* Legacy scheduler implementation based on CoreToolScheduler callbacks.
|
||||
*
|
||||
* This is currently the default implementation used by useGeminiStream.
|
||||
* It will be phased out once the event-driven scheduler migration is complete.
|
||||
*/
|
||||
export function useReactToolScheduler(
|
||||
onComplete: (tools: CompletedToolCall[]) => Promise<void>,
|
||||
config: Config,
|
||||
getPreferredEditor: () => EditorType | undefined,
|
||||
): [
|
||||
TrackedToolCall[],
|
||||
ScheduleFn,
|
||||
MarkToolsAsSubmittedFn,
|
||||
React.Dispatch<React.SetStateAction<TrackedToolCall[]>>,
|
||||
CancelAllFn,
|
||||
number,
|
||||
] {
|
||||
const [toolCallsForDisplay, setToolCallsForDisplay] = useState<
|
||||
TrackedToolCall[]
|
||||
>([]);
|
||||
const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0);
|
||||
|
||||
const onCompleteRef = useRef(onComplete);
|
||||
const getPreferredEditorRef = useRef(getPreferredEditor);
|
||||
|
||||
useEffect(() => {
|
||||
onCompleteRef.current = onComplete;
|
||||
}, [onComplete]);
|
||||
|
||||
useEffect(() => {
|
||||
getPreferredEditorRef.current = getPreferredEditor;
|
||||
}, [getPreferredEditor]);
|
||||
|
||||
const outputUpdateHandler: OutputUpdateHandler = useCallback(
|
||||
(toolCallId, outputChunk) => {
|
||||
setLastToolOutputTime(Date.now());
|
||||
setToolCallsForDisplay((prevCalls) =>
|
||||
prevCalls.map((tc) => {
|
||||
if (tc.request.callId === toolCallId && tc.status === 'executing') {
|
||||
const executingTc = tc;
|
||||
return { ...executingTc, liveOutput: outputChunk };
|
||||
}
|
||||
return tc;
|
||||
}),
|
||||
);
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
const allToolCallsCompleteHandler: AllToolCallsCompleteHandler = useCallback(
|
||||
async (completedToolCalls) => {
|
||||
await onCompleteRef.current(completedToolCalls);
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
const toolCallsUpdateHandler: ToolCallsUpdateHandler = useCallback(
|
||||
(allCoreToolCalls: ToolCall[]) => {
|
||||
setToolCallsForDisplay((prevTrackedCalls) => {
|
||||
const prevCallsMap = new Map(
|
||||
prevTrackedCalls.map((c) => [c.request.callId, c]),
|
||||
);
|
||||
|
||||
return allCoreToolCalls.map((coreTc): TrackedToolCall => {
|
||||
const existingTrackedCall = prevCallsMap.get(coreTc.request.callId);
|
||||
|
||||
const responseSubmittedToGemini =
|
||||
existingTrackedCall?.responseSubmittedToGemini ?? false;
|
||||
|
||||
if (coreTc.status === 'executing') {
|
||||
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
|
||||
const liveOutput = (existingTrackedCall as TrackedExecutingToolCall)
|
||||
?.liveOutput;
|
||||
return {
|
||||
...coreTc,
|
||||
responseSubmittedToGemini,
|
||||
liveOutput,
|
||||
};
|
||||
} else if (
|
||||
coreTc.status === 'success' ||
|
||||
coreTc.status === 'error' ||
|
||||
coreTc.status === 'cancelled'
|
||||
) {
|
||||
return {
|
||||
...coreTc,
|
||||
responseSubmittedToGemini,
|
||||
};
|
||||
} else {
|
||||
return {
|
||||
...coreTc,
|
||||
responseSubmittedToGemini,
|
||||
};
|
||||
}
|
||||
});
|
||||
});
|
||||
},
|
||||
[setToolCallsForDisplay],
|
||||
);
|
||||
|
||||
const stableGetPreferredEditor = useCallback(
|
||||
() => getPreferredEditorRef.current(),
|
||||
[],
|
||||
);
|
||||
|
||||
const scheduler = useMemo(
|
||||
() =>
|
||||
new CoreToolScheduler({
|
||||
outputUpdateHandler,
|
||||
onAllToolCallsComplete: allToolCallsCompleteHandler,
|
||||
onToolCallsUpdate: toolCallsUpdateHandler,
|
||||
getPreferredEditor: stableGetPreferredEditor,
|
||||
config,
|
||||
}),
|
||||
[
|
||||
config,
|
||||
outputUpdateHandler,
|
||||
allToolCallsCompleteHandler,
|
||||
toolCallsUpdateHandler,
|
||||
stableGetPreferredEditor,
|
||||
],
|
||||
);
|
||||
|
||||
const schedule: ScheduleFn = useCallback(
|
||||
(
|
||||
request: ToolCallRequestInfo | ToolCallRequestInfo[],
|
||||
signal: AbortSignal,
|
||||
) => {
|
||||
setToolCallsForDisplay([]);
|
||||
return scheduler.schedule(request, signal);
|
||||
},
|
||||
[scheduler, setToolCallsForDisplay],
|
||||
);
|
||||
|
||||
const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback(
|
||||
(callIdsToMark: string[]) => {
|
||||
setToolCallsForDisplay((prevCalls) =>
|
||||
prevCalls.map((tc) =>
|
||||
callIdsToMark.includes(tc.request.callId)
|
||||
? { ...tc, responseSubmittedToGemini: true }
|
||||
: tc,
|
||||
),
|
||||
);
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
const cancelAllToolCalls = useCallback(
|
||||
(signal: AbortSignal) => {
|
||||
scheduler.cancelAll(signal);
|
||||
},
|
||||
[scheduler],
|
||||
);
|
||||
|
||||
return [
|
||||
toolCallsForDisplay,
|
||||
schedule,
|
||||
markToolsAsSubmitted,
|
||||
setToolCallsForDisplay,
|
||||
cancelAllToolCalls,
|
||||
lastToolOutputTime,
|
||||
];
|
||||
}
|
||||
@@ -12,7 +12,7 @@ import {
|
||||
SHELL_SILENT_WORKING_TITLE_DELAY_MS,
|
||||
} from '../constants.js';
|
||||
import type { StreamingState } from '../types.js';
|
||||
import { type TrackedToolCall } from './useReactToolScheduler.js';
|
||||
import { type TrackedToolCall } from './useToolScheduler.js';
|
||||
|
||||
interface ShellInactivityStatusProps {
|
||||
activePtyId: number | string | null | undefined;
|
||||
|
||||
@@ -1,525 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest';
|
||||
import { act } from 'react';
|
||||
import { renderHook } from '../../test-utils/render.js';
|
||||
import { useToolExecutionScheduler } from './useToolExecutionScheduler.js';
|
||||
import {
|
||||
MessageBusType,
|
||||
ToolConfirmationOutcome,
|
||||
Scheduler,
|
||||
type Config,
|
||||
type MessageBus,
|
||||
type CompletedToolCall,
|
||||
type ToolCallConfirmationDetails,
|
||||
type ToolCallsUpdateMessage,
|
||||
type AnyDeclarativeTool,
|
||||
type AnyToolInvocation,
|
||||
ROOT_SCHEDULER_ID,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js';
|
||||
|
||||
// Mock Core Scheduler
|
||||
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
|
||||
const actual =
|
||||
await importOriginal<typeof import('@google/gemini-cli-core')>();
|
||||
return {
|
||||
...actual,
|
||||
Scheduler: vi.fn().mockImplementation(() => ({
|
||||
schedule: vi.fn().mockResolvedValue([]),
|
||||
cancelAll: vi.fn(),
|
||||
})),
|
||||
};
|
||||
});
|
||||
|
||||
const createMockTool = (
|
||||
overrides: Partial<AnyDeclarativeTool> = {},
|
||||
): AnyDeclarativeTool =>
|
||||
({
|
||||
name: 'test_tool',
|
||||
displayName: 'Test Tool',
|
||||
description: 'A test tool',
|
||||
kind: 'function',
|
||||
parameterSchema: {},
|
||||
isOutputMarkdown: false,
|
||||
build: vi.fn(),
|
||||
...overrides,
|
||||
}) as AnyDeclarativeTool;
|
||||
|
||||
const createMockInvocation = (
|
||||
overrides: Partial<AnyToolInvocation> = {},
|
||||
): AnyToolInvocation =>
|
||||
({
|
||||
getDescription: () => 'Executing test tool',
|
||||
shouldConfirmExecute: vi.fn(),
|
||||
execute: vi.fn(),
|
||||
params: {},
|
||||
toolLocations: [],
|
||||
...overrides,
|
||||
}) as AnyToolInvocation;
|
||||
|
||||
describe('useToolExecutionScheduler', () => {
|
||||
let mockConfig: Config;
|
||||
let mockMessageBus: MessageBus;
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockMessageBus = createMockMessageBus() as unknown as MessageBus;
|
||||
mockConfig = {
|
||||
getMessageBus: () => mockMessageBus,
|
||||
} as unknown as Config;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('initializes with empty tool calls', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
const [toolCalls] = result.current;
|
||||
expect(toolCalls).toEqual([]);
|
||||
});
|
||||
|
||||
it('updates tool calls when MessageBus emits TOOL_CALLS_UPDATE', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
|
||||
const mockToolCall = {
|
||||
status: 'executing' as const,
|
||||
request: {
|
||||
callId: 'call-1',
|
||||
name: 'test_tool',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
tool: createMockTool(),
|
||||
invocation: createMockInvocation(),
|
||||
liveOutput: 'Loading...',
|
||||
};
|
||||
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [mockToolCall],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
const [toolCalls] = result.current;
|
||||
expect(toolCalls).toHaveLength(1);
|
||||
// Expect Core Object structure, not Display Object
|
||||
expect(toolCalls[0]).toMatchObject({
|
||||
request: { callId: 'call-1', name: 'test_tool' },
|
||||
status: 'executing', // Core status
|
||||
liveOutput: 'Loading...',
|
||||
responseSubmittedToGemini: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('injects onConfirm callback for awaiting_approval tools (Adapter Pattern)', async () => {
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
|
||||
const mockToolCall = {
|
||||
status: 'awaiting_approval' as const,
|
||||
request: {
|
||||
callId: 'call-1',
|
||||
name: 'test_tool',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
tool: createMockTool(),
|
||||
invocation: createMockInvocation({
|
||||
getDescription: () => 'Confirming test tool',
|
||||
}),
|
||||
confirmationDetails: { type: 'info', title: 'Confirm', prompt: 'Sure?' },
|
||||
correlationId: 'corr-123',
|
||||
};
|
||||
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [mockToolCall],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
const [toolCalls] = result.current;
|
||||
const call = toolCalls[0];
|
||||
if (call.status !== 'awaiting_approval') {
|
||||
throw new Error('Expected status to be awaiting_approval');
|
||||
}
|
||||
const confirmationDetails =
|
||||
call.confirmationDetails as ToolCallConfirmationDetails;
|
||||
|
||||
expect(confirmationDetails).toBeDefined();
|
||||
expect(typeof confirmationDetails.onConfirm).toBe('function');
|
||||
|
||||
// Test that onConfirm publishes to MessageBus
|
||||
const publishSpy = vi.spyOn(mockMessageBus, 'publish');
|
||||
await confirmationDetails.onConfirm(ToolConfirmationOutcome.ProceedOnce);
|
||||
|
||||
expect(publishSpy).toHaveBeenCalledWith({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-123',
|
||||
confirmed: true,
|
||||
requiresUserConfirmation: false,
|
||||
outcome: ToolConfirmationOutcome.ProceedOnce,
|
||||
payload: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it('injects onConfirm with payload (Inline Edit support)', async () => {
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
|
||||
const mockToolCall = {
|
||||
status: 'awaiting_approval' as const,
|
||||
request: {
|
||||
callId: 'call-1',
|
||||
name: 'test_tool',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
tool: createMockTool(),
|
||||
invocation: createMockInvocation(),
|
||||
confirmationDetails: { type: 'edit', title: 'Edit', filePath: 'test.ts' },
|
||||
correlationId: 'corr-edit',
|
||||
};
|
||||
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [mockToolCall],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
const [toolCalls] = result.current;
|
||||
const call = toolCalls[0];
|
||||
if (call.status !== 'awaiting_approval') {
|
||||
throw new Error('Expected awaiting_approval');
|
||||
}
|
||||
const confirmationDetails =
|
||||
call.confirmationDetails as ToolCallConfirmationDetails;
|
||||
|
||||
const publishSpy = vi.spyOn(mockMessageBus, 'publish');
|
||||
const mockPayload = { newContent: 'updated code' };
|
||||
await confirmationDetails.onConfirm(
|
||||
ToolConfirmationOutcome.ProceedOnce,
|
||||
mockPayload,
|
||||
);
|
||||
|
||||
expect(publishSpy).toHaveBeenCalledWith({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId: 'corr-edit',
|
||||
confirmed: true,
|
||||
requiresUserConfirmation: false,
|
||||
outcome: ToolConfirmationOutcome.ProceedOnce,
|
||||
payload: mockPayload,
|
||||
});
|
||||
});
|
||||
|
||||
it('preserves responseSubmittedToGemini flag across updates', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
|
||||
const mockToolCall = {
|
||||
status: 'success' as const,
|
||||
request: {
|
||||
callId: 'call-1',
|
||||
name: 'test',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
tool: createMockTool(),
|
||||
invocation: createMockInvocation(),
|
||||
response: {
|
||||
callId: 'call-1',
|
||||
resultDisplay: 'OK',
|
||||
responseParts: [],
|
||||
error: undefined,
|
||||
errorType: undefined,
|
||||
},
|
||||
};
|
||||
|
||||
// 1. Initial success
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [mockToolCall],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
// 2. Mark as submitted
|
||||
act(() => {
|
||||
const [, , markAsSubmitted] = result.current;
|
||||
markAsSubmitted(['call-1']);
|
||||
});
|
||||
|
||||
expect(result.current[0][0].responseSubmittedToGemini).toBe(true);
|
||||
|
||||
// 3. Receive another update (should preserve the true flag)
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [mockToolCall],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
expect(result.current[0][0].responseSubmittedToGemini).toBe(true);
|
||||
});
|
||||
|
||||
it('updates lastToolOutputTime when tools are executing', () => {
|
||||
vi.useFakeTimers();
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
|
||||
const startTime = Date.now();
|
||||
vi.advanceTimersByTime(1000);
|
||||
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [
|
||||
{
|
||||
status: 'executing' as const,
|
||||
request: {
|
||||
callId: 'call-1',
|
||||
name: 'test',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
tool: createMockTool(),
|
||||
invocation: createMockInvocation(),
|
||||
},
|
||||
],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
const [, , , , , lastOutputTime] = result.current;
|
||||
expect(lastOutputTime).toBeGreaterThan(startTime);
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it('delegates cancelAll to the Core Scheduler', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
|
||||
const [, , , , cancelAll] = result.current;
|
||||
const signal = new AbortController().signal;
|
||||
|
||||
// We need to find the mock instance of Scheduler
|
||||
// Since we used vi.mock at top level, we can get it from vi.mocked(Scheduler)
|
||||
const schedulerInstance = vi.mocked(Scheduler).mock.results[0].value;
|
||||
|
||||
cancelAll(signal);
|
||||
|
||||
expect(schedulerInstance.cancelAll).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('resolves the schedule promise when scheduler resolves', async () => {
|
||||
const onComplete = vi.fn().mockResolvedValue(undefined);
|
||||
|
||||
const completedToolCall = {
|
||||
status: 'success' as const,
|
||||
request: {
|
||||
callId: 'call-1',
|
||||
name: 'test',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
tool: createMockTool(),
|
||||
invocation: createMockInvocation(),
|
||||
response: {
|
||||
callId: 'call-1',
|
||||
responseParts: [],
|
||||
resultDisplay: 'Success',
|
||||
error: undefined,
|
||||
errorType: undefined,
|
||||
},
|
||||
};
|
||||
|
||||
// Mock the specific return value for this test
|
||||
const { Scheduler } = await import('@google/gemini-cli-core');
|
||||
vi.mocked(Scheduler).mockImplementation(
|
||||
() =>
|
||||
({
|
||||
schedule: vi.fn().mockResolvedValue([completedToolCall]),
|
||||
cancelAll: vi.fn(),
|
||||
}) as unknown as Scheduler,
|
||||
);
|
||||
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(onComplete, mockConfig, () => undefined),
|
||||
);
|
||||
|
||||
const [, schedule] = result.current;
|
||||
const signal = new AbortController().signal;
|
||||
|
||||
let completedResult: CompletedToolCall[] = [];
|
||||
await act(async () => {
|
||||
completedResult = await schedule(
|
||||
{
|
||||
callId: 'call-1',
|
||||
name: 'test',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
signal,
|
||||
);
|
||||
});
|
||||
|
||||
expect(completedResult).toEqual([completedToolCall]);
|
||||
expect(onComplete).toHaveBeenCalledWith([completedToolCall]);
|
||||
});
|
||||
|
||||
it('setToolCallsForDisplay re-groups tools by schedulerId (Multi-Scheduler support)', () => {
|
||||
const { result } = renderHook(() =>
|
||||
useToolExecutionScheduler(
|
||||
vi.fn().mockResolvedValue(undefined),
|
||||
mockConfig,
|
||||
() => undefined,
|
||||
),
|
||||
);
|
||||
|
||||
const callRoot = {
|
||||
status: 'success' as const,
|
||||
request: {
|
||||
callId: 'call-root',
|
||||
name: 'test',
|
||||
args: {},
|
||||
isClientInitiated: false,
|
||||
prompt_id: 'p1',
|
||||
},
|
||||
tool: createMockTool(),
|
||||
invocation: createMockInvocation(),
|
||||
response: {
|
||||
callId: 'call-root',
|
||||
responseParts: [],
|
||||
resultDisplay: 'OK',
|
||||
error: undefined,
|
||||
errorType: undefined,
|
||||
},
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
};
|
||||
|
||||
const callSub = {
|
||||
...callRoot,
|
||||
request: { ...callRoot.request, callId: 'call-sub' },
|
||||
schedulerId: 'subagent-1',
|
||||
};
|
||||
|
||||
// 1. Populate state with multiple schedulers
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [callRoot],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [callSub],
|
||||
schedulerId: 'subagent-1',
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
let [toolCalls] = result.current;
|
||||
expect(toolCalls).toHaveLength(2);
|
||||
expect(
|
||||
toolCalls.find((t) => t.request.callId === 'call-root')?.schedulerId,
|
||||
).toBe(ROOT_SCHEDULER_ID);
|
||||
expect(
|
||||
toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
|
||||
).toBe('subagent-1');
|
||||
|
||||
// 2. Call setToolCallsForDisplay (e.g., simulate a manual update or clear)
|
||||
act(() => {
|
||||
const [, , , setToolCalls] = result.current;
|
||||
setToolCalls((prev) =>
|
||||
prev.map((t) => ({ ...t, responseSubmittedToGemini: true })),
|
||||
);
|
||||
});
|
||||
|
||||
// 3. Verify that tools are still present and maintain their scheduler IDs
|
||||
// The internal map should have been re-grouped.
|
||||
[toolCalls] = result.current;
|
||||
expect(toolCalls).toHaveLength(2);
|
||||
expect(toolCalls.every((t) => t.responseSubmittedToGemini)).toBe(true);
|
||||
|
||||
const updatedRoot = toolCalls.find((t) => t.request.callId === 'call-root');
|
||||
const updatedSub = toolCalls.find((t) => t.request.callId === 'call-sub');
|
||||
|
||||
expect(updatedRoot?.schedulerId).toBe(ROOT_SCHEDULER_ID);
|
||||
expect(updatedSub?.schedulerId).toBe('subagent-1');
|
||||
|
||||
// 4. Verify that a subsequent update to ONE scheduler doesn't wipe the other
|
||||
act(() => {
|
||||
void mockMessageBus.publish({
|
||||
type: MessageBusType.TOOL_CALLS_UPDATE,
|
||||
toolCalls: [{ ...callRoot, status: 'executing' }],
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
} as ToolCallsUpdateMessage);
|
||||
});
|
||||
|
||||
[toolCalls] = result.current;
|
||||
expect(toolCalls).toHaveLength(2);
|
||||
expect(
|
||||
toolCalls.find((t) => t.request.callId === 'call-root')?.status,
|
||||
).toBe('executing');
|
||||
expect(
|
||||
toolCalls.find((t) => t.request.callId === 'call-sub')?.schedulerId,
|
||||
).toBe('subagent-1');
|
||||
});
|
||||
});
|
||||
@@ -1,253 +0,0 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2026 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
type Config,
|
||||
type MessageBus,
|
||||
type ToolCallRequestInfo,
|
||||
type ToolCall,
|
||||
type CompletedToolCall,
|
||||
type ToolConfirmationPayload,
|
||||
MessageBusType,
|
||||
ToolConfirmationOutcome,
|
||||
Scheduler,
|
||||
type EditorType,
|
||||
type ToolCallsUpdateMessage,
|
||||
ROOT_SCHEDULER_ID,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
|
||||
|
||||
// Re-exporting types compatible with legacy hook expectations
|
||||
export type ScheduleFn = (
|
||||
request: ToolCallRequestInfo | ToolCallRequestInfo[],
|
||||
signal: AbortSignal,
|
||||
) => Promise<CompletedToolCall[]>;
|
||||
|
||||
export type MarkToolsAsSubmittedFn = (callIds: string[]) => void;
|
||||
export type CancelAllFn = (signal: AbortSignal) => void;
|
||||
|
||||
/**
|
||||
* The shape expected by useGeminiStream.
|
||||
* It matches the Core ToolCall structure + the UI metadata flag.
|
||||
*/
|
||||
export type TrackedToolCall = ToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* Modern tool scheduler hook using the event-driven Core Scheduler.
|
||||
*
|
||||
* This hook acts as an Adapter between the new MessageBus-driven Core
|
||||
* and the legacy callback-based UI components.
|
||||
*/
|
||||
export function useToolExecutionScheduler(
|
||||
onComplete: (tools: CompletedToolCall[]) => Promise<void>,
|
||||
config: Config,
|
||||
getPreferredEditor: () => EditorType | undefined,
|
||||
): [
|
||||
TrackedToolCall[],
|
||||
ScheduleFn,
|
||||
MarkToolsAsSubmittedFn,
|
||||
React.Dispatch<React.SetStateAction<TrackedToolCall[]>>,
|
||||
CancelAllFn,
|
||||
number,
|
||||
] {
|
||||
// State stores tool calls organized by their originating schedulerId
|
||||
const [toolCallsMap, setToolCallsMap] = useState<
|
||||
Record<string, TrackedToolCall[]>
|
||||
>({});
|
||||
const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0);
|
||||
|
||||
const messageBus = useMemo(() => config.getMessageBus(), [config]);
|
||||
|
||||
const onCompleteRef = useRef(onComplete);
|
||||
useEffect(() => {
|
||||
onCompleteRef.current = onComplete;
|
||||
}, [onComplete]);
|
||||
|
||||
const getPreferredEditorRef = useRef(getPreferredEditor);
|
||||
useEffect(() => {
|
||||
getPreferredEditorRef.current = getPreferredEditor;
|
||||
}, [getPreferredEditor]);
|
||||
|
||||
const scheduler = useMemo(
|
||||
() =>
|
||||
new Scheduler({
|
||||
config,
|
||||
messageBus,
|
||||
getPreferredEditor: () => getPreferredEditorRef.current(),
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
}),
|
||||
[config, messageBus],
|
||||
);
|
||||
|
||||
const internalAdaptToolCalls = useCallback(
|
||||
(coreCalls: ToolCall[], prevTracked: TrackedToolCall[]) =>
|
||||
adaptToolCalls(coreCalls, prevTracked, messageBus),
|
||||
[messageBus],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
const handler = (event: ToolCallsUpdateMessage) => {
|
||||
// Update output timer for UI spinners (Side Effect)
|
||||
if (event.toolCalls.some((tc) => tc.status === 'executing')) {
|
||||
setLastToolOutputTime(Date.now());
|
||||
}
|
||||
|
||||
setToolCallsMap((prev) => {
|
||||
const adapted = internalAdaptToolCalls(
|
||||
event.toolCalls,
|
||||
prev[event.schedulerId] ?? [],
|
||||
);
|
||||
|
||||
return {
|
||||
...prev,
|
||||
[event.schedulerId]: adapted,
|
||||
};
|
||||
});
|
||||
};
|
||||
|
||||
messageBus.subscribe(MessageBusType.TOOL_CALLS_UPDATE, handler);
|
||||
return () => {
|
||||
messageBus.unsubscribe(MessageBusType.TOOL_CALLS_UPDATE, handler);
|
||||
};
|
||||
}, [messageBus, internalAdaptToolCalls]);
|
||||
|
||||
const schedule: ScheduleFn = useCallback(
|
||||
async (request, signal) => {
|
||||
// Clear state for new run
|
||||
setToolCallsMap({});
|
||||
|
||||
// 1. Await Core Scheduler directly
|
||||
const results = await scheduler.schedule(request, signal);
|
||||
|
||||
// 2. Trigger legacy reinjection logic (useGeminiStream loop)
|
||||
// Since this hook instance owns the "root" scheduler, we always trigger
|
||||
// onComplete when it finishes its batch.
|
||||
await onCompleteRef.current(results);
|
||||
|
||||
return results;
|
||||
},
|
||||
[scheduler],
|
||||
);
|
||||
|
||||
const cancelAll: CancelAllFn = useCallback(
|
||||
(_signal) => {
|
||||
scheduler.cancelAll();
|
||||
},
|
||||
[scheduler],
|
||||
);
|
||||
|
||||
const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback(
|
||||
(callIdsToMark: string[]) => {
|
||||
setToolCallsMap((prevMap) => {
|
||||
const nextMap = { ...prevMap };
|
||||
for (const [sid, calls] of Object.entries(nextMap)) {
|
||||
nextMap[sid] = calls.map((tc) =>
|
||||
callIdsToMark.includes(tc.request.callId)
|
||||
? { ...tc, responseSubmittedToGemini: true }
|
||||
: tc,
|
||||
);
|
||||
}
|
||||
return nextMap;
|
||||
});
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
// Flatten the map for the UI components that expect a single list of tools.
|
||||
const toolCalls = useMemo(
|
||||
() => Object.values(toolCallsMap).flat(),
|
||||
[toolCallsMap],
|
||||
);
|
||||
|
||||
// Provide a setter that maintains compatibility with legacy [].
|
||||
const setToolCallsForDisplay = useCallback(
|
||||
(action: React.SetStateAction<TrackedToolCall[]>) => {
|
||||
setToolCallsMap((prev) => {
|
||||
const currentFlattened = Object.values(prev).flat();
|
||||
const nextFlattened =
|
||||
typeof action === 'function' ? action(currentFlattened) : action;
|
||||
|
||||
if (nextFlattened.length === 0) {
|
||||
return {};
|
||||
}
|
||||
|
||||
// Re-group by schedulerId to preserve multi-scheduler state
|
||||
const nextMap: Record<string, TrackedToolCall[]> = {};
|
||||
for (const call of nextFlattened) {
|
||||
// All tool calls should have a schedulerId from the core.
|
||||
// Default to ROOT_SCHEDULER_ID as a safeguard.
|
||||
const sid = call.schedulerId ?? ROOT_SCHEDULER_ID;
|
||||
if (!nextMap[sid]) {
|
||||
nextMap[sid] = [];
|
||||
}
|
||||
nextMap[sid].push(call);
|
||||
}
|
||||
return nextMap;
|
||||
});
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
return [
|
||||
toolCalls,
|
||||
schedule,
|
||||
markToolsAsSubmitted,
|
||||
setToolCallsForDisplay,
|
||||
cancelAll,
|
||||
lastToolOutputTime,
|
||||
];
|
||||
}
|
||||
|
||||
/**
|
||||
* ADAPTER: Merges UI metadata (submitted flag) and injects legacy callbacks.
|
||||
*/
|
||||
function adaptToolCalls(
|
||||
coreCalls: ToolCall[],
|
||||
prevTracked: TrackedToolCall[],
|
||||
messageBus: MessageBus,
|
||||
): TrackedToolCall[] {
|
||||
const prevMap = new Map(prevTracked.map((t) => [t.request.callId, t]));
|
||||
|
||||
return coreCalls.map((coreCall): TrackedToolCall => {
|
||||
const prev = prevMap.get(coreCall.request.callId);
|
||||
const responseSubmittedToGemini = prev?.responseSubmittedToGemini ?? false;
|
||||
|
||||
// Inject onConfirm adapter for tools awaiting approval.
|
||||
// The Core provides data-only (serializable) confirmationDetails. We must
|
||||
// inject the legacy callback function that proxies responses back to the
|
||||
// MessageBus.
|
||||
if (coreCall.status === 'awaiting_approval' && coreCall.correlationId) {
|
||||
const correlationId = coreCall.correlationId;
|
||||
return {
|
||||
...coreCall,
|
||||
confirmationDetails: {
|
||||
...coreCall.confirmationDetails,
|
||||
onConfirm: async (
|
||||
outcome: ToolConfirmationOutcome,
|
||||
payload?: ToolConfirmationPayload,
|
||||
) => {
|
||||
await messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId,
|
||||
confirmed: outcome !== ToolConfirmationOutcome.Cancel,
|
||||
requiresUserConfirmation: false,
|
||||
outcome,
|
||||
payload,
|
||||
});
|
||||
},
|
||||
},
|
||||
responseSubmittedToGemini,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
...coreCall,
|
||||
responseSubmittedToGemini,
|
||||
};
|
||||
});
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,67 +4,273 @@
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import type {
|
||||
Config,
|
||||
EditorType,
|
||||
CompletedToolCall,
|
||||
ToolCallRequestInfo,
|
||||
import {
|
||||
type Config,
|
||||
type MessageBus,
|
||||
type ToolCallRequestInfo,
|
||||
type ToolCall,
|
||||
type CompletedToolCall,
|
||||
type ToolConfirmationPayload,
|
||||
MessageBusType,
|
||||
ToolConfirmationOutcome,
|
||||
Scheduler,
|
||||
type EditorType,
|
||||
type ToolCallsUpdateMessage,
|
||||
ROOT_SCHEDULER_ID,
|
||||
} from '@google/gemini-cli-core';
|
||||
import {
|
||||
type TrackedScheduledToolCall,
|
||||
type TrackedValidatingToolCall,
|
||||
type TrackedWaitingToolCall,
|
||||
type TrackedExecutingToolCall,
|
||||
type TrackedCompletedToolCall,
|
||||
type TrackedCancelledToolCall,
|
||||
type MarkToolsAsSubmittedFn,
|
||||
type CancelAllFn,
|
||||
} from './useReactToolScheduler.js';
|
||||
import {
|
||||
useToolExecutionScheduler,
|
||||
type TrackedToolCall,
|
||||
} from './useToolExecutionScheduler.js';
|
||||
import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
|
||||
|
||||
// Re-export specific state types from Legacy, as the structures are compatible
|
||||
// and useGeminiStream relies on them for narrowing.
|
||||
export type {
|
||||
TrackedToolCall,
|
||||
TrackedScheduledToolCall,
|
||||
TrackedValidatingToolCall,
|
||||
TrackedWaitingToolCall,
|
||||
TrackedExecutingToolCall,
|
||||
TrackedCompletedToolCall,
|
||||
TrackedCancelledToolCall,
|
||||
MarkToolsAsSubmittedFn,
|
||||
CancelAllFn,
|
||||
};
|
||||
|
||||
// Unified Schedule function (Promise<void> | Promise<CompletedToolCall[]>)
|
||||
// Re-exporting types compatible with legacy hook expectations
|
||||
export type ScheduleFn = (
|
||||
request: ToolCallRequestInfo | ToolCallRequestInfo[],
|
||||
signal: AbortSignal,
|
||||
) => Promise<void | CompletedToolCall[]>;
|
||||
) => Promise<CompletedToolCall[]>;
|
||||
|
||||
export type UseToolSchedulerReturn = [
|
||||
export type MarkToolsAsSubmittedFn = (callIds: string[]) => void;
|
||||
export type CancelAllFn = (signal: AbortSignal) => void;
|
||||
|
||||
/**
|
||||
* The shape expected by useGeminiStream.
|
||||
* It matches the Core ToolCall structure + the UI metadata flag.
|
||||
*/
|
||||
export type TrackedToolCall = ToolCall & {
|
||||
responseSubmittedToGemini?: boolean;
|
||||
};
|
||||
|
||||
// Narrowed types for specific statuses (used by useGeminiStream)
|
||||
export type TrackedScheduledToolCall = Extract<
|
||||
TrackedToolCall,
|
||||
{ status: 'scheduled' }
|
||||
>;
|
||||
export type TrackedValidatingToolCall = Extract<
|
||||
TrackedToolCall,
|
||||
{ status: 'validating' }
|
||||
>;
|
||||
export type TrackedWaitingToolCall = Extract<
|
||||
TrackedToolCall,
|
||||
{ status: 'awaiting_approval' }
|
||||
>;
|
||||
export type TrackedExecutingToolCall = Extract<
|
||||
TrackedToolCall,
|
||||
{ status: 'executing' }
|
||||
>;
|
||||
export type TrackedCompletedToolCall = Extract<
|
||||
TrackedToolCall,
|
||||
{ status: 'success' | 'error' }
|
||||
>;
|
||||
export type TrackedCancelledToolCall = Extract<
|
||||
TrackedToolCall,
|
||||
{ status: 'cancelled' }
|
||||
>;
|
||||
|
||||
/**
|
||||
* Modern tool scheduler hook using the event-driven Core Scheduler.
|
||||
*/
|
||||
export function useToolScheduler(
|
||||
onComplete: (tools: CompletedToolCall[]) => Promise<void>,
|
||||
config: Config,
|
||||
getPreferredEditor: () => EditorType | undefined,
|
||||
): [
|
||||
TrackedToolCall[],
|
||||
ScheduleFn,
|
||||
MarkToolsAsSubmittedFn,
|
||||
React.Dispatch<React.SetStateAction<TrackedToolCall[]>>,
|
||||
CancelAllFn,
|
||||
number,
|
||||
];
|
||||
] {
|
||||
// State stores tool calls organized by their originating schedulerId
|
||||
const [toolCallsMap, setToolCallsMap] = useState<
|
||||
Record<string, TrackedToolCall[]>
|
||||
>({});
|
||||
const [lastToolOutputTime, setLastToolOutputTime] = useState<number>(0);
|
||||
|
||||
const messageBus = useMemo(() => config.getMessageBus(), [config]);
|
||||
|
||||
const onCompleteRef = useRef(onComplete);
|
||||
useEffect(() => {
|
||||
onCompleteRef.current = onComplete;
|
||||
}, [onComplete]);
|
||||
|
||||
const getPreferredEditorRef = useRef(getPreferredEditor);
|
||||
useEffect(() => {
|
||||
getPreferredEditorRef.current = getPreferredEditor;
|
||||
}, [getPreferredEditor]);
|
||||
|
||||
const scheduler = useMemo(
|
||||
() =>
|
||||
new Scheduler({
|
||||
config,
|
||||
messageBus,
|
||||
getPreferredEditor: () => getPreferredEditorRef.current(),
|
||||
schedulerId: ROOT_SCHEDULER_ID,
|
||||
}),
|
||||
[config, messageBus],
|
||||
);
|
||||
|
||||
const internalAdaptToolCalls = useCallback(
|
||||
(coreCalls: ToolCall[], prevTracked: TrackedToolCall[]) =>
|
||||
adaptToolCalls(coreCalls, prevTracked, messageBus),
|
||||
[messageBus],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
const handler = (event: ToolCallsUpdateMessage) => {
|
||||
// Update output timer for UI spinners (Side Effect)
|
||||
if (event.toolCalls.some((tc) => tc.status === 'executing')) {
|
||||
setLastToolOutputTime(Date.now());
|
||||
}
|
||||
|
||||
setToolCallsMap((prev) => {
|
||||
const adapted = internalAdaptToolCalls(
|
||||
event.toolCalls,
|
||||
prev[event.schedulerId] ?? [],
|
||||
);
|
||||
|
||||
return {
|
||||
...prev,
|
||||
[event.schedulerId]: adapted,
|
||||
};
|
||||
});
|
||||
};
|
||||
|
||||
messageBus.subscribe(MessageBusType.TOOL_CALLS_UPDATE, handler);
|
||||
return () => {
|
||||
messageBus.unsubscribe(MessageBusType.TOOL_CALLS_UPDATE, handler);
|
||||
};
|
||||
}, [messageBus, internalAdaptToolCalls]);
|
||||
|
||||
const schedule: ScheduleFn = useCallback(
|
||||
async (request, signal) => {
|
||||
// Clear state for new run
|
||||
setToolCallsMap({});
|
||||
|
||||
// 1. Await Core Scheduler directly
|
||||
const results = await scheduler.schedule(request, signal);
|
||||
|
||||
// 2. Trigger legacy reinjection logic (useGeminiStream loop)
|
||||
// Since this hook instance owns the "root" scheduler, we always trigger
|
||||
// onComplete when it finishes its batch.
|
||||
await onCompleteRef.current(results);
|
||||
|
||||
return results;
|
||||
},
|
||||
[scheduler],
|
||||
);
|
||||
|
||||
const cancelAll: CancelAllFn = useCallback(
|
||||
(_signal) => {
|
||||
scheduler.cancelAll();
|
||||
},
|
||||
[scheduler],
|
||||
);
|
||||
|
||||
const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback(
|
||||
(callIdsToMark: string[]) => {
|
||||
setToolCallsMap((prevMap) => {
|
||||
const nextMap = { ...prevMap };
|
||||
for (const [sid, calls] of Object.entries(nextMap)) {
|
||||
nextMap[sid] = calls.map((tc) =>
|
||||
callIdsToMark.includes(tc.request.callId)
|
||||
? { ...tc, responseSubmittedToGemini: true }
|
||||
: tc,
|
||||
);
|
||||
}
|
||||
return nextMap;
|
||||
});
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
// Flatten the map for the UI components that expect a single list of tools.
|
||||
const toolCalls = useMemo(
|
||||
() => Object.values(toolCallsMap).flat(),
|
||||
[toolCallsMap],
|
||||
);
|
||||
|
||||
// Provide a setter that maintains compatibility with legacy [].
|
||||
const setToolCallsForDisplay = useCallback(
|
||||
(action: React.SetStateAction<TrackedToolCall[]>) => {
|
||||
setToolCallsMap((prev) => {
|
||||
const currentFlattened = Object.values(prev).flat();
|
||||
const nextFlattened =
|
||||
typeof action === 'function' ? action(currentFlattened) : action;
|
||||
|
||||
if (nextFlattened.length === 0) {
|
||||
return {};
|
||||
}
|
||||
|
||||
// Re-group by schedulerId to preserve multi-scheduler state
|
||||
const nextMap: Record<string, TrackedToolCall[]> = {};
|
||||
for (const call of nextFlattened) {
|
||||
// All tool calls should have a schedulerId from the core.
|
||||
// Default to ROOT_SCHEDULER_ID as a safeguard.
|
||||
const sid = call.schedulerId ?? ROOT_SCHEDULER_ID;
|
||||
if (!nextMap[sid]) {
|
||||
nextMap[sid] = [];
|
||||
}
|
||||
nextMap[sid].push(call);
|
||||
}
|
||||
return nextMap;
|
||||
});
|
||||
},
|
||||
[],
|
||||
);
|
||||
|
||||
return [
|
||||
toolCalls,
|
||||
schedule,
|
||||
markToolsAsSubmitted,
|
||||
setToolCallsForDisplay,
|
||||
cancelAll,
|
||||
lastToolOutputTime,
|
||||
];
|
||||
}
|
||||
|
||||
/**
|
||||
* Hook that uses the Event-Driven scheduler for tool execution.
|
||||
* ADAPTER: Merges UI metadata (submitted flag) and injects legacy callbacks.
|
||||
*/
|
||||
export function useToolScheduler(
|
||||
onComplete: (tools: CompletedToolCall[]) => Promise<void>,
|
||||
config: Config,
|
||||
getPreferredEditor: () => EditorType | undefined,
|
||||
): UseToolSchedulerReturn {
|
||||
return useToolExecutionScheduler(
|
||||
onComplete,
|
||||
config,
|
||||
getPreferredEditor,
|
||||
) as UseToolSchedulerReturn;
|
||||
function adaptToolCalls(
|
||||
coreCalls: ToolCall[],
|
||||
prevTracked: TrackedToolCall[],
|
||||
messageBus: MessageBus,
|
||||
): TrackedToolCall[] {
|
||||
const prevMap = new Map(prevTracked.map((t) => [t.request.callId, t]));
|
||||
|
||||
return coreCalls.map((coreCall): TrackedToolCall => {
|
||||
const prev = prevMap.get(coreCall.request.callId);
|
||||
const responseSubmittedToGemini = prev?.responseSubmittedToGemini ?? false;
|
||||
|
||||
// Inject onConfirm adapter for tools awaiting approval.
|
||||
// The Core provides data-only (serializable) confirmationDetails. We must
|
||||
// inject the legacy callback function that proxies responses back to the
|
||||
// MessageBus.
|
||||
if (coreCall.status === 'awaiting_approval' && coreCall.correlationId) {
|
||||
const correlationId = coreCall.correlationId;
|
||||
return {
|
||||
...coreCall,
|
||||
confirmationDetails: {
|
||||
...coreCall.confirmationDetails,
|
||||
onConfirm: async (
|
||||
outcome: ToolConfirmationOutcome,
|
||||
payload?: ToolConfirmationPayload,
|
||||
) => {
|
||||
await messageBus.publish({
|
||||
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
|
||||
correlationId,
|
||||
confirmed: outcome !== ToolConfirmationOutcome.Cancel,
|
||||
requiresUserConfirmation: false,
|
||||
outcome,
|
||||
payload,
|
||||
});
|
||||
},
|
||||
},
|
||||
responseSubmittedToGemini,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
...coreCall,
|
||||
responseSubmittedToGemini,
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
@@ -9,7 +9,7 @@ import { renderHook } from '../../test-utils/render.js';
|
||||
import { useTurnActivityMonitor } from './useTurnActivityMonitor.js';
|
||||
import { StreamingState } from '../types.js';
|
||||
import { hasRedirection } from '@google/gemini-cli-core';
|
||||
import { type TrackedToolCall } from './useReactToolScheduler.js';
|
||||
import { type TrackedToolCall } from './useToolScheduler.js';
|
||||
|
||||
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
|
||||
const actual = await importOriginal<Record<string, unknown>>();
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
import { useState, useEffect, useRef, useMemo } from 'react';
|
||||
import { StreamingState } from '../types.js';
|
||||
import { hasRedirection } from '@google/gemini-cli-core';
|
||||
import { type TrackedToolCall } from './useReactToolScheduler.js';
|
||||
import { type TrackedToolCall } from './useToolScheduler.js';
|
||||
|
||||
export interface TurnActivityStatus {
|
||||
operationStartTime: number;
|
||||
|
||||
Reference in New Issue
Block a user