feat: Claude OAuth token rotation on rate-limit
- New token-rotation module: stores multiple tokens from CLAUDE_CODE_OAUTH_TOKENS env var (comma-separated) - On rate-limit, rotates to next healthy token before falling back to Kimi provider - Also fixes: intermediate text routes to progress message when no subagents active (prevents message flooding)
This commit is contained in:
@@ -5,6 +5,7 @@ import path from 'path';
|
||||
import { GROUPS_DIR, TIMEZONE } from './config.js';
|
||||
import { isPairedRoomJid } from './db.js';
|
||||
import { readEnvFile } from './env.js';
|
||||
import { getCurrentToken } from './token-rotation.js';
|
||||
import {
|
||||
resolveGroupFolderPath,
|
||||
resolveGroupIpcPath,
|
||||
@@ -120,14 +121,14 @@ function prepareClaudeEnvironment(args: {
|
||||
args.env.ANTHROPIC_BASE_URL =
|
||||
args.envVars.ANTHROPIC_BASE_URL || process.env.ANTHROPIC_BASE_URL || '';
|
||||
}
|
||||
if (
|
||||
args.envVars.CLAUDE_CODE_OAUTH_TOKEN ||
|
||||
process.env.CLAUDE_CODE_OAUTH_TOKEN
|
||||
) {
|
||||
args.env.CLAUDE_CODE_OAUTH_TOKEN =
|
||||
{
|
||||
const oauthToken =
|
||||
args.envVars.CLAUDE_CODE_OAUTH_TOKEN ||
|
||||
process.env.CLAUDE_CODE_OAUTH_TOKEN ||
|
||||
'';
|
||||
getCurrentToken() ||
|
||||
process.env.CLAUDE_CODE_OAUTH_TOKEN;
|
||||
if (oauthToken) {
|
||||
args.env.CLAUDE_CODE_OAUTH_TOKEN = oauthToken;
|
||||
}
|
||||
}
|
||||
for (const key of [
|
||||
'CLAUDE_MODEL',
|
||||
|
||||
@@ -63,11 +63,15 @@ import { startUnifiedDashboard } from './unified-dashboard.js';
|
||||
import { Channel, NewMessage, RegisteredGroup } from './types.js';
|
||||
import { logger } from './logger.js';
|
||||
import { normalizeStoredSeqCursor } from './message-cursor.js';
|
||||
import { initTokenRotation } from './token-rotation.js';
|
||||
|
||||
// Re-export for backwards compatibility during refactor
|
||||
export { escapeXml, formatMessages } from './router.js';
|
||||
export { composeDashboardContent } from './dashboard-render.js';
|
||||
|
||||
// Initialize token rotation early (reads CLAUDE_CODE_OAUTH_TOKENS from env)
|
||||
initTokenRotation();
|
||||
|
||||
export async function sendFormattedChannelMessage(
|
||||
channels: Channel[],
|
||||
jid: string,
|
||||
|
||||
@@ -21,6 +21,7 @@ import {
|
||||
markPrimaryCooldown,
|
||||
} from './provider-fallback.js';
|
||||
import { shouldResetSessionOnAgentFailure } from './session-recovery.js';
|
||||
import { rotateToken, getTokenCount, markTokenHealthy } from './token-rotation.js';
|
||||
import type { RegisteredGroup } from './types.js';
|
||||
|
||||
export interface MessageAgentExecutorDeps {
|
||||
@@ -299,6 +300,18 @@ export async function runAgentForGroup(
|
||||
}
|
||||
: detectFallbackTrigger(errMsg);
|
||||
if (trigger.shouldFallback) {
|
||||
// Try rotating token before falling back to another provider
|
||||
if (getTokenCount() > 1 && rotateToken()) {
|
||||
logger.info(
|
||||
{ chatJid, group: group.name, runId, reason: trigger.reason },
|
||||
'Rate-limited, retrying with rotated token',
|
||||
);
|
||||
const retryAttempt = await runAttempt('claude');
|
||||
if (!retryAttempt.error) {
|
||||
markTokenHealthy();
|
||||
return 'success';
|
||||
}
|
||||
}
|
||||
return runFallbackAttempt(trigger.reason, trigger.retryAfterMs);
|
||||
}
|
||||
}
|
||||
@@ -362,6 +375,17 @@ export async function runAgentForGroup(
|
||||
}
|
||||
: detectFallbackTrigger(output.error);
|
||||
if (trigger.shouldFallback) {
|
||||
if (getTokenCount() > 1 && rotateToken()) {
|
||||
logger.info(
|
||||
{ chatJid, group: group.name, runId, reason: trigger.reason },
|
||||
'Rate-limited (output error), retrying with rotated token',
|
||||
);
|
||||
const retryAttempt = await runAttempt('claude');
|
||||
if (!retryAttempt.error) {
|
||||
markTokenHealthy();
|
||||
return 'success';
|
||||
}
|
||||
}
|
||||
return runFallbackAttempt(trigger.reason, trigger.retryAfterMs);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -113,10 +113,22 @@ export class MessageTurnController {
|
||||
}
|
||||
|
||||
if (result.phase === 'intermediate') {
|
||||
// Send as standalone message without touching progress state
|
||||
if (text) {
|
||||
this.lastIntermediateText = text;
|
||||
await this.options.channel.sendMessage(this.options.chatJid, text);
|
||||
if (this.subagents.size > 0) {
|
||||
// Subagents active — standalone to not disrupt progress
|
||||
this.lastIntermediateText = text;
|
||||
await this.options.channel.sendMessage(this.options.chatJid, text);
|
||||
} else if (this.progressMessageId) {
|
||||
// Progress exists — update heading
|
||||
this.previousProgressText = this.latestProgressText;
|
||||
this.latestProgressText = text;
|
||||
this.latestProgressTextForFinal = text;
|
||||
this.toolActivities = [];
|
||||
void this.syncTrackedProgressMessage();
|
||||
} else {
|
||||
// No progress yet — buffer (creates on next event)
|
||||
this.bufferProgress(text);
|
||||
}
|
||||
}
|
||||
if (!this.poisonedSessionDetected) {
|
||||
this.resetIdleTimer();
|
||||
|
||||
@@ -105,10 +105,7 @@ export function evaluateTaskSuspension(
|
||||
/**
|
||||
* Apply suspension to a task in the DB.
|
||||
*/
|
||||
export function suspendTask(
|
||||
taskId: string,
|
||||
suspendedUntil: string,
|
||||
): void {
|
||||
export function suspendTask(taskId: string, suspendedUntil: string): void {
|
||||
updateTask(taskId, { suspended_until: suspendedUntil });
|
||||
logger.info(
|
||||
{ taskId, suspendedUntil },
|
||||
|
||||
116
src/token-rotation.ts
Normal file
116
src/token-rotation.ts
Normal file
@@ -0,0 +1,116 @@
|
||||
/**
|
||||
* OAuth Token Rotation
|
||||
*
|
||||
* Rotates between multiple CLAUDE_CODE_OAUTH_TOKEN values when
|
||||
* rate-limited. Tokens are stored as comma-separated values in
|
||||
* CLAUDE_CODE_OAUTH_TOKENS env var. Falls through to the single
|
||||
* CLAUDE_CODE_OAUTH_TOKEN if multi-token is not configured.
|
||||
*
|
||||
* On rate-limit: rotate to next token
|
||||
* All exhausted: fall through to provider fallback (Kimi etc.)
|
||||
*/
|
||||
|
||||
import { logger } from './logger.js';
|
||||
|
||||
interface TokenState {
|
||||
token: string;
|
||||
rateLimitedUntil: number | null;
|
||||
}
|
||||
|
||||
const tokens: TokenState[] = [];
|
||||
let currentIndex = 0;
|
||||
let initialized = false;
|
||||
|
||||
export function initTokenRotation(): void {
|
||||
if (initialized) return;
|
||||
initialized = true;
|
||||
|
||||
const multi = process.env.CLAUDE_CODE_OAUTH_TOKENS;
|
||||
const single = process.env.CLAUDE_CODE_OAUTH_TOKEN;
|
||||
|
||||
const raw = multi
|
||||
? multi.split(',').map((t) => t.trim()).filter(Boolean)
|
||||
: single
|
||||
? [single]
|
||||
: [];
|
||||
|
||||
for (const token of raw) {
|
||||
tokens.push({ token, rateLimitedUntil: null });
|
||||
}
|
||||
|
||||
if (tokens.length > 1) {
|
||||
logger.info(
|
||||
{ count: tokens.length },
|
||||
`Token rotation initialized with ${tokens.length} tokens`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/** Get the current active token. */
|
||||
export function getCurrentToken(): string | undefined {
|
||||
if (tokens.length === 0) return process.env.CLAUDE_CODE_OAUTH_TOKEN;
|
||||
return tokens[currentIndex % tokens.length]?.token;
|
||||
}
|
||||
|
||||
/**
|
||||
* Try to rotate to the next available (non-rate-limited) token.
|
||||
* Returns true if a fresh token was found, false if all are exhausted.
|
||||
*/
|
||||
export function rotateToken(): boolean {
|
||||
if (tokens.length <= 1) return false;
|
||||
|
||||
const now = Date.now();
|
||||
// Mark current as rate-limited (default 1 hour)
|
||||
tokens[currentIndex].rateLimitedUntil = now + 3_600_000;
|
||||
|
||||
// Find next available token
|
||||
for (let i = 1; i < tokens.length; i++) {
|
||||
const idx = (currentIndex + i) % tokens.length;
|
||||
const state = tokens[idx];
|
||||
if (!state.rateLimitedUntil || state.rateLimitedUntil <= now) {
|
||||
state.rateLimitedUntil = null;
|
||||
currentIndex = idx;
|
||||
logger.info(
|
||||
{ tokenIndex: currentIndex, totalTokens: tokens.length },
|
||||
`Rotated to token #${currentIndex + 1}/${tokens.length}`,
|
||||
);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
logger.warn(
|
||||
{ totalTokens: tokens.length },
|
||||
'All tokens are rate-limited, falling through to provider fallback',
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
/** Clear rate-limit flag for the current token (on successful response). */
|
||||
export function markTokenHealthy(): void {
|
||||
if (tokens.length === 0) return;
|
||||
const state = tokens[currentIndex];
|
||||
if (state?.rateLimitedUntil) {
|
||||
state.rateLimitedUntil = null;
|
||||
}
|
||||
}
|
||||
|
||||
/** Number of configured tokens. */
|
||||
export function getTokenCount(): number {
|
||||
return tokens.length;
|
||||
}
|
||||
|
||||
/** Diagnostic info. */
|
||||
export function getTokenRotationInfo(): {
|
||||
total: number;
|
||||
currentIndex: number;
|
||||
rateLimited: number;
|
||||
} {
|
||||
const now = Date.now();
|
||||
return {
|
||||
total: tokens.length,
|
||||
currentIndex,
|
||||
rateLimited: tokens.filter(
|
||||
(t) => t.rateLimitedUntil && t.rateLimitedUntil > now,
|
||||
).length,
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user