@machina.ai/cell-cli 1.36.0-rc1 → 1.38.1-rc2

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

@@ -12,7 +12,7 @@ import { useGeminiStream } from './useGeminiStream.js';
 import { useKeypress } from './useKeypress.js';
 import * as atCommandProcessor from './atCommandProcessor.js';
 import { useToolScheduler, } from './useToolScheduler.js';
-import { CoreToolCallStatus, ApprovalMode, AuthType, GeminiEventType as ServerGeminiEventType, ToolErrorType, ToolConfirmationOutcome, MessageBusType, tokenLimit, debugLogger, coreEvents, CoreEvent, MCPDiscoveryState, GeminiCliOperation, getPlanModeExitMessage, } from '@google/gemini-cli-core';
+import { CoreToolCallStatus, ApprovalMode, AuthType, GeminiEventType as ServerGeminiEventType, ToolErrorType, ToolConfirmationOutcome, MessageBusType, tokenLimit, debugLogger, coreEvents, CoreEvent, MCPDiscoveryState, GeminiCliOperation, getPlanModeExitMessage, UPDATE_TOPIC_TOOL_NAME, } from '@google/gemini-cli-core';
 import { MessageType, StreamingState } from '../types.js';
 import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js';
 import { theme } from '../semantic-colors.js';
