EJClaw/src/usage-primer.ts

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

import { msUntilNextPrimerSlotKST } from './codex-primer-alignment.js';
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 warm-up primer.
 *
 * Fires a tiny inference call at fixed KST slots (08, 13, 18, 23) so each
 * provider account is kept warm and any usage-window issues surface in the
 * logs at predictable times.
 *
 * Reset alignment: the primer alone cannot pin the reset — empirically the
 * Codex 5h limit slides forward with continued usage (observed reset moving
 * 18:00 → 18:33 over ~33 min of activity), so any off-slot Codex call by a
 * paired-room reviewer/arbiter on the same account would anchor the window
 * early. To make the slot the *anchor*, the primer is paired with a hold gate
 * (`shouldHoldCodexForPrimerAlignment` in codex-primer-alignment.ts) that keeps
 * non-primer Codex turns/warm-ups quiet during the dawn gap (04–08 KST) and the
 * minutes around each slot, so the primer wins the first-request race. The
 * primer itself passes `isPrimer` so it is never held. The primer is also a
 * best-effort warm-up + success/failure record.
 */

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 subPaths = [
    '@anthropic-ai/claude-agent-sdk-linux-x64/claude',
    '@anthropic-ai/claude-agent-sdk-linux-x64-musl/claude',
  ];
  // The SDK's platform package is usually hoisted to the repo-root node_modules
  // (the bot runs `bun src/index.ts` from the repo root), but can also live
  // under the agent-runner workspace when deps are not hoisted. Check both so
  // the Claude primer doesn't silently fail with binary_missing.
  const roots = [
    path.resolve(process.cwd(), 'node_modules'),
    path.resolve(process.cwd(), 'runners/agent-runner/node_modules'),
  ];
  const candidates = roots.flatMap((root) =>
    subPaths.map((sub) => path.join(root, sub)),
  );
  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, isPrimer: true },
    );
    // `isPrimer` exempts this call from the alignment hold (the hold blocks
    // every *other* Codex path during the slot windows) so the primer is the
    // first Codex request and anchors the window at the slot. A slot may still
    // land while the trailing 5h window is exhausted, in which case this records
    // `failed`; we do not chase the reported reset time, since a trailing window
    // can keep sliding with later usage.
    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()]);
}

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