@centerforagenticai/pi-multi-account 0.1.1

package/packages/pi-antigravity/src/usage/usage.ts

@@ -0,0 +1,416 @@
+import type { ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
+import {
+  antigravityHeaders,
+  endpointCandidates,
+  extractProjectId,
+  fetchAvailableModelsCatalog,
+  parseApiKey,
+  resolveProjectId,
+} from "../client/client.js";
+import {
+  setLastEndpoint,
+  setLastError,
+  setLastProjectId,
+  setLastStatus,
+} from "../diagnostics/diagnostics.js";
+import { isRecord } from "../utils/util.js";
+import { safeError } from "../utils/security.js";
+import { antigravityFetch } from "../utils/http.js";
+import type {
+  AccountUsage,
+  ApiErrorBody,
+  AvailableModelsRaw,
+  LoadCodeAssistRaw,
+  ModelQuotaRow,
+  QuotaBucket,
+  QuotaGroup,
+  QuotaSummaryRaw,
+  TierInfo,
+  TierRaw,
+} from "../types/types.js";
+
+function clampFraction(value: unknown): number | undefined {
+  if (typeof value !== "number" || !Number.isFinite(value)) return undefined;
+  if (value < 0) return 0;
+  if (value > 1) return 1;
+  return value;
+}
+
+function remainingPercent(remaining?: number): number | undefined {
+  if (remaining === undefined) return undefined;
+  return Math.round(remaining * 1000) / 10;
+}
+
+function progressBar(remaining?: number, width = 20): string {
+  if (remaining === undefined) return `[${"?".repeat(width)}]`;
+  const filled = Math.max(0, Math.min(width, Math.round(remaining * width)));
+  return `[${"#".repeat(filled)}${"-".repeat(width - filled)}]`;
+}
+
+function formatReset(resetTime?: string): string {
+  if (!resetTime) return "n/a";
+  const ts = Date.parse(resetTime);
+  if (!Number.isFinite(ts)) return resetTime;
+  const delta = ts - Date.now();
+  if (delta <= 0) return "now";
+  const totalMin = Math.round(delta / 60000);
+  const days = Math.floor(totalMin / (60 * 24));
+  const hours = Math.floor((totalMin % (60 * 24)) / 60);
+  const mins = totalMin % 60;
+  if (days > 0) return `${days}d ${hours}h`;
+  if (hours > 0) return `${hours}h ${mins}m`;
+  return `${mins}m`;
+}
+
+function jsonHeaders(token: string): Record<string, string> {
+  return {
+    ...antigravityHeaders(token),
+    Accept: "application/json",
+  };
+}
+
+async function postJson(
+  path: string,
+  token: string,
+  body: Record<string, unknown>,
+  signal?: AbortSignal,
+): Promise<{ endpoint: string; status: number; data: unknown }> {
+  let lastErrorText = "";
+  for (const endpoint of endpointCandidates()) {
+    try {
+      const res = await antigravityFetch(`${endpoint}${path}`, {
+        method: "POST",
+        headers: jsonHeaders(token),
+        body: JSON.stringify(body),
+        signal,
+      });
+      setLastEndpoint(endpoint);
+      setLastStatus(res.status);
+      const text = await res.text();
+      let data: unknown;
+      try {
+        data = JSON.parse(text) as unknown;
+      } catch {
+        data = { raw: text } satisfies ApiErrorBody;
+      }
+      if (!res.ok) {
+        const errorBody = isRecord(data) ? (data as ApiErrorBody) : undefined;
+        lastErrorText =
+          typeof errorBody?.error?.message === "string" ? errorBody.error.message : text;
+        if (![403, 404, 429, 500, 502, 503, 504].includes(res.status)) {
+          throw new Error(`${path} failed (${String(res.status)}): ${lastErrorText.slice(0, 300)}`);
+        }
+        continue;
+      }
+      return { endpoint, status: res.status, data };
+    } catch (error) {
+      // An abort must stop this endpoint-candidate loop immediately instead of being
+      // treated as "this endpoint failed, try the next one." Check the signal that was
+      // passed into *this* call, not the shape of the caught error, so it works
+      // uniformly regardless of how a given transport raises the abort.
+      if (signal?.aborted) throw error;
+      lastErrorText = safeError(error);
+      setLastError(lastErrorText);
+    }
+  }
+  throw new Error(`${path} failed: ${lastErrorText || "no endpoint available"}`);
+}
+
+function parseQuotaSummary(data: unknown): { groups: QuotaGroup[]; description?: string } {
+  const summary = (isRecord(data) ? data : {}) as QuotaSummaryRaw;
+  const groups: QuotaGroup[] = [];
+  for (const group of summary.groups || []) {
+    const buckets: QuotaBucket[] = [];
+    for (const bucket of group.buckets || []) {
+      const remaining = clampFraction(bucket.remainingFraction);
+      if (remaining === undefined && !bucket.bucketId) continue;
+      buckets.push({
+        bucketId: String(bucket.bucketId || bucket.displayName || "unknown"),
+        displayName: String(bucket.displayName || bucket.bucketId || "Limit"),
+        window: bucket.window ? String(bucket.window) : undefined,
+        resetTime: bucket.resetTime ? String(bucket.resetTime) : undefined,
+        description: bucket.description ? String(bucket.description) : undefined,
+        remainingFraction: remaining ?? 0,
+      });
+    }
+    if (!buckets.length && !group.displayName) continue;
+    groups.push({
+      displayName: String(group.displayName || "Quota group"),
+      description: group.description ? String(group.description) : undefined,
+      buckets,
+    });
+  }
+  return {
+    groups,
+    description: summary.description ? String(summary.description) : undefined,
+  };
+}
+
+function parseModels(data: unknown): {
+  models: ModelQuotaRow[];
+  defaultAgentModelId?: string;
+} {
+  const raw = (isRecord(data) ? data : {}) as AvailableModelsRaw;
+  const modelsObj = raw.models && isRecord(raw.models) ? raw.models : {};
+  const models: ModelQuotaRow[] = [];
+  for (const [modelId, info] of Object.entries(modelsObj)) {
+    if (!info || !isRecord(info)) continue;
+    if (info.isInternal || String(modelId).startsWith("chat_")) continue;
+    const qi = isRecord(info.quotaInfo) ? info.quotaInfo : {};
+    models.push({
+      modelId,
+      displayName:
+        typeof info.displayName === "string"
+          ? info.displayName
+          : typeof info.label === "string"
+            ? info.label
+            : typeof info.modelName === "string"
+              ? info.modelName
+              : undefined,
+      remainingFraction: clampFraction(qi.remainingFraction),
+      resetTime: qi.resetTime ? String(qi.resetTime) : undefined,
+      modelProvider:
+        typeof info.modelProvider === "string"
+          ? info.modelProvider
+          : typeof info.apiProvider === "string"
+            ? info.apiProvider
+            : undefined,
+      supportsThinking: !!info.supportsThinking,
+      supportsImages: !!info.supportsImages,
+      recommended: !!info.recommended,
+    });
+  }
+  models.sort((a, b) => a.modelId.localeCompare(b.modelId));
+  return {
+    models,
+    defaultAgentModelId:
+      raw.defaultAgentModelId || raw.defaultAgentModel
+        ? String(raw.defaultAgentModelId || raw.defaultAgentModel)
+        : undefined,
+  };
+}
+
+function parseTier(value: unknown): TierInfo | undefined {
+  if (!isRecord(value)) return undefined;
+  const tier = value as TierRaw;
+  if (!tier.id && !tier.name) return undefined;
+  return {
+    id: tier.id ? String(tier.id) : undefined,
+    name: tier.name ? String(tier.name) : undefined,
+    description: tier.description ? String(tier.description) : undefined,
+  };
+}
+
+async function loadCodeAssistSafe(token: string, signal?: AbortSignal) {
+  try {
+    return await postJson(
+      "/v1internal:loadCodeAssist",
+      token,
+      {
+        metadata: {
+          ideType: "ANTIGRAVITY",
+          platform: "PLATFORM_UNSPECIFIED",
+          pluginType: "GEMINI",
+        },
+      },
+      signal,
+    );
+  } catch (error) {
+    // An abort must not be folded into "this leg just came back empty": that would let
+    // fetchAccountUsage resolve with a partial, silently-truncated result instead of
+    // rejecting.
+    if (signal?.aborted) throw error;
+    return null;
+  }
+}
+
+/**
+ * The user-quota-summary RPC is gated behind a paid subscription: free-tier
+ * accounts get 403 SUBSCRIPTION_REQUIRED (#3501). It is best-effort diagnostics
+ * only — never let it block the rest of the account data (models, tier, project).
+ */
+async function fetchQuotaSummarySafe(
+  token: string,
+  signal?: AbortSignal,
+): Promise<
+  | { ok: true; result: { endpoint: string; status: number; data: unknown } }
+  | {
+      ok: false;
+      error: string;
+    }
+> {
+  try {
+    return {
+      ok: true,
+      result: await postJson("/v1internal:retrieveUserQuotaSummary", token, {}, signal),
+    };
+  } catch (error) {
+    // An abort must not be folded into "this leg just came back empty": see
+    // loadCodeAssistSafe above for the same reasoning.
+    if (signal?.aborted) throw error;
+    const msg = safeError(error);
+    setLastError(msg);
+    return { ok: false, error: msg };
+  }
+}
+
+export async function fetchAccountUsage(
+  apiKeyRaw?: string,
+  options?: { signal?: AbortSignal },
+): Promise<AccountUsage> {
+  options?.signal?.throwIfAborted();
+
+  const creds = parseApiKey(apiKeyRaw);
+  const initialProjectId =
+    creds.projectId ||
+    resolveProjectId({
+      token: creds.token,
+      credentialProjectId: creds.projectId,
+    });
+
+  // Fetch loadCodeAssist, quota summary, and available models all in parallel
+  // to minimize command execution latency. Drain every settlement (not just the
+  // first rejection): Promise.all's first-reject alone would not guarantee the
+  // other two legs' own transport calls have actually stopped, which matters once
+  // an abort is in flight.
+  const [assistSettled, summarySettled, availableSettled] = await Promise.allSettled([
+    loadCodeAssistSafe(creds.token, options?.signal),
+    fetchQuotaSummarySafe(creds.token, options?.signal),
+    fetchAvailableModelsCatalog(creds.token, initialProjectId, options?.signal),
+  ]);
+
+  // An abort takes priority over a leg that merely came back empty: any genuine
+  // rejection here (an abort, or the catalog leg's own "no endpoint available"
+  // once every candidate failed) must reject the exported promise rather than be
+  // masked by a sibling leg that happened to settle first.
+  if (assistSettled.status === "rejected") throw assistSettled.reason as unknown;
+  if (summarySettled.status === "rejected") throw summarySettled.reason as unknown;
+  if (availableSettled.status === "rejected") throw availableSettled.reason as unknown;
+
+  const assistResult = assistSettled.value;
+  const summaryRes = summarySettled.value;
+  const available = availableSettled.value;
+
+  // Derive project ID from the loadCodeAssist response or stored project ID.
+  const discoveredProject = assistResult ? extractProjectId(assistResult.data) : undefined;
+  const projectId = resolveProjectId({
+    token: creds.token,
+    warmedProject: discoveredProject ?? null,
+    credentialProjectId: creds.projectId,
+  });
+  setLastProjectId(projectId);
+
+  const summary = summaryRes.ok ? summaryRes.result : null;
+  const quotaSummaryError = summaryRes.ok ? undefined : summaryRes.error;
+  const { groups, description } = summary
+    ? parseQuotaSummary(summary.data)
+    : { groups: [], description: undefined };
+  const { models, defaultAgentModelId } = parseModels(available.data);
+
+  const assistData = (isRecord(assistResult?.data) ? assistResult.data : {}) as LoadCodeAssistRaw;
+  const productTier = parseTier(assistData.currentTier);
+  const paidTier = parseTier(assistData.paidTier);
+
+  // Google returns currentTier=free-tier even for Google AI Pro accounts.
+  // The real subscription lives in paidTier (e.g. g1-pro-tier / Google AI Pro).
+  const planLabel = paidTier?.name
+    ? `${paidTier.name}${paidTier.id ? ` (${paidTier.id})` : ""}`
+    : productTier?.name
+      ? `${productTier.name}${productTier.id ? ` (${productTier.id})` : ""}`
+      : undefined;
+
+  return {
+    projectId,
+    endpoint: summary?.endpoint ?? available.endpoint ?? assistResult?.endpoint,
+    productTier,
+    paidTier,
+    planLabel,
+    groups,
+    groupDescription: description,
+    quotaSummaryError,
+    models,
+    defaultAgentModelId,
+    fetchedAt: Date.now(),
+  };
+}
+
+function quotaErrorNote(msg: string): string {
+  if (/SUBSCRIPTION_REQUIRED|#3501|(?:lack|missing).*license/i.test(msg)) {
+    return "Aggregate quota summary needs a paid subscription (free-tier can't use that endpoint). Per-model usage is still available via /antigravity.models.";
+  }
+  return `Aggregate quota summary unavailable: ${msg.slice(0, 160)}`;
+}
+
+export function formatUsageSummary(usage: AccountUsage): string {
+  const lines: string[] = [];
+
+  if (usage.planLabel) lines.push(usage.planLabel);
+
+  if (!usage.groups.length) {
+    if (usage.quotaSummaryError) {
+      lines.push(quotaErrorNote(usage.quotaSummaryError));
+    } else {
+      lines.push("No quota groups returned.");
+    }
+    return lines.join("\n");
+  }
+
+  for (const group of usage.groups) {
+    if (lines.length) lines.push("");
+    lines.push(group.displayName);
+    for (const bucket of group.buckets) {
+      const rem = remainingPercent(bucket.remainingFraction);
+      lines.push(
+        `  ${progressBar(bucket.remainingFraction)} ${bucket.displayName}: ${rem ?? "?"}% left · resets ${formatReset(bucket.resetTime)}`,
+      );
+    }
+  }
+
+  return lines.join("\n").trimEnd();
+}
+
+export function formatModelsList(usage: AccountUsage, opts?: { all?: boolean }): string {
+  const lines: string[] = [];
+  lines.push("Antigravity available models");
+  lines.push(`project=${usage.projectId}`);
+  if (usage.defaultAgentModelId) lines.push(`defaultAgentModel=${usage.defaultAgentModelId}`);
+  lines.push("");
+
+  const rows = opts?.all
+    ? usage.models
+    : usage.models.filter((m) => !/tab_|chat_/i.test(m.modelId));
+
+  if (!rows.length) {
+    lines.push("No models returned.");
+    return lines.join("\n");
+  }
+
+  const maxId = Math.max(...rows.map((m) => m.modelId.length), 8);
+  for (const m of rows) {
+    const rem = remainingPercent(m.remainingFraction);
+    const flags = [
+      m.recommended ? "recommended" : "",
+      m.supportsThinking ? "thinking" : "",
+      m.supportsImages ? "images" : "",
+    ]
+      .filter(Boolean)
+      .join(",");
+    const name = m.displayName && m.displayName !== m.modelId ? `  ${m.displayName}` : "";
+    lines.push(
+      `${m.modelId.padEnd(maxId)}  rem ${rem === undefined ? "  ?" : String(rem).padStart(5)}%  reset ${formatReset(m.resetTime).padEnd(8)}${flags ? `  [${flags}]` : ""}${name}`,
+    );
+  }
+  lines.push("");
+  lines.push("Note: remaining % is pool-shared (not a private per-model budget).");
+  return lines.join("\n");
+}
+
+export async function resolveApiKeyFromContext(
+  ctx: ExtensionCommandContext,
+): Promise<string | undefined> {
+  try {
+    return await ctx.modelRegistry.getApiKeyForProvider("antigravity");
+  } catch {
+    return undefined;
+  }
+}

package/packages/pi-antigravity/src/utils/http.ts

@@ -0,0 +1,91 @@
+import { antigravityEnv } from "./util.js";
+
+/**
+ * Node's built-in fetch keeps an idle socket for only 4 seconds unless the server
+ * advertises a longer `Keep-Alive: timeout=`. The Cloud Code Assist endpoint sends no
+ * such header, and interactive coding turns are almost always more than 4 seconds
+ * apart, so without a dedicated pool every message pays a fresh DNS + TCP + TLS
+ * handshake. A long-lived dispatcher removes that per-turn setup cost.
+ *
+ * The dispatcher is scoped to this provider's requests rather than installed with
+ * `setGlobalDispatcher`, so it never changes HTTP behaviour for the rest of the host
+ * process or for other extensions.
+ *
+ * Node 22+ global `fetch` cannot take an npm `undici.Agent` (handler mismatch).
+ * Those runtimes already pool connections natively, so we skip the custom Agent.
+ */
+const KEEP_ALIVE_TIMEOUT_MS = 60_000;
+const KEEP_ALIVE_MAX_TIMEOUT_MS = 5 * 60_000;
+const CONNECT_TIMEOUT_MS = 10_000;
+const PREWARM_TIMEOUT_MS = 5_000;
+
+/** `dispatcher` is an undici extension to RequestInit that Node's fetch honours. */
+type DispatcherInit = RequestInit & { dispatcher?: unknown };
+
+let dispatcherPromise: Promise<unknown> | undefined;
+
+function hasProxyConfiguration(): boolean {
+  return ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "http_proxy", "https_proxy", "all_proxy"].some(
+    (name) => Boolean(process.env[name]?.trim()),
+  );
+}
+
+async function getDispatcher(): Promise<unknown> {
+  dispatcherPromise ??= (async () => {
+    // Pi configures a proxy-aware global dispatcher. Passing a private Agent here
+    // would bypass it and make Antigravity requests connect directly instead.
+    if (antigravityEnv("NO_KEEPALIVE") === "1" || hasProxyConfiguration()) {
+      return undefined;
+    }
+    const nodeMajor = Number(process.versions.node?.split(".")[0]);
+    if (!Number.isNaN(nodeMajor) && nodeMajor >= 22) {
+      return undefined;
+    }
+    try {
+      const { Agent } = await import("undici");
+      return new Agent({
+        keepAliveTimeout: KEEP_ALIVE_TIMEOUT_MS,
+        keepAliveMaxTimeout: KEEP_ALIVE_MAX_TIMEOUT_MS,
+        connections: 8,
+        connect: { timeout: CONNECT_TIMEOUT_MS },
+        // HTTP/2 is opt-in: the endpoint negotiates it, but moving SSE onto h2 is a
+        // transport change we do not want to force on every user.
+        allowH2: antigravityEnv("HTTP2") === "1",
+      });
+    } catch {
+      // undici unavailable — fall back to Node's default fetch behaviour.
+      return undefined;
+    }
+  })();
+  return dispatcherPromise;
+}
+
+/** fetch() bound to this provider's keep-alive connection pool when available. */
+export async function antigravityFetch(
+  input: string | URL,
+  init: RequestInit = {},
+): Promise<Response> {
+  const dispatcher = await getDispatcher();
+  if (!dispatcher) return fetch(input, init);
+  return fetch(input, { ...init, dispatcher } as DispatcherInit);
+}
+
+/**
+ * Open the TLS connection when the extension loads so the first message of a session
+ * does not pay the handshake either. Best-effort: failures are ignored.
+ */
+export function prewarmConnection(url: string): void {
+  if (antigravityEnv("NO_PREWARM") === "1") return;
+  void (async () => {
+    try {
+      const res = await antigravityFetch(url, {
+        method: "HEAD",
+        signal: AbortSignal.timeout(PREWARM_TIMEOUT_MS),
+      });
+      // Release the socket back to the pool even though HEAD carries no body.
+      await res.arrayBuffer();
+    } catch {
+      // Warm-up only; the real request will establish the connection instead.
+    }
+  })();
+}

