@bitkyc08/opencodex 2.1.8 → 2.1.10

package/src/codex-quota.ts

@@ -0,0 +1,130 @@
+export type StoredAccountQuota = {
+  weeklyPercent?: number;
+  fiveHourPercent?: number;
+  monthlyPercent?: number;
+  weeklyResetAt?: number;
+  fiveHourResetAt?: number;
+  monthlyResetAt?: number;
+  updatedAt: number;
+};
+
+export type WhamUsageResponse = {
+  email?: string | null;
+  plan_type?: string | null;
+  rate_limit?: {
+    primary_window?: { used_percent?: number; reset_at?: number };
+    secondary_window?: { used_percent?: number; reset_at?: number };
+    tertiary_window?: { used_percent?: number; reset_at?: number };
+  };
+};
+
+const accountQuota = new Map<string, StoredAccountQuota>();
+
+export const CODEX_UNKNOWN_USAGE_SCORE = 100;
+
+export function normalizeUsagePercent(value: unknown): number | undefined {
+  const numeric = typeof value === "number"
+    ? value
+    : typeof value === "string" && value.trim() !== ""
+      ? Number(value)
+      : undefined;
+  if (typeof numeric !== "number" || !Number.isFinite(numeric)) return undefined;
+  return Math.max(0, Math.min(100, numeric));
+}
+
+function normalizeResetAt(value: unknown): number | undefined {
+  const numeric = typeof value === "number"
+    ? value
+    : typeof value === "string" && value.trim() !== ""
+      ? Number(value)
+      : undefined;
+  if (typeof numeric !== "number" || !Number.isFinite(numeric) || numeric < 0) return undefined;
+  return numeric;
+}
+
+function hasKnownQuotaValue(quota: Omit<StoredAccountQuota, "updatedAt">): boolean {
+  return [quota.weeklyPercent, quota.fiveHourPercent, quota.monthlyPercent]
+    .some(value => typeof value === "number" && Number.isFinite(value));
+}
+
+export function updateAccountQuota(
+  accountId: string,
+  weekly: unknown,
+  fiveHour: unknown,
+  weeklyResetAt?: unknown,
+  fiveHourResetAt?: unknown,
+  monthly?: unknown,
+  monthlyResetAt?: unknown,
+): void {
+  const existing = accountQuota.get(accountId);
+  const nextWeekly = normalizeUsagePercent(weekly);
+  const nextFiveHour = normalizeUsagePercent(fiveHour);
+  const nextMonthly = normalizeUsagePercent(monthly);
+  if (nextWeekly === undefined && nextFiveHour === undefined && nextMonthly === undefined) return;
+
+  const quota: StoredAccountQuota = {
+    ...(existing?.weeklyPercent !== undefined ? { weeklyPercent: existing.weeklyPercent } : {}),
+    ...(existing?.fiveHourPercent !== undefined ? { fiveHourPercent: existing.fiveHourPercent } : {}),
+    ...(existing?.monthlyPercent !== undefined ? { monthlyPercent: existing.monthlyPercent } : {}),
+    ...(existing?.weeklyResetAt !== undefined ? { weeklyResetAt: existing.weeklyResetAt } : {}),
+    ...(existing?.fiveHourResetAt !== undefined ? { fiveHourResetAt: existing.fiveHourResetAt } : {}),
+    ...(existing?.monthlyResetAt !== undefined ? { monthlyResetAt: existing.monthlyResetAt } : {}),
+    updatedAt: Date.now(),
+  };
+
+  const nextWeeklyResetAt = normalizeResetAt(weeklyResetAt);
+  const nextFiveHourResetAt = normalizeResetAt(fiveHourResetAt);
+  const nextMonthlyResetAt = normalizeResetAt(monthlyResetAt);
+  if (nextWeekly !== undefined) {
+    quota.weeklyPercent = nextWeekly;
+    if (nextWeeklyResetAt !== undefined) quota.weeklyResetAt = nextWeeklyResetAt;
+  }
+  if (nextFiveHour !== undefined) {
+    quota.fiveHourPercent = nextFiveHour;
+    if (nextFiveHourResetAt !== undefined) quota.fiveHourResetAt = nextFiveHourResetAt;
+  }
+  if (nextMonthly !== undefined) {
+    quota.monthlyPercent = nextMonthly;
+    if (nextMonthlyResetAt !== undefined) quota.monthlyResetAt = nextMonthlyResetAt;
+  }
+
+  accountQuota.set(accountId, quota);
+}
+
+export function getAccountQuota(accountId: string): StoredAccountQuota | null {
+  return accountQuota.get(accountId) ?? null;
+}
+
+export function listAccountQuotas(): IterableIterator<[string, StoredAccountQuota]> {
+  return accountQuota.entries();
+}
+
+export function clearAccountQuota(accountId?: string): void {
+  if (accountId) accountQuota.delete(accountId);
+  else accountQuota.clear();
+}
+
+export function parseUsageQuota(data: WhamUsageResponse): Omit<StoredAccountQuota, "updatedAt"> | null {
+  if (!data.rate_limit) return null;
+  const quota: Omit<StoredAccountQuota, "updatedAt"> = {};
+  const weeklyPercent = normalizeUsagePercent(data.rate_limit.secondary_window?.used_percent);
+  const fiveHourPercent = normalizeUsagePercent(data.rate_limit.primary_window?.used_percent);
+  const monthlyPercent = normalizeUsagePercent(data.rate_limit.tertiary_window?.used_percent);
+  const weeklyResetAt = normalizeResetAt(data.rate_limit.secondary_window?.reset_at);
+  const fiveHourResetAt = normalizeResetAt(data.rate_limit.primary_window?.reset_at);
+  const monthlyResetAt = normalizeResetAt(data.rate_limit.tertiary_window?.reset_at);
+  if (weeklyPercent !== undefined) {
+    quota.weeklyPercent = weeklyPercent;
+    if (weeklyResetAt !== undefined) quota.weeklyResetAt = weeklyResetAt;
+  }
+  if (fiveHourPercent !== undefined) {
+    quota.fiveHourPercent = fiveHourPercent;
+    if (fiveHourResetAt !== undefined) quota.fiveHourResetAt = fiveHourResetAt;
+  }
+  if (monthlyPercent !== undefined) {
+    quota.monthlyPercent = monthlyPercent;
+    if (monthlyResetAt !== undefined) quota.monthlyResetAt = monthlyResetAt;
+  }
+
+  return hasKnownQuotaValue(quota) ? quota : null;
+}

