pi-openai-usage 0.1.0

package/src/status-controller.ts

@@ -0,0 +1,280 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+import {
+  resolveCodexOAuthCredentials,
+  type CodexCredentialResolution,
+  type CodexModelRegistry,
+} from "./auth";
+import { type LoadedUsageConfig, type UsageConfig, loadUsageConfig } from "./config";
+import type { UsageColorTheme } from "./color";
+import {
+  appendUsageRefreshFailureMarker,
+  formatUsageAuthFailedStatusLine,
+  formatUsageLoginRequiredStatusLine,
+  formatUsageRefreshFailedStatusLine,
+  formatUsageStatusLine,
+} from "./format";
+import { fetchCodexUsage, type UsageClientPort, type UsageFetchError } from "./usage-client";
+import { createUsageRefreshCoordinator, type UsageRefreshCoordinator } from "./usage-refresh-coordinator";
+import { createUsageStateStore, type UsageStateStore } from "./usage-state";
+import { decideUsageVisibility, type UsageModel, type UsageModelRegistry } from "./visibility";
+
+export const USAGE_STATUS_KEY = "openai-usage";
+
+export type StatusLineContext = {
+  model?: UsageModel;
+  modelRegistry?: UsageModelRegistry & CodexModelRegistry;
+  signal?: AbortSignal;
+  ui: {
+    setStatus(key: string, text: string | undefined): void;
+    theme?: UsageColorTheme;
+  };
+};
+
+type TimerHandle = unknown;
+
+export type TimerApi = {
+  setInterval(handler: () => void, delayMs: number): TimerHandle;
+  clearInterval(handle: TimerHandle): void;
+};
+
+export type UsageStatusControllerDependencies = {
+  loadConfig?: () => LoadedUsageConfig;
+  timerApi?: TimerApi;
+  resolveCredentials?: (
+    ctx: StatusLineContext,
+  ) => CodexCredentialResolution | Promise<CodexCredentialResolution>;
+  usageClient?: UsageClientPort;
+  usageState?: UsageStateStore;
+  usageRefreshCoordinator?: UsageRefreshCoordinator;
+};
+
+export type UsageStatusController = {
+  reapply(ctx: StatusLineContext): Promise<void>;
+};
+
+const defaultTimerApi: TimerApi = {
+  setInterval(handler, delayMs) {
+    const handle = setInterval(handler, delayMs);
+    handle.unref?.();
+    return handle;
+  },
+  clearInterval(handle) {
+    clearInterval(handle as ReturnType<typeof setInterval>);
+  },
+};
+
+const defaultUsageClient: UsageClientPort = {
+  fetchUsage: fetchCodexUsage,
+};
+
+/**
+ * Registers lifecycle ownership for the Usage Status Extension's status-line entry.
+ */
+export function registerUsageStatusController(
+  pi: ExtensionAPI,
+  dependencies: UsageStatusControllerDependencies = {},
+): UsageStatusController {
+  const loadConfig = dependencies.loadConfig ?? loadUsageConfig;
+  const timerApi = dependencies.timerApi ?? defaultTimerApi;
+  const resolveCredentials = dependencies.resolveCredentials ?? defaultResolveCredentials;
+  const usageClient = dependencies.usageClient ?? defaultUsageClient;
+  const usageState = dependencies.usageState ?? createUsageStateStore();
+  const refreshCoordinator =
+    dependencies.usageRefreshCoordinator ??
+    createUsageRefreshCoordinator({ usageClient, usageState });
+  let refreshTimer: TimerHandle | undefined;
+  let refreshTimerDelayMs: number | undefined;
+  let activeApplyId = 0;
+
+  pi.on("session_start", async (_event, ctx) => {
+    await applyCurrentConfiguration(ctx);
+  });
+
+  pi.on("model_select", async (_event, ctx) => {
+    await applyCurrentConfiguration(ctx);
+  });
+
+  pi.on("turn_end", async (_event, ctx) => {
+    await applyCurrentConfiguration(ctx);
+  });
+
+  pi.on("session_shutdown", (_event, ctx) => {
+    stopRefreshTimer();
+    refreshCoordinator.abortInFlightRefresh();
+    activeApplyId += 1;
+    clearUsageStatus(ctx);
+  });
+
+  async function applyCurrentConfiguration(ctx: StatusLineContext): Promise<void> {
+    const applyId = ++activeApplyId;
+
+    const { effective: initialConfig } = loadConfig();
+
+    if (!isCurrentApply(applyId)) {
+      return;
+    }
+    if (!initialConfig.enabled) {
+      stopRefreshTimer();
+      clearUsageStatus(ctx);
+      return;
+    }
+
+    const initialVisibility = decideUsageVisibility({
+      showAlways: initialConfig.display.showAlways,
+      model: ctx.model,
+      modelRegistry: ctx.modelRegistry,
+    });
+
+    if (!initialVisibility.visible) {
+      stopRefreshTimer();
+      clearUsageStatus(ctx);
+      return;
+    }
+
+    const credentials = await resolveCredentials(ctx);
+    if (!isCurrentApply(applyId)) {
+      return;
+    }
+    if (!credentials.ok) {
+      stopRefreshTimer();
+      publishLoginRequiredStatus(ctx);
+      return;
+    }
+
+    const hadCachedSnapshot = usageState.getSnapshot(ctx.model?.id) !== undefined;
+    if (hadCachedSnapshot && isCurrentApply(applyId)) {
+      publishUsageStatus(ctx, initialConfig, usageState);
+    }
+
+    const refreshResult = await refreshCoordinator.refresh({
+      credentials: credentials.credentials,
+      signal: ctx.signal,
+      modelId: ctx.model?.id,
+      staleAfterMs: initialConfig.refreshIntervalMs,
+    });
+
+    if (!isCurrentApply(applyId)) {
+      return;
+    }
+
+    const latestConfig = loadConfig().effective;
+
+    const latestVisibility = decideUsageVisibility({
+      showAlways: latestConfig.display.showAlways,
+      model: ctx.model,
+      modelRegistry: ctx.modelRegistry,
+    });
+
+    if (!latestConfig.enabled || !latestVisibility.visible) {
+      stopRefreshTimer();
+      clearUsageStatus(ctx);
+      return;
+    }
+
+    if (refreshResult.status === "failed") {
+      publishRefreshFailureStatus(ctx, latestConfig, refreshResult.error, hadCachedSnapshot);
+    } else if (refreshResult.status !== "skipped" || hadCachedSnapshot) {
+      publishUsageStatus(ctx, latestConfig, usageState);
+    }
+
+    ensureRefreshTimer(ctx, latestConfig.refreshIntervalMs);
+  }
+
+  function publishRefreshFailureStatus(
+    ctx: StatusLineContext,
+    config: UsageConfig,
+    error: UsageFetchError,
+    hadCachedSnapshot: boolean,
+  ): void {
+    if (error.kind === "auth") {
+      publishUsageAuthFailedStatus(ctx);
+      return;
+    }
+
+    if (hadCachedSnapshot) {
+      const failedCachedText = formatUsageStatusLine({
+        snapshot: usageState.getSnapshot(ctx.model?.id),
+        config,
+        theme: ctx.ui.theme,
+      });
+      if (failedCachedText !== undefined) {
+        publishRefreshFailedStatus(ctx, failedCachedText);
+        return;
+      }
+    }
+
+    publishUsageRefreshFailedStatus(ctx);
+  }
+
+  function isCurrentApply(applyId: number): boolean {
+    return applyId === activeApplyId;
+  }
+
+  function ensureRefreshTimer(ctx: StatusLineContext, delayMs: number): void {
+    if (refreshTimer !== undefined && refreshTimerDelayMs === delayMs) return;
+
+    stopRefreshTimer();
+    refreshTimer = timerApi.setInterval(() => {
+      void applyCurrentConfiguration(ctx);
+    }, delayMs);
+    refreshTimerDelayMs = delayMs;
+  }
+
+  function stopRefreshTimer(): void {
+    if (refreshTimer === undefined) return;
+
+    timerApi.clearInterval(refreshTimer);
+    refreshTimer = undefined;
+    refreshTimerDelayMs = undefined;
+  }
+
+  return {
+    reapply: applyCurrentConfiguration,
+  };
+}
+
+async function defaultResolveCredentials(
+  ctx: StatusLineContext,
+): Promise<CodexCredentialResolution> {
+  return resolveCodexOAuthCredentials({ modelRegistry: ctx.modelRegistry });
+}
+
+function publishUsageStatus(
+  ctx: StatusLineContext,
+  config: UsageConfig,
+  usageState: UsageStateStore,
+): void {
+  const statusText = formatUsageStatusLine({
+    snapshot: usageState.getSnapshot(ctx.model?.id),
+    config,
+    theme: ctx.ui.theme,
+  });
+
+  if (statusText === undefined || statusText.length === 0) {
+    clearUsageStatus(ctx);
+    return;
+  }
+
+  ctx.ui.setStatus(USAGE_STATUS_KEY, statusText);
+}
+
+function publishLoginRequiredStatus(ctx: StatusLineContext): void {
+  ctx.ui.setStatus(USAGE_STATUS_KEY, formatUsageLoginRequiredStatusLine(ctx.ui.theme));
+}
+
+function publishUsageAuthFailedStatus(ctx: StatusLineContext): void {
+  ctx.ui.setStatus(USAGE_STATUS_KEY, formatUsageAuthFailedStatusLine(ctx.ui.theme));
+}
+
+function publishRefreshFailedStatus(ctx: StatusLineContext, statusText: string): void {
+  ctx.ui.setStatus(USAGE_STATUS_KEY, appendUsageRefreshFailureMarker(statusText, ctx.ui.theme));
+}
+
+function publishUsageRefreshFailedStatus(ctx: StatusLineContext): void {
+  ctx.ui.setStatus(USAGE_STATUS_KEY, formatUsageRefreshFailedStatusLine(ctx.ui.theme));
+}
+
+function clearUsageStatus(ctx: StatusLineContext): void {
+  ctx.ui.setStatus(USAGE_STATUS_KEY, undefined);
+}

