diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 6a2ab307ce..b5f645430e 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -402,5 +402,183 @@ describe('Gemini Client (client.ts)', () => { // Assert expect(finalResult).toBeInstanceOf(Turn); }); + + it('should stop infinite loop after MAX_TURNS when nextSpeaker always returns model', async () => { + // Get the mocked checkNextSpeaker function and configure it to trigger infinite loop + const { checkNextSpeaker } = await import( + '../utils/nextSpeakerChecker.js' + ); + const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker); + mockCheckNextSpeaker.mockResolvedValue({ + next_speaker: 'model', + reasoning: 'Test case - always continue', + }); + + // Mock Turn to have no pending tool calls (which would allow nextSpeaker check) + const mockStream = (async function* () { + yield { type: 'content', value: 'Continue...' }; + })(); + mockTurnRunFn.mockReturnValue(mockStream); + + const mockChat: Partial = { + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([]), + }; + client['chat'] = mockChat as GeminiChat; + + const mockGenerator: Partial = { + countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }), + }; + client['contentGenerator'] = mockGenerator as ContentGenerator; + + // Use a signal that never gets aborted + const abortController = new AbortController(); + const signal = abortController.signal; + + // Act - Start the stream that should loop + const stream = client.sendMessageStream( + [{ text: 'Start conversation' }], + signal, + ); + + // Count how many stream events we get + let eventCount = 0; + let finalResult: Turn | undefined; + + // Consume the stream and count iterations + while (true) { + const result = await stream.next(); + if (result.done) { + finalResult = result.value; + break; + } + eventCount++; + + // Safety check to prevent actual infinite loop in test + if (eventCount > 200) { + abortController.abort(); + throw new Error( + 'Test exceeded expected event limit - possible actual infinite loop', + ); + } + } + + // Assert + expect(finalResult).toBeInstanceOf(Turn); + + // Debug: Check how many times checkNextSpeaker was called + const callCount = mockCheckNextSpeaker.mock.calls.length; + + // If infinite loop protection is working, checkNextSpeaker should be called many times + // but stop at MAX_TURNS (100). Since each recursive call should trigger checkNextSpeaker, + // we expect it to be called multiple times before hitting the limit + expect(mockCheckNextSpeaker).toHaveBeenCalled(); + + // The test should demonstrate that the infinite loop protection works: + // - If checkNextSpeaker is called many times (close to MAX_TURNS), it shows the loop was happening + // - If it's only called once, the recursive behavior might not be triggered + if (callCount === 0) { + throw new Error( + 'checkNextSpeaker was never called - the recursive condition was not met', + ); + } else if (callCount === 1) { + // This might be expected behavior if the turn has pending tool calls or other conditions prevent recursion + console.log( + 'checkNextSpeaker called only once - no infinite loop occurred', + ); + } else { + console.log( + `checkNextSpeaker called ${callCount} times - infinite loop protection worked`, + ); + // If called multiple times, we expect it to be stopped before MAX_TURNS + expect(callCount).toBeLessThanOrEqual(100); // Should not exceed MAX_TURNS + } + + // The stream should produce events and eventually terminate + expect(eventCount).toBeGreaterThanOrEqual(1); + expect(eventCount).toBeLessThan(200); // Should not exceed our safety limit + }); + + it('should respect MAX_TURNS limit even when turns parameter is set to a large value', async () => { + // This test verifies that the infinite loop protection works even when + // someone tries to bypass it by calling with a very large turns value + + // Get the mocked checkNextSpeaker function and configure it to trigger infinite loop + const { checkNextSpeaker } = await import( + '../utils/nextSpeakerChecker.js' + ); + const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker); + mockCheckNextSpeaker.mockResolvedValue({ + next_speaker: 'model', + reasoning: 'Test case - always continue', + }); + + // Mock Turn to have no pending tool calls (which would allow nextSpeaker check) + const mockStream = (async function* () { + yield { type: 'content', value: 'Continue...' }; + })(); + mockTurnRunFn.mockReturnValue(mockStream); + + const mockChat: Partial = { + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([]), + }; + client['chat'] = mockChat as GeminiChat; + + const mockGenerator: Partial = { + countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }), + }; + client['contentGenerator'] = mockGenerator as ContentGenerator; + + // Use a signal that never gets aborted + const abortController = new AbortController(); + const signal = abortController.signal; + + // Act - Start the stream with an extremely high turns value + // This simulates a case where the turns protection is bypassed + const stream = client.sendMessageStream( + [{ text: 'Start conversation' }], + signal, + Number.MAX_SAFE_INTEGER, // Bypass the MAX_TURNS protection + ); + + // Count how many stream events we get + let eventCount = 0; + const maxTestIterations = 1000; // Higher limit to show the loop continues + + // Consume the stream and count iterations + try { + while (true) { + const result = await stream.next(); + if (result.done) { + break; + } + eventCount++; + + // This test should hit this limit, demonstrating the infinite loop + if (eventCount > maxTestIterations) { + abortController.abort(); + // This is the expected behavior - we hit the infinite loop + break; + } + } + } catch (error) { + // If the test framework times out, that also demonstrates the infinite loop + console.error('Test timed out or errored:', error); + } + + // Assert that the fix works - the loop should stop at MAX_TURNS + const callCount = mockCheckNextSpeaker.mock.calls.length; + + // With the fix: even when turns is set to a very high value, + // the loop should stop at MAX_TURNS (100) + expect(callCount).toBeLessThanOrEqual(100); // Should not exceed MAX_TURNS + expect(eventCount).toBeLessThanOrEqual(200); // Should have reasonable number of events + + console.log( + `Infinite loop protection working: checkNextSpeaker called ${callCount} times, ` + + `${eventCount} events generated (properly bounded by MAX_TURNS)`, + ); + }); }); }); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 9a74add624..b00a689b29 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -219,7 +219,9 @@ export class GeminiClient { signal: AbortSignal, turns: number = this.MAX_TURNS, ): AsyncGenerator { - if (!turns) { + // Ensure turns never exceeds MAX_TURNS to prevent infinite loops + const boundedTurns = Math.min(turns, this.MAX_TURNS); + if (!boundedTurns) { return new Turn(this.getChat()); } @@ -242,7 +244,7 @@ export class GeminiClient { const nextRequest = [{ text: 'Please continue.' }]; // This recursive call's events will be yielded out, but the final // turn object will be from the top-level call. - yield* this.sendMessageStream(nextRequest, signal, turns - 1); + yield* this.sendMessageStream(nextRequest, signal, boundedTurns - 1); } } return turn;