mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-04-24 22:55:13 +00:00
Merge branch 'main' into feat-copy-command-for-history-copy
This commit is contained in:
@@ -13,6 +13,7 @@ import { clearCommand } from '../ui/commands/clearCommand.js';
|
||||
import { copyCommand } from '../ui/commands/copyCommand.js';
|
||||
import { authCommand } from '../ui/commands/authCommand.js';
|
||||
import { themeCommand } from '../ui/commands/themeCommand.js';
|
||||
import { privacyCommand } from '../ui/commands/privacyCommand.js';
|
||||
|
||||
// Mock the command modules to isolate the service from the command implementations.
|
||||
vi.mock('../ui/commands/memoryCommand.js', () => ({
|
||||
@@ -33,6 +34,9 @@ vi.mock('../ui/commands/themeCommand.js', () => ({
|
||||
vi.mock('../ui/commands/copyCommand.js', () => ({
|
||||
copyCommand: { name: 'copy', description: 'Mock Copy' },
|
||||
}));
|
||||
vi.mock('../ui/commands/privacyCommand.js', () => ({
|
||||
privacyCommand: { name: 'privacy', description: 'Mock Privacy' },
|
||||
}));
|
||||
|
||||
describe('CommandService', () => {
|
||||
describe('when using default production loader', () => {
|
||||
@@ -67,6 +71,7 @@ describe('CommandService', () => {
|
||||
expect(commandNames).toContain('clear');
|
||||
expect(commandNames).toContain('copy');
|
||||
expect(commandNames).toContain('theme');
|
||||
expect(commandNames).toContain('privacy');
|
||||
});
|
||||
|
||||
it('should overwrite any existing commands when called again', async () => {
|
||||
@@ -98,6 +103,7 @@ describe('CommandService', () => {
|
||||
copyCommand,
|
||||
helpCommand,
|
||||
memoryCommand,
|
||||
privacyCommand,
|
||||
themeCommand,
|
||||
]);
|
||||
});
|
||||
|
||||
@@ -11,6 +11,7 @@ import { clearCommand } from '../ui/commands/clearCommand.js';
|
||||
import { copyCommand } from '../ui/commands/copyCommand.js';
|
||||
import { authCommand } from '../ui/commands/authCommand.js';
|
||||
import { themeCommand } from '../ui/commands/themeCommand.js';
|
||||
import { privacyCommand } from '../ui/commands/privacyCommand.js';
|
||||
|
||||
const loadBuiltInCommands = async (): Promise<SlashCommand[]> => [
|
||||
authCommand,
|
||||
@@ -18,6 +19,7 @@ const loadBuiltInCommands = async (): Promise<SlashCommand[]> => [
|
||||
copyCommand,
|
||||
helpCommand,
|
||||
memoryCommand,
|
||||
privacyCommand,
|
||||
themeCommand,
|
||||
];
|
||||
|
||||
|
||||
38
packages/cli/src/ui/commands/privacyCommand.test.ts
Normal file
38
packages/cli/src/ui/commands/privacyCommand.test.ts
Normal file
@@ -0,0 +1,38 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach } from 'vitest';
|
||||
import { privacyCommand } from './privacyCommand.js';
|
||||
import { type CommandContext } from './types.js';
|
||||
import { createMockCommandContext } from '../../test-utils/mockCommandContext.js';
|
||||
|
||||
describe('privacyCommand', () => {
|
||||
let mockContext: CommandContext;
|
||||
|
||||
beforeEach(() => {
|
||||
mockContext = createMockCommandContext();
|
||||
});
|
||||
|
||||
it('should return a dialog action to open the privacy dialog', () => {
|
||||
// Ensure the command has an action to test.
|
||||
if (!privacyCommand.action) {
|
||||
throw new Error('The privacy command must have an action.');
|
||||
}
|
||||
|
||||
const result = privacyCommand.action(mockContext, '');
|
||||
|
||||
// Assert that the action returns the correct object to trigger the privacy dialog.
|
||||
expect(result).toEqual({
|
||||
type: 'dialog',
|
||||
dialog: 'privacy',
|
||||
});
|
||||
});
|
||||
|
||||
it('should have the correct name and description', () => {
|
||||
expect(privacyCommand.name).toBe('privacy');
|
||||
expect(privacyCommand.description).toBe('display the privacy notice');
|
||||
});
|
||||
});
|
||||
16
packages/cli/src/ui/commands/privacyCommand.ts
Normal file
16
packages/cli/src/ui/commands/privacyCommand.ts
Normal file
@@ -0,0 +1,16 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { OpenDialogActionReturn, SlashCommand } from './types.js';
|
||||
|
||||
export const privacyCommand: SlashCommand = {
|
||||
name: 'privacy',
|
||||
description: 'display the privacy notice',
|
||||
action: (): OpenDialogActionReturn => ({
|
||||
type: 'dialog',
|
||||
dialog: 'privacy',
|
||||
}),
|
||||
};
|
||||
@@ -66,7 +66,7 @@ export interface MessageActionReturn {
|
||||
export interface OpenDialogActionReturn {
|
||||
type: 'dialog';
|
||||
// TODO: Add 'theme' | 'auth' | 'editor' | 'privacy' as migration happens.
|
||||
dialog: 'help' | 'auth' | 'theme';
|
||||
dialog: 'help' | 'auth' | 'theme' | 'privacy';
|
||||
}
|
||||
|
||||
export type SlashCommandActionReturn =
|
||||
|
||||
@@ -247,11 +247,6 @@ export const useSlashCommandProcessor = (
|
||||
description: 'set external editor preference',
|
||||
action: (_mainCommand, _subCommand, _args) => openEditorDialog(),
|
||||
},
|
||||
{
|
||||
name: 'privacy',
|
||||
description: 'display the privacy notice',
|
||||
action: (_mainCommand, _subCommand, _args) => openPrivacyNotice(),
|
||||
},
|
||||
{
|
||||
name: 'stats',
|
||||
altName: 'usage',
|
||||
@@ -1023,7 +1018,6 @@ export const useSlashCommandProcessor = (
|
||||
}, [
|
||||
addMessage,
|
||||
openEditorDialog,
|
||||
openPrivacyNotice,
|
||||
toggleCorgiMode,
|
||||
savedChatTags,
|
||||
config,
|
||||
@@ -1125,6 +1119,9 @@ export const useSlashCommandProcessor = (
|
||||
case 'theme':
|
||||
openThemeDialog();
|
||||
return { type: 'handled' };
|
||||
case 'privacy':
|
||||
openPrivacyNotice();
|
||||
return { type: 'handled' };
|
||||
default: {
|
||||
const unhandled: never = result.dialog;
|
||||
throw new Error(
|
||||
@@ -1208,6 +1205,7 @@ export const useSlashCommandProcessor = (
|
||||
commandContext,
|
||||
addMessage,
|
||||
openThemeDialog,
|
||||
openPrivacyNotice,
|
||||
],
|
||||
);
|
||||
|
||||
|
||||
@@ -141,6 +141,8 @@ export const useGeminiStream = (
|
||||
[toolCalls],
|
||||
);
|
||||
|
||||
const loopDetectedRef = useRef(false);
|
||||
|
||||
const onExec = useCallback(async (done: Promise<void>) => {
|
||||
setIsResponding(true);
|
||||
await done;
|
||||
@@ -450,6 +452,16 @@ export const useGeminiStream = (
|
||||
[addItem, config],
|
||||
);
|
||||
|
||||
const handleLoopDetectedEvent = useCallback(() => {
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
}, [addItem]);
|
||||
|
||||
const processGeminiStreamEvents = useCallback(
|
||||
async (
|
||||
stream: AsyncIterable<GeminiEvent>,
|
||||
@@ -489,6 +501,11 @@ export const useGeminiStream = (
|
||||
case ServerGeminiEventType.MaxSessionTurns:
|
||||
handleMaxSessionTurnsEvent();
|
||||
break;
|
||||
case ServerGeminiEventType.LoopDetected:
|
||||
// handle later because we want to move pending history to history
|
||||
// before we add loop detected message to history
|
||||
loopDetectedRef.current = true;
|
||||
break;
|
||||
default: {
|
||||
// enforces exhaustive switch-case
|
||||
const unreachable: never = event;
|
||||
@@ -579,6 +596,10 @@ export const useGeminiStream = (
|
||||
addItem(pendingHistoryItemRef.current, userMessageTimestamp);
|
||||
setPendingHistoryItem(null);
|
||||
}
|
||||
if (loopDetectedRef.current) {
|
||||
loopDetectedRef.current = false;
|
||||
handleLoopDetectedEvent();
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
if (error instanceof UnauthorizedError) {
|
||||
onAuthError();
|
||||
@@ -616,6 +637,7 @@ export const useGeminiStream = (
|
||||
config,
|
||||
startNewPrompt,
|
||||
getPromptCount,
|
||||
handleLoopDetectedEvent,
|
||||
],
|
||||
);
|
||||
|
||||
|
||||
@@ -40,6 +40,7 @@ import {
|
||||
} from './contentGenerator.js';
|
||||
import { ProxyAgent, setGlobalDispatcher } from 'undici';
|
||||
import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
|
||||
import { LoopDetectionService } from '../services/loopDetectionService.js';
|
||||
|
||||
function isThinkingSupported(model: string) {
|
||||
if (model.startsWith('gemini-2.5')) return true;
|
||||
@@ -100,12 +101,16 @@ export class GeminiClient {
|
||||
*/
|
||||
private readonly COMPRESSION_PRESERVE_THRESHOLD = 0.3;
|
||||
|
||||
private readonly loopDetector: LoopDetectionService;
|
||||
private lastPromptId?: string;
|
||||
|
||||
constructor(private config: Config) {
|
||||
if (config.getProxy()) {
|
||||
setGlobalDispatcher(new ProxyAgent(config.getProxy() as string));
|
||||
}
|
||||
|
||||
this.embeddingModel = config.getEmbeddingModel();
|
||||
this.loopDetector = new LoopDetectionService(config);
|
||||
}
|
||||
|
||||
async initialize(contentGeneratorConfig: ContentGeneratorConfig) {
|
||||
@@ -272,6 +277,10 @@ export class GeminiClient {
|
||||
turns: number = this.MAX_TURNS,
|
||||
originalModel?: string,
|
||||
): AsyncGenerator<ServerGeminiStreamEvent, Turn> {
|
||||
if (this.lastPromptId !== prompt_id) {
|
||||
this.loopDetector.reset();
|
||||
this.lastPromptId = prompt_id;
|
||||
}
|
||||
this.sessionTurnCount++;
|
||||
if (
|
||||
this.config.getMaxSessionTurns() > 0 &&
|
||||
@@ -297,6 +306,10 @@ export class GeminiClient {
|
||||
const turn = new Turn(this.getChat(), prompt_id);
|
||||
const resultStream = turn.run(request, signal);
|
||||
for await (const event of resultStream) {
|
||||
if (this.loopDetector.addAndCheck(event)) {
|
||||
yield { type: GeminiEventType.LoopDetected };
|
||||
return turn;
|
||||
}
|
||||
yield event;
|
||||
}
|
||||
if (!turn.pendingToolCalls.length && signal && !signal.aborted) {
|
||||
|
||||
@@ -49,6 +49,7 @@ export enum GeminiEventType {
|
||||
ChatCompressed = 'chat_compressed',
|
||||
Thought = 'thought',
|
||||
MaxSessionTurns = 'max_session_turns',
|
||||
LoopDetected = 'loop_detected',
|
||||
}
|
||||
|
||||
export interface StructuredError {
|
||||
@@ -133,6 +134,10 @@ export type ServerGeminiMaxSessionTurnsEvent = {
|
||||
type: GeminiEventType.MaxSessionTurns;
|
||||
};
|
||||
|
||||
export type ServerGeminiLoopDetectedEvent = {
|
||||
type: GeminiEventType.LoopDetected;
|
||||
};
|
||||
|
||||
// The original union type, now composed of the individual types
|
||||
export type ServerGeminiStreamEvent =
|
||||
| ServerGeminiContentEvent
|
||||
@@ -143,7 +148,8 @@ export type ServerGeminiStreamEvent =
|
||||
| ServerGeminiErrorEvent
|
||||
| ServerGeminiChatCompressedEvent
|
||||
| ServerGeminiThoughtEvent
|
||||
| ServerGeminiMaxSessionTurnsEvent;
|
||||
| ServerGeminiMaxSessionTurnsEvent
|
||||
| ServerGeminiLoopDetectedEvent;
|
||||
|
||||
// A turn manages the agentic loop turn within the server context.
|
||||
export class Turn {
|
||||
|
||||
311
packages/core/src/services/loopDetectionService.test.ts
Normal file
311
packages/core/src/services/loopDetectionService.test.ts
Normal file
@@ -0,0 +1,311 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, beforeEach, vi } from 'vitest';
|
||||
import { LoopDetectionService } from './loopDetectionService.js';
|
||||
import {
|
||||
GeminiEventType,
|
||||
ServerGeminiContentEvent,
|
||||
ServerGeminiToolCallRequestEvent,
|
||||
} from '../core/turn.js';
|
||||
import { ServerGeminiStreamEvent } from '../core/turn.js';
|
||||
import { Config } from '../config/config.js';
|
||||
import * as loggers from '../telemetry/loggers.js';
|
||||
|
||||
vi.mock('../telemetry/loggers.js', () => ({
|
||||
logLoopDetected: vi.fn(),
|
||||
}));
|
||||
|
||||
const TOOL_CALL_LOOP_THRESHOLD = 5;
|
||||
const CONTENT_LOOP_THRESHOLD = 10;
|
||||
|
||||
describe('LoopDetectionService', () => {
|
||||
let service: LoopDetectionService;
|
||||
let mockConfig: Config;
|
||||
|
||||
beforeEach(() => {
|
||||
mockConfig = {
|
||||
getTelemetryEnabled: () => true,
|
||||
} 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,
|
||||
});
|
||||
|
||||
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)).toBe(false);
|
||||
}
|
||||
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)).toBe(true);
|
||||
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)).toBe(true);
|
||||
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
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)).toBe(false);
|
||||
expect(service.addAndCheck(event2)).toBe(false);
|
||||
expect(service.addAndCheck(event3)).toBe(false);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Content Loop Detection', () => {
|
||||
it(`should not detect a loop for fewer than CONTENT_LOOP_THRESHOLD identical content strings`, () => {
|
||||
const event = createContentEvent('This is a test sentence.');
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
|
||||
expect(service.addAndCheck(event)).toBe(false);
|
||||
}
|
||||
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it(`should detect a loop on the CONTENT_LOOP_THRESHOLD-th identical content string`, () => {
|
||||
const event = createContentEvent('This is a test sentence.');
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
|
||||
service.addAndCheck(event);
|
||||
}
|
||||
expect(service.addAndCheck(event)).toBe(true);
|
||||
expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should not detect a loop for different content strings', () => {
|
||||
const event1 = createContentEvent('Sentence A');
|
||||
const event2 = createContentEvent('Sentence B');
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 2; i++) {
|
||||
expect(service.addAndCheck(event1)).toBe(false);
|
||||
expect(service.addAndCheck(event2)).toBe(false);
|
||||
}
|
||||
expect(loggers.logLoopDetected).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Sentence Extraction and Punctuation', () => {
|
||||
it('should not check for loops when content has no sentence-ending punctuation', () => {
|
||||
const eventNoPunct = createContentEvent('This has no punctuation');
|
||||
expect(service.addAndCheck(eventNoPunct)).toBe(false);
|
||||
|
||||
const eventWithPunct = createContentEvent('This has punctuation!');
|
||||
expect(service.addAndCheck(eventWithPunct)).toBe(false);
|
||||
});
|
||||
|
||||
it('should not treat function calls or method calls as sentence endings', () => {
|
||||
// These should not trigger sentence detection, so repeating them many times should never cause a loop
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 2; i++) {
|
||||
expect(service.addAndCheck(createContentEvent('console.log()'))).toBe(
|
||||
false,
|
||||
);
|
||||
}
|
||||
|
||||
service.reset();
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 2; i++) {
|
||||
expect(service.addAndCheck(createContentEvent('obj.method()'))).toBe(
|
||||
false,
|
||||
);
|
||||
}
|
||||
|
||||
service.reset();
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 2; i++) {
|
||||
expect(
|
||||
service.addAndCheck(createContentEvent('arr.filter().map()')),
|
||||
).toBe(false);
|
||||
}
|
||||
|
||||
service.reset();
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD + 2; i++) {
|
||||
expect(
|
||||
service.addAndCheck(
|
||||
createContentEvent('if (condition) { return true; }'),
|
||||
),
|
||||
).toBe(false);
|
||||
}
|
||||
});
|
||||
|
||||
it('should correctly identify actual sentence endings and trigger loop detection', () => {
|
||||
// These should trigger sentence detection, so repeating them should eventually cause a loop
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
|
||||
expect(
|
||||
service.addAndCheck(createContentEvent('This is a sentence.')),
|
||||
).toBe(false);
|
||||
}
|
||||
expect(
|
||||
service.addAndCheck(createContentEvent('This is a sentence.')),
|
||||
).toBe(true);
|
||||
|
||||
service.reset();
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
|
||||
expect(
|
||||
service.addAndCheck(createContentEvent('Is this a question? ')),
|
||||
).toBe(false);
|
||||
}
|
||||
expect(
|
||||
service.addAndCheck(createContentEvent('Is this a question? ')),
|
||||
).toBe(true);
|
||||
|
||||
service.reset();
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
|
||||
expect(
|
||||
service.addAndCheck(createContentEvent('What excitement!\n')),
|
||||
).toBe(false);
|
||||
}
|
||||
expect(
|
||||
service.addAndCheck(createContentEvent('What excitement!\n')),
|
||||
).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle content with mixed punctuation', () => {
|
||||
service.addAndCheck(createContentEvent('Question?'));
|
||||
service.addAndCheck(createContentEvent('Exclamation!'));
|
||||
service.addAndCheck(createContentEvent('Period.'));
|
||||
|
||||
// Repeat one of them multiple times
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
|
||||
service.addAndCheck(createContentEvent('Period.'));
|
||||
}
|
||||
expect(service.addAndCheck(createContentEvent('Period.'))).toBe(true);
|
||||
});
|
||||
|
||||
it('should handle empty sentences after trimming', () => {
|
||||
service.addAndCheck(createContentEvent(' .'));
|
||||
expect(service.addAndCheck(createContentEvent('Normal sentence.'))).toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
|
||||
it('should require at least two sentences for loop detection', () => {
|
||||
const event = createContentEvent('Only one sentence.');
|
||||
expect(service.addAndCheck(event)).toBe(false);
|
||||
|
||||
// Even repeating the same single sentence shouldn't trigger detection
|
||||
for (let i = 0; i < 5; i++) {
|
||||
expect(service.addAndCheck(event)).toBe(false);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('Performance Optimizations', () => {
|
||||
it('should cache sentence extraction and only re-extract when content grows significantly', () => {
|
||||
// Add initial content
|
||||
service.addAndCheck(createContentEvent('First sentence.'));
|
||||
service.addAndCheck(createContentEvent('Second sentence.'));
|
||||
|
||||
// Add small amounts of content (shouldn't trigger re-extraction)
|
||||
for (let i = 0; i < 10; i++) {
|
||||
service.addAndCheck(createContentEvent('X'));
|
||||
}
|
||||
service.addAndCheck(createContentEvent('.'));
|
||||
|
||||
// Should still work correctly
|
||||
expect(service.addAndCheck(createContentEvent('Test.'))).toBe(false);
|
||||
});
|
||||
|
||||
it('should re-extract sentences when content grows by more than 100 characters', () => {
|
||||
service.addAndCheck(createContentEvent('Initial sentence.'));
|
||||
|
||||
// Add enough content to trigger re-extraction
|
||||
const longContent = 'X'.repeat(101);
|
||||
service.addAndCheck(createContentEvent(longContent + '.'));
|
||||
|
||||
// Should work correctly after re-extraction
|
||||
expect(service.addAndCheck(createContentEvent('Test.'))).toBe(false);
|
||||
});
|
||||
|
||||
it('should use indexOf for efficient counting instead of regex', () => {
|
||||
const repeatedSentence = 'This is a repeated sentence.';
|
||||
|
||||
// Build up content with the sentence repeated
|
||||
for (let i = 0; i < CONTENT_LOOP_THRESHOLD - 1; i++) {
|
||||
service.addAndCheck(createContentEvent(repeatedSentence));
|
||||
}
|
||||
|
||||
// The threshold should be reached
|
||||
expect(service.addAndCheck(createContentEvent(repeatedSentence))).toBe(
|
||||
true,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Edge Cases', () => {
|
||||
it('should handle empty content', () => {
|
||||
const event = createContentEvent('');
|
||||
expect(service.addAndCheck(event)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
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.'))).toBe(
|
||||
false,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('General Behavior', () => {
|
||||
it('should return false for unhandled event types', () => {
|
||||
const otherEvent = {
|
||||
type: 'unhandled_event',
|
||||
} as unknown as ServerGeminiStreamEvent;
|
||||
expect(service.addAndCheck(otherEvent)).toBe(false);
|
||||
expect(service.addAndCheck(otherEvent)).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
140
packages/core/src/services/loopDetectionService.ts
Normal file
140
packages/core/src/services/loopDetectionService.ts
Normal file
@@ -0,0 +1,140 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { createHash } from 'crypto';
|
||||
import { GeminiEventType, ServerGeminiStreamEvent } from '../core/turn.js';
|
||||
import { logLoopDetected } from '../telemetry/loggers.js';
|
||||
import { LoopDetectedEvent, LoopType } from '../telemetry/types.js';
|
||||
import { Config } from '../config/config.js';
|
||||
|
||||
const TOOL_CALL_LOOP_THRESHOLD = 5;
|
||||
const CONTENT_LOOP_THRESHOLD = 10;
|
||||
const SENTENCE_ENDING_PUNCTUATION_REGEX = /[.!?]+(?=\s|$)/;
|
||||
|
||||
/**
|
||||
* Service for detecting and preventing infinite loops in AI responses.
|
||||
* Monitors tool call repetitions and content sentence repetitions.
|
||||
*/
|
||||
export class LoopDetectionService {
|
||||
// Tool call tracking
|
||||
private lastToolCallKey: string | null = null;
|
||||
private toolCallRepetitionCount: number = 0;
|
||||
|
||||
// Content streaming tracking
|
||||
private lastRepeatedSentence: string = '';
|
||||
private sentenceRepetitionCount: number = 0;
|
||||
private partialContent: string = '';
|
||||
private config: Config;
|
||||
|
||||
constructor(config: Config) {
|
||||
this.config = config;
|
||||
}
|
||||
|
||||
private getToolCallKey(toolCall: { name: string; args: object }): string {
|
||||
const argsString = JSON.stringify(toolCall.args);
|
||||
const keyString = `${toolCall.name}:${argsString}`;
|
||||
return createHash('sha256').update(keyString).digest('hex');
|
||||
}
|
||||
|
||||
/**
|
||||
* Processes a stream event and checks for loop conditions.
|
||||
* @param event - The stream event to process
|
||||
* @returns true if a loop is detected, false otherwise
|
||||
*/
|
||||
addAndCheck(event: ServerGeminiStreamEvent): boolean {
|
||||
switch (event.type) {
|
||||
case GeminiEventType.ToolCallRequest:
|
||||
// content chanting only happens in one single stream, reset if there
|
||||
// is a tool call in between
|
||||
this.resetSentenceCount();
|
||||
return this.checkToolCallLoop(event.value);
|
||||
case GeminiEventType.Content:
|
||||
return this.checkContentLoop(event.value);
|
||||
default:
|
||||
this.reset();
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
private checkToolCallLoop(toolCall: { name: string; args: object }): boolean {
|
||||
const key = this.getToolCallKey(toolCall);
|
||||
if (this.lastToolCallKey === key) {
|
||||
this.toolCallRepetitionCount++;
|
||||
} else {
|
||||
this.lastToolCallKey = key;
|
||||
this.toolCallRepetitionCount = 1;
|
||||
}
|
||||
if (this.toolCallRepetitionCount >= TOOL_CALL_LOOP_THRESHOLD) {
|
||||
logLoopDetected(
|
||||
this.config,
|
||||
new LoopDetectedEvent(LoopType.CONSECUTIVE_IDENTICAL_TOOL_CALLS),
|
||||
);
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
private checkContentLoop(content: string): boolean {
|
||||
this.partialContent += content;
|
||||
|
||||
if (!SENTENCE_ENDING_PUNCTUATION_REGEX.test(this.partialContent)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const completeSentences =
|
||||
this.partialContent.match(/[^.!?]+[.!?]+(?=\s|$)/g) || [];
|
||||
if (completeSentences.length === 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const lastSentence = completeSentences[completeSentences.length - 1];
|
||||
const lastCompleteIndex = this.partialContent.lastIndexOf(lastSentence);
|
||||
const endOfLastSentence = lastCompleteIndex + lastSentence.length;
|
||||
this.partialContent = this.partialContent.slice(endOfLastSentence);
|
||||
|
||||
for (const sentence of completeSentences) {
|
||||
const trimmedSentence = sentence.trim();
|
||||
if (trimmedSentence === '') {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (this.lastRepeatedSentence === trimmedSentence) {
|
||||
this.sentenceRepetitionCount++;
|
||||
} else {
|
||||
this.lastRepeatedSentence = trimmedSentence;
|
||||
this.sentenceRepetitionCount = 1;
|
||||
}
|
||||
|
||||
if (this.sentenceRepetitionCount >= CONTENT_LOOP_THRESHOLD) {
|
||||
logLoopDetected(
|
||||
this.config,
|
||||
new LoopDetectedEvent(LoopType.CHANTING_IDENTICAL_SENTENCES),
|
||||
);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Resets all loop detection state.
|
||||
*/
|
||||
reset(): void {
|
||||
this.resetToolCallCount();
|
||||
this.resetSentenceCount();
|
||||
}
|
||||
|
||||
private resetToolCallCount(): void {
|
||||
this.lastToolCallKey = null;
|
||||
this.toolCallRepetitionCount = 0;
|
||||
}
|
||||
|
||||
private resetSentenceCount(): void {
|
||||
this.lastRepeatedSentence = '';
|
||||
this.sentenceRepetitionCount = 0;
|
||||
this.partialContent = '';
|
||||
}
|
||||
}
|
||||
@@ -15,6 +15,7 @@ import {
|
||||
ApiResponseEvent,
|
||||
ApiErrorEvent,
|
||||
FlashFallbackEvent,
|
||||
LoopDetectedEvent,
|
||||
} from '../types.js';
|
||||
import { EventMetadataKey } from './event-metadata-key.js';
|
||||
import { Config } from '../../config/config.js';
|
||||
@@ -33,6 +34,7 @@ const api_response_event_name = 'api_response';
|
||||
const api_error_event_name = 'api_error';
|
||||
const end_session_event_name = 'end_session';
|
||||
const flash_fallback_event_name = 'flash_fallback';
|
||||
const loop_detected_event_name = 'loop_detected';
|
||||
|
||||
export interface LogResponse {
|
||||
nextRequestWaitMs?: number;
|
||||
@@ -448,6 +450,18 @@ export class ClearcutLogger {
|
||||
});
|
||||
}
|
||||
|
||||
logLoopDetectedEvent(event: LoopDetectedEvent): void {
|
||||
const data = [
|
||||
{
|
||||
gemini_cli_key: EventMetadataKey.GEMINI_CLI_LOOP_DETECTED_TYPE,
|
||||
value: JSON.stringify(event.loop_type),
|
||||
},
|
||||
];
|
||||
|
||||
this.enqueueLogEvent(this.createLogEvent(loop_detected_event_name, data));
|
||||
this.flushIfNeeded();
|
||||
}
|
||||
|
||||
logEndSessionEvent(event: EndSessionEvent): void {
|
||||
const data = [
|
||||
{
|
||||
|
||||
@@ -150,6 +150,13 @@ export enum EventMetadataKey {
|
||||
|
||||
// Logs the total number of Google accounts ever used.
|
||||
GEMINI_CLI_GOOGLE_ACCOUNTS_COUNT = 37,
|
||||
|
||||
// ==========================================================================
|
||||
// Loop Detected Event Keys
|
||||
// ===========================================================================
|
||||
|
||||
// Logs the type of loop detected.
|
||||
GEMINI_CLI_LOOP_DETECTED_TYPE = 38,
|
||||
}
|
||||
|
||||
export function getEventMetadataKey(
|
||||
|
||||
@@ -25,6 +25,7 @@ import {
|
||||
ToolCallEvent,
|
||||
UserPromptEvent,
|
||||
FlashFallbackEvent,
|
||||
LoopDetectedEvent,
|
||||
} from './types.js';
|
||||
import {
|
||||
recordApiErrorMetrics,
|
||||
@@ -288,3 +289,23 @@ export function logApiResponse(config: Config, event: ApiResponseEvent): void {
|
||||
);
|
||||
recordTokenUsageMetrics(config, event.model, event.tool_token_count, 'tool');
|
||||
}
|
||||
|
||||
export function logLoopDetected(
|
||||
config: Config,
|
||||
event: LoopDetectedEvent,
|
||||
): void {
|
||||
ClearcutLogger.getInstance(config)?.logLoopDetectedEvent(event);
|
||||
if (!isTelemetrySdkInitialized()) return;
|
||||
|
||||
const attributes: LogAttributes = {
|
||||
...getCommonAttributes(config),
|
||||
...event,
|
||||
};
|
||||
|
||||
const logger = logs.getLogger(SERVICE_NAME);
|
||||
const logRecord: LogRecord = {
|
||||
body: `Loop detected. Type: ${event.loop_type}.`,
|
||||
attributes,
|
||||
};
|
||||
logger.emit(logRecord);
|
||||
}
|
||||
|
||||
@@ -246,6 +246,23 @@ export class FlashFallbackEvent {
|
||||
}
|
||||
}
|
||||
|
||||
export enum LoopType {
|
||||
CONSECUTIVE_IDENTICAL_TOOL_CALLS = 'consecutive_identical_tool_calls',
|
||||
CHANTING_IDENTICAL_SENTENCES = 'chanting_identical_sentences',
|
||||
}
|
||||
|
||||
export class LoopDetectedEvent {
|
||||
'event.name': 'loop_detected';
|
||||
'event.timestamp': string; // ISO 8601
|
||||
loop_type: LoopType;
|
||||
|
||||
constructor(loop_type: LoopType) {
|
||||
this['event.name'] = 'loop_detected';
|
||||
this['event.timestamp'] = new Date().toISOString();
|
||||
this.loop_type = loop_type;
|
||||
}
|
||||
}
|
||||
|
||||
export type TelemetryEvent =
|
||||
| StartSessionEvent
|
||||
| EndSessionEvent
|
||||
@@ -254,4 +271,5 @@ export type TelemetryEvent =
|
||||
| ApiRequestEvent
|
||||
| ApiErrorEvent
|
||||
| ApiResponseEvent
|
||||
| FlashFallbackEvent;
|
||||
| FlashFallbackEvent
|
||||
| LoopDetectedEvent;
|
||||
|
||||
Reference in New Issue
Block a user