mirror of
https://github.com/QwenLM/qwen-code.git
synced 2026-05-05 23:42:03 +00:00
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:
parent
2394d732c3
commit
3fb641ca1a
12 changed files with 796 additions and 42 deletions
|
|
@ -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} />
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
};
|
||||
|
|
@ -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([
|
||||
[],
|
||||
|
|
|
|||
|
|
@ -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(() => {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue