@monotykamary/pi-better-openai 0.1.23

package/src/usage-controller.ts

@@ -0,0 +1,304 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { ResolvedConfig } from "./config.ts";
+import { maskIdentifier, sanitizeDiagnosticError } from "./format.ts";
+import {
+  AUTH_FILE,
+  type UsageSnapshot,
+  formatResetCountdown,
+  formatUsageSnapshot,
+  parseUsageSnapshot,
+  readCodexAuth,
+  requestCodexUsage,
+  usageScopeForModel,
+} from "./usage.ts";
+import { currentModelKey } from "./fast-controller.ts";
+
+export function isOpenAISubscriptionModel(
+  ctx: ExtensionContext,
+  cfg: ResolvedConfig,
+  isUsingOAuth?: boolean,
+): boolean {
+  const model = ctx.model;
+  if (!model || (model.provider !== "openai" && model.provider !== "openai-codex")) return false;
+  return (
+    !cfg.usage.showOnlyOnSubscriptionModels ||
+    (isUsingOAuth ?? ctx.modelRegistry.isUsingOAuth(model))
+  );
+}
+
+const STALE_EXTENSION_CONTEXT_MESSAGE = "This extension ctx is stale";
+
+function isStaleExtensionContextError(error: unknown): boolean {
+  return error instanceof Error && error.message.includes(STALE_EXTENSION_CONTEXT_MESSAGE);
+}
+
+type UsageRefreshOptions = { notify?: boolean; force?: boolean };
+
+type QueuedUsageRefresh = {
+  ctx: ExtensionContext;
+  generation: number;
+  modelId?: string;
+  notify?: boolean;
+  force?: boolean;
+};
+
+export class UsageController {
+  private usageSnapshot: UsageSnapshot | undefined;
+  private usageUpdatedAt: number | undefined;
+  private usageError: string | undefined;
+  private usageLastFetchAt: number | undefined;
+  private usageTimer: ReturnType<typeof setInterval> | undefined;
+  private usageRefreshInFlight = false;
+  private queuedUsageRefresh: QueuedUsageRefresh | undefined;
+  private shuttingDown = false;
+  private usageAbortController: AbortController | undefined;
+  private sessionAbortSignal: AbortSignal | undefined;
+  private sessionAbortHandler: (() => void) | undefined;
+  private sessionGeneration = 0;
+  private readonly getConfig: (ctx: ExtensionContext) => ResolvedConfig;
+  private readonly updateFooter: (ctx: ExtensionContext) => void;
+
+  constructor(
+    getConfig: (ctx: ExtensionContext) => ResolvedConfig,
+    updateFooter: (ctx: ExtensionContext) => void,
+  ) {
+    this.getConfig = getConfig;
+    this.updateFooter = updateFooter;
+  }
+
+  get snapshot(): UsageSnapshot | undefined {
+    return this.usageSnapshot;
+  }
+
+  statusLine(
+    ctx: ExtensionContext,
+    cfg = this.getConfig(ctx),
+    isUsingOAuth?: boolean,
+  ): string | undefined {
+    return this.usageSnapshot &&
+      !this.usageError &&
+      this.usageSnapshot.scope === usageScopeForModel(ctx.model?.id) &&
+      cfg.usage.enabled &&
+      isOpenAISubscriptionModel(ctx, cfg, isUsingOAuth)
+      ? formatUsageSnapshot(this.usageSnapshot, cfg.usage)
+      : undefined;
+  }
+
+  formatStatus(ctx: ExtensionContext): string {
+    const cfg = this.getConfig(ctx);
+    if (!cfg.usage.enabled) return "Usage display is disabled.";
+    if (!isOpenAISubscriptionModel(ctx, cfg))
+      return "Usage hidden: current model is not an OpenAI subscription model.";
+    if (this.usageError) return `Usage unavailable: ${this.usageError}`;
+    if (!this.usageSnapshot || this.usageSnapshot.scope !== usageScopeForModel(ctx.model?.id))
+      return "Usage unavailable.";
+    const stale =
+      this.usageUpdatedAt && Date.now() - this.usageUpdatedAt > cfg.usage.refreshIntervalMs * 2
+        ? ` | stale ${formatResetCountdown((Date.now() - this.usageUpdatedAt) / 1000)}`
+        : "";
+    return `${formatUsageSnapshot(this.usageSnapshot, cfg.usage)}${stale}`;
+  }
+
+  formatDebug(ctx: ExtensionContext): string {
+    const cfg = this.getConfig(ctx);
+    const auth = readCodexAuth();
+    return [
+      `Usage enabled: ${cfg.usage.enabled}`,
+      `Current model: ${currentModelKey(ctx)}`,
+      `Current model eligible: ${isOpenAISubscriptionModel(ctx, cfg)}`,
+      `Requires subscription model: ${cfg.usage.showOnlyOnSubscriptionModels}`,
+      `Auth: ${auth ? "found" : "missing"}`,
+      `Account ID: ${maskIdentifier(auth?.accountId) ?? "none"}`,
+      `Last fetch: ${this.usageLastFetchAt ? new Date(this.usageLastFetchAt).toLocaleTimeString() : "never"}`,
+      `Last successful update: ${this.usageUpdatedAt ? new Date(this.usageUpdatedAt).toLocaleTimeString() : "never"}`,
+      `Last error: ${this.usageError ?? "none"}`,
+      `Refresh interval: ${cfg.usage.refreshIntervalMs}ms`,
+      `Endpoint: https://chatgpt.com/backend-api/wham/usage`,
+    ].join("\n");
+  }
+
+  private isGenerationCurrent(generation: number): boolean {
+    return !this.shuttingDown && generation === this.sessionGeneration;
+  }
+
+  private deactivateGeneration(generation: number): void {
+    if (generation !== this.sessionGeneration) return;
+    this.shuttingDown = true;
+    this.sessionGeneration++;
+    this.queuedUsageRefresh = undefined;
+    this.usageAbortController?.abort();
+    this.usageAbortController = undefined;
+    this.stopTimer();
+  }
+
+  private handleStaleContextError(error: unknown, generation: number): boolean {
+    if (!isStaleExtensionContextError(error)) return false;
+    this.deactivateGeneration(generation);
+    return true;
+  }
+
+  async refresh(
+    ctx: ExtensionContext,
+    modelId?: string,
+    options?: UsageRefreshOptions,
+    generation = this.sessionGeneration,
+  ): Promise<void> {
+    if (!this.isGenerationCurrent(generation)) return;
+
+    let resolvedModelId = modelId;
+    try {
+      if (!ctx.hasUI) return;
+      resolvedModelId ??= ctx.model?.id;
+    } catch (error) {
+      this.handleStaleContextError(error, generation);
+      return;
+    }
+
+    if (this.usageRefreshInFlight) {
+      const queued =
+        this.queuedUsageRefresh?.generation === generation ? this.queuedUsageRefresh : undefined;
+      this.queuedUsageRefresh = {
+        ctx,
+        generation,
+        modelId: resolvedModelId,
+        notify: queued?.notify || options?.notify,
+        force: queued?.force || options?.force,
+      };
+      return;
+    }
+
+    this.usageRefreshInFlight = true;
+    try {
+      const cfg = this.getConfig(ctx);
+      if (!this.isGenerationCurrent(generation)) return;
+
+      if (!cfg.usage.enabled) {
+        this.usageSnapshot = undefined;
+        this.usageError = "Usage display is disabled.";
+        this.updateFooter(ctx);
+        if (options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
+        return;
+      }
+      if (!isOpenAISubscriptionModel(ctx, cfg)) {
+        this.updateFooter(ctx);
+        if (options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
+        return;
+      }
+      const shouldThrottle =
+        !options?.force &&
+        !options?.notify &&
+        this.usageLastFetchAt !== undefined &&
+        Date.now() - this.usageLastFetchAt < cfg.usage.refreshIntervalMs;
+      if (shouldThrottle) return;
+      this.usageLastFetchAt = Date.now();
+      this.usageAbortController = new AbortController();
+      const timeoutSignal = AbortSignal.timeout(10_000);
+      const signal = ctx.signal
+        ? AbortSignal.any([ctx.signal, timeoutSignal, this.usageAbortController.signal])
+        : AbortSignal.any([timeoutSignal, this.usageAbortController.signal]);
+      const data = await requestCodexUsage(ctx, signal);
+      if (!this.isGenerationCurrent(generation)) return;
+
+      this.usageSnapshot = data ? parseUsageSnapshot(data, resolvedModelId) : undefined;
+      this.usageUpdatedAt = this.usageSnapshot ? Date.now() : undefined;
+      this.usageError = data
+        ? undefined
+        : `Missing openai-codex OAuth credentials in ${AUTH_FILE}.`;
+      this.updateFooter(ctx);
+      if (options?.notify)
+        ctx.ui.notify(this.formatStatus(ctx), this.usageSnapshot ? "info" : "warning");
+    } catch (error) {
+      if (this.handleStaleContextError(error, generation) || !this.isGenerationCurrent(generation))
+        return;
+      this.usageError = sanitizeDiagnosticError(
+        error instanceof Error ? error.message : String(error),
+      );
+      try {
+        this.updateFooter(ctx);
+        if (options?.notify) ctx.ui.notify(this.formatStatus(ctx), "warning");
+      } catch (secondaryError) {
+        if (!this.handleStaleContextError(secondaryError, generation)) {
+          this.usageError = sanitizeDiagnosticError(
+            secondaryError instanceof Error ? secondaryError.message : String(secondaryError),
+          );
+        }
+      }
+    } finally {
+      this.usageAbortController = undefined;
+      this.usageRefreshInFlight = false;
+      const next = this.queuedUsageRefresh;
+      this.queuedUsageRefresh = undefined;
+      if (next && !this.shuttingDown && next.generation === this.sessionGeneration) {
+        void this.refresh(
+          next.ctx,
+          next.modelId,
+          { notify: next.notify, force: next.force },
+          next.generation,
+        );
+      }
+    }
+  }
+
+  private stopTimer(): void {
+    if (this.usageTimer) clearInterval(this.usageTimer);
+    this.usageTimer = undefined;
+    if (this.sessionAbortSignal && this.sessionAbortHandler) {
+      this.sessionAbortSignal.removeEventListener("abort", this.sessionAbortHandler);
+    }
+    this.sessionAbortSignal = undefined;
+    this.sessionAbortHandler = undefined;
+  }
+
+  start(ctx: ExtensionContext): void {
+    this.usageAbortController?.abort();
+    this.queuedUsageRefresh = undefined;
+    this.stopTimer();
+    const generation = ++this.sessionGeneration;
+    this.shuttingDown = false;
+
+    const cfg = this.getConfig(ctx);
+    if (!cfg.usage.enabled) return;
+    const sessionSignal = ctx.signal;
+    if (sessionSignal?.aborted) {
+      this.deactivateGeneration(generation);
+      return;
+    }
+    this.sessionAbortSignal = sessionSignal;
+    this.sessionAbortHandler = () => {
+      this.deactivateGeneration(generation);
+    };
+    sessionSignal?.addEventListener("abort", this.sessionAbortHandler, { once: true });
+    void this.refresh(ctx, undefined, { force: true }, generation);
+    this.usageTimer = setInterval(() => {
+      if (!this.isGenerationCurrent(generation)) return;
+      if (sessionSignal?.aborted) {
+        this.deactivateGeneration(generation);
+        return;
+      }
+      void this.refresh(ctx, undefined, undefined, generation);
+    }, cfg.usage.refreshIntervalMs);
+    this.usageTimer.unref?.();
+  }
+
+  restartAfterSettingsChange(ctx: ExtensionContext, cfg: ResolvedConfig): void {
+    this.usageAbortController?.abort();
+    this.queuedUsageRefresh = undefined;
+    this.stopTimer();
+    this.sessionGeneration++;
+    this.shuttingDown = false;
+    if (cfg.usage.enabled) this.start(ctx);
+    else {
+      this.usageSnapshot = undefined;
+      this.usageError = "Usage display is disabled.";
+    }
+  }
+
+  shutdown(): void {
+    this.shuttingDown = true;
+    this.sessionGeneration++;
+    this.queuedUsageRefresh = undefined;
+    this.usageAbortController?.abort();
+    this.usageAbortController = undefined;
+    this.stopTimer();
+  }
+}

package/src/usage.ts

@@ -0,0 +1,264 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { getCodexCredentials } from "./codex-auth.ts";
+export { AUTH_FILE, readCodexAuth } from "./codex-auth.ts";
+
+export type UsageWindow = {
+  used_percent?: number | null;
+  reset_after_seconds?: number | null;
+  reset_at?: number | null;
+};
+
+export type RateLimitBucket = {
+  allowed?: boolean;
+  limit_reached?: boolean;
+  primary_window?: UsageWindow | null;
+  secondary_window?: UsageWindow | null;
+};
+
+export type CodexUsageResponse = {
+  rate_limit?: RateLimitBucket | null;
+  additional_rate_limits?: Record<string, unknown> | unknown[] | null;
+};
+
+export type UsageSnapshot = {
+  capturedAt: number;
+  scope: UsageScope;
+  fiveHourLeftPercent: number | null;
+  sevenDayLeftPercent: number | null;
+  fiveHourResetInSeconds: number | null;
+  sevenDayResetInSeconds: number | null;
+  isLimited: boolean;
+};
+
+export type UsageScope = "default" | "spark";
+
+export const USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
+const SPARK_MODEL_ID = "gpt-5.3-codex-spark";
+const SPARK_LIMIT_NAME = "GPT-5.3-Codex-Spark";
+type ResetClockFormatters = {
+  time: Intl.DateTimeFormat;
+  weekday: Intl.DateTimeFormat;
+  date: Intl.DateTimeFormat;
+};
+const RESET_CLOCK_FORMATTER_CACHE_LIMIT = 4;
+const resetClockFormatters = new Map<string, ResetClockFormatters>();
+
+function currentTimeZoneKey(date: Date): string {
+  const zoneLabel = /\(([^)]+)\)$/.exec(date.toString())?.[1] ?? "";
+  return `${process.env.TZ ?? ""}:${date.getTimezoneOffset()}:${zoneLabel}`;
+}
+
+function getResetClockFormatters(now: Date, reset: Date): ResetClockFormatters {
+  const timeZoneKey = `${currentTimeZoneKey(now)}:${reset.getTimezoneOffset()}`;
+  let formatters = resetClockFormatters.get(timeZoneKey);
+  if (!formatters) {
+    formatters = {
+      time: new Intl.DateTimeFormat(undefined, { hour: "numeric", minute: "2-digit" }),
+      weekday: new Intl.DateTimeFormat(undefined, { weekday: "short" }),
+      date: new Intl.DateTimeFormat(undefined, { month: "numeric", day: "numeric" }),
+    };
+    resetClockFormatters.set(timeZoneKey, formatters);
+    while (resetClockFormatters.size > RESET_CLOCK_FORMATTER_CACHE_LIMIT) {
+      const oldestKey = resetClockFormatters.keys().next().value;
+      if (oldestKey === undefined) break;
+      resetClockFormatters.delete(oldestKey);
+    }
+  }
+  return formatters;
+}
+
+function asObject(value: unknown): Record<string, unknown> | null {
+  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
+  return value as Record<string, unknown>;
+}
+
+function clampPercent(value: number): number {
+  return Math.min(100, Math.max(0, value));
+}
+
+function usedToLeftPercent(value: number | null | undefined): number | null {
+  if (typeof value !== "number" || !Number.isFinite(value)) return null;
+  return clampPercent(100 - value);
+}
+
+export function formatResetCountdown(seconds: number | null): string | null {
+  if (typeof seconds !== "number" || !Number.isFinite(seconds)) return null;
+  const total = Math.max(0, Math.round(seconds));
+  const days = Math.floor(total / 86_400);
+  const hours = Math.floor((total % 86_400) / 3_600);
+  const minutes = Math.floor((total % 3_600) / 60);
+  const secs = total % 60;
+  if (days > 0) return `${days}d${hours}h`;
+  if (hours > 0) return `${hours}h${minutes}m`;
+  if (minutes > 0) return `${minutes}m`;
+  return `${secs}s`;
+}
+
+function formatResetClock(
+  seconds: number | null,
+  options?: { includeDate?: boolean },
+  now = Date.now(),
+): string | null {
+  if (typeof seconds !== "number" || !Number.isFinite(seconds)) return null;
+  const resetDate = new Date(now + seconds * 1000);
+  const currentDate = new Date(now);
+  const formatters = getResetClockFormatters(currentDate, resetDate);
+  const time = formatters.time.format(resetDate);
+  if (!options?.includeDate && resetDate.toDateString() === currentDate.toDateString()) return time;
+  const weekday = formatters.weekday.format(resetDate);
+  if (!options?.includeDate) return `${weekday} ${time}`;
+  const date = formatters.date.format(resetDate);
+  return `${weekday} ${date} ${time}`;
+}
+
+function formatCompactReset(
+  label: string,
+  seconds: number | null,
+  options?: { includeDate?: boolean },
+  now = Date.now(),
+): string | null {
+  const countdown = formatResetCountdown(seconds);
+  const clock = formatResetClock(seconds, options, now);
+  return countdown && clock ? `${label} ↺ ${countdown} - ${clock}` : null;
+}
+
+function isAbortSignal(value: unknown): value is AbortSignal {
+  return typeof value === "object" && value !== null && "aborted" in value;
+}
+
+export async function requestCodexUsage(
+  ctxOrSignal?: ExtensionContext | AbortSignal,
+  signal?: AbortSignal,
+): Promise<CodexUsageResponse | undefined> {
+  const ctx = isAbortSignal(ctxOrSignal) ? undefined : ctxOrSignal;
+  const requestSignal = isAbortSignal(ctxOrSignal) ? ctxOrSignal : signal;
+  const credentials = await getCodexCredentials(ctx, requestSignal);
+  if (!credentials) return undefined;
+  const response = await fetch(USAGE_URL, {
+    headers: {
+      accept: "*/*",
+      authorization: `Bearer ${credentials.accessToken}`,
+      "chatgpt-account-id": credentials.accountId,
+    },
+    signal: requestSignal,
+  });
+  if (!response.ok) throw new Error(`Codex usage request failed (${response.status})`);
+  return (await response.json()) as CodexUsageResponse;
+}
+
+function normalizeRateLimitBucket(value: unknown): RateLimitBucket | null {
+  const record = asObject(value);
+  if (!record) return null;
+  if (
+    !(
+      "primary_window" in record ||
+      "secondary_window" in record ||
+      "limit_reached" in record ||
+      "allowed" in record
+    )
+  )
+    return null;
+  return record as RateLimitBucket;
+}
+
+function extractSparkRateLimitFromEntry(value: unknown): RateLimitBucket | null {
+  const record = asObject(value);
+  if (!record || record.limit_name !== SPARK_LIMIT_NAME) return null;
+  return normalizeRateLimitBucket(record.rate_limit);
+}
+
+function findSparkRateLimitBucket(data: CodexUsageResponse): RateLimitBucket | null {
+  const additional = data.additional_rate_limits;
+  if (Array.isArray(additional)) {
+    for (const entry of additional) {
+      const bucket = extractSparkRateLimitFromEntry(entry);
+      if (bucket) return bucket;
+    }
+  } else {
+    const map = asObject(additional);
+    if (map) {
+      for (const value of Object.values(map)) {
+        const bucket = extractSparkRateLimitFromEntry(value);
+        if (bucket) return bucket;
+      }
+    }
+  }
+  return null;
+}
+
+function getResetSeconds(window: UsageWindow | null | undefined, now: number): number | null {
+  if (
+    typeof window?.reset_after_seconds === "number" &&
+    Number.isFinite(window.reset_after_seconds)
+  )
+    return window.reset_after_seconds;
+  if (typeof window?.reset_at !== "number" || !Number.isFinite(window.reset_at)) return null;
+  const resetAtSeconds =
+    window.reset_at > 100_000_000_000 ? window.reset_at / 1000 : window.reset_at;
+  return Math.max(0, resetAtSeconds - now / 1000);
+}
+
+export function usageScopeForModel(modelId: string | undefined): UsageScope {
+  return modelId === SPARK_MODEL_ID ? "spark" : "default";
+}
+
+export function parseUsageSnapshot(
+  data: CodexUsageResponse,
+  modelId: string | undefined,
+  now = Date.now(),
+): UsageSnapshot {
+  const scope = usageScopeForModel(modelId);
+  const bucket =
+    scope === "spark"
+      ? (findSparkRateLimitBucket(data) ?? normalizeRateLimitBucket(data.rate_limit))
+      : normalizeRateLimitBucket(data.rate_limit);
+  return {
+    capturedAt: now,
+    scope,
+    fiveHourLeftPercent: usedToLeftPercent(bucket?.primary_window?.used_percent),
+    sevenDayLeftPercent: usedToLeftPercent(bucket?.secondary_window?.used_percent),
+    fiveHourResetInSeconds: getResetSeconds(bucket?.primary_window, now),
+    sevenDayResetInSeconds: getResetSeconds(bucket?.secondary_window, now),
+    isLimited: bucket?.limit_reached === true || bucket?.allowed === false,
+  };
+}
+
+export function formatPercent(value: number | null): string {
+  return typeof value === "number" && Number.isFinite(value)
+    ? `${Math.round(clampPercent(value))}%`
+    : "--";
+}
+
+export function formatUsageSnapshot(
+  snapshot: UsageSnapshot,
+  options: { showResetTimes: boolean },
+  now = Date.now(),
+): string {
+  const fiveHour = formatPercent(snapshot.fiveHourLeftPercent);
+  const sevenDay = formatPercent(snapshot.sevenDayLeftPercent);
+  const resets = options.showResetTimes
+    ? [
+        formatCompactReset(
+          "5h",
+          remainingResetSeconds(snapshot.fiveHourResetInSeconds, snapshot.capturedAt, now),
+          undefined,
+          now,
+        ),
+        formatCompactReset(
+          "7d",
+          remainingResetSeconds(snapshot.sevenDayResetInSeconds, snapshot.capturedAt, now),
+          { includeDate: true },
+          now,
+        ),
+      ].filter((value): value is string => value !== null)
+    : [];
+  return `Usage: 5h: ${fiveHour} | 7d: ${sevenDay}${resets.length ? ` | ${resets.join(" | ")}` : ""}`;
+}
+
+function remainingResetSeconds(
+  seconds: number | null,
+  capturedAt: number,
+  now: number,
+): number | null {
+  return seconds === null ? null : seconds - (now - capturedAt) / 1000;
+}