@centerforagenticai/pi-multi-account 0.1.1

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

@@ -0,0 +1,280 @@
+import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
+import { StringEnum, Type } from "@earendil-works/pi-ai";
+import { registerApiProvider } from "@earendil-works/pi-ai/compat";
+import { getApiKey, loginAntigravity, refreshAntigravityToken } from "./auth/index.js";
+import { DEFAULT_ENDPOINT, endpointCandidates } from "./client/index.js";
+import { getLastDiagnostics, runWithDiagnostics } from "./diagnostics/index.js";
+import {
+  DEFAULT_IMAGE_MODEL,
+  generateAntigravityImage,
+  IMAGE_ASPECT_RATIOS,
+  parseImageCommandArgs,
+} from "./image/index.js";
+import {
+  applyAntigravityCatalog,
+  discoverAntigravityModels,
+  getCurrentAntigravityCatalog,
+  PROVIDER_ID,
+  PROVIDER_NAME,
+  refreshAntigravityModels,
+  resolvedCatalog,
+} from "./models/index.js";
+import { ANTIGRAVITY_API, streamAntigravity } from "./stream/index.js";
+import {
+  fetchAccountUsage,
+  formatModelsList,
+  formatUsageSummary,
+  resolveApiKeyFromContext,
+} from "./usage/index.js";
+import { prewarmConnection, redactSecrets } from "./utils/index.js";
+
+/**
+ * Pi's interactive `notify` writes into the chat transcript. `console.log` in that
+ * mode prints to the raw terminal and paints over the TUI. Use one channel only.
+ */
+function emitCommandOutput(
+  ctx: ExtensionCommandContext,
+  text: string,
+  type: "info" | "warning" | "error" = "info",
+): void {
+  if (ctx.hasUI) {
+    ctx.ui.notify(text, type);
+    return;
+  }
+  if (type === "warning" || type === "error") console.error(text);
+  else console.log(text);
+}
+
+async function withUsage(
+  ctx: ExtensionCommandContext,
+  fn: (usage: Awaited<ReturnType<typeof fetchAccountUsage>>) => string,
+): Promise<void> {
+  try {
+    const apiKey = await resolveApiKeyFromContext(ctx);
+    if (!apiKey) {
+      emitCommandOutput(
+        ctx,
+        "No Antigravity credentials. Run /login antigravity first.",
+        "warning",
+      );
+      return;
+    }
+    if (ctx.hasUI) ctx.ui.notify("Fetching Antigravity usage…", "info");
+    const usage = await runWithDiagnostics(() => fetchAccountUsage(apiKey));
+    emitCommandOutput(ctx, fn(usage));
+  } catch (error) {
+    const msg = error instanceof Error ? error.message : String(error);
+    emitCommandOutput(ctx, `Antigravity usage failed: ${msg}`, "warning");
+  }
+}
+
+export default function (pi: ExtensionAPI): void {
+  // Open the TLS connection up front so the first message of a session does not pay
+  // the handshake. Opt out with ANTIGRAVITY_NO_PREWARM=1.
+  const primaryEndpoint = endpointCandidates()[0];
+  if (primaryEndpoint) prewarmConnection(primaryEndpoint);
+
+  registerApiProvider({
+    api: ANTIGRAVITY_API,
+    stream: streamAntigravity,
+    streamSimple: streamAntigravity,
+  });
+
+  const initialCatalog = getCurrentAntigravityCatalog();
+
+  pi.registerProvider(PROVIDER_ID, {
+    name: PROVIDER_NAME,
+    baseUrl: DEFAULT_ENDPOINT,
+    api: ANTIGRAVITY_API,
+    models: initialCatalog.models,
+    refreshModels: refreshAntigravityModels,
+    oauth: {
+      name: PROVIDER_NAME,
+      login: loginAntigravity,
+      refreshToken: refreshAntigravityToken,
+      getApiKey,
+    },
+    streamSimple: streamAntigravity,
+  });
+
+  pi.registerCommand("antigravity.usage", {
+    description: "Show Antigravity shared quota pools (Gemini / Claude+GPT, 5h + weekly)",
+    handler: async (_args, ctx) => {
+      await withUsage(ctx, formatUsageSummary);
+    },
+  });
+
+  pi.registerCommand("antigravity.models", {
+    description: "List Antigravity runtime models + remaining pool fraction",
+    handler: async (args, ctx) => {
+      const all = /\ball\b/i.test(args || "");
+      await withUsage(ctx, (usage) => formatModelsList(usage, { all }));
+    },
+  });
+
+  pi.registerCommand("antigravity.refresh", {
+    description: "Force refresh Antigravity dynamic model catalog",
+    handler: async (_args, ctx) => {
+      const apiKey = await resolveApiKeyFromContext(ctx);
+      if (!apiKey) {
+        emitCommandOutput(
+          ctx,
+          "No Antigravity credentials. Run /login antigravity first.",
+          "warning",
+        );
+        return;
+      }
+      if (ctx.hasUI) ctx.ui.notify("Refreshing Antigravity models…", "info");
+      try {
+        if (typeof ctx.modelRegistry?.refresh === "function") {
+          const result = await ctx.modelRegistry.refresh({
+            force: true,
+            providers: [PROVIDER_ID],
+          });
+          if (result?.errors?.has(PROVIDER_ID)) {
+            throw result.errors.get(PROVIDER_ID)!;
+          }
+        } else {
+          const discovered = await discoverAntigravityModels(apiKey);
+          const next = resolvedCatalog(discovered, getCurrentAntigravityCatalog());
+          if (discovered.models.length > 0) {
+            applyAntigravityCatalog(next);
+          }
+        }
+        const catalog = getCurrentAntigravityCatalog();
+        const count = catalog.models.length;
+        const sample = catalog.models
+          .slice(0, 4)
+          .map((m) => m.name || m.id)
+          .join(", ");
+        emitCommandOutput(
+          ctx,
+          `Antigravity models refreshed (${count} available: ${sample}${count > 4 ? ", …" : ""})`,
+          "info",
+        );
+      } catch (error) {
+        const msg = error instanceof Error ? error.message : String(error);
+        emitCommandOutput(ctx, `Antigravity model refresh failed: ${redactSecrets(msg)}`, "error");
+      }
+    },
+  });
+
+  pi.registerCommand("antigravity.doctor", {
+    description: "Show sanitized Antigravity provider diagnostics",
+    handler: async (_args, ctx) => {
+      const d = getLastDiagnostics();
+      const lines = [
+        `provider=${PROVIDER_ID}`,
+        `lastResolvedRuntimeModel=${d.resolvedRuntimeModel || "none"}`,
+        `availableModels=${d.availableModels || "none"}`,
+        `matchedModel=${d.matchedModelDebug || "none"}`,
+        `lastEndpoint=${d.endpoint || "none"}`,
+        `lastStatus=${d.status ?? "none"}`,
+        `lastProjectId=${d.projectId || "none"}`,
+        ...(d.latencyMs !== undefined ? [`lastLatencyMs=${d.latencyMs}`] : []),
+        `toolSchemaWarnings=${d.toolSchemaWarnings || "none"}`,
+        `lastError=${d.error ? redactSecrets(d.error) : "none"}`,
+        "transport=native-streamSimple",
+        "runtimeCli=not-used",
+        "commands=/antigravity.usage /antigravity.models /antigravity.refresh /antigravity.doctor /antigravity.image",
+      ];
+      emitCommandOutput(ctx, `Antigravity doctor\n${lines.join("\n")}`);
+    },
+  });
+
+  pi.registerCommand("antigravity.image", {
+    description:
+      "Generate an image via Antigravity (usage: /antigravity.image [--ratio 16:9] <prompt>)",
+    handler: async (args, ctx) => {
+      const parsed = parseImageCommandArgs(args || "");
+      if (!parsed.prompt) {
+        emitCommandOutput(
+          ctx,
+          "Usage: /antigravity.image [--ratio 16:9] [--model gemini-3-pro-image] [--path file.png] <prompt>",
+          "warning",
+        );
+        return;
+      }
+      try {
+        const apiKey = await resolveApiKeyFromContext(ctx);
+        if (!apiKey) {
+          emitCommandOutput(
+            ctx,
+            "No Antigravity credentials. Run /login antigravity first.",
+            "warning",
+          );
+          return;
+        }
+        if (ctx.hasUI) ctx.ui.notify("Generating Antigravity image…", "info");
+        const result = await generateAntigravityImage({
+          apiKey,
+          cwd: ctx.cwd,
+          prompt: parsed.prompt,
+          aspectRatio: parsed.aspectRatio,
+          model: parsed.model,
+          path: parsed.path,
+        });
+        emitCommandOutput(ctx, `Saved image to ${result.savedPaths.join(", ")}`);
+      } catch (error) {
+        const msg = error instanceof Error ? error.message : String(error);
+        emitCommandOutput(ctx, `Antigravity image failed: ${redactSecrets(msg)}`, "warning");
+      }
+    },
+  });
+
+  pi.registerTool({
+    name: "generate_image",
+    label: "Generate image",
+    description:
+      "Generate an image via Antigravity using the signed-in Google account. Saves under .pi/generated-images/ unless path is set.",
+    promptSnippet: "Generate images via Antigravity OAuth (Gemini image models)",
+    promptGuidelines: [
+      "Use generate_image when the user asks to create, draw, or generate an image.",
+    ],
+    parameters: Type.Object({
+      prompt: Type.String({ description: "Image description." }),
+      aspectRatio: Type.Optional(StringEnum(IMAGE_ASPECT_RATIOS)),
+      model: Type.Optional(
+        Type.String({
+          description: `Image model id. Default: ${DEFAULT_IMAGE_MODEL}.`,
+        }),
+      ),
+      path: Type.Optional(
+        Type.String({
+          description: "Project-relative file or directory to save the image.",
+        }),
+      ),
+    }),
+    async execute(_toolCallId, params, signal, onUpdate, ctx) {
+      const apiKey = await ctx.modelRegistry.getApiKeyForProvider("antigravity");
+      if (!apiKey) {
+        throw new Error("No Antigravity credentials. Run /login antigravity first.");
+      }
+      onUpdate?.({ content: [{ type: "text", text: "Generating image…" }], details: {} });
+      const result = await generateAntigravityImage({
+        apiKey,
+        cwd: ctx.cwd,
+        prompt: params.prompt,
+        aspectRatio: params.aspectRatio,
+        model: params.model,
+        path: params.path,
+        signal,
+      });
+      const notes = result.text.join(" ").trim();
+      return {
+        content: [
+          {
+            type: "text" as const,
+            text: `Saved image to ${result.savedPaths.join(", ")}${notes ? `. ${notes}` : ""}`,
+          },
+          ...result.images.map((image) => ({
+            type: "image" as const,
+            data: image.data,
+            mimeType: image.mimeType,
+          })),
+        ],
+        details: { model: result.model, savedPaths: result.savedPaths },
+      };
+    },
+  });
+}

