@latentminds/pi-quotas 0.1.0

package/src/extensions/usage-status/index.ts

@@ -0,0 +1,163 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+} from "@mariozechner/pi-coding-agent";
+import {
+  QUOTAS_CONFIG_UPDATED_EVENT,
+  QUOTAS_EXTENSIONS_REGISTER_EVENT,
+  QUOTAS_EXTENSIONS_REQUEST_EVENT,
+  type QuotasConfigUpdatedPayload,
+  configLoader,
+} from "../../config.js";
+import {
+  fetchProviderQuotas,
+  formatResetTime,
+  isSupportedProvider,
+} from "../../lib/quotas.js";
+import {
+  assessWindow,
+} from "../../utils/quotas-severity.js";
+import { formatWindowStatus, type WindowStatus } from "./format-status.js";
+
+const EXTENSION_ID = "pi-quotas-usage";
+const REFRESH_INTERVAL_MS = 60_000;
+
+function formatStatus(ctx: ExtensionContext, windows: WindowStatus[]): string {
+  const theme = ctx.ui.theme;
+  return windows
+    .map((w) => {
+      const core = formatWindowStatus(theme, w);
+      const reset = w.resetsAt ? theme.fg("dim", ` (↺${formatResetTime(w.resetsAt)})`) : "";
+      return `${core}${reset}`;
+    })
+    .join(" ");
+}
+
+function createStatusRefresher() {
+  let refreshTimer: ReturnType<typeof setInterval> | undefined;
+  let activeContext: ExtensionContext | undefined;
+  let activeProvider: string | undefined;
+  let lastStatus: WindowStatus[] | undefined;
+  let inFlight = false;
+  let queued = false;
+
+  async function update(ctx: ExtensionContext): Promise<void> {
+    if (!ctx.hasUI || !activeProvider || !isSupportedProvider(activeProvider)) return;
+    if (inFlight) {
+      queued = true;
+      return;
+    }
+    inFlight = true;
+    try {
+      const result = await fetchProviderQuotas(ctx.modelRegistry.authStorage, activeProvider);
+      if (!result.success) {
+        ctx.ui.setStatus(EXTENSION_ID, ctx.ui.theme.fg("warning", "usage unavailable"));
+        return;
+      }
+      const windows: WindowStatus[] = result.data.windows.map((window) => ({
+        label: window.label,
+        usedPercent: window.usedPercent,
+        severity: assessWindow(window).severity,
+        resetsAt: window.resetsAt.toISOString(),
+        limited: window.limited ?? false,
+        isCurrency: window.isCurrency,
+        usedValue: window.usedValue,
+        limitValue: window.limitValue,
+      }));
+      lastStatus = windows;
+      ctx.ui.setStatus(EXTENSION_ID, formatStatus(ctx, windows));
+    } catch {
+      ctx.ui.setStatus(EXTENSION_ID, ctx.ui.theme.fg("warning", "usage unavailable"));
+    } finally {
+      inFlight = false;
+      if (queued) {
+        queued = false;
+        void update(ctx);
+      }
+    }
+  }
+
+  return {
+    async refreshFor(ctx: ExtensionContext): Promise<void> {
+      activeContext = ctx;
+      activeProvider = ctx.model?.provider;
+      if (!activeProvider || !isSupportedProvider(activeProvider)) {
+        ctx.ui.setStatus(EXTENSION_ID, undefined);
+        return;
+      }
+      await update(ctx);
+    },
+    start(): void {
+      if (refreshTimer) clearInterval(refreshTimer);
+      refreshTimer = setInterval(() => {
+        if (activeContext) void update(activeContext);
+      }, REFRESH_INTERVAL_MS);
+      refreshTimer.unref?.();
+    },
+    stop(ctx?: ExtensionContext): void {
+      if (refreshTimer) clearInterval(refreshTimer);
+      refreshTimer = undefined;
+      activeContext = undefined;
+      activeProvider = undefined;
+      lastStatus = undefined;
+      ctx?.ui.setStatus(EXTENSION_ID, undefined);
+    },
+    renderLast(ctx: ExtensionContext): boolean {
+      if (!lastStatus || !ctx.hasUI) return false;
+      ctx.ui.setStatus(EXTENSION_ID, formatStatus(ctx, lastStatus));
+      return true;
+    },
+  };
+}
+
+export default async function (pi: ExtensionAPI) {
+  await configLoader.load();
+  const refresher = createStatusRefresher();
+  let enabled = configLoader.getConfig().usageStatus;
+  let currentContext: ExtensionContext | undefined;
+
+  pi.events.on(QUOTAS_CONFIG_UPDATED_EVENT, (data: unknown) => {
+    enabled = (data as QuotasConfigUpdatedPayload).config.usageStatus;
+    if (!enabled) {
+      refresher.stop(currentContext);
+      return;
+    }
+    if (currentContext) {
+      refresher.start();
+      void refresher.refreshFor(currentContext);
+    }
+  });
+
+  pi.on("session_start", async (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) return;
+    refresher.start();
+    await refresher.refreshFor(ctx);
+  });
+
+  pi.on("turn_end", async (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) return;
+    await refresher.refreshFor(ctx);
+  });
+
+  pi.on("model_select", async (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) {
+      refresher.stop(ctx);
+      return;
+    }
+    await refresher.refreshFor(ctx);
+  });
+
+  pi.on("session_shutdown", async (_event, ctx) => {
+    currentContext = undefined;
+    refresher.stop(ctx);
+  });
+
+  pi.events.on(QUOTAS_EXTENSIONS_REQUEST_EVENT, () => {
+    if (configLoader.getConfig().usageStatus) {
+      pi.events.emit(QUOTAS_EXTENSIONS_REGISTER_EVENT, { feature: "usageStatus" });
+    }
+  });
+}

