@diousk/pi-usage-fast 0.1.0

package/src/settings.ts

@@ -0,0 +1,366 @@
+import { randomUUID } from "node:crypto";
+import { constants } from "node:fs";
+import { chmod, mkdir, open, rename, rm, writeFile } from "node:fs/promises";
+import { basename, dirname, join } from "node:path";
+import { getAgentDir } from "@earendil-works/pi-coding-agent";
+import { isFireworksAccountId } from "./providers/fireworks.js";
+import { isBoundedTargetId } from "./usage-targets.js";
+
+export const USAGE_SETTINGS_FILE = "pi-usage.json";
+export const MAX_USAGE_SETTINGS_BYTES = 64 * 1024;
+
+export type CodexStatusPercentage = "remaining" | "used";
+
+export interface UsageSettings {
+  codexFastMode: boolean;
+  codexStatusResetCountdown: boolean;
+  codexStatusPercentage: CodexStatusPercentage;
+  selectedTargets: Record<string, string>;
+}
+
+export const DEFAULT_USAGE_SETTINGS: Readonly<UsageSettings> = Object.freeze({
+  codexFastMode: false,
+  codexStatusResetCountdown: true,
+  codexStatusPercentage: "remaining",
+  selectedTargets: Object.freeze({}),
+});
+
+export interface UsageSettingsState {
+  kind: "missing" | "loaded" | "invalid";
+  path: string;
+  settings: UsageSettings;
+  document?: Record<string, unknown>;
+  issue?: string;
+}
+
+export type UsageTargetPublicationCheck = () => Promise<void>;
+
+export interface UsageSettingsRuntime {
+  get(): Readonly<UsageSettingsState>;
+  reload(signal?: AbortSignal): Promise<Readonly<UsageSettingsState>>;
+  update(patch: Partial<UsageSettings>, signal?: AbortSignal): Promise<Readonly<UsageSettingsState>>;
+  updateSelectedTarget(
+    providerId: string,
+    targetId: string,
+    signal?: AbortSignal,
+    checkPublishedSelection?: UsageTargetPublicationCheck,
+  ): Promise<Readonly<UsageSettingsState>>;
+  flush(): Promise<void>;
+}
+
+interface UsageSettingsFileOperations {
+  rename: typeof rename;
+  writeFile: typeof writeFile;
+}
+
+interface UsageSettingsRuntimeOptions {
+  operations?: Partial<UsageSettingsFileOperations>;
+  path?: string;
+}
+
+export function usageSettingsPath(): string {
+  return join(getAgentDir(), USAGE_SETTINGS_FILE);
+}
+
+export function normalizeUsageSettings(value: unknown): UsageSettings | undefined {
+  if (!isRecord(value)) return undefined;
+  if (Object.hasOwn(value, "codexFastMode") && typeof value.codexFastMode !== "boolean") {
+    return undefined;
+  }
+  if (Object.hasOwn(value, "codexStatusResetCountdown") && typeof value.codexStatusResetCountdown !== "boolean") {
+    return undefined;
+  }
+  if (
+    Object.hasOwn(value, "codexStatusPercentage") &&
+    value.codexStatusPercentage !== "remaining" &&
+    value.codexStatusPercentage !== "used"
+  ) {
+    return undefined;
+  }
+  if (Object.hasOwn(value, "fireworksAccountId") && !isFireworksAccountId(value.fireworksAccountId)) {
+    return undefined;
+  }
+  const selectedTargets = normalizeSelectedTargets(value.selectedTargets);
+  if (Object.hasOwn(value, "selectedTargets") && !selectedTargets) return undefined;
+  const effectiveTargets = { ...(selectedTargets ?? {}) };
+  if (!effectiveTargets.fireworks && isFireworksAccountId(value.fireworksAccountId)) {
+    effectiveTargets.fireworks = value.fireworksAccountId;
+  }
+  return {
+    codexFastMode:
+      typeof value.codexFastMode === "boolean" ? value.codexFastMode : DEFAULT_USAGE_SETTINGS.codexFastMode,
+    codexStatusResetCountdown:
+      typeof value.codexStatusResetCountdown === "boolean"
+        ? value.codexStatusResetCountdown
+        : DEFAULT_USAGE_SETTINGS.codexStatusResetCountdown,
+    codexStatusPercentage:
+      value.codexStatusPercentage === "used" ? "used" : DEFAULT_USAGE_SETTINGS.codexStatusPercentage,
+    selectedTargets: effectiveTargets,
+  };
+}
+
+export async function loadUsageSettings(path = usageSettingsPath(), signal?: AbortSignal): Promise<UsageSettingsState> {
+  throwIfAborted(signal);
+  try {
+    const handle = await open(path, constants.O_RDONLY | constants.O_NOFOLLOW);
+    let text: string;
+    try {
+      const stats = await handle.stat();
+      throwIfAborted(signal);
+      if (!stats.isFile()) throw new Error("settings path is not a regular file");
+      if (stats.size > MAX_USAGE_SETTINGS_BYTES) {
+        throw new Error("settings file exceeds 64 KiB");
+      }
+      text = await handle.readFile("utf8");
+    } finally {
+      await handle.close();
+    }
+    throwIfAborted(signal);
+    const document = JSON.parse(text) as unknown;
+    const settings = normalizeUsageSettings(document);
+    if (!settings || !isRecord(document)) throw new Error("invalid settings shape");
+    return { kind: "loaded", path, settings, document };
+  } catch (error) {
+    if (signal?.aborted) throw error;
+    if (isNodeError(error) && error.code === "ENOENT") {
+      return {
+        kind: "missing",
+        path,
+        settings: { ...DEFAULT_USAGE_SETTINGS },
+        document: {},
+      };
+    }
+    return {
+      kind: "invalid",
+      path,
+      settings: { ...DEFAULT_USAGE_SETTINGS },
+      issue:
+        isNodeError(error) && error.code === "ELOOP"
+          ? "symbolic links are not accepted"
+          : error instanceof Error
+            ? error.message
+            : String(error),
+    };
+  }
+}
+
+export function createUsageSettingsRuntime(options: UsageSettingsRuntimeOptions | string = {}): UsageSettingsRuntime {
+  const path = typeof options === "string" ? options : (options.path ?? usageSettingsPath());
+  const operations: UsageSettingsFileOperations = {
+    rename,
+    writeFile,
+    ...(typeof options === "string" ? undefined : options.operations),
+  };
+  let state: UsageSettingsState = {
+    kind: "missing",
+    path,
+    settings: { ...DEFAULT_USAGE_SETTINGS },
+    document: {},
+  };
+  let queue = Promise.resolve();
+  const enqueue = <T>(operation: () => Promise<T>): Promise<T> => {
+    const result = queue.then(operation, operation);
+    queue = result.then(
+      () => undefined,
+      () => undefined,
+    );
+    return result;
+  };
+  return {
+    get: () => structuredClone(state),
+    reload: (signal) =>
+      enqueue(async () => {
+        const loaded = await loadUsageSettings(path, signal);
+        state = loaded;
+        return structuredClone(state);
+      }),
+    update: (patch, signal) =>
+      enqueue(async () => {
+        const saved = await saveUsageSettingsPatch(path, patch, operations, signal);
+        state = saved;
+        return structuredClone(state);
+      }),
+    updateSelectedTarget: (providerId, targetId, signal, checkPublishedSelection) =>
+      enqueue(async () => {
+        const transaction = await saveUsageTargetSelection(path, providerId, targetId, operations, signal);
+        try {
+          await checkPublishedSelection?.();
+          throwIfAborted(signal);
+        } catch (error) {
+          try {
+            await restoreUsageSettingsState(path, transaction.saved, transaction.previous, operations);
+            state = transaction.previous;
+          } catch (rollbackError) {
+            state = await loadUsageSettings(path);
+            throw new AggregateError(
+              [error, rollbackError],
+              "Target selection changed after publication and pi-usage.json rollback failed",
+            );
+          }
+          throw error;
+        }
+        state = transaction.saved;
+        return structuredClone(state);
+      }),
+    flush: () => queue,
+  };
+}
+
+async function saveUsageSettingsPatch(
+  path: string,
+  patch: Partial<UsageSettings>,
+  operations: UsageSettingsFileOperations,
+  signal?: AbortSignal,
+): Promise<UsageSettingsState> {
+  return saveUsageSettingsDocument(
+    path,
+    (document) => {
+      for (const [key, value] of Object.entries(patch)) {
+        if (value === undefined) delete document[key];
+        else document[key] = value;
+      }
+    },
+    operations,
+    signal,
+  );
+}
+
+async function saveUsageTargetSelection(
+  path: string,
+  providerId: string,
+  targetId: string,
+  operations: UsageSettingsFileOperations,
+  signal?: AbortSignal,
+): Promise<{ saved: UsageSettingsState; previous: UsageSettingsState }> {
+  if (!isProviderId(providerId) || !isBoundedTargetId(targetId)) {
+    throw new Error("Refusing to save an invalid usage target selection");
+  }
+  const previous = await loadUsageSettings(path, signal);
+  const saved = await saveUsageSettingsDocument(
+    path,
+    (document) => {
+      document.selectedTargets = {
+        ...(normalizeSelectedTargets(document.selectedTargets) ?? {}),
+        [providerId]: targetId,
+      };
+      if (providerId === "fireworks") delete document.fireworksAccountId;
+    },
+    operations,
+    signal,
+    previous,
+  );
+  return { saved, previous };
+}
+
+async function saveUsageSettingsDocument(
+  path: string,
+  mutate: (document: Record<string, unknown>) => void,
+  operations: UsageSettingsFileOperations,
+  signal?: AbortSignal,
+  expected?: UsageSettingsState,
+): Promise<UsageSettingsState> {
+  const latest = await loadUsageSettings(path, signal);
+  if (latest.kind === "invalid") {
+    throw new Error("Cannot overwrite an invalid pi-usage.json; repair it and reload first");
+  }
+  if (expected && !sameUsageSettingsDocument(latest, expected)) {
+    throw new Error("pi-usage.json changed while saving; retry the action");
+  }
+  const document = { ...latest.document };
+  mutate(document);
+  const settings = normalizeUsageSettings(document);
+  if (!settings) throw new Error("Refusing to save invalid pi-usage settings");
+  const directory = dirname(path);
+  const temporaryPath = join(directory, `.${basename(path)}.${randomUUID()}.tmp`);
+  await mkdir(directory, { recursive: true, mode: 0o700 });
+  throwIfAborted(signal);
+  try {
+    await operations.writeFile(temporaryPath, `${JSON.stringify(document, null, 2)}\n`, {
+      encoding: "utf8",
+      flag: "wx",
+      mode: 0o600,
+    });
+    if (process.platform !== "win32") await chmodPrivate(temporaryPath);
+    throwIfAborted(signal);
+    const current = await loadUsageSettings(path, signal);
+    if (
+      current.kind === "invalid" ||
+      current.kind !== latest.kind ||
+      JSON.stringify(current.document) !== JSON.stringify(latest.document)
+    ) {
+      throw new Error("pi-usage.json changed while saving; retry the action");
+    }
+    throwIfAborted(signal);
+    await operations.rename(temporaryPath, path);
+  } finally {
+    await rm(temporaryPath, { force: true }).catch(() => undefined);
+  }
+  return { kind: "loaded", path, settings, document };
+}
+
+async function restoreUsageSettingsState(
+  path: string,
+  published: UsageSettingsState,
+  previous: UsageSettingsState,
+  operations: UsageSettingsFileOperations,
+): Promise<void> {
+  if (previous.kind === "missing") {
+    const current = await loadUsageSettings(path);
+    if (!sameUsageSettingsDocument(current, published)) {
+      throw new Error("pi-usage.json changed before target selection rollback");
+    }
+    await rm(path);
+    return;
+  }
+  if (previous.kind !== "loaded" || !previous.document) {
+    throw new Error("Cannot restore invalid prior pi-usage.json settings");
+  }
+  await saveUsageSettingsDocument(
+    path,
+    (document) => {
+      for (const key of Object.keys(document)) delete document[key];
+      Object.assign(document, previous.document);
+    },
+    operations,
+    undefined,
+    published,
+  );
+}
+
+function sameUsageSettingsDocument(
+  left: Pick<UsageSettingsState, "kind" | "document">,
+  right: Pick<UsageSettingsState, "kind" | "document">,
+): boolean {
+  return left.kind === right.kind && JSON.stringify(left.document) === JSON.stringify(right.document);
+}
+
+async function chmodPrivate(path: string): Promise<void> {
+  await chmod(path, 0o600);
+}
+
+function throwIfAborted(signal?: AbortSignal): void {
+  if (signal?.aborted) throw new DOMException("Settings operation aborted", "AbortError");
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function isNodeError(error: unknown): error is NodeJS.ErrnoException {
+  return error instanceof Error && "code" in error;
+}
+
+function normalizeSelectedTargets(value: unknown): Record<string, string> | undefined {
+  if (value === undefined) return {};
+  if (!isRecord(value)) return undefined;
+  const targets: Record<string, string> = {};
+  for (const [providerId, targetId] of Object.entries(value)) {
+    if (!isProviderId(providerId) || !isBoundedTargetId(targetId)) return undefined;
+    targets[providerId] = targetId;
+  }
+  return targets;
+}
+
+function isProviderId(value: string): boolean {
+  return /^[A-Za-z0-9][A-Za-z0-9._~-]{0,127}$/u.test(value);
+}

package/src/types.ts

@@ -0,0 +1,235 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+
+export type PiModel = NonNullable<ExtensionContext["model"]>;
+export type UsageModel = Pick<PiModel, "id" | "name" | "provider">;
+
+export type UsageSemanticsKind = "consumer-subscription" | "api-key" | "project";
+export type UsageUnit = "percent" | "usd" | "currency" | "count";
+export type UsageDisplayState = "current" | "configured";
+
+export interface UsageSemantics {
+  kind: UsageSemanticsKind;
+  label: string;
+}
+
+export interface UsageBucket {
+  id: string;
+  label: string;
+  groupId?: string;
+  groupLabel?: string;
+  modelKeys?: string[];
+  used?: number;
+  remaining?: number;
+  limit?: number;
+  unit: UsageUnit;
+  period?: string;
+  windowMinutes?: number;
+  resetsAt?: number;
+}
+
+export interface UsageMetric {
+  id: string;
+  label: string;
+  value: number | string;
+  unit?: UsageUnit;
+  currency?: string;
+}
+
+export interface UsageReport {
+  providerId: string;
+  providerName: string;
+  capturedAt: number;
+  source: string;
+  semantics: UsageSemantics;
+  accountLabel?: string;
+  buckets: UsageBucket[];
+  metrics: UsageMetric[];
+  notes?: string[];
+}
+
+export interface ResolvedUsageAuth {
+  apiKey?: string;
+  headers: Record<string, string>;
+  fingerprint: string;
+  secrets: string[];
+  model: PiModel;
+  auth?: {
+    apiKey?: string;
+    headers?: Record<string, string | null>;
+    baseUrl?: string;
+  };
+  env?: Record<string, string>;
+  source?: string;
+  effectiveBaseUrl?: string;
+}
+
+export interface UsageQuerySettings {
+  fireworksAccountId?: string;
+}
+
+export type UsageRequestGuard = () => Promise<void>;
+
+export interface UsageProviderTarget {
+  id: string;
+  label: string;
+  description?: string;
+}
+
+export interface UsageTargetResolver {
+  singularLabel: string;
+  pluralLabel: string;
+  list(
+    auth: ResolvedUsageAuth,
+    signal: AbortSignal,
+    timeoutMs: number,
+    guard: UsageRequestGuard,
+  ): Promise<readonly UsageProviderTarget[]>;
+}
+
+export interface UsageProviderAdapter {
+  id: string;
+  displayName: string;
+  semantics: UsageSemantics;
+  publishesStatusline?: boolean;
+  /** Invalidate the matching ready report when the latest query fails; default preserves it. */
+  invalidateCacheOnFailure?: boolean;
+  targets?: UsageTargetResolver;
+  query(
+    auth: ResolvedUsageAuth,
+    signal: AbortSignal,
+    timeoutMs: number,
+    guard?: UsageRequestGuard,
+    targetId?: string,
+  ): Promise<UsageReport>;
+}
+
+export type ProviderUsageState =
+  | {
+      providerId: string;
+      providerName: string;
+      displayState: UsageDisplayState;
+      status: "ready";
+      report: UsageReport;
+    }
+  | {
+      providerId: string;
+      providerName: string;
+      displayState: UsageDisplayState;
+      status: "selection-required";
+      singularLabel: string;
+      pluralLabel: string;
+      choices: readonly UsageProviderTarget[];
+    }
+  | {
+      providerId: string;
+      providerName: string;
+      displayState: UsageDisplayState;
+      status: "unsupported" | "auth-unavailable" | "query-failed";
+      message: string;
+    };
+
+export type BasetenBillingUsagePayload = {
+  dedicated_usage?: unknown;
+  model_apis_usage?: unknown;
+  training_usage?: unknown;
+};
+
+export type DeepSeekBalancePayload = {
+  is_available?: unknown;
+  balance_infos?: unknown;
+};
+
+export type FireworksAccountsPayload = {
+  accounts?: unknown;
+  nextPageToken?: unknown;
+};
+
+export type FireworksBillingSummaryPayload = {
+  lineItems?: unknown;
+};
+
+export type GitHubCopilotUsagePayload = {
+  login?: unknown;
+  copilot_plan?: unknown;
+  access_type_sku?: unknown;
+  limited_user_quotas?: unknown;
+  limited_user_reset_date?: unknown;
+  monthly_quotas?: unknown;
+  quota_reset_date?: unknown;
+  quota_reset_date_utc?: unknown;
+  quota_snapshots?: unknown;
+};
+
+export type OpenRouterKeyPayload = {
+  data?: unknown;
+};
+
+export type MiniMaxUsagePayload = {
+  available_amount?: unknown;
+  balance_alert_switch?: unknown;
+  balance_alert_threshold?: unknown;
+  base_resp?: unknown;
+  cash_balance?: unknown;
+  credit_balance?: unknown;
+  model_remains?: unknown;
+  owed_amount?: unknown;
+  voucher_balance?: unknown;
+};
+
+export type MoonshotBalancePayload = {
+  code?: unknown;
+  data?: unknown;
+  scode?: unknown;
+  status?: unknown;
+};
+
+export type VercelAIGatewayCreditsPayload = {
+  balance?: unknown;
+  total_used?: unknown;
+};
+
+export type OpenCodeZenPayload = {
+  usage?: unknown;
+};
+
+export type ZaiQuotaPayload = {
+  error?: unknown;
+  code?: unknown;
+  success?: unknown;
+  data?: unknown;
+};
+
+export type ZaiSubscriptionPayload = {
+  error?: unknown;
+  code?: unknown;
+  success?: unknown;
+  data?: unknown;
+};
+
+export interface ZaiPlanInfo {
+  name: string;
+  renewsAt?: string;
+}
+
+export type KimiCodingUsagePayload = {
+  usage?: unknown;
+  limits?: unknown;
+  boosterWallet?: unknown;
+};
+
+export type XaiUserPayload = {
+  userId?: unknown;
+  subscriptionTier?: unknown;
+};
+
+export type XaiBillingPayload = {
+  config?: unknown;
+};
+
+export type CodexBackendPayload = {
+  plan_type?: unknown;
+  rate_limit?: unknown;
+  additional_rate_limits?: unknown;
+  credits?: unknown;
+  rate_limit_reset_credits?: unknown;
+};

package/src/usage-helpers.ts

@@ -0,0 +1,40 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { sanitizeDisplayText } from "./core.js";
+import { providerIsConfigured, usageAdapters } from "./query.js";
+import type { PiModel, UsageProviderAdapter } from "./types.js";
+
+export function configuredAdapters(ctx: ExtensionContext): UsageProviderAdapter[] {
+  return usageAdapters().filter(
+    (adapter) => adapter.id === ctx.model?.provider || providerIsConfigured(ctx, adapter.id),
+  );
+}
+
+export function providerDisplayName(ctx: ExtensionContext, providerId: string): string {
+  try {
+    return sanitizeDisplayText(ctx.modelRegistry.getProviderDisplayName(providerId), 80);
+  } catch {
+    return sanitizeDisplayText(providerId, 80);
+  }
+}
+
+export function setBoundedMap<T>(map: Map<string, T>, key: string, value: T, limit: number): void {
+  map.delete(key);
+  while (map.size >= limit) {
+    const oldest = map.keys().next().value;
+    if (oldest === undefined) break;
+    map.delete(oldest);
+  }
+  map.set(key, value);
+}
+
+export function modelIdentity(model: PiModel | undefined): string | undefined {
+  return model ? `${model.provider}/${model.id}` : undefined;
+}
+
+export function isAbortError(error: unknown): boolean {
+  return error instanceof Error && error.name === "AbortError";
+}
+
+export function isTimeoutError(error: unknown): boolean {
+  return error instanceof Error && error.name === "TimeoutError";
+}