@wayner6/pi-usage 0.1.0

package/src/modules/provider/adapters/openai-codex.ts

@@ -0,0 +1,196 @@
+import { getAgentDir } from "@earendil-works/pi-coding-agent";
+import { readFile } from "node:fs/promises";
+import { join } from "node:path";
+import type { Metric, UsageAdapter, UsageSnapshot } from "../../../core/types.ts";
+import { safeError, sameOriginFetch } from "../../../core/security.ts";
+
+const CODEX_BASE_ORIGIN = "https://chatgpt.com";
+const CODEX_USAGE_PATH = "/backend-api/wham/usage";
+
+interface WhamWindow {
+  used_percent?: number;
+  limit_window_seconds?: number;
+  reset_after_seconds?: number;
+  reset_at?: number;
+}
+
+interface WhamUsageResponse {
+  user_id?: string;
+  account_id?: string;
+  email?: string;
+  plan_type?: string;
+  rate_limit?: {
+    allowed?: boolean;
+    limit_reached?: boolean;
+    primary_window?: WhamWindow | null;
+    secondary_window?: WhamWindow | null;
+  } | null;
+  credits?: {
+    has_credits?: boolean;
+    balance?: string;
+  } | null;
+}
+
+function parseWindow(
+  window: WhamWindow | null | undefined,
+  defaultId: string,
+  defaultLabel: string,
+): Metric | undefined {
+  if (!window || typeof window.used_percent !== "number") return undefined;
+  const used = Math.min(100, Math.max(0, window.used_percent));
+  const remainingFraction = Math.max(0, (100 - used) / 100);
+
+  let resetAt: string | undefined;
+  if (typeof window.reset_at === "number" && window.reset_at > 0) {
+    resetAt = new Date(window.reset_at * 1000).toISOString();
+  } else if (typeof window.reset_after_seconds === "number" && window.reset_after_seconds > 0) {
+    resetAt = new Date(Date.now() + window.reset_after_seconds * 1000).toISOString();
+  }
+
+  return {
+    kind: "quota-window",
+    id: defaultId,
+    label: defaultLabel,
+    remainingFraction,
+    ...(resetAt ? { resetAt } : {}),
+  };
+}
+
+async function resolveLocalCodexAuth(): Promise<{ accessToken?: string | undefined; accountId?: string | undefined }> {
+  try {
+    const authPath = join(getAgentDir(), "auth.json");
+    const raw = await readFile(authPath, "utf8");
+    const parsed = JSON.parse(raw) as Record<string, { access?: string; apiKey?: string; accountId?: string; chatgpt_account_id?: string }>;
+    const codex = parsed["openai-codex"];
+    if (codex) {
+      return {
+        accessToken: codex.access ?? codex.apiKey,
+        accountId: codex.accountId ?? codex.chatgpt_account_id,
+      };
+    }
+  } catch {
+    // Ignore fallback errors
+  }
+  return {};
+}
+
+export const openAICodexAdapter: UsageAdapter = {
+  id: "openai-codex",
+  label: "OpenAI Codex (ChatGPT)",
+  canHandle(target) {
+    const pid = target.providerId.toLowerCase();
+    if (pid === "openai-codex") return true;
+    if (target.baseUrl) {
+      try {
+        const origin = new URL(target.baseUrl).origin;
+        if (origin.includes("chatgpt.com")) return true;
+      } catch {
+        return false;
+      }
+    }
+    return false;
+  },
+  async fetch({ target, signal, fetchFn }): Promise<UsageSnapshot> {
+    const fetchedAt = new Date().toISOString();
+    const authRecord = target.auth?.auth as Record<string, unknown> | undefined;
+    let accessToken = (authRecord?.apiKey ?? authRecord?.access) as string | undefined;
+    let accountId = (authRecord?.accountId ?? authRecord?.chatgpt_account_id) as string | undefined;
+
+    if (!accessToken || !accountId) {
+      const local = await resolveLocalCodexAuth();
+      accessToken = accessToken ?? local.accessToken;
+      accountId = accountId ?? local.accountId;
+    }
+
+    if (!accessToken) {
+      return {
+        adapterId: this.id,
+        sourceProviderId: target.providerId,
+        displayName: "OpenAI Codex",
+        state: "unauthorized",
+        fetchedAt,
+        accounts: [],
+        error: "No access token found in Pi auth for openai-codex",
+      };
+    }
+
+    try {
+      const headers: Record<string, string> = {
+        Authorization: `Bearer ${accessToken}`,
+        Accept: "application/json",
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
+      };
+      if (accountId) {
+        headers["ChatGPT-Account-Id"] = accountId;
+      }
+
+      const response = await sameOriginFetch(
+        new URL(CODEX_USAGE_PATH, CODEX_BASE_ORIGIN),
+        { method: "GET", headers, signal },
+        fetchFn,
+        CODEX_BASE_ORIGIN,
+      );
+
+      if (response.status === 401 || response.status === 403) {
+        return {
+          adapterId: this.id,
+          sourceProviderId: target.providerId,
+          displayName: "OpenAI Codex",
+          state: "unauthorized",
+          fetchedAt,
+          accounts: [],
+          error: `ChatGPT returned HTTP ${response.status} (token may need refresh)`,
+        };
+      }
+
+      if (!response.ok) {
+        throw new Error(`ChatGPT wham/usage returned HTTP ${response.status}`);
+      }
+
+      const data = (await response.json()) as WhamUsageResponse;
+      const metrics: Metric[] = [];
+
+      const primary = parseWindow(data.rate_limit?.primary_window, "primary-window", "Codex 5h");
+      if (primary) metrics.push(primary);
+
+      const secondary = parseWindow(data.rate_limit?.secondary_window, "secondary-window", "Codex 7d");
+      if (secondary) metrics.push(secondary);
+
+      const planLabel = data.plan_type ? `ChatGPT ${data.plan_type.toUpperCase()}` : "ChatGPT Plus/Pro";
+      const accountLabel = data.email || data.user_id || planLabel;
+
+      const rawGroups = metrics.map((m) => {
+        const qw = m as Extract<Metric, { kind: "quota-window" }>;
+        return {
+          id: qw.id,
+          label: qw.label,
+          remainingFraction: qw.remainingFraction,
+          ...(qw.resetAt ? { resetTime: qw.resetAt } : {}),
+          models: [{ id: qw.id }],
+        };
+      });
+
+      const accounts = [
+        {
+          id: data.account_id || data.user_id || "openai-codex-account",
+          provider: "openai-codex",
+          label: accountLabel,
+          status: data.rate_limit?.limit_reached ? "limit_reached" : "available",
+          metrics,
+          rawGroups,
+        },
+      ];
+
+      return {
+        adapterId: this.id,
+        sourceProviderId: target.providerId,
+        displayName: "OpenAI Codex",
+        state: metrics.length ? "ok" : "empty",
+        fetchedAt,
+        accounts,
+      };
+    } catch (error) {
+      throw new Error(safeError(error));
+    }
+  },
+};