package/src/usage-client.ts

@@ -0,0 +1,144 @@
+export const USAGE_ENDPOINT = "https://chatgpt.com/backend-api/wham/usage";
+export const DEFAULT_USAGE_REQUEST_TIMEOUT_MS = 10_000;
+
+export type UsageRequestCredentials = {
+  accessToken: string;
+  accountId: string;
+};
+
+export type UsageFetchErrorKind = "auth" | "http" | "network" | "timeout" | "aborted" | "parse";
+
+export type UsageFetchError = {
+  kind: UsageFetchErrorKind;
+  message: string;
+  status?: number;
+};
+
+export type UsageFetchResult =
+  | { ok: true; raw: unknown; status: number }
+  | { ok: false; error: UsageFetchError };
+
+export type FetchCodexUsageOptions = {
+  fetchImpl?: typeof fetch;
+  signal?: AbortSignal;
+  timeoutMs?: number;
+};
+
+export type UsageClientPort = {
+  fetchUsage(
+    credentials: UsageRequestCredentials,
+    options?: Pick<FetchCodexUsageOptions, "signal">,
+  ): Promise<UsageFetchResult>;
+};
+
+const TIMEOUT_ABORT_MARKER = Symbol("usage request timeout");
+const EXTERNAL_ABORT_MARKER = Symbol("usage request aborted");
+
+export async function fetchCodexUsage(
+  credentials: UsageRequestCredentials,
+  options: FetchCodexUsageOptions = {},
+): Promise<UsageFetchResult> {
+  const fetchImpl = options.fetchImpl ?? fetch;
+  const timeoutMs = options.timeoutMs ?? DEFAULT_USAGE_REQUEST_TIMEOUT_MS;
+  const controller = new AbortController();
+  let abortReason: typeof TIMEOUT_ABORT_MARKER | typeof EXTERNAL_ABORT_MARKER | undefined;
+
+  const timeoutHandle = setTimeout(() => {
+    abortReason ??= TIMEOUT_ABORT_MARKER;
+    controller.abort(TIMEOUT_ABORT_MARKER);
+  }, timeoutMs);
+  timeoutHandle.unref?.();
+
+  const removeExternalAbortListener = pipeExternalAbort(options.signal, controller, () => {
+    abortReason ??= EXTERNAL_ABORT_MARKER;
+  });
+
+  try {
+    const response = await fetchImpl(USAGE_ENDPOINT, {
+      headers: {
+        accept: "*/*",
+        authorization: `Bearer ${credentials.accessToken}`,
+        "chatgpt-account-id": credentials.accountId,
+      },
+      signal: controller.signal,
+    });
+
+    if (!response.ok) {
+      return {
+        ok: false,
+        error: classifyHttpError(response.status),
+      };
+    }
+
+    try {
+      return {
+        ok: true,
+        raw: await response.json(),
+        status: response.status,
+      };
+    } catch {
+      return {
+        ok: false,
+        error: { kind: "parse", message: "Codex usage response was not valid JSON" },
+      };
+    }
+  } catch (error) {
+    return { ok: false, error: classifyThrownError(error, abortReason) };
+  } finally {
+    clearTimeout(timeoutHandle);
+    removeExternalAbortListener();
+  }
+}
+
+function classifyHttpError(status: number): UsageFetchError {
+  if (status === 401 || status === 403) {
+    return { kind: "auth", status, message: `Codex usage authentication failed (${status})` };
+  }
+
+  return { kind: "http", status, message: `Codex usage request failed (${status})` };
+}
+
+function classifyThrownError(
+  error: unknown,
+  abortReason: typeof TIMEOUT_ABORT_MARKER | typeof EXTERNAL_ABORT_MARKER | undefined,
+): UsageFetchError {
+  if (abortReason === TIMEOUT_ABORT_MARKER) {
+    return { kind: "timeout", message: "Codex usage request timed out" };
+  }
+
+  if (abortReason === EXTERNAL_ABORT_MARKER || isAbortError(error)) {
+    return { kind: "aborted", message: "Codex usage request was aborted" };
+  }
+
+  return { kind: "network", message: messageFromUnknownError(error) };
+}
+
+function pipeExternalAbort(
+  externalSignal: AbortSignal | undefined,
+  controller: AbortController,
+  onAbort: () => void,
+): () => void {
+  if (externalSignal === undefined) return () => undefined;
+
+  const abort = () => {
+    onAbort();
+    controller.abort(EXTERNAL_ABORT_MARKER);
+  };
+
+  if (externalSignal.aborted) {
+    abort();
+    return () => undefined;
+  }
+
+  externalSignal.addEventListener("abort", abort, { once: true });
+  return () => externalSignal.removeEventListener("abort", abort);
+}
+
+function isAbortError(error: unknown): boolean {
+  return error instanceof DOMException && error.name === "AbortError";
+}
+
+function messageFromUnknownError(error: unknown): string {
+  if (error instanceof Error && error.message.trim().length > 0) return error.message;
+  return "Codex usage request failed";
+}

