diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 6fd637d780..4bef9d7531 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -505,6 +505,7 @@ Logging in with Google... Please restart Gemini CLI to continue. pendingHistoryItems: pendingGeminiHistoryItems, thought, cancelOngoingRequest, + loopDetectionConfirmationRequest, } = useGeminiStream( config.getGeminiClient(), historyManager.history, @@ -896,35 +897,20 @@ Logging in with Google... Please restart Gemini CLI to continue. const nightly = props.version.includes('nightly'); - const dialogsVisible = useMemo( - () => - showWorkspaceMigrationDialog || - shouldShowIdePrompt || - isFolderTrustDialogOpen || - !!shellConfirmationRequest || - !!confirmationRequest || - isThemeDialogOpen || - isSettingsDialogOpen || - isAuthenticating || - isAuthDialogOpen || - isEditorDialogOpen || - showPrivacyNotice || - !!proQuotaRequest, - [ - showWorkspaceMigrationDialog, - shouldShowIdePrompt, - isFolderTrustDialogOpen, - shellConfirmationRequest, - confirmationRequest, - isThemeDialogOpen, - isSettingsDialogOpen, - isAuthenticating, - isAuthDialogOpen, - isEditorDialogOpen, - showPrivacyNotice, - proQuotaRequest, - ], - ); + const dialogsVisible = + showWorkspaceMigrationDialog || + shouldShowIdePrompt || + isFolderTrustDialogOpen || + !!shellConfirmationRequest || + !!confirmationRequest || + !!loopDetectionConfirmationRequest || + isThemeDialogOpen || + isSettingsDialogOpen || + isAuthenticating || + isAuthDialogOpen || + isEditorDialogOpen || + showPrivacyNotice || + !!proQuotaRequest; const pendingHistoryItems = useMemo( () => [...pendingSlashCommandHistoryItems, ...pendingGeminiHistoryItems], @@ -952,6 +938,7 @@ Logging in with Google... Please restart Gemini CLI to continue. commandContext, shellConfirmationRequest, confirmationRequest, + loopDetectionConfirmationRequest, geminiMdFileCount, streamingState, initError, @@ -1024,6 +1011,7 @@ Logging in with Google... Please restart Gemini CLI to continue. commandContext, shellConfirmationRequest, confirmationRequest, + loopDetectionConfirmationRequest, geminiMdFileCount, streamingState, initError, diff --git a/packages/cli/src/ui/components/DialogManager.tsx b/packages/cli/src/ui/components/DialogManager.tsx index fff3e2f1cc..aee6191d5a 100644 --- a/packages/cli/src/ui/components/DialogManager.tsx +++ b/packages/cli/src/ui/components/DialogManager.tsx @@ -6,6 +6,7 @@ import { Box, Text } from 'ink'; import { IdeIntegrationNudge } from '../IdeIntegrationNudge.js'; +import { LoopDetectionConfirmation } from './LoopDetectionConfirmation.js'; import { FolderTrustDialog } from './FolderTrustDialog.js'; import { ShellConfirmationDialog } from './ShellConfirmationDialog.js'; import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; @@ -83,6 +84,13 @@ export const DialogManager = () => { ); } + if (uiState.loopDetectionConfirmationRequest) { + return ( + + ); + } if (uiState.confirmationRequest) { return ( diff --git a/packages/cli/src/ui/components/LoopDetectionConfirmation.test.tsx b/packages/cli/src/ui/components/LoopDetectionConfirmation.test.tsx new file mode 100644 index 0000000000..87b57033ee --- /dev/null +++ b/packages/cli/src/ui/components/LoopDetectionConfirmation.test.tsx @@ -0,0 +1,34 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { renderWithProviders } from '../../test-utils/render.js'; +import { describe, it, expect, vi } from 'vitest'; +import { LoopDetectionConfirmation } from './LoopDetectionConfirmation.js'; + +describe('LoopDetectionConfirmation', () => { + const onComplete = vi.fn(); + + it('renders correctly', () => { + const { lastFrame } = renderWithProviders( + , + ); + expect(lastFrame()).toMatchSnapshot(); + }); + + it('contains the expected options', () => { + const { lastFrame } = renderWithProviders( + , + ); + const output = lastFrame()!.toString(); + + expect(output).toContain('A potential loop was detected'); + expect(output).toContain('Keep loop detection enabled (esc)'); + expect(output).toContain('Disable loop detection for this session'); + expect(output).toContain( + 'This can happen due to repetitive tool calls or other model behavior', + ); + }); +}); diff --git a/packages/cli/src/ui/components/LoopDetectionConfirmation.tsx b/packages/cli/src/ui/components/LoopDetectionConfirmation.tsx new file mode 100644 index 0000000000..c644c8866a --- /dev/null +++ b/packages/cli/src/ui/components/LoopDetectionConfirmation.tsx @@ -0,0 +1,88 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { Box, Text } from 'ink'; +import type { RadioSelectItem } from './shared/RadioButtonSelect.js'; +import { RadioButtonSelect } from './shared/RadioButtonSelect.js'; +import { useKeypress } from '../hooks/useKeypress.js'; +import { theme } from '../semantic-colors.js'; + +export type LoopDetectionConfirmationResult = { + userSelection: 'disable' | 'keep'; +}; + +interface LoopDetectionConfirmationProps { + onComplete: (result: LoopDetectionConfirmationResult) => void; +} + +export function LoopDetectionConfirmation({ + onComplete, +}: LoopDetectionConfirmationProps) { + useKeypress( + (key) => { + if (key.name === 'escape') { + onComplete({ + userSelection: 'keep', + }); + } + }, + { isActive: true }, + ); + + const OPTIONS: Array> = [ + { + label: 'Keep loop detection enabled (esc)', + value: { + userSelection: 'keep', + }, + }, + { + label: 'Disable loop detection for this session', + value: { + userSelection: 'disable', + }, + }, + ]; + + return ( + + + + + + ? + + + + + + A potential loop was detected + {' '} + + + + + + + This can happen due to repetitive tool calls or other model + behavior. Do you want to keep loop detection enabled or disable it + for this session? + + + + + + + + + ); +} diff --git a/packages/cli/src/ui/components/__snapshots__/LoopDetectionConfirmation.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/LoopDetectionConfirmation.test.tsx.snap new file mode 100644 index 0000000000..1686c4bc13 --- /dev/null +++ b/packages/cli/src/ui/components/__snapshots__/LoopDetectionConfirmation.test.tsx.snap @@ -0,0 +1,13 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[`LoopDetectionConfirmation > renders correctly 1`] = ` +" ╭──────────────────────────────────────────────────────────────────────────────────────────────────╮ + │ ? A potential loop was detected │ + │ │ + │ This can happen due to repetitive tool calls or other model behavior. Do you want to keep loop │ + │ detection enabled or disable it for this session? │ + │ │ + │ ● 1. Keep loop detection enabled (esc) │ + │ 2. Disable loop detection for this session │ + ╰──────────────────────────────────────────────────────────────────────────────────────────────────╯" +`; diff --git a/packages/cli/src/ui/contexts/UIStateContext.tsx b/packages/cli/src/ui/contexts/UIStateContext.tsx index 0c62ebaf57..f2f713dee7 100644 --- a/packages/cli/src/ui/contexts/UIStateContext.tsx +++ b/packages/cli/src/ui/contexts/UIStateContext.tsx @@ -11,6 +11,7 @@ import type { ConsoleMessageItem, ShellConfirmationRequest, ConfirmationRequest, + LoopDetectionConfirmationRequest, HistoryItemWithoutId, StreamingState, } from '../types.js'; @@ -53,6 +54,7 @@ export interface UIState { commandContext: CommandContext; shellConfirmationRequest: ShellConfirmationRequest | null; confirmationRequest: ConfirmationRequest | null; + loopDetectionConfirmationRequest: LoopDetectionConfirmationRequest | null; geminiMdFileCount: number; streamingState: StreamingState; initError: string | null; diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 46869cf887..f863916513 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -1789,4 +1789,262 @@ describe('useGeminiStream', () => { ); }); }); + + describe('Loop Detection Confirmation', () => { + beforeEach(() => { + // Add mock for getLoopDetectionService to the config + const mockLoopDetectionService = { + disableForSession: vi.fn(), + }; + mockConfig.getGeminiClient = vi.fn().mockReturnValue({ + ...new MockedGeminiClientClass(mockConfig), + getLoopDetectionService: () => mockLoopDetectionService, + }); + }); + + it('should set loopDetectionConfirmationRequest when LoopDetected event is received', async () => { + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'Some content', + }; + yield { + type: ServerGeminiEventType.LoopDetected, + }; + })(), + ); + + const { result } = renderTestHook(); + + await act(async () => { + await result.current.submitQuery('test query'); + }); + + await waitFor(() => { + expect(result.current.loopDetectionConfirmationRequest).not.toBeNull(); + expect( + typeof result.current.loopDetectionConfirmationRequest?.onComplete, + ).toBe('function'); + }); + }); + + it('should disable loop detection and show message when user selects "disable"', async () => { + const mockLoopDetectionService = { + disableForSession: vi.fn(), + }; + const mockClient = { + ...new MockedGeminiClientClass(mockConfig), + getLoopDetectionService: () => mockLoopDetectionService, + }; + mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient); + + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.LoopDetected, + }; + })(), + ); + + const { result } = renderTestHook(); + + await act(async () => { + await result.current.submitQuery('test query'); + }); + + // Wait for confirmation request to be set + await waitFor(() => { + expect(result.current.loopDetectionConfirmationRequest).not.toBeNull(); + }); + + // Simulate user selecting "disable" + await act(async () => { + result.current.loopDetectionConfirmationRequest?.onComplete({ + userSelection: 'disable', + }); + }); + + // Verify loop detection was disabled + expect(mockLoopDetectionService.disableForSession).toHaveBeenCalledTimes( + 1, + ); + + // Verify confirmation request was cleared + expect(result.current.loopDetectionConfirmationRequest).toBeNull(); + + // Verify appropriate message was added + expect(mockAddItem).toHaveBeenCalledWith( + { + type: 'info', + text: 'Loop detection has been disabled for this session. Please try your request again.', + }, + expect.any(Number), + ); + }); + + it('should keep loop detection enabled and show message when user selects "keep"', async () => { + const mockLoopDetectionService = { + disableForSession: vi.fn(), + }; + const mockClient = { + ...new MockedGeminiClientClass(mockConfig), + getLoopDetectionService: () => mockLoopDetectionService, + }; + mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient); + + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.LoopDetected, + }; + })(), + ); + + const { result } = renderTestHook(); + + await act(async () => { + await result.current.submitQuery('test query'); + }); + + // Wait for confirmation request to be set + await waitFor(() => { + expect(result.current.loopDetectionConfirmationRequest).not.toBeNull(); + }); + + // Simulate user selecting "keep" + await act(async () => { + result.current.loopDetectionConfirmationRequest?.onComplete({ + userSelection: 'keep', + }); + }); + + // Verify loop detection was NOT disabled + expect(mockLoopDetectionService.disableForSession).not.toHaveBeenCalled(); + + // Verify confirmation request was cleared + expect(result.current.loopDetectionConfirmationRequest).toBeNull(); + + // Verify appropriate message was added + expect(mockAddItem).toHaveBeenCalledWith( + { + 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.', + }, + expect.any(Number), + ); + }); + + it('should handle multiple loop detection events properly', async () => { + const { result } = renderTestHook(); + + // First loop detection - set up fresh mock for first call + mockSendMessageStream.mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.LoopDetected, + }; + })(), + ); + + // First loop detection + await act(async () => { + await result.current.submitQuery('first query'); + }); + + await waitFor(() => { + expect(result.current.loopDetectionConfirmationRequest).not.toBeNull(); + }); + + // Simulate user selecting "keep" for first request + await act(async () => { + result.current.loopDetectionConfirmationRequest?.onComplete({ + userSelection: 'keep', + }); + }); + + expect(result.current.loopDetectionConfirmationRequest).toBeNull(); + + // Verify first message was added + expect(mockAddItem).toHaveBeenCalledWith( + { + 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.', + }, + expect.any(Number), + ); + + // Second loop detection - set up fresh mock for second call + mockSendMessageStream.mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.LoopDetected, + }; + })(), + ); + + // Second loop detection + await act(async () => { + await result.current.submitQuery('second query'); + }); + + await waitFor(() => { + expect(result.current.loopDetectionConfirmationRequest).not.toBeNull(); + }); + + // Simulate user selecting "disable" for second request + await act(async () => { + result.current.loopDetectionConfirmationRequest?.onComplete({ + userSelection: 'disable', + }); + }); + + expect(result.current.loopDetectionConfirmationRequest).toBeNull(); + + // Verify second message was added + expect(mockAddItem).toHaveBeenCalledWith( + { + type: 'info', + text: 'Loop detection has been disabled for this session. Please try your request again.', + }, + expect.any(Number), + ); + }); + + it('should process LoopDetected event after moving pending history to history', async () => { + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'Some response content', + }; + yield { + type: ServerGeminiEventType.LoopDetected, + }; + })(), + ); + + const { result } = renderTestHook(); + + await act(async () => { + await result.current.submitQuery('test query'); + }); + + // Verify that the content was added to history before the loop detection dialog + await waitFor(() => { + expect(mockAddItem).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'gemini', + text: 'Some response content', + }), + expect.any(Number), + ); + }); + + // Then verify loop detection confirmation request was set + await waitFor(() => { + expect(result.current.loopDetectionConfirmationRequest).not.toBeNull(); + }); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index ac92068edb..f60b6aeb82 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -153,6 +153,12 @@ export const useGeminiStream = ( ); const loopDetectedRef = useRef(false); + const [ + loopDetectionConfirmationRequest, + setLoopDetectionConfirmationRequest, + ] = useState<{ + onComplete: (result: { userSelection: 'disable' | 'keep' }) => void; + } | null>(null); const onExec = useCallback(async (done: Promise) => { setIsResponding(true); @@ -588,15 +594,38 @@ export const useGeminiStream = ( [addItem, config], ); + const handleLoopDetectionConfirmation = useCallback( + (result: { userSelection: 'disable' | 'keep' }) => { + setLoopDetectionConfirmationRequest(null); + + if (result.userSelection === 'disable') { + config.getGeminiClient().getLoopDetectionService().disableForSession(); + addItem( + { + type: 'info', + text: `Loop detection has been disabled for this session. Please try your request again.`, + }, + Date.now(), + ); + } else { + 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(), + ); + } + }, + [config, addItem], + ); + 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]); + // Show the confirmation dialog to choose whether to disable loop detection + setLoopDetectionConfirmationRequest({ + onComplete: handleLoopDetectionConfirmation, + }); + }, [handleLoopDetectionConfirmation]); const processGeminiStreamEvents = useCallback( async ( @@ -1045,5 +1074,6 @@ export const useGeminiStream = ( pendingHistoryItems, thought, cancelOngoingRequest, + loopDetectionConfirmationRequest, }; }; diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 1de5853979..424e743835 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -284,3 +284,7 @@ export interface ConfirmationRequest { prompt: ReactNode; onConfirm: (confirm: boolean) => void; } + +export interface LoopDetectionConfirmationRequest { + onComplete: (result: { userSelection: 'disable' | 'keep' }) => void; +} diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 2ce2e93fcc..db331ad26a 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -186,6 +186,10 @@ export class GeminiClient { return this.chat?.getChatRecordingService(); } + getLoopDetectionService(): LoopDetectionService { + return this.loopDetector; + } + async addDirectoryContext(): Promise { if (!this.chat) { return; diff --git a/packages/core/src/services/loopDetectionService.test.ts b/packages/core/src/services/loopDetectionService.test.ts index 542ea6ce54..0b842d6054 100644 --- a/packages/core/src/services/loopDetectionService.test.ts +++ b/packages/core/src/services/loopDetectionService.test.ts @@ -130,6 +130,15 @@ describe('LoopDetectionService', () => { expect(service.addAndCheck(toolCallEvent)).toBe(true); expect(loggers.logLoopDetected).toHaveBeenCalledTimes(1); }); + + it('should not detect a loop when disabled for session', () => { + service.disableForSession(); + const event = createToolCallRequestEvent('testTool', { param: 'value' }); + for (let i = 0; i < TOOL_CALL_LOOP_THRESHOLD; i++) { + expect(service.addAndCheck(event)).toBe(false); + } + expect(loggers.logLoopDetected).not.toHaveBeenCalled(); + }); }); describe('Content Loop Detection', () => { @@ -719,4 +728,12 @@ describe('LoopDetectionService LLM Checks', () => { expect(result).toBe(false); expect(loggers.logLoopDetected).not.toHaveBeenCalled(); }); + + it('should not trigger LLM check when disabled for session', async () => { + service.disableForSession(); + await advanceTurns(30); + const result = await service.turnStarted(abortController.signal); + expect(result).toBe(false); + expect(mockGeminiClient.generateJson).not.toHaveBeenCalled(); + }); }); diff --git a/packages/core/src/services/loopDetectionService.ts b/packages/core/src/services/loopDetectionService.ts index 97cb94dcb6..d2c531b291 100644 --- a/packages/core/src/services/loopDetectionService.ts +++ b/packages/core/src/services/loopDetectionService.ts @@ -74,10 +74,20 @@ export class LoopDetectionService { private llmCheckInterval = DEFAULT_LLM_CHECK_INTERVAL; private lastCheckTurn = 0; + // Session-level disable flag + private disabledForSession = false; + constructor(config: Config) { this.config = config; } + /** + * Disables loop detection for the current session. + */ + disableForSession(): void { + this.disabledForSession = true; + } + private getToolCallKey(toolCall: { name: string; args: object }): string { const argsString = JSON.stringify(toolCall.args); const keyString = `${toolCall.name}:${argsString}`; @@ -90,8 +100,8 @@ export class LoopDetectionService { * @returns true if a loop is detected, false otherwise */ addAndCheck(event: ServerGeminiStreamEvent): boolean { - if (this.loopDetected) { - return true; + if (this.loopDetected || this.disabledForSession) { + return this.loopDetected; } switch (event.type) { @@ -121,6 +131,9 @@ export class LoopDetectionService { * @returns A promise that resolves to `true` if a loop is detected, and `false` otherwise. */ async turnStarted(signal: AbortSignal) { + if (this.disabledForSession) { + return false; + } this.turnsInCurrentPrompt++; if (