EJClaw/src/codex-warmup.ts

import { ChildProcess, spawn } from 'child_process';
import fs from 'fs';
import { createRequire } from 'module';
import os from 'os';
import path from 'path';
import { fileURLToPath } from 'url';

import { DATA_DIR } from './config.js';
import type { AppConfig } from './config/schema.js';
import {
  getAllCodexAccounts,
  getCodexAuthPath,
} from './codex-token-rotation.js';
import { logger } from './logger.js';

export type CodexWarmupConfig = AppConfig['codexWarmup'];

interface CodexWarmupAccountState {
  lastWarmupAt?: string;
  lastAttemptAt?: string;
  lastErrorAt?: string;
  zeroUsageWarmupUntil?: string;
  failures?: number;
}

interface CodexWarmupState {
  lastWarmupAt?: string;
  lastAttemptAt?: string;
  disabledUntil?: string;
  consecutiveFailures?: number;
  accounts?: Record<string, CodexWarmupAccountState>;
}

interface CodexWarmupRuntimeOptions {
  nowMs?: number;
  statePath?: string;
  shouldSkip?: () => boolean;
  ignoreZeroUsageWindow?: boolean;
  /**
   * Fixed-slot primer mode: bypass the opportunistic warm-up skip gates
   * (stagger, post-failure cooldown, min-interval, per-account usage/rate-limit
   * filters) so a real Codex command is attempted at every 08/13/18/23 slot and
   * its success/failure recorded. Without this, two slots exactly 5h apart fall
   * just under minIntervalMs (the primer fires a few hundred ms after the slot)
   * and the later slot is silently skipped as `no_eligible_accounts`.
   */
  forceAttempt?: boolean;
}

export type CodexWarmupCycleResult =
  | { status: 'disabled' }
  | { status: 'skipped'; reason: string }
  | { status: 'warmed'; accountIndex: number }
  | { status: 'failed'; accountIndex: number; reason: string };

const DEFAULT_STATE_FILE = path.join(DATA_DIR, 'codex-warmup-state.json');
const DEFAULT_ZERO_USAGE_WARMUP_WINDOW_MS = 7 * 24 * 60 * 60 * 1000;

function parseTimestamp(value?: string): number | null {
  if (!value) return null;
  const parsed = Date.parse(value);
  return Number.isFinite(parsed) ? parsed : null;
}

function readWarmupState(statePath: string): CodexWarmupState {
  try {
    if (!fs.existsSync(statePath)) return { accounts: {} };
    const parsed = JSON.parse(
      fs.readFileSync(statePath, 'utf8'),
    ) as CodexWarmupState;
    if (!parsed.accounts || typeof parsed.accounts !== 'object') {
      parsed.accounts = {};
    }
    return parsed;
  } catch (err) {
    logger.warn(
      { err, statePath },
      'Failed to read Codex warm-up state; starting fresh',
    );
    return { accounts: {} };
  }
}

function writeWarmupState(statePath: string, state: CodexWarmupState): void {
  fs.mkdirSync(path.dirname(statePath), { recursive: true });
  fs.writeFileSync(`${statePath}.tmp`, `${JSON.stringify(state, null, 2)}\n`);
  fs.renameSync(`${statePath}.tmp`, statePath);
}

function getPreferredCodexPathEntries(): string[] {
  const entries = [
    path.dirname(process.execPath),
    path.join(os.homedir(), '.npm-global', 'bin'),
  ];
  if (process.versions.bun || path.basename(process.execPath) === 'bun') {
    entries.push(path.join(os.homedir(), '.hermes', 'node', 'bin'));
  }
  return [...new Set(entries)];
}

const require = createRequire(import.meta.url);

interface CodexLauncher {
  command: string;
  prefixArgs: string[];
}

/**
 * Resolve how to spawn the Codex CLI for warm-up.
 *
 * The bot runs under bun/systemd where a global `codex` binary usually is NOT
 * on PATH, so the old bare-`codex` fallback failed with ENOENT and the primer
 * never fired. Prefer the bundled `@openai/codex` JS launcher (the same path
 * the codex-runner spawns) and run it through the current JS runtime, falling
 * back to a globally installed binary only when the package is unavailable.
 */
