MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-05-12 12:54:07 -07:00
Code Issues Packages Projects Releases Wiki Activity

refactor(cli): consolidate useToolScheduler and delete legacy implementation (#18567)

Browse Source
This commit is contained in:
Abhi
2026-02-11 20:49:30 -05:00
committed by GitHub
parent a1148ea1f1
commit fad9f46273
14 changed files with 721 additions and 2302 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/cli/src/test-utils/mockConfig.ts
+1
View File
@@ -18,6 +18,7 @@ export const createMockConfig = (overrides: Partial<Config> = {}): Config =>
getSandbox: vi.fn(() => undefined), getSandbox: vi.fn(() => undefined),
getQuestion: vi.fn(() => ''), getQuestion: vi.fn(() => ''),
isInteractive: vi.fn(() => false), isInteractive: vi.fn(() => false),
isInitialized: vi.fn(() => true),
setTerminalBackground: vi.fn(), setTerminalBackground: vi.fn(),
storage: { storage: {
getProjectTempDir: vi.fn().mockReturnValue('/tmp/gemini-test'), getProjectTempDir: vi.fn().mockReturnValue('/tmp/gemini-test'),
packages/cli/src/ui/App.test.tsx
-4
View File
@@ -220,10 +220,6 @@ describe('App', () => {
} as UIState; } as UIState;
const configWithExperiment = makeFakeConfig(); const configWithExperiment = makeFakeConfig();
vi.spyOn(
configWithExperiment,
'isEventDrivenSchedulerEnabled',
).mockReturnValue(true);
vi.spyOn(configWithExperiment, 'isTrustedFolder').mockReturnValue(true); vi.spyOn(configWithExperiment, 'isTrustedFolder').mockReturnValue(true);
vi.spyOn(configWithExperiment, 'getIdeMode').mockReturnValue(false); vi.spyOn(configWithExperiment, 'getIdeMode').mockReturnValue(false);
packages/cli/src/ui/AppContainer.test.tsx
+1 -1
View File
@@ -20,7 +20,7 @@ import { cleanup } from 'ink-testing-library';
import { act, useContext, type ReactElement } from 'react'; import { act, useContext, type ReactElement } from 'react';
import { AppContainer } from './AppContainer.js'; import { AppContainer } from './AppContainer.js';
import { SettingsContext } from './contexts/SettingsContext.js'; import { SettingsContext } from './contexts/SettingsContext.js';
import { type TrackedToolCall } from './hooks/useReactToolScheduler.js'; import { type TrackedToolCall } from './hooks/useToolScheduler.js';
import { import {
type Config, type Config,
makeFakeConfig, makeFakeConfig,
packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap
-97
View File
@@ -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",
}
`;
packages/cli/src/ui/hooks/useGeminiStream.test.tsx
-1
View File
@@ -246,7 +246,6 @@ describe('useGeminiStream', () => {
getContentGenerator: vi.fn(), getContentGenerator: vi.fn(),
isInteractive: () => false, isInteractive: () => false,
getExperiments: () => {}, getExperiments: () => {},
isEventDrivenSchedulerEnabled: vi.fn(() => false),
getMaxSessionTurns: vi.fn(() => 100), getMaxSessionTurns: vi.fn(() => 100),
isJitContextEnabled: vi.fn(() => false), isJitContextEnabled: vi.fn(() => false),
getGlobalMemory: vi.fn(() => ''), getGlobalMemory: vi.fn(() => ''),
packages/cli/src/ui/hooks/useReactToolScheduler.test.ts
-77
View File
@@ -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);
});
});
packages/cli/src/ui/hooks/useReactToolScheduler.ts
-221
View File
@@ -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,
];
}
packages/cli/src/ui/hooks/useShellInactivityStatus.ts
+1 -1
View File
@@ -12,7 +12,7 @@ import {
SHELL_SILENT_WORKING_TITLE_DELAY_MS, SHELL_SILENT_WORKING_TITLE_DELAY_MS,
} from '../constants.js'; } from '../constants.js';
import type { StreamingState } from '../types.js'; import type { StreamingState } from '../types.js';
import { type TrackedToolCall } from './useReactToolScheduler.js'; import { type TrackedToolCall } from './useToolScheduler.js';
interface ShellInactivityStatusProps { interface ShellInactivityStatusProps {
activePtyId: number | string | null | undefined; activePtyId: number | string | null | undefined;
packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts
-525
View File
@@ -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');
});
});
packages/cli/src/ui/hooks/useToolExecutionScheduler.ts
-253
View File
@@ -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,
};
});
}
packages/cli/src/ui/hooks/useToolScheduler.test.ts
+462 -1072
View File
File diff suppressed because it is too large Load Diff
packages/cli/src/ui/hooks/useToolScheduler.ts
+254 -48
View File
@@ -4,67 +4,273 @@
* SPDX-License-Identifier: Apache-2.0 * SPDX-License-Identifier: Apache-2.0
*/ */
import type { import {
Config, type Config,
EditorType, type MessageBus,
CompletedToolCall, type ToolCallRequestInfo,
ToolCallRequestInfo, type ToolCall,
type CompletedToolCall,
type ToolConfirmationPayload,
MessageBusType,
ToolConfirmationOutcome,
Scheduler,
type EditorType,
type ToolCallsUpdateMessage,
ROOT_SCHEDULER_ID,
} from '@google/gemini-cli-core'; } from '@google/gemini-cli-core';
import { import { useCallback, useState, useMemo, useEffect, useRef } from 'react';
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';
// Re-export specific state types from Legacy, as the structures are compatible // Re-exporting types compatible with legacy hook expectations
// 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[]>)
export type ScheduleFn = ( export type ScheduleFn = (
request: ToolCallRequestInfo | ToolCallRequestInfo[], request: ToolCallRequestInfo | ToolCallRequestInfo[],
signal: AbortSignal, 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[], TrackedToolCall[],
ScheduleFn, ScheduleFn,
MarkToolsAsSubmittedFn, MarkToolsAsSubmittedFn,
React.Dispatch<React.SetStateAction<TrackedToolCall[]>>, React.Dispatch<React.SetStateAction<TrackedToolCall[]>>,
CancelAllFn, CancelAllFn,
number, 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( function adaptToolCalls(
onComplete: (tools: CompletedToolCall[]) => Promise<void>, coreCalls: ToolCall[],
config: Config, prevTracked: TrackedToolCall[],
getPreferredEditor: () => EditorType | undefined, messageBus: MessageBus,
): UseToolSchedulerReturn { ): TrackedToolCall[] {
return useToolExecutionScheduler( const prevMap = new Map(prevTracked.map((t) => [t.request.callId, t]));
onComplete,
config, return coreCalls.map((coreCall): TrackedToolCall => {
getPreferredEditor, const prev = prevMap.get(coreCall.request.callId);
) as UseToolSchedulerReturn; 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,
};
});
} }
packages/cli/src/ui/hooks/useTurnActivityMonitor.test.ts
+1 -1
View File
@@ -9,7 +9,7 @@ import { renderHook } from '../../test-utils/render.js';
import { useTurnActivityMonitor } from './useTurnActivityMonitor.js'; import { useTurnActivityMonitor } from './useTurnActivityMonitor.js';
import { StreamingState } from '../types.js'; import { StreamingState } from '../types.js';
import { hasRedirection } from '@google/gemini-cli-core'; 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) => { vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const actual = await importOriginal<Record<string, unknown>>(); const actual = await importOriginal<Record<string, unknown>>();
packages/cli/src/ui/hooks/useTurnActivityMonitor.ts
+1 -1
View File
@@ -7,7 +7,7 @@
import { useState, useEffect, useRef, useMemo } from 'react'; import { useState, useEffect, useRef, useMemo } from 'react';
import { StreamingState } from '../types.js'; import { StreamingState } from '../types.js';
import { hasRedirection } from '@google/gemini-cli-core'; import { hasRedirection } from '@google/gemini-cli-core';
import { type TrackedToolCall } from './useReactToolScheduler.js'; import { type TrackedToolCall } from './useToolScheduler.js';
export interface TurnActivityStatus { export interface TurnActivityStatus {
operationStartTime: number; operationStartTime: number;