token pools

src/token-pool.ts

@@ -0,0 +1,65 @@
+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);
+	}
+}