File size: 43,948 Bytes

7a1ad33

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type { Content } from '@google/genai';
import type { Config } from '../config/config.js';
import type { GeminiClient } from '../core/client.js';
import type { BaseLlmClient } from '../core/baseLlmClient.js';
import {
  GeminiEventType,
  type ServerGeminiContentEvent,
  type ServerGeminiStreamEvent,
  type ServerGeminiToolCallRequestEvent,
} from '../core/turn.js';
import * as loggers from '../telemetry/loggers.js';
import { LoopType } from '../telemetry/types.js';
import { LoopDetectionService } from './loopDetectionService.js';
import { createAvailabilityServiceMock } from '../availability/testUtils.js';

vi.mock('../telemetry/loggers.js', () => ({
  logLoopDetected: vi.fn(),
  logLoopDetectionDisabled: vi.fn(),
  logLlmLoopCheck: vi.fn(),
}));

const TOOL_CALL_LOOP_THRESHOLD = 5;
const CONTENT_LOOP_THRESHOLD = 10;
const CONTENT_CHUNK_SIZE = 50;

describe('LoopDetectionService', () => {
  let service: LoopDetectionService;
  let mockConfig: Config;

  beforeEach(() => {
    mockConfig = {
      get config() {
        return this;
      },
      getTelemetryEnabled: () => true,
      isInteractive: () => false,
      getDisableLoopDetection: () => false,
      getModelAvailabilityService: vi
        .fn()
        .mockReturnValue(createAvailabilityServiceMock()),
    } as unknown as Config;
    service = new LoopDetectionService(mockConfig);
    vi.clearAllMocks();
  });

  const createToolCallRequestEvent = (
    name: string,
    args: Record<string, unknown>,
  ): ServerGeminiToolCallRequestEvent => ({
    type: GeminiEventType.ToolCallRequest,
    value: {
      name,
      args,
      callId: 'test-id',
      isClientInitiated: false,
      prompt_id: 'test-prompt-id',
    },
  });

  const createContentEvent = (content: string): ServerGeminiContentEvent => ({
    type: GeminiEventType.Content,
    value: content,
  });

  const createRepetitiveContent = (id: number, length: number): string => {
    const baseString = `This is a unique sentence, id=${id}. `;
    let content = '';
    while (content.length < length) {
      content += baseString;
    }
    return content.slice(0, length);
  };

  describe('Tool Call Loop Detection', () => {
    it(`should not detect a loop for fewer than TOOL_CALL_LOOP_THRESHOLD identical calls`, () => {
      const event = createToolCallRequestEvent('testTool', { param: 'value' });
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) {
        expect(service.addAndCheck(event).count).toBe(0);
      }
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it(`should detect a loop on the TOOL_CALL_LOOP_THRESHOLD-th identical call`, () => {
      const event = createToolCallRequestEvent('testTool', { param: 'value' });
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) {
        service.addAndCheck(event);
      }
      expect(service.addAndCheck(event).count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should detect a loop on subsequent identical calls', () => {
      const event = createToolCallRequestEvent('testTool', { param: 'value' });
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
        service.addAndCheck(event);
      }
      expect(service.addAndCheck(event).count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should not detect a loop for different tool calls', () => {
      const event1 = createToolCallRequestEvent('testTool', {
        param: 'value1',
      });
      const event2 = createToolCallRequestEvent('testTool', {
        param: 'value2',
      });
      const event3 = createToolCallRequestEvent('anotherTool', {
        param: 'value1',
      });

      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 2; i++) {
        expect(service.addAndCheck(event1).count).toBe(0);
        expect(service.addAndCheck(event2).count).toBe(0);
        expect(service.addAndCheck(event3).count).toBe(0);
      }
    });

    it('should not reset tool call counter for other event types', () => {
      const toolCallEvent = createToolCallRequestEvent('testTool', {
        param: 'value',
      });
      const otherEvent = {
        type: 'thought',
      } as unknown as ServerGeminiStreamEvent;

      // Send events just below the threshold
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD - 1; i++) {
        expect(service.addAndCheck(toolCallEvent).count).toBe(0);
      }

      // Send a different event type
      expect(service.addAndCheck(otherEvent).count).toBe(0);

      // Send the tool call event again, which should now trigger the loop
      expect(service.addAndCheck(toolCallEvent).count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should not detect a loop when disabled for session', () => {
      service.disableForSession();
      expect(loggers.logLoopDetectionDisabled).toHaveBeenCalledTimes(1);
      const event = createToolCallRequestEvent('testTool', { param: 'value' });
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
        expect(service.addAndCheck(event).count).toBe(0);
      }
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should stop reporting a loop if disabled after detection', () => {
      const event = createToolCallRequestEvent('testTool', { param: 'value' });
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
        service.addAndCheck(event);
      }
      expect(service.addAndCheck(event).count).toBe(1);

      service.disableForSession();

      // Should now return 0 even though a loop was previously detected
      expect(service.addAndCheck(event).count).toBe(0);
    });

    it('should skip loop detection if disabled in config', () => {
      vi.spyOn(mockConfig, 'getDisableLoopDetection').mockReturnValue(true);
      const event = createToolCallRequestEvent('testTool', { param: 'value' });
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD + 2; i++) {
        expect(service.addAndCheck(event).count).toBe(0);
      }
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should detect an alternating tool call loop (cycle of length 2)', () => {
      const eventA = createToolCallRequestEvent('read_file', {
        file_path: 'loop_a.txt',
      });
      const eventB = createToolCallRequestEvent('read_file', {
        file_path: 'loop_b.txt',
      });

      let loopCount = 0;
      for (let i = 0; i < 15; i++) {
        const currentEvent = i % 2 === 0 ? eventA : eventB;
        const result = service.addAndCheck(currentEvent);
        if (result.count > 0) {
          loopCount = result.count;
        }
      }

      expect(loopCount).toBeGreaterThan(0);
    });

    it('should detect a cyclic tool call loop of length 3', () => {
      const eventA = createToolCallRequestEvent('read_file', {
        file_path: 'loop_a.txt',
      });
      const eventB = createToolCallRequestEvent('read_file', {
        file_path: 'loop_b.txt',
      });
      const eventC = createToolCallRequestEvent('read_file', {
        file_path: 'loop_c.txt',
      });

      let loopCount = 0;
      const sequence = [eventA, eventB, eventC];
      for (let i = 0; i < 20; i++) {
        const result = service.addAndCheck(sequence[i % 3]);
        if (result.count > 0) {
          loopCount = result.count;
        }
      }

      expect(loopCount).toBeGreaterThan(0);
    });

    it('should detect a cyclic tool call loop of length 5', () => {
      const e1 = createToolCallRequestEvent('t', { id: 1 });
      const e2 = createToolCallRequestEvent('t', { id: 2 });
      const e3 = createToolCallRequestEvent('t', { id: 3 });
      const e4 = createToolCallRequestEvent('t', { id: 4 });
      const e5 = createToolCallRequestEvent('t', { id: 5 });

      let loopCount = 0;
      const sequence = [e1, e2, e3, e4, e5];
      for (let i = 0; i < 30; i++) {
        const result = service.addAndCheck(sequence[i % 5]);
        if (result.count > 0) {
          loopCount = result.count;
        }
      }

      expect(loopCount).toBeGreaterThan(0);
    });

    it('should not detect loops for non-looping alternating sequences that vary', () => {
      const eventA = createToolCallRequestEvent('read_file', {
        file_path: 'loop_a.txt',
      });
      const eventB = createToolCallRequestEvent('read_file', {
        file_path: 'loop_b.txt',
      });
      const eventC = createToolCallRequestEvent('read_file', {
        file_path: 'loop_c.txt',
      });

      // Run some alternating calls, then break the pattern with eventC
      for (let i = 0; i < 4; i++) {
        expect(service.addAndCheck(i % 2 === 0 ? eventA : eventB).count).toBe(
          0,
        );
      }
      expect(service.addAndCheck(eventC).count).toBe(0);
      for (let i = 0; i < 4; i++) {
        expect(service.addAndCheck(i % 2 === 0 ? eventA : eventB).count).toBe(
          0,
        );
      }
    });
  });

  describe('Content Loop Detection', () => {
    const generateRandomString = (length: number) => {
      let result = '';
      const characters =
        'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
      const charactersLength = characters.length;
      for (let i = 0; i < length; i++) {
        result += characters.charAt(
          Math.floor(Math.random() * charactersLength),
        );
      }
      return result;
    };

    it('should not detect a loop for random content', () => {
      service.reset('');
      for (let i = 0; i < 1000; i++) {
        const content = generateRandomString(10);
        const result = service.addAndCheck(createContentEvent(content));
        expect(result.count).toBe(0);
      }
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should detect a loop when a chunk of content repeats consecutively', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
        result = service.addAndCheck(createContentEvent(repeatedContent));
      }
      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should not detect a loop for a list with a long shared prefix', () => {
      service.reset('');
      let result = { count: 0 };
      const longPrefix =
        'projects/my-google-cloud-project-12345/locations/us-central1/services/';

      let listContent = '';
      for (let i = 0; i < 15; i++) {
        listContent += `- ${longPrefix}${i}\n`;
      }

      // Simulate receiving the list in a single large chunk or a few chunks
      // This is the specific case where the issue occurs, as list boundaries might not reset tracking properly
      result = service.addAndCheck(createContentEvent(listContent));

      expect(result.count).toBe(0);
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should not detect a loop if repetitions are very far apart', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);
      const fillerContent = generateRandomString(500);

      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
        result = service.addAndCheck(createContentEvent(repeatedContent));
        result = service.addAndCheck(createContentEvent(fillerContent));
      }
      expect(result.count).toBe(0);
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should detect a loop with longer repeating patterns (e.g. ~150 chars)', () => {
      service.reset('');
      const longPattern = createRepetitiveContent(1, 150);
      expect(longPattern.length).toBe(150);

      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 2; i++) {
        result = service.addAndCheck(createContentEvent(longPattern));
        if (result.count > 0) break;
      }
      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should detect the specific user-provided loop example', () => {
      service.reset('');
      const userPattern = `I will not output any text.
  I will just end the turn.
  I am done.
  I will not do anything else.
  I will wait for the user's next command.
`;

      let result = { count: 0 };
      // Loop enough times to trigger the threshold
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
        result = service.addAndCheck(createContentEvent(userPattern));
        if (result.count > 0) break;
      }
      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should detect the second specific user-provided loop example', () => {
      service.reset('');
      const userPattern =
        'I have added all the requested logs and verified the test file. I will now mark the task as complete.\n  ';

      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
        result = service.addAndCheck(createContentEvent(userPattern));
        if (result.count > 0) break;
      }
      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should detect a loop of alternating short phrases', () => {
      service.reset('');
      const alternatingPattern = 'Thinking... Done. ';

      let result = { count: 0 };
      // Needs more iterations because the pattern is short relative to chunk size,
      // so it takes a few slides of the window to find the exact alignment.
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD * 3; i++) {
        result = service.addAndCheck(createContentEvent(alternatingPattern));
        if (result.count > 0) break;
      }
      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should detect a loop of repeated complex thought processes', () => {
      service.reset('');
      const thoughtPattern =
        'I need to check the file. The file does not exist. I will create the file. ';

      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
        result = service.addAndCheck(createContentEvent(thoughtPattern));
        if (result.count > 0) break;
      }
      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });
  });

  describe('Content Loop Detection with Code Blocks', () => {
    it('should not detect a loop when repetitive content is inside a code block', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      service.addAndCheck(createContentEvent('```\n'));

      for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
        const result = service.addAndCheck(createContentEvent(repeatedContent));
        expect(result.count).toBe(0);
      }

      const result = service.addAndCheck(createContentEvent('\n```'));
      expect(result.count).toBe(0);
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should not detect loops when content transitions into a code block', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      // Add some repetitive content outside of code block
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 2; i++) {
        service.addAndCheck(createContentEvent(repeatedContent));
      }

      // Now transition into a code block - this should prevent loop detection
      // even though we were already close to the threshold
      const codeBlockStart = '```javascript\n';
      const result = service.addAndCheck(createContentEvent(codeBlockStart));
      expect(result.count).toBe(0);

      // Continue adding repetitive content inside the code block - should not trigger loop
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
        const resultInside = service.addAndCheck(
          createContentEvent(repeatedContent),
        );
        expect(resultInside.count).toBe(0);
      }

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should skip loop detection when already inside a code block (this.inCodeBlock)', () => {
      service.reset('');

      // Start with content that puts us inside a code block
      service.addAndCheck(createContentEvent('Here is some code:\n```\n'));

      // Verify we are now inside a code block and any content should be ignored for loop detection
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
        const result = service.addAndCheck(createContentEvent(repeatedContent));
        expect(result.count).toBe(0);
      }

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should correctly track inCodeBlock state with multiple fence transitions', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      // Outside code block - should track content
      service.addAndCheck(createContentEvent('Normal text '));

      // Enter code block (1 fence) - should stop tracking
      const enterResult = service.addAndCheck(createContentEvent('```\n'));
      expect(enterResult.count).toBe(0);

      // Inside code block - should not track loops
      for (let i = 0; i < 5; i++) {
        const insideResult = service.addAndCheck(
          createContentEvent(repeatedContent),
        );
        expect(insideResult.count).toBe(0);
      }

      // Exit code block (2nd fence) - should reset tracking but still return false
      const exitResult = service.addAndCheck(createContentEvent('```\n'));
      expect(exitResult.count).toBe(0);

      // Enter code block again (3rd fence) - should stop tracking again
      const reenterResult = service.addAndCheck(
        createContentEvent('```python\n'),
      );
      expect(reenterResult.count).toBe(0);

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should detect a loop when repetitive content is outside a code block', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      service.addAndCheck(createContentEvent('```'));
      service.addAndCheck(createContentEvent('\nsome code\n'));
      service.addAndCheck(createContentEvent('```'));

      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
        result = service.addAndCheck(createContentEvent(repeatedContent));
      }
      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should handle content with multiple code blocks and no loops', () => {
      service.reset('');
      service.addAndCheck(createContentEvent('```\ncode1\n```'));
      service.addAndCheck(createContentEvent('\nsome text\n'));
      const result = service.addAndCheck(createContentEvent('```\ncode2\n```'));

      expect(result.count).toBe(0);
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should handle content with mixed code blocks and looping text', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      service.addAndCheck(createContentEvent('```'));
      service.addAndCheck(createContentEvent('\ncode1\n'));
      service.addAndCheck(createContentEvent('```'));

      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
        result = service.addAndCheck(createContentEvent(repeatedContent));
      }

      expect(result.count).toBe(1);
      expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
    });

    it('should not detect a loop for a long code block with some repeating tokens', () => {
      service.reset('');
      const repeatingTokens =
        'for (let i = 0; i < 10; i++) { console.log(i); }';

      service.addAndCheck(createContentEvent('```\n'));

      for (let i = 0; i < 20; i++) {
        const result = service.addAndCheck(createContentEvent(repeatingTokens));
        expect(result.count).toBe(0);
      }

      const result = service.addAndCheck(createContentEvent('\n```'));
      expect(result.count).toBe(0);
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should reset tracking when a code fence is found', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        service.addAndCheck(createContentEvent(repeatedContent));
      }

      // This should not trigger a loop because of the reset
      service.addAndCheck(createContentEvent('```'));

      // We are now in a code block, so loop detection should be off.
      // Let's add the repeated content again, it should not trigger a loop.
      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD; i++) {
        result = service.addAndCheck(createContentEvent(repeatedContent));
        expect(result.count).toBe(0);
      }

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });
    it('should reset tracking when a table is detected', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        service.addAndCheck(createContentEvent(repeatedContent));
      }

      // This should reset tracking and not trigger a loop
      service.addAndCheck(createContentEvent('| Column 1 | Column 2 |'));

      // Add more repeated content after table - should not trigger loop
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        const result = service.addAndCheck(createContentEvent(repeatedContent));
        expect(result.count).toBe(0);
      }

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should reset tracking when a list item is detected', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        service.addAndCheck(createContentEvent(repeatedContent));
      }

      // This should reset tracking and not trigger a loop
      service.addAndCheck(createContentEvent('* List item'));

      // Add more repeated content after list - should not trigger loop
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        const result = service.addAndCheck(createContentEvent(repeatedContent));
        expect(result.count).toBe(0);
      }

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should reset tracking when a heading is detected', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        service.addAndCheck(createContentEvent(repeatedContent));
      }

      // This should reset tracking and not trigger a loop
      service.addAndCheck(createContentEvent('## Heading'));

      // Add more repeated content after heading - should not trigger loop
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        const result = service.addAndCheck(createContentEvent(repeatedContent));
        expect(result.count).toBe(0);
      }

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should reset tracking when a blockquote is detected', () => {
      service.reset('');
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        service.addAndCheck(createContentEvent(repeatedContent));
      }

      // This should reset tracking and not trigger a loop
      service.addAndCheck(createContentEvent('> Quote text'));

      // Add more repeated content after blockquote - should not trigger loop
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
        const result = service.addAndCheck(createContentEvent(repeatedContent));
        expect(result.count).toBe(0);
      }

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should reset tracking for various list item formats', () => {
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      // Test different list formats - make sure they start at beginning of line
      const listFormats = [
        '* Bullet item',
        '- Dash item',
        '+ Plus item',
        '1. Numbered item',
        '42. Another numbered item',
      ];

      listFormats.forEach((listFormat, index) => {
        service.reset('');

        // Build up to near threshold
        for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
          service.addAndCheck(createContentEvent(repeatedContent));
        }

        // Reset should occur with list item - add newline to ensure it starts at beginning
        service.addAndCheck(createContentEvent('\n' + listFormat));

        // Should not trigger loop after reset - use different content to avoid any cached state issues
        const newRepeatedContent = createRepetitiveContent(
          index + 100,
          CONTENT_CHUNK_SIZE,
        );
        for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
          const result = service.addAndCheck(
            createContentEvent(newRepeatedContent),
          );
          expect(result.count).toBe(0);
        }
      });

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should reset tracking for various table formats', () => {
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      const tableFormats = [
        '| Column 1 | Column 2 |',
        '|---|---|',
        '|++|++|',
        '+---+---+',
      ];

      tableFormats.forEach((tableFormat, index) => {
        service.reset('');

        // Build up to near threshold
        for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
          service.addAndCheck(createContentEvent(repeatedContent));
        }

        // Reset should occur with table format - add newline to ensure it starts at beginning
        service.addAndCheck(createContentEvent('\n' + tableFormat));

        // Should not trigger loop after reset - use different content to avoid any cached state issues
        const newRepeatedContent = createRepetitiveContent(
          index + 200,
          CONTENT_CHUNK_SIZE,
        );
        for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
          const result = service.addAndCheck(
            createContentEvent(newRepeatedContent),
          );
          expect(result.count).toBe(0);
        }
      });

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should reset tracking for various heading levels', () => {
      const repeatedContent = createRepetitiveContent(1, CONTENT_CHUNK_SIZE);

      const headingFormats = [
        '# H1 Heading',
        '## H2 Heading',
        '### H3 Heading',
        '#### H4 Heading',
        '##### H5 Heading',
        '###### H6 Heading',
      ];

      headingFormats.forEach((headingFormat, index) => {
        service.reset('');

        // Build up to near threshold
        for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
          service.addAndCheck(createContentEvent(repeatedContent));
        }

        // Reset should occur with heading - add newline to ensure it starts at beginning
        service.addAndCheck(createContentEvent('\n' + headingFormat));

        // Should not trigger loop after reset - use different content to avoid any cached state issues
        const newRepeatedContent = createRepetitiveContent(
          index + 300,
          CONTENT_CHUNK_SIZE,
        );
        for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
          const result = service.addAndCheck(
            createContentEvent(newRepeatedContent),
          );
          expect(result.count).toBe(0);
        }
      });

      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });
  });

  describe('Edge Cases', () => {
    it('should handle empty content', () => {
      const event = createContentEvent('');
      expect(service.addAndCheck(event).count).toBe(0);
    });
  });

  describe('Divider Content Detection', () => {
    it('should not detect a loop for repeating divider-like content', () => {
      service.reset('');
      const dividerContent = '-'.repeat(CONTENT_CHUNK_SIZE);
      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
        result = service.addAndCheck(createContentEvent(dividerContent));
        expect(result.count).toBe(0);
      }
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });

    it('should not detect a loop for repeating complex box-drawing dividers', () => {
      service.reset('');
      const dividerContent = '╭─'.repeat(CONTENT_CHUNK_SIZE / 2);
      let result = { count: 0 };
      for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 5; i++) {
        result = service.addAndCheck(createContentEvent(dividerContent));
        expect(result.count).toBe(0);
      }
      expect(loggers.logLoopDetected).not.toHaveBeenCalled();
    });
  });

  describe('Strike Management', () => {
    it('should increment strike count for repeated detections', () => {
      const event = createToolCallRequestEvent('testTool', { param: 'value' });

      // First strike
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
        service.addAndCheck(event);
      }
      expect(service.addAndCheck(event).count).toBe(1);

      // Recovery simulated by caller calling clearDetection()
      service.clearDetection();

      // Second strike
      expect(service.addAndCheck(event).count).toBe(2);
    });

    it('should allow recovery turn to proceed after clearDetection', () => {
      const event = createToolCallRequestEvent('testTool', { param: 'value' });

      // Trigger loop
      for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) {
        service.addAndCheck(event);
      }
      expect(service.addAndCheck(event).count).toBe(1);

      // Caller clears detection to allow recovery
      service.clearDetection();

      // Subsequent call in the same turn (or next turn before it repeats) should be 0
      // In reality, addAndCheck is called per event.
      // If the model sends a NEW event, it should not immediately trigger.
      const newEvent = createContentEvent('Recovery text');
      expect(service.addAndCheck(newEvent).count).toBe(0);
    });
  });

  describe('Reset Functionality', () => {
    it('tool call should reset content count', () => {
      const contentEvent = createContentEvent('Some content.');
      const toolEvent = createToolCallRequestEvent('testTool', {
        param: 'value',
      });
      for (let i = 0; i < 9; i++) {
        service.addAndCheck(contentEvent);
      }

      service.addAndCheck(toolEvent);

      // Should start fresh
      expect(
        service.addAndCheck(createContentEvent('Fresh content.')).count,
      ).toBe(0);
    });
  });

  describe('General Behavior', () => {
    it('should return 0 count for unhandled event types', () => {
      const otherEvent = {
        type: 'unhandled_event',
      } as unknown as ServerGeminiStreamEvent;
      expect(service.addAndCheck(otherEvent).count).toBe(0);
      expect(service.addAndCheck(otherEvent).count).toBe(0);
    });
  });
});

