66 lines
2.1 KiB
TypeScript
66 lines
2.1 KiB
TypeScript
const DEFAULT_COOLDOWN = 15 * 60 * 1000;
|
|
|
|
type TokenState = {
|
|
token: string;
|
|
cooldownUntil: number; // 0 = available now
|
|
lastError?: {code: number, message: string};
|
|
};
|
|
|
|
export class TokenPoolExhaustedError extends Error {
|
|
constructor(public tokens: Record<string, {code: number, message: string}>) {
|
|
super(`All tokens exhausted:\n${Object.entries(tokens).map(([t, e]) => `${t}: [${e.code}] ${e.message}`).join('\n')}`);
|
|
this.name = 'TokenPoolExhaustedError';
|
|
}
|
|
}
|
|
|
|
export class TokenPool {
|
|
private states: TokenState[];
|
|
|
|
constructor(...tokens: string[]) {
|
|
this.states = tokens.map(token => ({token, cooldownUntil: 0}));
|
|
}
|
|
|
|
private preview(token: string): string {
|
|
return token.length <= 8 ? '****' : `${token.slice(0, 4)}...${token.slice(-4)}`;
|
|
}
|
|
|
|
/** Anthropic & OpenAI SDKs both attach `status` to thrown errors */
|
|
private statusCode(err: any): number {
|
|
return err?.status ?? err?.response?.status ?? err?.statusCode;
|
|
}
|
|
|
|
private retryAfter(err: any): number {
|
|
const headers = err?.headers || err?.response?.headers;
|
|
const raw = headers?.get?.('retry-after') ?? headers?.['retry-after'];
|
|
if(raw) {
|
|
const seconds = Number(raw);
|
|
if(!isNaN(seconds)) return Date.now() + seconds * 1000;
|
|
const date = new Date(raw).getTime();
|
|
if(!isNaN(date)) return date;
|
|
}
|
|
return Date.now() + DEFAULT_COOLDOWN;
|
|
}
|
|
|
|
async run<T>(fn: (token: string) => Promise<T>): Promise<T> {
|
|
const now = Date.now();
|
|
for(const state of this.states) {
|
|
if(state.cooldownUntil > now) continue;
|
|
try {
|
|
const result = await fn(state.token);
|
|
state.cooldownUntil = 0;
|
|
state.lastError = undefined;
|
|
return result;
|
|
} catch(err: any) {
|
|
const code = this.statusCode(err);
|
|
if(![401, 403, 429].includes(code)) throw err;
|
|
state.cooldownUntil = code === 429 ? this.retryAfter(err) : Date.now() + DEFAULT_COOLDOWN;
|
|
state.lastError = {code, message: err?.message || 'Unknown error'};
|
|
}
|
|
}
|
|
|
|
const failures: Record<string, {code: number, message: string}> = {};
|
|
this.states.forEach(s => { if(s.lastError) failures[this.preview(s.token)] = s.lastError; });
|
|
throw new TokenPoolExhaustedError(failures);
|
|
}
|
|
}
|