pi-better-openai 0.1.18 → 0.1.21

package/src/usage-controller.ts

@@ -0,0 +1,299 @@
+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;
+
+  constructor(
+    private readonly getConfig: (ctx: ExtensionContext) => ResolvedConfig,
+    private readonly updateFooter: (ctx: ExtensionContext) => void,
+  ) {}
+
+  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

@@ -1,6 +1,6 @@
-import { readFileSync } from "node:fs";
-import { homedir } from "node:os";
-import { join } from "node:path";
+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;
@@ -21,6 +21,8 @@ export type CodexUsageResponse = {
 };
 
 export type UsageSnapshot = {
+  capturedAt: number;
+  scope: UsageScope;
   fiveHourLeftPercent: number | null;
   sevenDayLeftPercent: number | null;
   fiveHourResetInSeconds: number | null;
@@ -28,11 +30,42 @@ export type UsageSnapshot = {
   isLimited: boolean;
 };
 
-const AGENT_DIR = process.env.PI_CODING_AGENT_DIR?.trim() || join(homedir(), ".pi", "agent");
-export const AUTH_FILE = join(AGENT_DIR, "auth.json");
+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;
@@ -44,12 +77,12 @@ function clampPercent(value: number): number {
 }
 
 function usedToLeftPercent(value: number | null | undefined): number | null {
-  if (typeof value !== "number" || Number.isNaN(value)) return 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.isNaN(seconds)) return 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);
@@ -64,15 +97,17 @@ export function formatResetCountdown(seconds: number | null): string | null {
 function formatResetClock(
   seconds: number | null,
   options?: { includeDate?: boolean },
+  now = Date.now(),
 ): string | null {
-  if (typeof seconds !== "number" || Number.isNaN(seconds)) return null;
-  const resetDate = new Date(Date.now() + Math.max(0, seconds) * 1000);
-  const now = new Date();
-  const time = resetDate.toLocaleTimeString(undefined, { hour: "numeric", minute: "2-digit" });
-  if (!options?.includeDate && resetDate.toDateString() === now.toDateString()) return time;
-  const weekday = resetDate.toLocaleDateString(undefined, { weekday: "short" });
+  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 = resetDate.toLocaleDateString(undefined, { month: "numeric", day: "numeric" });
+  const date = formatters.date.format(resetDate);
   return `${weekday} ${date} ${time}`;
 }
 
@@ -80,38 +115,24 @@ function formatCompactReset(
   label: string,
   seconds: number | null,
   options?: { includeDate?: boolean },
+  now = Date.now(),
 ): string | null {
   const countdown = formatResetCountdown(seconds);
-  const clock = formatResetClock(seconds, options);
+  const clock = formatResetClock(seconds, options, now);
   return countdown && clock ? `${label} ↺ ${countdown} - ${clock}` : null;
 }
 
-export function readCodexAuth(): { accessToken: string; accountId: string } | undefined {
-  try {
-    const auth = JSON.parse(readFileSync(AUTH_FILE, "utf8")) as Record<
-      string,
-      | {
-          type?: string;
-          access?: string | null;
-          accountId?: string | null;
-          account_id?: string | null;
-        }
-      | undefined
-    >;
-    const entry = auth["openai-codex"];
-    if (entry?.type !== "oauth") return undefined;
-    const accessToken = entry.access?.trim();
-    const accountId = (entry.accountId ?? entry.account_id)?.trim();
-    return accessToken && accountId ? { accessToken, accountId } : undefined;
-  } catch {
-    return undefined;
-  }
+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 credentials = readCodexAuth();
+  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: {
@@ -119,7 +140,7 @@ export async function requestCodexUsage(
       authorization: `Bearer ${credentials.accessToken}`,
       "chatgpt-account-id": credentials.accountId,
     },
-    signal,
+    signal: requestSignal,
   });
   if (!response.ok) throw new Error(`Codex usage request failed (${response.status})`);
   return (await response.json()) as CodexUsageResponse;
@@ -165,34 +186,45 @@ function findSparkRateLimitBucket(data: CodexUsageResponse): RateLimitBucket | n
   return null;
 }
 
-function getResetSeconds(window: UsageWindow | null | undefined): number | null {
-  if (typeof window?.reset_after_seconds === "number" && !Number.isNaN(window.reset_after_seconds))
+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.isNaN(window.reset_at)) return null;
+  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 - Date.now() / 1000);
+  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 =
-    modelId === SPARK_MODEL_ID
-      ? findSparkRateLimitBucket(data)
+    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),
-    sevenDayResetInSeconds: getResetSeconds(bucket?.secondary_window),
+    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.isNaN(value)
+  return typeof value === "number" && Number.isFinite(value)
     ? `${Math.round(clampPercent(value))}%`
     : "--";
 }
@@ -200,14 +232,33 @@ export function formatPercent(value: number | null): string {
 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", snapshot.fiveHourResetInSeconds),
-        formatCompactReset("7d", snapshot.sevenDayResetInSeconds, { includeDate: true }),
+        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;
+}