feat(core): add rate limit error detection utility

- Extract rate-limit detection into dedicated rateLimit.ts module
- Support detection from ApiError, StructuredError, HttpError, and JSON strings
- Handle common rate-limit codes: 429, 503, 1302 (GLM)
- Simplify retry.ts by removing duplicated detection logic
This commit is contained in:
yiliang114 2026-02-13 15:32:35 +08:00
parent 0d2e394ef1
commit 6eb6812f5e
11 changed files with 229 additions and 209 deletions

View file

@ -67,7 +67,12 @@ const MockedUserPromptEvent = vi.hoisted(() =>
const MockedApiCancelEvent = vi.hoisted(() =>
vi.fn().mockImplementation(() => {}),
);
const mockParseAndFormatApiError = vi.hoisted(() => vi.fn());
const mockParseAndFormatApiError = vi.hoisted(() =>
vi.fn(
(msg: unknown) =>
`[API Error: ${typeof msg === 'string' ? msg : 'An unknown error occurred.'}]`,
),
);
const mockLogApiCancel = vi.hoisted(() => vi.fn());
// Vision auto-switch mocks (hoisted)
@ -123,22 +128,6 @@ vi.mock('../utils/markdownUtilities.js', () => ({
findLastSafeSplitPoint: vi.fn((s: string) => s.length),
}));
vi.mock('./useStateAndRef.js', () => ({
useStateAndRef: vi.fn((initial) => {
let val = initial;
const ref = { current: val };
const setVal = vi.fn((updater) => {
if (typeof updater === 'function') {
val = updater(val);
} else {
val = updater;
}
ref.current = val;
});
return [val, ref, setVal];
}),
}));
vi.mock('./useLogger.js', () => ({
useLogger: vi.fn().mockReturnValue({
logMessage: vi.fn().mockResolvedValue(undefined),
@ -2305,12 +2294,15 @@ describe('useGeminiStream', () => {
yield {
type: ServerGeminiEventType.Retry,
retryInfo: {
reason: 'Rate limit exceeded',
message: '[API Error: Rate limit exceeded]',
attempt: 1,
maxRetries: 3,
delayMs: 3000,
},
};
yield {
type: ServerGeminiEventType.Retry,
};
await new Promise<void>((resolve) => {
resolveStream = resolve;
});
@ -2353,16 +2345,33 @@ describe('useGeminiStream', () => {
await Promise.resolve();
});
// Error line should be rendered as ERROR type
const errorItem = result.current.pendingHistoryItems.find(
(item) => item.type === MessageType.ERROR,
);
const findErrorItem = () =>
result.current.pendingHistoryItems.find(
(item) => item.type === MessageType.ERROR,
);
const findCountdownItem = () =>
result.current.pendingHistoryItems.find(
(item) => item.type === 'retry_countdown',
);
let errorItem = findErrorItem();
let countdownItem = findCountdownItem();
for (
let attempts = 0;
attempts < 5 && (!errorItem || !countdownItem);
attempts++
) {
await act(async () => {
await Promise.resolve();
});
errorItem = findErrorItem();
countdownItem = findCountdownItem();
}
// Error line should be rendered as ERROR type (wrapped by parseAndFormatApiError)
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 () => {
@ -2370,7 +2379,7 @@ describe('useGeminiStream', () => {
});
const countdownAfterOneSecond = result.current.pendingHistoryItems.find(
(item) => item.type === ('retry_countdown' as MessageType),
(item) => item.type === 'retry_countdown',
);
expect(countdownAfterOneSecond?.text).toContain(
'Retrying in 2 seconds',
@ -2388,7 +2397,7 @@ describe('useGeminiStream', () => {
(item) => item.type === MessageType.ERROR,
);
const remainingCountdown = result.current.pendingHistoryItems.find(
(item) => item.type === ('retry_countdown' as MessageType),
(item) => item.type === 'retry_countdown',
);
expect(remainingError).toBeUndefined();
expect(remainingCountdown).toBeUndefined();

View file

@ -128,8 +128,11 @@ export const useGeminiStream = (
useStateAndRef<HistoryItemWithoutId | null>(null);
const [pendingRetryErrorItem, setPendingRetryErrorItem] =
useState<HistoryItemWithoutId | null>(null);
const [pendingRetryCountdownItem, setPendingRetryCountdownItem] =
useState<HistoryItemWithoutId | null>(null);
const [
pendingRetryCountdownItem,
pendingRetryCountdownItemRef,
setPendingRetryCountdownItem,
] = useStateAndRef<HistoryItemWithoutId | null>(null);
const retryCountdownTimerRef = useRef<ReturnType<typeof setInterval> | null>(
null,
);
@ -208,23 +211,25 @@ export const useGeminiStream = (
stopRetryCountdownTimer();
setPendingRetryErrorItem(null);
setPendingRetryCountdownItem(null);
}, [stopRetryCountdownTimer]);
}, [setPendingRetryCountdownItem, stopRetryCountdownTimer]);
const startRetryCountdown = useCallback(
(retryInfo: {
reason: string;
message?: string;
attempt: number;
maxRetries: number;
delayMs: number;
}) => {
stopRetryCountdownTimer();
const startTime = Date.now();
const { reason, attempt, maxRetries, delayMs } = retryInfo;
const { message, attempt, maxRetries, delayMs } = retryInfo;
const retryReasonText =
message ?? t('Rate limit exceeded. Please wait and try again.');
// Error line stays static (red with ✕ prefix)
setPendingRetryErrorItem({
type: MessageType.ERROR,
text: t('Rate limit error: {{reason}}', { reason }),
text: retryReasonText,
});
// Countdown line updates every second (dim/secondary color)
@ -253,7 +258,7 @@ export const useGeminiStream = (
updateCountdown();
retryCountdownTimerRef.current = setInterval(updateCountdown, 1000);
},
[stopRetryCountdownTimer],
[setPendingRetryCountdownItem, stopRetryCountdownTimer],
);
useEffect(() => () => stopRetryCountdownTimer(), [stopRetryCountdownTimer]);
@ -947,7 +952,7 @@ export const useGeminiStream = (
// Show retry info if available (rate-limit / throttling errors)
if (event.retryInfo) {
startRetryCountdown(event.retryInfo);
} else {
} else if (!pendingRetryCountdownItemRef.current) {
clearRetryCountdown();
}
break;
@ -979,6 +984,7 @@ export const useGeminiStream = (
setThought,
pendingHistoryItemRef,
setPendingHistoryItem,
pendingRetryCountdownItemRef,
],
);