package/src/modules/provider/adapters/xai.ts

@@ -0,0 +1,178 @@
+import { getAgentDir } from "@earendil-works/pi-coding-agent";
+import { readFile } from "node:fs/promises";
+import { join } from "node:path";
+import type { Metric, UsageAdapter, UsageSnapshot } from "../../../core/types.ts";
+import { safeError, sameOriginFetch } from "../../../core/security.ts";
+
+const XAI_AUTH_ORIGIN = "https://auth.x.ai";
+const XAI_USERINFO_PATH = "/oauth2/userinfo";
+const XAI_API_ORIGIN = "https://api.x.ai";
+
+interface XAIUserInfo {
+  sub?: string;
+  name?: string;
+  email?: string;
+  email_verified?: boolean;
+}
+
+interface XAIErrorResponse {
+  code?: string;
+  error?: string;
+}
+
+async function resolveLocalXAIAuth(): Promise<{ accessToken?: string | undefined }> {
+  try {
+    const authPath = join(getAgentDir(), "auth.json");
+    const raw = await readFile(authPath, "utf8");
+    const parsed = JSON.parse(raw) as Record<string, { access?: string; apiKey?: string }>;
+    const xai = parsed.xai;
+    if (xai) {
+      return { accessToken: xai.access ?? xai.apiKey };
+    }
+  } catch {
+    // Ignore fallback errors
+  }
+  return {};
+}
+
+export const xaiAdapter: UsageAdapter = {
+  id: "xai",
+  label: "xAI / Grok",
+  canHandle(target) {
+    const pid = target.providerId.toLowerCase();
+    if (pid === "xai" || pid === "grok") return true;
+    if (target.baseUrl) {
+      try {
+        const origin = new URL(target.baseUrl).origin;
+        if (origin.includes("x.ai")) return true;
+      } catch {
+        return false;
+      }
+    }
+    return false;
+  },
+  async fetch({ target, signal, fetchFn }): Promise<UsageSnapshot> {
+    const fetchedAt = new Date().toISOString();
+    const authRecord = target.auth?.auth as Record<string, unknown> | undefined;
+    let accessToken = (authRecord?.apiKey ?? authRecord?.access) as string | undefined;
+
+    if (!accessToken) {
+      const local = await resolveLocalXAIAuth();
+      accessToken = local.accessToken;
+    }
+
+    if (!accessToken) {
+      return {
+        adapterId: this.id,
+        sourceProviderId: target.providerId,
+        displayName: "xAI Grok",
+        state: "unauthorized",
+        fetchedAt,
+        accounts: [],
+        error: "No access token found in Pi auth for xai",
+      };
+    }
+
+    try {
+      // 1. Verify OAuth token identity with auth.x.ai
+      const userinfoRes = await sameOriginFetch(
+        new URL(XAI_USERINFO_PATH, XAI_AUTH_ORIGIN),
+        {
+          method: "GET",
+          headers: { Authorization: `Bearer ${accessToken}`, Accept: "application/json" },
+          signal,
+        },
+        fetchFn,
+        XAI_AUTH_ORIGIN,
+      );
+
+      if (userinfoRes.status === 401 || userinfoRes.status === 403) {
+        return {
+          adapterId: this.id,
+          sourceProviderId: target.providerId,
+          displayName: "xAI Grok",
+          state: "unauthorized",
+          fetchedAt,
+          accounts: [],
+          error: `xAI returned HTTP ${userinfoRes.status} (token expired or invalid)`,
+        };
+      }
+
+      let userLabel = "Grok Account";
+      let userId = "xai-user";
+      if (userinfoRes.ok) {
+        const info = (await userinfoRes.json()) as XAIUserInfo;
+        userLabel = info.email || info.name || userLabel;
+        userId = info.sub || userId;
+      }
+
+      // 2. Query billing / spending limit status via lightweight probe
+      // xAI returns HTTP 402 with code: 'personal-team-blocked:spending-limit' when credit limit is hit or subscription is required
+      let statusText = "Active";
+      let limitHit = false;
+
+      try {
+        const probeRes = await sameOriginFetch(
+          new URL("/v1/chat/completions", XAI_API_ORIGIN),
+          {
+            method: "POST",
+            headers: {
+              Authorization: `Bearer ${accessToken}`,
+              "Content-Type": "application/json",
+            },
+            body: JSON.stringify({
+              model: "grok-4.6",
+              messages: [{ role: "user", content: "" }],
+              max_tokens: 1,
+            }),
+            signal,
+          },
+          fetchFn,
+          XAI_API_ORIGIN,
+        );
+
+        if (probeRes.status === 402) {
+          limitHit = true;
+          const errData = (await probeRes.json().catch(() => ({}))) as XAIErrorResponse;
+          if (errData.code?.includes("spending-limit") || errData.error?.includes("credits")) {
+            statusText = "Limit Reached";
+          }
+        }
+      } catch {
+        // Probe is best-effort
+      }
+
+      const metrics: Metric[] = [
+        {
+          kind: "status",
+          id: "grok-subscription",
+          label: "Subscription",
+          value: statusText,
+          detail: limitHit ? "Out of credits or subscription needed" : "Connected",
+        },
+      ];
+
+      const accounts = [
+        {
+          id: userId,
+          provider: "xai",
+          label: userLabel,
+          status: limitHit ? "limit_reached" : "available",
+          metrics,
+        },
+      ];
+
+      return {
+        adapterId: this.id,
+        sourceProviderId: target.providerId,
+        displayName: "xAI Grok",
+        state: "ok",
+        fetchedAt,
+        summary: `Grok · ${statusText}`,
+        accounts,
+      };
+    } catch (error) {
+      throw new Error(safeError(error));
+    }
+  },
+};

