feat(core, cli): add rate limit throttling retry with countdown UI

- Refactor retry utility to support GLM rate limit errors (code 1302) and TPM throttling
- Add getRateLimitRetryInfo() for unified rate-limit error detection
- Add exponential backoff for non-TPM rate limit errors
- Extend StreamEventType.RETRY with RetryInfo payload for UI feedback
- Add RetryCountdownMessage component for visual retry countdown
- Update useGeminiStream hook to handle retry events with countdown timer
- Add i18n support for rate limit messages (en/zh)
This commit is contained in:
yiliang114 2026-02-12 16:21:10 +08:00
parent 2394d732c3
commit 3fb641ca1a
12 changed files with 796 additions and 42 deletions

View file

@ -20,6 +20,7 @@ import { GeminiThoughtMessageContent } from './messages/GeminiThoughtMessageCont
import { CompressionMessage } from './messages/CompressionMessage.js';
import { SummaryMessage } from './messages/SummaryMessage.js';
import { WarningMessage } from './messages/WarningMessage.js';
import { RetryCountdownMessage } from './messages/RetryCountdownMessage.js';
import { Box } from 'ink';
import { AboutBox } from './AboutBox.js';
import { StatsDisplay } from './StatsDisplay.js';
@ -126,6 +127,9 @@ const HistoryItemDisplayComponent: React.FC<HistoryItemDisplayProps> = ({
{itemForDisplay.type === 'error' && (
<ErrorMessage text={itemForDisplay.text} />
)}
{itemForDisplay.type === 'retry_countdown' && (
<RetryCountdownMessage text={itemForDisplay.text} />
)}
{itemForDisplay.type === 'about' && (
<AboutBox {...itemForDisplay.systemInfo} width={boxWidth} />
)}

View file

@ -0,0 +1,41 @@
/**
* @license
* Copyright 2025 Qwen
* SPDX-License-Identifier: Apache-2.0
*/
import type React from 'react';
import { Text, Box } from 'ink';
import { theme } from '../../semantic-colors.js';
interface RetryCountdownMessageProps {
text: string;
}
/**
* Displays a retry countdown message in a dimmed/secondary style
* to visually distinguish it from error messages.
*/
export const RetryCountdownMessage: React.FC<RetryCountdownMessageProps> = ({
text,
}) => {
if (!text || text.trim() === '') {
return null;
}
const prefix = '↻ ';
const prefixWidth = prefix.length;
return (
<Box flexDirection="row">
<Box width={prefixWidth}>
<Text color={theme.text.secondary}>{prefix}</Text>
</Box>
<Box flexGrow={1}>
<Text wrap="wrap" color={theme.text.secondary}>
{text}
</Text>
</Box>
</Box>
);
};

View file

@ -2296,6 +2296,107 @@ describe('useGeminiStream', () => {
});
});
it('should show a retry countdown and update pending history over time', async () => {
vi.useFakeTimers();
try {
let resolveStream: (() => void) | undefined;
mockSendMessageStream.mockReturnValue(
(async function* () {
yield {
type: ServerGeminiEventType.Retry,
retryInfo: {
reason: 'Rate limit exceeded',
attempt: 1,
maxRetries: 3,
delayMs: 3000,
},
};
await new Promise<void>((resolve) => {
resolveStream = resolve;
});
yield {
type: ServerGeminiEventType.Finished,
value: { reason: 'STOP', usageMetadata: undefined },
};
})(),
);
const { result } = renderHook(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
mockAddItem,
mockConfig,
mockLoadedSettings,
mockOnDebugMessage,
mockHandleSlashCommand,
false,
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
false, // visionModelPreviewEnabled
() => {},
80,
24,
),
);
act(() => {
void result.current.submitQuery('Trigger retry');
});
await act(async () => {
await Promise.resolve();
});
// Error line should be rendered as ERROR type
const errorItem = result.current.pendingHistoryItems.find(
(item) => item.type === MessageType.ERROR,
);
expect(errorItem?.text).toContain('Rate limit exceeded');
// Countdown line should be rendered as retry_countdown type
const countdownItem = result.current.pendingHistoryItems.find(
(item) => item.type === ('retry_countdown' as MessageType),
);
expect(countdownItem?.text).toContain('Retrying in 3 seconds');
await act(async () => {
await vi.advanceTimersByTimeAsync(1000);
});
const countdownAfterOneSecond = result.current.pendingHistoryItems.find(
(item) => item.type === ('retry_countdown' as MessageType),
);
expect(countdownAfterOneSecond?.text).toContain(
'Retrying in 2 seconds',
);
resolveStream?.();
await act(async () => {
await Promise.resolve();
await vi.runAllTimersAsync();
});
// Both error and countdown should be cleared after retry succeeds
const remainingError = result.current.pendingHistoryItems.find(
(item) => item.type === MessageType.ERROR,
);
const remainingCountdown = result.current.pendingHistoryItems.find(
(item) => item.type === ('retry_countdown' as MessageType),
);
expect(remainingError).toBeUndefined();
expect(remainingCountdown).toBeUndefined();
} finally {
vi.useRealTimers();
}
});
it('should memoize pendingHistoryItems', () => {
mockUseReactToolScheduler.mockReturnValue([
[],

View file

@ -65,6 +65,7 @@ import path from 'node:path';
import { useSessionStats } from '../contexts/SessionContext.js';
import { useKeypress } from './useKeypress.js';
import type { LoadedSettings } from '../../config/settings.js';
import { t } from '../../i18n/index.js';
const debugLogger = createDebugLogger('GEMINI_STREAM');
@ -125,6 +126,13 @@ export const useGeminiStream = (
const [thought, setThought] = useState<ThoughtSummary | null>(null);
const [pendingHistoryItem, pendingHistoryItemRef, setPendingHistoryItem] =
useStateAndRef<HistoryItemWithoutId | null>(null);
const [pendingRetryErrorItem, setPendingRetryErrorItem] =
useState<HistoryItemWithoutId | null>(null);
const [pendingRetryCountdownItem, setPendingRetryCountdownItem] =
useState<HistoryItemWithoutId | null>(null);
const retryCountdownTimerRef = useRef<ReturnType<typeof setInterval> | null>(
null,
);
const processedMemoryToolsRef = useRef<Set<string>>(new Set());
const {
startNewPrompt,
@ -189,6 +197,67 @@ export const useGeminiStream = (
onComplete: (result: { userSelection: 'disable' | 'keep' }) => void;
} | null>(null);
const stopRetryCountdownTimer = useCallback(() => {
if (retryCountdownTimerRef.current) {
clearInterval(retryCountdownTimerRef.current);
retryCountdownTimerRef.current = null;
}
}, []);
const clearRetryCountdown = useCallback(() => {
stopRetryCountdownTimer();
setPendingRetryErrorItem(null);
setPendingRetryCountdownItem(null);
}, [stopRetryCountdownTimer]);
const startRetryCountdown = useCallback(
(retryInfo: {
reason: string;
attempt: number;
maxRetries: number;
delayMs: number;
}) => {
stopRetryCountdownTimer();
const startTime = Date.now();
const { reason, attempt, maxRetries, delayMs } = retryInfo;
// Error line stays static (red with ✕ prefix)
setPendingRetryErrorItem({
type: MessageType.ERROR,
text: t('Rate limit error: {{reason}}', { reason }),
});
// Countdown line updates every second (dim/secondary color)
const updateCountdown = () => {
const elapsedMs = Date.now() - startTime;
const remainingMs = Math.max(0, delayMs - elapsedMs);
const remainingSec = Math.ceil(remainingMs / 1000);
setPendingRetryCountdownItem({
type: 'retry_countdown',
text: t(
'Retrying in {{seconds}} seconds… (attempt {{attempt}}/{{maxRetries}})',
{
seconds: String(remainingSec),
attempt: String(attempt),
maxRetries: String(maxRetries),
},
),
} as HistoryItemWithoutId);
if (remainingMs <= 0) {
stopRetryCountdownTimer();
}
};
updateCountdown();
retryCountdownTimerRef.current = setInterval(updateCountdown, 1000);
},
[stopRetryCountdownTimer],
);
useEffect(() => () => stopRetryCountdownTimer(), [stopRetryCountdownTimer]);
const onExec = useCallback(async (done: Promise<void>) => {
setIsResponding(true);
await done;
@ -295,6 +364,7 @@ export const useGeminiStream = (
Date.now(),
);
setPendingHistoryItem(null);
clearRetryCountdown();
onCancelSubmit();
setIsResponding(false);
setShellInputFocused(false);
@ -305,6 +375,7 @@ export const useGeminiStream = (
onCancelSubmit,
pendingHistoryItemRef,
setShellInputFocused,
clearRetryCountdown,
config,
getPromptCount,
]);
@ -609,10 +680,17 @@ export const useGeminiStream = (
{ type: MessageType.INFO, text: 'User cancelled the request.' },
userMessageTimestamp,
);
clearRetryCountdown();
setIsResponding(false);
setThought(null); // Reset thought when user cancels
},
[addItem, pendingHistoryItemRef, setPendingHistoryItem, setThought],
[
addItem,
pendingHistoryItemRef,
setPendingHistoryItem,
setThought,
clearRetryCountdown,
],
);
const handleErrorEvent = useCallback(
@ -631,9 +709,17 @@ export const useGeminiStream = (
},
userMessageTimestamp,
);
clearRetryCountdown();
setThought(null); // Reset thought when there's an error
},
[addItem, pendingHistoryItemRef, setPendingHistoryItem, config, setThought],
[
addItem,
pendingHistoryItemRef,
setPendingHistoryItem,
config,
setThought,
clearRetryCountdown,
],
);
const handleCitationEvent = useCallback(
@ -693,8 +779,9 @@ export const useGeminiStream = (
userMessageTimestamp,
);
}
clearRetryCountdown();
},
[addItem],
[addItem, clearRetryCountdown],
);
const handleChatCompressionEvent = useCallback(
@ -853,7 +940,16 @@ export const useGeminiStream = (
loopDetectedRef.current = true;
break;
case ServerGeminiEventType.Retry:
// Will add the missing logic later
// Clear any pending partial content from the failed attempt
if (pendingHistoryItemRef.current) {
setPendingHistoryItem(null);
}
// Show retry info if available (rate-limit / throttling errors)
if (event.retryInfo) {
startRetryCountdown(event.retryInfo);
} else {
clearRetryCountdown();
}
break;
default: {
// enforces exhaustive switch-case
@ -878,7 +974,11 @@ export const useGeminiStream = (
handleMaxSessionTurnsEvent,
handleSessionTokenLimitExceededEvent,
handleCitationEvent,
startRetryCountdown,
clearRetryCountdown,
setThought,
pendingHistoryItemRef,
setPendingHistoryItem,
],
);
@ -1216,10 +1316,18 @@ export const useGeminiStream = (
const pendingHistoryItems = useMemo(
() =>
[pendingHistoryItem, pendingToolCallGroupDisplay].filter(
(i) => i !== undefined && i !== null,
),
[pendingHistoryItem, pendingToolCallGroupDisplay],
[
pendingHistoryItem,
pendingRetryErrorItem,
pendingRetryCountdownItem,
pendingToolCallGroupDisplay,
].filter((i) => i !== undefined && i !== null),
[
pendingHistoryItem,
pendingRetryErrorItem,
pendingRetryCountdownItem,
pendingToolCallGroupDisplay,
],
);
useEffect(() => {

View file

@ -128,6 +128,11 @@ export type HistoryItemWarning = HistoryItemBase & {
text: string;
};
export type HistoryItemRetryCountdown = HistoryItemBase & {
type: 'retry_countdown';
text: string;
};
export type HistoryItemAbout = HistoryItemBase & {
type: 'about';
systemInfo: {
@ -265,6 +270,7 @@ export type HistoryItemWithoutId =
| HistoryItemInfo
| HistoryItemError
| HistoryItemWarning
| HistoryItemRetryCountdown
| HistoryItemAbout
| HistoryItemHelp
| HistoryItemToolGroup