EJClaw/src/usage-primer.ts

import { spawn } from 'child_process';
import fs from 'fs';
import path from 'path';

import { CODEX_WARMUP_CONFIG } from './config.js';
import {
  refreshActiveCodexUsage,
  refreshAllCodexAccountUsage,
} from './codex-usage-collector.js';
import { runCodexWarmupCycle } from './codex-warmup.js';
import { logger } from './logger.js';
import { getAllTokens } from './token-rotation.js';

/**
 * Usage-window alignment primer.
 *
 * Anthropic/OpenAI enforce 5h rolling usage windows that start on first
 * inference call. We anchor windows to fixed KST times by firing a tiny
 * primer call at 08, 13, 18, 23 KST (4 windows × 5h = 20h, leaving a
 * 04:00–08:00 gap blocked by message-runtime-gating).
 *
 * Daily reset point ends up being 08:00 KST (start of first slot after gap).
 */

const PRIMER_HOURS_KST = [8, 13, 18, 23];
const KST_OFFSET_MS = 9 * 60 * 60 * 1000;
const CLAUDE_PRIMER_TIMEOUT_MS = 60_000;
const CODEX_PRIMER_MIN_INTERVAL_MS = 5 * 60 * 60 * 1000;
// Fire the Codex primer unconditionally at every slot, mirroring the Claude
// primer (which always fires unless the account is rate-limited). The 5h usage
// level no longer gates the call (maxUsagePct=100), so a slot is never skipped
// just because Codex happens to be partway through its window. The d7 gate
// stays at 100 so we still skip an account whose weekly quota is exhausted,
// which is the Codex analogue of Claude's rate-limit skip.
const CODEX_PRIMER_MAX_USAGE_PCT = 100;
const CODEX_PRIMER_MAX_D7_USAGE_PCT = 100;

function resolveClaudeBinary(): string {
  const candidates = [
    path.resolve(
      process.cwd(),
      'runners/agent-runner/node_modules/@anthropic-ai/claude-agent-sdk-linux-x64/claude',
    ),
    path.resolve(
      process.cwd(),
      'runners/agent-runner/node_modules/@anthropic-ai/claude-agent-sdk-linux-x64-musl/claude',
    ),
  ];
  for (const p of candidates) {
    if (fs.existsSync(p)) return p;
  }
  return candidates[0];
}

function runClaudePrimerOnce(
  token: string,
): Promise<{ ok: boolean; reason: string }> {
  const binary = resolveClaudeBinary();
  if (!fs.existsSync(binary)) {
    return Promise.resolve({ ok: false, reason: 'binary_missing' });
  }
  const args = [
    '-p',
    'ok',
    '--model',
    'haiku',
    '--output-format',
    'text',
    '--max-turns',
    '1',
    '--dangerously-skip-permissions',
  ];

  return new Promise((resolve) => {
    let done = false;
    let stderr = '';
    const finish = (result: { ok: boolean; reason: string }): void => {
      if (done) return;
      done = true;
      clearTimeout(timer);
      resolve(result);
    };
    const timer = setTimeout(() => {
      try {
        proc.kill();
      } catch {
        /* ignore */
      }
      finish({ ok: false, reason: 'timeout' });
    }, CLAUDE_PRIMER_TIMEOUT_MS);

    const proc = spawn(binary, args, {
      stdio: ['ignore', 'pipe', 'pipe'],
      env: {
        ...process.env,
        CLAUDE_CODE_OAUTH_TOKEN: token,
      },
    });
    proc.stderr?.on('data', (chunk: Buffer) => {
      stderr += chunk.toString();
      if (stderr.length > 2000) stderr = stderr.slice(-2000);
    });
    proc.on('error', (err) => {
      logger.warn({ err }, 'Claude primer spawn error');
      finish({ ok: false, reason: 'spawn_error' });
    });
    proc.on('close', (code) => {
      if (code === 0) {
        finish({ ok: true, reason: 'ok' });
      } else {
        logger.warn(
          { exitCode: code, stderr: stderr.trim().slice(-500) || undefined },
          'Claude primer command failed',
        );
        finish({ ok: false, reason: `exit_${code ?? 'unknown'}` });
      }
    });
  });
}

async function runClaudePrimerCycle(): Promise<void> {
  const tokens = getAllTokens();
  const candidates = tokens.filter((t) => !t.isRateLimited);
  if (candidates.length === 0) {
    logger.info('Claude primer skipped: no eligible tokens');
    return;
  }
  const target = candidates[0];
  logger.info(
    { tokenIndex: target.index },
    'Starting Claude primer (haiku, prompt=ok)',
  );
  const result = await runClaudePrimerOnce(target.token);
  if (result.ok) {
    logger.info({ tokenIndex: target.index }, 'Claude primer completed');
  } else {
    logger.warn(
      { tokenIndex: target.index, reason: result.reason },
      'Claude primer failed',
    );
  }
}

export async function runCodexPrimerCycle(): Promise<void> {
  try {
    await refreshAllCodexAccountUsage();
    await refreshActiveCodexUsage();
    const result = await runCodexWarmupCycle(
      {
        ...CODEX_WARMUP_CONFIG,
        enabled: true,
        minIntervalMs: CODEX_PRIMER_MIN_INTERVAL_MS,
        maxUsagePct: CODEX_PRIMER_MAX_USAGE_PCT,
        maxD7UsagePct: CODEX_PRIMER_MAX_D7_USAGE_PCT,
      },
      { ignoreZeroUsageWindow: true },
    );
    logger.info({ result }, 'Codex primer cycle finished');
  } catch (err) {
    logger.warn({ err }, 'Codex primer cycle threw');
  }
}

async function runPrimerCycle(): Promise<void> {
  logger.info('Running scheduled usage primer cycle');
  await Promise.allSettled([runClaudePrimerCycle(), runCodexPrimerCycle()]);
}

export function msUntilNextPrimerSlotKST(nowMs: number = Date.now()): number {
  // KST = UTC+9 (no DST). Shift so KST clock can be read via getUTC*.
  const kstNowShifted = nowMs + KST_OFFSET_MS;
  const kstDate = new Date(kstNowShifted);
  const kstH = kstDate.getUTCHours();

  let nextH = PRIMER_HOURS_KST.find((h) => h > kstH);
  let dayOffset = 0;
  if (nextH === undefined) {
    nextH = PRIMER_HOURS_KST[0];
    dayOffset = 1;
  }

  const target = new Date(kstNowShifted);
  target.setUTCHours(nextH, 0, 0, 0);
  if (dayOffset === 1) target.setUTCDate(target.getUTCDate() + 1);

  return target.getTime() - kstNowShifted;
}

let primerTimer: ReturnType<typeof setTimeout> | null = null;

export function startUsagePrimer(): void {
  if (primerTimer) return;
  const scheduleNext = (): void => {
    const delay = msUntilNextPrimerSlotKST();
    primerTimer = setTimeout(() => {
      void runPrimerCycle().finally(scheduleNext);
    }, delay);
    const fireAt = new Date(Date.now() + delay).toISOString();
    logger.info(
      { delayMinutes: Math.round(delay / 60_000), fireAt },
      'Next usage primer scheduled',
    );
  };
  scheduleNext();
}

export function stopUsagePrimer(): void {
  if (primerTimer) {
    clearTimeout(primerTimer);
    primerTimer = null;
  }
}