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:
Eyejoker
2026-03-23 22:01:06 +09:00
parent ee65330e66
commit d378598f2e
6 changed files with 168 additions and 14 deletions

View File

@@ -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',

View File

@@ -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,

View File

@@ -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);
}
}

View File

@@ -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();

View File

@@ -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
View 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,
};
}