pi-ark-usage 0.1.0

package/extensions/lib/arkcli.ts

@@ -0,0 +1,129 @@
+/**
+ * `arkcli` process wrapper: locate the binary, invoke `arkcli usage plan`,
+ * parse and validate its JSON output.
+ */
+
+import { execFile } from "node:child_process";
+import { promisify } from "node:util";
+import type { ArkPlanOutput, ArkPlanItem, ArkPeriod } from "./types.js";
+
+const execFileAsync = promisify(execFile);
+
+const DEFAULT_TIMEOUT_MS = 10_000;
+const MAX_JSON_BYTES = 1024 * 1024;
+
+const PRODUCT_IDS = new Set([
+  "coding-plan",
+  "agent-plan",
+  "coding-plan-team",
+  "agent-plan-team",
+]);
+
+/** Locate arkcli: explicit env override, else plain name (resolved via PATH). */
+export function arkcliBin(): string {
+  return process.env.PI_ARK_USAGE_ARKCLI || "arkcli";
+}
+
+export async function checkArkcliAvailable(): Promise<boolean> {
+  try {
+    await execFileAsync(arkcliBin(), ["--help"], { timeout: 5_000 });
+    return true;
+  } catch {
+    // --help may exit non-zero in odd wrappers; fall back to version check.
+    try {
+      await execFileAsync(arkcliBin(), ["version"], { timeout: 5_000 });
+      return true;
+    } catch {
+      return false;
+    }
+  }
+}
+
+function isFiniteNumber(v: unknown): v is number {
+  return typeof v === "number" && Number.isFinite(v);
+}
+
+function validatePeriod(raw: unknown): ArkPeriod | null {
+  if (!raw || typeof raw !== "object") return null;
+  const p = raw as Record<string, unknown>;
+  if (typeof p.label !== "string" || p.label.length === 0) return null;
+  if (!isFiniteNumber(p.percent) || p.percent < 0) return null;
+  const out: ArkPeriod = {
+    label: p.label,
+    percent: Math.min(p.percent, 100),
+  };
+  if (typeof p.reset_at === "string" && p.reset_at.length > 0) {
+    out.reset_at = p.reset_at;
+  }
+  return out;
+}
+
+function validateItem(raw: unknown): ArkPlanItem | null {
+  if (!raw || typeof raw !== "object") return null;
+  const it = raw as Record<string, unknown>;
+  if (typeof it.product !== "string" || !PRODUCT_IDS.has(it.product)) return null;
+  if (typeof it.edition !== "string") return null;
+  const out: ArkPlanItem = {
+    product: it.product,
+    edition: it.edition,
+    subscribed: it.subscribed === true,
+  };
+  if (typeof it.error === "string") out.error = it.error;
+  if (isFiniteNumber(it.updated_at)) out.updated_at = it.updated_at;
+  if (Array.isArray(it.periods)) {
+    const periods = it.periods
+      .map(validatePeriod)
+      .filter((p): p is ArkPeriod => p !== null);
+    if (periods.length > 0) out.periods = periods;
+  }
+  return out;
+}
+
+/** Parse and validate raw arkcli JSON. Throws on malformed payload. */
+export function parsePlanOutput(raw: string): ArkPlanOutput {
+  if (raw.length > MAX_JSON_BYTES) throw new Error("arkcli output too large");
+  const json: unknown = JSON.parse(raw);
+  if (!json || typeof json !== "object") throw new Error("invalid arkcli output");
+  const root = json as Record<string, unknown>;
+  if (!Array.isArray(root.items)) throw new Error("invalid arkcli output: missing items");
+  const items = root.items
+    .map(validateItem)
+    .filter((it): it is ArkPlanItem => it !== null);
+  if (items.length === 0) throw new Error("arkcli returned no usable items");
+  const out: ArkPlanOutput = { items };
+  if (root.viewer && typeof root.viewer === "object") {
+    out.viewer = root.viewer as ArkPlanOutput["viewer"];
+  }
+  return out;
+}
+
+export type FetchPlanOptions = {
+  signal?: AbortSignal;
+  timeoutMs?: number;
+  /** specific product id; undefined = auto-discover all subscriptions */
+  product?: string;
+  seat?: string;
+};
+
+/**
+ * Run `arkcli usage plan --format json`.
+ * Authentication is fully handled by arkcli itself (SSO profiles etc.).
+ */
+export async function fetchPlan(options: FetchPlanOptions = {}): Promise<ArkPlanOutput> {
+  const args = ["usage", "plan", "--format", "json"];
+  if (options.product) {
+    args.push("--product", options.product);
+  }
+  if (options.seat) {
+    args.push("--seat", options.seat);
+  }
+
+  const { stdout } = await execFileAsync(arkcliBin(), args, {
+    timeout: options.timeoutMs ?? DEFAULT_TIMEOUT_MS,
+    maxBuffer: MAX_JSON_BYTES,
+    signal: options.signal,
+    // arkcli may print warnings on stderr; they don't fail the call unless exit code != 0.
+    windowsHide: true,
+  });
+  return parsePlanOutput(stdout);
+}

