@herbertgao/pi-extensions 2026.9.2 → 2026.9.4

package/node_modules/pi-antigravity/src/image/image.ts

@@ -0,0 +1,336 @@
+import { mkdir, writeFile } from "node:fs/promises";
+import { dirname, extname, isAbsolute, join, relative, resolve, sep } from "node:path";
+import {
+  antigravityHeaders,
+  endpointCandidates,
+  jsonOrTextError,
+  parseApiKey,
+} from "../client/client.js";
+import { AntigravityRequestType, AntigravityUserAgent, GeminiRole } from "../types/enums.js";
+import { antigravityFetch } from "../utils/http.js";
+import { safeError } from "../utils/security.js";
+import { antigravityRequestEnvelope, sanitizeText } from "../utils/util.js";
+
+export const DEFAULT_IMAGE_MODEL = "gemini-3-pro-image";
+export const IMAGE_ASPECT_RATIOS = [
+  "1:1",
+  "2:3",
+  "3:2",
+  "3:4",
+  "4:3",
+  "4:5",
+  "5:4",
+  "9:16",
+  "16:9",
+  "21:9",
+] as const;
+export type ImageAspectRatio = (typeof IMAGE_ASPECT_RATIOS)[number];
+
+const IMAGE_MODEL_FALLBACKS = [
+  DEFAULT_IMAGE_MODEL,
+  "gemini-3.1-flash-image",
+  "gemini-3-pro-image-preview",
+];
+const IMAGE_SYSTEM_INSTRUCTION =
+  "You are an AI image generator. Generate images based on user descriptions. Focus on creating high-quality, visually appealing images that match the user's request.";
+const DEFAULT_IMAGE_DIR = join(".pi", "generated-images");
+const MAX_PROMPT_CHARS = 8000;
+
+export type GeneratedImage = { data: string; mimeType: string };
+
+export type ImageGenerateRequest = {
+  project: string;
+  model: string;
+  request: {
+    contents: Array<{ role: GeminiRole.User; parts: Array<{ text: string }> }>;
+    systemInstruction: { role: GeminiRole.User; parts: Array<{ text: string }> };
+    generationConfig: {
+      imageConfig: { aspectRatio: string };
+      candidateCount: number;
+    };
+  };
+  requestType: AntigravityRequestType.Agent;
+  userAgent: AntigravityUserAgent.Antigravity;
+  requestId: string;
+};
+
+export type ImageCommandArgs = {
+  prompt: string;
+  aspectRatio?: string;
+  model?: string;
+  path?: string;
+};
+
+export type GenerateImageOptions = ImageCommandArgs & {
+  apiKey: string;
+  cwd: string;
+  signal?: AbortSignal;
+};
+
+export type GenerateImageResult = {
+  images: GeneratedImage[];
+  savedPaths: string[];
+  text: string[];
+  model: string;
+};
+
+type ImageStreamChunk = {
+  error?: { message?: string };
+  response?: {
+    candidates?: Array<{
+      content?: {
+        parts?: Array<{ text?: string; inlineData?: { mimeType?: string; data?: string } }>;
+      };
+    }>;
+  };
+  candidates?: Array<{
+    content?: {
+      parts?: Array<{ text?: string; inlineData?: { mimeType?: string; data?: string } }>;
+    };
+  }>;
+};
+
+function imageExtension(mimeType: string): string {
+  const lower = mimeType.toLowerCase();
+  if (lower.includes("jpeg") || lower.includes("jpg")) return "jpg";
+  if (lower.includes("webp")) return "webp";
+  if (lower.includes("gif")) return "gif";
+  return "png";
+}
+
+export function assertSafeImageModel(modelId: string): string {
+  const id = modelId.trim();
+  if (id.length === 0 || id.length > 80) {
+    throw new Error("Unsupported image model id.");
+  }
+  if (!/^(gemini-[a-z0-9.+-]*image[a-z0-9.+-]*|imagen-[a-z0-9.+-]+)$/i.test(id)) {
+    throw new Error(`Unsupported image model: ${id}`);
+  }
+  return id;
+}
+
+export function assertSafeAspectRatio(ratio: string): ImageAspectRatio {
+  const value = ratio.trim();
+  for (const allowed of IMAGE_ASPECT_RATIOS) {
+    if (allowed === value) return allowed;
+  }
+  throw new Error(
+    `Unsupported aspect ratio: ${value}. Use one of ${IMAGE_ASPECT_RATIOS.join(", ")}.`,
+  );
+}
+
+export function parseImageCommandArgs(args: string): ImageCommandArgs {
+  const tokens = args.trim().split(/\s+/).filter(Boolean);
+  const out: ImageCommandArgs = { prompt: "" };
+  const rest: string[] = [];
+  for (let i = 0; i < tokens.length; i++) {
+    const token = tokens[i];
+    if (token === undefined) continue;
+    const next = tokens[i + 1];
+    if ((token === "--ratio" || token === "--aspect-ratio") && next) {
+      out.aspectRatio = next;
+      i += 1;
+      continue;
+    }
+    if (token === "--model" && next) {
+      out.model = next;
+      i += 1;
+      continue;
+    }
+    if (token === "--path" && next) {
+      out.path = next;
+      i += 1;
+      continue;
+    }
+    rest.push(token);
+  }
+  out.prompt = rest.join(" ");
+  return out;
+}
+
+export function resolveImageSavePath(
+  cwd: string,
+  requested?: string,
+  mimeType = "image/png",
+  index?: number,
+): string {
+  const ext = imageExtension(mimeType);
+  const stamp = new Date().toISOString().replace(/[:.]/g, "-");
+  const suffix = index === undefined ? "" : `-${index + 1}`;
+  const defaultName = `image-${stamp}${suffix}.${ext}`;
+  const root = resolve(cwd);
+  const target = requested?.trim()
+    ? resolve(root, requested.trim())
+    : resolve(root, DEFAULT_IMAGE_DIR, defaultName);
+  const rel = relative(root, target);
+  if (!rel || rel === ".." || rel.startsWith(`..${sep}`) || isAbsolute(rel)) {
+    throw new Error("Image save path must be inside the working directory.");
+  }
+  if (!extname(target)) return join(target, defaultName);
+  if (index === undefined) return target;
+  const currentExt = extname(target);
+  return `${target.slice(0, -currentExt.length)}${suffix}${currentExt}`;
+}
+
+export function buildImageGenerateRequest(
+  prompt: string,
+  model: string,
+  projectId: string,
+  aspectRatio: string,
+): ImageGenerateRequest {
+  const envelope = antigravityRequestEnvelope(model, false);
+  return {
+    project: projectId,
+    model,
+    request: {
+      contents: [{ role: GeminiRole.User, parts: [{ text: sanitizeText(prompt) }] }],
+      systemInstruction: {
+        role: GeminiRole.User,
+        parts: [{ text: IMAGE_SYSTEM_INSTRUCTION }],
+      },
+      generationConfig: {
+        imageConfig: { aspectRatio },
+        candidateCount: 1,
+      },
+    },
+    requestType: AntigravityRequestType.Agent,
+    userAgent: AntigravityUserAgent.Antigravity,
+    requestId: envelope.requestId,
+  };
+}
+
+function collectImagesFromParts(
+  parts: Array<{ text?: string; inlineData?: { mimeType?: string; data?: string } }> | undefined,
+  images: GeneratedImage[],
+  text: string[],
+): void {
+  for (const part of parts || []) {
+    if (part.text) text.push(part.text);
+    if (part.inlineData?.data) {
+      images.push({
+        data: part.inlineData.data,
+        mimeType: part.inlineData.mimeType || "image/png",
+      });
+    }
+  }
+}
+
+export async function collectImagesFromSse(
+  response: Response,
+  signal?: AbortSignal,
+): Promise<{ images: GeneratedImage[]; text: string[] }> {
+  if (!response.body) throw new Error("No response body");
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+  const images: GeneratedImage[] = [];
+  const text: string[] = [];
+  try {
+    while (true) {
+      if (signal?.aborted) throw new Error("Request was aborted");
+      const result = await reader.read();
+      if (result.done) break;
+      if (!(result.value instanceof Uint8Array)) continue;
+      buffer += decoder.decode(result.value, { stream: true });
+      const lines = buffer.split("\n");
+      buffer = lines.pop() || "";
+      for (const line of lines) {
+        if (!line.startsWith("data:")) continue;
+        const json = line.slice(5).trim();
+        if (!json || json === "[DONE]") continue;
+        let chunk: ImageStreamChunk;
+        try {
+          chunk = JSON.parse(json) as ImageStreamChunk;
+        } catch {
+          continue;
+        }
+        if (chunk.error?.message) throw new Error(chunk.error.message);
+        const responseData = chunk.response || chunk;
+        for (const candidate of responseData.candidates || []) {
+          collectImagesFromParts(candidate.content?.parts, images, text);
+        }
+      }
+    }
+  } finally {
+    reader.releaseLock();
+  }
+  return { images, text };
+}
+
+async function writeImage(filePath: string, image: GeneratedImage): Promise<string> {
+  await mkdir(dirname(filePath), { recursive: true });
+  await writeFile(filePath, Buffer.from(image.data, "base64"));
+  return filePath;
+}
+
+export async function generateAntigravityImage(
+  options: GenerateImageOptions,
+): Promise<GenerateImageResult> {
+  const prompt = options.prompt.trim();
+  if (!prompt) throw new Error("Image prompt is required.");
+  if (prompt.length > MAX_PROMPT_CHARS) {
+    throw new Error(`Image prompt is too long (max ${MAX_PROMPT_CHARS} characters).`);
+  }
+  const aspectRatio = assertSafeAspectRatio(options.aspectRatio || "1:1");
+  const preferred = assertSafeImageModel(options.model || DEFAULT_IMAGE_MODEL);
+  const models = [preferred, ...IMAGE_MODEL_FALLBACKS.filter((id) => id !== preferred)];
+  const creds = parseApiKey(options.apiKey);
+  const headers = antigravityHeaders(creds.token);
+
+  let lastError = "no endpoint available";
+  for (const model of models) {
+    const body = JSON.stringify(
+      buildImageGenerateRequest(prompt, model, creds.projectId, aspectRatio),
+    );
+    for (const endpoint of endpointCandidates()) {
+      if (options.signal?.aborted) throw new Error("Request was aborted");
+      try {
+        const response = await antigravityFetch(
+          `${endpoint}/v1internal:streamGenerateContent?alt=sse`,
+          {
+            method: "POST",
+            headers,
+            body,
+            signal: options.signal,
+          },
+        );
+        if (!response.ok) {
+          lastError = jsonOrTextError(await response.text()).slice(0, 400);
+          if (response.status === 404 || [403, 429, 500, 502, 503, 504].includes(response.status)) {
+            continue;
+          }
+          throw new Error(
+            `Antigravity image request failed (${response.status}): ${safeError(lastError)}`,
+          );
+        }
+        const parsed = await collectImagesFromSse(response, options.signal);
+        if (!parsed.images.length) {
+          lastError = parsed.text.join(" ").trim() || "No image data returned.";
+          continue;
+        }
+        const savedPaths: string[] = [];
+        const many = parsed.images.length > 1;
+        for (const [index, image] of parsed.images.entries()) {
+          savedPaths.push(
+            await writeImage(
+              resolveImageSavePath(
+                options.cwd,
+                options.path,
+                image.mimeType,
+                many ? index : undefined,
+              ),
+              image,
+            ),
+          );
+        }
+        return { images: parsed.images, savedPaths, text: parsed.text, model };
+      } catch (error) {
+        lastError = safeError(error);
+        if (options.signal?.aborted) {
+          throw new Error("Request was aborted", { cause: error });
+        }
+      }
+    }
+  }
+  throw new Error(`Antigravity image generation failed: ${safeError(lastError)}`);
+}

package/node_modules/pi-antigravity/src/image/index.ts

@@ -0,0 +1 @@
+export * from "./image.js";

package/node_modules/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 },
+      };
+    },
+  });
+}