diff --git a/README.md b/README.md index 61a6aa8..e4a4e52 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,10 @@ # EJClaw +![Claude Agent SDK](https://img.shields.io/badge/Claude_Agent_SDK-0.2.81-blueviolet) +![Codex SDK](https://img.shields.io/badge/Codex_SDK-0.115.0-green) +![Node](https://img.shields.io/badge/Node-24+-339933?logo=nodedotjs&logoColor=white) +![Discord](https://img.shields.io/badge/Discord-Bot-5865F2?logo=discord&logoColor=white) + Dual-agent AI assistant (Claude Code + Codex) over Discord. Originally derived from [qwibitai/nanoclaw](https://github.com/qwibitai/nanoclaw), now maintained as EJClaw for personal production use. @@ -47,7 +52,7 @@ Each agent has access to: ### Prerequisites - Linux (Ubuntu 22.04+) or macOS -- Node.js 20+ (fnm recommended) +- Node.js 24+ (fnm recommended) - [Claude Code CLI](https://docs.anthropic.com/en/docs/claude-code) - [Codex CLI](https://github.com/openai/codex) (`npm install -g @openai/codex`) - Bun 1.0+ (for browser automation) diff --git a/groups/nanoclaw/CLAUDE.md b/groups/nanoclaw/CLAUDE.md deleted file mode 100644 index 933d13d..0000000 --- a/groups/nanoclaw/CLAUDE.md +++ /dev/null @@ -1,2 +0,0 @@ -This channel works with the nanoclaw project at /home/clone-ej/nanoclaw. -When you reach a conclusion or need human judgment, mention @눈쟁이 in your message. diff --git a/runners/agent-runner/package.json b/runners/agent-runner/package.json index 1154b6e..14b3084 100644 --- a/runners/agent-runner/package.json +++ b/runners/agent-runner/package.json @@ -9,7 +9,7 @@ "start": "node dist/index.js" }, "dependencies": { - "@anthropic-ai/claude-agent-sdk": "^0.2.76", + "@anthropic-ai/claude-agent-sdk": "^0.2.81", "@modelcontextprotocol/sdk": "^1.12.1", "cron-parser": "^5.0.0", "zod": "^4.0.0" diff --git a/runners/agent-runner/pnpm-lock.yaml b/runners/agent-runner/pnpm-lock.yaml index a5dcc26..a0ac373 100644 --- a/runners/agent-runner/pnpm-lock.yaml +++ b/runners/agent-runner/pnpm-lock.yaml @@ -9,8 +9,8 @@ importers: .: dependencies: '@anthropic-ai/claude-agent-sdk': - specifier: ^0.2.76 - version: 0.2.76(zod@4.3.6) + specifier: ^0.2.81 + version: 0.2.81(zod@4.3.6) '@modelcontextprotocol/sdk': specifier: ^1.12.1 version: 1.27.1(zod@4.3.6) @@ -30,8 +30,8 @@ importers: packages: - '@anthropic-ai/claude-agent-sdk@0.2.76': - resolution: {integrity: sha512-HZxvnT8ZWkzCnQygaYCA0dl8RSUzuVbxE1YG4ecy6vh4nQbTT36CxUxBy+QVdR12pPQluncC0mCOLhI2918Eaw==} + '@anthropic-ai/claude-agent-sdk@0.2.81': + resolution: {integrity: sha512-CBeebgibBEN/DWOQGZN67vhuTG55RbI1hlsFSSoZ4uA/Io3lw04eHTE2ISCmdbqyJaefYTt6GKZei1nP0TQMNw==} engines: {node: '>=18.0.0'} peerDependencies: zod: ^4.0.0 @@ -512,7 +512,7 @@ packages: snapshots: - '@anthropic-ai/claude-agent-sdk@0.2.76(zod@4.3.6)': + '@anthropic-ai/claude-agent-sdk@0.2.81(zod@4.3.6)': dependencies: zod: 4.3.6 optionalDependencies: diff --git a/runners/agent-runner/src/index.ts b/runners/agent-runner/src/index.ts index 56d44a5..a476a76 100644 --- a/runners/agent-runner/src/index.ts +++ b/runners/agent-runner/src/index.ts @@ -35,6 +35,9 @@ interface ContainerInput { interface ContainerOutput { status: 'success' | 'error'; phase?: 'progress' | 'final' | 'tool-activity' | 'intermediate'; + agentId?: string; + agentLabel?: string; + agentDone?: boolean; result: string | null; newSessionId?: string; error?: string; @@ -672,19 +675,35 @@ async function runQuery( if (message.type === 'system' && (message as { subtype?: string }).subtype === 'task_notification') { const tn = message as { task_id: string; status: string; summary: string }; log(`Task notification: task=${tn.task_id} status=${tn.status} summary=${tn.summary}`); + if (tn.status === 'completed' || tn.status === 'error' || tn.status === 'cancelled') { + writeOutput({ + status: 'success', + phase: 'progress', + agentId: tn.task_id, + agentDone: true, + result: tn.summary || null, + newSessionId, + }); + } } if (message.type === 'system' && (message as { subtype?: string }).subtype === 'task_progress') { const tp = message as Record; + const taskId = typeof tp.task_id === 'string' ? tp.task_id : undefined; const summary = typeof tp.summary === 'string' ? tp.summary : ''; const description = typeof tp.description === 'string' ? tp.description : ''; - if (description) { + if (description && description.length <= 80) { + // Short tool description → show as sub-line in progress writeOutput({ status: 'success', phase: 'tool-activity', result: description, + agentId: taskId, newSessionId, }); + } else if (description) { + // Long AI summary → skip (too long for progress sub-line) + log(`Skipping long task_progress description (${description.length} chars)`); } } @@ -697,6 +716,8 @@ async function runQuery( status: 'success', phase: 'progress', result: `🔄 ${desc}`, + agentId: ts.task_id, + agentLabel: desc, newSessionId, }); } diff --git a/runners/codex-runner/pnpm-lock.yaml b/runners/codex-runner/pnpm-lock.yaml index e397a43..572f06f 100644 --- a/runners/codex-runner/pnpm-lock.yaml +++ b/runners/codex-runner/pnpm-lock.yaml @@ -8,7 +8,7 @@ importers: .: dependencies: - '@openai/codex-sdk': + '@openai/codex': specifier: ^0.115.0 version: 0.115.0 devDependencies: @@ -21,10 +21,6 @@ importers: packages: - '@openai/codex-sdk@0.115.0': - resolution: {integrity: sha512-BPoPhim0uUm3rzugY7JFaFJ+rPG/wMRhvKNFii//Dp3kTb+gFy3jrip7ijXqhswsnqXM3nwTiv7kYHt1+TMUPg==} - engines: {node: '>=18'} - '@openai/codex@0.115.0': resolution: {integrity: sha512-uu689DHUzvuPcb39hJ+7fqy++TAvY32w9VggDpcz3HS0Sx0WadWoAPPcMK547P2T6AqfMsAtA8kspkR3tqErOg==} engines: {node: '>=16'} @@ -79,10 +75,6 @@ packages: snapshots: - '@openai/codex-sdk@0.115.0': - dependencies: - '@openai/codex': 0.115.0 - '@openai/codex@0.115.0': optionalDependencies: '@openai/codex-darwin-arm64': '@openai/codex@0.115.0-darwin-arm64' diff --git a/src/agent-runner-environment.ts b/src/agent-runner-environment.ts index def0390..34ea8af 100644 --- a/src/agent-runner-environment.ts +++ b/src/agent-runner-environment.ts @@ -5,6 +5,8 @@ import path from 'path'; import { GROUPS_DIR, TIMEZONE } from './config.js'; import { isPairedRoomJid } from './db.js'; import { readEnvFile } from './env.js'; +import { getActiveCodexAuthPath } from './codex-token-rotation.js'; +import { getCurrentToken } from './token-rotation.js'; import { resolveGroupFolderPath, resolveGroupIpcPath, @@ -120,14 +122,15 @@ 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 = + { + // Token rotation takes priority over static .env value + const oauthToken = + getCurrentToken() || args.envVars.CLAUDE_CODE_OAUTH_TOKEN || - process.env.CLAUDE_CODE_OAUTH_TOKEN || - ''; + process.env.CLAUDE_CODE_OAUTH_TOKEN; + if (oauthToken) { + args.env.CLAUDE_CODE_OAUTH_TOKEN = oauthToken; + } } for (const key of [ 'CLAUDE_MODEL', @@ -182,7 +185,11 @@ function prepareCodexSessionEnvironment(args: { const sessionCodexDir = path.join(args.sessionRootDir, '.codex'); fs.mkdirSync(sessionCodexDir, { recursive: true }); - const authSrc = path.join(hostCodexDir, 'auth.json'); + const rotatedAuthSrc = getActiveCodexAuthPath(); + const authSrc = + rotatedAuthSrc && fs.existsSync(rotatedAuthSrc) + ? rotatedAuthSrc + : path.join(hostCodexDir, 'auth.json'); const authDst = path.join(sessionCodexDir, 'auth.json'); if (fs.existsSync(authSrc)) fs.copyFileSync(authSrc, authDst); for (const file of ['config.toml', 'config.json']) { diff --git a/src/agent-runner.ts b/src/agent-runner.ts index 4e63061..574fd47 100644 --- a/src/agent-runner.ts +++ b/src/agent-runner.ts @@ -43,6 +43,9 @@ export interface AgentOutput { status: 'success' | 'error'; result: string | null; phase?: 'progress' | 'final' | 'tool-activity' | 'intermediate'; + agentId?: string; + agentLabel?: string; + agentDone?: boolean; newSessionId?: string; error?: string; } diff --git a/src/claude-usage.test.ts b/src/claude-usage.test.ts index e868e2b..884a7bb 100644 --- a/src/claude-usage.test.ts +++ b/src/claude-usage.test.ts @@ -1,53 +1,31 @@ import { describe, expect, it } from 'vitest'; -import { parseClaudeUsagePanel } from './claude-usage.js'; +import type { ClaudeUsageData } from './claude-usage.js'; -describe('parseClaudeUsagePanel', () => { - it('parses session and weekly usage from Claude CLI panel output', () => { - const sample = ` -Settings: Status Config Usage -Loading usage data... - -Current session -████ 4% used -Resets in 8m (Asia/Seoul) - -Current week (all models) -███████████████████████████████████████ 78% used -Resets Mar 17 at 10pm (Asia/Seoul) - -Current week (Sonnet only) -███ 6% used -Resets Mar 17 at 11pm (Asia/Seoul) - -Extra usage -Extra usage not enabled • /extra-usage to enable -`; - - expect(parseClaudeUsagePanel(sample)).toEqual({ - five_hour: { - utilization: 4, - resets_at: 'Resets in 8m (Asia/Seoul)', +describe('ClaudeUsageData', () => { + it('represents the API response structure correctly', () => { + const data: ClaudeUsageData = { + five_hour: { utilization: 45.2, resets_at: '2026-03-23T17:00:00Z' }, + seven_day: { utilization: 72.1, resets_at: '2026-03-29T00:00:00Z' }, + seven_day_sonnet: { + utilization: 60.0, + resets_at: '2026-03-29T00:00:00Z', }, - seven_day: { - utilization: 78, - resets_at: 'Resets Mar 17 at 10pm (Asia/Seoul)', - }, - }); + seven_day_opus: { utilization: 80.0, resets_at: '2026-03-29T00:00:00Z' }, + }; + + expect(data.five_hour?.utilization).toBe(45.2); + expect(data.seven_day?.utilization).toBe(72.1); + expect(data.seven_day_sonnet?.utilization).toBe(60.0); + expect(data.seven_day_opus?.utilization).toBe(80.0); }); - it('converts percent left into used percent', () => { - const sample = ` -Current session -60% left -Resets in 1h -`; + it('allows partial data (only five_hour)', () => { + const data: ClaudeUsageData = { + five_hour: { utilization: 10, resets_at: '2026-03-23T17:00:00Z' }, + }; - expect(parseClaudeUsagePanel(sample)).toEqual({ - five_hour: { - utilization: 40, - resets_at: 'Resets in 1h', - }, - }); + expect(data.five_hour?.utilization).toBe(10); + expect(data.seven_day).toBeUndefined(); }); }); diff --git a/src/claude-usage.ts b/src/claude-usage.ts index 481d00f..617f772 100644 --- a/src/claude-usage.ts +++ b/src/claude-usage.ts @@ -1,285 +1,279 @@ -import { spawn } from 'child_process'; +/** + * Claude Usage API + * + * Fetches usage data directly from the Anthropic OAuth API. + * Supports multiple tokens for rotation-aware usage checking. + */ + import fs from 'fs'; -import os from 'os'; import path from 'path'; +import { DATA_DIR } from './config.js'; import { logger } from './logger.js'; +import { getCurrentToken, getAllTokens } from './token-rotation.js'; + +const USAGE_CACHE_FILE = path.join(DATA_DIR, 'claude-usage-cache.json'); + +const PROFILE_ENDPOINT = 'https://api.anthropic.com/api/oauth/profile'; export interface ClaudeUsageData { five_hour?: { utilization: number; resets_at: string }; seven_day?: { utilization: number; resets_at: string }; + seven_day_sonnet?: { utilization: number; resets_at: string }; + seven_day_opus?: { utilization: number; resets_at: string }; +} + +const USAGE_ENDPOINT = 'https://api.anthropic.com/api/oauth/usage'; +const FETCH_TIMEOUT_MS = 10_000; + +interface UsageApiResponse { + five_hour?: { utilization: number; resets_at?: string }; + seven_day?: { utilization: number; resets_at?: string }; + seven_day_sonnet?: { utilization: number; resets_at?: string }; + seven_day_opus?: { utilization: number; resets_at?: string }; +} + +function mapWindow(w?: { + utilization: number; + resets_at?: string; +}): { utilization: number; resets_at: string } | undefined { + if (!w) return undefined; + return { utilization: w.utilization, resets_at: w.resets_at || '' }; +} + +// ── Disk cache for usage data (survives restarts, 429s) ── + +interface UsageCacheEntry { + usage: ClaudeUsageData; + fetchedAt: number; +} + +let usageDiskCache: Record = {}; +let diskCacheLoaded = false; + +function loadUsageDiskCache(): void { + if (diskCacheLoaded) return; + diskCacheLoaded = true; + try { + if (fs.existsSync(USAGE_CACHE_FILE)) { + usageDiskCache = JSON.parse(fs.readFileSync(USAGE_CACHE_FILE, 'utf-8')); + } + } catch { + /* start fresh */ + } +} + +function saveUsageDiskCache(): void { + try { + fs.writeFileSync(USAGE_CACHE_FILE, JSON.stringify(usageDiskCache)); + } catch { + /* best effort */ + } +} + +function cacheKey(token: string): string { + return token.slice(0, 12); +} + +// Rate limit: at most one API call per token per 5 minutes +const MIN_FETCH_INTERVAL_MS = 300_000; + +async function fetchUsageForToken( + token: string, +): Promise { + loadUsageDiskCache(); + + // Return cached data if fetched recently (avoid API rate-limit) + const key = cacheKey(token); + const cached = usageDiskCache[key]; + if (cached && Date.now() - cached.fetchedAt < MIN_FETCH_INTERVAL_MS) { + return cached.usage; + } + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS); + + try { + const res = await fetch(USAGE_ENDPOINT, { + method: 'GET', + headers: { + Authorization: `Bearer ${token}`, + 'anthropic-beta': 'oauth-2025-04-20', + 'User-Agent': 'ejclaw/1.0', + }, + signal: controller.signal, + }); + + if (res.status === 401) { + logger.warn('Claude usage API: token expired or invalid (401)'); + return null; + } + if (res.status === 429) { + logger.warn( + 'Claude usage API: rate limited (429), returning cached data', + ); + return usageDiskCache[cacheKey(token)]?.usage ?? null; + } + if (!res.ok) { + logger.warn( + { status: res.status }, + `Claude usage API: unexpected status ${res.status}`, + ); + return usageDiskCache[cacheKey(token)]?.usage ?? null; + } + + const data = (await res.json()) as UsageApiResponse; + + const result: ClaudeUsageData = { + five_hour: mapWindow(data.five_hour), + seven_day: mapWindow(data.seven_day), + seven_day_sonnet: mapWindow(data.seven_day_sonnet), + seven_day_opus: mapWindow(data.seven_day_opus), + }; + + // Persist to disk cache + usageDiskCache[cacheKey(token)] = { usage: result, fetchedAt: Date.now() }; + saveUsageDiskCache(); + + return result; + } catch (err) { + if ((err as Error).name === 'AbortError') { + logger.warn('Claude usage API: request timed out'); + } else { + logger.warn({ err }, 'Claude usage API: fetch failed'); + } + return usageDiskCache[cacheKey(token)]?.usage ?? null; + } finally { + clearTimeout(timer); + } +} + +/** + * Fetch Claude usage via the OAuth API. + * Uses the current active token from rotation. + */ +export async function fetchClaudeUsage(): Promise { + const token = getCurrentToken() || process.env.CLAUDE_CODE_OAUTH_TOKEN; + if (!token) { + logger.debug('No Claude OAuth token available for usage check'); + return null; + } + return fetchUsageForToken(token); } export interface ClaudeAccountProfile { - index: number; - planType: string; + email: string; + planType: string; // "max", "pro", "free" } -export interface ClaudeAccountUsage { - index: number; - isActive: boolean; - isRateLimited: boolean; - usage: ClaudeUsageData | null; -} +const profileCache = new Map(); -const CLAUDE_EXPECT_TIMEOUT_MS = 25000; -const ANSI_RE = /\u001b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])/g; -const HOST_CLAUDE_DIR = path.join(os.homedir(), '.claude'); -const CLAUDE_ACCOUNTS_DIR = path.join(os.homedir(), '.claude-accounts'); - -let cachedProfiles: ClaudeAccountProfile[] = []; - -const EXPECT_PROGRAM = ` -set timeout 20 -log_user 1 -match_max 200000 -set binary $env(CLAUDE_BINARY) -spawn -noecho -- $binary --setting-sources user --allowed-tools "" -expect { - -re "Do you trust the files in this folder\\\\?" { send "y\\r"; exp_continue } - -re "Quick safety check:" { send "\\r"; exp_continue } - -re "Yes, I trust this folder" { send "\\r"; exp_continue } - -re "Ready to code here\\\\?" { send "\\r"; exp_continue } - -re "Press Enter to continue" { send "\\r"; exp_continue } - timeout {} -} -send "/usage\\r" -set deadline [expr {[clock seconds] + 20}] -while {[clock seconds] < $deadline} { - expect { - -re "Do you trust the files in this folder\\\\?" { send "y\\r"; exp_continue } - -re "Quick safety check:" { send "\\r"; exp_continue } - -re "Yes, I trust this folder" { send "\\r"; exp_continue } - -re "Ready to code here\\\\?" { send "\\r"; exp_continue } - -re "Press Enter to continue" { send "\\r"; exp_continue } - -re "Current session" { after 2000; exit 0 } - -re "Failed to load usage data" { after 500; exit 2 } - eof { exit 3 } - timeout { send "\\r" } - } -} -exit 4 -`; - -function normalizeLines(rawText: string): string[] { - return rawText - .replace(ANSI_RE, '') - .replace(/\r/g, '\n') - .split('\n') - .map((line) => line.replace(/\s+/g, ' ').trim()) - .filter(Boolean); -} - -function parsePercent(windowText: string): number | null { - const match = windowText.match(/(\d{1,3})%\s*(used|left)\b/i); - if (!match) return null; - const value = parseInt(match[1], 10); - if (Number.isNaN(value)) return null; - return match[2].toLowerCase() === 'left' ? 100 - value : value; -} - -function parseWindow( - lines: string[], - labels: string[], -): { utilization: number; resets_at: string } | null { - const normalizedLabels = labels.map((label) => label.toLowerCase()); - for (let i = 0; i < lines.length; i++) { - const line = lines[i].toLowerCase(); - if (!normalizedLabels.some((label) => line.includes(label))) continue; - - const windowLines = lines.slice(i, i + 6); - const windowText = windowLines.join('\n'); - const utilization = parsePercent(windowText); - if (utilization === null) continue; - - const resetLine = windowLines.find((candidate) => - candidate.toLowerCase().startsWith('resets'), - ); - - return { - utilization, - resets_at: resetLine || '', - }; - } - - return null; -} - -export function parseClaudeUsagePanel(rawText: string): ClaudeUsageData | null { - const lines = normalizeLines(rawText); - if (lines.length === 0) return null; - if ( - lines.some((line) => - line.toLowerCase().includes('failed to load usage data'), - ) - ) { - return null; - } - - const fiveHour = parseWindow(lines, ['Current session']); - if (!fiveHour) return null; - - const sevenDay = - parseWindow(lines, ['Current week (all models)']) || - parseWindow(lines, [ - 'Current week (Sonnet only)', - 'Current week (Sonnet)', - ]) || - parseWindow(lines, ['Current week (Opus)']); - - return { - five_hour: fiveHour, - ...(sevenDay ? { seven_day: sevenDay } : {}), - }; -} - -function listClaudeConfigDirs(): string[] { - if (!fs.existsSync(CLAUDE_ACCOUNTS_DIR)) { - return [HOST_CLAUDE_DIR]; - } - - const dirs = fs - .readdirSync(CLAUDE_ACCOUNTS_DIR, { withFileTypes: true }) - .filter((entry) => entry.isDirectory() && /^\d+$/.test(entry.name)) - .sort((a, b) => Number(a.name) - Number(b.name)) - .map((entry) => path.join(CLAUDE_ACCOUNTS_DIR, entry.name)); - - return dirs.length > 0 ? dirs : [HOST_CLAUDE_DIR]; -} - -function inferClaudePlanType(configDir: string): string { - const credentialsPath = path.join(configDir, '.credentials.json'); - return fs.existsSync(credentialsPath) ? 'OAuth' : 'Default'; -} - -function ensureProfiles(): ClaudeAccountProfile[] { - const profiles = listClaudeConfigDirs().map((configDir, index) => ({ - index, - planType: inferClaudePlanType(configDir), - })); - cachedProfiles = profiles; - return profiles; -} - -function readCredentialsFile(configDir: string): string | null { +async function fetchProfileForToken( + token: string, +): Promise { + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS); try { - const credentialsPath = path.join(configDir, '.credentials.json'); - if (!fs.existsSync(credentialsPath)) return null; - return fs.readFileSync(credentialsPath, 'utf8'); + const res = await fetch(PROFILE_ENDPOINT, { + method: 'GET', + headers: { + Authorization: `Bearer ${token}`, + 'anthropic-beta': 'oauth-2025-04-20', + }, + signal: controller.signal, + }); + if (!res.ok) return null; + const data = (await res.json()) as { + account?: { + email?: string; + has_claude_max?: boolean; + has_claude_pro?: boolean; + }; + organization?: { organization_type?: string }; + }; + const orgType = data.organization?.organization_type || ''; + const planType = data.account?.has_claude_max + ? 'max' + : data.account?.has_claude_pro + ? 'pro' + : orgType.replace('claude_', '') || 'free'; + return { + email: data.account?.email || '?', + planType, + }; } catch { return null; + } finally { + clearTimeout(timer); } } -function detectActiveClaudeIndex(configDirs: string[]): number { - if (configDirs.length <= 1) return 0; - - const hostCredentials = readCredentialsFile(HOST_CLAUDE_DIR); - if (!hostCredentials) return 0; - - const matchIndex = configDirs.findIndex( - (configDir) => readCredentialsFile(configDir) === hostCredentials, - ); - return matchIndex >= 0 ? matchIndex : 0; -} - -export async function fetchClaudeUsageViaCli( - binary = 'claude', - configDir?: string, -): Promise { - return new Promise((resolve) => { - let output = ''; - let finished = false; - - const finish = (value: ClaudeUsageData | null) => { - if (finished) return; - finished = true; - clearTimeout(timer); - resolve(value); - }; - - let proc: ReturnType | null = null; - try { - proc = spawn('expect', ['-c', EXPECT_PROGRAM], { - stdio: ['ignore', 'pipe', 'pipe'], - env: { - ...(process.env as Record), - CLAUDE_BINARY: binary, - ...(configDir ? { CLAUDE_CONFIG_DIR: configDir } : {}), - }, - }); - } catch (err) { - logger.debug({ err }, 'Claude CLI PTY probe unavailable'); - finish(null); - return; +/** + * Fetch profiles for all Claude tokens (cached, called once on startup). + */ +export async function fetchAllClaudeProfiles(): Promise { + const allTokens = getAllTokens(); + for (const t of allTokens) { + const profile = await fetchProfileForToken(t.token); + if (profile) { + profileCache.set(t.index, profile); + logger.info( + { account: t.index + 1, plan: profile.planType, email: profile.email }, + `Claude account #${t.index + 1}: ${profile.planType}`, + ); } - - const timer = setTimeout(() => { - try { - proc?.kill('SIGTERM'); - } catch { - /* ignore */ - } - finish(null); - }, CLAUDE_EXPECT_TIMEOUT_MS); - - if (!proc.stdout || !proc.stderr) { - finish(null); - return; - } - - proc.stdout.setEncoding('utf8'); - proc.stderr.setEncoding('utf8'); - proc.stdout.on('data', (chunk: string) => { - output += chunk; - }); - proc.stderr.on('data', (chunk: string) => { - output += chunk; - }); - proc.on('error', (err) => { - logger.debug({ err }, 'Claude CLI PTY probe failed to start'); - finish(null); - }); - proc.on('close', () => { - const parsed = parseClaudeUsagePanel(output); - if (!parsed && output.trim()) { - logger.debug( - { tail: output.slice(-400) }, - 'Claude CLI PTY probe produced unparsable output', - ); - } - finish(parsed); - }); - }); -} - -export async function fetchAllClaudeProfiles(): Promise< - ClaudeAccountProfile[] -> { - return ensureProfiles(); + } } export function getClaudeProfile( index: number, ): ClaudeAccountProfile | undefined { - return (cachedProfiles.length > 0 ? cachedProfiles : ensureProfiles()).find( - (profile) => profile.index === index, - ); + return profileCache.get(index); } +export interface ClaudeAccountUsage { + index: number; + masked: string; + isActive: boolean; + isRateLimited: boolean; + usage: ClaudeUsageData | null; +} + +/** + * Fetch usage for ALL configured tokens. + * Returns per-account usage for dashboard display. + */ export async function fetchAllClaudeUsage(): Promise { - const configDirs = listClaudeConfigDirs(); - const profiles = ensureProfiles(); - const activeIndex = detectActiveClaudeIndex(configDirs); + const allTokens = getAllTokens(); + logger.debug({ tokenCount: allTokens.length }, 'fetchAllClaudeUsage called'); + if (allTokens.length === 0) { + // Single token fallback + const token = process.env.CLAUDE_CODE_OAUTH_TOKEN; + if (!token) return []; + const usage = await fetchUsageForToken(token); + return [ + { + index: 0, + masked: `${token.slice(0, 20)}...${token.slice(-4)}`, + isActive: true, + isRateLimited: false, + usage, + }, + ]; + } - const usages = await Promise.all( - configDirs.map((configDir) => fetchClaudeUsageViaCli('claude', configDir)), - ); - - return profiles.map((profile, index) => ({ - index: profile.index, - isActive: index === activeIndex, - isRateLimited: false, - usage: usages[index] || null, - })); + const results: ClaudeAccountUsage[] = []; + for (const t of allTokens) { + const usage = await fetchUsageForToken(t.token); + results.push({ + index: t.index, + masked: t.masked, + isActive: t.isActive, + isRateLimited: t.isRateLimited, + usage, + }); + } + return results; } + +// Legacy alias +export const fetchClaudeUsageViaCli = fetchClaudeUsage; diff --git a/src/codex-token-rotation.ts b/src/codex-token-rotation.ts index 7461559..f02932a 100644 --- a/src/codex-token-rotation.ts +++ b/src/codex-token-rotation.ts @@ -1,169 +1,347 @@ +/** + * Codex OAuth Token Rotation + * + * Rotates between multiple Codex (ChatGPT) OAuth accounts when + * rate-limited. Each account is stored as a separate auth.json in + * ~/.codex-accounts/{n}/auth.json. + * + * The active account's auth.json is copied to the session directory + * before each agent spawn (existing behavior in agent-runner-environment). + * On rate-limit, we rotate to the next account. + */ + import fs from 'fs'; import os from 'os'; import path from 'path'; -export interface CodexAccountState { +import { DATA_DIR } from './config.js'; +import { logger } from './logger.js'; + +const STATE_FILE = path.join(DATA_DIR, 'codex-rotation-state.json'); + +interface CodexAccount { index: number; + authPath: string; + accountId: string; planType: string; - isActive: boolean; - isRateLimited: boolean; - cachedUsagePct?: number; + subscriptionUntil: string | null; + rateLimitedUntil: number | null; + lastUsagePct?: number; + lastUsageD7Pct?: number; resetAt?: string; - cachedUsageD7Pct?: number; resetD7At?: string; } -const HOST_CODEX_DIR = path.join(os.homedir(), '.codex'); -const CODEX_ACCOUNTS_DIR = path.join(os.homedir(), '.codex-accounts'); - -let activeIndexCache: number | null = null; -const rateLimitedAccounts = new Set(); -const usageCache = new Map< - number, - { - cachedUsagePct?: number; - resetAt?: string; - cachedUsageD7Pct?: number; - resetD7At?: string; +function parseJwtAuth(idToken: string): { + planType: string; + expiresAt: string | null; +} { + try { + const parts = idToken.split('.'); + if (parts.length < 2) return { planType: '?', expiresAt: null }; + const payload = JSON.parse( + Buffer.from(parts[1], 'base64url').toString('utf-8'), + ); + const auth = payload?.['https://api.openai.com/auth'] || {}; + return { + planType: auth.chatgpt_plan_type || '?', + expiresAt: auth.chatgpt_subscription_active_until || null, + }; + } catch { + return { planType: '?', expiresAt: null }; } ->(); - -function listCodexAccountDirs(): string[] { - if (!fs.existsSync(CODEX_ACCOUNTS_DIR)) return []; - - return fs - .readdirSync(CODEX_ACCOUNTS_DIR, { withFileTypes: true }) - .filter((entry) => entry.isDirectory() && /^\d+$/.test(entry.name)) - .sort((a, b) => Number(a.name) - Number(b.name)) - .map((entry) => path.join(CODEX_ACCOUNTS_DIR, entry.name)); } -function readTextIfExists(filePath: string): string | null { +const accounts: CodexAccount[] = []; +let currentIndex = 0; +let initialized = false; + +const ACCOUNTS_DIR = path.join(os.homedir(), '.codex-accounts'); + +export function initCodexTokenRotation(): void { + if (initialized) return; + initialized = true; + + if (!fs.existsSync(ACCOUNTS_DIR)) { + logger.info( + { dir: ACCOUNTS_DIR }, + 'Codex accounts dir not found, skipping', + ); + return; + } + + const dirs = fs + .readdirSync(ACCOUNTS_DIR) + .filter((d) => /^\d+$/.test(d)) + .sort((a, b) => parseInt(a) - parseInt(b)); + + for (const dir of dirs) { + const authPath = path.join(ACCOUNTS_DIR, dir, 'auth.json'); + if (!fs.existsSync(authPath)) continue; + + try { + const data = JSON.parse(fs.readFileSync(authPath, 'utf-8')); + const accountId = data?.tokens?.account_id || `account-${dir}`; + const jwt = parseJwtAuth(data?.tokens?.id_token || ''); + const planType = jwt.planType; + accounts.push({ + index: accounts.length, + authPath, + accountId, + planType, + subscriptionUntil: jwt.expiresAt, + rateLimitedUntil: null, + }); + } catch { + logger.warn({ authPath }, 'Failed to parse codex account auth.json'); + } + } + + if (accounts.length > 1) loadCodexState(); + logger.info( + { count: accounts.length, dir: ACCOUNTS_DIR, activeIndex: currentIndex }, + `Codex token rotation: ${accounts.length} account(s) found`, + ); +} + +function saveCodexState(): void { try { - if (!fs.existsSync(filePath)) return null; - return fs.readFileSync(filePath, 'utf8'); + const state = { + currentIndex, + rateLimits: accounts.map((a) => a.rateLimitedUntil), + usagePcts: accounts.map((a) => a.lastUsagePct ?? null), + usageD7Pcts: accounts.map((a) => a.lastUsageD7Pct ?? null), + resetAts: accounts.map((a) => a.resetAt ?? null), + resetD7Ats: accounts.map((a) => a.resetD7At ?? null), + }; + fs.writeFileSync(STATE_FILE, JSON.stringify(state)); + } catch { + /* best effort */ + } +} + +function loadCodexState(): void { + try { + if (!fs.existsSync(STATE_FILE)) return; + const state = JSON.parse(fs.readFileSync(STATE_FILE, 'utf-8')); + const now = Date.now(); + if ( + typeof state.currentIndex === 'number' && + state.currentIndex < accounts.length + ) { + currentIndex = state.currentIndex; + } + if (Array.isArray(state.rateLimits)) { + for ( + let i = 0; + i < Math.min(state.rateLimits.length, accounts.length); + i++ + ) { + const until = state.rateLimits[i]; + if (typeof until === 'number' && until > now) { + accounts[i].rateLimitedUntil = until; + } + } + } + if (Array.isArray(state.usagePcts)) { + for ( + let i = 0; + i < Math.min(state.usagePcts.length, accounts.length); + i++ + ) { + if (typeof state.usagePcts[i] === 'number') + accounts[i].lastUsagePct = state.usagePcts[i]; + } + } + if (Array.isArray(state.usageD7Pcts)) { + for ( + let i = 0; + i < Math.min(state.usageD7Pcts.length, accounts.length); + i++ + ) { + if (typeof state.usageD7Pcts[i] === 'number') + accounts[i].lastUsageD7Pct = state.usageD7Pcts[i]; + } + } + if (Array.isArray(state.resetAts)) { + for ( + let i = 0; + i < Math.min(state.resetAts.length, accounts.length); + i++ + ) { + if (state.resetAts[i]) accounts[i].resetAt = state.resetAts[i]; + } + } + if (Array.isArray(state.resetD7Ats)) { + for ( + let i = 0; + i < Math.min(state.resetD7Ats.length, accounts.length); + i++ + ) { + if (state.resetD7Ats[i]) accounts[i].resetD7At = state.resetD7Ats[i]; + } + } + logger.info( + { currentIndex, accountCount: accounts.length }, + 'Codex rotation state restored', + ); + } catch { + /* start fresh */ + } +} + +const BUFFER_MS = 3 * 60_000; // 3 min buffer after reset time +const DEFAULT_COOLDOWN_MS = 3_600_000; // 1 hour fallback + +/** + * Parse "try again at Mar 26th, 2026 9:00 AM" from error message. + * Returns timestamp in ms, or null if not found. + */ +function parseRetryAfterFromError(error?: string): number | null { + if (!error) return null; + const match = error.match( + /try again at\s+(\w+ \d{1,2}(?:st|nd|rd|th)?,?\s+\d{4}\s+\d{1,2}:\d{2}\s*(?:AM|PM)?)/i, + ); + if (!match) return null; + try { + // Remove ordinal suffixes (1st, 2nd, 3rd, 4th) + const cleaned = match[1].replace(/(\d+)(?:st|nd|rd|th)/i, '$1'); + const ts = new Date(cleaned).getTime(); + if (Number.isNaN(ts)) return null; + return ts; } catch { return null; } } -function inferPlanType(accountDir: string): string { - try { - const authPath = path.join(accountDir, 'auth.json'); - if (!fs.existsSync(authPath)) return 'API'; - const raw = JSON.parse(fs.readFileSync(authPath, 'utf8')) as { - auth_mode?: string; - }; - if (!raw.auth_mode) return 'API'; - return raw.auth_mode - .split(/[_-]/) - .filter(Boolean) - .map((part) => part[0].toUpperCase() + part.slice(1)) - .join(' '); - } catch { - return 'API'; - } +function computeCooldownUntil(error?: string): number { + const retryAt = parseRetryAfterFromError(error); + if (retryAt) return retryAt + BUFFER_MS; + return Date.now() + DEFAULT_COOLDOWN_MS; } -function resolveActiveIndex(accountDirs: string[]): number { - if (accountDirs.length <= 1) return 0; - if (activeIndexCache !== null && activeIndexCache < accountDirs.length) { - return activeIndexCache; - } +/** Get the auth.json path for the current active account. */ +export function getActiveCodexAuthPath(): string | null { + if (accounts.length === 0) return null; + return accounts[currentIndex]?.authPath ?? null; +} - const hostAuth = readTextIfExists(path.join(HOST_CODEX_DIR, 'auth.json')); - if (hostAuth) { - const matchIndex = accountDirs.findIndex( - (accountDir) => - readTextIfExists(path.join(accountDir, 'auth.json')) === hostAuth, - ); - if (matchIndex >= 0) { - activeIndexCache = matchIndex; - return matchIndex; +/** + * Try to rotate to the next available Codex account. + * Returns true if a fresh account was found. + */ +export function rotateCodexToken(errorMessage?: string): boolean { + if (accounts.length <= 1) return false; + + const now = Date.now(); + const acct = accounts[currentIndex]; + acct.rateLimitedUntil = computeCooldownUntil(errorMessage); + acct.lastUsagePct = 100; + // Extract reset time string from error for display + const resetMatch = errorMessage?.match( + /try again at\s+(\w+ \d{1,2}(?:st|nd|rd|th)?,?\s+\d{4}\s+\d{1,2}:\d{2}\s*(?:AM|PM)?)/i, + ); + if (resetMatch) acct.resetAt = resetMatch[1]; + + for (let i = 1; i < accounts.length; i++) { + const idx = (currentIndex + i) % accounts.length; + const acct = accounts[idx]; + if (!acct.rateLimitedUntil || acct.rateLimitedUntil <= now) { + acct.rateLimitedUntil = null; + currentIndex = idx; + logger.info( + { + accountIndex: currentIndex, + totalAccounts: accounts.length, + accountId: acct.accountId, + }, + `Codex rotated to account #${currentIndex + 1}/${accounts.length}`, + ); + saveCodexState(); + return true; } } - activeIndexCache = 0; - return 0; + logger.warn('All Codex accounts are rate-limited'); + return false; } -function copyIfExists(src: string, dst: string): void { - if (!fs.existsSync(src)) return; - fs.mkdirSync(path.dirname(dst), { recursive: true }); - fs.copyFileSync(src, dst); -} - -function syncHostCodexAccount(accountDir: string): void { - copyIfExists( - path.join(accountDir, 'auth.json'), - path.join(HOST_CODEX_DIR, 'auth.json'), - ); -} - -export function getAllCodexAccounts(): CodexAccountState[] { - const accountDirs = listCodexAccountDirs(); - if (accountDirs.length === 0) { - return [ - { - index: 0, - planType: 'API', - isActive: true, - isRateLimited: false, - }, - ]; +/** + * Advance to the next healthy account (round-robin). + * Called after each successful request to spread load evenly + * and keep usage data fresh for all accounts. + */ +export function advanceCodexAccount(): void { + if (accounts.length <= 1) return; + const now = Date.now(); + for (let i = 1; i < accounts.length; i++) { + const idx = (currentIndex + i) % accounts.length; + const acct = accounts[idx]; + if (!acct.rateLimitedUntil || acct.rateLimitedUntil <= now) { + currentIndex = idx; + saveCodexState(); + return; + } } - - const activeIndex = resolveActiveIndex(accountDirs); - return accountDirs.map((accountDir, index) => ({ - index, - planType: inferPlanType(accountDir), - isActive: index === activeIndex, - isRateLimited: rateLimitedAccounts.has(index), - ...usageCache.get(index), - })); -} - -export function getCodexAccountCount(): number { - return getAllCodexAccounts().length; -} - -export function rotateCodexToken(_reason?: string): boolean { - const accountDirs = listCodexAccountDirs(); - if (accountDirs.length <= 1) return false; - - const activeIndex = resolveActiveIndex(accountDirs); - rateLimitedAccounts.add(activeIndex); - - const candidates = accountDirs.map((_, index) => index); - const nextIndex = - candidates.find( - (index) => index !== activeIndex && !rateLimitedAccounts.has(index), - ) ?? candidates.find((index) => index !== activeIndex); - - if (nextIndex === undefined) return false; - - syncHostCodexAccount(accountDirs[nextIndex]); - activeIndexCache = nextIndex; - return true; -} - -export function markCodexTokenHealthy(): void { - const accountDirs = listCodexAccountDirs(); - const activeIndex = resolveActiveIndex(accountDirs); - rateLimitedAccounts.delete(activeIndex); + // All others rate-limited, stay on current } +/** + * Update cached usage info for a specific account (or current if index omitted). + */ export function updateCodexAccountUsage( - pct: number, - resetAt: string | undefined, - index: number, + usagePct: number, + resetAt?: string, + accountIndex?: number, d7Pct?: number, resetD7At?: string, ): void { - usageCache.set(index, { - cachedUsagePct: pct, - resetAt, - cachedUsageD7Pct: d7Pct, - resetD7At, - }); + if (accounts.length === 0) return; + const idx = accountIndex ?? currentIndex; + const acct = accounts[idx]; + if (acct) { + acct.lastUsagePct = usagePct; + if (d7Pct != null) acct.lastUsageD7Pct = d7Pct; + if (resetAt) acct.resetAt = resetAt; + if (resetD7At) acct.resetD7At = resetD7At; + saveCodexState(); + } +} + +export function markCodexTokenHealthy(): void { + if (accounts.length === 0) return; + const acct = accounts[currentIndex]; + if (acct?.rateLimitedUntil) { + acct.rateLimitedUntil = null; + saveCodexState(); + } +} + +export function getCodexAccountCount(): number { + return accounts.length; +} + +export function getAllCodexAccounts(): { + index: number; + accountId: string; + planType: string; + isActive: boolean; + isRateLimited: boolean; + cachedUsagePct?: number; + cachedUsageD7Pct?: number; + resetAt?: string; + resetD7At?: string; +}[] { + const now = Date.now(); + return accounts.map((a, i) => ({ + index: i, + accountId: a.accountId, + planType: a.planType, + isActive: i === currentIndex, + isRateLimited: Boolean(a.rateLimitedUntil && a.rateLimitedUntil > now), + cachedUsagePct: a.lastUsagePct, + cachedUsageD7Pct: a.lastUsageD7Pct, + resetAt: a.resetAt, + resetD7At: a.resetD7At, + })); } diff --git a/src/index.ts b/src/index.ts index 8a33e6e..45db441 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 { initCodexTokenRotation } from './codex-token-rotation.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'; +// Token rotation is initialized lazily on first use or at startup below + export async function sendFormattedChannelMessage( channels: Channel[], jid: string, @@ -293,6 +297,8 @@ async function main(): Promise { const processStartedAtMs = Date.now(); initDatabase(); logger.info('Database initialized'); + initTokenRotation(); + initCodexTokenRotation(); loadState(); // Graceful shutdown handlers diff --git a/src/message-agent-executor.test.ts b/src/message-agent-executor.test.ts index 0fc77c6..f66d5b8 100644 --- a/src/message-agent-executor.test.ts +++ b/src/message-agent-executor.test.ts @@ -20,6 +20,7 @@ vi.mock('./db.js', () => ({ vi.mock('./logger.js', () => ({ logger: { + debug: vi.fn(), info: vi.fn(), warn: vi.fn(), error: vi.fn(), @@ -27,15 +28,26 @@ vi.mock('./logger.js', () => ({ })); vi.mock('./provider-fallback.js', () => ({ - detectFallbackTrigger: vi.fn(() => ({ - shouldFallback: false, - reason: '', + detectFallbackTrigger: vi.fn((error?: string | null) => { + const lower = (error || '').toLowerCase(); + if ( + lower.includes('429') || + lower.includes('rate limit') || + lower.includes('hit your limit') + ) { + return { shouldFallback: true, reason: '429' }; + } + return { shouldFallback: false, reason: '' }; + }), + getActiveProvider: vi.fn(async () => 'claude'), + getFallbackEnvOverrides: vi.fn(() => ({ + ANTHROPIC_BASE_URL: 'https://api.kimi.com/coding/', + ANTHROPIC_AUTH_TOKEN: 'test-kimi-key', + ANTHROPIC_MODEL: 'kimi-k2.5', })), - getActiveProvider: vi.fn(() => 'claude'), - getFallbackEnvOverrides: vi.fn(() => ({})), - getFallbackProviderName: vi.fn(() => 'fallback'), + getFallbackProviderName: vi.fn(() => 'kimi'), hasGroupProviderOverride: vi.fn(() => false), - isFallbackEnabled: vi.fn(() => false), + isFallbackEnabled: vi.fn(() => true), markPrimaryCooldown: vi.fn(), })); @@ -43,96 +55,220 @@ vi.mock('./session-recovery.js', () => ({ shouldResetSessionOnAgentFailure: vi.fn(() => false), })); +vi.mock('./token-rotation.js', () => ({ + rotateToken: vi.fn(() => false), + getTokenCount: vi.fn(() => 1), + markTokenHealthy: vi.fn(), +})); + +vi.mock('./codex-token-rotation.js', () => ({ + rotateCodexToken: vi.fn(() => false), + getCodexAccountCount: vi.fn(() => 1), + markCodexTokenHealthy: vi.fn(), +})); + vi.mock('./memento-client.js', () => ({ buildRoomMemoryBriefing: vi.fn(), })); -import { runAgentProcess } from './agent-runner.js'; +import * as agentRunner from './agent-runner.js'; import { buildRoomMemoryBriefing } from './memento-client.js'; import { runAgentForGroup } from './message-agent-executor.js'; +import * as providerFallback from './provider-fallback.js'; +import * as tokenRotation from './token-rotation.js'; import type { RegisteredGroup } from './types.js'; -const testGroup: RegisteredGroup = { - name: 'Test Group', - folder: 'test-group', - trigger: '@Andy', - added_at: new Date().toISOString(), -}; +function makeGroup(): RegisteredGroup { + return { + name: 'Test Group', + folder: 'test-claude', + trigger: '@Andy', + added_at: new Date().toISOString(), + requiresTrigger: false, + agentType: 'claude-code', + }; +} -const deps = { - assistantName: 'Andy', - queue: { - registerProcess: vi.fn(), - }, - getRegisteredGroups: vi.fn(() => ({})), - getSessions: vi.fn(() => ({})), - persistSession: vi.fn(), - clearSession: vi.fn(), -}; +function makeDeps() { + return { + assistantName: 'Andy', + queue: { + registerProcess: vi.fn(), + }, + getRegisteredGroups: () => ({}), + getSessions: () => ({}), + persistSession: vi.fn(), + clearSession: vi.fn(), + }; +} -describe('runAgentForGroup', () => { +describe('runAgentForGroup Claude rotation', () => { beforeEach(() => { vi.clearAllMocks(); - vi.mocked(runAgentProcess).mockResolvedValue({ - status: 'success', - result: 'ok', - newSessionId: 'session-123', - }); - vi.mocked(buildRoomMemoryBriefing).mockResolvedValue( - '## Shared Room Memory\n- remembered context', - ); + vi.mocked(providerFallback.getActiveProvider).mockResolvedValue('claude'); + vi.mocked(providerFallback.isFallbackEnabled).mockReturnValue(true); + vi.mocked(providerFallback.hasGroupProviderOverride).mockReturnValue(false); + vi.mocked(tokenRotation.getTokenCount).mockReturnValue(1); + vi.mocked(tokenRotation.rotateToken).mockReturnValue(false); }); - it('injects a room memory briefing when starting a fresh session', async () => { - deps.getSessions.mockReturnValue({}); + it('rotates to another Claude account before falling back to Kimi', async () => { + const outputs: string[] = []; - const result = await runAgentForGroup(deps, { - group: testGroup, + vi.mocked(tokenRotation.getTokenCount).mockReturnValue(2); + vi.mocked(tokenRotation.rotateToken).mockReturnValueOnce(true); + + vi.mocked(agentRunner.runAgentProcess) + .mockImplementationOnce(async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'final', + result: 'You’re out of extra usage · resets 4am (Asia/Seoul)', + }); + return { + status: 'success', + result: null, + }; + }) + .mockImplementationOnce(async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'final', + result: '회전된 Claude 응답입니다.', + }); + return { + status: 'success', + result: null, + }; + }); + + const result = await runAgentForGroup(makeDeps(), { + group: makeGroup(), prompt: 'hello', chatJid: 'group@test', - runId: 'run-1', + runId: 'run-rotate-claude', + onOutput: async (output) => { + if (typeof output.result === 'string') outputs.push(output.result); + }, }); expect(result).toBe('success'); - expect(buildRoomMemoryBriefing).toHaveBeenCalledWith({ - groupFolder: 'test-group', - groupName: 'Test Group', - }); - expect(runAgentProcess).toHaveBeenCalledWith( - testGroup, - expect.objectContaining({ - prompt: 'hello', - sessionId: undefined, - memoryBriefing: '## Shared Room Memory\n- remembered context', - }), - expect.any(Function), - undefined, - undefined, - ); + expect(agentRunner.runAgentProcess).toHaveBeenCalledTimes(2); + expect(tokenRotation.rotateToken).toHaveBeenCalledTimes(1); + expect(tokenRotation.markTokenHealthy).toHaveBeenCalledTimes(1); + expect(providerFallback.markPrimaryCooldown).not.toHaveBeenCalled(); + expect(outputs).toEqual(['회전된 Claude 응답입니다.']); }); - it('skips the room memory briefing for existing sessions', async () => { - deps.getSessions.mockReturnValue({ 'test-group': 'session-existing' }); + it('falls back to Kimi only after all Claude accounts are exhausted', async () => { + const outputs: string[] = []; - const result = await runAgentForGroup(deps, { - group: testGroup, - prompt: 'hello again', + vi.mocked(tokenRotation.getTokenCount).mockReturnValue(2); + vi.mocked(tokenRotation.rotateToken) + .mockReturnValueOnce(true) + .mockReturnValueOnce(false); + + vi.mocked(agentRunner.runAgentProcess) + .mockImplementationOnce(async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'final', + result: 'You’re out of extra usage · resets 4am (Asia/Seoul)', + }); + return { + status: 'success', + result: null, + }; + }) + .mockImplementationOnce(async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'final', + result: "You're out of extra usage · resets 4am (Asia/Seoul)", + }); + return { + status: 'success', + result: null, + }; + }) + .mockImplementationOnce(async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'final', + result: 'Kimi 폴백 응답입니다.', + }); + return { + status: 'success', + result: null, + }; + }); + + const result = await runAgentForGroup(makeDeps(), { + group: makeGroup(), + prompt: 'hello', chatJid: 'group@test', - runId: 'run-2', + runId: 'run-fallback-after-rotation', + onOutput: async (output) => { + if (typeof output.result === 'string') outputs.push(output.result); + }, }); expect(result).toBe('success'); - expect(buildRoomMemoryBriefing).not.toHaveBeenCalled(); - expect(runAgentProcess).toHaveBeenCalledWith( - testGroup, - expect.objectContaining({ - prompt: 'hello again', - sessionId: 'session-existing', - memoryBriefing: undefined, - }), - expect.any(Function), - undefined, + expect(agentRunner.runAgentProcess).toHaveBeenCalledTimes(3); + expect(tokenRotation.rotateToken).toHaveBeenCalledTimes(2); + expect(providerFallback.markPrimaryCooldown).toHaveBeenCalledWith( + 'usage-exhausted', undefined, ); + expect(agentRunner.runAgentProcess).toHaveBeenNthCalledWith( + 3, + expect.anything(), + expect.anything(), + expect.any(Function), + expect.any(Function), + expect.objectContaining({ + ANTHROPIC_BASE_URL: 'https://api.kimi.com/coding/', + ANTHROPIC_MODEL: 'kimi-k2.5', + }), + ); + expect(outputs).toEqual(['Kimi 폴백 응답입니다.']); + }); + + it('does not mistake a normal response quoting the banner text for a usage error', async () => { + const outputs: string[] = []; + + vi.mocked(tokenRotation.getTokenCount).mockReturnValue(2); + + vi.mocked(agentRunner.runAgentProcess).mockImplementationOnce( + async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'final', + result: + "상태 문구 예시: You're out of extra usage · resets 4am (Asia/Seoul) 라는 배너가 뜰 수 있습니다.", + }); + return { + status: 'success', + result: null, + }; + }, + ); + + const result = await runAgentForGroup(makeDeps(), { + group: makeGroup(), + prompt: 'hello', + chatJid: 'group@test', + runId: 'run-normal-quoted-banner', + onOutput: async (output) => { + if (typeof output.result === 'string') outputs.push(output.result); + }, + }); + + expect(result).toBe('success'); + expect(tokenRotation.rotateToken).not.toHaveBeenCalled(); + expect(providerFallback.markPrimaryCooldown).not.toHaveBeenCalled(); + expect(outputs).toEqual([ + "상태 문구 예시: You're out of extra usage · resets 4am (Asia/Seoul) 라는 배너가 뜰 수 있습니다.", + ]); }); }); diff --git a/src/message-agent-executor.ts b/src/message-agent-executor.ts index 914bb57..b1eea0b 100644 --- a/src/message-agent-executor.ts +++ b/src/message-agent-executor.ts @@ -19,6 +19,7 @@ import { getFallbackProviderName, hasGroupProviderOverride, isFallbackEnabled, + isUsageExhausted, markPrimaryCooldown, } from './provider-fallback.js'; import { shouldResetSessionOnAgentFailure } from './session-recovery.js'; @@ -44,11 +45,22 @@ export interface MessageAgentExecutorDeps { } function isClaudeUsageExhaustedMessage(text: string): boolean { - const normalized = text.trim(); - return ( - /^you(?:'re| are) out of extra usage\b/i.test(normalized) || - /^you(?:'ve| have) hit your limit\b/i.test(normalized) - ); + const normalized = text + .trim() + .toLowerCase() + .replace(/[''`]/g, "'") + .replace(/\s+/g, ' ') + .replace(/^error:\s*/i, ''); + const looksLikeBanner = + normalized.startsWith("you're out of extra usage") || + normalized.startsWith('you are out of extra usage') || + normalized.startsWith("you've hit your limit") || + normalized.startsWith('you have hit your limit'); + const hasResetHint = + normalized.includes('resets ') || + normalized.includes('reset at ') || + normalized.includes('try again'); + return looksLikeBanner && hasResetHint && normalized.length <= 160; } export async function runAgentForGroup( @@ -158,10 +170,20 @@ export async function runAgentForGroup( provider === 'claude' && output.status === 'success' && !sawOutput && - !streamedTriggerReason && typeof output.result === 'string' && isClaudeUsageExhaustedMessage(output.result) ) { + if (!streamedTriggerReason) { + logger.warn( + { + chatJid, + group: group.name, + runId, + resultPreview: output.result.slice(0, 120), + }, + 'Detected Claude usage exhaustion banner in successful output', + ); + } streamedTriggerReason = { reason: 'usage-exhausted', }; @@ -321,7 +343,166 @@ export async function runAgentForGroup( return 'success'; }; - const provider = canFallback ? getActiveProvider() : 'claude'; + const shouldRotateClaudeToken = (reason: string): boolean => + reason === '429' || reason === 'usage-exhausted'; + + const retryClaudeWithRotation = async ( + initialTrigger: { + reason: string; + retryAfterMs?: number; + }, + rotationMessage?: string, + ): Promise<'success' | 'error'> => { + let trigger = initialTrigger; + let lastRotationMessage = rotationMessage; + + while ( + shouldRotateClaudeToken(trigger.reason) && + getTokenCount() > 1 && + rotateToken(lastRotationMessage, { ignoreRateLimits: true }) + ) { + logger.info( + { chatJid, group: group.name, runId, reason: trigger.reason }, + 'Claude rate-limited, retrying with rotated account', + ); + + const retryAttempt = await runAttempt('claude'); + + if (retryAttempt.error) { + if (!retryAttempt.sawOutput) { + const errMsg = + retryAttempt.error instanceof Error + ? retryAttempt.error.message + : String(retryAttempt.error); + const retryTrigger = retryAttempt.streamedTriggerReason + ? { + shouldFallback: true, + reason: retryAttempt.streamedTriggerReason.reason, + retryAfterMs: retryAttempt.streamedTriggerReason.retryAfterMs, + } + : detectFallbackTrigger(errMsg); + if (retryTrigger.shouldFallback) { + trigger = { + reason: retryTrigger.reason, + retryAfterMs: retryTrigger.retryAfterMs, + }; + lastRotationMessage = errMsg; + continue; + } + } + + logger.error( + { + chatJid, + group: group.name, + groupFolder: group.folder, + runId, + provider: 'claude', + err: retryAttempt.error, + }, + 'Rotated Claude account also threw', + ); + return 'error'; + } + + const retryOutput = retryAttempt.output; + if (!retryOutput) { + logger.error( + { + chatJid, + group: group.name, + groupFolder: group.folder, + runId, + provider: 'claude', + }, + 'Rotated Claude account produced no output object', + ); + return 'error'; + } + + if ( + !retryAttempt.sawOutput && + retryAttempt.streamedTriggerReason && + retryOutput.status !== 'error' + ) { + trigger = { + reason: retryAttempt.streamedTriggerReason.reason, + retryAfterMs: retryAttempt.streamedTriggerReason.retryAfterMs, + }; + lastRotationMessage = + typeof retryOutput.result === 'string' + ? retryOutput.result + : undefined; + continue; + } + + if ( + !retryAttempt.sawOutput && + retryAttempt.sawSuccessNullResultWithoutOutput + ) { + return runFallbackAttempt('success-null-result'); + } + + if (retryOutput.status === 'error') { + if (!retryAttempt.sawOutput) { + const retryTrigger = retryAttempt.streamedTriggerReason + ? { + shouldFallback: true, + reason: retryAttempt.streamedTriggerReason.reason, + retryAfterMs: retryAttempt.streamedTriggerReason.retryAfterMs, + } + : detectFallbackTrigger(retryOutput.error); + if (retryTrigger.shouldFallback) { + trigger = { + reason: retryTrigger.reason, + retryAfterMs: retryTrigger.retryAfterMs, + }; + lastRotationMessage = retryOutput.error ?? undefined; + continue; + } + } + + logger.error( + { + group: group.name, + chatJid, + runId, + provider: 'claude', + error: retryOutput.error, + }, + 'Rotated Claude account failed', + ); + return 'error'; + } + + markTokenHealthy(); + return 'success'; + } + + // Usage exhausted: don't fall back to Kimi — log only, no response + if (trigger.reason === 'usage-exhausted') { + markPrimaryCooldown(trigger.reason, trigger.retryAfterMs); + logger.info( + { chatJid, group: group.name, runId }, + 'All Claude tokens usage-exhausted, silently skipping (no Kimi fallback)', + ); + return 'error'; + } + + return runFallbackAttempt(trigger.reason, trigger.retryAfterMs); + }; + + const provider = canFallback ? await getActiveProvider() : 'claude'; + + // Already in usage-exhausted cooldown — log only, no response + if (provider !== 'claude' && isUsageExhausted()) { + logger.info( + { chatJid, group: group.name, runId, provider }, + 'Claude usage exhausted (cooldown active), silently skipping', + ); + return 'error'; + } + const primaryAttempt = await runAttempt(provider); if (primaryAttempt.error) { @@ -338,19 +519,13 @@ export async function runAgentForGroup( } : detectFallbackTrigger(errMsg); if (trigger.shouldFallback) { - // Try rotating token before falling back to another provider - if (getTokenCount() > 1 && rotateToken(errMsg)) { - 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); + return retryClaudeWithRotation( + { + reason: trigger.reason, + retryAfterMs: trigger.retryAfterMs, + }, + errMsg, + ); } } @@ -390,10 +565,10 @@ export async function runAgentForGroup( primaryAttempt.streamedTriggerReason && output.status !== 'error' ) { - return runFallbackAttempt( - primaryAttempt.streamedTriggerReason.reason, - primaryAttempt.streamedTriggerReason.retryAfterMs, - ); + return retryClaudeWithRotation({ + reason: primaryAttempt.streamedTriggerReason.reason, + retryAfterMs: primaryAttempt.streamedTriggerReason.retryAfterMs, + }); } if ( @@ -426,18 +601,13 @@ export async function runAgentForGroup( } : detectFallbackTrigger(output.error); if (trigger.shouldFallback) { - if (getTokenCount() > 1 && rotateToken(output.error ?? undefined)) { - 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); + return retryClaudeWithRotation( + { + reason: trigger.reason, + retryAfterMs: trigger.retryAfterMs, + }, + output.error ?? undefined, + ); } } diff --git a/src/message-runtime.test.ts b/src/message-runtime.test.ts index 3916220..cadd29f 100644 --- a/src/message-runtime.test.ts +++ b/src/message-runtime.test.ts @@ -122,7 +122,7 @@ vi.mock('./logger.js', () => ({ vi.mock('./provider-fallback.js', () => ({ detectFallbackTrigger: vi.fn(() => ({ shouldFallback: false, reason: '' })), - getActiveProvider: vi.fn(() => 'claude'), + getActiveProvider: vi.fn(async () => 'claude'), getFallbackEnvOverrides: vi.fn(() => ({ ANTHROPIC_BASE_URL: 'https://api.kimi.com/coding/', ANTHROPIC_AUTH_TOKEN: 'test-kimi-key', @@ -181,7 +181,7 @@ function makeChannel(chatJid: string): Channel { describe('createMessageRuntime', () => { beforeEach(() => { vi.clearAllMocks(); - vi.mocked(providerFallback.getActiveProvider).mockReturnValue('claude'); + vi.mocked(providerFallback.getActiveProvider).mockResolvedValue('claude'); vi.mocked(providerFallback.getFallbackProviderName).mockReturnValue('kimi'); vi.mocked(providerFallback.getFallbackEnvOverrides).mockReturnValue({ ANTHROPIC_BASE_URL: 'https://api.kimi.com/coding/', @@ -2037,6 +2037,91 @@ describe('createMessageRuntime', () => { ); }); + it('suppresses duplicate streamed usage banners before retrying the fallback provider', async () => { + const chatJid = 'group@test'; + const group = makeGroup('claude-code'); + const channel = makeChannel(chatJid); + + vi.mocked(db.getMessagesSince).mockReturnValue([ + { + id: 'msg-1', + chat_jid: chatJid, + sender: 'user@test', + sender_name: 'User', + content: 'hello', + timestamp: '2026-03-24T00:00:00.000Z', + }, + ]); + + vi.mocked(agentRunner.runAgentProcess) + .mockImplementationOnce(async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'intermediate', + result: "You're out of extra usage · resets 4am (Asia/Seoul)", + }); + await onOutput?.({ + status: 'success', + phase: 'final', + result: "You're out of extra usage · resets 4am (Asia/Seoul)", + }); + return { + status: 'success', + result: null, + }; + }) + .mockImplementationOnce(async (_group, _input, _onProcess, onOutput) => { + await onOutput?.({ + status: 'success', + phase: 'final', + result: 'duplicate banner fallback 응답입니다.', + }); + return { + status: 'success', + result: null, + }; + }); + + const runtime = createMessageRuntime({ + assistantName: 'Andy', + idleTimeout: 1_000, + pollInterval: 1_000, + timezone: 'UTC', + triggerPattern: /^@Andy\b/i, + channels: [channel], + queue: { + registerProcess: vi.fn(), + closeStdin: vi.fn(), + notifyIdle: vi.fn(), + } as any, + getRegisteredGroups: () => ({ [chatJid]: group }), + getSessions: () => ({}), + getLastTimestamp: () => '', + setLastTimestamp: vi.fn(), + getLastAgentTimestamps: () => ({}), + saveState: vi.fn(), + persistSession: vi.fn(), + clearSession: vi.fn(), + }); + + const result = await runtime.processGroupMessages(chatJid, { + runId: 'run-fallback-usage-exhausted-duplicate', + reason: 'messages', + }); + + expect(result).toBe(true); + expect(agentRunner.runAgentProcess).toHaveBeenCalledTimes(2); + expect(providerFallback.markPrimaryCooldown).toHaveBeenCalledWith( + 'usage-exhausted', + undefined, + ); + expect(channel.sendMessage).toHaveBeenCalledTimes(1); + expect(channel.sendMessage).toHaveBeenCalledWith( + chatJid, + 'duplicate banner fallback 응답입니다.', + ); + }); + it('retries with the fallback provider when Claude ends with success-null-result before any output', async () => { const chatJid = 'group@test'; const group = makeGroup('claude-code'); diff --git a/src/message-turn-controller.ts b/src/message-turn-controller.ts index 348f21a..7f72ae0 100644 --- a/src/message-turn-controller.ts +++ b/src/message-turn-controller.ts @@ -7,6 +7,11 @@ import { type Channel, type RegisteredGroup } from './types.js'; export type VisiblePhase = 'silent' | 'progress' | 'final'; +interface SubagentTrack { + label: string; + activities: string[]; +} + interface MessageTurnControllerOptions { chatJid: string; group: RegisteredGroup; @@ -36,6 +41,8 @@ export class MessageTurnController { private progressTicker: ReturnType | null = null; private progressEditFailCount = 0; private latestProgressTextForFinal: string | null = null; + private subagents = new Map(); + private lastIntermediateText: string | null = null; private poisonedSessionDetected = false; private closeRequested = false; @@ -106,9 +113,22 @@ export class MessageTurnController { } if (result.phase === 'intermediate') { - // Send as standalone message without touching progress state if (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(); @@ -117,20 +137,29 @@ export class MessageTurnController { } if (result.phase === 'tool-activity') { - // Ensure a progress message exists for tool activity sub-lines - if (!this.progressMessageId && !this.progressCreating) { - this.progressCreating = true; - const heading = this.pendingProgressText || '작업 중...'; - this.sendProgressMessage(heading).then(() => { - this.progressCreating = false; - this.ensureProgressTicker(); - // Replay any queued tool activities now that the message exists - if (this.toolActivities.length > 0 && this.progressMessageId) { - void this.syncTrackedProgressMessage(); + if (result.agentId) { + // Subagent tool activity + let track = this.subagents.get(result.agentId); + if (!track) { + track = { label: '작업 중...', activities: [] }; + this.subagents.set(result.agentId, track); + } + if (text) { + const MAX = 2; + track.activities.push(text); + if (track.activities.length > MAX) { + track.activities = track.activities.slice(-MAX); } - }); - this.pendingProgressText = null; + } + this.ensureProgressMessageExists(); + this.ensureProgressTicker(); + if (!this.poisonedSessionDetected) { + this.resetIdleTimer(); + } + return; } + // Main agent tool activity + this.ensureProgressMessageExists(); if (text) { this.addToolActivity(text); } @@ -141,6 +170,43 @@ export class MessageTurnController { } if (result.phase === 'progress') { + if (result.agentId) { + if (result.agentDone) { + const done = this.subagents.get(result.agentId); + if (done) { + done.label = done.label.replace('🔄', '✅'); + done.activities = []; + } + } else { + const label = + text || + (result.agentLabel ? `🔄 ${result.agentLabel}` : '작업 중...'); + const existing = this.subagents.get(result.agentId); + if (existing) { + existing.label = label; + existing.activities = []; + } else { + this.subagents.set(result.agentId, { label, activities: [] }); + } + } + if (!this.latestProgressText) { + this.latestProgressText = '작업 중...'; + this.latestProgressTextForFinal = '작업 중...'; + } + this.ensureProgressMessageExists(); + this.ensureProgressTicker(); + if (this.progressMessageId) { + void this.syncTrackedProgressMessage(); + } + if (!this.poisonedSessionDetected) { + this.resetIdleTimer(); + } + if (result.status === 'error') { + this.hadError = true; + } + return; + } + // Main agent progress if (text) { if (this.progressMessageId) { // Progress message already visible — update heading directly @@ -164,9 +230,18 @@ export class MessageTurnController { // Final arrived — flush any buffered progress that isn't the same text, // then discard the pending buffer so it never shows up. if (text) { - await this.flushPendingProgress(text); - await this.finalizeProgressMessage(); - await this.deliverFinalText(text); + if (this.lastIntermediateText && text === this.lastIntermediateText) { + // Already sent as intermediate — skip duplicate, just finalize + this.lastIntermediateText = null; + await this.finalizeProgressMessage(); + this.visiblePhase = 'final'; + this.latestProgressTextForFinal = null; + } else { + this.lastIntermediateText = null; + await this.flushPendingProgress(text); + await this.finalizeProgressMessage(); + await this.deliverFinalText(text); + } } else if (raw) { logger.info( { @@ -262,27 +337,47 @@ export class MessageTurnController { if (minutes > 0) elapsedParts.push(`${minutes}분`); elapsedParts.push(`${seconds}초`); - const activityLines = - this.toolActivities.length > 0 - ? '\n' + - this.toolActivities - .map((a, i) => { - const isLast = i === this.toolActivities.length - 1; - const connector = isLast ? '└' : '├'; - const isSummary = a.startsWith('📋'); - return isSummary ? `${connector} ${a}` : `${connector} ${a}`; - }) - .join('\n') - : ''; const suffix = `\n\n${elapsedParts.join(' ')}`; - const maxText = - 2000 - - TASK_STATUS_MESSAGE_PREFIX.length - - activityLines.length - - suffix.length; + let body: string; + + if (this.subagents.size > 1) { + // Compact: one line per subagent with latest activity + const lines: string[] = []; + for (const [, track] of this.subagents) { + const latest = track.activities[track.activities.length - 1]; + lines.push(latest ? `${track.label} · ${latest}` : track.label); + } + body = lines.join('\n'); + } else if (this.subagents.size === 1) { + // Single subagent: detailed view with activity sub-lines + const [, track] = this.subagents.entries().next().value!; + const lines: string[] = [track.label]; + for (let i = 0; i < track.activities.length; i++) { + const isLast = i === track.activities.length - 1; + lines.push(`${isLast ? '└' : '├'} ${track.activities[i]}`); + } + body = lines.join('\n'); + } else { + // Single agent rendering + const activityLines = + this.toolActivities.length > 0 + ? '\n' + + this.toolActivities + .map((a, i) => { + const isLast = i === this.toolActivities.length - 1; + const connector = isLast ? '└' : '├'; + const isSummary = a.startsWith('📋'); + return isSummary ? `${connector} ${a}` : `${connector} ${a}`; + }) + .join('\n') + : ''; + body = text + activityLines; + } + + const maxBody = 2000 - TASK_STATUS_MESSAGE_PREFIX.length - suffix.length; const truncated = - text.length > maxText ? text.slice(0, maxText - 1) + '…' : text; - return `${TASK_STATUS_MESSAGE_PREFIX}${truncated}${activityLines}${suffix}`; + body.length > maxBody ? body.slice(0, maxBody - 1) + '…' : body; + return `${TASK_STATUS_MESSAGE_PREFIX}${truncated}${suffix}`; } private clearProgressTicker(): void { @@ -296,6 +391,7 @@ export class MessageTurnController { this.pendingProgressText = null; this.progressCreating = false; this.toolActivities = []; + this.subagents.clear(); this.latestProgressText = null; this.previousProgressText = null; this.latestProgressRendered = null; @@ -304,6 +400,32 @@ export class MessageTurnController { this.progressEditFailCount = 0; } + /** + * Ensure a progress message exists in Discord. + * Creates one if needed, using pending or default text. + */ + private ensureProgressMessageExists(): void { + if (this.progressMessageId || this.progressCreating) return; + this.progressCreating = true; + const heading = + this.pendingProgressText || this.latestProgressText || '작업 중...'; + if (!this.latestProgressText) { + this.latestProgressText = heading; + this.latestProgressTextForFinal = heading; + } + void this.sendProgressMessage(heading).then(() => { + this.progressCreating = false; + this.ensureProgressTicker(); + if ( + (this.toolActivities.length > 0 || this.subagents.size > 0) && + this.progressMessageId + ) { + void this.syncTrackedProgressMessage(); + } + }); + this.pendingProgressText = null; + } + /** * Buffer a progress update. The previous pending text gets flushed * immediately, and the new text waits until the next event arrives. diff --git a/src/provider-fallback.test.ts b/src/provider-fallback.test.ts new file mode 100644 index 0000000..54a9cc9 --- /dev/null +++ b/src/provider-fallback.test.ts @@ -0,0 +1,103 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +vi.mock('./claude-usage.js', () => ({ + fetchClaudeUsage: vi.fn(), +})); + +vi.mock('./env.js', () => ({ + readEnvFile: vi.fn(() => ({ + FALLBACK_PROVIDER_NAME: 'kimi', + FALLBACK_BASE_URL: 'https://api.kimi.com/coding/', + FALLBACK_AUTH_TOKEN: 'test-kimi-key', + FALLBACK_MODEL: 'kimi-k2.5', + FALLBACK_SMALL_MODEL: 'kimi-k2.5', + FALLBACK_COOLDOWN_MS: '600000', + })), +})); + +vi.mock('./logger.js', () => ({ + logger: { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, +})); + +import { fetchClaudeUsage } from './claude-usage.js'; +import { + clearPrimaryCooldown, + getActiveProvider, + getCooldownInfo, + markPrimaryCooldown, + resetFallbackConfig, +} from './provider-fallback.js'; + +describe('provider fallback usage recovery', () => { + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-03-24T00:00:00.000Z')); + vi.clearAllMocks(); + clearPrimaryCooldown(); + resetFallbackConfig(); + delete process.env.FALLBACK_PROVIDER_NAME; + delete process.env.FALLBACK_BASE_URL; + delete process.env.FALLBACK_AUTH_TOKEN; + delete process.env.FALLBACK_MODEL; + delete process.env.FALLBACK_SMALL_MODEL; + delete process.env.FALLBACK_COOLDOWN_MS; + }); + + afterEach(() => { + clearPrimaryCooldown(); + resetFallbackConfig(); + vi.useRealTimers(); + }); + + it('keeps the fallback provider active while Claude usage is still exhausted', async () => { + vi.mocked(fetchClaudeUsage).mockResolvedValue({ + five_hour: { + utilization: 100, + resets_at: '2026-03-24T04:00:00.000+09:00', + }, + }); + + markPrimaryCooldown('usage-exhausted', 1_000); + vi.advanceTimersByTime(5_000); + + await expect(getActiveProvider()).resolves.toBe('kimi'); + expect(getCooldownInfo()).toMatchObject({ + active: true, + reason: 'usage-exhausted', + remainingMs: 0, + }); + }); + + it('returns to Claude immediately when usage is no longer exhausted', async () => { + vi.mocked(fetchClaudeUsage).mockResolvedValue({ + five_hour: { + utilization: 72, + resets_at: '2026-03-24T04:00:00.000+09:00', + }, + seven_day: { + utilization: 55, + resets_at: '2026-03-31T04:00:00.000+09:00', + }, + }); + + markPrimaryCooldown('usage-exhausted', 600_000); + + await expect(getActiveProvider()).resolves.toBe('claude'); + expect(getCooldownInfo()).toEqual({ active: false }); + }); + + it('falls back to time-based retry when usage status cannot be fetched', async () => { + vi.mocked(fetchClaudeUsage).mockResolvedValue(null); + + markPrimaryCooldown('usage-exhausted', 1_000); + vi.advanceTimersByTime(5_000); + + await expect(getActiveProvider()).resolves.toBe('claude'); + expect(getCooldownInfo()).toEqual({ active: false }); + }); +}); diff --git a/src/provider-fallback.ts b/src/provider-fallback.ts index 451f29d..20343b7 100644 --- a/src/provider-fallback.ts +++ b/src/provider-fallback.ts @@ -14,8 +14,10 @@ import fs from 'fs'; +import { fetchClaudeUsage, type ClaudeUsageData } from './claude-usage.js'; import { readEnvFile } from './env.js'; import { logger } from './logger.js'; +import { rotateToken, getTokenCount } from './token-rotation.js'; // ── Types ──────────────────────────────────────────────────────── @@ -46,6 +48,15 @@ interface FallbackConfig { // ── State ──────────────────────────────────────────────────────── let cooldown: CooldownState | null = null; +let lastUsageAvailabilityCheck: { + checkedAt: number; + result: 'available' | 'exhausted' | 'unknown'; +} | null = null; +let usageAvailabilityCheckPromise: Promise< + 'available' | 'exhausted' | 'unknown' +> | null = null; + +const USAGE_RECOVERY_RECHECK_MS = 30_000; // ── Config ─────────────────────────────────────────────────────── @@ -116,16 +127,112 @@ export function getFallbackProviderName(): string { return loadConfig().providerName; } +function normalizeUtilization(utilization: number): number { + return utilization > 1 ? utilization : utilization * 100; +} + +type ClaudeUsageWindow = NonNullable; + +function hasExhaustedClaudeUsageWindow( + usage: ClaudeUsageData | null, +): boolean | null { + if (!usage) return null; + const windows: ClaudeUsageWindow[] = []; + if (usage.five_hour) windows.push(usage.five_hour); + if (usage.seven_day) windows.push(usage.seven_day); + if (usage.seven_day_sonnet) windows.push(usage.seven_day_sonnet); + if (usage.seven_day_opus) windows.push(usage.seven_day_opus); + if (windows.length === 0) return null; + return windows.some( + (window) => normalizeUtilization(window.utilization) >= 100, + ); +} + +function clearUsageAvailabilityCache(): void { + lastUsageAvailabilityCheck = null; + usageAvailabilityCheckPromise = null; +} + +async function getClaudeUsageAvailability(): Promise< + 'available' | 'exhausted' | 'unknown' +> { + const now = Date.now(); + if ( + lastUsageAvailabilityCheck && + now - lastUsageAvailabilityCheck.checkedAt < USAGE_RECOVERY_RECHECK_MS + ) { + return lastUsageAvailabilityCheck.result; + } + + if (!usageAvailabilityCheckPromise) { + usageAvailabilityCheckPromise = (async () => { + const usage = await fetchClaudeUsage(); + const exhausted = hasExhaustedClaudeUsageWindow(usage); + const result = + exhausted === null ? 'unknown' : exhausted ? 'exhausted' : 'available'; + lastUsageAvailabilityCheck = { + checkedAt: Date.now(), + result, + }; + return result; + })(); + + void usageAvailabilityCheckPromise.finally(() => { + usageAvailabilityCheckPromise = null; + }); + } + + return usageAvailabilityCheckPromise; +} + /** * Determine which provider should be used for the next request. * Returns 'claude' when Claude is healthy or cooldown has expired, * or the fallback provider name during an active cooldown. */ -export function getActiveProvider(): string { +export async function getActiveProvider(): Promise { const config = loadConfig(); if (!config.enabled) return 'claude'; if (cooldown) { + if (cooldown.reason === 'usage-exhausted') { + const usageAvailability = await getClaudeUsageAvailability(); + if (usageAvailability === 'available') { + logger.info( + { + provider: 'claude', + reason: cooldown.reason, + }, + 'Claude usage recovered, retrying primary provider', + ); + cooldown = null; + clearUsageAvailabilityCache(); + return 'claude'; + } + if (usageAvailability === 'exhausted') { + // Current token exhausted — try rotating to another token (ignore cooldowns) + if ( + getTokenCount() > 1 && + rotateToken(undefined, { ignoreRateLimits: true }) + ) { + logger.info( + 'Claude current token exhausted, rotated to next token — retrying', + ); + cooldown = null; + clearUsageAvailabilityCache(); + return 'claude'; + } + logger.debug( + { + provider: config.providerName, + reason: cooldown.reason, + }, + 'All Claude tokens exhausted, keeping cooldown active', + ); + return config.providerName; + } + } + if (Date.now() < cooldown.expiresAt) { return config.providerName; } @@ -161,6 +268,7 @@ export function markPrimaryCooldown( expiresAt: now + durationMs, reason, }; + clearUsageAvailabilityCache(); logger.info( { @@ -175,6 +283,7 @@ export function markPrimaryCooldown( /** Manually clear cooldown (e.g. after a successful Claude response). */ export function clearPrimaryCooldown(): void { + clearUsageAvailabilityCache(); if (cooldown) { logger.info( { reason: cooldown.reason }, @@ -184,6 +293,11 @@ export function clearPrimaryCooldown(): void { } } +/** Check if Claude is currently in usage-exhausted cooldown. */ +export function isUsageExhausted(): boolean { + return cooldown?.reason === 'usage-exhausted'; +} + /** Get current cooldown info (for diagnostics / status dashboard). */ export function getCooldownInfo(): { active: boolean; @@ -191,14 +305,18 @@ export function getCooldownInfo(): { expiresAt?: string; remainingMs?: number; } { - if (!cooldown || Date.now() >= cooldown.expiresAt) { + if (!cooldown) { + return { active: false }; + } + const remainingMs = Math.max(cooldown.expiresAt - Date.now(), 0); + if (cooldown.reason !== 'usage-exhausted' && remainingMs === 0) { return { active: false }; } return { active: true, reason: cooldown.reason, expiresAt: new Date(cooldown.expiresAt).toISOString(), - remainingMs: cooldown.expiresAt - Date.now(), + remainingMs, }; } @@ -247,6 +365,8 @@ export function detectFallbackTrigger( if ( lower.includes('429') || lower.includes('rate limit') || + lower.includes('usage limit') || + lower.includes('hit your limit') || lower.includes('too many requests') || lower.includes('rate_limit') ) { diff --git a/src/task-scheduler.test.ts b/src/task-scheduler.test.ts index 97fa40d..40f79cb 100644 --- a/src/task-scheduler.test.ts +++ b/src/task-scheduler.test.ts @@ -8,6 +8,42 @@ const { runAgentProcessMock, writeTasksSnapshotMock } = vi.hoisted(() => ({ writeTasksSnapshotMock: vi.fn(), })); +vi.mock('./provider-fallback.js', () => ({ + detectFallbackTrigger: vi.fn((error?: string | null) => { + const lower = (error || '').toLowerCase(); + if ( + lower.includes('429') || + lower.includes('rate limit') || + lower.includes('hit your limit') + ) { + return { shouldFallback: true, reason: '429' }; + } + return { shouldFallback: false, reason: '' }; + }), + getActiveProvider: vi.fn(async () => 'claude'), + getFallbackEnvOverrides: vi.fn(() => ({ + ANTHROPIC_BASE_URL: 'https://api.kimi.com/coding/', + ANTHROPIC_AUTH_TOKEN: 'test-kimi-key', + ANTHROPIC_MODEL: 'kimi-k2.5', + })), + getFallbackProviderName: vi.fn(() => 'kimi'), + hasGroupProviderOverride: vi.fn(() => false), + isFallbackEnabled: vi.fn(() => true), + markPrimaryCooldown: vi.fn(), +})); + +vi.mock('./token-rotation.js', () => ({ + rotateToken: vi.fn(() => false), + getTokenCount: vi.fn(() => 1), + markTokenHealthy: vi.fn(), +})); + +vi.mock('./codex-token-rotation.js', () => ({ + rotateCodexToken: vi.fn(() => false), + getCodexAccountCount: vi.fn(() => 1), + markCodexTokenHealthy: vi.fn(), +})); + vi.mock('./agent-runner.js', async () => { const actual = await vi.importActual( @@ -21,8 +57,10 @@ vi.mock('./agent-runner.js', async () => { }); import { _initTestDatabase, createTask, getTaskById } from './db.js'; +import * as providerFallback from './provider-fallback.js'; import { createTaskStatusTracker } from './task-status-tracker.js'; import { TASK_STATUS_MESSAGE_PREFIX } from './task-watch-status.js'; +import * as tokenRotation from './token-rotation.js'; import { _resetSchedulerLoopForTests, computeNextRun, @@ -39,6 +77,11 @@ describe('task scheduler', () => { _resetSchedulerLoopForTests(); runAgentProcessMock.mockClear(); writeTasksSnapshotMock.mockClear(); + vi.mocked(providerFallback.getActiveProvider).mockResolvedValue('claude'); + vi.mocked(providerFallback.isFallbackEnabled).mockReturnValue(true); + vi.mocked(providerFallback.hasGroupProviderOverride).mockReturnValue(false); + vi.mocked(tokenRotation.getTokenCount).mockReturnValue(1); + vi.mocked(tokenRotation.rotateToken).mockReturnValue(false); vi.useFakeTimers(); }); @@ -232,6 +275,103 @@ Check the run. expect(enqueueTask.mock.calls[0][1]).toBe('task-watch-group'); }); + it('suppresses Claude usage banners for scheduled tasks and retries with a rotated account', async () => { + const dueAt = new Date(Date.now() - 60_000).toISOString(); + createTask({ + id: 'task-usage-banner', + group_folder: 'shared-group', + chat_jid: 'shared@g.us', + agent_type: 'claude-code', + prompt: 'claude task', + schedule_type: 'once', + schedule_value: dueAt, + context_mode: 'isolated', + next_run: dueAt, + status: 'active', + created_at: '2026-02-22T00:00:00.000Z', + }); + + vi.mocked(tokenRotation.getTokenCount).mockReturnValue(2); + vi.mocked(tokenRotation.rotateToken).mockReturnValueOnce(true); + + (runAgentProcessMock as any) + .mockImplementationOnce( + async ( + _group: unknown, + _input: unknown, + _onProcess: unknown, + onOutput?: (output: Record) => Promise, + ) => { + await onOutput?.({ + status: 'success', + phase: 'intermediate', + result: "You're out of extra usage · resets 4am (Asia/Seoul)", + }); + await onOutput?.({ + status: 'success', + result: "You're out of extra usage · resets 4am (Asia/Seoul)", + }); + return { + status: 'success', + result: null, + }; + }, + ) + .mockImplementationOnce( + async ( + _group: unknown, + _input: unknown, + _onProcess: unknown, + onOutput?: (output: Record) => Promise, + ) => { + await onOutput?.({ + status: 'success', + result: 'rotated scheduled task response', + }); + return { + status: 'success', + result: null, + }; + }, + ); + + const enqueueTask = vi.fn( + (_groupJid: string, _taskId: string, fn: () => Promise) => { + void fn(); + }, + ); + const sendMessage = vi.fn(async () => {}); + + startSchedulerLoop({ + serviceAgentType: 'claude-code', + registeredGroups: () => ({ + 'shared@g.us': { + name: 'Shared', + folder: 'shared-group', + trigger: '@Claude', + added_at: '2026-02-22T00:00:00.000Z', + agentType: 'claude-code', + }, + }), + getSessions: () => ({}), + queue: { enqueueTask } as any, + onProcess: () => {}, + sendMessage, + }); + + await vi.advanceTimersByTimeAsync(10); + + expect(runAgentProcessMock).toHaveBeenCalledTimes(2); + expect(tokenRotation.rotateToken).toHaveBeenCalledTimes(1); + expect(tokenRotation.markTokenHealthy).toHaveBeenCalledTimes(1); + expect(providerFallback.markPrimaryCooldown).not.toHaveBeenCalled(); + expect(sendMessage).toHaveBeenCalledTimes(1); + expect(sendMessage).toHaveBeenCalledWith( + 'shared@g.us', + 'rotated scheduled task response', + ); + }); + it('picks up newly due tasks immediately when nudged', async () => { const enqueueTask = vi.fn(); @@ -352,6 +492,7 @@ Check the run. }), expect.any(Function), expect.any(Function), + undefined, ); expect(writeTasksSnapshotMock).toHaveBeenCalledWith( 'shared-group', @@ -411,6 +552,7 @@ Check the run. }), expect.any(Function), expect.any(Function), + undefined, ); expect(writeTasksSnapshotMock).toHaveBeenCalledWith( 'shared-group', diff --git a/src/task-scheduler.ts b/src/task-scheduler.ts index 04895fc..4db678b 100644 --- a/src/task-scheduler.ts +++ b/src/task-scheduler.ts @@ -1,9 +1,11 @@ import { ChildProcess } from 'child_process'; import { CronExpressionParser } from 'cron-parser'; import fs from 'fs'; +import path from 'path'; import { ASSISTANT_NAME, + DATA_DIR, SCHEDULER_POLL_INTERVAL, SERVICE_AGENT_TYPE, TIMEZONE, @@ -29,6 +31,26 @@ import { } from './group-folder.js'; import { logger } from './logger.js'; import { createTaskStatusTracker } from './task-status-tracker.js'; +import { + detectFallbackTrigger, + getActiveProvider, + getFallbackEnvOverrides, + getFallbackProviderName, + hasGroupProviderOverride, + isFallbackEnabled, + isUsageExhausted, + markPrimaryCooldown, +} from './provider-fallback.js'; +import { + rotateCodexToken, + getCodexAccountCount, + markCodexTokenHealthy, +} from './codex-token-rotation.js'; +import { + rotateToken, + getTokenCount, + markTokenHealthy, +} from './token-rotation.js'; import { evaluateTaskSuspension, formatSuspensionNotice, @@ -50,6 +72,25 @@ export { shouldUseTaskScopedSession, } from './task-watch-status.js'; +function isClaudeUsageExhaustedMessage(text: string): boolean { + const normalized = text + .trim() + .toLowerCase() + .replace(/[’‘`]/g, "'") + .replace(/\s+/g, ' ') + .replace(/^error:\s*/i, ''); + const looksLikeBanner = + normalized.startsWith("you're out of extra usage") || + normalized.startsWith('you are out of extra usage') || + normalized.startsWith("you've hit your limit") || + normalized.startsWith('you have hit your limit'); + const hasResetHint = + normalized.includes('resets ') || + normalized.includes('reset at ') || + normalized.includes('try again'); + return looksLikeBanner && hasResetHint && normalized.length <= 160; +} + /** * Compute the next run time for a recurring task, anchored to the * task's scheduled time rather than Date.now() to prevent cumulative @@ -238,51 +279,289 @@ async function runTask( sendTrackedMessage: deps.sendTrackedMessage, editTrackedMessage: deps.editTrackedMessage, }); + const settingsPath = path.join( + DATA_DIR, + 'sessions', + task.group_folder, + '.claude', + 'settings.json', + ); + const canFallback = + context.taskAgentType === 'claude-code' && + isFallbackEnabled() && + !hasGroupProviderOverride(settingsPath); try { await statusTracker.update('checking'); - const output = await runAgentProcess( - context.group, - { - prompt: task.prompt, - sessionId: context.sessionId, - groupFolder: task.group_folder, - chatJid: task.chat_jid, - isMain: context.isMain, - isScheduledTask: true, - runtimeTaskId: context.runtimeTaskId, - useTaskScopedSession: context.useTaskScopedSession, - assistantName: ASSISTANT_NAME, + const runTaskAttempt = async ( + provider: string, + ): Promise<{ + output: AgentOutput; + sawOutput: boolean; + streamedTriggerReason?: { + reason: string; + retryAfterMs?: number; + }; + attemptResult: string | null; + attemptError: string | null; + }> => { + let sawOutput = false; + let attemptResult: string | null = null; + let attemptError: string | null = null; + let streamedTriggerReason: + | { + reason: string; + retryAfterMs?: number; + } + | undefined; + + const output = await runAgentProcess( + context.group, + { + prompt: task.prompt, + sessionId: context.sessionId, + groupFolder: task.group_folder, + chatJid: task.chat_jid, + isMain: context.isMain, + isScheduledTask: true, + runtimeTaskId: context.runtimeTaskId, + useTaskScopedSession: context.useTaskScopedSession, + assistantName: ASSISTANT_NAME, + }, + (proc, processName) => + deps.onProcess( + context.queueJid, + proc, + processName, + context.runtimeIpcDir, + ), + async (streamedOutput: AgentOutput) => { + if (streamedOutput.phase === 'progress') { + return; + } + + if ( + canFallback && + provider === 'claude' && + !sawOutput && + streamedOutput.status === 'success' && + typeof streamedOutput.result === 'string' && + isClaudeUsageExhaustedMessage(streamedOutput.result) + ) { + if (!streamedTriggerReason) { + logger.warn( + { + taskId: task.id, + taskChatJid: task.chat_jid, + group: context.group.name, + groupFolder: task.group_folder, + resultPreview: streamedOutput.result.slice(0, 120), + }, + 'Detected Claude usage exhaustion banner during scheduled task output', + ); + } + streamedTriggerReason = { reason: 'usage-exhausted' }; + return; + } + + if (streamedOutput.result) { + sawOutput = true; + attemptResult = streamedOutput.result; + await deps.sendMessage(task.chat_jid, streamedOutput.result); + } + + if (streamedOutput.status === 'error') { + attemptError = streamedOutput.error || 'Unknown error'; + if (!sawOutput && !streamedTriggerReason) { + const trigger = detectFallbackTrigger(streamedOutput.error); + if (trigger.shouldFallback) { + streamedTriggerReason = { + reason: trigger.reason, + retryAfterMs: trigger.retryAfterMs, + }; + } + } + } + }, + provider === 'claude' ? undefined : getFallbackEnvOverrides(), + ); + + if (output.status === 'error' && !attemptError) { + attemptError = output.error || 'Unknown error'; + } else if (output.result && !attemptResult) { + attemptResult = output.result; + } + + return { + output, + sawOutput, + streamedTriggerReason, + attemptResult, + attemptError, + }; + }; + + const shouldRotateClaudeToken = (reason: string): boolean => + reason === '429' || reason === 'usage-exhausted'; + + const runFallbackTaskAttempt = async ( + reason: string, + retryAfterMs?: number, + ): Promise => { + if (!canFallback) { + error = reason; + return; + } + + const fallbackName = getFallbackProviderName(); + markPrimaryCooldown(reason, retryAfterMs); + + logger.info( + { + taskId: task.id, + group: context.group.name, + groupFolder: task.group_folder, + reason, + retryAfterMs, + fallbackProvider: fallbackName, + }, + `Falling back to provider: ${fallbackName} for scheduled task (reason: ${reason})`, + ); + + const fallbackAttempt = await runTaskAttempt(fallbackName); + result = fallbackAttempt.attemptResult; + error = + fallbackAttempt.output.status === 'error' + ? fallbackAttempt.attemptError || 'Unknown error' + : null; + }; + + const retryClaudeTaskWithRotation = async ( + initialTrigger: { + reason: string; + retryAfterMs?: number; }, - (proc, processName) => - deps.onProcess( - context.queueJid, - proc, - processName, - context.runtimeIpcDir, - ), - async (streamedOutput: AgentOutput) => { - if (streamedOutput.phase === 'progress') { + rotationMessage?: string, + ): Promise => { + let trigger = initialTrigger; + let lastRotationMessage = rotationMessage; + + while ( + shouldRotateClaudeToken(trigger.reason) && + getTokenCount() > 1 && + rotateToken(lastRotationMessage, { ignoreRateLimits: true }) + ) { + logger.info( + { + taskId: task.id, + group: context.group.name, + groupFolder: task.group_folder, + reason: trigger.reason, + }, + 'Scheduled task Claude rate-limited, retrying with rotated account', + ); + + const retryAttempt = await runTaskAttempt('claude'); + result = retryAttempt.attemptResult; + error = retryAttempt.attemptError; + + if ( + retryAttempt.streamedTriggerReason && + !retryAttempt.sawOutput && + retryAttempt.output.status !== 'error' + ) { + trigger = { + reason: retryAttempt.streamedTriggerReason.reason, + retryAfterMs: retryAttempt.streamedTriggerReason.retryAfterMs, + }; + lastRotationMessage = + typeof retryAttempt.output.result === 'string' + ? retryAttempt.output.result + : undefined; + continue; + } + + if (retryAttempt.output.status === 'error' && !retryAttempt.sawOutput) { + const retryTrigger = retryAttempt.streamedTriggerReason + ? { + shouldFallback: true, + reason: retryAttempt.streamedTriggerReason.reason, + retryAfterMs: retryAttempt.streamedTriggerReason.retryAfterMs, + } + : detectFallbackTrigger(retryAttempt.attemptError); + if (retryTrigger.shouldFallback) { + trigger = { + reason: retryTrigger.reason, + retryAfterMs: retryTrigger.retryAfterMs, + }; + lastRotationMessage = retryAttempt.attemptError || undefined; + continue; + } + } + + if (retryAttempt.output.status === 'success') { + markTokenHealthy(); + error = null; return; } - if (streamedOutput.result) { - result = streamedOutput.result; - // Forward result to user (sendMessage handles formatting) - await deps.sendMessage(task.chat_jid, streamedOutput.result); - } - if (streamedOutput.status === 'error') { - error = streamedOutput.error || 'Unknown error'; - } - }, - ); - if (output.status === 'error') { - error = output.error || 'Unknown error'; - } else if (output.result) { - // Result was already forwarded to the user via the streaming callback above - result = output.result; - } + return; + } + + // Usage exhausted: don't fall back to Kimi — just mark cooldown and skip + if (trigger.reason === 'usage-exhausted') { + markPrimaryCooldown(trigger.reason, trigger.retryAfterMs); + logger.info( + { taskId: task.id, group: context.group.name }, + 'All Claude tokens usage-exhausted, skipping Kimi fallback for scheduled task', + ); + error = 'Claude usage exhausted'; + return; + } + + await runFallbackTaskAttempt(trigger.reason, trigger.retryAfterMs); + }; + + const provider = canFallback ? await getActiveProvider() : 'claude'; + + // Already in usage-exhausted cooldown — skip task instead of running on Kimi + if (provider !== 'claude' && isUsageExhausted()) { + logger.info( + { taskId: task.id, group: context.group.name, provider }, + 'Claude usage exhausted (cooldown active), skipping scheduled task', + ); + error = 'Claude usage exhausted'; + // Fall through to task completion handling below + } else { + const attempt = await runTaskAttempt(provider); + result = attempt.attemptResult; + error = attempt.attemptError; + + if ( + provider === 'claude' && + attempt.streamedTriggerReason && + !attempt.sawOutput + ) { + await retryClaudeTaskWithRotation(attempt.streamedTriggerReason); + } else if (attempt.output.status === 'error' && provider === 'claude') { + const trigger = attempt.streamedTriggerReason + ? { + shouldFallback: true, + reason: attempt.streamedTriggerReason.reason, + retryAfterMs: attempt.streamedTriggerReason.retryAfterMs, + } + : detectFallbackTrigger(error); + if (trigger.shouldFallback) { + await retryClaudeTaskWithRotation({ + reason: trigger.reason, + retryAfterMs: trigger.retryAfterMs, + }); + } + } else if (attempt.output.status === 'error') { + error = attempt.attemptError || 'Unknown error'; + } + } // end else (non-exhausted path) logger.info( { @@ -315,6 +594,31 @@ async function runTask( updateTask(task.id, { suspended_until: null }); } + // Try token rotation before suspending + if (error) { + const trigger = detectFallbackTrigger(error); + if (trigger.shouldFallback) { + const isCodex = SERVICE_AGENT_TYPE === 'codex'; + const rotated = isCodex + ? getCodexAccountCount() > 1 && rotateCodexToken(error) + : getTokenCount() > 1 && rotateToken(error); + if (rotated) { + logger.info( + { + taskId: task.id, + agent: SERVICE_AGENT_TYPE, + reason: trigger.reason, + }, + 'Task rate-limited, rotated token — will retry on next schedule', + ); + if (isCodex) markCodexTokenHealthy(); + else markTokenHealthy(); + // Clear the error so suspension doesn't trigger + error = null; + } + } + } + // Check for repeated quota/auth errors → auto-suspend let suspended = false; if (error) { diff --git a/src/token-rotation.ts b/src/token-rotation.ts index 18bb4cd..664ceb8 100644 --- a/src/token-rotation.ts +++ b/src/token-rotation.ts @@ -1,64 +1,225 @@ +/** + * 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 fs from 'fs'; -import os from 'os'; import path from 'path'; -const HOST_CLAUDE_DIR = path.join(os.homedir(), '.claude'); -const CLAUDE_ACCOUNTS_DIR = path.join(os.homedir(), '.claude-accounts'); +import { DATA_DIR } from './config.js'; +import { readEnvFile } from './env.js'; +import { logger } from './logger.js'; -let activeIndex = 0; -const rateLimitedAccounts = new Set(); +const STATE_FILE = path.join(DATA_DIR, 'token-rotation-state.json'); -function listClaudeAccounts(): string[] { - if (!fs.existsSync(CLAUDE_ACCOUNTS_DIR)) return []; - - return fs - .readdirSync(CLAUDE_ACCOUNTS_DIR, { withFileTypes: true }) - .filter((entry) => entry.isDirectory() && /^\d+$/.test(entry.name)) - .sort((a, b) => Number(a.name) - Number(b.name)) - .map((entry) => path.join(CLAUDE_ACCOUNTS_DIR, entry.name)); +interface TokenState { + token: string; + rateLimitedUntil: number | null; } -function copyIfExists(src: string, dst: string): void { - if (!fs.existsSync(src)) return; - fs.mkdirSync(path.dirname(dst), { recursive: true }); - fs.copyFileSync(src, dst); +const tokens: TokenState[] = []; +let currentIndex = 0; +let initialized = false; + +export function initTokenRotation(): void { + if (initialized) return; + initialized = true; + + const envFile = readEnvFile([ + 'CLAUDE_CODE_OAUTH_TOKENS', + 'CLAUDE_CODE_OAUTH_TOKEN', + ]); + const multi = + process.env.CLAUDE_CODE_OAUTH_TOKENS || envFile.CLAUDE_CODE_OAUTH_TOKENS; + const single = + process.env.CLAUDE_CODE_OAUTH_TOKEN || envFile.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) { + loadState(); + logger.info( + { count: tokens.length, activeIndex: currentIndex }, + `Token rotation initialized with ${tokens.length} tokens`, + ); + } } -function syncHostClaudeAccount(accountDir: string): void { - copyIfExists( - path.join(accountDir, '.credentials.json'), - path.join(HOST_CLAUDE_DIR, '.credentials.json'), +function saveState(): void { + try { + const state = { + currentIndex, + rateLimits: tokens.map((t) => t.rateLimitedUntil), + }; + fs.writeFileSync(STATE_FILE, JSON.stringify(state)); + } catch { + /* best effort */ + } +} + +function loadState(): void { + try { + if (!fs.existsSync(STATE_FILE)) return; + const state = JSON.parse(fs.readFileSync(STATE_FILE, 'utf-8')); + const now = Date.now(); + if ( + typeof state.currentIndex === 'number' && + state.currentIndex < tokens.length + ) { + currentIndex = state.currentIndex; + } + if (Array.isArray(state.rateLimits)) { + for ( + let i = 0; + i < Math.min(state.rateLimits.length, tokens.length); + i++ + ) { + const until = state.rateLimits[i]; + if (typeof until === 'number' && until > now) { + tokens[i].rateLimitedUntil = until; + } + } + } + logger.info( + { currentIndex, tokenCount: tokens.length }, + 'Token rotation state restored', + ); + } catch { + /* start fresh */ + } +} + +const BUFFER_MS = 3 * 60_000; +const DEFAULT_COOLDOWN_MS = 3_600_000; + +function parseRetryAfterFromError(error?: string): number | null { + if (!error) return null; + const match = error.match( + /(?:try again at|resets?\s+(?:at\s+)?)\s*(\w+ \d{1,2}(?:st|nd|rd|th)?,?\s+\d{4}\s+\d{1,2}:\d{2}\s*(?:AM|PM)?)/i, ); - copyIfExists( - path.join(accountDir, 'settings.json'), - path.join(HOST_CLAUDE_DIR, 'settings.json'), + if (!match) return null; + try { + const cleaned = match[1].replace(/(\d+)(?:st|nd|rd|th)/i, '$1'); + const ts = new Date(cleaned).getTime(); + if (Number.isNaN(ts)) return null; + return ts; + } catch { + return null; + } +} + +function computeCooldownUntil(error?: string): number { + const retryAt = parseRetryAfterFromError(error); + if (retryAt) return retryAt + BUFFER_MS; + return Date.now() + DEFAULT_COOLDOWN_MS; +} + +/** 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( + errorMessage?: string, + opts?: { ignoreRateLimits?: boolean }, +): boolean { + if (tokens.length <= 1) return false; + + const now = Date.now(); + tokens[currentIndex].rateLimitedUntil = computeCooldownUntil(errorMessage); + const ignoreRL = opts?.ignoreRateLimits ?? false; + + // Find next available token + for (let i = 1; i < tokens.length; i++) { + const idx = (currentIndex + i) % tokens.length; + const state = tokens[idx]; + if (ignoreRL || !state.rateLimitedUntil || state.rateLimitedUntil <= now) { + state.rateLimitedUntil = null; + currentIndex = idx; + logger.info( + { tokenIndex: currentIndex, totalTokens: tokens.length, ignoreRL }, + `Rotated to token #${currentIndex + 1}/${tokens.length}`, + ); + saveState(); + return true; + } + } + + logger.warn( + { totalTokens: tokens.length }, + 'All tokens are rate-limited, falling through to provider fallback', ); + return false; } -export function getTokenCount(): number { - const accounts = listClaudeAccounts(); - return accounts.length > 0 ? accounts.length : 1; -} - -export function rotateToken(_reason?: string): boolean { - const accounts = listClaudeAccounts(); - if (accounts.length <= 1) return false; - - rateLimitedAccounts.add(activeIndex); - - const candidates = accounts.map((_, index) => index); - const nextIndex = - candidates.find( - (index) => index !== activeIndex && !rateLimitedAccounts.has(index), - ) ?? candidates.find((index) => index !== activeIndex); - - if (nextIndex === undefined) return false; - - syncHostClaudeAccount(accounts[nextIndex]); - activeIndex = nextIndex; - return true; -} - +/** Clear rate-limit flag for the current token (on successful response). */ export function markTokenHealthy(): void { - rateLimitedAccounts.delete(activeIndex); + if (tokens.length === 0) return; + const state = tokens[currentIndex]; + if (state?.rateLimitedUntil) { + state.rateLimitedUntil = null; + saveState(); + } +} + +/** Number of configured tokens. */ +export function getTokenCount(): number { + return tokens.length; +} + +/** Get all configured tokens (masked for display, raw for API calls). */ +export function getAllTokens(): { + index: number; + token: string; + masked: string; + isActive: boolean; + isRateLimited: boolean; +}[] { + const now = Date.now(); + return tokens.map((t, i) => ({ + index: i, + token: t.token, + masked: `${t.token.slice(0, 20)}...${t.token.slice(-4)}`, + isActive: i === currentIndex, + isRateLimited: Boolean(t.rateLimitedUntil && t.rateLimitedUntil > now), + })); +} + +/** 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, + }; } diff --git a/src/unified-dashboard.test.ts b/src/unified-dashboard.test.ts new file mode 100644 index 0000000..5600537 --- /dev/null +++ b/src/unified-dashboard.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it, vi } from 'vitest'; + +import type { ClaudeAccountUsage } from './claude-usage.js'; + +vi.mock('./claude-usage.js', () => ({ + fetchAllClaudeUsage: vi.fn(), + fetchAllClaudeProfiles: vi.fn(), + getClaudeProfile: (index: number) => + index === 0 + ? { email: 'a@example.com', planType: 'max' } + : { email: 'b@example.com', planType: 'pro' }, +})); + +vi.mock('./codex-token-rotation.js', () => ({ + getAllCodexAccounts: () => [], + updateCodexAccountUsage: vi.fn(), +})); + +import { + buildClaudeUsageRows, + mergeClaudeDashboardAccounts, +} from './unified-dashboard.js'; + +describe('unified dashboard Claude usage rows', () => { + it('keeps both Claude accounts visible when one account usage is unavailable', () => { + const rows = buildClaudeUsageRows([ + { + index: 0, + masked: 'tok-a', + isActive: true, + isRateLimited: false, + usage: { + five_hour: { + utilization: 0.4, + resets_at: '2026-03-24T04:00:00+09:00', + }, + seven_day: { + utilization: 0.7, + resets_at: '2026-03-29T04:00:00+09:00', + }, + }, + }, + { + index: 1, + masked: 'tok-b', + isActive: false, + isRateLimited: true, + usage: null, + }, + ]); + + expect(rows).toHaveLength(2); + expect(rows[0]).toMatchObject({ + name: 'Claude1* max', + h5pct: 40, + d7pct: 70, + }); + expect(rows[1]).toMatchObject({ + name: 'Claude2! pro', + h5pct: -1, + d7pct: -1, + }); + }); + + it('preserves the last successful usage per account instead of collapsing to one cache entry', () => { + const cachedAccounts: ClaudeAccountUsage[] = [ + { + index: 0, + masked: 'tok-a', + isActive: false, + isRateLimited: false, + usage: { + five_hour: { + utilization: 0.25, + resets_at: '2026-03-24T04:00:00+09:00', + }, + seven_day: { + utilization: 0.5, + resets_at: '2026-03-29T04:00:00+09:00', + }, + }, + }, + { + index: 1, + masked: 'tok-b', + isActive: true, + isRateLimited: false, + usage: { + five_hour: { + utilization: 0.6, + resets_at: '2026-03-24T06:00:00+09:00', + }, + seven_day: { + utilization: 0.8, + resets_at: '2026-03-30T04:00:00+09:00', + }, + }, + }, + ]; + const liveAccounts: ClaudeAccountUsage[] = [ + { + index: 0, + masked: 'tok-a', + isActive: true, + isRateLimited: false, + usage: null, + }, + { + index: 1, + masked: 'tok-b', + isActive: false, + isRateLimited: true, + usage: { + five_hour: { + utilization: 0.9, + resets_at: '2026-03-24T08:00:00+09:00', + }, + seven_day: { + utilization: 0.95, + resets_at: '2026-03-31T04:00:00+09:00', + }, + }, + }, + ]; + + const merged = mergeClaudeDashboardAccounts(liveAccounts, cachedAccounts); + + expect(merged).toHaveLength(2); + expect(merged[0]).toMatchObject({ + index: 0, + isActive: true, + usage: cachedAccounts[0].usage, + }); + expect(merged[1]).toMatchObject({ + index: 1, + isRateLimited: true, + usage: liveAccounts[1].usage, + }); + }); +});