package/src/index.ts

@@ -0,0 +1,4 @@
+export { default as coreExtension } from "./extensions/core/index.js";
+export { default as commandQuotasExtension } from "./extensions/command-quotas/index.js";
+export { default as usageStatusExtension } from "./extensions/usage-status/index.js";
+export { default as quotaWarningsExtension } from "./extensions/quota-warnings/index.js";

package/src/lib/quotas.ts

@@ -0,0 +1,96 @@
+import type { AuthStorage } from "@mariozechner/pi-coding-agent";
+import { PROVIDER_FETCHERS } from "../providers/fetch.js";
+import type { QuotasResult, SupportedQuotaProvider } from "../types/quotas.js";
+
+export const SUPPORTED_PROVIDERS: SupportedQuotaProvider[] = [
+  "anthropic",
+  "openai-codex",
+  "github-copilot",
+];
+
+export const PROVIDER_LABELS: Record<SupportedQuotaProvider, string> = {
+  anthropic: "Anthropic",
+  "openai-codex": "OpenAI Codex",
+  "github-copilot": "GitHub Copilot",
+};
+
+const PROVIDER_TTLS_MS: Record<SupportedQuotaProvider, number> = {
+  anthropic: 5 * 60_000,
+  "openai-codex": 60_000,
+  "github-copilot": 5 * 60_000,
+};
+
+type CacheEntry = {
+  result?: QuotasResult;
+  fetchedAt?: number;
+  inFlight?: Promise<QuotasResult>;
+};
+
+const cache = new Map<SupportedQuotaProvider, CacheEntry>();
+
+export function isSupportedProvider(
+  provider: string | undefined,
+): provider is SupportedQuotaProvider {
+  return SUPPORTED_PROVIDERS.includes(provider as SupportedQuotaProvider);
+}
+
+export function clearQuotaCache(provider?: SupportedQuotaProvider): void {
+  if (provider) cache.delete(provider);
+  else cache.clear();
+}
+
+export async function fetchProviderQuotas(
+  authStorage: AuthStorage,
+  provider: SupportedQuotaProvider,
+  options?: { force?: boolean; signal?: AbortSignal },
+): Promise<QuotasResult> {
+  const entry = cache.get(provider) ?? {};
+  const now = Date.now();
+  const ttl = PROVIDER_TTLS_MS[provider];
+
+  if (!options?.force && entry.result && entry.fetchedAt && now - entry.fetchedAt < ttl) {
+    return entry.result;
+  }
+  if (!options?.force && entry.inFlight) return entry.inFlight;
+
+  const promise = PROVIDER_FETCHERS[provider](authStorage, options?.signal)
+    .then((result: QuotasResult) => {
+      cache.set(provider, { result, fetchedAt: Date.now() });
+      return result;
+    })
+    .finally(() => {
+      const current = cache.get(provider) ?? {};
+      delete current.inFlight;
+      cache.set(provider, current);
+    });
+
+  cache.set(provider, { ...entry, inFlight: promise });
+  return promise;
+}
+
+export async function fetchAllProviderQuotas(
+  authStorage: AuthStorage,
+  options?: { force?: boolean; signal?: AbortSignal },
+): Promise<Array<{ provider: SupportedQuotaProvider; result: QuotasResult }>> {
+  return Promise.all(
+    SUPPORTED_PROVIDERS.map(async (provider) => ({
+      provider,
+      result: await fetchProviderQuotas(authStorage, provider, options),
+    })),
+  );
+}
+
+export function formatResetTime(renewsAt: string): string {
+  const date = new Date(renewsAt);
+  const now = new Date();
+  const diffMs = date.getTime() - now.getTime();
+
+  if (diffMs <= 0) return "soon";
+
+  const diffHours = Math.ceil(diffMs / (1000 * 60 * 60));
+  const diffDays = Math.ceil(diffMs / (1000 * 60 * 60 * 24));
+
+  if (diffHours < 24) return `in ${diffHours}h`;
+  if (diffDays < 7) return `in ${diffDays}d`;
+  return date.toLocaleDateString("en-US", { month: "short", day: "numeric" });
+}