describe('LoopDetectionService LLM Checks', () => {
  let service: LoopDetectionService;
  let mockConfig: Config;
  let mockGeminiClient: GeminiClient;
  let mockBaseLlmClient: BaseLlmClient;
  let abortController: AbortController;

  beforeEach(() => {
    mockGeminiClient = {
      getHistory: vi.fn().mockReturnValue([]),
    } as unknown as GeminiClient;

    mockBaseLlmClient = {
      generateJson: vi.fn(),
    } as unknown as BaseLlmClient;

    const mockAvailability = createAvailabilityServiceMock();
    vi.mocked(mockAvailability.snapshot).mockReturnValue({ available: true });

    mockConfig = {
      get config() {
        return this;
      },
      getGeminiClient: () => mockGeminiClient,
      get geminiClient() {
        return mockGeminiClient;
      },
      getBaseLlmClient: () => mockBaseLlmClient,
      getDisableLoopDetection: () => false,
      getDebugMode: () => false,
      getTelemetryEnabled: () => true,
      getModel: vi.fn().mockReturnValue('cognitive-loop-v1'),
      modelConfigService: {
        getResolvedConfig: vi.fn().mockImplementation((key) => {
          if (key.model === 'loop-detection') {
            return { model: 'gemini-2.5-flash', generateContentConfig: {} };
          }
          return {
            model: 'cognitive-loop-v1',
            generateContentConfig: {},
          };
        }),
      },
      isInteractive: () => false,
      getModelAvailabilityService: vi.fn().mockReturnValue(mockAvailability),
    } as unknown as Config;

    service = new LoopDetectionService(mockConfig);
    abortController = new AbortController();
    vi.clearAllMocks();
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  const advanceTurns = async (count: number) => {
    for (let i = 0; i < count; i++) {
      await service.turnStarted(abortController.signal);
    }
  };

  it('should not trigger LLM check before LLM_CHECK_AFTER_TURNS (30)', async () => {
    await advanceTurns(29);
    expect(mockBaseLlmClient.generateJson).not.toHaveBeenCalled();
  });

  it('should trigger LLM check on the 30th turn', async () => {
    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockResolvedValue({ unproductive_state_confidence: 0.1 });
    await advanceTurns(30);
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledWith(
      expect.objectContaining({
        modelConfigKey: { model: 'loop-detection' },
        systemInstruction: expect.any(String),
        contents: expect.any(Array),
        schema: expect.any(Object),
        promptId: expect.any(String),
      }),
    );
  });

  it('should detect a cognitive loop when confidence is high', async () => {
    // First check at turn 30
    mockBaseLlmClient.generateJson = vi.fn().mockResolvedValue({
      unproductive_state_confidence: 0.85,
      unproductive_state_analysis: 'Repetitive actions',
    });
    await advanceTurns(30);
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledWith(
      expect.objectContaining({
        modelConfigKey: { model: 'loop-detection' },
      }),
    );

    // The confidence of 0.85 will result in a low interval.
    // The interval will be: 5 + (15 - 5) * (1 - 0.85) = 5 + 10 * 0.15 = 6.5 -> rounded to 7
    await advanceTurns(6); // advance to turn 36

    mockBaseLlmClient.generateJson = vi.fn().mockResolvedValue({
      unproductive_state_confidence: 0.95,
      unproductive_state_analysis: 'Repetitive actions',
    });
    const finalResult = await service.turnStarted(abortController.signal); // This is turn 37

    expect(finalResult.count).toBe(1);
    expect(loggers.logLoopDetected).toHaveBeenCalledWith(
      mockConfig,
      expect.objectContaining({
        'event.name': 'loop_detected',
        loop_type: LoopType.LLM_DETECTED_LOOP,
        confirmed_by_model: 'cognitive-loop-v1',
      }),
    );
  });

  it('should not detect a loop when confidence is low', async () => {
    mockBaseLlmClient.generateJson = vi.fn().mockResolvedValue({
      unproductive_state_confidence: 0.5,
      unproductive_state_analysis: 'Looks okay',
    });
    await advanceTurns(30);
    const result = await service.turnStarted(abortController.signal);
    expect(result.count).toBe(0);
    expect(loggers.logLoopDetected).not.toHaveBeenCalled();
  });

  it('should adjust the check interval based on confidence', async () => {
    // Confidence is 0.0, so interval should be MAX_LLM_CHECK_INTERVAL (15)
    // Interval = 5 + (15 - 5) * (1 - 0.0) = 15
    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockResolvedValue({ unproductive_state_confidence: 0.0 });
    await advanceTurns(30); // First check at turn 30
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);

    await advanceTurns(14); // Advance to turn 44
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);

    await service.turnStarted(abortController.signal); // Turn 45
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(2);
  });

  it('should handle errors from generateJson gracefully', async () => {
    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockRejectedValue(new Error('API error'));
    await advanceTurns(30);
    const result = await service.turnStarted(abortController.signal);
    expect(result.count).toBe(0);
    expect(loggers.logLoopDetected).not.toHaveBeenCalled();
  });

  it('should not trigger LLM check when disabled for session', async () => {
    service.disableForSession();
    expect(loggers.logLoopDetectionDisabled).toHaveBeenCalledTimes(1);
    await advanceTurns(30);
    const result = await service.turnStarted(abortController.signal);
    expect(result.count).toBe(0);
    expect(mockBaseLlmClient.generateJson).not.toHaveBeenCalled();
  });

  it('should prepend user message if history starts with a function call', async () => {
    const functionCallHistory: Content[] = [
      {
        role: 'model',
        parts: [{ functionCall: { name: 'someTool', args: {} } }],
      },
      {
        role: 'model',
        parts: [{ text: 'Some follow up text' }],
      },
    ];
    vi.mocked(mockGeminiClient.getHistory).mockReturnValue(functionCallHistory);

    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockResolvedValue({ unproductive_state_confidence: 0.1 });

    await advanceTurns(30);

    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
    const calledArg = vi.mocked(mockBaseLlmClient.generateJson).mock
      .calls[0][0];
    expect(calledArg.contents[0]).toEqual({
      role: 'user',
      parts: [{ text: 'Recent conversation history:' }],
    });
    // Verify the original history follows
    expect(calledArg.contents[1]).toEqual(functionCallHistory[0]);
  });

  it('should detect a loop when confidence is exactly equal to the threshold (0.9)', async () => {
    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockResolvedValueOnce({
        unproductive_state_confidence: 0.9,
        unproductive_state_analysis: 'Flash says loop',
      })
      .mockResolvedValueOnce({
        unproductive_state_confidence: 0.9,
        unproductive_state_analysis: 'Main says loop',
      });

    await advanceTurns(30);

    // It should have called generateJson twice
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(2);
    expect(mockBaseLlmClient.generateJson).toHaveBeenNthCalledWith(
      1,
      expect.objectContaining({
        modelConfigKey: { model: 'loop-detection' },
      }),
    );
    expect(mockBaseLlmClient.generateJson).toHaveBeenNthCalledWith(
      2,
      expect.objectContaining({
        modelConfigKey: { model: 'loop-detection-double-check' },
      }),
    );

    // And it should have detected a loop
    expect(loggers.logLoopDetected).toHaveBeenCalledWith(
      mockConfig,
      expect.objectContaining({
        'event.name': 'loop_detected',
        loop_type: LoopType.LLM_DETECTED_LOOP,
        confirmed_by_model: 'cognitive-loop-v1',
      }),
    );
  });

  it('should not detect a loop when Flash is confident (0.9) but Main model is not (0.89)', async () => {
    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockResolvedValueOnce({
        unproductive_state_confidence: 0.9,
        unproductive_state_analysis: 'Flash says loop',
      })
      .mockResolvedValueOnce({
        unproductive_state_confidence: 0.89,
        unproductive_state_analysis: 'Main says no loop',
      });

    await advanceTurns(30);

    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(2);
    expect(mockBaseLlmClient.generateJson).toHaveBeenNthCalledWith(
      1,
      expect.objectContaining({
        modelConfigKey: { model: 'loop-detection' },
      }),
    );
    expect(mockBaseLlmClient.generateJson).toHaveBeenNthCalledWith(
      2,
      expect.objectContaining({
        modelConfigKey: { model: 'loop-detection-double-check' },
      }),
    );

    // Should NOT have detected a loop
    expect(loggers.logLoopDetected).not.toHaveBeenCalled();

    // But should have updated the interval based on the main model's confidence (0.89)
    // Interval = 5 + (15-5) * (1 - 0.89) = 5 + 10 * 0.11 = 5 + 1.1 = 6.1 -> 6

    // Advance by 5 turns
    await advanceTurns(5);

    // Next turn (36) should trigger another check
    await service.turnStarted(abortController.signal);
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(3);
  });

  it('should only call Flash model if main model is unavailable', async () => {
    // Mock availability to return unavailable for the main model
    const availability = mockConfig.getModelAvailabilityService();
    vi.mocked(availability.snapshot).mockReturnValue({
      available: false,
      reason: 'quota',
    });

    mockBaseLlmClient.generateJson = vi.fn().mockResolvedValueOnce({
      unproductive_state_confidence: 0.9,
      unproductive_state_analysis: 'Flash says loop',
    });

    await advanceTurns(30);

    // It should have called generateJson only once
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledWith(
      expect.objectContaining({
        modelConfigKey: { model: 'loop-detection' },
      }),
    );

    // And it should have detected a loop
    expect(loggers.logLoopDetected).toHaveBeenCalledWith(
      mockConfig,
      expect.objectContaining({
        confirmed_by_model: 'gemini-2.5-flash',
      }),
    );
  });

  it('should include user prompt in LLM check contents when provided', async () => {
    service.reset('test-prompt-id', 'Add license headers to all files');

    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockResolvedValue({ unproductive_state_confidence: 0.1 });

    await advanceTurns(30);

    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
    const calledArg = vi.mocked(mockBaseLlmClient.generateJson).mock
      .calls[0][0];
    // First content should be the user prompt context wrapped in XML
    expect(calledArg.contents[0]).toEqual({
      role: 'user',
      parts: [
        {
          text: '<original_user_request>\nAdd license headers to all files\n</original_user_request>',
        },
      ],
    });
  });

  it('should not include user prompt in contents when not provided', async () => {
    service.reset('test-prompt-id');

    vi.mocked(mockGeminiClient.getHistory).mockReturnValue([
      {
        role: 'model',
        parts: [{ text: 'Some response' }],
      },
    ]);

    mockBaseLlmClient.generateJson = vi
      .fn()
      .mockResolvedValue({ unproductive_state_confidence: 0.1 });

    await advanceTurns(30);

    expect(mockBaseLlmClient.generateJson).toHaveBeenCalledTimes(1);
    const calledArg = vi.mocked(mockBaseLlmClient.generateJson).mock
      .calls[0][0];
    // First content should be the history, not a user prompt message
    expect(calledArg.contents[0]).toEqual({
      role: 'model',
      parts: [{ text: 'Some response' }],
    });
  });
});