package/extensions/lib/cache.ts

@@ -0,0 +1,196 @@
+/**
+ * Durable cache: persists settings and latest snapshots across restarts.
+ *
+ * File: ~/.pi/agent/pi-ark-usage/cache.json (0600), written via temp-file
+ * rename to survive crashes, mirroring pi-check-agent-quota's safe-write
+ * approach but simplified.
+ */
+
+import {
+  closeSync,
+  constants as fsConstants,
+  fchmodSync,
+  fstatSync,
+  openSync,
+  readFileSync,
+} from "node:fs";
+import { chmod, lstat, mkdir, open, rename, unlink } from "node:fs/promises";
+import { randomUUID } from "node:crypto";
+import { homedir } from "node:os";
+import { join } from "node:path";
+import { DEFAULT_SETTINGS, getSettings, patchSettings } from "./widget.js";
+import type {
+  ArkSettings,
+  DiskCache,
+  Language,
+  PlanSnapshot,
+} from "./types.js";
+
+export const CACHE_DIR = join(homedir(), ".pi", "agent", "pi-ark-usage");
+export const CACHE_FILE = join(CACHE_DIR, "cache.json");
+
+const FILE_MODE = 0o600;
+const DIR_MODE = 0o700;
+const MAX_BYTES = 256 * 1024;
+
+function isFiniteNumber(v: unknown): v is number {
+  return typeof v === "number" && Number.isFinite(v);
+}
+
+function validateSnapshot(raw: unknown): PlanSnapshot | null {
+  if (!raw || typeof raw !== "object") return null;
+  const s = raw as Record<string, unknown>;
+  if (typeof s.product !== "string" || typeof s.edition !== "string") return null;
+  if (!isFiniteNumber(s.fetchedAt)) return null;
+  if (!Array.isArray(s.periods) || s.periods.length === 0) return null;
+  const periods = s.periods
+    .filter((p): p is Record<string, unknown> => !!p && typeof p === "object")
+    .filter((p) => typeof p.label === "string" && isFiniteNumber(p.percent))
+    .map((p) => ({
+      label: p.label as string,
+      percent: p.percent as number,
+      reset_at: typeof p.reset_at === "string" ? p.reset_at : undefined,
+    }));
+  if (periods.length === 0) return null;
+  return {
+    product: s.product,
+    edition: s.edition,
+    fetchedAt: s.fetchedAt,
+    periods,
+    updatedAt: isFiniteNumber(s.updatedAt) ? s.updatedAt : undefined,
+  };
+}
+
+function validateSettings(raw: unknown): Partial<ArkSettings> {
+  if (!raw || typeof raw !== "object") return {};
+  const r = raw as Record<string, unknown>;
+  const out: Partial<ArkSettings> = {};
+  if (isFiniteNumber(r.pctYellow) && r.pctYellow > 0 && r.pctYellow < 100) {
+    out.pctYellow = r.pctYellow;
+  }
+  if (isFiniteNumber(r.pctRed) && r.pctRed > 0 && r.pctRed < 100) {
+    out.pctRed = r.pctRed;
+  }
+  if (isFiniteNumber(r.autoRefreshMinutes) && r.autoRefreshMinutes >= 0) {
+    out.autoRefreshMinutes = Math.min(30, Math.round(r.autoRefreshMinutes));
+  }
+  if (
+    r.product === "auto" ||
+    r.product === "coding-plan" ||
+    r.product === "agent-plan" ||
+    r.product === "coding-plan-team" ||
+    r.product === "agent-plan-team"
+  ) {
+    out.product = r.product;
+  }
+  if (typeof r.seat === "string" && r.seat.length <= 128) {
+    out.seat = r.seat;
+  }
+  return out;
+}
+
+export type LoadedCache = {
+  language: Language;
+  snapshots: Record<string, PlanSnapshot>;
+};
+
+/** Read & validate cache; applies settings on success. Never throws. */
+export function loadCache(): LoadedCache | null {
+  let fd: number | undefined;
+  try {
+    fd = openSync(
+      CACHE_FILE,
+      fsConstants.O_RDONLY | (fsConstants.O_NOFOLLOW ?? 0),
+    );
+    const info = fstatSync(fd);
+    if (!info.isFile() || info.nlink !== 1 || info.size > MAX_BYTES) return null;
+    fchmodSync(fd, FILE_MODE);
+    const disk = JSON.parse(readFileSync(fd, "utf8")) as DiskCache;
+    if (!disk || disk.version !== 1) return null;
+
+    const validated = validateSettings(disk.settings);
+    patchSettings({ ...DEFAULT_SETTINGS, ...validated });
+
+    const snapshots: Record<string, PlanSnapshot> = {};
+    if (disk.snapshots && typeof disk.snapshots === "object") {
+      for (const [k, v] of Object.entries(disk.snapshots)) {
+        const s = validateSnapshot(v);
+        if (s) snapshots[k] = s;
+      }
+    }
+    return {
+      language: disk.language === "en" ? "en" : "zh",
+      snapshots,
+    };
+  } catch {
+    return null;
+  } finally {
+    if (fd !== undefined) {
+      try {
+        closeSync(fd);
+      } catch {}
+    }
+  }
+}
+
+let writeQueue: Promise<void> = Promise.resolve();
+let pending: string | null = null;
+let scheduled = false;
+let shuttingDown = false;
+
+/** Schedule an async cache write; multiple rapid calls coalesce into one. */
+export function saveCache(
+  language: Language,
+  snapshots: Record<string, PlanSnapshot>,
+): void {
+  if (shuttingDown) return;
+  pending = JSON.stringify({
+    version: 1,
+    language,
+    settings: getSettings(),
+    snapshots,
+  } satisfies DiskCache);
+  if (scheduled) return;
+  scheduled = true;
+  writeQueue = writeQueue
+    .then(async () => {
+      scheduled = false;
+      while (pending !== null) {
+        const data = pending;
+        pending = null;
+        await writeFile(data);
+      }
+    })
+    .catch(() => {});
+}
+
+export async function flushWrites(): Promise<void> {
+  await writeQueue;
+}
+
+export function markShuttingDown(): void {
+  shuttingDown = true;
+}
+
+async function writeFile(data: string): Promise<void> {
+  await mkdir(CACHE_DIR, { recursive: true, mode: DIR_MODE });
+  const dir = await lstat(CACHE_DIR);
+  if (!dir.isDirectory() || dir.isSymbolicLink()) {
+    throw new Error("unsafe cache directory");
+  }
+  await chmod(CACHE_DIR, DIR_MODE);
+
+  const tmp = `${CACHE_FILE}.${randomUUID()}.tmp`;
+  try {
+    const handle = await open(tmp, "wx", FILE_MODE);
+    try {
+      await handle.writeFile(data, { encoding: "utf8" });
+    } finally {
+      await handle.close();
+    }
+    await rename(tmp, CACHE_FILE);
+    await chmod(CACHE_FILE, FILE_MODE);
+  } finally {
+    await unlink(tmp).catch(() => {});
+  }
+}