function resolveCodexLauncher(): CodexLauncher {
  const launcherCandidates: string[] = [];
  try {
    launcherCandidates.push(
      path.join(
        path.dirname(require.resolve('@openai/codex/package.json')),
        'bin',
        'codex.js',
      ),
    );
  } catch {
    /* @openai/codex not resolvable from this module — try explicit paths */
  }
  // The codex-runner workspace always vendors @openai/codex.
  const repoRoot = path.resolve(
    path.dirname(fileURLToPath(import.meta.url)),
    '..',
  );
  launcherCandidates.push(
    path.join(
      repoRoot,
      'runners',
      'codex-runner',
      'node_modules',
      '@openai',
      'codex',
      'bin',
      'codex.js',
    ),
  );
  for (const launcher of launcherCandidates) {
    if (fs.existsSync(launcher)) {
      return { command: process.execPath, prefixArgs: [launcher] };
    }
  }
  // Fall back to a globally installed codex binary, then bare PATH lookup.
  const npmGlobalBin = path.join(os.homedir(), '.npm-global', 'bin', 'codex');
  if (fs.existsSync(npmGlobalBin)) {
    return { command: npmGlobalBin, prefixArgs: [] };
  }
  return { command: 'codex', prefixArgs: [] };
}

function ensureAccountState(
  state: CodexWarmupState,
  accountIndex: number,
): CodexWarmupAccountState {
  state.accounts ??= {};
  const key = String(accountIndex);
  state.accounts[key] ??= {};
  return state.accounts[key];
}

function selectWarmupCandidate(
  config: CodexWarmupConfig,
  state: CodexWarmupState,
  nowMs: number,
  options: { ignoreZeroUsageWindow?: boolean; forceAttempt?: boolean } = {},
): { accountIndex: number; zeroUsageWarmupUntil: string } | { reason: string } {
  const disabledUntilMs = parseTimestamp(state.disabledUntil);
  if (
    !options.forceAttempt &&
    disabledUntilMs != null &&
    disabledUntilMs > nowMs
  ) {
    return { reason: 'disabled_cooldown' };
  }

  const lastGlobalWarmupMs = parseTimestamp(state.lastWarmupAt);
  if (
    !options.forceAttempt &&
    lastGlobalWarmupMs != null &&
    config.staggerMs > 0 &&
    nowMs - lastGlobalWarmupMs < config.staggerMs
  ) {
    return { reason: 'stagger_wait' };
  }

  const accounts = getAllCodexAccounts();
  if (accounts.length === 0) return { reason: 'no_accounts' };

  // Fixed-slot primer: attempt a real command on the active account every slot,
  // bypassing the per-account usage/rate-limit/min-interval filters below.
  if (options.forceAttempt) {
    const active = accounts.find((a) => a.isActive) ?? accounts[0];
    return {
      accountIndex: active.index,
      zeroUsageWarmupUntil: new Date(nowMs).toISOString(),
    };
  }

  for (const account of accounts) {
    if (account.isRateLimited) continue;
    if (typeof account.cachedUsagePct !== 'number') continue;
    if (typeof account.cachedUsageD7Pct !== 'number') continue;
    if (account.cachedUsagePct > config.maxUsagePct) continue;
    if (account.cachedUsageD7Pct > config.maxD7UsagePct) continue;

    const accountState = state.accounts?.[String(account.index)];
    const zeroUsageWarmupUntilMs = parseTimestamp(
      accountState?.zeroUsageWarmupUntil,
    );
    if (
      !options.ignoreZeroUsageWindow &&
      zeroUsageWarmupUntilMs != null &&
      zeroUsageWarmupUntilMs > nowMs
    ) {
      continue;
    }

    const lastWarmupMs = parseTimestamp(accountState?.lastWarmupAt);
    if (lastWarmupMs != null && nowMs - lastWarmupMs < config.minIntervalMs) {
      continue;
    }

    const resetD7Ms = parseTimestamp(account.resetD7At);
    let zeroUsageWarmupUntilMsForState = nowMs;
    if (!options.ignoreZeroUsageWindow) {
      zeroUsageWarmupUntilMsForState =
        resetD7Ms != null && resetD7Ms > nowMs
          ? resetD7Ms
          : nowMs + DEFAULT_ZERO_USAGE_WARMUP_WINDOW_MS;
    }

    return {
      accountIndex: account.index,
      zeroUsageWarmupUntil: new Date(
        zeroUsageWarmupUntilMsForState,
      ).toISOString(),
    };
  }

  return { reason: 'no_eligible_accounts' };
}

