@machina.ai/cell-cli 1.20.2-rc1 → 1.25.0-rc3

package/dist/src/ui/hooks/useGeminiStream.test.js

@@ -11,7 +11,7 @@ import { useGeminiStream } from './useGeminiStream.js';
 import { useKeypress } from './useKeypress.js';
 import * as atCommandProcessor from './atCommandProcessor.js';
 import { useReactToolScheduler } from './useReactToolScheduler.js';
-import { ApprovalMode, AuthType, GeminiEventType as ServerGeminiEventType, ToolErrorType, ToolConfirmationOutcome, tokenLimit, debugLogger, } from '@google/gemini-cli-core';
+import { ApprovalMode, AuthType, GeminiEventType as ServerGeminiEventType, ToolErrorType, ToolConfirmationOutcome, tokenLimit, debugLogger, coreEvents, CoreEvent, } from '@google/gemini-cli-core';
 import { MessageType, StreamingState } from '../types.js';
 // --- MOCKS ---
 const mockSendMessageStream = vi
@@ -62,6 +62,8 @@ vi.mock('./useKeypress.js', () => ({
 vi.mock('./shellCommandProcessor.js', () => ({
     useShellCommandProcessor: vi.fn().mockReturnValue({
         handleShellCommand: vi.fn(),
+        activeShellPtyId: null,
+        lastShellOutputTime: 0,
     }),
 }));
 vi.mock('./atCommandProcessor.js');
@@ -167,7 +169,6 @@ describe('useGeminiStream', () => {
             getContentGeneratorConfig: vi
                 .fn()
                 .mockReturnValue(contentGeneratorConfig),
-            getUseSmartEdit: () => false,
             isInteractive: () => false,
             getExperiments: () => { },
         };
@@ -184,6 +185,7 @@ describe('useGeminiStream', () => {
             mockMarkToolsAsSubmitted,
             vi.fn(), // setToolCallsForDisplay
             mockCancelAllToolCalls,
+            0, // lastToolOutputTime
         ]);
         // Reset mocks for GeminiClient instance methods (startChat and sendMessageStream)
         // The GeminiClient constructor itself is mocked at the module level.
@@ -486,6 +488,66 @@ describe('useGeminiStream', () => {
             expect(mockSendMessageStream).not.toHaveBeenCalled();
         });
     });
+    it('should stop agent execution immediately when a tool call returns STOP_EXECUTION error', async () => {
+        const stopExecutionToolCalls = [
+            {
+                request: {
+                    callId: 'stop-call',
+                    name: 'stopTool',
+                    args: {},
+                    isClientInitiated: false,
+                    prompt_id: 'prompt-id-stop',
+                },
+                status: 'error',
+                response: {
+                    callId: 'stop-call',
+                    responseParts: [{ text: 'error occurred' }],
+                    errorType: ToolErrorType.STOP_EXECUTION,
+                    error: new Error('Stop reason from hook'),
+                    resultDisplay: undefined,
+                },
+                responseSubmittedToGemini: false,
+                tool: {
+                    displayName: 'stop tool',
+                },
+                invocation: {
+                    getDescription: () => `Mock description`,
+                },
+            },
+        ];
+        const client = new MockedGeminiClientClass(mockConfig);
+        // Capture the onComplete callback
+        let capturedOnComplete = null;
+        mockUseReactToolScheduler.mockImplementation((onComplete) => {
+            capturedOnComplete = onComplete;
+            return [
+                [],
+                mockScheduleToolCalls,
+                mockMarkToolsAsSubmitted,
+                vi.fn(),
+                mockCancelAllToolCalls,
+            ];
+        });
+        const { result } = renderHook(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, 80, 24));
+        // Trigger the onComplete callback with STOP_EXECUTION tool
+        await act(async () => {
+            if (capturedOnComplete) {
+                await capturedOnComplete(stopExecutionToolCalls);
+            }
+        });
+        await waitFor(() => {
+            expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['stop-call']);
+            // Should add an info message to history
+            expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
+                type: MessageType.INFO,
+                text: expect.stringContaining('Agent execution stopped: Stop reason from hook'),
+            }));
+            // Ensure we do NOT call back to the API
+            expect(mockSendMessageStream).not.toHaveBeenCalled();
+            // Streaming state should be Idle
+            expect(result.current.streamingState).toBe(StreamingState.Idle);
+        });
+    });
     it('should group multiple cancelled tool call responses into a single history entry', async () => {
         const cancelledToolCall1 = {
             request: {
@@ -694,6 +756,7 @@ describe('useGeminiStream', () => {
             const { result } = renderTestHook();
             // Start a query
             await act(async () => {
+                // eslint-disable-next-line @typescript-eslint/no-floating-promises
                 result.current.submitQuery('test query');
             });
             // Wait for the first part of the response
@@ -707,7 +770,7 @@ describe('useGeminiStream', () => {
                 expect(mockAddItem).toHaveBeenCalledWith({
                     type: MessageType.INFO,
                     text: 'Request cancelled.',
-                }, expect.any(Number));
+                });
             });
             // Verify state is reset
             expect(result.current.streamingState).toBe(StreamingState.Idle);
@@ -723,6 +786,7 @@ describe('useGeminiStream', () => {
             const { result } = renderHook(() => useGeminiStream(mockConfig.getGeminiClient(), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, cancelSubmitSpy, () => { }, 80, 24));
             // Start a query
             await act(async () => {
+                // eslint-disable-next-line @typescript-eslint/no-floating-promises
                 result.current.submitQuery('test query');
             });
             simulateEscapeKeyPress();
@@ -739,6 +803,7 @@ describe('useGeminiStream', () => {
             80, 24));
             // Start a query
             await act(async () => {
+                // eslint-disable-next-line @typescript-eslint/no-floating-promises
                 result.current.submitQuery('test query');
             });
             simulateEscapeKeyPress();
@@ -752,7 +817,7 @@ describe('useGeminiStream', () => {
             // No change should happen, no cancellation message
             expect(mockAddItem).not.toHaveBeenCalledWith(expect.objectContaining({
                 text: 'Request cancelled.',
-            }), expect.any(Number));
+            }));
         });
         it('should prevent further processing after cancellation', async () => {
             let continueStream;
@@ -767,6 +832,7 @@ describe('useGeminiStream', () => {
             mockSendMessageStream.mockReturnValue(mockStream);
             const { result } = renderTestHook();
             await act(async () => {
+                // eslint-disable-next-line @typescript-eslint/no-floating-promises
                 result.current.submitQuery('long running query');
             });
             await waitFor(() => {
@@ -860,12 +926,59 @@ describe('useGeminiStream', () => {
             await waitFor(() => {
                 expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
                     text: 'Request cancelled.',
-                }), expect.any(Number));
+                }));
             });
             // The final state should be idle
             expect(result.current.streamingState).toBe(StreamingState.Idle);
         });
     });
+    describe('Retry Handling', () => {
+        it('should update retryStatus when CoreEvent.RetryAttempt is emitted', async () => {
+            const { result } = renderHookWithDefaults();
+            const retryPayload = {
+                model: 'gemini-2.5-pro',
+                attempt: 2,
+                maxAttempts: 3,
+                delayMs: 1000,
+            };
+            await act(async () => {
+                coreEvents.emit(CoreEvent.RetryAttempt, retryPayload);
+            });
+            expect(result.current.retryStatus).toEqual(retryPayload);
+        });
+        it('should reset retryStatus when isResponding becomes false', async () => {
+            const { result } = renderTestHook();
+            const retryPayload = {
+                model: 'gemini-2.5-pro',
+                attempt: 2,
+                maxAttempts: 3,
+                delayMs: 1000,
+            };
+            // Start a query to make isResponding true
+            const mockStream = (async function* () {
+                yield { type: ServerGeminiEventType.Content, value: 'Part 1' };
+                await new Promise(() => { }); // Keep stream open
+            })();
+            mockSendMessageStream.mockReturnValue(mockStream);
+            await act(async () => {
+                // eslint-disable-next-line @typescript-eslint/no-floating-promises
+                result.current.submitQuery('test query');
+            });
+            await waitFor(() => {
+                expect(result.current.streamingState).toBe(StreamingState.Responding);
+            });
+            // Emit retry event
+            await act(async () => {
+                coreEvents.emit(CoreEvent.RetryAttempt, retryPayload);
+            });
+            expect(result.current.retryStatus).toEqual(retryPayload);
+            // Cancel to make isResponding false
+            await act(async () => {
+                result.current.cancelOngoingRequest();
+            });
+            expect(result.current.retryStatus).toBeNull();
+        });
+    });
     describe('Slash Command Handling', () => {
         it('should schedule a tool call when the command processor returns a schedule_tool action', async () => {
             const clientToolRequest = {
@@ -1307,7 +1420,7 @@ describe('useGeminiStream', () => {
                     expect(mockAddItem).toHaveBeenCalledWith({
                         type: 'info',
                         text: expectedMessage,
-                    }, expect.any(Number));
+                    });
                 });
             });
         });
