diff --git a/packages/cli/src/test-utils/mockConfig.ts b/packages/cli/src/test-utils/mockConfig.ts index ac2176c0e3..0a02e01889 100644 --- a/packages/cli/src/test-utils/mockConfig.ts +++ b/packages/cli/src/test-utils/mockConfig.ts @@ -18,6 +18,7 @@ export const createMockConfig = (overrides: Partial = {}): 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'), diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 6a19d80184..7d817f44f5 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -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); diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index b6fdd53325..0c333176e0 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -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, diff --git a/packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap b/packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap deleted file mode 100644 index 3195316980..0000000000 --- a/packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap +++ /dev/null @@ -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", -} -`; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 294c537af4..ed7168667a 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -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(() => ''), diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.test.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.test.ts deleted file mode 100644 index ed2c64d212..0000000000 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.test.ts +++ /dev/null @@ -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(); - 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); - }); -}); diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts deleted file mode 100644 index cd17b305b5..0000000000 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ /dev/null @@ -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; -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, - config: Config, - getPreferredEditor: () => EditorType | undefined, -): [ - TrackedToolCall[], - ScheduleFn, - MarkToolsAsSubmittedFn, - React.Dispatch>, - CancelAllFn, - number, -] { - const [toolCallsForDisplay, setToolCallsForDisplay] = useState< - TrackedToolCall[] - >([]); - const [lastToolOutputTime, setLastToolOutputTime] = useState(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, - ]; -} diff --git a/packages/cli/src/ui/hooks/useShellInactivityStatus.ts b/packages/cli/src/ui/hooks/useShellInactivityStatus.ts index d0e5c0706d..092e58baae 100644 --- a/packages/cli/src/ui/hooks/useShellInactivityStatus.ts +++ b/packages/cli/src/ui/hooks/useShellInactivityStatus.ts @@ -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; diff --git a/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts b/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts deleted file mode 100644 index 797109499b..0000000000 --- a/packages/cli/src/ui/hooks/useToolExecutionScheduler.test.ts +++ /dev/null @@ -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(); - 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; - }); - - 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'); - }); -}); diff --git a/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts b/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts deleted file mode 100644 index 0c58e7fc41..0000000000 --- a/packages/cli/src/ui/hooks/useToolExecutionScheduler.ts +++ /dev/null @@ -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; - -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 tool calls organized by their originating schedulerId - const [toolCallsMap, setToolCallsMap] = useState< - Record - >({}); - 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(), - 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) => { - 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 = {}; - 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, - }; - }); -} diff --git a/packages/cli/src/ui/hooks/useToolScheduler.test.ts b/packages/cli/src/ui/hooks/useToolScheduler.test.ts index 81cafb4f34..4a04d6225c 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.test.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.test.ts @@ -1,1135 +1,525 @@ /** * @license - * Copyright 2025 Google LLC + * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -/* eslint-disable @typescript-eslint/no-explicit-any */ -import type { Mock } from 'vitest'; -import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { describe, it, expect, vi, beforeEach } from 'vitest'; import { act } from 'react'; import { renderHook } from '../../test-utils/render.js'; -import { useReactToolScheduler } from './useReactToolScheduler.js'; -import { mapToDisplay } from './toolMapping.js'; -import type { PartUnion, FunctionResponse } from '@google/genai'; -import type { - Config, - ToolCallRequestInfo, - ToolRegistry, - ToolResult, - ToolCallConfirmationDetails, - ToolCallResponseInfo, - ToolCall, // Import from core - Status as ToolCallStatusType, - AnyDeclarativeTool, - AnyToolInvocation, -} from '@google/gemini-cli-core'; +import { useToolScheduler } from './useToolScheduler.js'; import { - DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, + MessageBusType, ToolConfirmationOutcome, - ApprovalMode, - HookSystem, - PREVIEW_GEMINI_MODEL, - PolicyDecision, + Scheduler, + type Config, + type MessageBus, + type CompletedToolCall, + type ToolCallConfirmationDetails, + type ToolCallsUpdateMessage, + type AnyDeclarativeTool, + type AnyToolInvocation, + ROOT_SCHEDULER_ID, } from '@google/gemini-cli-core'; -import { MockTool } from '@google/gemini-cli-core/src/test-utils/mock-tool.js'; import { createMockMessageBus } from '@google/gemini-cli-core/src/test-utils/mock-message-bus.js'; -import { ToolCallStatus } from '../types.js'; -// Mocks -vi.mock('@google/gemini-cli-core', async () => { - const actual = await vi.importActual('@google/gemini-cli-core'); - // Patch CoreToolScheduler to have cancelAll if it's missing in the test environment - if ( - actual.CoreToolScheduler && - !actual.CoreToolScheduler.prototype.cancelAll - ) { - actual.CoreToolScheduler.prototype.cancelAll = vi.fn(); - } +// Mock Core Scheduler +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const actual = + await importOriginal(); return { ...actual, - ToolRegistry: vi.fn(), - Config: vi.fn(), + Scheduler: vi.fn().mockImplementation(() => ({ + schedule: vi.fn().mockResolvedValue([]), + cancelAll: vi.fn(), + })), }; }); -const mockToolRegistry = { - getTool: vi.fn(), - getAllToolNames: vi.fn(() => ['mockTool', 'anotherTool']), -}; +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 mockConfig = { - getToolRegistry: vi.fn(() => mockToolRegistry as unknown as ToolRegistry), - getApprovalMode: vi.fn(() => ApprovalMode.DEFAULT), - getSessionId: () => 'test-session-id', - getUsageStatisticsEnabled: () => true, - getDebugMode: () => false, - getWorkingDir: () => '/working/dir', - storage: { - getProjectTempDir: () => '/tmp', - }, - getTruncateToolOutputThreshold: () => DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD, - getAllowedTools: vi.fn(() => []), - getActiveModel: () => PREVIEW_GEMINI_MODEL, - getContentGeneratorConfig: () => ({ - model: 'test-model', - authType: 'oauth-personal', - }), - getGeminiClient: () => null, // No client needed for these tests - getShellExecutionConfig: () => ({ terminalWidth: 80, terminalHeight: 24 }), - getMessageBus: () => null, - isInteractive: () => false, - getExperiments: () => {}, - getEnableHooks: () => false, -} as unknown as Config; -mockConfig.getMessageBus = vi.fn().mockReturnValue(createMockMessageBus()); -mockConfig.getHookSystem = vi.fn().mockReturnValue(new HookSystem(mockConfig)); -mockConfig.getPolicyEngine = vi.fn().mockReturnValue({ - check: async () => { - const mode = mockConfig.getApprovalMode(); - if (mode === ApprovalMode.YOLO) { - return { decision: PolicyDecision.ALLOW }; - } - return { decision: PolicyDecision.ASK_USER }; - }, -}); +const createMockInvocation = ( + overrides: Partial = {}, +): AnyToolInvocation => + ({ + getDescription: () => 'Executing test tool', + shouldConfirmExecute: vi.fn(), + execute: vi.fn(), + params: {}, + toolLocations: [], + ...overrides, + }) as AnyToolInvocation; -function createMockConfigOverride(overrides: Partial = {}): Config { - return { ...mockConfig, ...overrides } as Config; -} - -const mockTool = new MockTool({ - name: 'mockTool', - displayName: 'Mock Tool', - execute: vi.fn(), - shouldConfirmExecute: vi.fn(), -}); -const mockToolWithLiveOutput = new MockTool({ - name: 'mockToolWithLiveOutput', - displayName: 'Mock Tool With Live Output', - description: 'A mock tool for testing', - params: {}, - isOutputMarkdown: true, - canUpdateOutput: true, - execute: vi.fn(), - shouldConfirmExecute: vi.fn(), -}); -let mockOnUserConfirmForToolConfirmation: Mock; -const mockToolRequiresConfirmation = new MockTool({ - name: 'mockToolRequiresConfirmation', - displayName: 'Mock Tool Requires Confirmation', - execute: vi.fn(), - shouldConfirmExecute: vi.fn(), -}); - -describe('useReactToolScheduler in YOLO Mode', () => { - let onComplete: Mock; +describe('useToolScheduler', () => { + let mockConfig: Config; + let mockMessageBus: MessageBus; beforeEach(() => { - onComplete = vi.fn(); - mockToolRegistry.getTool.mockClear(); - (mockToolRequiresConfirmation.execute as Mock).mockClear(); - (mockToolRequiresConfirmation.shouldConfirmExecute as Mock).mockClear(); - - // IMPORTANT: Enable YOLO mode for this test suite - (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.YOLO); - - vi.useFakeTimers(); + vi.clearAllMocks(); + mockMessageBus = createMockMessageBus() as unknown as MessageBus; + mockConfig = { + getMessageBus: () => mockMessageBus, + } as unknown as Config; }); afterEach(() => { - vi.clearAllTimers(); - vi.useRealTimers(); - // IMPORTANT: Disable YOLO mode after this test suite - (mockConfig.getApprovalMode as Mock).mockReturnValue(ApprovalMode.DEFAULT); + vi.clearAllMocks(); }); - const renderSchedulerInYoloMode = () => - renderHook(() => - useReactToolScheduler( - onComplete, - mockConfig as unknown as Config, + it('initializes with empty tool calls', () => { + const { result } = renderHook(() => + useToolScheduler( + 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(() => + useToolScheduler( + vi.fn().mockResolvedValue(undefined), + mockConfig, () => undefined, ), ); - it('should skip confirmation and execute tool directly when yoloMode is true', async () => { - mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation); - const expectedOutput = 'YOLO Confirmed output'; - (mockToolRequiresConfirmation.execute as Mock).mockResolvedValue({ - llmContent: expectedOutput, - returnDisplay: 'YOLO Formatted tool output', - } as ToolResult); + 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...', + }; - const { result } = renderSchedulerInYoloMode(); - const schedule = result.current[1]; - const request: ToolCallRequestInfo = { - callId: 'yoloCall', - name: 'mockToolRequiresConfirmation', - args: { data: 'any data' }, - } as any; - - await act(async () => { - await schedule(request, new AbortController().signal); + act(() => { + void mockMessageBus.publish({ + type: MessageBusType.TOOL_CALLS_UPDATE, + toolCalls: [mockToolCall], + schedulerId: ROOT_SCHEDULER_ID, + } as ToolCallsUpdateMessage); }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); // Process validation - }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); // Process scheduling - }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); // Process execution + 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, }); + }); - // Check that execute WAS called - expect(mockToolRequiresConfirmation.execute).toHaveBeenCalledWith( - request.args, + it('injects onConfirm callback for awaiting_approval tools (Adapter Pattern)', async () => { + const { result } = renderHook(() => + useToolScheduler( + vi.fn().mockResolvedValue(undefined), + mockConfig, + () => undefined, + ), ); - // Check that onComplete was called with success - expect(onComplete).toHaveBeenCalledWith([ - expect.objectContaining({ - status: 'success', - request, - response: expect.objectContaining({ - resultDisplay: 'YOLO Formatted tool output', - responseParts: [ - { - functionResponse: { - id: 'yoloCall', - name: 'mockToolRequiresConfirmation', - response: { output: expectedOutput }, - }, - }, - ], - }), + 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, + }); }); -}); -describe('useReactToolScheduler', () => { - let onComplete: Mock; - let capturedOnConfirmForTest: - | ((outcome: ToolConfirmationOutcome) => void | Promise) - | undefined; - - const advanceAndSettle = async () => { - await act(async () => { - await vi.advanceTimersByTimeAsync(0); - }); - }; - - const scheduleAndWaitForExecution = async ( - schedule: ( - req: ToolCallRequestInfo | ToolCallRequestInfo[], - signal: AbortSignal, - ) => Promise, - request: ToolCallRequestInfo | ToolCallRequestInfo[], - ) => { - await act(async () => { - await schedule(request, new AbortController().signal); - }); - - await advanceAndSettle(); - await advanceAndSettle(); - await advanceAndSettle(); - }; - - beforeEach(() => { - onComplete = vi.fn(); - capturedOnConfirmForTest = undefined; - - mockToolRegistry.getTool.mockClear(); - (mockTool.execute as Mock).mockClear(); - (mockTool.shouldConfirmExecute as Mock).mockClear(); - (mockToolWithLiveOutput.execute as Mock).mockClear(); - (mockToolWithLiveOutput.shouldConfirmExecute as Mock).mockClear(); - (mockToolRequiresConfirmation.execute as Mock).mockClear(); - (mockToolRequiresConfirmation.shouldConfirmExecute as Mock).mockClear(); - - mockOnUserConfirmForToolConfirmation = vi.fn(); - ( - mockToolRequiresConfirmation.shouldConfirmExecute as Mock - ).mockImplementation( - async (): Promise => - ({ - onConfirm: mockOnUserConfirmForToolConfirmation, - fileName: 'mockToolRequiresConfirmation.ts', - fileDiff: 'Mock tool requires confirmation', - type: 'edit', - title: 'Mock Tool Requires Confirmation', - }) as any, + it('injects onConfirm with payload (Inline Edit support)', async () => { + const { result } = renderHook(() => + useToolScheduler( + vi.fn().mockResolvedValue(undefined), + mockConfig, + () => undefined, + ), ); - vi.useFakeTimers(); + 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, + }); }); - afterEach(() => { - vi.clearAllTimers(); + it('preserves responseSubmittedToGemini flag across updates', () => { + const { result } = renderHook(() => + useToolScheduler( + 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(() => + useToolScheduler( + 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(); }); - const renderScheduler = (config: Config = mockConfig) => - renderHook(() => - useReactToolScheduler(onComplete, config, () => undefined), - ); - - it('initial state should be empty', () => { - const { result } = renderScheduler(); - expect(result.current[0]).toEqual([]); - }); - - it('should schedule and execute a tool call successfully', async () => { - mockToolRegistry.getTool.mockReturnValue(mockTool); - (mockTool.execute as Mock).mockResolvedValue({ - llmContent: 'Tool output', - returnDisplay: 'Formatted tool output', - } as ToolResult); - (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null); - - const { result } = renderScheduler(); - const request: ToolCallRequestInfo = { - callId: 'call1', - name: 'mockTool', - args: { param: 'value' }, - } as any; - - let completedToolCalls: ToolCall[] = []; - onComplete.mockImplementation((calls) => { - completedToolCalls = calls; - }); - - await scheduleAndWaitForExecution(result.current[1], request); - - expect(mockTool.execute).toHaveBeenCalledWith(request.args); - expect(completedToolCalls).toHaveLength(1); - expect(completedToolCalls[0].status).toBe('success'); - expect(completedToolCalls[0].request).toBe(request); - - if ( - completedToolCalls[0].status === 'success' || - completedToolCalls[0].status === 'error' - ) { - expect(completedToolCalls[0].response).toMatchSnapshot(); - } - }); - - it('should clear previous tool calls when scheduling new ones', async () => { - mockToolRegistry.getTool.mockReturnValue(mockTool); - (mockTool.execute as Mock).mockImplementation(async () => { - await new Promise((r) => setTimeout(r, 10)); - return { - llmContent: 'Tool output', - returnDisplay: 'Formatted tool output', - }; - }); - - const { result } = renderScheduler(); - const schedule = result.current[1]; - const setToolCallsForDisplay = result.current[3]; - - // Manually set a tool call in the display. - const oldToolCall = { - request: { callId: 'oldCall' }, - status: 'success', - } as any; - act(() => { - setToolCallsForDisplay([oldToolCall]); - }); - expect(result.current[0]).toEqual([oldToolCall]); - - const newRequest: ToolCallRequestInfo = { - callId: 'newCall', - name: 'mockTool', - args: {}, - } as any; - let schedulePromise: Promise; - await act(async () => { - schedulePromise = schedule(newRequest, new AbortController().signal); - }); - - await advanceAndSettle(); - - // After scheduling, the old call should be gone, - // and the new one should be in the display in its initial state. - expect(result.current[0].length).toBe(1); - expect(result.current[0][0].request.callId).toBe('newCall'); - expect(result.current[0][0].request.callId).not.toBe('oldCall'); - - // Let the new call finish. - await act(async () => { - await vi.advanceTimersByTimeAsync(20); - }); - - await act(async () => { - await schedulePromise; - }); - - expect(onComplete).toHaveBeenCalled(); - }); - - it('should cancel all running tool calls', async () => { - mockToolRegistry.getTool.mockReturnValue(mockTool); - - let resolveExecute: (value: ToolResult) => void = () => {}; - const executePromise = new Promise((resolve) => { - resolveExecute = resolve; - }); - (mockTool.execute as Mock).mockReturnValue(executePromise); - (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null); - - const { result } = renderScheduler(); - const schedule = result.current[1]; - const cancelAllToolCalls = result.current[4]; - const request: ToolCallRequestInfo = { - callId: 'cancelCall', - name: 'mockTool', - args: {}, - } as any; - - let schedulePromise: Promise; - await act(async () => { - schedulePromise = schedule(request, new AbortController().signal); - }); - - await advanceAndSettle(); // validation - await advanceAndSettle(); // Process scheduling - - // At this point, the tool is 'executing' and waiting on the promise. - expect(result.current[0][0].status).toBe('executing'); - - const cancelController = new AbortController(); - act(() => { - cancelAllToolCalls(cancelController.signal); - }); - - await advanceAndSettle(); - - expect(onComplete).toHaveBeenCalledWith([ - expect.objectContaining({ - status: 'cancelled', - request, - }), - ]); - - // Clean up the pending promise to avoid open handles. - await act(async () => { - resolveExecute({ llmContent: 'output', returnDisplay: 'display' }); - }); - - // Now await the schedule promise - await act(async () => { - await schedulePromise; - }); - }); - - it.each([ - { - desc: 'tool not found', - setup: () => { - mockToolRegistry.getTool.mockReturnValue(undefined); - }, - request: { - callId: 'call1', - name: 'nonexistentTool', - args: {}, - } as any, - expectedErrorContains: [ - 'Tool "nonexistentTool" not found in registry', - 'Did you mean one of:', - ], - }, - { - desc: 'error during shouldConfirmExecute', - setup: () => { - mockToolRegistry.getTool.mockReturnValue(mockTool); - const confirmError = new Error('Confirmation check failed'); - (mockTool.shouldConfirmExecute as Mock).mockRejectedValue(confirmError); - }, - request: { - callId: 'call1', - name: 'mockTool', - args: {}, - } as any, - expectedError: new Error('Confirmation check failed'), - }, - { - desc: 'error during execute', - setup: () => { - mockToolRegistry.getTool.mockReturnValue(mockTool); - (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null); - const execError = new Error('Execution failed'); - (mockTool.execute as Mock).mockRejectedValue(execError); - }, - request: { - callId: 'call1', - name: 'mockTool', - args: {}, - } as any, - expectedError: new Error('Execution failed'), - }, - ])( - 'should handle $desc', - async ({ setup, request, expectedErrorContains, expectedError }) => { - setup(); - const { result } = renderScheduler(); - - let completedToolCalls: ToolCall[] = []; - onComplete.mockImplementation((calls) => { - completedToolCalls = calls; - }); - - await scheduleAndWaitForExecution(result.current[1], request); - - expect(completedToolCalls).toHaveLength(1); - expect(completedToolCalls[0].status).toBe('error'); - expect(completedToolCalls[0].request).toBe(request); - - if (expectedErrorContains) { - expectedErrorContains.forEach((errorText) => { - expect( - (completedToolCalls[0] as any).response.error.message, - ).toContain(errorText); - }); - } - - if (expectedError) { - expect((completedToolCalls[0] as any).response.error.message).toBe( - expectedError.message, - ); - } - }, - ); - - it('should handle tool requiring confirmation - approved', async () => { - mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation); - const config = createMockConfigOverride({ - isInteractive: () => true, - }); - const expectedOutput = 'Confirmed output'; - (mockToolRequiresConfirmation.execute as Mock).mockResolvedValue({ - llmContent: expectedOutput, - returnDisplay: 'Confirmed display', - } as ToolResult); - - const { result } = renderScheduler(config); - const schedule = result.current[1]; - const request: ToolCallRequestInfo = { - callId: 'callConfirm', - name: 'mockToolRequiresConfirmation', - args: { data: 'sensitive' }, - } as any; - - let schedulePromise: Promise; - await act(async () => { - schedulePromise = schedule(request, new AbortController().signal); - }); - await advanceAndSettle(); - - const waitingCall = result.current[0][0] as any; - expect(waitingCall.status).toBe('awaiting_approval'); - capturedOnConfirmForTest = waitingCall.confirmationDetails?.onConfirm; - expect(capturedOnConfirmForTest).toBeDefined(); - - await act(async () => { - await capturedOnConfirmForTest?.(ToolConfirmationOutcome.ProceedOnce); - }); - - await advanceAndSettle(); - - // Now await the schedule promise as it should complete - await act(async () => { - await schedulePromise; - }); - - expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith( - ToolConfirmationOutcome.ProceedOnce, - ); - expect(mockToolRequiresConfirmation.execute).toHaveBeenCalled(); - - const completedCalls = onComplete.mock.calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe('success'); - expect(completedCalls[0].request).toBe(request); - if ( - completedCalls[0].status === 'success' || - completedCalls[0].status === 'error' - ) { - expect(completedCalls[0].response).toMatchSnapshot(); - } - }); - - it('should handle tool requiring confirmation - cancelled by user', async () => { - mockToolRegistry.getTool.mockReturnValue(mockToolRequiresConfirmation); - const config = createMockConfigOverride({ - isInteractive: () => true, - }); - const { result } = renderScheduler(config); - const schedule = result.current[1]; - const request: ToolCallRequestInfo = { - callId: 'callConfirmCancel', - name: 'mockToolRequiresConfirmation', - args: {}, - } as any; - - let schedulePromise: Promise; - await act(async () => { - schedulePromise = schedule(request, new AbortController().signal); - }); - await advanceAndSettle(); - - const waitingCall = result.current[0][0] as any; - expect(waitingCall.status).toBe('awaiting_approval'); - capturedOnConfirmForTest = waitingCall.confirmationDetails?.onConfirm; - expect(capturedOnConfirmForTest).toBeDefined(); - - await act(async () => { - await capturedOnConfirmForTest?.(ToolConfirmationOutcome.Cancel); - }); - - await advanceAndSettle(); - - // Now await the schedule promise - await act(async () => { - await schedulePromise; - }); - - expect(mockOnUserConfirmForToolConfirmation).toHaveBeenCalledWith( - ToolConfirmationOutcome.Cancel, - ); - - const completedCalls = onComplete.mock.calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe('cancelled'); - expect(completedCalls[0].request).toBe(request); - if ( - completedCalls[0].status === 'success' || - completedCalls[0].status === 'error' || - completedCalls[0].status === 'cancelled' - ) { - expect(completedCalls[0].response).toMatchSnapshot(); - } - }); - - it('should handle live output updates', async () => { - mockToolRegistry.getTool.mockReturnValue(mockToolWithLiveOutput); - let liveUpdateFn: ((output: string) => void) | undefined; - let resolveExecutePromise: (value: ToolResult) => void; - const executePromise = new Promise((resolve) => { - resolveExecutePromise = resolve; - }); - - (mockToolWithLiveOutput.execute as Mock).mockImplementation( - async ( - _args: Record, - _signal: AbortSignal, - updateFn: ((output: string) => void) | undefined, - ) => { - liveUpdateFn = updateFn; - return executePromise; - }, - ); - (mockToolWithLiveOutput.shouldConfirmExecute as Mock).mockResolvedValue( - null, - ); - - const { result } = renderScheduler(); - const request: ToolCallRequestInfo = { - callId: 'liveCall', - name: 'mockToolWithLiveOutput', - args: {}, - } as any; - - let schedulePromise: Promise; - await act(async () => { - schedulePromise = result.current[1]( - request, - new AbortController().signal, - ); - }); - await advanceAndSettle(); - - expect(liveUpdateFn).toBeDefined(); - expect(result.current[0][0].status).toBe('executing'); - - await act(async () => { - liveUpdateFn?.('Live output 1'); - }); - await advanceAndSettle(); - - await act(async () => { - liveUpdateFn?.('Live output 2'); - }); - await advanceAndSettle(); - - act(() => { - resolveExecutePromise({ - llmContent: 'Final output', - returnDisplay: 'Final display', - } as ToolResult); - }); - await advanceAndSettle(); - - // Now await schedule - await act(async () => { - await schedulePromise; - }); - - const completedCalls = onComplete.mock.calls[0][0] as ToolCall[]; - expect(completedCalls[0].status).toBe('success'); - expect(completedCalls[0].request).toBe(request); - if ( - completedCalls[0].status === 'success' || - completedCalls[0].status === 'error' - ) { - expect(completedCalls[0].response).toMatchSnapshot(); - } - expect(result.current[0]).toEqual([]); - }); - - it('should schedule and execute multiple tool calls', async () => { - const tool1 = new MockTool({ - name: 'tool1', - displayName: 'Tool 1', - execute: vi.fn().mockResolvedValue({ - llmContent: 'Output 1', - returnDisplay: 'Display 1', - } as ToolResult), - }); - - const tool2 = new MockTool({ - name: 'tool2', - displayName: 'Tool 2', - execute: vi.fn().mockResolvedValue({ - llmContent: 'Output 2', - returnDisplay: 'Display 2', - } as ToolResult), - }); - - mockToolRegistry.getTool.mockImplementation((name) => { - if (name === 'tool1') return tool1; - if (name === 'tool2') return tool2; - return undefined; - }); - - const { result } = renderScheduler(); - const schedule = result.current[1]; - const requests: ToolCallRequestInfo[] = [ - { callId: 'multi1', name: 'tool1', args: { p: 1 } } as any, - { callId: 'multi2', name: 'tool2', args: { p: 2 } } as any, - ]; - - await act(async () => { - await schedule(requests, new AbortController().signal); - }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); - }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); - }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); - }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); - }); - - expect(onComplete).toHaveBeenCalledTimes(1); - const completedCalls = onComplete.mock.calls[0][0] as ToolCall[]; - expect(completedCalls.length).toBe(2); - - const call1Result = completedCalls.find( - (c) => c.request.callId === 'multi1', - ); - const call2Result = completedCalls.find( - (c) => c.request.callId === 'multi2', - ); - - expect(call1Result).toMatchObject({ - status: 'success', - request: requests[0], - response: expect.objectContaining({ - resultDisplay: 'Display 1', - responseParts: [ - { - functionResponse: { - id: 'multi1', - name: 'tool1', - response: { output: 'Output 1' }, - }, - }, - ], - }), - }); - expect(call2Result).toMatchObject({ - status: 'success', - request: requests[1], - response: expect.objectContaining({ - resultDisplay: 'Display 2', - responseParts: [ - { - functionResponse: { - id: 'multi2', - name: 'tool2', - response: { output: 'Output 2' }, - }, - }, - ], - }), - }); - - expect(completedCalls).toHaveLength(2); - expect(completedCalls.every((t) => t.status === 'success')).toBe(true); - }); - - it('should queue if scheduling while already running', async () => { - mockToolRegistry.getTool.mockReturnValue(mockTool); - const longExecutePromise = new Promise((resolve) => - setTimeout( - () => - resolve({ - llmContent: 'done', - returnDisplay: 'done display', - }), - 50, + it('delegates cancelAll to the Core Scheduler', () => { + const { result } = renderHook(() => + useToolScheduler( + vi.fn().mockResolvedValue(undefined), + mockConfig, + () => undefined, ), ); - (mockTool.execute as Mock).mockReturnValue(longExecutePromise); - (mockTool.shouldConfirmExecute as Mock).mockResolvedValue(null); - const { result } = renderScheduler(); - const schedule = result.current[1]; - const request1: ToolCallRequestInfo = { - callId: 'run1', - name: 'mockTool', - args: {}, - } as any; - const request2: ToolCallRequestInfo = { - callId: 'run2', - name: 'mockTool', - args: {}, - } as any; + const [, , , , cancelAll] = result.current; + const signal = new AbortController().signal; - let schedulePromise1: Promise; - let schedulePromise2: Promise; + // 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(() => + useToolScheduler(onComplete, mockConfig, () => undefined), + ); + + const [, schedule] = result.current; + const signal = new AbortController().signal; + + let completedResult: CompletedToolCall[] = []; await act(async () => { - schedulePromise1 = schedule(request1, new AbortController().signal); - }); - await act(async () => { - await vi.advanceTimersByTimeAsync(0); + completedResult = await schedule( + { + callId: 'call-1', + name: 'test', + args: {}, + isClientInitiated: false, + prompt_id: 'p1', + }, + signal, + ); }); - await act(async () => { - schedulePromise2 = schedule(request2, new AbortController().signal); + expect(completedResult).toEqual([completedToolCall]); + expect(onComplete).toHaveBeenCalledWith([completedToolCall]); + }); + + it('setToolCallsForDisplay re-groups tools by schedulerId (Multi-Scheduler support)', () => { + const { result } = renderHook(() => + useToolScheduler( + 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); }); - await act(async () => { - await vi.advanceTimersByTimeAsync(50); - await vi.advanceTimersByTimeAsync(0); + 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 })), + ); }); - // Wait for first to complete - await act(async () => { - await schedulePromise1; + // 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); }); - expect(onComplete).toHaveBeenCalledWith([ - expect.objectContaining({ - status: 'success', - request: request1, - response: expect.objectContaining({ resultDisplay: 'done display' }), - }), - ]); - - await act(async () => { - await vi.advanceTimersByTimeAsync(50); - await vi.advanceTimersByTimeAsync(0); - }); - - // Wait for second to complete - await act(async () => { - await schedulePromise2; - }); - - expect(onComplete).toHaveBeenCalledWith([ - expect.objectContaining({ - status: 'success', - request: request2, - response: expect.objectContaining({ resultDisplay: 'done display' }), - }), - ]); - const toolCalls = result.current[0]; - expect(toolCalls).toHaveLength(0); - }); -}); - -describe('mapToDisplay', () => { - const baseRequest: ToolCallRequestInfo = { - callId: 'testCallId', - name: 'testTool', - args: { foo: 'bar' }, - } as any; - - const baseTool = new MockTool({ - name: 'testTool', - displayName: 'Test Tool Display', - execute: vi.fn(), - shouldConfirmExecute: vi.fn(), - }); - - const baseResponse: ToolCallResponseInfo = { - callId: 'testCallId', - responseParts: [ - { - functionResponse: { - name: 'testTool', - id: 'testCallId', - response: { output: 'Test output' }, - } as FunctionResponse, - } as PartUnion, - ], - resultDisplay: 'Test display output', - error: undefined, - } as any; - - // Define a more specific type for extraProps for these tests - // This helps ensure that tool and confirmationDetails are only accessed when they are expected to exist. - type MapToDisplayExtraProps = - | { - tool?: AnyDeclarativeTool; - invocation?: AnyToolInvocation; - liveOutput?: string; - response?: ToolCallResponseInfo; - confirmationDetails?: ToolCallConfirmationDetails; - } - | { - tool: AnyDeclarativeTool; - invocation?: AnyToolInvocation; - response?: ToolCallResponseInfo; - confirmationDetails?: ToolCallConfirmationDetails; - } - | { - response: ToolCallResponseInfo; - tool?: undefined; - confirmationDetails?: ToolCallConfirmationDetails; - } - | { - confirmationDetails: ToolCallConfirmationDetails; - tool?: AnyDeclarativeTool; - invocation?: AnyToolInvocation; - response?: ToolCallResponseInfo; - }; - - const baseInvocation = baseTool.build(baseRequest.args); - const testCases: Array<{ - name: string; - status: ToolCallStatusType; - extraProps?: MapToDisplayExtraProps; - expectedStatus: ToolCallStatus; - expectedResultDisplay?: string; - expectedName?: string; - expectedDescription?: string; - }> = [ - { - name: 'validating', - status: 'validating', - extraProps: { tool: baseTool, invocation: baseInvocation }, - expectedStatus: ToolCallStatus.Pending, - expectedName: baseTool.displayName, - expectedDescription: baseInvocation.getDescription(), - }, - { - name: 'awaiting_approval', - status: 'awaiting_approval', - extraProps: { - tool: baseTool, - invocation: baseInvocation, - confirmationDetails: { - onConfirm: vi.fn(), - type: 'edit', - title: 'Test Tool Display', - serverName: 'testTool', - toolName: 'testTool', - toolDisplayName: 'Test Tool Display', - filePath: 'mock', - fileName: 'test.ts', - fileDiff: 'Test diff', - originalContent: 'Original content', - newContent: 'New content', - } as ToolCallConfirmationDetails, - }, - expectedStatus: ToolCallStatus.Confirming, - expectedName: baseTool.displayName, - expectedDescription: baseInvocation.getDescription(), - }, - { - name: 'scheduled', - status: 'scheduled', - extraProps: { tool: baseTool, invocation: baseInvocation }, - expectedStatus: ToolCallStatus.Pending, - expectedName: baseTool.displayName, - expectedDescription: baseInvocation.getDescription(), - }, - { - name: 'executing no live output', - status: 'executing', - extraProps: { tool: baseTool, invocation: baseInvocation }, - expectedStatus: ToolCallStatus.Executing, - expectedName: baseTool.displayName, - expectedDescription: baseInvocation.getDescription(), - }, - { - name: 'executing with live output', - status: 'executing', - extraProps: { - tool: baseTool, - invocation: baseInvocation, - liveOutput: 'Live test output', - }, - expectedStatus: ToolCallStatus.Executing, - expectedResultDisplay: 'Live test output', - expectedName: baseTool.displayName, - expectedDescription: baseInvocation.getDescription(), - }, - { - name: 'success', - status: 'success', - extraProps: { - tool: baseTool, - invocation: baseInvocation, - response: baseResponse, - }, - expectedStatus: ToolCallStatus.Success, - expectedResultDisplay: baseResponse.resultDisplay as any, - expectedName: baseTool.displayName, - expectedDescription: baseInvocation.getDescription(), - }, - { - name: 'error tool not found', - status: 'error', - extraProps: { - response: { - ...baseResponse, - error: new Error('Test error tool not found'), - resultDisplay: 'Error display tool not found', - }, - }, - expectedStatus: ToolCallStatus.Error, - expectedResultDisplay: 'Error display tool not found', - expectedName: baseRequest.name, - expectedDescription: JSON.stringify(baseRequest.args), - }, - { - name: 'error tool execution failed', - status: 'error', - extraProps: { - tool: baseTool, - response: { - ...baseResponse, - error: new Error('Tool execution failed'), - resultDisplay: 'Execution failed display', - }, - }, - expectedStatus: ToolCallStatus.Error, - expectedResultDisplay: 'Execution failed display', - expectedName: baseTool.displayName, // Changed from baseTool.name - expectedDescription: JSON.stringify(baseRequest.args), - }, - { - name: 'cancelled', - status: 'cancelled', - extraProps: { - tool: baseTool, - invocation: baseInvocation, - response: { - ...baseResponse, - resultDisplay: 'Cancelled display', - }, - }, - expectedStatus: ToolCallStatus.Canceled, - expectedResultDisplay: 'Cancelled display', - expectedName: baseTool.displayName, - expectedDescription: baseInvocation.getDescription(), - }, - ]; - - testCases.forEach( - ({ - name: testName, - status, - extraProps, - expectedStatus, - expectedResultDisplay, - expectedName, - expectedDescription, - }) => { - it(`should map ToolCall with status '${status}' (${testName}) correctly`, () => { - const toolCall: ToolCall = { - request: baseRequest, - status, - ...(extraProps || {}), - } as ToolCall; - - const display = mapToDisplay(toolCall); - expect(display.type).toBe('tool_group'); - expect(display.tools.length).toBe(1); - const toolDisplay = display.tools[0]; - - expect(toolDisplay.callId).toBe(baseRequest.callId); - expect(toolDisplay.status).toBe(expectedStatus); - expect(toolDisplay.resultDisplay).toBe(expectedResultDisplay); - - expect(toolDisplay.name).toBe(expectedName); - expect(toolDisplay.description).toBe(expectedDescription); - - expect(toolDisplay.renderOutputAsMarkdown).toBe( - extraProps?.tool?.isOutputMarkdown ?? false, - ); - if (status === 'awaiting_approval') { - expect(toolDisplay.confirmationDetails).toBe( - extraProps!.confirmationDetails, - ); - } else { - expect(toolDisplay.confirmationDetails).toBeUndefined(); - } - }); - }, - ); - - it('should map an array of ToolCalls correctly', () => { - const toolCall1: ToolCall = { - request: { ...baseRequest, callId: 'call1' }, - status: 'success', - tool: baseTool, - invocation: baseTool.build(baseRequest.args), - response: { ...baseResponse, callId: 'call1' }, - } as ToolCall; - const toolForCall2 = new MockTool({ - name: baseTool.name, - displayName: baseTool.displayName, - isOutputMarkdown: true, - execute: vi.fn(), - shouldConfirmExecute: vi.fn(), - }); - const toolCall2: ToolCall = { - request: { ...baseRequest, callId: 'call2' }, - status: 'executing', - tool: toolForCall2, - invocation: toolForCall2.build(baseRequest.args), - liveOutput: 'markdown output', - } as ToolCall; - - const display = mapToDisplay([toolCall1, toolCall2]); - expect(display.tools.length).toBe(2); - expect(display.tools[0].callId).toBe('call1'); - expect(display.tools[0].status).toBe(ToolCallStatus.Success); - expect(display.tools[0].renderOutputAsMarkdown).toBe(false); - expect(display.tools[1].callId).toBe('call2'); - expect(display.tools[1].status).toBe(ToolCallStatus.Executing); - expect(display.tools[1].resultDisplay).toBe('markdown output'); - expect(display.tools[1].renderOutputAsMarkdown).toBe(true); + [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'); }); }); diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts index b6835565e7..b50ed1b717 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.ts @@ -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 | Promise) +// Re-exporting types compatible with legacy hook expectations export type ScheduleFn = ( request: ToolCallRequestInfo | ToolCallRequestInfo[], signal: AbortSignal, -) => Promise; +) => Promise; -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, + config: Config, + getPreferredEditor: () => EditorType | undefined, +): [ TrackedToolCall[], ScheduleFn, MarkToolsAsSubmittedFn, React.Dispatch>, CancelAllFn, number, -]; +] { + // State stores tool calls organized by their originating schedulerId + const [toolCallsMap, setToolCallsMap] = useState< + Record + >({}); + 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(), + 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) => { + 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 = {}; + 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, - 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, + }; + }); } diff --git a/packages/cli/src/ui/hooks/useTurnActivityMonitor.test.ts b/packages/cli/src/ui/hooks/useTurnActivityMonitor.test.ts index 9ac44c3ebc..f77ab7504d 100644 --- a/packages/cli/src/ui/hooks/useTurnActivityMonitor.test.ts +++ b/packages/cli/src/ui/hooks/useTurnActivityMonitor.test.ts @@ -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>(); diff --git a/packages/cli/src/ui/hooks/useTurnActivityMonitor.ts b/packages/cli/src/ui/hooks/useTurnActivityMonitor.ts index cd6ee7ee8a..8cd7883007 100644 --- a/packages/cli/src/ui/hooks/useTurnActivityMonitor.ts +++ b/packages/cli/src/ui/hooks/useTurnActivityMonitor.ts @@ -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;