pi-statusline/usage/openai-codex.ts

import { Buffer } from "node:buffer";
import { CODEX_BASE_URL } from "../providers/openai-codex/constants";
import type {
  CredentialRankingStrategy,
  UsageAmount,
  UsageFetchContext,
  UsageFetchParams,
  UsageLimit,
  UsageProvider,
  UsageReport,
  UsageWindow,
} from "../usage";
import { isRecord, toNumber } from "../utils";

const CODEX_USAGE_PATH = "wham/usage";
const JWT_AUTH_CLAIM = "https://api.openai.com/auth";
const JWT_PROFILE_CLAIM = "https://api.openai.com/profile";

interface CodexUsageWindowPayload {
  used_percent?: number;
  limit_window_seconds?: number;
  reset_after_seconds?: number;
  reset_at?: number;
}

interface CodexUsageRateLimitPayload {
  allowed?: boolean;
  limit_reached?: boolean;
  primary_window?: CodexUsageWindowPayload | null;
  secondary_window?: CodexUsageWindowPayload | null;
}

interface CodexUsagePayload {
  plan_type?: string;
  rate_limit?: CodexUsageRateLimitPayload | null;
}

interface ParsedUsageWindow {
  usedPercent?: number;
  limitWindowSeconds?: number;
  resetAfterSeconds?: number;
  resetAt?: number;
}

interface ParsedUsage {
  planType?: string;
  allowed?: boolean;
  limitReached?: boolean;
  primary?: ParsedUsageWindow;
  secondary?: ParsedUsageWindow;
  raw: CodexUsagePayload;
}

interface JwtPayload {
  [JWT_AUTH_CLAIM]?: {
    chatgpt_account_id?: string;
  };
  [JWT_PROFILE_CLAIM]?: {
    email?: string;
  };
}

const toBoolean = (value: unknown): boolean | undefined => {
  if (typeof value === "boolean") return value;
  return undefined;
};

function base64UrlDecode(input: string): string {
  const base64 = input.replace(/-/g, "+").replace(/_/g, "/");
  const padLen = (4 - (base64.length % 4)) % 4;
  const padded = base64 + "=".repeat(padLen);
  return Buffer.from(padded, "base64").toString("utf8");
}

function parseJwt(token: string): JwtPayload | null {
  const parts = token.split(".");
  if (parts.length !== 3) return null;
  try {
    const payloadJson = base64UrlDecode(parts[1]);
    return JSON.parse(payloadJson) as JwtPayload;
  } catch {
    return null;
  }
}

function normalizeEmail(email: string | undefined): string | undefined {
  if (!email) return undefined;
  const normalized = email.trim().toLowerCase();
  return normalized || undefined;
}

function extractAccountId(token: string | undefined): string | undefined {
  if (!token) return undefined;
  const payload = parseJwt(token);
  return payload?.[JWT_AUTH_CLAIM]?.chatgpt_account_id ?? undefined;
}

function extractEmail(token: string | undefined): string | undefined {
  if (!token) return undefined;
  const payload = parseJwt(token);
  return normalizeEmail(payload?.[JWT_PROFILE_CLAIM]?.email);
}

function parseUsageWindow(payload: unknown): ParsedUsageWindow | undefined {
  if (!isRecord(payload)) return undefined;
  const usedPercent = toNumber(payload.used_percent);
  const limitWindowSeconds = toNumber(payload.limit_window_seconds);
  const resetAfterSeconds = toNumber(payload.reset_after_seconds);
  const resetAt = toNumber(payload.reset_at);
  if (
    usedPercent === undefined &&
    limitWindowSeconds === undefined &&
    resetAfterSeconds === undefined &&
    resetAt === undefined
  ) {
    return undefined;
  }
  return {
    usedPercent,
    limitWindowSeconds,
    resetAfterSeconds,
    resetAt,
  };
}

function parseUsagePayload(payload: unknown): ParsedUsage | null {
  if (!isRecord(payload)) return null;
  const planType =
    typeof payload.plan_type === "string" ? payload.plan_type : undefined;
  const rateLimit = isRecord(payload.rate_limit)
    ? payload.rate_limit
    : undefined;
  if (!rateLimit) return null;
  const parsed: ParsedUsage = {
    planType,
    allowed: toBoolean(rateLimit.allowed),
    limitReached: toBoolean(rateLimit.limit_reached),
    primary: parseUsageWindow(rateLimit.primary_window),
    secondary: parseUsageWindow(rateLimit.secondary_window),
    raw: payload as CodexUsagePayload,
  };
  if (
    !parsed.primary &&
    !parsed.secondary &&
    parsed.allowed === undefined &&
    parsed.limitReached === undefined
  ) {
    return null;
  }
  return parsed;
}