package/packages/pi-antigravity/src/utils/index.ts

@@ -0,0 +1,3 @@
+export * from "./util.js";
+export * from "./security.js";
+export * from "./http.js";

package/packages/pi-antigravity/src/utils/security.ts

@@ -0,0 +1,73 @@
+import { antigravityEnv } from "./util.js";
+
+const LOOPBACK_HOSTS = new Set(["127.0.0.1", "::1", "localhost"]);
+const ALLOWED_API_HOST_SUFFIXES = [".googleapis.com", ".sandbox.googleapis.com"];
+
+/** Only loopback binds are allowed so OAuth codes cannot be stolen off-machine. */
+export function resolveCallbackHost(raw = antigravityEnv("CALLBACK_HOST")): string {
+  const host = (raw || "127.0.0.1").trim().toLowerCase();
+  if (!LOOPBACK_HOSTS.has(host)) {
+    throw new Error(
+      `Unsafe ANTIGRAVITY_CALLBACK_HOST="${host}". Only loopback hosts are allowed: 127.0.0.1, ::1, localhost.`,
+    );
+  }
+  return host === "localhost" ? "127.0.0.1" : host;
+}
+
+/** Prevent token exfiltration via poisoned BASE_URL (SSRF / credential leak). */
+export function assertSafeApiBaseUrl(raw: string): string {
+  let url: URL;
+  try {
+    url = new URL(raw);
+  } catch {
+    throw new Error(`Invalid ANTIGRAVITY_BASE_URL: ${raw}`);
+  }
+  if (url.protocol !== "https:") {
+    throw new Error(`ANTIGRAVITY_BASE_URL must use https (got ${url.protocol})`);
+  }
+  if (url.username || url.password) {
+    throw new Error("ANTIGRAVITY_BASE_URL must not include credentials");
+  }
+  const host = url.hostname.toLowerCase();
+  const allowed =
+    host === "googleapis.com" || ALLOWED_API_HOST_SUFFIXES.some((suffix) => host.endsWith(suffix));
+  if (!allowed) {
+    throw new Error(
+      `ANTIGRAVITY_BASE_URL host "${host}" is not allowed. Use a *.googleapis.com endpoint.`,
+    );
+  }
+  const path = url.pathname.replace(/\/+$/, "");
+  return `${url.origin}${path === "/" ? "" : path}`;
+}
+
+/** Redact bearer tokens, refresh tokens, and similar secrets from diagnostics/errors. */
+export function redactSecrets(text: string): string {
+  return text
+    .replace(/\bya29\.[A-Za-z0-9._~+/-]+=*/g, "[redacted-access-token]")
+    .replace(/\b1\/[A-Za-z0-9_-]{20,}/g, "[redacted-refresh-token]")
+    .replace(/\bBearer\s+[A-Za-z0-9._~+/-]+=*/gi, "Bearer [redacted]")
+    .replace(
+      /("?(?:access_token|refresh_token|id_token|token|client_secret|code_verifier|authorization)"?\s*[:=]\s*")[^"]*(")/gi,
+      "$1[redacted]$2",
+    )
+    .replace(
+      /("?(?:access_token|refresh_token|id_token|token|client_secret|code_verifier|authorization)"?\s*[:=]\s*)[^\s&,}]+/gi,
+      "$1[redacted]",
+    );
+}
+
+export function maskEmail(email: string | undefined): string | undefined {
+  if (!email || typeof email !== "string") return undefined;
+  const parts = email.split("@");
+  if (parts.length !== 2 || !parts[0] || !parts[1]) return "[redacted-email]";
+  const name = parts[0];
+  const domain = parts[1];
+  const lastChar = name.at(-1) || "";
+  const maskedName = name.length > 2 ? `${name[0]}***${lastChar}` : `${name[0]}***`;
+  return `${maskedName}@${domain}`;
+}
+
+export function safeError(error: unknown): string {
+  const raw = error instanceof Error ? error.message : String(error);
+  return redactSecrets(raw);
+}