package/extensions/lib/types.ts

@@ -0,0 +1,73 @@
+/**
+ * Shared types for pi-ark-usage.
+ */
+
+export type Language = "zh" | "en";
+
+/** Render model, modeled after pi-check-agent-quota. */
+export type RenderItem =
+  | { kind: "text"; text: string }
+  | { kind: "pct"; pct: number; metric: string }
+  | { kind: "annotation"; text: string }
+  | { kind: "age"; text: string };
+
+/** User-tunable settings (persisted in the disk cache). */
+export type ArkSettings = {
+  /** yellow threshold, percent */
+  pctYellow: number;
+  /** red threshold, percent */
+  pctRed: number;
+  /** auto refresh interval in minutes; 0 = off */
+  autoRefreshMinutes: number;
+  /** forced product; "auto" = map from active provider */
+  product: "auto" | "coding-plan" | "agent-plan" | "coding-plan-team" | "agent-plan-team";
+  /** seat id for team products */
+  seat: string;
+};
+
+export type PeriodLabel = "session" | "5h" | "weekly" | "monthly";
+
+/** One billing period as returned by `arkcli usage plan`. */
+export type ArkPeriod = {
+  label: PeriodLabel | string;
+  percent: number;
+  reset_at?: string;
+};
+
+/** One product entry as returned by `arkcli usage plan`. */
+export type ArkPlanItem = {
+  product: string;
+  edition: "personal" | "team" | string;
+  subscribed: boolean;
+  periods?: ArkPeriod[];
+  error?: string;
+  updated_at?: number;
+};
+
+/** Shape of `arkcli usage plan --format json` output. */
+export type ArkPlanOutput = {
+  viewer?: {
+    auth_method?: string;
+    user_name?: string;
+    region?: string;
+    profile?: string;
+    [k: string]: unknown;
+  };
+  items: ArkPlanItem[];
+};
+
+/** Normalized per-product snapshot kept at runtime. */
+export type PlanSnapshot = {
+  product: string;
+  edition: string;
+  fetchedAt: number;
+  periods: ArkPeriod[];
+  updatedAt?: number;
+};
+
+export type DiskCache = {
+  version: 1;
+  language: Language;
+  settings: Partial<ArkSettings>;
+  snapshots?: Record<string, PlanSnapshot>;
+};