EJClaw/src/token-rotation.ts

/**
 * 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:  surface error to caller
 */

import fs from 'fs';
import path from 'path';

import { DATA_DIR } from './config.js';
import { getEnv } from './env.js';
import { logger } from './logger.js';
import {
  computeCooldownUntil,
  findNextAvailable,
} from './token-rotation-base.js';
import { readJsonFile, writeJsonFile } from './utils.js';

const STATE_FILE = path.join(DATA_DIR, 'token-rotation-state.json');

interface TokenState {
  token: string;
  rateLimitedUntil: number | null;
}

const tokens: TokenState[] = [];
let currentIndex = 0;
let initialized = false;

// ── Post-rotation callback ──────────────────────────────────────
// Invoked after a successful token rotation so that callers (e.g.
// index.ts) can tear down reviewer containers whose SDK process
// still holds the old token.
type TokenRotationCallback = () => void;
let onTokenRotatedCallback: TokenRotationCallback | null = null;

/** Register a callback to be invoked after a successful token rotation. */
export function onTokenRotated(cb: TokenRotationCallback): void {
  onTokenRotatedCallback = cb;
}

export function initTokenRotation(): void {
  if (initialized) return;
  initialized = true;

  const multi = getEnv('CLAUDE_CODE_OAUTH_TOKENS');
  const single = getEnv('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 saveState(): void {
  try {
    const state = {
      currentIndex,
      rateLimits: tokens.map((t) => t.rateLimitedUntil),
    };
    writeJsonFile(STATE_FILE, state);
  } catch {
    /* best effort */
  }
}

function loadState(): void {
  const state = readJsonFile<{
    currentIndex?: number;
    rateLimits?: (number | null)[];
  }>(STATE_FILE);
  if (!state) return;

  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',
  );
}

/** 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 previousIndex = currentIndex;
  const cooldownUntil = computeCooldownUntil(errorMessage);
  tokens[currentIndex].rateLimitedUntil = cooldownUntil;

  const nextIdx = findNextAvailable(tokens, currentIndex, opts);
  if (nextIdx !== null) {
    tokens[nextIdx].rateLimitedUntil = null;
    currentIndex = nextIdx;
    logger.info(
      {
        transition: 'rotation:execute',
        fromIndex: previousIndex,
        toIndex: currentIndex,
        totalTokens: tokens.length,
        ignoreRL: opts?.ignoreRateLimits ?? false,
        cooldownUntil:
          cooldownUntil != null ? new Date(cooldownUntil).toISOString() : null,
        reason: errorMessage ?? null,
      },
      `Rotated to token #${currentIndex + 1}/${tokens.length}`,
    );
    saveState();
    if (onTokenRotatedCallback) {
      try {
        onTokenRotatedCallback();
      } catch {
        /* best effort — don't break rotation flow */
      }
    }
    return true;
  }

  logger.warn(
    {
      transition: 'rotation:skip',
      fromIndex: previousIndex,
      totalTokens: tokens.length,
      ignoreRL: opts?.ignoreRateLimits ?? false,
      cooldownUntil:
        cooldownUntil != null ? new Date(cooldownUntil).toISOString() : null,
      reason: errorMessage ?? null,
    },
    'All tokens are rate-limited, no available tokens',
  );
  return false;
}

/** Clear rate-limit flag for the current token (on successful response). */
export function markTokenHealthy(): void {
  if (tokens.length === 0) return;
  const state = tokens[currentIndex];
  if (state?.rateLimitedUntil) {
    const previousCooldownUntil = state.rateLimitedUntil;
    state.rateLimitedUntil = null;
    logger.info(
      {
        transition: 'rotation:clear-rate-limit',
        tokenIndex: currentIndex,
        cooldownUntil: new Date(previousCooldownUntil).toISOString(),
      },
      'Cleared Claude token rate-limit state after successful response',
    );
    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),
  }));
}

/** Update the access token value for a specific index (after OAuth refresh). */
export function updateTokenValue(index: number, newAccessToken: string): void {
  if (index < 0 || index >= tokens.length) {
    logger.warn(
      { index, total: tokens.length },
      'updateTokenValue: index out of range',
    );
    return;
  }
  tokens[index].token = newAccessToken;
  logger.info(
    {
      index,
      masked: `${newAccessToken.slice(0, 20)}...${newAccessToken.slice(-4)}`,
    },
    'Token value updated after OAuth refresh',
  );
}

/** 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,
  };
}

export function hasAvailableClaudeToken(): boolean {
  if (tokens.length === 0) {
    return Boolean(process.env.CLAUDE_CODE_OAUTH_TOKEN);
  }
  const now = Date.now();
  return tokens.some(
    (token) => !token.rateLimitedUntil || token.rateLimitedUntil <= now,
  );
}