function normalizeCodexBaseUrl(baseUrl?: string): string {
  const fallback = CODEX_BASE_URL;
  const trimmed = baseUrl?.trim() ? baseUrl.trim() : fallback;
  const base = trimmed.replace(/\/+$/, "");
  const lower = base.toLowerCase();
  if (
    (lower.startsWith("https://chatgpt.com") ||
      lower.startsWith("https://chat.openai.com")) &&
    !lower.includes("/backend-api")
  ) {
    return `${base}/backend-api`;
  }
  return base;
}

function buildCodexUsageUrl(baseUrl: string): string {
  const normalized = baseUrl.endsWith("/") ? baseUrl : `${baseUrl}/`;
  return `${normalized}${CODEX_USAGE_PATH}`;
}

function formatWindowLabel(value: number, unit: "hour" | "day"): string {
  const rounded = Math.round(value);
  const suffix = rounded === 1 ? unit : `${unit}s`;
  return `${rounded} ${suffix}`;
}

function buildWindowLabel(seconds: number): { id: string; label: string } {
  const daySeconds = 86_400;
  if (seconds >= daySeconds) {
    const days = Math.round(seconds / daySeconds);
    return { id: `${days}d`, label: formatWindowLabel(days, "day") };
  }
  const hours = Math.max(1, Math.round(seconds / 3600));
  return { id: `${hours}h`, label: formatWindowLabel(hours, "hour") };
}

function resolveResetTime(
  window: ParsedUsageWindow,
  nowMs: number,
): number | undefined {
  const resetAt = window.resetAt;
  if (resetAt !== undefined) {
    const resetAtMs = resetAt > 1_000_000_000_000 ? resetAt : resetAt * 1000;
    if (Number.isFinite(resetAtMs)) return resetAtMs;
  }
  if (window.resetAfterSeconds !== undefined) {
    return nowMs + window.resetAfterSeconds * 1000;
  }
  return undefined;
}

function buildUsageWindow(
  window: ParsedUsageWindow,
  key: string,
  nowMs: number,
): UsageWindow {
  const resetsAt = resolveResetTime(window, nowMs);
  if (window.limitWindowSeconds !== undefined) {
    const { id, label } = buildWindowLabel(window.limitWindowSeconds);
    const durationMs = window.limitWindowSeconds * 1000;
    return {
      id,
      label,
      durationMs,
      ...(resetsAt !== undefined ? { resetsAt } : {}),
    };
  }
  const fallbackLabel =
    key === "primary" ? "Primary window" : "Secondary window";
  return {
    id: key,
    label: fallbackLabel,
    ...(resetsAt !== undefined ? { resetsAt } : {}),
  };
}

function buildUsageAmount(window: ParsedUsageWindow): UsageAmount {
  const usedPercent = window.usedPercent;
  if (usedPercent === undefined) {
    return { unit: "percent" };
  }
  const clamped = Math.min(Math.max(usedPercent, 0), 100);
  const usedFraction = clamped / 100;
  return {
    used: clamped,
    limit: 100,
    remaining: Math.max(0, 100 - clamped),
    usedFraction,
    remainingFraction: Math.max(0, 1 - usedFraction),
    unit: "percent",
  };
}

function buildUsageStatus(
  usedFraction?: number,
  limitReached?: boolean,
): UsageLimit["status"] {
  if (limitReached) return "exhausted";
  if (usedFraction === undefined) return "unknown";
  if (usedFraction >= 1) return "exhausted";
  if (usedFraction >= 0.9) return "warning";
  return "ok";
}

function buildUsageLimit(args: {
  key: "primary" | "secondary";
  window: ParsedUsageWindow;
  accountId?: string;
  planType?: string;
  limitReached?: boolean;
  nowMs: number;
}): UsageLimit {
  const usageWindow = buildUsageWindow(args.window, args.key, args.nowMs);
  const amount = buildUsageAmount(args.window);
  return {
    id: `openai-codex:${args.key}`,
    label: usageWindow.label,
    scope: {
      provider: "openai-codex",
      accountId: args.accountId,
      tier: args.planType,
      windowId: usageWindow.id,
      shared: true,
    },
    window: usageWindow,
    amount,
    status: buildUsageStatus(amount.usedFraction, args.limitReached),
  };
}

