diff --git a/src/agent-runner-environment.ts b/src/agent-runner-environment.ts index 1f3f204..c32587e 100644 --- a/src/agent-runner-environment.ts +++ b/src/agent-runner-environment.ts @@ -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', diff --git a/src/index.ts b/src/index.ts index 8a33e6e..8e5c216 100644 --- a/src/index.ts +++ b/src/index.ts @@ -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, diff --git a/src/message-agent-executor.ts b/src/message-agent-executor.ts index c1b4d8f..968bf1c 100644 --- a/src/message-agent-executor.ts +++ b/src/message-agent-executor.ts @@ -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); } } diff --git a/src/message-turn-controller.ts b/src/message-turn-controller.ts index 72623af..7f72ae0 100644 --- a/src/message-turn-controller.ts +++ b/src/message-turn-controller.ts @@ -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(); diff --git a/src/task-suspension.ts b/src/task-suspension.ts index 514f808..b3f4555 100644 --- a/src/task-suspension.ts +++ b/src/task-suspension.ts @@ -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 }, diff --git a/src/token-rotation.ts b/src/token-rotation.ts new file mode 100644 index 0000000..944c7a3 --- /dev/null +++ b/src/token-rotation.ts @@ -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, + }; +}