package/src/codex-routing.ts

@@ -0,0 +1,382 @@
+import { saveConfig } from "./config";
+import { isCodexAccountGenerationLive, readCodexAccountRecord } from "./codex-account-store";
+import { codexAccountLogLabel } from "./codex-account-label";
+import { isCodexAccountUsable } from "./codex-account-usability";
+import { isAccountNeedsReauth, markAccountNeedsReauth } from "./codex-account-runtime-state";
+import { CODEX_UNKNOWN_USAGE_SCORE, getAccountQuota } from "./codex-quota";
+import type { OcxConfig } from "./types";
+
+type ThreadAffinityEntry = {
+  accountId: string;
+  generation: number;
+  createdAt: number;
+  lastUsedAt: number;
+};
+
+export type CodexThreadResolution =
+  | { status: "selected"; accountId: string }
+  | { status: "none" }
+  | { status: "expired"; accountId: string };
+
+const threadAccountMap = new Map<string, ThreadAffinityEntry>();
+type CodexUpstreamHealth = {
+  consecutiveFailures: number;
+  lastFailureStatus?: number;
+  lastFailureAt?: number;
+  cooldownUntil?: number;
+};
+
+const CODEX_DEFAULT_QUOTA_COOLDOWN_MS = 60_000;
+const CODEX_MAX_QUOTA_COOLDOWN_MS = 24 * 60 * 60_000;
+export const CODEX_FAILURE_WINDOW_MS = 5 * 60_000;
+export const CODEX_THREAD_AFFINITY_IDLE_TTL_MS = 24 * 60 * 60_000;
+export const CODEX_THREAD_AFFINITY_MAX_ENTRIES = 2048;
+
+const upstreamHealth = new Map<string, CodexUpstreamHealth>();
+
+export type CodexUpstreamOutcome = number | "connect_error" | "timeout";
+export type CodexUpstreamOutcomeClass = "success" | "credential" | "quota" | "transient" | "caller" | "unknown";
+export type CodexUpstreamOutcomeMeta = {
+  retryAfter?: string | null;
+  resetAt?: unknown | unknown[];
+  now?: number;
+};
+
+function hasConfiguredPoolAccount(config: OcxConfig, accountId: string): boolean {
+  return (config.codexAccounts ?? []).some(account => !account.isMain && account.id === accountId);
+}
+
+export function clearThreadAccountMap(): void {
+  threadAccountMap.clear();
+}
+
+export function clearThreadAccountMapForAccount(accountId: string): void {
+  for (const [threadId, entry] of threadAccountMap) {
+    if (entry.accountId === accountId) threadAccountMap.delete(threadId);
+  }
+}
+
+export function clearCodexUpstreamHealth(): void {
+  upstreamHealth.clear();
+}
+
+export function clearCodexUpstreamHealthForAccount(accountId: string): void {
+  upstreamHealth.delete(accountId);
+}
+
+export function getCodexUpstreamHealth(
+  accountId: string,
+): CodexUpstreamHealth | null {
+  return upstreamHealth.get(accountId) ?? null;
+}
+
+export function computeCodexUsageScore(quota: {
+  weeklyPercent?: number;
+  fiveHourPercent?: number;
+  monthlyPercent?: number;
+} | null): number {
+  if (!quota) return CODEX_UNKNOWN_USAGE_SCORE;
+  const values = [quota.weeklyPercent, quota.fiveHourPercent, quota.monthlyPercent]
+    .filter((value): value is number => typeof value === "number" && Number.isFinite(value));
+  return values.length > 0 ? Math.max(...values) : CODEX_UNKNOWN_USAGE_SCORE;
+}
+
+export function classifyCodexUpstreamOutcome(outcome: CodexUpstreamOutcome): CodexUpstreamOutcomeClass {
+  if (outcome === "connect_error" || outcome === "timeout") return "transient";
+  if (!Number.isFinite(outcome)) return "unknown";
+  if (outcome >= 200 && outcome < 300) return "success";
+  if (outcome === 401 || outcome === 403) return "credential";
+  if (outcome === 429) return "quota";
+  if (outcome >= 400 && outcome < 500) return "caller";
+  if (outcome >= 500 && outcome < 600) return "transient";
+  return "unknown";
+}
+
+function clampCooldownMs(ms: number): number {
+  return Math.min(Math.max(ms, 1), CODEX_MAX_QUOTA_COOLDOWN_MS);
+}
+
+export function parseRetryAfterMs(value: string | null | undefined, now = Date.now()): number | undefined {
+  const text = value?.trim();
+  if (!text) return undefined;
+  if (/^\d+(?:\.\d+)?$/.test(text)) {
+    const seconds = Number(text);
+    if (Number.isFinite(seconds) && seconds > 0) return clampCooldownMs(Math.ceil(seconds * 1000));
+  }
+  const timestamp = Date.parse(text);
+  if (!Number.isFinite(timestamp)) return undefined;
+  const delay = timestamp - now;
+  return delay > 0 ? clampCooldownMs(delay) : undefined;
+}
+
+function resetTimestampMs(value: unknown): number | undefined {
+  const numeric = typeof value === "number"
+    ? value
+    : typeof value === "string" && value.trim() !== ""
+      ? Number(value)
+      : undefined;
+  if (typeof numeric !== "number" || !Number.isFinite(numeric) || numeric <= 0) return undefined;
+  return numeric < 1_000_000_000_000 ? numeric * 1000 : numeric;
+}
+
+export function parseResetCooldownMs(resetAt: unknown | unknown[] | undefined, now = Date.now()): number | undefined {
+  const values = Array.isArray(resetAt) ? resetAt : [resetAt];
+  let best: number | undefined;
+  for (const value of values) {
+    const timestamp = resetTimestampMs(value);
+    if (timestamp === undefined) continue;
+    const delay = timestamp - now;
+    if (delay <= 0) continue;
+    const clamped = clampCooldownMs(delay);
+    if (best === undefined || clamped < best) best = clamped;
+  }
+  return best;
+}
+
+export function computeQuotaCooldownUntil(meta: CodexUpstreamOutcomeMeta = {}): number {
+  const now = meta.now ?? Date.now();
+  const retryAfterMs = parseRetryAfterMs(meta.retryAfter, now);
+  const resetCooldownMs = retryAfterMs === undefined ? parseResetCooldownMs(meta.resetAt, now) : undefined;
+  return now + (retryAfterMs ?? resetCooldownMs ?? CODEX_DEFAULT_QUOTA_COOLDOWN_MS);
+}
+
+export function getCodexAccountCooldownUntil(accountId: string, now = Date.now()): number | null {
+  const cooldownUntil = upstreamHealth.get(accountId)?.cooldownUntil;
+  return typeof cooldownUntil === "number" && Number.isFinite(cooldownUntil) && cooldownUntil > now ? cooldownUntil : null;
+}
+
+export function isCodexAccountInCooldown(accountId: string, now = Date.now()): boolean {
+  return getCodexAccountCooldownUntil(accountId, now) !== null;
+}
+
+function isCodexAccountSelectable(config: OcxConfig, accountId: string, now: number): boolean {
+  return !isCodexAccountInCooldown(accountId, now) && isCodexAccountUsable(config, accountId);
+}
+
+function isThreadAffinityExpired(entry: ThreadAffinityEntry, now: number): boolean {
+  return now - entry.lastUsedAt > CODEX_THREAD_AFFINITY_IDLE_TTL_MS;
+}
+
+function isThreadAffinityGenerationLive(entry: ThreadAffinityEntry): boolean {
+  return isCodexAccountGenerationLive(entry.accountId, entry.generation);
+}
+
+function pruneExpiredThreadAffinities(now: number): void {
+  for (const [threadId, entry] of threadAccountMap) {
+    if (isThreadAffinityExpired(entry, now)) threadAccountMap.delete(threadId);
+  }
+}
+
+function pruneLruThreadAffinities(): void {
+  while (threadAccountMap.size > CODEX_THREAD_AFFINITY_MAX_ENTRIES) {
+    let oldestThreadId: string | null = null;
+    let oldestLastUsedAt = Number.POSITIVE_INFINITY;
+    for (const [threadId, entry] of threadAccountMap) {
+      if (entry.lastUsedAt < oldestLastUsedAt) {
+        oldestThreadId = threadId;
+        oldestLastUsedAt = entry.lastUsedAt;
+      }
+    }
+    if (!oldestThreadId) return;
+    threadAccountMap.delete(oldestThreadId);
+  }
+}
+
+function bindThreadAffinity(threadId: string, accountId: string, now: number): void {
+  const record = readCodexAccountRecord(accountId);
+  if (!record?.credential || record.deletedAt != null) return;
+  pruneExpiredThreadAffinities(now);
+  const previous = threadAccountMap.get(threadId);
+  threadAccountMap.set(threadId, {
+    accountId,
+    generation: record.generation,
+    createdAt: previous?.createdAt ?? now,
+    lastUsedAt: now,
+  });
+  pruneLruThreadAffinities();
+}
+
+function getEligiblePoolAccounts(config: OcxConfig, excludeId?: string, now = Date.now()): string[] {
+  return (config.codexAccounts ?? [])
+    .filter(account => !account.isMain && account.id !== excludeId && !isAccountNeedsReauth(account.id))
+    .filter(account => !isCodexAccountInCooldown(account.id, now))
+    .filter(account => isCodexAccountUsable(config, account.id))
+    .map(account => account.id);
+}
+
+function pickLowerUsageAccount(config: OcxConfig, active: string, activeUsage: number, now: number): string {
+  let best = active;
+  let bestUsage = activeUsage;
+  for (const id of getEligiblePoolAccounts(config, active, now)) {
+    const usage = computeCodexUsageScore(getAccountQuota(id));
+    if (usage < bestUsage) {
+      best = id;
+      bestUsage = usage;
+    }
+  }
+  return best;
+}
+
+export function pickLowestUsageCodexAccount(config: OcxConfig, excludeId?: string, now = Date.now()): string | null {
+  let best: string | null = null;
+  let bestUsage = Number.POSITIVE_INFINITY;
+  for (const id of getEligiblePoolAccounts(config, excludeId, now)) {
+    const usage = computeCodexUsageScore(getAccountQuota(id));
+    if (usage < bestUsage) {
+      best = id;
+      bestUsage = usage;
+    }
+  }
+  return best;
+}
+
+function setActiveCodexAccount(config: OcxConfig, accountId: string): void {
+  if (config.activeCodexAccountId === accountId) return;
+  config.activeCodexAccountId = accountId;
+  saveConfig(config);
+}
+
+function applyQuotaAutoSwitch(config: OcxConfig, active: string, now: number): string {
+  const threshold = config.autoSwitchThreshold ?? 80;
+  if (threshold <= 0) return active;
+  const quota = getAccountQuota(active);
+  const activeUsage = computeCodexUsageScore(quota);
+  if (activeUsage < threshold) return active;
+  const best = pickLowerUsageAccount(config, active, activeUsage, now);
+  if (best !== active) setActiveCodexAccount(config, best);
+  return best;
+}
+
+function shouldFailover(config: OcxConfig, accountId: string, now: number): boolean {
+  const threshold = config.upstreamFailoverThreshold ?? 3;
+  if (threshold <= 0) return false;
+  const health = upstreamHealth.get(accountId);
+  if (health?.lastFailureAt && now - health.lastFailureAt > CODEX_FAILURE_WINDOW_MS) return false;
+  return !!health && health.consecutiveFailures >= threshold;
+}
+
+function applyFailureFailover(config: OcxConfig, active: string, now: number): string {
+  if (!shouldFailover(config, active, now)) return active;
+  const best = pickLowestUsageCodexAccount(config, active, now);
+  if (best) {
+    setActiveCodexAccount(config, best);
+    return best;
+  }
+  return active;
+}
+
+export function resolveCodexAccountForThread(
+  threadId: string | null,
+  config: OcxConfig,
+  now = Date.now(),
+): string | null {
+  const resolution = resolveCodexAccountForThreadDetailed(threadId, config, now);
+  return resolution.status === "selected" ? resolution.accountId : null;
+}
+
+export function resolveCodexAccountForThreadDetailed(
+  threadId: string | null,
+  config: OcxConfig,
+  now = Date.now(),
+): CodexThreadResolution {
+  if (threadId && threadAccountMap.has(threadId)) {
+    const entry = threadAccountMap.get(threadId)!;
+    if (isThreadAffinityExpired(entry, now)) {
+      threadAccountMap.delete(threadId);
+      return { status: "expired", accountId: entry.accountId };
+    }
+    if (
+      isThreadAffinityGenerationLive(entry)
+      && isCodexAccountSelectable(config, entry.accountId, now)
+    ) {
+      entry.lastUsedAt = now;
+      return { status: "selected", accountId: entry.accountId };
+    }
+    threadAccountMap.delete(threadId);
+  }
+  let active = config.activeCodexAccountId;
+  if (!active) return { status: "none" };
+  if (!isCodexAccountSelectable(config, active, now)) {
+    const fallback = pickLowestUsageCodexAccount(config, active, now);
+    if (fallback) {
+      setActiveCodexAccount(config, fallback);
+      active = fallback;
+    } else if (hasConfiguredPoolAccount(config, active)) {
+      return { status: "selected", accountId: active };
+    } else {
+      return { status: "none" };
+    }
+  }
+  active = applyQuotaAutoSwitch(config, active, now);
+  active = applyFailureFailover(config, active, now);
+  if (!isCodexAccountUsable(config, active)) {
+    return hasConfiguredPoolAccount(config, active) ? { status: "selected", accountId: active } : { status: "none" };
+  }
+  if (isCodexAccountInCooldown(active, now)) {
+    return hasConfiguredPoolAccount(config, active) ? { status: "selected", accountId: active } : { status: "none" };
+  }
+  if (threadId) bindThreadAffinity(threadId, active, now);
+  return { status: "selected", accountId: active };
+}
+
+export function recordCodexUpstreamOutcome(
+  config: OcxConfig,
+  accountId: string | null,
+  outcome: CodexUpstreamOutcome,
+  meta: CodexUpstreamOutcomeMeta = {},
+): void {
+  if (!accountId) return;
+  const now = meta.now ?? Date.now();
+  const outcomeClass = classifyCodexUpstreamOutcome(outcome);
+  if (outcomeClass === "success") {
+    const cooldownUntil = getCodexAccountCooldownUntil(accountId, now);
+    if (cooldownUntil) upstreamHealth.set(accountId, { consecutiveFailures: 0, cooldownUntil });
+    else upstreamHealth.delete(accountId);
+    return;
+  }
+  if (outcomeClass === "caller") return;
+
+  const lastFailureStatus = typeof outcome === "number" ? outcome : 0;
+  if (outcomeClass === "credential") {
+    upstreamHealth.set(accountId, {
+      consecutiveFailures: 1,
+      lastFailureStatus,
+      lastFailureAt: now,
+    });
+    markAccountNeedsReauth(accountId);
+    clearThreadAccountMapForAccount(accountId);
+    return;
+  }
+
+  if (outcomeClass === "quota") {
+    upstreamHealth.set(accountId, {
+      consecutiveFailures: 0,
+      lastFailureStatus,
+      lastFailureAt: now,
+      cooldownUntil: computeQuotaCooldownUntil(meta),
+    });
+    clearThreadAccountMapForAccount(accountId);
+    if (config.activeCodexAccountId === accountId) {
+      const fallback = pickLowestUsageCodexAccount(config, accountId, now);
+      if (fallback) setActiveCodexAccount(config, fallback);
+    }
+    return;
+  }
+
+  const current = upstreamHealth.get(accountId);
+  const stale = current?.lastFailureAt ? now - current.lastFailureAt > CODEX_FAILURE_WINDOW_MS : false;
+  const cooldownUntil = getCodexAccountCooldownUntil(accountId, now) ?? undefined;
+  upstreamHealth.set(accountId, {
+    consecutiveFailures: stale ? 1 : (current?.consecutiveFailures ?? 0) + 1,
+    lastFailureStatus,
+    lastFailureAt: now,
+    ...(cooldownUntil ? { cooldownUntil } : {}),
+  });
+  if (config.activeCodexAccountId === accountId) applyFailureFailover(config, accountId, now);
+}
+
+export function formatCodexProviderForLog(providerName: string, accountId: string | null, config: OcxConfig): string {
+  if (!accountId) return providerName;
+  const account = (config.codexAccounts ?? []).find(a => !a.isMain && a.id === accountId);
+  return account ? `${providerName}-${codexAccountLogLabel(account)}` : providerName;
+}

