From 4afd3741df7c3322f7bce276876682d320fa7ae2 Mon Sep 17 00:00:00 2001 From: Sehoon Shon Date: Tue, 13 Jan 2026 23:03:19 -0500 Subject: [PATCH] feat(core/ui): enhance retry mechanism and UX (#16489) --- packages/cli/src/ui/AppContainer.tsx | 2 + .../cli/src/ui/hooks/useGeminiStream.test.tsx | 62 +++++++++++++++++++ packages/cli/src/ui/hooks/useGeminiStream.ts | 48 ++++++++++---- .../src/ui/hooks/useLoadingIndicator.test.tsx | 26 ++++++++ .../cli/src/ui/hooks/useLoadingIndicator.ts | 13 +++- packages/core/src/core/geminiChat.ts | 17 +++++ packages/core/src/utils/events.ts | 20 ++++++ packages/core/src/utils/retry.test.ts | 22 +++---- packages/core/src/utils/retry.ts | 16 ++++- 9 files changed, 200 insertions(+), 26 deletions(-) diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 8e05ac1fd7..10f5a54a1c 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -802,6 +802,7 @@ Logging in with Google... Restarting Gemini CLI to continue. activePtyId, loopDetectionConfirmationRequest, lastOutputTime, + retryStatus, } = useGeminiStream( config.getGeminiClient(), historyManager.history, @@ -1223,6 +1224,7 @@ Logging in with Google... Restarting Gemini CLI to continue. settings.merged.ui?.customWittyPhrases, !!activePtyId && !embeddedShellFocused, lastOutputTime, + retryStatus, ); const handleGlobalKeypress = useCallback( diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 21bb2191e9..c5a004a437 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -34,6 +34,8 @@ import { ToolConfirmationOutcome, tokenLimit, debugLogger, + coreEvents, + CoreEvent, } from '@google/gemini-cli-core'; import type { Part, PartListUnion } from '@google/genai'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; @@ -1333,6 +1335,66 @@ describe('useGeminiStream', () => { }); }); + describe('Retry Handling', () => { + it('should update retryStatus when CoreEvent.RetryAttempt is emitted', async () => { + const { result } = renderHookWithDefaults(); + + const retryPayload = { + model: 'gemini-2.5-pro', + attempt: 2, + maxAttempts: 3, + delayMs: 1000, + }; + + await act(async () => { + coreEvents.emit(CoreEvent.RetryAttempt, retryPayload); + }); + + expect(result.current.retryStatus).toEqual(retryPayload); + }); + + it('should reset retryStatus when isResponding becomes false', async () => { + const { result } = renderTestHook(); + + const retryPayload = { + model: 'gemini-2.5-pro', + attempt: 2, + maxAttempts: 3, + delayMs: 1000, + }; + + // Start a query to make isResponding true + const mockStream = (async function* () { + yield { type: ServerGeminiEventType.Content, value: 'Part 1' }; + await new Promise(() => {}); // Keep stream open + })(); + mockSendMessageStream.mockReturnValue(mockStream); + + await act(async () => { + // eslint-disable-next-line @typescript-eslint/no-floating-promises + result.current.submitQuery('test query'); + }); + + await waitFor(() => { + expect(result.current.streamingState).toBe(StreamingState.Responding); + }); + + // Emit retry event + await act(async () => { + coreEvents.emit(CoreEvent.RetryAttempt, retryPayload); + }); + + expect(result.current.retryStatus).toEqual(retryPayload); + + // Cancel to make isResponding false + await act(async () => { + result.current.cancelOngoingRequest(); + }); + + expect(result.current.retryStatus).toBeNull(); + }); + }); + describe('Slash Command Handling', () => { it('should schedule a tool call when the command processor returns a schedule_tool action', async () => { const clientToolRequest: SlashCommandProcessorResult = { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 253582359e..ab88962047 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -5,18 +5,6 @@ */ import { useState, useRef, useCallback, useEffect, useMemo } from 'react'; -import type { - Config, - EditorType, - GeminiClient, - ServerGeminiChatCompressedEvent, - ServerGeminiContentEvent as ContentEvent, - ServerGeminiFinishedEvent, - ServerGeminiStreamEvent as GeminiEvent, - ThoughtSummary, - ToolCallRequestInfo, - GeminiErrorEventValue, -} from '@google/gemini-cli-core'; import { GeminiEventType as ServerGeminiEventType, getErrorMessage, @@ -40,6 +28,21 @@ import { processRestorableToolCalls, recordToolCallInteractions, ToolErrorType, + coreEvents, + CoreEvent, +} from '@google/gemini-cli-core'; +import type { + Config, + EditorType, + GeminiClient, + ServerGeminiChatCompressedEvent, + ServerGeminiContentEvent as ContentEvent, + ServerGeminiFinishedEvent, + ServerGeminiStreamEvent as GeminiEvent, + ThoughtSummary, + ToolCallRequestInfo, + GeminiErrorEventValue, + RetryAttemptPayload, } from '@google/gemini-cli-core'; import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import type { @@ -113,6 +116,9 @@ export const useGeminiStream = ( isShellFocused?: boolean, ) => { const [initError, setInitError] = useState(null); + const [retryStatus, setRetryStatus] = useState( + null, + ); const abortControllerRef = useRef(null); const turnCancelledRef = useRef(false); const activeQueryIdRef = useRef(null); @@ -133,6 +139,16 @@ export const useGeminiStream = ( return new GitService(config.getProjectRoot(), storage); }, [config, storage]); + useEffect(() => { + const handleRetryAttempt = (payload: RetryAttemptPayload) => { + setRetryStatus(payload); + }; + coreEvents.on(CoreEvent.RetryAttempt, handleRetryAttempt); + return () => { + coreEvents.off(CoreEvent.RetryAttempt, handleRetryAttempt); + }; + }, []); + const [ toolCalls, scheduleToolCalls, @@ -297,6 +313,12 @@ export const useGeminiStream = ( } }, [streamingState, config, history]); + useEffect(() => { + if (!isResponding) { + setRetryStatus(null); + } + }, [isResponding]); + const cancelOngoingRequest = useCallback(() => { if ( streamingState !== StreamingState.Responding && @@ -527,6 +549,7 @@ export const useGeminiStream = ( currentGeminiMessageBuffer: string, userMessageTimestamp: number, ): string => { + setRetryStatus(null); if (turnCancelledRef.current) { // Prevents additional output after a user initiated cancel. return ''; @@ -1362,5 +1385,6 @@ export const useGeminiStream = ( activePtyId, loopDetectionConfirmationRequest, lastOutputTime, + retryStatus, }; }; diff --git a/packages/cli/src/ui/hooks/useLoadingIndicator.test.tsx b/packages/cli/src/ui/hooks/useLoadingIndicator.test.tsx index 14270019ac..300ef823da 100644 --- a/packages/cli/src/ui/hooks/useLoadingIndicator.test.tsx +++ b/packages/cli/src/ui/hooks/useLoadingIndicator.test.tsx @@ -15,6 +15,7 @@ import { } from './usePhraseCycler.js'; import { WITTY_LOADING_PHRASES } from '../constants/wittyPhrases.js'; import { INFORMATIVE_TIPS } from '../constants/tips.js'; +import type { RetryAttemptPayload } from '@google/gemini-cli-core'; describe('useLoadingIndicator', () => { beforeEach(() => { @@ -32,22 +33,26 @@ describe('useLoadingIndicator', () => { initialStreamingState: StreamingState, initialIsInteractiveShellWaiting: boolean = false, initialLastOutputTime: number = 0, + initialRetryStatus: RetryAttemptPayload | null = null, ) => { let hookResult: ReturnType; function TestComponent({ streamingState, isInteractiveShellWaiting, lastOutputTime, + retryStatus, }: { streamingState: StreamingState; isInteractiveShellWaiting?: boolean; lastOutputTime?: number; + retryStatus?: RetryAttemptPayload | null; }) { hookResult = useLoadingIndicator( streamingState, undefined, isInteractiveShellWaiting, lastOutputTime, + retryStatus, ); return null; } @@ -56,6 +61,7 @@ describe('useLoadingIndicator', () => { streamingState={initialStreamingState} isInteractiveShellWaiting={initialIsInteractiveShellWaiting} lastOutputTime={initialLastOutputTime} + retryStatus={initialRetryStatus} />, ); return { @@ -68,6 +74,7 @@ describe('useLoadingIndicator', () => { streamingState: StreamingState; isInteractiveShellWaiting?: boolean; lastOutputTime?: number; + retryStatus?: RetryAttemptPayload | null; }) => rerender(), }; }; @@ -206,4 +213,23 @@ describe('useLoadingIndicator', () => { }); expect(result.current.elapsedTime).toBe(0); }); + + it('should reflect retry status in currentLoadingPhrase when provided', () => { + const retryStatus = { + model: 'gemini-pro', + attempt: 2, + maxAttempts: 3, + delayMs: 1000, + }; + const { result } = renderLoadingIndicatorHook( + StreamingState.Responding, + false, + 0, + retryStatus, + ); + + expect(result.current.currentLoadingPhrase).toBe( + 'Trying to reach gemini-pro (Attempt 2/3)', + ); + }); }); diff --git a/packages/cli/src/ui/hooks/useLoadingIndicator.ts b/packages/cli/src/ui/hooks/useLoadingIndicator.ts index a39b0c0e29..9782752d46 100644 --- a/packages/cli/src/ui/hooks/useLoadingIndicator.ts +++ b/packages/cli/src/ui/hooks/useLoadingIndicator.ts @@ -7,13 +7,18 @@ import { StreamingState } from '../types.js'; import { useTimer } from './useTimer.js'; import { usePhraseCycler } from './usePhraseCycler.js'; -import { useState, useEffect, useRef } from 'react'; // Added useRef +import { useState, useEffect, useRef } from 'react'; +import { + getDisplayString, + type RetryAttemptPayload, +} from '@google/gemini-cli-core'; export const useLoadingIndicator = ( streamingState: StreamingState, customWittyPhrases?: string[], isInteractiveShellWaiting: boolean = false, lastOutputTime: number = 0, + retryStatus: RetryAttemptPayload | null = null, ) => { const [timerResetKey, setTimerResetKey] = useState(0); const isTimerActive = streamingState === StreamingState.Responding; @@ -55,11 +60,15 @@ export const useLoadingIndicator = ( prevStreamingStateRef.current = streamingState; }, [streamingState, elapsedTimeFromTimer]); + const retryPhrase = retryStatus + ? `Trying to reach ${getDisplayString(retryStatus.model)} (Attempt ${retryStatus.attempt}/${retryStatus.maxAttempts})` + : null; + return { elapsedTime: streamingState === StreamingState.WaitingForConfirmation ? retainedElapsedTime : elapsedTimeFromTimer, - currentLoadingPhrase, + currentLoadingPhrase: retryPhrase || currentLoadingPhrase, }; }; diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 2dff70c16d..25f32e09a7 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -54,6 +54,7 @@ import { fireBeforeModelHook, fireBeforeToolSelectionHook, } from './geminiChatHookTriggers.js'; +import { coreEvents } from '../utils/events.js'; export enum StreamEventType { /** A regular content chunk from the API. */ @@ -401,6 +402,13 @@ export class GeminiChat { this.config, new ContentRetryEvent(attempt, retryType, delayMs, model), ); + coreEvents.emitRetryAttempt({ + attempt: attempt + 1, + maxAttempts, + delayMs: delayMs * (attempt + 1), + error: error instanceof Error ? error.message : String(error), + model, + }); await new Promise((res) => setTimeout(res, delayMs * (attempt + 1)), ); @@ -601,6 +609,15 @@ export class GeminiChat { signal: abortSignal, maxAttempts: availabilityMaxAttempts, getAvailabilityContext, + onRetry: (attempt, error, delayMs) => { + coreEvents.emitRetryAttempt({ + attempt, + maxAttempts: availabilityMaxAttempts ?? 10, + delayMs, + error: error instanceof Error ? error.message : String(error), + model: lastModelToUse, + }); + }, }); // Store the original request for AfterModel hooks diff --git a/packages/core/src/utils/events.ts b/packages/core/src/utils/events.ts index 89dd02395f..e6a15c68ab 100644 --- a/packages/core/src/utils/events.ts +++ b/packages/core/src/utils/events.ts @@ -97,6 +97,17 @@ export interface HookEndPayload extends HookPayload { success: boolean; } +/** + * Payload for the 'retry-attempt' event. + */ +export interface RetryAttemptPayload { + attempt: number; + maxAttempts: number; + delayMs: number; + error?: string; + model: string; +} + export enum CoreEvent { UserFeedback = 'user-feedback', ModelChanged = 'model-changed', @@ -108,6 +119,7 @@ export enum CoreEvent { HookStart = 'hook-start', HookEnd = 'hook-end', AgentsRefreshed = 'agents-refreshed', + RetryAttempt = 'retry-attempt', } export interface CoreEvents { @@ -121,6 +133,7 @@ export interface CoreEvents { [CoreEvent.HookStart]: [HookStartPayload]; [CoreEvent.HookEnd]: [HookEndPayload]; [CoreEvent.AgentsRefreshed]: never[]; + [CoreEvent.RetryAttempt]: [RetryAttemptPayload]; } type EventBacklogItem = { @@ -229,6 +242,13 @@ export class CoreEventEmitter extends EventEmitter { this.emit(CoreEvent.AgentsRefreshed); } + /** + * Notifies subscribers that a retry attempt is happening. + */ + emitRetryAttempt(payload: RetryAttemptPayload): void { + this.emit(CoreEvent.RetryAttempt, payload); + } + /** * Flushes buffered messages. Call this immediately after primary UI listener * subscribes. diff --git a/packages/core/src/utils/retry.test.ts b/packages/core/src/utils/retry.test.ts index b4edf6a9ce..7a70b45bb4 100644 --- a/packages/core/src/utils/retry.test.ts +++ b/packages/core/src/utils/retry.test.ts @@ -101,33 +101,33 @@ describe('retryWithBackoff', () => { expect(mockFn).toHaveBeenCalledTimes(3); }); - it('should default to 3 maxAttempts if no options are provided', async () => { - // This function will fail more than 3 times to ensure all retries are used. - const mockFn = createFailingFunction(10); + it('should default to 10 maxAttempts if no options are provided', async () => { + // This function will fail more than 10 times to ensure all retries are used. + const mockFn = createFailingFunction(15); const promise = retryWithBackoff(mockFn); await Promise.all([ - expect(promise).rejects.toThrow('Simulated error attempt 3'), + expect(promise).rejects.toThrow('Simulated error attempt 10'), vi.runAllTimersAsync(), ]); - expect(mockFn).toHaveBeenCalledTimes(3); + expect(mockFn).toHaveBeenCalledTimes(10); }); - it('should default to 3 maxAttempts if options.maxAttempts is undefined', async () => { - // This function will fail more than 3 times to ensure all retries are used. - const mockFn = createFailingFunction(10); + it('should default to 10 maxAttempts if options.maxAttempts is undefined', async () => { + // This function will fail more than 10 times to ensure all retries are used. + const mockFn = createFailingFunction(15); const promise = retryWithBackoff(mockFn, { maxAttempts: undefined }); - // Expect it to fail with the error from the 5th attempt. + // Expect it to fail with the error from the 10th attempt. await Promise.all([ - expect(promise).rejects.toThrow('Simulated error attempt 3'), + expect(promise).rejects.toThrow('Simulated error attempt 10'), vi.runAllTimersAsync(), ]); - expect(mockFn).toHaveBeenCalledTimes(3); + expect(mockFn).toHaveBeenCalledTimes(10); }); it('should not retry if shouldRetry returns false', async () => { diff --git a/packages/core/src/utils/retry.ts b/packages/core/src/utils/retry.ts index 5716e48886..eb9e145c18 100644 --- a/packages/core/src/utils/retry.ts +++ b/packages/core/src/utils/retry.ts @@ -32,10 +32,11 @@ export interface RetryOptions { retryFetchErrors?: boolean; signal?: AbortSignal; getAvailabilityContext?: () => RetryAvailabilityContext | undefined; + onRetry?: (attempt: number, error: unknown, delayMs: number) => void; } const DEFAULT_RETRY_OPTIONS: RetryOptions = { - maxAttempts: 3, + maxAttempts: 10, initialDelayMs: 5000, maxDelayMs: 30000, // 30 seconds shouldRetryOnError: isRetryableError, @@ -149,6 +150,7 @@ export async function retryWithBackoff( retryFetchErrors, signal, getAvailabilityContext, + onRetry, } = { ...DEFAULT_RETRY_OPTIONS, shouldRetryOnError: isRetryableError, @@ -172,6 +174,9 @@ export async function retryWithBackoff( ) { const jitter = currentDelay * 0.3 * (Math.random() * 2 - 1); const delayWithJitter = Math.max(0, currentDelay + jitter); + if (onRetry) { + onRetry(attempt, new Error('Invalid content'), delayWithJitter); + } await delay(delayWithJitter, signal); currentDelay = Math.min(maxDelayMs, currentDelay * 2); continue; @@ -252,6 +257,9 @@ export async function retryWithBackoff( debugLogger.warn( `Attempt ${attempt} failed: ${classifiedError.message}. Retrying after ${classifiedError.retryDelayMs}ms...`, ); + if (onRetry) { + onRetry(attempt, error, classifiedError.retryDelayMs); + } await delay(classifiedError.retryDelayMs, signal); continue; } else { @@ -261,6 +269,9 @@ export async function retryWithBackoff( // Exponential backoff with jitter for non-quota errors const jitter = currentDelay * 0.3 * (Math.random() * 2 - 1); const delayWithJitter = Math.max(0, currentDelay + jitter); + if (onRetry) { + onRetry(attempt, error, delayWithJitter); + } await delay(delayWithJitter, signal); currentDelay = Math.min(maxDelayMs, currentDelay * 2); continue; @@ -281,6 +292,9 @@ export async function retryWithBackoff( // Exponential backoff with jitter for non-quota errors const jitter = currentDelay * 0.3 * (Math.random() * 2 - 1); const delayWithJitter = Math.max(0, currentDelay + jitter); + if (onRetry) { + onRetry(attempt, error, delayWithJitter); + } await delay(delayWithJitter, signal); currentDelay = Math.min(maxDelayMs, currentDelay * 2); }