function runCodexWarmupCommand(
  accountDir: string,
  config: CodexWarmupConfig,
): Promise<{ ok: boolean; reason: string }> {
  const args = [
    'exec',
    '--ephemeral',
    '--ignore-rules',
    '--skip-git-repo-check',
    '--sandbox',
    'read-only',
    '-C',
    os.tmpdir(),
    '-m',
    config.model,
    config.prompt,
  ];

  const spawnEnv: Record<string, string> = {
    ...(process.env as Record<string, string>),
    CODEX_HOME: accountDir,
    PATH: [...getPreferredCodexPathEntries(), process.env.PATH || '']
      .filter(Boolean)
      .join(path.delimiter),
  };

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

    const timer = setTimeout(
      () => finish({ ok: false, reason: 'timeout' }),
      config.commandTimeoutMs,
    );

    try {
      const launcher = resolveCodexLauncher();
      proc = spawn(launcher.command, [...launcher.prefixArgs, ...args], {
        stdio: ['ignore', 'pipe', 'pipe'],
        env: spawnEnv,
      });
    } catch (err) {
      logger.warn({ err }, 'Failed to spawn Codex warm-up command');
      finish({ ok: false, reason: 'spawn_error' });
      return;
    }

    proc.stderr?.on('data', (chunk: Buffer) => {
      stderr += chunk.toString();
      if (stderr.length > 2000) stderr = stderr.slice(-2000);
    });
    proc.on('error', (err) => {
      logger.warn({ err }, 'Codex warm-up process error');
      finish({ ok: false, reason: 'process_error' });
    });
    proc.on('close', (code) => {
      if (code === 0) {
        finish({ ok: true, reason: 'ok' });
        return;
      }
      logger.warn(
        { exitCode: code, stderr: stderr.trim() || undefined },
        'Codex warm-up command failed',
      );
      finish({ ok: false, reason: `exit_${code ?? 'unknown'}` });
    });
  });
}

export async function runCodexWarmupCycle(
  config: CodexWarmupConfig,
  runtime: CodexWarmupRuntimeOptions = {},
): Promise<CodexWarmupCycleResult> {
  if (!config.enabled) return { status: 'disabled' };
  if (runtime.shouldSkip?.())
    return { status: 'skipped', reason: 'runtime_busy' };

  const nowMs = runtime.nowMs ?? Date.now();
  const nowIso = new Date(nowMs).toISOString();
  const statePath = runtime.statePath ?? DEFAULT_STATE_FILE;
  const state = readWarmupState(statePath);
  const selected = selectWarmupCandidate(config, state, nowMs, {
    ignoreZeroUsageWindow: runtime.ignoreZeroUsageWindow,
    forceAttempt: runtime.forceAttempt,
  });
  if ('reason' in selected) {
    return { status: 'skipped', reason: selected.reason };
  }

  const authPath = getCodexAuthPath(selected.accountIndex);
  if (!authPath || !fs.existsSync(authPath)) {
    return { status: 'skipped', reason: 'missing_auth' };
  }
  const accountDir = path.dirname(authPath);
  const accountState = ensureAccountState(state, selected.accountIndex);
  state.lastAttemptAt = nowIso;
  accountState.lastAttemptAt = nowIso;

  logger.info(
    { account: selected.accountIndex + 1, model: config.model },
    'Starting Codex warm-up prompt',
  );
  const result = await runCodexWarmupCommand(accountDir, config);

  if (result.ok) {
    state.lastWarmupAt = nowIso;
    state.consecutiveFailures = 0;
    delete state.disabledUntil;
    accountState.lastWarmupAt = nowIso;
    accountState.zeroUsageWarmupUntil = selected.zeroUsageWarmupUntil;
    accountState.failures = 0;
    writeWarmupState(statePath, state);
    logger.info(
      { account: selected.accountIndex + 1 },
      'Codex warm-up prompt completed',
    );
    return { status: 'warmed', accountIndex: selected.accountIndex };
  }

  state.consecutiveFailures = (state.consecutiveFailures ?? 0) + 1;
  accountState.lastErrorAt = nowIso;
  accountState.failures = (accountState.failures ?? 0) + 1;
  if (state.consecutiveFailures >= config.maxConsecutiveFailures) {
    state.disabledUntil = new Date(
      nowMs + config.failureCooldownMs,
    ).toISOString();
  }
  writeWarmupState(statePath, state);
  logger.warn(
    {
      account: selected.accountIndex + 1,
      reason: result.reason,
      consecutiveFailures: state.consecutiveFailures,
      disabledUntil: state.disabledUntil,
    },
    'Codex warm-up prompt failed',
  );
  return {
    status: 'failed',
    accountIndex: selected.accountIndex,
    reason: result.reason,
  };
}