diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 78d0b0cf9b..84b558321b 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -19,8 +19,8 @@ import type { TrackedExecutingToolCall, TrackedCancelledToolCall, TrackedWaitingToolCall, -} from './useReactToolScheduler.js'; -import { useReactToolScheduler } from './useReactToolScheduler.js'; +} from './useToolScheduler.js'; +import { useToolScheduler } from './useToolScheduler.js'; import type { Config, EditorType, @@ -87,12 +87,12 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }; }); -const mockUseReactToolScheduler = useReactToolScheduler as Mock; -vi.mock('./useReactToolScheduler.js', async (importOriginal) => { +const mockUseToolScheduler = useToolScheduler as Mock; +vi.mock('./useToolScheduler.js', async (importOriginal) => { const actualSchedulerModule = (await importOriginal()) as any; return { ...(actualSchedulerModule || {}), - useReactToolScheduler: vi.fn(), + useToolScheduler: vi.fn(), }; }); @@ -243,7 +243,7 @@ describe('useGeminiStream', () => { mockMarkToolsAsSubmitted = vi.fn(); // Default mock for useReactToolScheduler to prevent toolCalls being undefined initially - mockUseReactToolScheduler.mockReturnValue([ + mockUseToolScheduler.mockReturnValue([ [], // Default to empty array for toolCalls mockScheduleToolCalls, mockMarkToolsAsSubmitted, @@ -334,7 +334,7 @@ describe('useGeminiStream', () => { rerender({ ...props, toolCalls: newToolCalls }); }); - mockUseReactToolScheduler.mockImplementation(() => [ + mockUseToolScheduler.mockImplementation(() => [ props.toolCalls, mockScheduleToolCalls, mockMarkToolsAsSubmitted, @@ -579,7 +579,7 @@ describe('useGeminiStream', () => { | ((completedTools: TrackedToolCall[]) => Promise) | null = null; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()]; }); @@ -661,7 +661,7 @@ describe('useGeminiStream', () => { | ((completedTools: TrackedToolCall[]) => Promise) | null = null; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()]; }); @@ -740,7 +740,7 @@ describe('useGeminiStream', () => { | ((completedTools: TrackedToolCall[]) => Promise) | null = null; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [ [], @@ -864,7 +864,7 @@ describe('useGeminiStream', () => { | ((completedTools: TrackedToolCall[]) => Promise) | null = null; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()]; }); @@ -972,7 +972,7 @@ describe('useGeminiStream', () => { | null = null; let currentToolCalls = initialToolCalls; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [ currentToolCalls, @@ -1009,7 +1009,7 @@ describe('useGeminiStream', () => { // 2. Update the tool calls to completed state and rerender currentToolCalls = completedToolCalls; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [ completedToolCalls, @@ -1616,7 +1616,7 @@ describe('useGeminiStream', () => { | ((completedTools: TrackedToolCall[]) => Promise) | null = null; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted, vi.fn()]; }); @@ -2306,9 +2306,8 @@ describe('useGeminiStream', () => { addItemOrder.push(`addItem:${item.type}`); }); - // We need to capture the onComplete callback from useReactToolScheduler - const mockUseReactToolScheduler = useReactToolScheduler as Mock; - mockUseReactToolScheduler.mockImplementation((onComplete) => { + // We need to capture the onComplete callback from useToolScheduler + mockUseToolScheduler.mockImplementation((onComplete) => { capturedOnComplete = onComplete; return [ [], // toolCalls @@ -2529,7 +2528,7 @@ describe('useGeminiStream', () => { }); it('should memoize pendingHistoryItems', () => { - mockUseReactToolScheduler.mockReturnValue([ + mockUseToolScheduler.mockReturnValue([ [], mockScheduleToolCalls, mockCancelAllToolCalls, @@ -2580,7 +2579,7 @@ describe('useGeminiStream', () => { } as unknown as TrackedExecutingToolCall, ]; - mockUseReactToolScheduler.mockReturnValue([ + mockUseToolScheduler.mockReturnValue([ newToolCalls, mockScheduleToolCalls, mockCancelAllToolCalls, diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 205e577fd2..16c088617b 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -66,12 +66,12 @@ import { useLogger } from './useLogger.js'; import { SHELL_COMMAND_NAME } from '../constants.js'; import { mapToDisplay as mapTrackedToolCallsToDisplay } from './toolMapping.js'; import { - useReactToolScheduler, + useToolScheduler, type TrackedToolCall, type TrackedCompletedToolCall, type TrackedCancelledToolCall, type TrackedWaitingToolCall, -} from './useReactToolScheduler.js'; +} from './useToolScheduler.js'; import { promises as fs } from 'node:fs'; import path from 'node:path'; import { useSessionStats } from '../contexts/SessionContext.js'; @@ -159,7 +159,7 @@ export const useGeminiStream = ( setToolCallsForDisplay, cancelAllToolCalls, lastToolOutputTime, - ] = useReactToolScheduler( + ] = useToolScheduler( async (completedToolCallsFromScheduler) => { // This onComplete is called when ALL scheduled tools for a given batch are done. if (completedToolCallsFromScheduler.length > 0) { diff --git a/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts b/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts new file mode 100644 index 0000000000..2a526150c3 --- /dev/null +++ b/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts @@ -0,0 +1,415 @@ +/** + * @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, +} 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(); + return { + ...actual, + Scheduler: vi.fn().mockImplementation(() => ({ + schedule: vi.fn().mockResolvedValue([]), + cancelAll: vi.fn(), + })), + }; +}); + +const createMockTool = ( + overrides: Partial = {}, +): 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 => + ({ + 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; + }); + + 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], + } 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], + } 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], + } 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], + } 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], + } 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(), + }, + ], + } 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]); + }); +}); diff --git a/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts b/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts new file mode 100644 index 0000000000..c68e414e9b --- /dev/null +++ b/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts @@ -0,0 +1,202 @@ +/** + * @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, +} 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; + +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, + config: Config, + getPreferredEditor: () => EditorType | undefined, +): [ + TrackedToolCall[], + ScheduleFn, + MarkToolsAsSubmittedFn, + React.Dispatch>, + CancelAllFn, + number, +] { + // State stores Core objects, not Display objects + const [toolCalls, setToolCalls] = useState([]); + const [lastToolOutputTime, setLastToolOutputTime] = useState(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(), + }), + [config, messageBus], + ); + + const internalAdaptToolCalls = useCallback( + (coreCalls: ToolCall[], prevTracked: TrackedToolCall[]) => + adaptToolCalls(coreCalls, prevTracked, messageBus), + [messageBus], + ); + + useEffect(() => { + const handler = (event: ToolCallsUpdateMessage) => { + setToolCalls((prev) => { + const adapted = internalAdaptToolCalls(event.toolCalls, prev); + + // Update output timer for UI spinners + if (event.toolCalls.some((tc) => tc.status === 'executing')) { + setLastToolOutputTime(Date.now()); + } + + return 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 + setToolCalls([]); + + // 1. Await Core Scheduler directly + const results = await scheduler.schedule(request, signal); + + // 2. Trigger legacy reinjection logic (useGeminiStream loop) + await onCompleteRef.current(results); + + return results; + }, + [scheduler], + ); + + const cancelAll: CancelAllFn = useCallback( + (_signal) => { + scheduler.cancelAll(); + }, + [scheduler], + ); + + const markToolsAsSubmitted: MarkToolsAsSubmittedFn = useCallback( + (callIdsToMark: string[]) => { + setToolCalls((prevCalls) => + prevCalls.map((tc) => + callIdsToMark.includes(tc.request.callId) + ? { ...tc, responseSubmittedToGemini: true } + : tc, + ), + ); + }, + [], + ); + + return [ + toolCalls, + schedule, + markToolsAsSubmitted, + setToolCalls, + 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, + }; + }); +} diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts new file mode 100644 index 0000000000..079e3f1327 --- /dev/null +++ b/packages/cli/src/ui/hooks/useToolScheduler.ts @@ -0,0 +1,86 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + Config, + EditorType, + CompletedToolCall, + ToolCallRequestInfo, +} from '@google/gemini-cli-core'; +import { + useReactToolScheduler, + type TrackedToolCall as LegacyTrackedToolCall, + type TrackedScheduledToolCall, + type TrackedValidatingToolCall, + type TrackedWaitingToolCall, + type TrackedExecutingToolCall, + type TrackedCompletedToolCall, + type TrackedCancelledToolCall, + type MarkToolsAsSubmittedFn, + type CancelAllFn, +} from './useReactToolScheduler.js'; +import { + useToolExecutionScheduler, + type TrackedToolCall as NewTrackedToolCall, +} from './useToolExecutionScheduler.js'; + +// Re-export specific state types from Legacy, as the structures are compatible +// and useGeminiStream relies on them for narrowing. +export type { + TrackedScheduledToolCall, + TrackedValidatingToolCall, + TrackedWaitingToolCall, + TrackedExecutingToolCall, + TrackedCompletedToolCall, + TrackedCancelledToolCall, +}; + +// Unified type that covers both implementations +export type TrackedToolCall = LegacyTrackedToolCall | NewTrackedToolCall; + +// Unified Schedule function (Promise | Promise) +export type ScheduleFn = ( + request: ToolCallRequestInfo | ToolCallRequestInfo[], + signal: AbortSignal, +) => Promise; + +export type UseToolSchedulerReturn = [ + TrackedToolCall[], + ScheduleFn, + MarkToolsAsSubmittedFn, + React.Dispatch>, + CancelAllFn, + number, +]; + +/** + * Facade hook that switches between the Legacy and Event-Driven schedulers + * based on configuration. + * + * Note: This conditionally calls hooks, which technically violates the standard + * Rules of Hooks linting. However, this is safe here because + * `config.isEventDrivenSchedulerEnabled()` is static for the lifetime of the + * application session (it essentially acts as a compile-time feature flag). + */ +export function useToolScheduler( + onComplete: (tools: CompletedToolCall[]) => Promise, + config: Config, + getPreferredEditor: () => EditorType | undefined, +): UseToolSchedulerReturn { + const isEventDriven = config.isEventDrivenSchedulerEnabled(); + + // Note: We return the hooks directly without casting. They return compatible + // tuple structures, but use explicit tuple signatures rather than the + // UseToolSchedulerReturn named type to avoid circular dependencies back to + // this facade. + if (isEventDriven) { + // eslint-disable-next-line react-hooks/rules-of-hooks + return useToolExecutionScheduler(onComplete, config, getPreferredEditor); + } + + // eslint-disable-next-line react-hooks/rules-of-hooks + return useReactToolScheduler(onComplete, config, getPreferredEditor); +} diff --git a/packages/cli/src/ui/hooks/useToolSchedulerFacade.test.ts b/packages/cli/src/ui/hooks/useToolSchedulerFacade.test.ts new file mode 100644 index 0000000000..112b7f34db --- /dev/null +++ b/packages/cli/src/ui/hooks/useToolSchedulerFacade.test.ts @@ -0,0 +1,70 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { renderHook } from '../../test-utils/render.js'; +import { useToolScheduler } from './useToolScheduler.js'; +import { useReactToolScheduler } from './useReactToolScheduler.js'; +import { useToolExecutionScheduler } from './useToolExecutionScheduler.js'; +import type { Config } from '@google/gemini-cli-core'; + +vi.mock('./useReactToolScheduler.js', () => ({ + useReactToolScheduler: vi.fn().mockReturnValue(['legacy']), +})); + +vi.mock('./useToolExecutionScheduler.js', () => ({ + useToolExecutionScheduler: vi.fn().mockReturnValue(['modern']), +})); + +describe('useToolScheduler (Facade)', () => { + let mockConfig: Config; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('delegates to useReactToolScheduler when event-driven scheduler is disabled', () => { + mockConfig = { + isEventDrivenSchedulerEnabled: () => false, + } as unknown as Config; + + const onComplete = vi.fn(); + const getPreferredEditor = vi.fn(); + + const { result } = renderHook(() => + useToolScheduler(onComplete, mockConfig, getPreferredEditor), + ); + + expect(result.current).toEqual(['legacy']); + expect(useReactToolScheduler).toHaveBeenCalledWith( + onComplete, + mockConfig, + getPreferredEditor, + ); + expect(useToolExecutionScheduler).not.toHaveBeenCalled(); + }); + + it('delegates to useToolExecutionScheduler when event-driven scheduler is enabled', () => { + mockConfig = { + isEventDrivenSchedulerEnabled: () => true, + } as unknown as Config; + + const onComplete = vi.fn(); + const getPreferredEditor = vi.fn(); + + const { result } = renderHook(() => + useToolScheduler(onComplete, mockConfig, getPreferredEditor), + ); + + expect(result.current).toEqual(['modern']); + expect(useToolExecutionScheduler).toHaveBeenCalledWith( + onComplete, + mockConfig, + getPreferredEditor, + ); + expect(useReactToolScheduler).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index bc1fd35f28..775223f4e5 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -36,6 +36,7 @@ export * from './core/tokenLimits.js'; export * from './core/turn.js'; export * from './core/geminiRequest.js'; export * from './core/coreToolScheduler.js'; +export * from './scheduler/scheduler.js'; export * from './scheduler/types.js'; export * from './scheduler/tool-executor.js'; export * from './core/nonInteractiveToolExecutor.js';