export const openaiCodexUsageProvider: UsageProvider = {
  id: "openai-codex",
  supports(params: UsageFetchParams): boolean {
    return (
      params.provider === "openai-codex" && params.credential.type === "oauth"
    );
  },
  async fetchUsage(
    params: UsageFetchParams,
    ctx: UsageFetchContext,
  ): Promise<UsageReport | null> {
    if (params.provider !== "openai-codex") return null;
    const { credential } = params;
    if (credential.type !== "oauth") return null;

    const accessToken = credential.accessToken;
    if (!accessToken) return null;

    const nowMs = Date.now();
    if (credential.expiresAt !== undefined && credential.expiresAt <= nowMs) {
      ctx.logger?.warn("Codex usage token expired", {
        provider: params.provider,
      });
      return null;
    }

    const baseUrl = normalizeCodexBaseUrl(params.baseUrl);
    const accountId = credential.accountId ?? extractAccountId(accessToken);
    const email = normalizeEmail(credential.email ?? extractEmail(accessToken));

    const headers: Record<string, string> = {
      Authorization: `Bearer ${accessToken}`,
      "User-Agent": "OpenCode-Status-Plugin/1.0",
    };
    if (accountId) {
      headers["ChatGPT-Account-Id"] = accountId;
    }

    const url = buildCodexUsageUrl(baseUrl);
    let payload: unknown;
    try {
      const response = await ctx.fetch(url, { headers, signal: params.signal });
      if (!response.ok) {
        ctx.logger?.warn("Codex usage request failed", {
          status: response.status,
          provider: params.provider,
        });
        return null;
      }
      payload = await response.json();
    } catch (error) {
      ctx.logger?.warn("Codex usage request error", {
        provider: params.provider,
        error: String(error),
      });
      return null;
    }

    const parsed = parseUsagePayload(payload);
    const planType =
      parsed?.planType ??
      (isRecord(payload) && typeof payload.plan_type === "string"
        ? payload.plan_type
        : undefined);

    const limits: UsageLimit[] = [];
    if (parsed?.primary) {
      limits.push(
        buildUsageLimit({
          key: "primary",
          window: parsed.primary,
          accountId,
          planType,
          limitReached: parsed.limitReached,
          nowMs,
        }),
      );
    }
    if (parsed?.secondary) {
      limits.push(
        buildUsageLimit({
          key: "secondary",
          window: parsed.secondary,
          accountId,
          planType,
          limitReached: parsed.limitReached,
          nowMs,
        }),
      );
    }

    const report: UsageReport = {
      provider: "openai-codex",
      fetchedAt: nowMs,
      limits,
      metadata: {
        planType,
        allowed: parsed?.allowed,
        limitReached: parsed?.limitReached,
        email,
        accountId,
      },
      raw: parsed?.raw ?? payload,
    };

    return report;
  },
};

const FIVE_HOUR_MS = 5 * 60 * 60 * 1000;

export const codexRankingStrategy: CredentialRankingStrategy = {
  findWindowLimits(report) {
    const findLimit = (
      key: "primary" | "secondary",
    ): UsageLimit | undefined => {
      const direct = report.limits.find((l) => l.id === `openai-codex:${key}`);
      if (direct) return direct;
      const byId = report.limits.find((l) => l.id.toLowerCase().includes(key));
      if (byId) return byId;
      const windowId = key === "secondary" ? "7d" : "1h";
      return report.limits.find(
        (l) => l.scope.windowId?.toLowerCase() === windowId,
      );
    };
    return { primary: findLimit("primary"), secondary: findLimit("secondary") };
  },
  windowDefaults: {
    primaryMs: 60 * 60 * 1000,
    secondaryMs: 7 * 24 * 60 * 60 * 1000,
  },
  hasPriorityBoost(primary) {
    if (!primary) return false;
    const windowId = primary.scope.windowId?.toLowerCase();
    const durationMs = primary.window?.durationMs;
    const isFiveHourWindow =
      windowId === "5h" ||
      (typeof durationMs === "number" &&
        Number.isFinite(durationMs) &&
        Math.abs(durationMs - FIVE_HOUR_MS) <= 60_000);
    if (!isFiveHourWindow) return false;
    const usedFraction = primary.amount.usedFraction;
    return (
      typeof usedFraction === "number" &&
      Number.isFinite(usedFraction) &&
      usedFraction === 0
    );
  },
};