package/src/codex-websocket-registry.ts

@@ -0,0 +1,57 @@
+import type { ServerWebSocket } from "bun";
+import type { WsData } from "./ws-bridge";
+
+const socketsByAccount = new Map<string, Set<ServerWebSocket<WsData>>>();
+
+function trackedAccountId(ws: ServerWebSocket<WsData>): string | null {
+  const ctx = ws.data.authContext;
+  return ctx?.kind === "pool" ? ctx.accountId : null;
+}
+
+export function registerCodexWebSocket(ws: ServerWebSocket<WsData>): void {
+  const accountId = trackedAccountId(ws);
+  if (!accountId) return;
+  let sockets = socketsByAccount.get(accountId);
+  if (!sockets) {
+    sockets = new Set();
+    socketsByAccount.set(accountId, sockets);
+  }
+  sockets.add(ws);
+}
+
+export function unregisterCodexWebSocket(ws: ServerWebSocket<WsData>): void {
+  const accountId = trackedAccountId(ws);
+  if (!accountId) return;
+  const sockets = socketsByAccount.get(accountId);
+  if (!sockets) return;
+  sockets.delete(ws);
+  if (sockets.size === 0) socketsByAccount.delete(accountId);
+}
+
+export function invalidateCodexWebSocketsForAccount(accountId: string): number {
+  const sockets = socketsByAccount.get(accountId);
+  if (!sockets) return 0;
+  const snapshot = [...sockets];
+  socketsByAccount.delete(accountId);
+  for (const ws of snapshot) {
+    try {
+      ws.data.cancel?.();
+    } catch {
+      /* ignore cancel callbacks during invalidation */
+    }
+    try {
+      ws.close(4001, "Codex account invalidated");
+    } catch {
+      /* socket may already be closing */
+    }
+  }
+  return snapshot.length;
+}
+
+export function getTrackedCodexWebSocketCountForAccount(accountId: string): number {
+  return socketsByAccount.get(accountId)?.size ?? 0;
+}
+
+export function clearCodexWebSocketRegistry(): void {
+  socketsByAccount.clear();
+}