package/src/providers/fetch.test.ts

@@ -0,0 +1,137 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import {
+  fetchAnthropicQuotasWithToken,
+  fetchCodexQuotasWithToken,
+  fetchGitHubCopilotQuotasWithToken,
+} from "./fetch.js";
+
+const originalFetch = globalThis.fetch;
+
+afterEach(() => {
+  globalThis.fetch = originalFetch;
+  vi.restoreAllMocks();
+});
+
+describe("fetchAnthropicQuotasWithToken", () => {
+  it("returns config error when token missing", async () => {
+    const result = await fetchAnthropicQuotasWithToken(undefined);
+    expect(result).toMatchObject({
+      success: false,
+      error: { kind: "config" },
+    });
+  });
+
+  it("fetches and parses quota windows", async () => {
+    globalThis.fetch = vi.fn().mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          five_hour: { utilization: 21, resets_at: "2026-04-22T18:30:00Z" },
+          seven_day: { utilization: 9, resets_at: "2026-04-25T08:30:00Z" },
+        }),
+        { status: 200 },
+      ),
+    ) as any;
+
+    const result = await fetchAnthropicQuotasWithToken("token");
+    expect(result.success).toBe(true);
+    if (result.success) {
+      expect(result.data.provider).toBe("anthropic");
+      expect(result.data.windows).toHaveLength(2);
+    }
+  });
+});
+
+describe("fetchCodexQuotasWithToken", () => {
+  it("returns config error when account id missing", async () => {
+    const result = await fetchCodexQuotasWithToken("token", undefined);
+    expect(result).toMatchObject({
+      success: false,
+      error: { kind: "config" },
+    });
+  });
+
+  it("fetches and parses codex windows", async () => {
+    globalThis.fetch = vi.fn().mockResolvedValue(
+      new Response(
+        JSON.stringify({
+          rate_limit: {
+            primary_window: {
+              used_percent: 44,
+              reset_at: 1776880800,
+              limit_window_seconds: 18000,
+            },
+            secondary_window: {
+              used_percent: 12,
+              reset_at: 1777485600,
+              limit_window_seconds: 604800,
+            },
+          },
+        }),
+        { status: 200 },
+      ),
+    ) as any;
+
+    const result = await fetchCodexQuotasWithToken("token", "acct_123");
+    expect(result.success).toBe(true);
+    if (result.success) {
+      expect(result.data.provider).toBe("openai-codex");
+      expect(result.data.windows).toHaveLength(2);
+    }
+  });
+});
+
+describe("fetchGitHubCopilotQuotasWithToken", () => {
+  it("exchanges token then fetches usage on happy path", async () => {
+    globalThis.fetch = vi
+      .fn()
+      .mockResolvedValueOnce(
+        new Response(JSON.stringify({ token: "copilot-token" }), { status: 200 }),
+      )
+      .mockResolvedValueOnce(
+        new Response(
+          JSON.stringify({
+            quota_reset_date: "2026-05-01T00:00:00Z",
+            quota_snapshots: {
+              premium_interactions: {
+                entitlement: 300,
+                remaining: 240,
+                percent_remaining: 80,
+              },
+            },
+          }),
+          { status: 200 },
+        ),
+      ) as any;
+
+    const result = await fetchGitHubCopilotQuotasWithToken("gh-token");
+    expect(result.success).toBe(true);
+    if (result.success) {
+      expect(result.data.provider).toBe("github-copilot");
+      expect(result.data.windows).toHaveLength(1);
+    }
+    expect(globalThis.fetch).toHaveBeenCalledTimes(2);
+  });
+
+  it("falls back to direct token when exchange returns 401", async () => {
+    globalThis.fetch = vi
+      .fn()
+      .mockResolvedValueOnce(
+        new Response(JSON.stringify({ message: "Bad credentials" }), { status: 401 }),
+      )
+      .mockResolvedValueOnce(
+        new Response(
+          JSON.stringify({
+            quota_reset_date: "2026-05-01T00:00:00Z",
+            quota_snapshots: {
+              premium_interactions: { entitlement: 300, remaining: 293 },
+            },
+          }),
+          { status: 200 },
+        ),
+      ) as any;
+
+    const result = await fetchGitHubCopilotQuotasWithToken("gh-token");
+    expect(result.success).toBe(true);
+    expect(globalThis.fetch).toHaveBeenCalledTimes(2);
+  });
+});

