MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-20 10:10:56 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat(core/ui): enhance retry mechanism and UX (#16489)

Browse Source
This commit is contained in:
Sehoon Shon
2026-01-13 23:03:19 -05:00
committed by GitHub
parent 428e602882
commit 4afd3741df
9 changed files with 200 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/cli/src/ui/AppContainer.tsx
+2
View File
@@ -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(
packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+62
View File
@@ -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 = {
packages/cli/src/ui/hooks/useGeminiStream.ts
+36 -12
View File
@@ -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<string | null>(null);
const [retryStatus, setRetryStatus] = useState<RetryAttemptPayload | null>(
null,
);
const abortControllerRef = useRef<AbortController | null>(null);
const turnCancelledRef = useRef(false);
const activeQueryIdRef = useRef<string | null>(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,
};
};
packages/cli/src/ui/hooks/useLoadingIndicator.test.tsx
+26
View File
@@ -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<typeof useLoadingIndicator>;
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(<TestComponent {...newProps} />),
};
};
@@ -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)',
);
});
});
packages/cli/src/ui/hooks/useLoadingIndicator.ts
+11 -2
View File
@@ -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,
};
};
packages/core/src/core/geminiChat.ts
+17
View File
@@ -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
packages/core/src/utils/events.ts
+20
View File
@@ -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<CoreEvents> {
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.
packages/core/src/utils/retry.test.ts
+11 -11
View File
@@ -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 () => {
packages/core/src/utils/retry.ts
+15 -1
View File
@@ -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<T>(
retryFetchErrors,
signal,
getAvailabilityContext,
onRetry,
} = {
...DEFAULT_RETRY_OPTIONS,
shouldRetryOnError: isRetryableError,
@@ -172,6 +174,9 @@ export async function retryWithBackoff<T>(
) {
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<T>(
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<T>(
// 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<T>(
// 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);
}