@machina.ai/cell-cli 1.51.0-rc2 → 1.54.4-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, SHELL_TOOL_NAME, MCPDiscoveryState, GeminiCliOperation, getPlanModeExitMessage, UPDATE_TOPIC_TOOL_NAME, } from '@google/gemini-cli-core';
+import { CoreToolCallStatus, ApprovalMode, DEFAULT_GEMINI_FLASH_MODEL, AuthType, GeminiEventType as ServerGeminiEventType, ToolErrorType, ToolConfirmationOutcome, MessageBusType, tokenLimit, debugLogger, coreEvents, CoreEvent, SHELL_TOOL_NAME, MCPDiscoveryState, GeminiCliOperation, getPlanModeExitMessage, UPDATE_TOPIC_TOOL_NAME, TRUE_EMPTY_RESPONSE_MESSAGE, } from '@google/gemini-cli-core';
 import { MessageType, StreamingState } from '../types.js';
 import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js';
 import { theme } from '../semantic-colors.js';
@@ -737,6 +737,77 @@ describe('useGeminiStream', () => {
             expect(mockSendMessageStream).not.toHaveBeenCalled();
         });
     });
+    it('should record tool responses in history when the model was switched due to a quota error', async () => {
+        // Regression test: returning early on a quota-triggered model switch
+        // without recording the responses leaves the already-recorded
+        // functionCall unpaired, which corrupts all subsequent requests.
+        const responseParts = [
+            {
+                functionResponse: {
+                    name: 'testTool',
+                    id: 'call1',
+                    response: { output: 'tool result' },
+                },
+            },
+        ];
+        const completedToolCalls = [
+            {
+                request: {
+                    callId: 'call1',
+                    name: 'testTool',
+                    args: {},
+                    isClientInitiated: false,
+                    prompt_id: 'prompt-id-quota',
+                },
+                status: CoreToolCallStatus.Success,
+                responseSubmittedToGemini: false,
+                response: {
+                    callId: 'call1',
+                    responseParts,
+                    errorType: undefined,
+                },
+                tool: { displayName: 'MockTool' },
+                invocation: {
+                    getDescription: () => `Mock description`,
+                },
+            },
+        ];
+        const client = new MockedGeminiClientClass(mockConfig);
+        const mockConsumeUserHint = vi.fn(() => 'switch to the nprd database');
+        let capturedOnComplete = null;
+        mockUseToolScheduler.mockImplementation((onComplete) => {
+            capturedOnComplete = onComplete;
+            return [
+                [],
+                mockScheduleToolCalls,
+                mockMarkToolsAsSubmitted,
+                vi.fn(),
+                mockCancelAllToolCalls,
+                0,
+            ];
+        });
+        await renderHookWithProviders(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), true, // modelSwitchedFromQuotaError
+        () => { }, () => { }, () => { }, 80, 24, false, mockConsumeUserHint));
+        await act(async () => {
+            if (capturedOnComplete) {
+                await new Promise((resolve) => setTimeout(resolve, 0));
+                await capturedOnComplete(completedToolCalls);
+            }
+        });
+        await waitFor(() => {
+            expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['call1']);
+            // The tool response must be paired with its functionCall in history,
+            // with no steering-hint text ahead of it...
+            expect(client.addHistory).toHaveBeenCalledWith({
+                role: 'user',
+                parts: responseParts,
+            });
+            // ...the turn must NOT auto-continue on the fallback model...
+            expect(mockSendMessageStream).not.toHaveBeenCalled();
+            // ...and the pending hint is left for the next real submit.
+            expect(mockConsumeUserHint).not.toHaveBeenCalled();
+        });
+    });
     it('should NOT stop responding when only update_topic is called', async () => {
         const topicToolCalls = [
             {
@@ -1248,6 +1319,106 @@ describe('useGeminiStream', () => {
             // The cancel function should be called
             expect(mockCancelAllToolCalls).toHaveBeenCalled();
         });
+        it('should transition to Idle state when cancelled while a tool call is in progress and completes', async () => {
+            const toolCalls = [
+                {
+                    request: { callId: 'call1', name: 'tool1', args: {} },
+                    status: CoreToolCallStatus.Executing,
+                    responseSubmittedToGemini: false,
+                    tool: {
+                        name: 'tool1',
+                        description: 'desc1',
+                        build: vi.fn().mockImplementation((_) => ({
+                            getDescription: () => `Mock description`,
+                        })),
+                    },
+                    invocation: {
+                        getDescription: () => `Mock description`,
+                    },
+                    startTime: Date.now(),
+                    liveOutput: '...',
+                },
+            ];
+            const { result } = await renderTestHook(toolCalls);
+            // State is `Responding` because a tool is running
+            expect(result.current.streamingState).toBe(StreamingState.Responding);
+            // Try to cancel
+            simulateEscapeKeyPress();
+            // Trigger the onComplete callback with the cancelled tool call
+            await act(async () => {
+                if (capturedOnComplete) {
+                    await capturedOnComplete([
+                        {
+                            ...toolCalls[0],
+                            status: CoreToolCallStatus.Cancelled,
+                            response: {
+                                callId: 'call1',
+                                responseParts: [],
+                            },
+                        },
+                    ]);
+                }
+            });
+            // The final state should be idle because the cancelled tool call was marked as submitted
+            expect(result.current.streamingState).toBe(StreamingState.Idle);
+        });
+        it('should append cancelled tool responses to history when cancelled while a tool call is in progress and completes with response parts', async () => {
+            const toolCalls = [
+                {
+                    request: { callId: 'call1', name: 'tool1', args: {} },
+                    status: CoreToolCallStatus.Executing,
+                    responseSubmittedToGemini: false,
+                    tool: {
+                        name: 'tool1',
+                        description: 'desc1',
+                        build: vi.fn().mockImplementation((_) => ({
+                            getDescription: () => `Mock description`,
+                        })),
+                    },
+                    invocation: {
+                        getDescription: () => `Mock description`,
+                    },
+                    startTime: Date.now(),
+                    liveOutput: '...',
+                },
+            ];
+            const { result, client } = await renderTestHook(toolCalls);
+            // State is `Responding` because a tool is running
+            expect(result.current.streamingState).toBe(StreamingState.Responding);
+            // Try to cancel
+            simulateEscapeKeyPress();
+            const expectedResponseParts = [
+                {
+                    functionResponse: {
+                        name: 'tool1',
+                        id: 'call1',
+                        response: { error: 'cancelled' },
+                    },
+                },
+            ];
+            // Trigger the onComplete callback with the cancelled tool call having non-empty response parts
+            await act(async () => {
+                if (capturedOnComplete) {
+                    await capturedOnComplete([
+                        {
+                            ...toolCalls[0],
+                            status: CoreToolCallStatus.Cancelled,
+                            response: {
+                                callId: 'call1',
+                                responseParts: expectedResponseParts,
+                            },
+                        },
+                    ]);
+                }
+            });
+            // Assert that addHistory was called with the combined response parts
+            expect(client.addHistory).toHaveBeenCalledWith({
+                role: 'user',
+                parts: expectedResponseParts,
+            });
+            // The final state should be idle because the cancelled tool call was marked as submitted
+            expect(result.current.streamingState).toBe(StreamingState.Idle);
+        });
         it('should cancel a request when a tool is awaiting confirmation', async () => {
             const mockOnConfirm = vi.fn().mockResolvedValue(undefined);
             const toolCalls = [
@@ -1616,7 +1787,51 @@ describe('useGeminiStream', () => {
             });
             // 3. Assertion
             await waitFor(() => {
-                expect(mockParseAndFormatApiError).toHaveBeenCalledWith('Rate limit exceeded', mockAuthType, undefined, 'gemini-2.5-pro', 'gemini-2.5-flash');
+                expect(mockParseAndFormatApiError).toHaveBeenCalledWith('Rate limit exceeded', mockAuthType, undefined, 'gemini-2.5-pro', DEFAULT_GEMINI_FLASH_MODEL);
+            });
+        });
+        it('should use TRUE_EMPTY_RESPONSE_MESSAGE when receiving an invalid stream event of type NO_RESPONSE_TEXT', async () => {
+            mockSendMessageStream.mockClear();
+            mockSendMessageStream.mockReturnValue((async function* () {
+                yield {
+                    type: ServerGeminiEventType.InvalidStream,
+                    value: {
+                        type: 'NO_RESPONSE_TEXT',
+                        message: 'empty response text',
+                    },
+                };
+            })());
+            const { result } = await renderTestHook();
+            await act(async () => {
+                await result.current.submitQuery('test query');
+            });
+            await waitFor(() => {
+                expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
+                    type: MessageType.ERROR,
+                    text: TRUE_EMPTY_RESPONSE_MESSAGE,
+                }), expect.any(Number));
+            });
+        });
+        it('should use the event message when receiving a non-NO_RESPONSE_TEXT invalid stream event', async () => {
+            mockSendMessageStream.mockClear();
+            mockSendMessageStream.mockReturnValue((async function* () {
+                yield {
+                    type: ServerGeminiEventType.InvalidStream,
+                    value: {
+                        type: 'MALFORMED_FUNCTION_CALL',
+                        message: 'Custom malformed function call message',
+                    },
+                };
+            })());
+            const { result } = await renderTestHook();
+            await act(async () => {
+                await result.current.submitQuery('test query');
+            });
+            await waitFor(() => {
+                expect(mockAddItem).toHaveBeenCalledWith(expect.objectContaining({
+                    type: MessageType.ERROR,
+                    text: 'Custom malformed function call message',
+                }), expect.any(Number));
             });
         });
     });
@@ -2369,7 +2584,7 @@ describe('useGeminiStream', () => {
                 }), expect.any(Number));
             });
             // Verify parseAndFormatApiError was called
-            expect(mockParseAndFormatApiError).toHaveBeenCalledWith({ message: 'Test error' }, expect.any(String), undefined, 'gemini-2.5-pro', 'gemini-2.5-flash');
+            expect(mockParseAndFormatApiError).toHaveBeenCalledWith({ message: 'Test error' }, expect.any(String), undefined, 'gemini-2.5-pro', DEFAULT_GEMINI_FLASH_MODEL);
         });
         it('should update lastOutputTime on Gemini thought and content events', async () => {
             vi.useFakeTimers();