package/src/providers/fetch.ts

@@ -0,0 +1,245 @@
+import { execFileSync } from "node:child_process";
+import { readFileSync } from "node:fs";
+import { homedir } from "node:os";
+import { join } from "node:path";
+import type { AuthStorage } from "@mariozechner/pi-coding-agent";
+import type { QuotasResult, SupportedQuotaProvider } from "../types/quotas.js";
+import {
+  parseAnthropicUsage,
+  parseCodexUsage,
+  parseGitHubCopilotUsage,
+} from "./providers.js";
+
+const FETCH_TIMEOUT_MS = 15_000;
+const COPILOT_VERSION = "0.35.0";
+const EDITOR_VERSION = "vscode/1.107.0";
+
+function isTimeoutReason(reason: unknown): boolean {
+  return (
+    (reason instanceof DOMException && reason.name === "TimeoutError") ||
+    (reason instanceof Error && reason.name === "TimeoutError")
+  );
+}
+
+async function providerAccessToken(
+  authStorage: AuthStorage,
+  provider: string,
+): Promise<string | undefined> {
+  return authStorage.getApiKey(provider);
+}
+
+function codexAccountId(authStorage: AuthStorage): string | undefined {
+  const credential = authStorage.get("openai-codex") as any;
+  if (typeof credential?.accountId === "string") return credential.accountId;
+  try {
+    const authPath = join(homedir(), ".codex", "auth.json");
+    const data = JSON.parse(readFileSync(authPath, "utf8")) as any;
+    return data?.tokens?.account_id ?? data?.tokens?.accountId;
+  } catch {
+    return undefined;
+  }
+}
+
+type FetchJsonResult =
+  | { ok: true; data: any }
+  | {
+      ok: false;
+      status?: number;
+      message: string;
+      kind: "timeout" | "cancelled" | "http" | "network";
+    };
+
+async function fetchJson(
+  url: string,
+  init: RequestInit,
+  signal?: AbortSignal,
+): Promise<FetchJsonResult> {
+  const signals: AbortSignal[] = [AbortSignal.timeout(FETCH_TIMEOUT_MS)];
+  if (signal) signals.push(signal);
+  const combined = AbortSignal.any(signals);
+
+  try {
+    const response = await fetch(url, { ...init, signal: combined });
+    if (!response.ok) {
+      const body = await response.text().catch(() => "");
+      return {
+        ok: false,
+        status: response.status,
+        message: body || response.statusText || `HTTP ${response.status}`,
+        kind: "http",
+      };
+    }
+    return { ok: true, data: await response.json() };
+  } catch (err: unknown) {
+    const isAbort =
+      combined.aborted ||
+      (err instanceof DOMException && err.name === "AbortError");
+    if (isAbort) {
+      if (isTimeoutReason(combined.reason)) {
+        return { ok: false, message: "Request timed out", kind: "timeout" };
+      }
+      return { ok: false, message: "Request cancelled", kind: "cancelled" };
+    }
+    const message = err instanceof Error ? err.message : "Unknown error";
+    return { ok: false, message, kind: "network" };
+  }
+}
+
+function success(provider: SupportedQuotaProvider, windows: ReturnType<typeof parseAnthropicUsage>): QuotasResult {
+  return { success: true, data: { provider, windows } };
+}
+
+function failure(message: string, kind: "cancelled" | "timeout" | "config" | "http" | "network"): QuotasResult {
+  return { success: false, error: { message, kind } };
+}
+
+export async function fetchAnthropicQuotasWithToken(
+  accessToken: string | undefined,
+  signal?: AbortSignal,
+): Promise<QuotasResult> {
+  if (!accessToken) return failure("No Anthropic OAuth token found", "config");
+  const result = await fetchJson(
+    "https://api.anthropic.com/api/oauth/usage",
+    {
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "anthropic-beta": "oauth-2025-04-20",
+        Accept: "application/json",
+      },
+    },
+    signal,
+  );
+  if (!result.ok) return failure(result.message, result.kind);
+  return success("anthropic", parseAnthropicUsage(result.data));
+}
+
+export async function fetchCodexQuotasWithToken(
+  accessToken: string | undefined,
+  accountId: string | undefined,
+  signal?: AbortSignal,
+): Promise<QuotasResult> {
+  if (!accessToken) return failure("No Codex access token found", "config");
+  if (!accountId) return failure("No Codex account id found", "config");
+  const result = await fetchJson(
+    "https://chatgpt.com/backend-api/wham/usage",
+    {
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "ChatGPT-Account-Id": accountId,
+        Accept: "application/json",
+        Origin: "https://chatgpt.com",
+        Referer: "https://chatgpt.com/",
+        "User-Agent": "Mozilla/5.0",
+      },
+    },
+    signal,
+  );
+  if (!result.ok) return failure(result.message, result.kind);
+  return success("openai-codex", parseCodexUsage(result.data));
+}
+
+function copilotHeaders(authHeader: string): Record<string, string> {
+  return {
+    Accept: "application/json",
+    Authorization: authHeader,
+    "User-Agent": `GitHubCopilotChat/${COPILOT_VERSION}`,
+    "Editor-Version": EDITOR_VERSION,
+    "Editor-Plugin-Version": `copilot-chat/${COPILOT_VERSION}`,
+    "Copilot-Integration-Id": "vscode-chat",
+    "Content-Type": "application/json",
+  };
+}
+
+/**
+ * Try to get a token from `gh auth token` CLI as fallback when the Pi-stored
+ * OAuth token is stale or the token exchange returns 401.
+ */
+function ghCliToken(): string | undefined {
+  try {
+    return execFileSync("gh", ["auth", "token"], {
+      timeout: 5000,
+      encoding: "utf8",
+      stdio: ["ignore", "pipe", "ignore"],
+    }).trim() || undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+async function tryGitHubUserEndpoint(
+  authHeader: string,
+  signal?: AbortSignal,
+): Promise<FetchJsonResult> {
+  return fetchJson(
+    "https://api.github.com/copilot_internal/user",
+    { headers: copilotHeaders(authHeader) },
+    signal,
+  );
+}
+
+export async function fetchGitHubCopilotQuotasWithToken(
+  accessToken: string | undefined,
+  signal?: AbortSignal,
+): Promise<QuotasResult> {
+  if (!accessToken) return failure("No GitHub Copilot OAuth token found", "config");
+
+  // 1) Try Copilot token exchange with stored Pi token
+  const exchange = await fetchJson(
+    "https://api.github.com/copilot_internal/v2/token",
+    { headers: copilotHeaders(`Bearer ${accessToken}`) },
+    signal,
+  );
+
+  if (exchange.ok && exchange.data?.token) {
+    const usage = await tryGitHubUserEndpoint(`Bearer ${exchange.data.token}`, signal);
+    if (usage.ok) return success("github-copilot", parseGitHubCopilotUsage(usage.data));
+  }
+
+  // 2) Try stored token directly
+  const directUsage = await tryGitHubUserEndpoint(`token ${accessToken}`, signal);
+  if (directUsage.ok) return success("github-copilot", parseGitHubCopilotUsage(directUsage.data));
+
+  // 3) Fallback: gh CLI token
+  const cliToken = ghCliToken();
+  if (cliToken && cliToken !== accessToken) {
+    const cliUsage = await tryGitHubUserEndpoint(`token ${cliToken}`, signal);
+    if (cliUsage.ok) return success("github-copilot", parseGitHubCopilotUsage(cliUsage.data));
+    return failure(cliUsage.message, cliUsage.kind);
+  }
+
+  return failure(directUsage.message, directUsage.kind);
+}
+
+export async function fetchAnthropicQuotas(
+  authStorage: AuthStorage,
+  signal?: AbortSignal,
+): Promise<QuotasResult> {
+  return fetchAnthropicQuotasWithToken(await providerAccessToken(authStorage, "anthropic"), signal);
+}
+
+export async function fetchCodexQuotas(
+  authStorage: AuthStorage,
+  signal?: AbortSignal,
+): Promise<QuotasResult> {
+  return fetchCodexQuotasWithToken(
+    await providerAccessToken(authStorage, "openai-codex"),
+    codexAccountId(authStorage),
+    signal,
+  );
+}
+
+export async function fetchGitHubCopilotQuotas(
+  authStorage: AuthStorage,
+  signal?: AbortSignal,
+): Promise<QuotasResult> {
+  return fetchGitHubCopilotQuotasWithToken(
+    await providerAccessToken(authStorage, "github-copilot"),
+    signal,
+  );
+}
+
+export const PROVIDER_FETCHERS = {
+  anthropic: fetchAnthropicQuotas,
+  "openai-codex": fetchCodexQuotas,
+  "github-copilot": fetchGitHubCopilotQuotas,
+} as const;