@@ -103,33 +103,57 @@ vi.mock('./useToolScheduler.js', async (importOriginal) => {
 vi.mock('./useKeypress.js', () => ({
     useKeypress: vi.fn(),
 }));
-vi.mock('./shellCommandProcessor.js', () => ({
-    useShellCommandProcessor: vi.fn().mockReturnValue({
+vi.mock('./useExecutionLifecycle.js', () => ({
+    useExecutionLifecycle: vi.fn().mockReturnValue({
         handleShellCommand: vi.fn(),
         activeShellPtyId: null,
         lastShellOutputTime: 0,
+        backgroundTaskCount: 0,
+        isBackgroundTaskVisible: false,
+        toggleBackgroundTasks: vi.fn(),
+        backgroundCurrentExecution: vi.fn(),
+        backgroundTasks: new Map(),
+        dismissBackgroundTask: vi.fn(),
+        registerBackgroundTask: vi.fn(),
     }),
 }));
 vi.mock('./atCommandProcessor.js');
 vi.mock('../utils/markdownUtilities.js', () => ({
     findLastSafeSplitPoint: vi.fn((s) => s.length),
 }));
-vi.mock('./useStateAndRef.js', () => ({
-    useStateAndRef: vi.fn((initial) => {
-        let val = initial;
-        const ref = { current: val };
-        const setVal = vi.fn((updater) => {
-            if (typeof updater === 'function') {
-                val = updater(val);
-            }
-            else {
-                val = updater;
+vi.mock('./useStateAndRef.js', async () => {
+    const React = await vi.importActual('react');
+    return {
+        useStateAndRef: vi.fn((initial) => {
+            // Keep the heavyweight test file lightweight, but still let
+            // `isResponding` participate in real rerenders.
+            if (initial === false) {
+                const [state, setState] = React.useState(initial);
+                const ref = React.useRef(initial);
+                const setStateInternal = (updater) => {
+                    const nextValue = typeof updater === 'function'
+                        ? updater(ref.current)
+                        : updater;
+                    ref.current = nextValue;
+                    setState(nextValue);
+                };
+                return [state, ref, setStateInternal];
             }
-            ref.current = val;
-        });
-        return [val, ref, setVal];
-    }),
-}));
+            let val = initial;
+            const ref = { current: val };
+            const setVal = vi.fn((updater) => {
+                if (typeof updater === 'function') {
+                    val = updater(val);
+                }
+                else {
+                    val = updater;
+                }
+                ref.current = val;
+            });
+            return [val, ref, setVal];
+        }),
+    };
+});
 vi.mock('./useLogger.js', () => ({
     useLogger: vi.fn().mockReturnValue({
         logMessage: vi.fn().mockResolvedValue(undefined),
@@ -625,6 +649,30 @@ describe('useGeminiStream', () => {
     });
     it('should handle all tool calls being cancelled', async () => {
         const cancelledToolCalls = [
+            {
+                request: {
+                    callId: 'topic1',
+                    name: UPDATE_TOPIC_TOOL_NAME,
+                    args: {},
+                    isClientInitiated: false,
+                    prompt_id: 'prompt-id-3',
+                },
+                status: CoreToolCallStatus.Success,
+                response: {
+                    callId: 'topic1',
+                    responseParts: [
+                        {
+                            functionResponse: {
+                                name: UPDATE_TOPIC_TOOL_NAME,
+                                id: 'topic1',
+                                response: {},
+                            },
+                        },
+                    ],
+                },
+                tool: { displayName: 'Update Topic Context' },
+                invocation: { getDescription: () => 'Updating topic' },
+            },
             {
                 request: {
                     callId: '1',
@@ -672,15 +720,80 @@ describe('useGeminiStream', () => {
             }
         });
         await waitFor(() => {
-            expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['1']);
+            expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['topic1', '1']);
             expect(client.addHistory).toHaveBeenCalledWith({
                 role: 'user',
-                parts: [{ text: CoreToolCallStatus.Cancelled }],
+                parts: [
+                    {
+                        functionResponse: {
+                            name: UPDATE_TOPIC_TOOL_NAME,
+                            id: 'topic1',
+                            response: {},
+                        },
+                    },
+                    { text: CoreToolCallStatus.Cancelled },
+                ],
             });
             // Ensure we do NOT call back to the API
             expect(mockSendMessageStream).not.toHaveBeenCalled();
         });
     });
+    it('should NOT stop responding when only update_topic is called', async () => {
+        const topicToolCalls = [
+            {
+                request: {
+                    callId: 'topic1',
+                    name: UPDATE_TOPIC_TOOL_NAME,
+                    args: {},
+                    isClientInitiated: false,
+                    prompt_id: 'prompt-id-3',
+                },
+                status: CoreToolCallStatus.Success,
+                response: {
+                    callId: 'topic1',
+                    responseParts: [
+                        {
+                            functionResponse: {
+                                name: UPDATE_TOPIC_TOOL_NAME,
+                                id: 'topic1',
+                                response: {},
+                            },
+                        },
+                    ],
+                },
+                tool: { displayName: 'Update Topic Context' },
+                invocation: { getDescription: () => 'Updating topic' },
+            },
+        ];
+        const client = new MockedGeminiClientClass(mockConfig);
+        // Capture the onComplete callback
+        let capturedOnComplete = null;
+        mockUseToolScheduler.mockImplementation((onComplete) => {
+            capturedOnComplete = onComplete;
+            return [
+                topicToolCalls,
+                vi.fn(),
+                mockMarkToolsAsSubmitted,
+                vi.fn(),
+                vi.fn(),
+                0,
+            ];
+        });
+        await renderHookWithProviders(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, 80, 24));
+        // Trigger the onComplete callback with the topic tool
+        await act(async () => {
+            if (capturedOnComplete) {
+                await capturedOnComplete(topicToolCalls);
+            }
+        });
+        await waitFor(() => {
+            // The streaming state should still be Responding because we didn't cancel anything important
+            // and we expect a continuation.
+            expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['topic1']);
+            // Should HAVE called back to the API for continuation
+            expect(mockSendMessageStream).toHaveBeenCalled();
+        });
+    });
     it('should stop agent execution immediately when a tool call returns STOP_EXECUTION error', async () => {
         const stopExecutionToolCalls = [
             {
@@ -1189,7 +1302,7 @@ describe('useGeminiStream', () => {
         });
     });
     describe('Retry Handling', () => {
-        it('should update retryStatus when CoreEvent.RetryAttempt is emitted', async () => {
+        it('should ignore retryStatus updates when not responding', async () => {
             const { result } = await renderHookWithDefaults();
             const retryPayload = {
                 model: 'gemini-2.5-pro',
@@ -1200,7 +1313,7 @@ describe('useGeminiStream', () => {
             await act(async () => {
                 coreEvents.emit(CoreEvent.RetryAttempt, retryPayload);
             });
-            expect(result.current.retryStatus).toEqual(retryPayload);
+            expect(result.current.retryStatus).toBeNull();
         });
         it('should reset retryStatus when isResponding becomes false', async () => {
             const { result } = await renderTestHook();
@@ -1234,6 +1347,47 @@ describe('useGeminiStream', () => {
             });
             expect(result.current.retryStatus).toBeNull();
         });
+        it('should ignore late retry events after cancellation', async () => {
+            const { result } = await renderTestHook();
+            const retryPayload = {
+                model: 'gemini-2.5-pro',
+                attempt: 2,
+                maxAttempts: 3,
+                delayMs: 1000,
+            };
+            const lateRetryPayload = {
+                model: 'gemini-2.5-pro',
+                attempt: 3,
+                maxAttempts: 3,
+                delayMs: 2000,
+            };
+            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);
+            });
+            await act(async () => {
+                coreEvents.emit(CoreEvent.RetryAttempt, retryPayload);
+            });
+            expect(result.current.retryStatus).toEqual(retryPayload);
+            await act(async () => {
+                result.current.cancelOngoingRequest();
+            });
+            await waitFor(() => {
+                expect(result.current.retryStatus).toBeNull();
+            });
+            await act(async () => {
+                coreEvents.emit(CoreEvent.RetryAttempt, lateRetryPayload);
+            });
+            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 () => {