diff options
Diffstat (limited to 'packages/cli/src/ui/hooks')
| -rw-r--r-- | packages/cli/src/ui/hooks/useConsoleMessages.ts | 6 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.test.tsx | 200 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useGeminiStream.ts | 88 | ||||
| -rw-r--r-- | packages/cli/src/ui/hooks/useReactToolScheduler.ts | 4 |
4 files changed, 162 insertions, 136 deletions
diff --git a/packages/cli/src/ui/hooks/useConsoleMessages.ts b/packages/cli/src/ui/hooks/useConsoleMessages.ts index 90dc6f81..52ffbd39 100644 --- a/packages/cli/src/ui/hooks/useConsoleMessages.ts +++ b/packages/cli/src/ui/hooks/useConsoleMessages.ts @@ -25,9 +25,12 @@ export function useConsoleMessages(): UseConsoleMessagesReturn { return; } + const newMessagesToAdd = messageQueueRef.current; + messageQueueRef.current = []; + setConsoleMessages((prevMessages) => { const newMessages = [...prevMessages]; - messageQueueRef.current.forEach((queuedMessage) => { + newMessagesToAdd.forEach((queuedMessage) => { if ( newMessages.length > 0 && newMessages[newMessages.length - 1].type === queuedMessage.type && @@ -42,7 +45,6 @@ export function useConsoleMessages(): UseConsoleMessagesReturn { return newMessages; }); - messageQueueRef.current = []; messageQueueTimeoutRef.current = null; // Allow next scheduling }, []); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 5ac9aaa7..0c8b261e 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -496,13 +496,17 @@ describe('useGeminiStream', () => { } as TrackedCompletedToolCall, // Treat error as a form of completion for submission ]; - // 1. On the first render, there are no tool calls. - mockUseReactToolScheduler.mockReturnValue([ - [], - mockScheduleToolCalls, - mockMarkToolsAsSubmitted, - ]); - const { rerender } = renderHook(() => + // Capture the onComplete callback + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise<void>) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => useGeminiStream( new MockedGeminiClientClass(mockConfig), [], @@ -518,16 +522,11 @@ describe('useGeminiStream', () => { ), ); - // 2. Before the second render, change the mock to return the completed tools. - mockUseReactToolScheduler.mockReturnValue([ - completedToolCalls, - mockScheduleToolCalls, - mockMarkToolsAsSubmitted, - ]); - - // 3. Trigger a re-render. The hook will now receive the completed tools, causing the effect to run. - act(() => { - rerender(); + // Trigger the onComplete callback with completed tools + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete(completedToolCalls); + } }); await waitFor(() => { @@ -561,13 +560,17 @@ describe('useGeminiStream', () => { ]; const client = new MockedGeminiClientClass(mockConfig); - // 1. First render: no tool calls. - mockUseReactToolScheduler.mockReturnValue([ - [], - mockScheduleToolCalls, - mockMarkToolsAsSubmitted, - ]); - const { rerender } = renderHook(() => + // Capture the onComplete callback + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise<void>) + | null = null; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => useGeminiStream( client, [], @@ -583,16 +586,11 @@ describe('useGeminiStream', () => { ), ); - // 2. Second render: tool calls are now cancelled. - mockUseReactToolScheduler.mockReturnValue([ - cancelledToolCalls, - mockScheduleToolCalls, - mockMarkToolsAsSubmitted, - ]); - - // 3. Trigger the re-render. - act(() => { - rerender(); + // Trigger the onComplete callback with cancelled tools + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete(cancelledToolCalls); + } }); await waitFor(() => { @@ -685,7 +683,12 @@ describe('useGeminiStream', () => { const initialToolCalls: TrackedToolCall[] = [ { - request: { callId: 'call1', name: 'tool1', args: {} }, + request: { + callId: 'call1', + name: 'tool1', + args: {}, + isClientInitiated: false, + }, status: 'executing', responseSubmittedToGemini: false, tool: { @@ -711,36 +714,67 @@ describe('useGeminiStream', () => { } as TrackedCompletedToolCall, ]; - const { result, rerender, client } = renderTestHook(initialToolCalls); + // Capture the onComplete callback + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise<void>) + | null = null; + let currentToolCalls = initialToolCalls; + + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [ + currentToolCalls, + mockScheduleToolCalls, + mockMarkToolsAsSubmitted, + ]; + }); + + const { result, rerender } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + ), + ); // 1. Initial state should be Responding because a tool is executing. expect(result.current.streamingState).toBe(StreamingState.Responding); - // 2. Rerender with the completed tool call. - // The useEffect should pick this up but hasn't called submitQuery yet. + // 2. Update the tool calls to completed state and rerender + currentToolCalls = completedToolCalls; + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [ + completedToolCalls, + mockScheduleToolCalls, + mockMarkToolsAsSubmitted, + ]; + }); + act(() => { - rerender({ - client, - history: [], - addItem: mockAddItem, - setShowHelp: mockSetShowHelp, - config: mockConfig, - onDebugMessage: mockOnDebugMessage, - handleSlashCommand: - mockHandleSlashCommand as unknown as typeof mockHandleSlashCommand, - shellModeActive: false, - loadedSettings: mockLoadedSettings, - // This is the key part of the test: update the toolCalls array - // to simulate the tool finishing. - toolCalls: completedToolCalls, - }); + rerender(); }); // 3. The state should *still* be Responding, not Idle. // This is because the completed tool's response has not been submitted yet. expect(result.current.streamingState).toBe(StreamingState.Responding); - // 4. Wait for the useEffect to call submitQuery. + // 4. Trigger the onComplete callback to simulate tool completion + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete(completedToolCalls); + } + }); + + // 5. Wait for submitQuery to be called await waitFor(() => { expect(mockSendMessageStream).toHaveBeenCalledWith( toolCallResponseParts, @@ -748,7 +782,7 @@ describe('useGeminiStream', () => { ); }); - // 5. After submission, the state should remain Responding. + // 6. After submission, the state should remain Responding until the stream completes. expect(result.current.streamingState).toBe(StreamingState.Responding); }); @@ -929,14 +963,17 @@ describe('useGeminiStream', () => { } as any, }; - // 1. Initial render state: no tool calls - mockUseReactToolScheduler.mockReturnValue([ - [], - mockScheduleToolCalls, - mockMarkToolsAsSubmitted, - ]); + // Capture the onComplete callback + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise<void>) + | null = null; - const { result, rerender } = renderHook(() => + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + const { result } = renderHook(() => useGeminiStream( new MockedGeminiClientClass(mockConfig), [], @@ -957,17 +994,11 @@ describe('useGeminiStream', () => { await result.current.submitQuery('/memory add "test fact"'); }); - // The command handler schedules the tool. Now we simulate the tool completing. - // 2. Before the next render, set the mock to return the completed tool. - mockUseReactToolScheduler.mockReturnValue([ - [completedToolCall], - mockScheduleToolCalls, - mockMarkToolsAsSubmitted, - ]); - - // 3. Trigger a re-render to process the completed tool. - act(() => { - rerender(); + // Trigger the onComplete callback with the completed client-initiated tool + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([completedToolCall]); + } }); // --- Assert the outcome --- @@ -1007,13 +1038,17 @@ describe('useGeminiStream', () => { } as any, }; - mockUseReactToolScheduler.mockReturnValue([ - [completedToolCall], - mockScheduleToolCalls, - mockMarkToolsAsSubmitted, - ]); + // Capture the onComplete callback + let capturedOnComplete: + | ((completedTools: TrackedToolCall[]) => Promise<void>) + | null = null; - const { rerender } = renderHook(() => + mockUseReactToolScheduler.mockImplementation((onComplete) => { + capturedOnComplete = onComplete; + return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted]; + }); + + renderHook(() => useGeminiStream( new MockedGeminiClientClass(mockConfig), [], @@ -1029,8 +1064,11 @@ describe('useGeminiStream', () => { ), ); - act(() => { - rerender(); + // Trigger the onComplete callback with the completed save_memory tool + await act(async () => { + if (capturedOnComplete) { + await capturedOnComplete([completedToolCall]); + } }); await waitFor(() => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index ddd7c185..3d24ede7 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -111,17 +111,21 @@ export const useGeminiStream = ( const [toolCalls, scheduleToolCalls, markToolsAsSubmitted] = useReactToolScheduler( - (completedToolCallsFromScheduler) => { + async (completedToolCallsFromScheduler) => { // This onComplete is called when ALL scheduled tools for a given batch are done. if (completedToolCallsFromScheduler.length > 0) { // Add the final state of these tools to the history for display. - // The new useEffect will handle submitting their responses. addItem( mapTrackedToolCallsToDisplay( completedToolCallsFromScheduler as TrackedToolCall[], ), Date.now(), ); + + // Handle tool response submission immediately when tools complete + await handleCompletedTools( + completedToolCallsFromScheduler as TrackedToolCall[], + ); } }, config, @@ -570,40 +574,33 @@ export const useGeminiStream = ( ], ); - /** - * Automatically submits responses for completed tool calls. - * This effect runs when `toolCalls` or `isResponding` changes. - * It ensures that tool responses are sent back to Gemini only when - * all processing for a given set of tools is finished and Gemini - * is not already generating a response. - */ - useEffect(() => { - const run = async () => { + const handleCompletedTools = useCallback( + async (completedToolCallsFromScheduler: TrackedToolCall[]) => { if (isResponding) { return; } - const completedAndReadyToSubmitTools = toolCalls.filter( - ( - tc: TrackedToolCall, - ): tc is TrackedCompletedToolCall | TrackedCancelledToolCall => { - const isTerminalState = - tc.status === 'success' || - tc.status === 'error' || - tc.status === 'cancelled'; + const completedAndReadyToSubmitTools = + completedToolCallsFromScheduler.filter( + ( + tc: TrackedToolCall, + ): tc is TrackedCompletedToolCall | TrackedCancelledToolCall => { + const isTerminalState = + tc.status === 'success' || + tc.status === 'error' || + tc.status === 'cancelled'; - if (isTerminalState) { - const completedOrCancelledCall = tc as - | TrackedCompletedToolCall - | TrackedCancelledToolCall; - return ( - !completedOrCancelledCall.responseSubmittedToGemini && - completedOrCancelledCall.response?.responseParts !== undefined - ); - } - return false; - }, - ); + if (isTerminalState) { + const completedOrCancelledCall = tc as + | TrackedCompletedToolCall + | TrackedCancelledToolCall; + return ( + completedOrCancelledCall.response?.responseParts !== undefined + ); + } + return false; + }, + ); // Finalize any client-initiated tools as soon as they are done. const clientTools = completedAndReadyToSubmitTools.filter( @@ -630,15 +627,6 @@ export const useGeminiStream = ( ); } - // Only proceed with submitting to Gemini if ALL tools are complete. - const allToolsAreComplete = - toolCalls.length > 0 && - toolCalls.length === completedAndReadyToSubmitTools.length; - - if (!allToolsAreComplete) { - return; - } - const geminiTools = completedAndReadyToSubmitTools.filter( (t) => !t.request.isClientInitiated, ); @@ -693,17 +681,15 @@ export const useGeminiStream = ( submitQuery(mergePartListUnions(responsesToSend), { isContinuation: true, }); - }; - void run(); - }, [ - toolCalls, - isResponding, - submitQuery, - markToolsAsSubmitted, - addItem, - geminiClient, - performMemoryRefresh, - ]); + }, + [ + isResponding, + submitQuery, + markToolsAsSubmitted, + geminiClient, + performMemoryRefresh, + ], + ); const pendingHistoryItems = [ pendingHistoryItemRef.current, diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index 5e1bbd25..22988fef 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -128,7 +128,7 @@ export function useReactToolScheduler( }), ); }, - [], + [setToolCallsForDisplay], ); const scheduler = useMemo( @@ -152,7 +152,7 @@ export function useReactToolScheduler( ); const schedule: ScheduleFn = useCallback( - async ( + ( request: ToolCallRequestInfo | ToolCallRequestInfo[], signal: AbortSignal, ) => { |