package/src/usage-command-facade.ts

@@ -0,0 +1,103 @@
+import type { ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
+
+import { resolveCodexOAuthCredentials, type CodexCredentialResolution } from "./auth";
+import { type LoadedUsageConfig, loadUsageConfig, type UsageConfig } from "./config";
+import {
+  appendUsageRefreshFailureMarker,
+  formatUsageAuthFailedStatusLine,
+  formatUsageLoginRequiredStatusLine,
+  formatUsageRefreshFailedStatusLine,
+  formatUsageStatusLine,
+} from "./format";
+import type { UsageRefreshCoordinator, UsageRefreshResult } from "./usage-refresh-coordinator";
+
+export type UsageCommandFacadeContext = Pick<
+  ExtensionCommandContext,
+  "ui" | "model" | "modelRegistry" | "signal"
+>;
+
+export type UsageCommandFacadeDependencies = {
+  loadConfig?: () => LoadedUsageConfig;
+  resolveCredentials?: (
+    ctx: UsageCommandFacadeContext,
+  ) => Promise<CodexCredentialResolution>;
+  usageRefreshCoordinator: UsageRefreshCoordinator;
+};
+
+export type ShowUsageOptions = {
+  forceRefresh?: boolean;
+};
+
+export type UsageCommandFacade = {
+  showUsage(ctx: UsageCommandFacadeContext, options?: ShowUsageOptions): Promise<void>;
+};
+
+export function createUsageCommandFacade(
+  dependencies: UsageCommandFacadeDependencies,
+): UsageCommandFacade {
+  const loadConfig = dependencies.loadConfig ?? loadUsageConfig;
+  const resolveCredentials = dependencies.resolveCredentials ?? defaultResolveCredentials;
+  const usageRefreshCoordinator = dependencies.usageRefreshCoordinator;
+
+  return {
+    async showUsage(ctx, options = {}) {
+      const initialConfig = loadConfig().effective;
+      const credentials = await resolveCredentials(ctx);
+      if (!credentials.ok) {
+        notify(ctx, formatUsageLoginRequiredStatusLine());
+        return;
+      }
+
+      const refreshResult = await usageRefreshCoordinator.refresh({
+        credentials: credentials.credentials,
+        signal: ctx.signal,
+        modelId: ctx.model?.id,
+        staleAfterMs: initialConfig.refreshIntervalMs,
+        force: options.forceRefresh,
+      });
+
+      const latestConfig = loadConfig().effective;
+      const statusText = resolveUsageStatusText(latestConfig, refreshResult, ctx.ui.theme);
+      notify(ctx, statusText ?? formatUsageRefreshFailedStatusLine());
+    },
+  };
+}
+
+function resolveUsageStatusText(
+  config: UsageConfig,
+  refreshResult: UsageRefreshResult,
+  theme?: UsageCommandFacadeContext["ui"]["theme"],
+): string | undefined {
+  if (refreshResult.status === "failed") {
+    if (refreshResult.error.kind === "auth") {
+      return formatUsageAuthFailedStatusLine();
+    }
+
+    const cachedText = formatUsageStatusLine({
+      snapshot: refreshResult.snapshot,
+      config,
+      theme,
+    });
+    if (cachedText !== undefined) {
+      return appendUsageRefreshFailureMarker(cachedText, theme);
+    }
+
+    return formatUsageRefreshFailedStatusLine();
+  }
+
+  return formatUsageStatusLine({
+    snapshot: refreshResult.snapshot,
+    config,
+    theme,
+  });
+}
+
+function notify(ctx: UsageCommandFacadeContext, message: string): void {
+  ctx.ui.notify(message, "info");
+}
+
+function defaultResolveCredentials(
+  ctx: UsageCommandFacadeContext,
+): Promise<CodexCredentialResolution> {
+  return resolveCodexOAuthCredentials({ modelRegistry: ctx.modelRegistry });
+}

package/src/usage-refresh-coordinator.ts

@@ -0,0 +1,193 @@
+import type { UsageClientPort, UsageFetchError, UsageRequestCredentials } from "./usage-client";
+import { createUsageStateStore, type UsageStateStore } from "./usage-state";
+import { parseUsageSnapshot, type UsageSnapshot } from "./usage-snapshot";
+
+export type UsageRefreshOptions = {
+  credentials: UsageRequestCredentials;
+  signal?: AbortSignal;
+  modelId?: string;
+  staleAfterMs?: number;
+  force?: boolean;
+};
+
+export type UsageRefreshResult =
+  | {
+      status: "refreshed";
+      joined: boolean;
+      snapshot: UsageSnapshot;
+      fetchedAt: Date;
+    }
+  | {
+      status: "skipped";
+      reason: "fresh";
+      joined: false;
+      snapshot: UsageSnapshot | undefined;
+    }
+  | {
+      status: "failed";
+      joined: boolean;
+      error: UsageFetchError;
+      snapshot: UsageSnapshot | undefined;
+      attemptedAt: Date;
+    };
+
+export type UsageRefreshCoordinator = {
+  refresh(options: UsageRefreshOptions): Promise<UsageRefreshResult>;
+  isRefreshing(): boolean;
+  abortInFlightRefresh(): void;
+};
+
+export type UsageRefreshCoordinatorDependencies = {
+  usageClient: UsageClientPort;
+  usageState?: UsageStateStore;
+  now?: () => Date;
+};
+
+type InFlightRefreshResult =
+  | {
+      status: "refreshed";
+      raw: unknown;
+      fetchedAt: Date;
+    }
+  | {
+      status: "failed";
+      error: UsageFetchError;
+      attemptedAt: Date;
+    };
+
+export function createUsageRefreshCoordinator(
+  dependencies: UsageRefreshCoordinatorDependencies,
+): UsageRefreshCoordinator {
+  const usageState = dependencies.usageState ?? createUsageStateStore();
+  const now = dependencies.now ?? (() => new Date());
+  let inFlightRefresh: Promise<InFlightRefreshResult> | undefined;
+  let inFlightAbortController: AbortController | undefined;
+
+  async function refresh(options: UsageRefreshOptions): Promise<UsageRefreshResult> {
+    if (inFlightRefresh !== undefined) {
+      return materializeRefreshResult(usageState, await inFlightRefresh, true, options.modelId);
+    }
+
+    if (!options.force && !isCachedSnapshotStale(usageState, options.staleAfterMs, now())) {
+      return {
+        status: "skipped",
+        reason: "fresh",
+        joined: false,
+        snapshot: usageState.getSnapshot(options.modelId),
+      };
+    }
+
+    const ownerAbort = new AbortController();
+    inFlightAbortController = ownerAbort;
+    const cleanupAbortForwarder = pipeExternalAbortSignal(options.signal, ownerAbort);
+    const ownerRefresh = runRefresh(options, ownerAbort.signal);
+    inFlightRefresh = ownerRefresh;
+
+    try {
+      return materializeRefreshResult(usageState, await ownerRefresh, false, options.modelId);
+    } finally {
+      inFlightAbortController = undefined;
+      cleanupAbortForwarder();
+      if (inFlightRefresh === ownerRefresh) {
+        inFlightRefresh = undefined;
+      }
+    }
+  }
+
+  async function runRefresh(
+    options: UsageRefreshOptions,
+    signal: AbortSignal,
+  ): Promise<InFlightRefreshResult> {
+    const attemptedAt = now();
+    usageState.recordFetchAttempt(attemptedAt);
+
+    const result = await dependencies.usageClient.fetchUsage(options.credentials, { signal });
+
+    if (!result.ok) {
+      usageState.recordFetchError(result.error, attemptedAt);
+      return {
+        status: "failed",
+        error: result.error,
+        attemptedAt,
+      };
+    }
+
+    const fetchedAt = now();
+    usageState.storeRawUsageResponse(result.raw, fetchedAt);
+
+    return {
+      status: "refreshed",
+      raw: result.raw,
+      fetchedAt,
+    };
+  }
+
+  return {
+    refresh,
+    isRefreshing() {
+      return inFlightRefresh !== undefined;
+    },
+    abortInFlightRefresh() {
+      inFlightAbortController?.abort();
+    },
+  };
+}
+
+function isCachedSnapshotStale(
+  usageState: UsageStateStore,
+  staleAfterMs: number | undefined,
+  now: Date,
+): boolean {
+  if (usageState.getSnapshot() === undefined) return true;
+
+  const lastSuccessAt = usageState.getLastSuccessAt();
+  if (lastSuccessAt === undefined) return true;
+  if (staleAfterMs === undefined) return true;
+
+  return now.getTime() - lastSuccessAt.getTime() >= staleAfterMs;
+}
+
+function materializeRefreshResult(
+  usageState: UsageStateStore,
+  result: InFlightRefreshResult,
+  joined: boolean,
+  modelId: string | undefined,
+): UsageRefreshResult {
+  if (result.status === "failed") {
+    return {
+      status: "failed",
+      joined,
+      error: result.error,
+      snapshot: usageState.getSnapshot(modelId),
+      attemptedAt: result.attemptedAt,
+    };
+  }
+
+  return {
+    status: "refreshed",
+    joined,
+    snapshot: parseUsageSnapshot(result.raw, { modelId }),
+    fetchedAt: result.fetchedAt,
+  };
+}
+
+function pipeExternalAbortSignal(
+  signal: AbortSignal | undefined,
+  internalController: AbortController,
+): () => void {
+  if (signal === undefined) {
+    return () => undefined;
+  }
+
+  const abort = (): void => {
+    internalController.abort(signal.reason);
+  };
+
+  if (signal.aborted) {
+    abort();
+    return () => undefined;
+  }
+
+  signal.addEventListener("abort", abort, { once: true });
+  return () => signal.removeEventListener("abort", abort);
+}