Skip to content

Commit e8aea5c

Browse files
aishaneeshahyashodipmore
aishaneeshah
authored and
yashodipmore
committed
feat(loop-reduction): implement iterative loop detection and model feedback (#20763)
1 parent cd56083 commit e8aea5c

5 files changed

Lines changed: 667 additions & 251 deletions

File tree

‎packages/core/src/core/client.test.ts‎

Lines changed: 247 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ import type {
4747
} from '../services/modelConfigService.js';
4848
import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
4949
import * as policyCatalog from '../availability/policyCatalog.js';
50-
import { LlmRole } from '../telemetry/types.js';
50+
import { LlmRole, LoopType } from '../telemetry/types.js';
5151
import { partToString } from '../utils/partUtils.js';
5252
import { coreEvents } from '../utils/events.js';
5353

@@ -2915,45 +2915,257 @@ ${JSON.stringify(
29152915
expect(mockCheckNextSpeaker).not.toHaveBeenCalled();
29162916
});
29172917

2918-
it('should abort linked signal when loop is detected', async () => {
2919-
// Arrange
2920-
vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue(false);
2921-
vi.spyOn(client['loopDetector'], 'addAndCheck')
2922-
.mockReturnValueOnce(false)
2923-
.mockReturnValueOnce(true);
2924-
2925-
let capturedSignal: AbortSignal;
2926-
mockTurnRunFn.mockImplementation((_modelConfigKey, _request, signal) => {
2927-
capturedSignal = signal;
2928-
return (async function* () {
2929-
yield { type: 'content', value: 'First event' };
2930-
yield { type: 'content', value: 'Second event' };
2931-
})();
2918+
describe('Loop Recovery (Two-Strike)', () => {
2919+
beforeEach(() => {
2920+
const mockChat: Partial<GeminiChat> = {
2921+
addHistory: vi.fn(),
2922+
setTools: vi.fn(),
2923+
getHistory: vi.fn().mockReturnValue([]),
2924+
getLastPromptTokenCount: vi.fn(),
2925+
};
2926+
client['chat'] = mockChat as GeminiChat;
2927+
vi.spyOn(client['loopDetector'], 'clearDetection');
2928+
vi.spyOn(client['loopDetector'], 'reset');
29322929
});
29332930

2934-
const mockChat: Partial<GeminiChat> = {
2935-
addHistory: vi.fn(),
2936-
setTools: vi.fn(),
2937-
getHistory: vi.fn().mockReturnValue([]),
2938-
getLastPromptTokenCount: vi.fn(),
2939-
};
2940-
client['chat'] = mockChat as GeminiChat;
2931+
it('should trigger recovery (Strike 1) and continue', async () => {
2932+
// Arrange
2933+
vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue({
2934+
count: 0,
2935+
});
2936+
vi.spyOn(client['loopDetector'], 'addAndCheck')
2937+
.mockReturnValueOnce({ count: 0 })
2938+
.mockReturnValueOnce({ count: 1, detail: 'Repetitive tool call' });
29412939

2942-
// Act
2943-
const stream = client.sendMessageStream(
2944-
[{ text: 'Hi' }],
2945-
new AbortController().signal,
2946-
'prompt-id-loop',
2947-
);
2940+
const sendMessageStreamSpy = vi.spyOn(client, 'sendMessageStream');
29482941

2949-
const events = [];
2950-
for await (const event of stream) {
2951-
events.push(event);
2952-
}
2942+
mockTurnRunFn.mockImplementation(() =>
2943+
(async function* () {
2944+
yield { type: GeminiEventType.Content, value: 'First event' };
2945+
yield { type: GeminiEventType.Content, value: 'Second event' };
2946+
})(),
2947+
);
29532948

2954-
// Assert
2955-
expect(events).toContainEqual({ type: GeminiEventType.LoopDetected });
2956-
expect(capturedSignal!.aborted).toBe(true);
2949+
// Act
2950+
const stream = client.sendMessageStream(
2951+
[{ text: 'Hi' }],
2952+
new AbortController().signal,
2953+
'prompt-id-loop-1',
2954+
);
2955+
2956+
const events = [];
2957+
for await (const event of stream) {
2958+
events.push(event);
2959+
}
2960+
2961+
// Assert
2962+
// sendMessageStream should be called twice (original + recovery)
2963+
expect(sendMessageStreamSpy).toHaveBeenCalledTimes(2);
2964+
2965+
// Verify recovery call parameters
2966+
const recoveryCall = sendMessageStreamSpy.mock.calls[1];
2967+
expect((recoveryCall[0] as Part[])[0].text).toContain(
2968+
'System: Potential loop detected',
2969+
);
2970+
expect((recoveryCall[0] as Part[])[0].text).toContain(
2971+
'Repetitive tool call',
2972+
);
2973+
2974+
// Verify loopDetector.clearDetection was called
2975+
expect(client['loopDetector'].clearDetection).toHaveBeenCalled();
2976+
});
2977+
2978+
it('should terminate (Strike 2) after recovery fails', async () => {
2979+
// Arrange
2980+
vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue({
2981+
count: 0,
2982+
});
2983+
2984+
// First call triggers Strike 1, Second call triggers Strike 2
2985+
vi.spyOn(client['loopDetector'], 'addAndCheck')
2986+
.mockReturnValueOnce({ count: 0 })
2987+
.mockReturnValueOnce({ count: 1, detail: 'Strike 1' }) // Triggers recovery in turn 1
2988+
.mockReturnValueOnce({ count: 2, detail: 'Strike 2' }); // Triggers termination in turn 2 (recovery turn)
2989+
2990+
const sendMessageStreamSpy = vi.spyOn(client, 'sendMessageStream');
2991+
2992+
mockTurnRunFn.mockImplementation(() =>
2993+
(async function* () {
2994+
yield { type: GeminiEventType.Content, value: 'Event' };
2995+
yield { type: GeminiEventType.Content, value: 'Event' };
2996+
})(),
2997+
);
2998+
2999+
// Act
3000+
const stream = client.sendMessageStream(
3001+
[{ text: 'Hi' }],
3002+
new AbortController().signal,
3003+
'prompt-id-loop-2',
3004+
);
3005+
3006+
const events = [];
3007+
for await (const event of stream) {
3008+
events.push(event);
3009+
}
3010+
3011+
// Assert
3012+
expect(events).toContainEqual({ type: GeminiEventType.LoopDetected });
3013+
expect(sendMessageStreamSpy).toHaveBeenCalledTimes(2); // One original, one recovery
3014+
});
3015+
3016+
it('should respect boundedTurns during recovery', async () => {
3017+
// Arrange
3018+
vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue({
3019+
count: 0,
3020+
});
3021+
vi.spyOn(client['loopDetector'], 'addAndCheck').mockReturnValue({
3022+
count: 1,
3023+
detail: 'Loop',
3024+
});
3025+
3026+
const sendMessageStreamSpy = vi.spyOn(client, 'sendMessageStream');
3027+
3028+
mockTurnRunFn.mockImplementation(() =>
3029+
(async function* () {
3030+
yield { type: GeminiEventType.Content, value: 'Event' };
3031+
})(),
3032+
);
3033+
3034+
// Act
3035+
const stream = client.sendMessageStream(
3036+
[{ text: 'Hi' }],
3037+
new AbortController().signal,
3038+
'prompt-id-loop-3',
3039+
1, // Only 1 turn allowed
3040+
);
3041+
3042+
const events = [];
3043+
for await (const event of stream) {
3044+
events.push(event);
3045+
}
3046+
3047+
// Assert
3048+
// Should NOT trigger recovery because boundedTurns would reach 0
3049+
expect(events).toContainEqual({
3050+
type: GeminiEventType.MaxSessionTurns,
3051+
});
3052+
expect(sendMessageStreamSpy).toHaveBeenCalledTimes(1);
3053+
});
3054+
3055+
it('should suppress LoopDetected event on Strike 1', async () => {
3056+
// Arrange
3057+
vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue({
3058+
count: 0,
3059+
});
3060+
vi.spyOn(client['loopDetector'], 'addAndCheck')
3061+
.mockReturnValueOnce({ count: 0 })
3062+
.mockReturnValueOnce({ count: 1, detail: 'Strike 1' });
3063+
3064+
const sendMessageStreamSpy = vi.spyOn(client, 'sendMessageStream');
3065+
3066+
mockTurnRunFn.mockImplementation(() =>
3067+
(async function* () {
3068+
yield { type: GeminiEventType.Content, value: 'Event' };
3069+
yield { type: GeminiEventType.Content, value: 'Event 2' };
3070+
})(),
3071+
);
3072+
3073+
// Act
3074+
const stream = client.sendMessageStream(
3075+
[{ text: 'Hi' }],
3076+
new AbortController().signal,
3077+
'prompt-telemetry',
3078+
);
3079+
3080+
const events = [];
3081+
for await (const event of stream) {
3082+
events.push(event);
3083+
}
3084+
3085+
// Assert
3086+
// Strike 1 should trigger recovery call but NOT emit LoopDetected event
3087+
expect(events).not.toContainEqual({
3088+
type: GeminiEventType.LoopDetected,
3089+
});
3090+
expect(sendMessageStreamSpy).toHaveBeenCalledTimes(2);
3091+
});
3092+
3093+
it('should escalate Strike 2 even if loop type changes', async () => {
3094+
// Arrange
3095+
vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue({
3096+
count: 0,
3097+
});
3098+
3099+
// Strike 1: Tool Call Loop, Strike 2: LLM Detected Loop
3100+
vi.spyOn(client['loopDetector'], 'addAndCheck')
3101+
.mockReturnValueOnce({ count: 0 })
3102+
.mockReturnValueOnce({
3103+
count: 1,
3104+
type: LoopType.TOOL_CALL_LOOP,
3105+
detail: 'Repetitive tool',
3106+
})
3107+
.mockReturnValueOnce({
3108+
count: 2,
3109+
type: LoopType.LLM_DETECTED_LOOP,
3110+
detail: 'LLM loop',
3111+
});
3112+
3113+
const sendMessageStreamSpy = vi.spyOn(client, 'sendMessageStream');
3114+
3115+
mockTurnRunFn.mockImplementation(() =>
3116+
(async function* () {
3117+
yield { type: GeminiEventType.Content, value: 'Event' };
3118+
yield { type: GeminiEventType.Content, value: 'Event 2' };
3119+
})(),
3120+
);
3121+
3122+
// Act
3123+
const stream = client.sendMessageStream(
3124+
[{ text: 'Hi' }],
3125+
new AbortController().signal,
3126+
'prompt-escalate',
3127+
);
3128+
3129+
const events = [];
3130+
for await (const event of stream) {
3131+
events.push(event);
3132+
}
3133+
3134+
// Assert
3135+
expect(events).toContainEqual({ type: GeminiEventType.LoopDetected });
3136+
expect(sendMessageStreamSpy).toHaveBeenCalledTimes(2);
3137+
});
3138+
3139+
it('should reset loop detector on new prompt', async () => {
3140+
// Arrange
3141+
vi.spyOn(client['loopDetector'], 'turnStarted').mockResolvedValue({
3142+
count: 0,
3143+
});
3144+
vi.spyOn(client['loopDetector'], 'addAndCheck').mockReturnValue({
3145+
count: 0,
3146+
});
3147+
mockTurnRunFn.mockImplementation(() =>
3148+
(async function* () {
3149+
yield { type: GeminiEventType.Content, value: 'Event' };
3150+
})(),
3151+
);
3152+
3153+
// Act
3154+
const stream = client.sendMessageStream(
3155+
[{ text: 'Hi' }],
3156+
new AbortController().signal,
3157+
'prompt-id-new',
3158+
);
3159+
for await (const _ of stream) {
3160+
// Consume stream
3161+
}
3162+
3163+
// Assert
3164+
expect(client['loopDetector'].reset).toHaveBeenCalledWith(
3165+
'prompt-id-new',
3166+
'Hi',
3167+
);
3168+
});
29573169
});
29583170
});
29593171

‎packages/core/src/core/client.ts‎

Lines changed: 70 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -642,10 +642,23 @@ export class GeminiClient {
642642
const controller = new AbortController();
643643
const linkedSignal = AbortSignal.any([signal, controller.signal]);
644644

645-
const loopDetected = await this.loopDetector.turnStarted(signal);
646-
if (loopDetected) {
645+
const loopResult = await this.loopDetector.turnStarted(signal);
646+
if (loopResult.count > 1) {
647647
yield { type: GeminiEventType.LoopDetected };
648648
return turn;
649+
} else if (loopResult.count === 1) {
650+
if (boundedTurns <= 1) {
651+
yield { type: GeminiEventType.MaxSessionTurns };
652+
return turn;
653+
}
654+
return yield* this._recoverFromLoop(
655+
loopResult,
656+
signal,
657+
prompt_id,
658+
boundedTurns,
659+
isInvalidStreamRetry,
660+
displayContent,
661+
);
649662
}
650663

651664
const routingContext: RoutingContext = {
@@ -696,10 +709,26 @@ export class GeminiClient {
696709
let isInvalidStream = false;
697710

698711
for await (const event of resultStream) {
699-
if (this.loopDetector.addAndCheck(event)) {
712+
const loopResult = this.loopDetector.addAndCheck(event);
713+
if (loopResult.count > 1) {
700714
yield { type: GeminiEventType.LoopDetected };
701715
controller.abort();
702716
return turn;
717+
} else if (loopResult.count === 1) {
718+
if (boundedTurns <= 1) {
719+
yield { type: GeminiEventType.MaxSessionTurns };
720+
controller.abort();
721+
return turn;
722+
}
723+
return yield* this._recoverFromLoop(
724+
loopResult,
725+
signal,
726+
prompt_id,
727+
boundedTurns,
728+
isInvalidStreamRetry,
729+
displayContent,
730+
controller,
731+
);
703732
}
704733
yield event;
705734

@@ -1128,4 +1157,42 @@ export class GeminiClient {
11281157
this.getChat().setHistory(result.newHistory);
11291158
}
11301159
}
1160+
1161+
/**
1162+
* Handles loop recovery by providing feedback to the model and initiating a new turn.
1163+
*/
1164+
private _recoverFromLoop(
1165+
loopResult: { detail?: string },
1166+
signal: AbortSignal,
1167+
prompt_id: string,
1168+
boundedTurns: number,
1169+
isInvalidStreamRetry: boolean,
1170+
displayContent?: PartListUnion,
1171+
controllerToAbort?: AbortController,
1172+
): AsyncGenerator<ServerGeminiStreamEvent, Turn> {
1173+
controllerToAbort?.abort();
1174+
1175+
// Clear the detection flag so the recursive turn can proceed, but the count remains 1.
1176+
this.loopDetector.clearDetection();
1177+
1178+
const feedbackText = `System: Potential loop detected. Details: ${loopResult.detail || 'Repetitive patterns identified'}. Please take a step back and confirm you're making forward progress. If not, take a step back, analyze your previous actions and rethink how you're approaching the problem. Avoid repeating the same tool calls or responses without new results.`;
1179+
1180+
if (this.config.getDebugMode()) {
1181+
debugLogger.warn(
1182+
'Iterative Loop Recovery: Injecting feedback message to model.',
1183+
);
1184+
}
1185+
1186+
const feedback = [{ text: feedbackText }];
1187+
1188+
// Recursive call with feedback
1189+
return this.sendMessageStream(
1190+
feedback,
1191+
signal,
1192+
prompt_id,
1193+
boundedTurns - 1,
1194+
isInvalidStreamRetry,
1195+
displayContent,
1196+
);
1197+
}
11311198
}

0 commit comments

Comments
 (0)