@herbertgao/pi-extensions 2026.10.0 → 2026.10.2

package/node_modules/pi-antigravity/src/features/google-search.ts

@@ -0,0 +1,71 @@
+/**
+ * Optional model-facing search tool.
+ * Disable this extension resource in `pi config` to keep the Antigravity provider
+ * and `/antigravity.search` without registering `google_search`.
+ */
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import { Type } from "@earendil-works/pi-ai";
+import { executeAntigravitySearch } from "../search/index.js";
+import { isExtraToolEnabled, redactSecrets } from "../utils/index.js";
+
+export default function (pi: ExtensionAPI): void {
+  if (!isExtraToolEnabled("SEARCH")) return;
+  pi.registerTool({
+    name: "google_search",
+    label: "Google Search",
+    description:
+      "Search the web using Google Search via Antigravity Grounding. Returns a compact evidence brief with source citations. Use this whenever current, real-time information or web research is needed.",
+    promptSnippet: "Real-time Google Search grounding via Antigravity",
+    promptGuidelines: [
+      "Use google_search when you need real-time, up-to-date web information, latest news, documentation, or fact checking.",
+      "Pass specific URLs into the urls array if you want the search engine to fetch and analyze specific pages.",
+      "Set instruction to provide special research directives from the lead agent (e.g. focus on issues, compare benchmarks, restrict time range).",
+    ],
+    parameters: Type.Object({
+      query: Type.String({ description: "Search query or question." }),
+      instruction: Type.Optional(
+        Type.String({
+          description:
+            "Specific directive or focus from the lead agent (e.g., '重点关注英文社区近七天的讨论', 'focus on GitHub issues and benchmarks', 'ignore marketing announcements').",
+        }),
+      ),
+      urls: Type.Optional(
+        Type.Array(Type.String(), {
+          description: "Optional specific URLs to fetch and analyze alongside the search.",
+        }),
+      ),
+      thinking: Type.Optional(
+        Type.Boolean({
+          description: "Enable deeper thinking for complex analysis (default: false).",
+        }),
+      ),
+    }),
+    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: `Searching Google for: "${params.query}"…` }],
+        details: {},
+      });
+      try {
+        const result = await executeAntigravitySearch({
+          apiKey,
+          query: params.query,
+          instruction: params.instruction,
+          urls: params.urls,
+          thinking: params.thinking,
+          signal,
+        });
+        return {
+          content: [{ type: "text", text: result }],
+          details: {},
+        };
+      } catch (error) {
+        const msg = error instanceof Error ? error.message : String(error);
+        throw new Error(redactSecrets(msg), { cause: error });
+      }
+    },
+  });
+}

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

@@ -1,3 +1,4 @@
+import { randomUUID } from "node:crypto";
 import { mkdir, writeFile } from "node:fs/promises";
 import { dirname, extname, isAbsolute, join, relative, resolve, sep } from "node:path";
 import {
@@ -9,9 +10,9 @@ import {
 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";
+import { sanitizeText } from "../utils/util.js";
 
-export const DEFAULT_IMAGE_MODEL = "gemini-3-pro-image";
+export const DEFAULT_IMAGE_MODEL = "gemini-3.1-flash-image";
 export const IMAGE_ASPECT_RATIOS = [
   "1:1",
   "2:3",
@@ -28,11 +29,9 @@ export type ImageAspectRatio = (typeof IMAGE_ASPECT_RATIOS)[number];
 
 const IMAGE_MODEL_FALLBACKS = [
   DEFAULT_IMAGE_MODEL,
-  "gemini-3.1-flash-image",
+  "gemini-3-pro-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;
 
@@ -42,14 +41,13 @@ export type ImageGenerateRequest = {
   project: string;
   model: string;
   request: {
-    contents: Array<{ role: "user"; parts: Array<{ text: string }> }>;
-    systemInstruction: { role: "user"; parts: Array<{ text: string }> };
+    contents: Array<{ role: "user"; parts: Array<{ text?: string; inlineData?: GeneratedImage }> }>;
     generationConfig: {
       imageConfig: { aspectRatio: string };
       candidateCount: number;
     };
   };
-  requestType: "agent";
+  requestType: "image_gen";
   userAgent: "antigravity";
   requestId: string;
 };
@@ -177,25 +175,29 @@ export function buildImageGenerateRequest(
   model: string,
   projectId: string,
   aspectRatio: string,
+  inputImages: GeneratedImage[] = [],
 ): 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 }],
-      },
+      contents: [
+        {
+          role: GeminiRole.User,
+          parts: [
+            { text: sanitizeText(prompt) },
+            ...inputImages.map((inlineData) => ({ inlineData })),
+          ],
+        },
+      ],
       generationConfig: {
         imageConfig: { aspectRatio },
         candidateCount: 1,
       },
     },
