@machina.ai/cell-cli 1.60.0-rc1 → 1.64.0-rc2

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

@@ -36,6 +36,17 @@ const MockedGeminiClientClass = vi.hoisted(() => vi.fn().mockImplementation(func
     this.setHistory = vi.fn().mockImplementation((newHistory) => {
         mockHistory = [...newHistory];
     });
+    this.discardTrailingUnansweredToolCallTurn = vi
+        .fn()
+        .mockImplementation(() => {
+        const last = mockHistory[mockHistory.length - 1];
+        if (last?.role !== 'model' ||
+            !last.parts?.some((part) => !!part.functionCall)) {
+            return false;
+        }
+        mockHistory = mockHistory.slice(0, -1);
+        return true;
+    });
     this.generateContent = vi.fn().mockResolvedValue({
         candidates: [
             { content: { parts: [{ text: 'Got it. Focusing on tests only.' }] } },
@@ -788,7 +799,7 @@ describe('useGeminiStream', () => {
             ];
         });
         const { result } = await renderHookWithProviders(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, 80, 24));
-        // Call submitQuery to populate the user turn and set historyLengthAfterUserPromptRef
+        // Call submitQuery to populate the user turn
         await act(async () => {
             // eslint-disable-next-line @typescript-eslint/no-floating-promises
             result.current.submitQuery('User prompt');
@@ -819,6 +830,127 @@ describe('useGeminiStream', () => {
             expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
         });
     });
+    it('should keep the user prompt and completed tool rounds when a later tool batch is declined', async () => {
+        const cancelledToolCalls = [
+            {
+                request: {
+                    callId: '2',
+                    name: 'testTool',
+                    args: {},
+                    isClientInitiated: false,
+                    prompt_id: 'prompt-id-4',
+                },
+                status: CoreToolCallStatus.Cancelled,
+                response: {
+                    callId: '2',
+                    responseParts: [{ text: CoreToolCallStatus.Cancelled }],
+                    errorType: undefined,
+                },
+                responseSubmittedToGemini: false,
+                tool: {
+                    displayName: 'mock tool',
+                },
+                invocation: {
+                    getDescription: () => `Mock description`,
+                },
+            },
+        ];
+        const client = new MockedGeminiClientClass(mockConfig);
+        const priorTurn = [
+            { role: 'user', parts: [{ text: 'Earlier prompt' }] },
+            { role: 'model', parts: [{ text: 'Earlier answer' }] },
+        ];
+        client.setHistory(priorTurn);
+        // Model the real sendMessageStream contract: the user turn is recorded in
+        // the client history when the request is sent, not before submitQuery.
+        mockSendMessageStream.mockImplementation((query) => {
+            const parts = (Array.isArray(query) ? query : [query]).map((part) => typeof part === 'string' ? { text: part } : part);
+            client.setHistory([...client.getHistory(), { role: 'user', parts }]);
+            return (async function* () {
+                yield { type: ServerGeminiEventType.Content, value: 'Working on it' };
+            })();
+        });
+        let capturedOnComplete = null;
+        mockUseToolScheduler.mockImplementation((onComplete) => {
+            capturedOnComplete = onComplete;
+            return [
+                [],
+                mockScheduleToolCalls,
+                mockMarkToolsAsSubmitted,
+                vi.fn(),
+                mockCancelAllToolCalls,
+                0,
+            ];
+        });
+        const { result } = await renderHookWithProviders(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, 80, 24));
+        // Turn 1 starts: sendMessageStream records the user prompt.
+        await act(async () => {
+            await result.current.submitQuery('User prompt');
+        });
+        // The model answers with a first tool call, which completes and is sent
+        // back to the model as a continuation turn.
+        const firstRoundResponse = [
+            {
+                functionResponse: {
+                    name: 'testTool',
+                    id: '1',
+                    response: { output: 'first result' },
+                },
+            },
+        ];
+        client.setHistory([
+            ...client.getHistory(),
+            {
+                role: 'model',
+                parts: [{ functionCall: { name: 'testTool', id: '1', args: {} } }],
+            },
+        ]);
+        await act(async () => {
+            await result.current.submitQuery(firstRoundResponse, {
+                isContinuation: true,
+            });
+        });
+        // The model requests a second tool call within the same chain.
+        client.setHistory([
+            ...client.getHistory(),
+            {
+                role: 'model',
+                parts: [{ functionCall: { name: 'testTool', id: '2', args: {} } }],
+            },
+        ]);
+        vi.mocked(client.setHistory).mockClear();
+        // The second tool call is declined.
+        await act(async () => {
+            if (capturedOnComplete) {
+                await new Promise((resolve) => setTimeout(resolve, 0));
+                await capturedOnComplete(cancelledToolCalls);
+            }
+        });
+        try {
+            await waitFor(() => {
+                expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['2']);
+                expect(client.addHistory).not.toHaveBeenCalled();
+                // Only the trailing unanswered call is removed, without re-setting
+                // the whole history (which would re-record every turn).
+                expect(client.discardTrailingUnansweredToolCallTurn).toHaveBeenCalledTimes(1);
+                expect(client.setHistory).not.toHaveBeenCalled();
+                // The previous turn, the user prompt and the completed first round
+                // are preserved.
+                expect(client.getHistory()).toEqual([
+                    ...priorTurn,
+                    { role: 'user', parts: [{ text: 'User prompt' }] },
+                    {
+                        role: 'model',
+                        parts: [{ functionCall: { name: 'testTool', id: '1', args: {} } }],
+                    },
+                    { role: 'user', parts: firstRoundResponse },
+                ]);
+            });
+        }
+        finally {
+            mockSendMessageStream.mockImplementation(() => (async function* () { })());
+        }
+    });
     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
@@ -1129,7 +1261,7 @@ describe('useGeminiStream', () => {
             ];
         });
         const { result } = await renderHookWithProviders(() => useGeminiStream(client, [], mockAddItem, mockConfig, mockLoadedSettings, mockOnDebugMessage, mockHandleSlashCommand, false, () => 'vscode', () => { }, () => Promise.resolve(), false, () => { }, () => { }, () => { }, 80, 24));
-        // Call submitQuery to populate the user turn and set historyLengthAfterUserPromptRef
+        // Call submitQuery to populate the user turn
         await act(async () => {
             // eslint-disable-next-line @typescript-eslint/no-floating-promises
             result.current.submitQuery('User prompt');