summaryrefslogtreecommitdiff
path: root/packages/core/src
diff options
context:
space:
mode:
Diffstat (limited to 'packages/core/src')
-rw-r--r--packages/core/src/core/client.test.ts178
-rw-r--r--packages/core/src/core/client.ts6
2 files changed, 182 insertions, 2 deletions
diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 6a2ab307..b5f64543 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<GeminiChat> = {
+ addHistory: vi.fn(),
+ getHistory: vi.fn().mockReturnValue([]),
+ };
+ client['chat'] = mockChat as GeminiChat;
+
+ const mockGenerator: Partial<ContentGenerator> = {
+ 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<GeminiChat> = {
+ addHistory: vi.fn(),
+ getHistory: vi.fn().mockReturnValue([]),
+ };
+ client['chat'] = mockChat as GeminiChat;
+
+ const mockGenerator: Partial<ContentGenerator> = {
+ 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 9a74add6..b00a689b 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<ServerGeminiStreamEvent, Turn> {
- 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;