package/src/modules/provider/controller.ts

@@ -0,0 +1,147 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { Api, Model } from "@earendil-works/pi-ai";
+import type { UsageConfig } from "../../core/config.ts";
+import { UsageCache } from "../../core/cache.ts";
+import type { Metric, ProviderTarget, UsageAdapter, UsageSnapshot } from "../../core/types.ts";
+import { anthropicAdapter } from "./adapters/anthropic.ts";
+import { cliProxyBridgeAdapter } from "./adapters/cliproxy-pi-bridge.ts";
+import { deepSeekAdapter } from "./adapters/deepseek.ts";
+import { glmAdapter } from "./adapters/glm.ts";
+import { openAICodexAdapter } from "./adapters/openai-codex.ts";
+import { xaiAdapter } from "./adapters/xai.ts";
+import { chooseAdapter, matchModelAcrossAccounts } from "./matching.ts";
+import { relativeTime } from "../../ui/format.ts";
+
+export class ProviderUsageController {
+  readonly cache = new UsageCache();
+  private adapters: UsageAdapter[];
+
+  constructor(private config: UsageConfig, private fetchFn: typeof fetch = fetch) {
+    this.adapters = [deepSeekAdapter, openAICodexAdapter, xaiAdapter, anthropicAdapter, glmAdapter, cliProxyBridgeAdapter];
+  }
+
+  setConfig(config: UsageConfig): void { this.config = config; }
+
+  async target(ctx: ExtensionContext, providerId: string, model = ctx.model): Promise<ProviderTarget> {
+    const provider = ctx.modelRegistry.getProvider(providerId);
+    const auth = await ctx.modelRegistry.getProviderAuth(providerId);
+    const baseUrl = auth?.auth.baseUrl ?? model?.baseUrl ?? provider?.baseUrl;
+    return {
+      providerId,
+      ...(model?.provider === providerId ? { model } : {}),
+      ...(provider ? { provider } : {}),
+      ...(auth ? { auth } : {}),
+      ...(baseUrl ? { baseUrl } : {}),
+    };
+  }
+
+  private enabled(adapter: UsageAdapter): boolean {
+    if (adapter.id === "deepseek") return this.config.adapters.deepseek.enabled;
+    if (adapter.id === "cliproxy-pi-bridge") return this.config.adapters.cliproxyPiBridge.enabled;
+    if (adapter.id === "openai-codex") return this.config.adapters.openaiCodex.enabled;
+    if (adapter.id === "xai") return this.config.adapters.xai.enabled;
+    if (adapter.id === "anthropic") return this.config.adapters.anthropic.enabled;
+    if (adapter.id === "glm") return this.config.adapters.glm.enabled;
+    return true;
+  }
+
+  async fetchTarget(target: ProviderTarget, force = false): Promise<UsageSnapshot> {
+    const adapter = chooseAdapter(target, this.adapters.filter((item) => this.enabled(item)), this.config);
+    if (!adapter) return { adapterId: "none", sourceProviderId: target.providerId, displayName: target.providerId, state: "unsupported", fetchedAt: new Date().toISOString(), accounts: [], error: "No enabled usage adapter matched this provider" };
+    const key = `${target.providerId}:${adapter.id}`;
+    return this.cache.coalesce(key, async () => {
+      const timeout = AbortSignal.timeout(this.config.refresh.timeoutSeconds * 1000);
+      return adapter.fetch({ target, signal: timeout, force, fetchFn: this.fetchFn });
+    });
+  }
+
+  async refreshCurrent(ctx: ExtensionContext, force = false, model: Model<Api> | undefined = ctx.model): Promise<UsageSnapshot | undefined> {
+    if (!model) return undefined;
+    return this.fetchTarget(await this.target(ctx, model.provider, model), force);
+  }
+
+  async refreshAll(ctx: ExtensionContext, force = false): Promise<UsageSnapshot[]> {
+    const providerIds = new Set(ctx.modelRegistry.getAvailable().map((model) => model.provider));
+    for (const id of ctx.modelRegistry.getRegisteredProviderIds()) providerIds.add(id);
+    if (ctx.modelRegistry.getProviderAuthStatus("deepseek").configured) providerIds.add("deepseek");
+    if (ctx.modelRegistry.getProviderAuthStatus("openai-codex").configured) providerIds.add("openai-codex");
+    if (ctx.modelRegistry.getProviderAuthStatus("xai").configured) providerIds.add("xai");
+    if (ctx.modelRegistry.getProviderAuthStatus("anthropic").configured) providerIds.add("anthropic");
+    if (ctx.modelRegistry.getProviderAuthStatus("zai-coding-cn").configured) providerIds.add("zai-coding-cn");
+    if (ctx.modelRegistry.getProviderAuthStatus("zai").configured) providerIds.add("zai");
+    if (ctx.modelRegistry.getProviderAuthStatus("glm").configured) providerIds.add("glm");
+    for (const id of Object.keys(this.config.providerOverrides)) providerIds.add(id);
+    if (ctx.model) providerIds.add(ctx.model.provider);
+    const targets = await Promise.all([...providerIds].map((id) => this.target(ctx, id)));
+    const snapshots = await Promise.all(targets.map((target) => this.fetchTarget(target, force)));
+    return snapshots.filter((snapshot) => snapshot.state !== "unsupported");
+  }
+
+  /**
+   * Derive a view tailored to the active model using universal cross-account group matching.
+   */
+  currentView(ctx: ExtensionContext, snapshot?: UsageSnapshot, model: Model<Api> | undefined = ctx.model): UsageSnapshot | undefined {
+    if (!snapshot) return undefined;
+    if (snapshot.adapterId !== "cliproxy-pi-bridge" && snapshot.adapterId !== "openai-codex") return snapshot;
+
+    const matched = matchModelAcrossAccounts(snapshot.accounts, model?.id);
+    if (matched) {
+      let summary: string;
+
+      if (matched.quota.multiWindows && matched.quota.multiWindows.length > 1) {
+        // Special case: Multiple time windows for the same model (e.g. Codex 5h and 7d)
+        const parts = matched.quota.multiWindows.map((q) => {
+          const sub = q.label.replace(/^Codex\s+/, "");
+          const reset = q.resetAt ? relativeTime(q.resetAt) : undefined;
+          return `${sub} ${Math.round(q.remainingFraction * 100)}%${reset ? ` (${reset})` : ""}`;
+        });
+        summary = `Codex ${parts.join(" · ")}`;
+      } else {
+        const reset = matched.quota.resetAt ? relativeTime(matched.quota.resetAt) : undefined;
+        summary = `${matched.quota.label} ${Math.round(matched.quota.remainingFraction * 100)}%${reset ? ` (${reset})` : ""}`;
+      }
+
+      return {
+        ...snapshot,
+        accounts: [matched.account as never],
+        state: snapshot.state,
+        summary,
+      };
+    }
+
+    // Fallback: If snapshot has accounts with metrics, check if it's an openai-codex or time-window multi metric
+    const activeAccounts = snapshot.accounts.filter((a) => !a.disabled && !a.unavailable);
+    const firstAccount = activeAccounts[0];
+    if (firstAccount && firstAccount.metrics.length > 1) {
+      const windowMetrics = firstAccount.metrics.filter((m): m is Extract<Metric, { kind: "quota-window" }> => m.kind === "quota-window");
+      const allCodex = windowMetrics.length > 1 && windowMetrics.every((m) => m.label.startsWith("Codex"));
+      if (allCodex) {
+        const parts = windowMetrics.map((m) => {
+          const sub = m.label.replace(/^Codex\s+/, "");
+          const reset = m.resetAt ? relativeTime(m.resetAt) : undefined;
+          return `${sub} ${Math.round(m.remainingFraction * 100)}%${reset ? ` (${reset})` : ""}`;
+        });
+        return {
+          ...snapshot,
+          state: snapshot.state,
+          summary: `Codex ${parts.join(" · ")}`,
+        };
+      }
+    }
+
+    // Standard fallback: pick the lowest metric from all active accounts
+    const worst = activeAccounts
+      .flatMap((a) => a.metrics)
+      .filter((m): m is Extract<Metric, { kind: "quota-window" }> => m.kind === "quota-window")
+      .sort((a, b) => a.remainingFraction - b.remainingFraction)[0];
+
+    const reset = worst?.resetAt ? relativeTime(worst.resetAt) : undefined;
+    const summary = worst ? `${worst.label} ${Math.round(worst.remainingFraction * 100)}%${reset ? ` (${reset})` : ""}` : undefined;
+
+    return {
+      ...snapshot,
+      state: snapshot.accounts.length ? snapshot.state : "empty",
+      ...(summary ? { summary } : {}),
+    };
+  }
+}