package/packages/pi-antigravity/src/models/discovery.ts

@@ -0,0 +1,154 @@
+import type {
+  Api,
+  Credential,
+  Model,
+  ModelsStoreEntry,
+  RefreshModelsContext,
+} from "@earendil-works/pi-ai";
+import type { ProviderModelConfig } from "@earendil-works/pi-coding-agent";
+import { getApiKey } from "../auth/index.js";
+import { DEFAULT_ENDPOINT, fetchAvailableModelsCatalog, parseApiKey } from "../client/index.js";
+import { ANTIGRAVITY_API } from "../types/types.js";
+import { antigravityEnv, isRecord } from "../utils/util.js";
+import { buildAntigravityCatalog, resolvedCatalog, type AntigravityCatalog } from "./grouping.js";
+import {
+  ANTIGRAVITY_MODELS,
+  ANTIGRAVITY_ROUTING,
+  applyAntigravityCatalog,
+  getCurrentAntigravityCatalog,
+  registerDiscoveredModelEnums,
+  restoreDynamicModelEnums,
+  snapshotDynamicModelEnums,
+} from "./models.js";
+
+export const DEFAULT_CATALOG_REFRESH_INTERVAL_MS = 4 * 60 * 60 * 1000;
+export const ANTIGRAVITY_PERSIST_KEY = "pi-antigravity";
+
+type PersistedAntigravityCatalog = {
+  catalog: AntigravityCatalog;
+  checkedAt: number;
+  modelEnums: Record<string, string>;
+};
+
+export function getCatalogRefreshIntervalMs(): number {
+  const envVal =
+    antigravityEnv("CATALOG_REFRESH_INTERVAL_MS") ?? antigravityEnv("REFRESH_INTERVAL_MS");
+  if (envVal) {
+    const parsed = Number.parseInt(envVal, 10);
+    if (!Number.isNaN(parsed) && parsed >= 0) return parsed;
+  }
+  return DEFAULT_CATALOG_REFRESH_INTERVAL_MS;
+}
+
+const fallbackCatalog = (): AntigravityCatalog => ({
+  models: ANTIGRAVITY_MODELS,
+  routing: { ...ANTIGRAVITY_ROUTING },
+});
+
+/** Restore provider state supplied by Pi before checking offline mode or refresh TTL. */
+export function hydrateAntigravityCatalog(stored: unknown): number {
+  if (!isRecord(stored)) return 0;
+  const persisted = stored[ANTIGRAVITY_PERSIST_KEY];
+  if (!isRecord(persisted)) return 0;
+
+  if (isStringMap(persisted.modelEnums)) restoreDynamicModelEnums(persisted.modelEnums);
+  if (isCatalog(persisted.catalog)) applyAntigravityCatalog(persisted.catalog);
+  return typeof persisted.checkedAt === "number" && persisted.checkedAt > 0
+    ? persisted.checkedAt
+    : 0;
+}
+
+export async function discoverAntigravityModels(
+  apiKey: string,
+  signal?: AbortSignal,
+): Promise<AntigravityCatalog> {
+  const creds = parseApiKey(apiKey);
+  const available = await fetchAvailableModelsCatalog(creds.token, creds.projectId, signal);
+  const models = available.data.models;
+  if (!models || Object.keys(models).length === 0) {
+    return { models: [], routing: {} };
+  }
+  registerDiscoveredModelEnums(models);
+  return buildAntigravityCatalog(models, fallbackCatalog());
+}
+
+export async function refreshAntigravityModels(
+  context: RefreshModelsContext,
+): Promise<ProviderModelConfig[]> {
+  const checkedAt = hydrateAntigravityCatalog(context.stored);
+  const current = getCurrentAntigravityCatalog();
+  if (!context.allowNetwork) return current.models;
+
+  const apiKey = apiKeyFromCredential(context.credential);
+  if (!apiKey || context.signal.aborted) return current.models;
+
+  const now = Date.now();
+  if (
+    !context.force &&
+    checkedAt > 0 &&
+    now >= checkedAt &&
+    now - checkedAt < getCatalogRefreshIntervalMs()
+  ) {
+    return current.models;
+  }
+
+  try {
+    const discovered = await discoverAntigravityModels(apiKey, context.signal);
+    if (context.signal.aborted) return current.models;
+    const next = resolvedCatalog(discovered, current);
+    if (next.models.length > 0 && discovered.models.length > 0) {
+      applyAntigravityCatalog(next);
+      const refreshedAt = Date.now();
+      await context.publish({
+        persist: {
+          models: toStoredModels(next.models),
+          [ANTIGRAVITY_PERSIST_KEY]: {
+            catalog: next,
+            checkedAt: refreshedAt,
+            modelEnums: snapshotDynamicModelEnums(),
+          } satisfies PersistedAntigravityCatalog,
+        } as unknown as ModelsStoreEntry,
+      });
+      return next.models;
+    }
+  } catch (error) {
+    // Keep last-known-good models; a failed refresh must not wipe the catalog.
+    // Forced/manual refreshes must still report the failure to their caller.
+    if (context.force) throw error;
+  }
+
+  return getCurrentAntigravityCatalog().models;
+}
+
+function isCatalog(value: unknown): value is AntigravityCatalog {
+  return (
+    isRecord(value) &&
+    Array.isArray(value.models) &&
+    value.models.length > 0 &&
+    isRecord(value.routing)
+  );
+}
+
+function isStringMap(value: unknown): value is Record<string, string> {
+  return isRecord(value) && Object.values(value).every((entry) => typeof entry === "string");
+}
+
+function apiKeyFromCredential(credential: Credential | undefined): string | undefined {
+  if (!credential) return undefined;
+  if (credential.type === "api_key") {
+    return typeof credential.key === "string" && credential.key ? credential.key : undefined;
+  }
+  if (credential.type === "oauth" && typeof credential.access === "string") {
+    return getApiKey(credential);
+  }
+  return undefined;
+}
+
+function toStoredModels(models: ProviderModelConfig[]): Model<Api>[] {
+  return models.map((model) => ({
+    ...model,
+    api: ANTIGRAVITY_API,
+    provider: "antigravity",
+    baseUrl: DEFAULT_ENDPOINT,
+  }));
+}