@@ -1415,6 +1528,68 @@ describe('useGeminiStream', () => {
             }
         });
     });
+    it('should flush pending text rationale before scheduling tool calls to ensure correct history order', async () => {
+        const addItemOrder = [];
+        let capturedOnComplete;
+        const mockScheduleToolCalls = vi.fn(async (requests) => {
+            addItemOrder.push('scheduleToolCalls_START');
+            // Simulate tools completing and triggering onComplete immediately.
+            // This mimics the behavior that caused the regression where tool results
+            // were added to history during the await scheduleToolCalls(...) block.
+            const tools = requests.map((r) => ({
+                request: r,
+                status: 'success',
+                tool: { displayName: r.name, name: r.name },
+                invocation: { getDescription: () => 'desc' },
+                response: { responseParts: [], resultDisplay: 'done' },
+                startTime: Date.now(),
+                endTime: Date.now(),
+            }));
+            await capturedOnComplete(tools);
+            addItemOrder.push('scheduleToolCalls_END');
+        });
+        mockAddItem.mockImplementation((item) => {
+            addItemOrder.push(`addItem:${item.type}`);
+        });
+        // We need to capture the onComplete callback from useReactToolScheduler
+        const mockUseReactToolScheduler = useReactToolScheduler;
+        mockUseReactToolScheduler.mockImplementation((onComplete) => {
+            capturedOnComplete = onComplete;
+            return [
+                [], // toolCalls
+                mockScheduleToolCalls,
+                vi.fn(), // markToolsAsSubmitted
+                vi.fn(), // setToolCallsForDisplay
+                vi.fn(), // cancelAllToolCalls
+                0, // lastToolOutputTime
+            ];
+        });
+        const { result } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, vi.fn(), vi.fn(), false, () => 'vscode', vi.fn(), vi.fn(), false, vi.fn(), vi.fn(), vi.fn(), 80, 24));
+        const mockStream = (async function* () {
+            yield {
+                type: ServerGeminiEventType.Content,
+                value: 'Rationale rationale.',
+            };
+            yield {
+                type: ServerGeminiEventType.ToolCallRequest,
+                value: { callId: '1', name: 'test_tool', args: {} },
+            };
+        })();
+        mockSendMessageStream.mockReturnValue(mockStream);
+        await act(async () => {
+            await result.current.submitQuery('test input');
+        });
+        // Expectation: addItem:gemini (rationale) MUST happen before scheduleToolCalls_START
+        const rationaleIndex = addItemOrder.indexOf('addItem:gemini');
+        const scheduleIndex = addItemOrder.indexOf('scheduleToolCalls_START');
+        const toolGroupIndex = addItemOrder.indexOf('addItem:tool_group');
+        expect(rationaleIndex).toBeGreaterThan(-1);
+        expect(scheduleIndex).toBeGreaterThan(-1);
+        expect(toolGroupIndex).toBeGreaterThan(-1);
+        // This is the core fix validation: Rationale comes before tools are even scheduled (awaited)
+        expect(rationaleIndex).toBeLessThan(scheduleIndex);
+        expect(rationaleIndex).toBeLessThan(toolGroupIndex);
+    });
     it('should process @include commands, adding user turn after processing to prevent race conditions', async () => {
         const rawQuery = '@include file.txt Summarize this.';
         const processedQueryParts = [
@@ -1597,6 +1772,33 @@ describe('useGeminiStream', () => {
             // Verify parseAndFormatApiError was called
             expect(mockParseAndFormatApiError).toHaveBeenCalledWith({ message: 'Test error' }, expect.any(String), undefined, 'gemini-2.5-pro', 'gemini-2.5-flash');
         });
+        it('should update lastOutputTime on Gemini thought and content events', async () => {
+            vi.useFakeTimers();
+            const startTime = 1000000;
+            vi.setSystemTime(startTime);
+            // Mock a stream that yields a thought then content
+            mockSendMessageStream.mockReturnValue((async function* () {
+                yield {
+                    type: ServerGeminiEventType.Thought,
+                    value: { subject: 'Thinking...', description: '' },
+                };
+                // Advance time for the next event
+                vi.advanceTimersByTime(1000);
+                yield {
+                    type: ServerGeminiEventType.Content,
+                    value: 'Hello',
+                };
+            })());
+            const { result } = renderHook(() => useGeminiStream(new MockedGeminiClientClass(mockConfig), [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, 80, 24));
+            // Submit query
+            await act(async () => {
+                await result.current.submitQuery('Test query');
+            });
+            // Verify lastOutputTime was updated
+            // It should be the time of the last event (startTime + 1000)
+            expect(result.current.lastOutputTime).toBe(startTime + 1000);
+            vi.useRealTimers();
+        });
     });
     describe('Loop Detection Confirmation', () => {
         beforeEach(() => {
@@ -1676,7 +1878,7 @@ describe('useGeminiStream', () => {
             expect(mockAddItem).toHaveBeenCalledWith({
                 type: 'info',
                 text: 'Loop detection has been disabled for this session. Retrying request...',
-            }, expect.any(Number));
+            });
             // Verify that the request was retried
             await waitFor(() => {
                 expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
@@ -1719,7 +1921,7 @@ describe('useGeminiStream', () => {
             expect(mockAddItem).toHaveBeenCalledWith({
                 type: 'info',
                 text: 'A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.',
-            }, expect.any(Number));
+            });
             // Verify that the request was NOT retried
             expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
         });
@@ -1749,7 +1951,7 @@ describe('useGeminiStream', () => {
             expect(mockAddItem).toHaveBeenCalledWith({
                 type: 'info',
                 text: 'A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.',
-            }, expect.any(Number));
+            });
             // Second loop detection - set up fresh mock for second call
             mockSendMessageStream.mockReturnValueOnce((async function* () {
                 yield {
@@ -1785,7 +1987,7 @@ describe('useGeminiStream', () => {
             expect(mockAddItem).toHaveBeenCalledWith({
                 type: 'info',
                 text: 'Loop detection has been disabled for this session. Retrying request...',
-            }, expect.any(Number));
+            });
             // Verify that the request was retried
             await waitFor(() => {
                 expect(mockSendMessageStream).toHaveBeenCalledTimes(3); // 1st query, 2nd query, retry of 2nd query
@@ -1819,5 +2021,87 @@ describe('useGeminiStream', () => {
             });
         });
     });
+    describe('Agent Execution Events', () => {
+        it('should handle AgentExecutionStopped event with systemMessage', async () => {
+            mockSendMessageStream.mockReturnValue((async function* () {
+                yield {
+                    type: ServerGeminiEventType.AgentExecutionStopped,
+                    value: {
+                        reason: 'hook-reason',
+                        systemMessage: 'Custom stop message',
+                    },
+                };
+            })());
+            const { result } = renderTestHook();
+            await act(async () => {
+                await result.current.submitQuery('test stop');
+            });
+            await waitFor(() => {
+                expect(mockAddItem).toHaveBeenCalledWith({
+                    type: MessageType.INFO,
+                    text: 'Agent execution stopped: Custom stop message',
+                }, expect.any(Number));
+                expect(result.current.streamingState).toBe(StreamingState.Idle);
+            });
+        });
+        it('should handle AgentExecutionStopped event by falling back to reason when systemMessage is missing', async () => {
+            mockSendMessageStream.mockReturnValue((async function* () {
+                yield {
+                    type: ServerGeminiEventType.AgentExecutionStopped,
+                    value: { reason: 'Stopped by hook' },
+                };
+            })());
+            const { result } = renderTestHook();
+            await act(async () => {
+                await result.current.submitQuery('test stop');
+            });
+            await waitFor(() => {
+                expect(mockAddItem).toHaveBeenCalledWith({
+                    type: MessageType.INFO,
+                    text: 'Agent execution stopped: Stopped by hook',
+                }, expect.any(Number));
+                expect(result.current.streamingState).toBe(StreamingState.Idle);
+            });
+        });
+        it('should handle AgentExecutionBlocked event with systemMessage', async () => {
+            mockSendMessageStream.mockReturnValue((async function* () {
+                yield {
+                    type: ServerGeminiEventType.AgentExecutionBlocked,
+                    value: {
+                        reason: 'hook-reason',
+                        systemMessage: 'Custom block message',
+                    },
+                };
+            })());
+            const { result } = renderTestHook();
+            await act(async () => {
+                await result.current.submitQuery('test block');
+            });
+            await waitFor(() => {
+                expect(mockAddItem).toHaveBeenCalledWith({
+                    type: MessageType.WARNING,
+                    text: 'Agent execution blocked: Custom block message',
+                }, expect.any(Number));
+            });
+        });
+        it('should handle AgentExecutionBlocked event by falling back to reason when systemMessage is missing', async () => {
+            mockSendMessageStream.mockReturnValue((async function* () {
+                yield {
+                    type: ServerGeminiEventType.AgentExecutionBlocked,
+                    value: { reason: 'Blocked by hook' },
+                };
+            })());
+            const { result } = renderTestHook();
+            await act(async () => {
+                await result.current.submitQuery('test block');
+            });
+            await waitFor(() => {
+                expect(mockAddItem).toHaveBeenCalledWith({
+                    type: MessageType.WARNING,
+                    text: 'Agent execution blocked: Blocked by hook',
+                }, expect.any(Number));
+            });
+        });
+    });
 });
 //# sourceMappingURL=useGeminiStream.test.js.map