EJClaw/src/message-runtime-gating.ts

import {
  CLAUDE_EXHAUSTION_THRESHOLD_PCT,
  getClaudeUsageExhaustion,
} from './claude-usage.js';
import {
  CODEX_EXHAUSTION_THRESHOLD_PCT,
  getCodexUsageExhaustion,
} from './codex-usage-collector.js';
import { getAllCodexAccounts } from './codex-token-rotation.js';
import { logger } from './logger.js';
import {
  handleSessionCommand,
  type SessionCommandDeps,
} from './session-commands.js';
import type { AgentType, NewMessage, RegisteredGroup } from './types.js';

/**
 * Test seam — overridable from unit tests so we don't need a real disk
 * cache or rotation state to exercise the exhaustion gate.
 */
export const exhaustionChecks = {
  claude: getClaudeUsageExhaustion,
  codex: getCodexUsageExhaustion,
};

/** Render an ISO timestamp as "약 N분 뒤" / "약 Nh Nm 뒤" relative to now. */
function formatResetCountdown(iso: string | null): string | null {
  if (!iso) return null;
  const target = Date.parse(iso);
  if (!Number.isFinite(target)) return null;
  const diffMs = target - Date.now();
  if (diffMs <= 0) return '곧';
  const totalMinutes = Math.ceil(diffMs / 60_000);
  if (totalMinutes < 60) return `약 ${totalMinutes}분 뒤`;
  const hours = Math.floor(totalMinutes / 60);
  const minutes = totalMinutes % 60;
  if (hours < 24) {
    return minutes > 0
      ? `약 ${hours}시간 ${minutes}분 뒤`
      : `약 ${hours}시간 뒤`;
  }
  const days = Math.floor(hours / 24);
  const remH = hours % 24;
  return remH > 0 ? `약 ${days}일 ${remH}시간 뒤` : `약 ${days}일 뒤`;
}

function exhaustionMessageFor(
  agentType: AgentType,
  nextResetAt: string | null,
): string {
  const label = agentType === 'codex' ? 'Codex' : 'Claude';
  const pct =
    agentType === 'codex'
      ? CODEX_EXHAUSTION_THRESHOLD_PCT
      : CLAUDE_EXHAUSTION_THRESHOLD_PCT;
  const countdown = formatResetCountdown(nextResetAt);
  const head = `토큰 부족으로 종료됨 (${label} 모든 계정 사용량 ${pct}% 초과)`;
  return countdown ? `${head} — ${countdown} 초기화 예정` : head;
}

// ── Usage-window alignment gap ──
// Bot primer fires at 08/13/18/23 KST to anchor 5h windows. The 04:00–08:00
// KST gap before the day's first primer must be free of inference calls,
// otherwise an off-cycle window starts and the daily reset alignment drifts.
// During the gap, send an auto-reply instead of running the model.
// Urgent bypass: any message containing an @-mention (bot's trigger pattern)
// is processed anyway — user explicitly tagged the bot, accept the drift.
const GAP_START_HOUR_KST = 4;
const GAP_END_HOUR_KST = 8;
const PRIMER_HOURS_KST = [8, 13, 18, 23];
const KST_OFFSET_MS = 9 * 60 * 60 * 1000;
const CODEX_ALIGNMENT_FRESH_USAGE_PCT = 1;
const PRIMER_GRACE_MS = 2 * 60 * 1000;

export function isInUsageAlignmentGap(nowMs: number = Date.now()): boolean {
  const kstHour = new Date(nowMs + KST_OFFSET_MS).getUTCHours();
  return kstHour >= GAP_START_HOUR_KST && kstHour < GAP_END_HOUR_KST;
}

function msUntilNextPrimerSlotKST(nowMs: number): number {
  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;
}

function formatPrimerHourKST(nowMs: number): string {
  const nextMs = nowMs + msUntilNextPrimerSlotKST(nowMs);
  const kst = new Date(nextMs + KST_OFFSET_MS);
  return `${String(kst.getUTCHours()).padStart(2, '0')}:00 KST`;
}