-    requestType: AntigravityRequestType.Agent,
+    requestType: AntigravityRequestType.ImageGen,
     userAgent: AntigravityUserAgent.Antigravity,
-    requestId: envelope.requestId,
+    requestId: `image_gen/${Date.now()}/${randomUUID()}/1`,
   };
 }
 
@@ -257,17 +259,48 @@ export async function collectImagesFromSse(
   return { images, text };
 }
 
+/** Uses the direct image endpoint; only a 404 enables the legacy SSE endpoint. */
+export async function requestImageResponse(
+  endpoint: string,
+  init: RequestInit,
+  fetcher: (url: string, init: RequestInit) => Promise<Response> = antigravityFetch,
+): Promise<{ response: Response; streaming: boolean }> {
+  const response = await fetcher(`${endpoint}/v1internal:generateContent`, init);
+  if (response.status !== 404) return { response, streaming: false };
+  await response.text();
+  return {
+    response: await fetcher(`${endpoint}/v1internal:streamGenerateContent?alt=sse`, init),
+    streaming: true,
+  };
+}
+
+/** Extracts images and notes from wrapped or unwrapped direct responses. */
+export function parseImageResponse(chunk: ImageStreamChunk): {
+  images: GeneratedImage[];
+  text: string[];
+} {
+  if (chunk.error?.message) throw new Error(chunk.error.message);
+  const images: GeneratedImage[] = [];
+  const text: string[] = [];
+  for (const candidate of (chunk.response || chunk).candidates || []) {
+    collectImagesFromParts(candidate.content?.parts, images, text);
+  }
+  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> {
+/** Generates in-memory content shared by the native API and file-saving wrapper. */
+export async function generateAntigravityImageContent(
+  options: Omit<GenerateImageOptions, "cwd" | "path"> & { inputImages?: GeneratedImage[] },
+): Promise<Omit<GenerateImageResult, "savedPaths">> {
   const prompt = options.prompt.trim();
-  if (!prompt) throw new Error("Image prompt is required.");
+  if (!prompt && !options.inputImages?.length)
+    throw new Error("Image prompt or input image is required.");
   if (prompt.length > MAX_PROMPT_CHARS) {
     throw new Error(`Image prompt is too long (max ${MAX_PROMPT_CHARS} characters).`);
   }
@@ -280,20 +313,17 @@ export async function generateAntigravityImage(
   let lastError = "no endpoint available";
   for (const model of models) {
     const body = JSON.stringify(
-      buildImageGenerateRequest(prompt, model, creds.projectId, aspectRatio),
+      buildImageGenerateRequest(prompt, model, creds.projectId, aspectRatio, options.inputImages),
     );
     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,
-          },
-        );
+        const { response, streaming } = await requestImageResponse(endpoint, {
+          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)) {
@@ -303,27 +333,14 @@ export async function generateAntigravityImage(
             `Antigravity image request failed (${response.status}): ${safeError(lastError)}`,
           );
         }
-        const parsed = await collectImagesFromSse(response, options.signal);
+        const parsed = streaming
+          ? await collectImagesFromSse(response, options.signal)
+          : parseImageResponse((await response.json()) as ImageStreamChunk);
         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 };
+        return { images: parsed.images, text: parsed.text, model };
       } catch (error) {
         lastError = safeError(error);
         if (options.signal?.aborted) {
@@ -334,3 +351,25 @@ export async function generateAntigravityImage(
   }
   throw new Error(`Antigravity image generation failed: ${safeError(lastError)}`);
 }
+
+/** Saves generated images for the existing tool and slash command. */
+export async function generateAntigravityImage(
+  options: GenerateImageOptions,
+): Promise<GenerateImageResult> {
+  const result = await generateAntigravityImageContent(options);
+  const savedPaths: string[] = [];
+  for (const [index, image] of result.images.entries()) {
+    savedPaths.push(
+      await writeImage(
+        resolveImageSavePath(
+          options.cwd,
+          options.path,
+          image.mimeType,
+          result.images.length > 1 ? index : undefined,
+        ),
+        image,
+      ),
+    );
+  }
+  return { ...result, savedPaths };
+}

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

@@ -0,0 +1,69 @@
+import type { AssistantImages, ImagesContext, ImagesOptions } from "@earendil-works/pi-ai";
+import { DEFAULT_IMAGE_MODEL, generateAntigravityImageContent } from "./image.js";
+import { safeError } from "../utils/security.js";
+
+export const ANTIGRAVITY_IMAGE_API = "antigravity-images";
+export const ANTIGRAVITY_IMAGE_MODELS = [
+  {
+    type: "image" as const,
+    id: DEFAULT_IMAGE_MODEL,
+    name: "Gemini 3.1 Flash Image (Antigravity)",
+    api: ANTIGRAVITY_IMAGE_API,
+    input: ["text", "image"] as ("text" | "image")[],
+    output: ["text", "image"] as ("text" | "image")[],
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+  },
+];
+
+type ImageModelIdentity = { id: string; api: string; provider: string };
+
+/** Implements Pi's in-memory image API using the existing Antigravity credentials and transport. */
+export async function generateAntigravityImages(
+  model: ImageModelIdentity,
+  context: ImagesContext,
+  options?: ImagesOptions,
+): Promise<AssistantImages> {
+  const identity = { api: model.api, provider: model.provider, model: model.id };
+  try {
+    if (options?.signal?.aborted) throw new Error("Request was aborted");
+    if (!options?.apiKey)
+      throw new Error("No Antigravity credentials. Run /login antigravity first.");
+    const ratio = options.metadata?.aspectRatio;
+    if (ratio !== undefined && typeof ratio !== "string")
+      throw new Error("aspectRatio must be a string.");
+    const result = await generateAntigravityImageContent({
+      apiKey: options.apiKey,
+      model: model.id,
+      signal: options.signal,
+      aspectRatio: ratio,
+      prompt: context.input
+        .filter((item) => item.type === "text")
+        .map((item) => item.text)
+        .join("\n\n"),
+      inputImages: context.input
+        .filter((item) => item.type === "image")
+        .map((item) => ({
+          data: item.data,
+          mimeType: item.mimeType,
+        })),
+    });
+    return {
+      ...identity,
+      model: result.model,
+      stopReason: "stop",
+      timestamp: Date.now(),
+      output: [
+        ...result.text.map((text) => ({ type: "text" as const, text })),
+        ...result.images.map((image) => ({ type: "image" as const, ...image })),
+      ],
+    };
+  } catch (error) {
+    return {
+      ...identity,
+      output: [],
+      timestamp: Date.now(),
+      stopReason: options?.signal?.aborted ? "aborted" : "error",
+      errorMessage: safeError(error),
+    };
+  }
+}

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

@@ -1,5 +1,10 @@
 import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
-import { StringEnum, Type } from "@earendil-works/pi-ai";
+import {
+  ANTIGRAVITY_IMAGE_API,
+  ANTIGRAVITY_IMAGE_MODELS,
+  generateAntigravityImages,
+} from "./image/native.js";
+
 // Namespace import: Oh My Pi rewrites this specifier onto bundled pi-ai, which
 // does not export registerApiProvider. A static named import fails plugin load.
 import * as piAiCompat from "@earendil-works/pi-ai/compat";
@@ -15,12 +20,7 @@ import {
 } from "./auth/index.js";
 import { DEFAULT_ENDPOINT } 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 { generateAntigravityImage, parseImageCommandArgs } from "./image/index.js";
 import { executeAntigravitySearch, parseSearchCommandArgs } from "./search/index.js";
 import {
   applyAntigravityCatalog,
@@ -38,7 +38,7 @@ import {
   formatUsageSummary,
   resolveApiKeyFromContext,
 } from "./usage/index.js";
-import { isExtraToolEnabled, redactSecrets, maskEmail } from "./utils/index.js";
+import { redactSecrets, maskEmail } from "./utils/index.js";
 
 /**
  * Pi's interactive `notify` writes into the chat transcript. `console.log` in that
@@ -102,6 +102,11 @@ type CompatApiProviderRegistrar = (provider: {
   streamSimple: typeof streamAntigravity;
 }) => void;
 
+type CompatImagesProviderRegistrar = (provider: {
+  api: typeof ANTIGRAVITY_IMAGE_API;
+  generateImages: typeof generateAntigravityImages;
+}) => void;
+
 /**
  * Pi dispatches custom APIs through the compat registry. Oh My Pi does not
  * export `registerApiProvider` and registers the stream inside `registerProvider`.
@@ -109,12 +114,24 @@ type CompatApiProviderRegistrar = (provider: {
 function registerCompatApiProvider(): void {
   const register = (piAiCompat as { registerApiProvider?: CompatApiProviderRegistrar })
     .registerApiProvider;
-  if (typeof register !== "function") return;
-  register({
-    api: ANTIGRAVITY_API,
-    stream: streamAntigravity,
-    streamSimple: streamAntigravity,
-  });
+  if (typeof register === "function") {
+    register({
+      api: ANTIGRAVITY_API,
+      stream: streamAntigravity,
+      streamSimple: streamAntigravity,
+    });
+  }
+  const registerImages = (
+    piAiCompat as unknown as {
+      registerImagesApiProvider?: CompatImagesProviderRegistrar;
+    }
+  ).registerImagesApiProvider;
+  if (typeof registerImages === "function") {
+    registerImages({
+      api: ANTIGRAVITY_IMAGE_API,
+      generateImages: generateAntigravityImages,
+    });
+  }
 }
 
 export default function (pi: ExtensionAPI): void {
@@ -122,20 +139,37 @@ export default function (pi: ExtensionAPI): void {
 
   const initialCatalog = getCurrentAntigravityCatalog();
 
-  pi.registerProvider(PROVIDER_ID, {
+  const providerConfig = {
     name: PROVIDER_NAME,
     baseUrl: DEFAULT_ENDPOINT,
     api: ANTIGRAVITY_API,
-    models: initialCatalog.models,
-    refreshModels: refreshAntigravityModels,
+    models: [...initialCatalog.models, ...ANTIGRAVITY_IMAGE_MODELS],
+    refreshModels: async (context: Parameters<typeof refreshAntigravityModels>[0]) => {
+      const refreshed = await refreshAntigravityModels(context);
+      return [...refreshed, ...ANTIGRAVITY_IMAGE_MODELS];
+    },
     oauth: {
       name: PROVIDER_NAME,
+      // Antigravity access is quota-backed by the Google account, not metered per token.
+      // Pi uses this for the footer's (sub) marker and subscription-aware cost display.
+      isSubscription: true,
       login: loginAndRemember,
       refreshToken: refreshAndRemember,
       getApiKey,
     },
     streamSimple: streamAntigravity,
-  });
+    images: {
+      [ANTIGRAVITY_IMAGE_API]: {
+        generateImages: generateAntigravityImages,
+      },
+    },
+  };
+
+  // Pi 0.86 types are chat-only; newer Pi accepts mixed-operation models and images.
+  pi.registerProvider(
+    PROVIDER_ID,
+    providerConfig as unknown as Parameters<typeof pi.registerProvider>[1],
+  );
 
   pi.registerCommand("antigravity.usage", {
     description: "Show Antigravity shared quota pools (Gemini / Claude+GPT, 5h + weekly)",
@@ -352,123 +386,4 @@ export default function (pi: ExtensionAPI): void {
       }
     },
   });
-
-  if (isExtraToolEnabled("IMAGE")) {
-    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 },
-        };
-      },
-    });
-  }
-
-  if (isExtraToolEnabled("SEARCH")) {
-    pi.registerTool({
-      name: "google_search",
-      label: "Google Search",
-      description:
-        "Search the web using Google Search via Antigravity Grounding powered by Gemini 3 Flash. Returns real-time web results with multi-angle deep search and source citations. Use this whenever current, real-time information or web research is needed.",
-      promptSnippet: "Real-time Google Search grounding via Antigravity (Gemini 3 Flash)",
-      promptGuidelines: [
-        "Use google_search when you need real-time, up-to-date web information, latest news, documentation, or fact checking.",
-        "Pass specific URLs into the urls array if you want the search engine to fetch and analyze specific pages.",
-        "Set instruction to provide special research directives from the lead agent (e.g. focus on issues, compare benchmarks, restrict time range).",
-      ],
-      parameters: Type.Object({
-        query: Type.String({ description: "Search query or question." }),
-        instruction: Type.Optional(
-          Type.String({
-            description:
-              "Specific directive or focus from the lead agent (e.g., '重点关注英文社区近七天的讨论', 'focus on GitHub issues and benchmarks', 'ignore marketing announcements').",
-          }),
-        ),
-        urls: Type.Optional(
-          Type.Array(Type.String(), {
-            description: "Optional specific URLs to fetch and analyze alongside the search.",
-          }),
-        ),
-        thinking: Type.Optional(
-          Type.Boolean({
-            description: "Enable deeper thinking for complex analysis (default: false).",
-          }),
-        ),
-      }),
-      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: `Searching Google for: "${params.query}"…` }],
-          details: {},
-        });
-        try {
-          const result = await executeAntigravitySearch({
-            apiKey,
-            query: params.query,
-            instruction: params.instruction,
-            urls: params.urls,
-            thinking: params.thinking,
-            signal,
-          });
-          return {
-            content: [{ type: "text", text: result }],
-            details: {},
-          };
-        } catch (error) {
-          const msg = error instanceof Error ? error.message : String(error);
-          throw new Error(redactSecrets(msg), { cause: error });
-        }
-      },
-    });
-  }
 }