export function shouldHoldFreshCodexUntilPrimer(
  nowMs: number = Date.now(),
): boolean {
  const accounts = getAllCodexAccounts();
  const active = accounts.find((a) => a.isActive) ?? accounts[0];
  if (!active || active.isRateLimited) return false;
  const h5 = active.cachedUsagePct;
  if (h5 == null || h5 < 0 || h5 > CODEX_ALIGNMENT_FRESH_USAGE_PCT) {
    return false;
  }

  const elapsedFromPreviousSlot =
    (5 * 60 * 60 * 1000 - msUntilNextPrimerSlotKST(nowMs)) %
    (5 * 60 * 60 * 1000);
  return elapsedFromPreviousSlot > PRIMER_GRACE_MS;
}

const GAP_AUTO_REPLY =
  '사용량 정렬 시간(04~08 KST)입니다. 08시 이후 다시 보내주세요. 긴급한 경우 @태그하면 처리됩니다.';

const CODEX_ALIGNMENT_AUTO_REPLY = (nowMs: number) =>
  `Codex 사용량 정렬 대기 중입니다. 다음 프라이머 시각(${formatPrimerHourKST(
    nowMs,
  )}) 이후 다시 보내주세요. 긴급한 경우 @태그하면 처리됩니다.`;

export async function handleQueuedRunGates(args: {
  chatJid: string;
  group: RegisteredGroup;
  runId: string;
  missedMessages: NewMessage[];
  triggerPattern: RegExp;
  timezone: string;
  hasImplicitContinuationWindow: (
    chatJid: string,
    messages: NewMessage[],
  ) => boolean;
  sessionCommandDeps: SessionCommandDeps;
}): Promise<{ handled: true; success: boolean } | { handled: false }> {
  const cmdResult = await handleSessionCommand({
    missedMessages: args.missedMessages,
    isMainGroup: args.group.isMain === true,
    groupName: args.group.name,
    runId: args.runId,
    triggerPattern: args.triggerPattern,
    timezone: args.timezone,
    deps: args.sessionCommandDeps,
  });
  if (cmdResult.handled) {
    return cmdResult;
  }

  // ── Usage-window alignment gap gate ──
  if (isInUsageAlignmentGap()) {
    const hasUrgentMention = args.missedMessages.some((m) =>
      args.triggerPattern.test(m.content),
    );
    if (!hasUrgentMention) {
      logger.info(
        {
          chatJid: args.chatJid,
          groupName: args.group.name,
          runId: args.runId,
        },
        'Blocking new turn — usage alignment gap (04-08 KST)',
      );
      try {
        await args.sessionCommandDeps.sendMessage(GAP_AUTO_REPLY);
      } catch (err) {
        logger.warn({ err }, 'Failed to send alignment gap notice');
      }
      return { handled: true, success: true };
    }
  }

  const agentType: AgentType = args.group.agentType ?? 'claude-code';
  if (agentType === 'codex' && shouldHoldFreshCodexUntilPrimer()) {
    const hasUrgentMention = args.missedMessages.some((m) =>
      args.triggerPattern.test(m.content),
    );
    if (!hasUrgentMention) {
      logger.info(
        {
          chatJid: args.chatJid,
          groupName: args.group.name,
          runId: args.runId,
        },
        'Blocking new Codex turn — waiting for fixed primer slot',
      );
      try {
        await args.sessionCommandDeps.sendMessage(
          CODEX_ALIGNMENT_AUTO_REPLY(Date.now()),
        );
      } catch (err) {
        logger.warn({ err }, 'Failed to send Codex alignment notice');
      }
      return { handled: true, success: true };
    }
  }

  // ── Usage-exhaustion gate ──
  // Block new turns when the agent provider that would handle this room has
  // every account at ≥95% on either window (or rate-limited). In-flight
  // runs are not affected; this only stops *new* turns from starting.
  const info =
    agentType === 'codex'
      ? exhaustionChecks.codex()
      : exhaustionChecks.claude();
  if (info.exhausted) {
    const text = exhaustionMessageFor(agentType, info.nextResetAt);
    logger.warn(
      {
        chatJid: args.chatJid,
        groupName: args.group.name,
        agentType,
        runId: args.runId,
        nextResetAt: info.nextResetAt,
      },
      'Blocking new turn — all accounts exhausted',
    );
    try {
      await args.sessionCommandDeps.sendMessage(text);
    } catch (err) {
      logger.warn({ err }, 'Failed to send exhaustion notice');
    }
    return { handled: true, success: true };
  }

  return { handled: false };
}