pi-provider-antigravity 0.8.1 → 0.9.0

package/README.md

@@ -9,14 +9,16 @@ Use Google Antigravity models directly in [Pi](https://pi.dev).
 
 ## Features
 
-- **Verified wire compatibility**: headers, request envelopes, and tier/thinking mapping
-  strictly validated against the Antigravity protocol specification — not guessed.
+- **Wire parity**: headers, request envelopes, and thinking/tier mapping
+  matched to the official Antigravity client.
 - **Dynamic model catalog**: exposes exactly the models your account can use,
   synthesized from the live catalog and cached for offline startup.
   Newly released models are usable immediately without waiting for an extension update.
 - **OAuth with auto-refresh**: log in once via `/login antigravity`; tokens refresh silently.
 - **Quota visibility**: 5-hour and weekly remaining quota per model pool, with reset times
   and a live footer slot. See [Quota footer](#quota-footer).
+- **Search grounding**: built-in `search_web` tool and `/antigravity search` command
+  with source citations (no extra API keys required).
 
 ## Requirements
 
@@ -70,6 +72,7 @@ Extension subcommands (`/antigravity ...`):
 | `models` | View available models and remaining fraction (alias: `model`) |
 | `refresh` | Fetch the latest model list from Antigravity |
 | `settings` | Configure provider settings (alias: `setting`) |
+| `search` | Search the web using Google Search Grounding |
 | `login` | Shortcut: fills in `/login antigravity` for you |
 
 Related Pi builtins:

package/dist/index.js

@@ -1,3 +1,6 @@
+// src/index.ts
+import { Type } from "@earendil-works/pi-ai";
+
 // src/auth.ts
 import { createHash, randomBytes } from "node:crypto";
 
@@ -251,8 +254,7 @@ function parseStoredCredentials(rawApiKey) {
   return { token: rawApiKey, projectId: FALLBACK_PROJECT_ID };
 }
 
-// src/model-catalog.ts
-var PRIVATE_SNAPSHOT_KEY = "pi-provider-antigravity";
+// src/models.ts
 var CANONICAL_TIER_SUFFIXES = Object.freeze({
   low: "-low",
   medium: "-medium",
@@ -282,6 +284,27 @@ function extractBaseModelId(runtimeId) {
   if (runtimeId.startsWith("gemini-3.1-pro-")) return "gemini-3.1-pro";
   return runtimeId.replace(TIER_SUFFIX_PATTERN, "");
 }
+var TIER_FALLBACKS = Object.freeze({
+  minimal: Object.freeze(["-low", "-extra-low", ""]),
+  low: Object.freeze(["-extra-low", ""]),
+  medium: Object.freeze(["", "-high"]),
+  // Gemini 3.1 Pro lists no -medium: up to high, never down
+  high: Object.freeze(["-thinking", "-agent", ""]),
+  // Claude's high tier is -thinking
+  xhigh: Object.freeze(["-high", "-thinking", "-agent", ""]),
+  max: Object.freeze(["-high", "-thinking", "-agent", ""])
+});
+var DEFAULT_TIER_ORDER = Object.freeze([
+  CANONICAL_TIER_SUFFIXES.high,
+  ...TIER_ALIASES.high,
+  ""
+]);
+function tierCandidateOrder(effort) {
+  if (!effort) return DEFAULT_TIER_ORDER;
+  const canonical = effort in CANONICAL_TIER_SUFFIXES ? [CANONICAL_TIER_SUFFIXES[effort]] : [`-${effort}`];
+  const fallbacks = TIER_FALLBACKS[effort] ?? DEFAULT_TIER_ORDER;
+  return [.../* @__PURE__ */ new Set([...canonical, ...fallbacks])];
+}
 function classifyModelFamily(modelId) {
   const bare = ((modelId || "").split("/").pop() || "").toLowerCase();
   if (bare.startsWith("gemini-")) return "gemini";
@@ -296,6 +319,9 @@ function isCompatibleFamily(msgModel, targetModelId) {
   if (msgFamily !== "unknown" && msgFamily === classifyModelFamily(targetModelId)) return true;
   return extractBaseModelId(msgModel) === extractBaseModelId(targetModelId);
 }
+
+// src/model-catalog.ts
+var PRIVATE_SNAPSHOT_KEY = "pi-provider-antigravity";
 function buildThinkingMap(models) {
   return Object.fromEntries(
     models.map((m) => [
@@ -326,62 +352,10 @@ function copySnapshot(snapshot) {
   };
 }
 var EMPTY_SNAPSHOT = () => ({ enums: {}, runtimeIds: [], thinking: {}, deprecated: {} });
-function createModelCatalog() {
-  let generation = { snapshot: EMPTY_SNAPSHOT(), version: 0 };
-  const record = (catalog) => {
-    generation = {
-      snapshot: snapshotFromCatalog(catalog),
-      items: catalog,
-      version: generation.version + 1
-    };
-  };
-  const restore = (snapshot) => {
-    if (generation.version > 0) return;
-    generation = { snapshot: copySnapshot(snapshot), version: 0 };
-  };
-  const refresh = async (context) => {
-    const persisted = fromPersistedSnapshot(context.stored?.[PRIVATE_SNAPSHOT_KEY]);
-    if (persisted) restore(persisted);
-    const storedModels = () => [...context.stored?.models ?? []];
-    if (!context.allowNetwork) return storedModels();
-    try {
-      const creds = await resolveCredentials(context);
-      if (!creds) return storedModels();
-      const catalog = await fetchAvailableModelsCatalog(creds, creds.projectId, context.signal);
-      record(catalog);
-      const dynamicModels = buildDynamicPublicModels(catalog);
-      if (context.publish) {
-        await context.publish({
-          persist: {
-            models: dynamicModels,
-            checkedAt: Date.now(),
-            [PRIVATE_SNAPSHOT_KEY]: toPersistedSnapshot(generation.snapshot)
-          }
-        });
-      }
-      return dynamicModels;
-    } catch {
-      return storedModels();
-    }
-  };
-  const refreshGen = async (doRefresh) => {
-    const versionBefore = generation.version;
-    await doRefresh();
-    const { items, version } = generation;
-    if (!items) return { status: "failed" };
-    return { status: version === versionBefore ? "stale" : "fresh", catalog: items };
-  };
-  return {
-    generation: () => ({ ...generation, snapshot: copySnapshot(generation.snapshot) }),
-    record,
-    restore,
-    refresh,
-    refreshGeneration: refreshGen,
-    resolvePlan: (publicModelId, effort) => resolveModelPlan(publicModelId, effort, generation.snapshot),
-    formatList: () => generation.items ? formatModelsList(generation.items) : "No models available."
-  };
-}
 function parseAvailableModels(data) {
+  if (!data || typeof data !== "object") {
+    return { models: [], modelEnums: {} };
+  }
   const models = [];
   const modelEnums = {};
   const modelsObj = data.models || {};
@@ -468,27 +442,6 @@ function fromPersistedSnapshot(raw) {
   if (!enums && !runtimeIds && !thinking && !deprecated) return void 0;
   return { enums: enums ?? {}, runtimeIds: runtimeIds ?? [], thinking: thinking ?? {}, deprecated: deprecated ?? {} };
 }
-var TIER_FALLBACKS = Object.freeze({
-  minimal: Object.freeze(["-low", "-extra-low", ""]),
-  low: Object.freeze(["-extra-low", ""]),
-  medium: Object.freeze(["", "-high"]),
-  // Gemini 3.1 Pro lists no -medium: up to high, never down
-  high: Object.freeze(["-thinking", "-agent", ""]),
-  // Claude's high tier is -thinking
-  xhigh: Object.freeze(["-high", "-thinking", "-agent", ""]),
-  max: Object.freeze(["-high", "-thinking", "-agent", ""])
-});
-var DEFAULT_TIER_ORDER = Object.freeze([
-  CANONICAL_TIER_SUFFIXES.high,
-  ...TIER_ALIASES.high,
-  ""
-]);
-function tierCandidateOrder(effort) {
-  if (!effort) return DEFAULT_TIER_ORDER;
-  const canonical = effort in CANONICAL_TIER_SUFFIXES ? [CANONICAL_TIER_SUFFIXES[effort]] : [`-${effort}`];
-  const fallbacks = TIER_FALLBACKS[effort] ?? DEFAULT_TIER_ORDER;
-  return [.../* @__PURE__ */ new Set([...canonical, ...fallbacks])];
-}
 function resolveRuntimeModelId(modelId, effort, availableRuntimeIds) {
   if (ALL_TIER_SUFFIXES.some((suffix) => modelId.endsWith(suffix))) {
     return modelId;
@@ -507,6 +460,21 @@ function resolveRuntimeModelId(modelId, effort, availableRuntimeIds) {
   }
   return modelId;
 }
+function followRenames(runtimeModelId, snapshot) {
+  const renamed = snapshot.deprecated;
+  if (!renamed) return runtimeModelId;
+  let current = runtimeModelId;
+  const seen = /* @__PURE__ */ new Set([current]);
+  while (renamed[current] && !seen.has(renamed[current])) {
+    current = renamed[current];
+    seen.add(current);
+  }
+  return current;
+}
+function resolveThinkingConfig(runtimeModelId, snapshot) {
+  const budget = snapshot.thinking[runtimeModelId]?.budget;
+  return typeof budget === "number" ? { includeThoughts: true, thinkingBudget: budget } : { includeThoughts: false, thinkingBudget: 0 };
+}
 function resolveModelPlan(publicModelId, effort, snapshot) {
   const runtimeModelId = followRenames(
     resolveRuntimeModelId(publicModelId, effort, snapshot.runtimeIds),
@@ -527,24 +495,22 @@ function resolveModelPlan(publicModelId, effort, snapshot) {
     isClaude: family === "claude"
   };
 }
-function followRenames(runtimeModelId, snapshot) {
-  const renamed = snapshot.deprecated;
-  if (!renamed) return runtimeModelId;
-  let current = runtimeModelId;
-  const seen = /* @__PURE__ */ new Set([current]);
-  while (renamed[current] && !seen.has(renamed[current])) {
-    current = renamed[current];
-    seen.add(current);
-  }
-  return current;
-}
-function resolveThinkingConfig(runtimeModelId, snapshot) {
-  const budget = snapshot.thinking[runtimeModelId]?.budget;
-  return typeof budget === "number" ? { includeThoughts: true, thinkingBudget: budget } : { includeThoughts: false, thinkingBudget: 0 };
-}
 function estimateModelCost(_baseId) {
   return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
 }
+function formatModelDisplayName(baseId, rawDisplayName) {
+  if (rawDisplayName) {
+    const cleaned = rawDisplayName.replace(/\s*\([^)]*\)/g, "").trim();
+    if (cleaned.length > 0) {
+      return cleaned;
+    }
+  }
+  const words = baseId.split("-").map((w) => {
+    if (/^\d+(\.\d+)?$/.test(w)) return w;
+    return w.charAt(0).toUpperCase() + w.slice(1);
+  });
+  return words.join(" ");
+}
 function synthesizeDynamicModel(baseId, items) {
   const repItem = items.find((it) => it.id === `${baseId}-high` || it.id === baseId) || items[0];
   const hasVariant = (suffixes) => items.some((it) => suffixes.some((suffix) => it.id.endsWith(suffix)));
@@ -604,19 +570,6 @@ function buildDynamicPublicModels(catalog) {
   }
   return orderedBaseIds.map((baseId) => synthesizeDynamicModel(baseId, runtimeGroups.get(baseId)));
 }
-function formatModelDisplayName(baseId, rawDisplayName) {
-  if (rawDisplayName) {
-    const cleaned = rawDisplayName.replace(/\s*\([^)]*\)/g, "").trim();
-    if (cleaned.length > 0) {
-      return cleaned;
-    }
-  }
-  const words = baseId.split("-").map((w) => {
-    if (/^\d+(\.\d+)?$/.test(w)) return w;
-    return w.charAt(0).toUpperCase() + w.slice(1);
-  });
-  return words.join(" ");
-}
 function formatTokenCount(n) {
   if (typeof n !== "number") return "N/A";
   if (n % 1048576 === 0) return `${n / 1048576}M`;
@@ -661,6 +614,82 @@ function formatModelsList(catalog) {
   }
   return lines.join("\n");
 }
+function createModelCatalog() {
+  let generation = { snapshot: EMPTY_SNAPSHOT(), version: 0 };
+  const record = (catalogOrRaw) => {
+    let catalog;
+    if (catalogOrRaw && typeof catalogOrRaw === "object" && "models" in catalogOrRaw && Array.isArray(catalogOrRaw.models) && "modelEnums" in catalogOrRaw && typeof catalogOrRaw.modelEnums === "object") {
+      catalog = catalogOrRaw;
+    } else {
+      catalog = parseAvailableModels(catalogOrRaw);
+    }
+    generation = {
+      snapshot: snapshotFromCatalog(catalog),
+      items: catalog,
+      version: generation.version + 1
+    };
+  };
+  const restore = (persistedOrSnapshot) => {
+    if (generation.version > 0) return;
+    if (!persistedOrSnapshot || typeof persistedOrSnapshot !== "object") return;
+    if ("enums" in persistedOrSnapshot && "runtimeIds" in persistedOrSnapshot && Array.isArray(persistedOrSnapshot.runtimeIds)) {
+      generation = { snapshot: copySnapshot(persistedOrSnapshot), version: 0 };
+      return;
+    }
+    const snapshot = fromPersistedSnapshot(persistedOrSnapshot);
+    if (snapshot) {
+      generation = { snapshot: copySnapshot(snapshot), version: 0 };
+    }
+  };
+  const toPersisted = () => {
+    if (generation.version === 0 && generation.snapshot.runtimeIds.length === 0) {
+      return void 0;
+    }
+    return toPersistedSnapshot(generation.snapshot);
+  };
+  const refresh = async (context) => {
+    const persisted = context.stored?.[PRIVATE_SNAPSHOT_KEY];
+    if (persisted) restore(persisted);
+    const storedModels = () => [...context.stored?.models ?? []];
+    if (!context.allowNetwork) return storedModels();
+    try {
+      const creds = await resolveCredentials(context);
+      if (!creds) return storedModels();
+      const catalog = await fetchAvailableModelsCatalog(creds, creds.projectId, context.signal);
+      record(catalog);
+      const dynamicModels = buildDynamicPublicModels(catalog);
+      if (context.publish) {
+        await context.publish({
+          persist: {
+            models: dynamicModels,
+            checkedAt: Date.now(),
+            [PRIVATE_SNAPSHOT_KEY]: toPersisted()
+          }
+        });
+      }
+      return dynamicModels;
+    } catch {
+      return storedModels();
+    }
+  };
+  const refreshGen = async (doRefresh) => {
+    const versionBefore = generation.version;
+    await doRefresh();
+    const { items, version } = generation;
+    if (!items) return { status: "failed" };
+    return { status: version === versionBefore ? "stale" : "fresh", catalog: items };
+  };
+  return {
+    generation: () => ({ ...generation, snapshot: copySnapshot(generation.snapshot) }),
+    record,
+    restore,
+    toPersisted,
+    refresh,
+    refreshGeneration: refreshGen,
+    resolvePlan: (publicModelId, effort) => resolveModelPlan(publicModelId, effort, generation.snapshot),
+    formatList: () => generation.items ? formatModelsList(generation.items) : "No models available."
+  };
+}
 
 // src/quota-status.ts
 import * as fs from "node:fs";
@@ -1876,6 +1905,102 @@ function streamAntigravity(model, context, options, catalog) {
   return stream;
 }
 
+// src/search.ts
+var SEARCH_SYSTEM_INSTRUCTION = "You are a search engine bot. You will be given a query from a user. Your task is to search the web for relevant information that will help the user. You MUST perform a web search. Do not respond or interact with the user, please respond as if they typed the query into a search bar.";
+var SEARCH_MODEL = "gemini-3.1-flash-lite";
+function formatSearchResults(rawText, metadata, query, createdAt, completedAt) {
+  const sources = (metadata?.groundingChunks || []).map((chunk, i) => ({
+    index: i + 1,
+    title: chunk.web?.title || "Web Source",
+    url: chunk.web?.uri || ""
+  })).filter((s) => Boolean(s.url));
+  const queries = metadata?.webSearchQueries || [];
+  let annotatedText = rawText;
+  const supports = (metadata?.groundingSupports || []).filter(
+    (s) => typeof s.segment?.endIndex === "number" && Array.isArray(s.groundingChunkIndices) && s.groundingChunkIndices.length > 0
+  ).slice().sort((a, b) => (b.segment.endIndex ?? 0) - (a.segment.endIndex ?? 0));
+  for (const sup of supports) {
+    const end = sup.segment.endIndex;
+    if (end >= 0 && end <= annotatedText.length) {
+      const citations = (sup.groundingChunkIndices || []).map((idx) => `[${idx + 1}]`).join("");
+      annotatedText = annotatedText.slice(0, end) + citations + annotatedText.slice(end);
+    }
+  }
+  const sourceLines = sources.map((s) => `[${s.index}] [${s.title}](${s.url})`);
+  let formattedOutput = `Created At: ${createdAt}
+Completed At: ${completedAt}
+The search for "${query}" returned the following summary:
+${annotatedText.trim() || "No content returned."}`;
+  if (sourceLines.length > 0) {
+    formattedOutput += `
+
+Sources:
+${sourceLines.join("\n")}`;
+  }
+  return {
+    query,
+    rawText,
+    formattedOutput,
+    sources,
+    queries,
+    createdAt,
+    completedAt
+  };
+}
+async function performWebSearch(creds, options) {
+  const createdAt = (/* @__PURE__ */ new Date()).toISOString();
+  let effectiveQuery = options.query.trim();
+  if (options.domain && options.domain.trim()) {
+    const d = options.domain.trim();
+    if (!effectiveQuery.toLowerCase().includes(`site:${d.toLowerCase()}`)) {
+      effectiveQuery = `${effectiveQuery} site:${d}`;
+    }
+  }
+  const reqBody = {
+    project: creds.projectId || FALLBACK_PROJECT_ID,
+    request: {
+      contents: [
+        {
+          role: "user",
+          parts: [{ text: effectiveQuery }]
+        }
+      ],
+      systemInstruction: {
+        role: "user",
+        parts: [{ text: SEARCH_SYSTEM_INSTRUCTION }]
+      },
+      tools: [
+        {
+          googleSearch: {
+            enhancedContent: {
+              imageSearch: {
+                maxResultCount: 5
+              }
+            }
+          }
+        }
+      ],
+      generationConfig: {
+        candidateCount: 1
+      }
+    },
+    model: SEARCH_MODEL,
+    userAgent: "antigravity",
+    requestType: "web_search"
+  };
+  const resp = await postAntigravityJson({
+    auth: creds.token,
+    path: "v1internal:generateContent",
+    body: reqBody,
+    signal: options.signal
+  });
+  const completedAt = (/* @__PURE__ */ new Date()).toISOString();
+  const candidate = resp.response?.candidates?.[0];
+  const rawText = candidate?.content?.parts?.filter((p) => typeof p.text === "string").map((p) => p.text).join("") || "";
+  const metadata = candidate?.groundingMetadata;
+  return formatSearchResults(rawText, metadata, effectiveQuery, createdAt, completedAt);
+}
+
 // src/commands.ts
 function emitOutput(ctx, text, type = "info") {
   if (ctx.hasUI) {
@@ -1908,6 +2033,11 @@ var SUBCOMMANDS = Object.freeze([
     aliases: Object.freeze(["setting"]),
     run: async ({ ctx, quotaStatus }) => openSettings(ctx, quotaStatus)
   },
+  {
+    name: "search",
+    description: "Search the web using Google Search Grounding",
+    run: async ({ ctx, subArgs }) => runSearchSubcommand(ctx, subArgs)
+  },
   {
     name: "login",
     description: "Run /login antigravity",
@@ -2013,12 +2143,36 @@ function runLoginSubcommand(ctx) {
     emitOutput(ctx, "Please run /login antigravity to authenticate.");
   }
 }
+async function runSearchSubcommand(ctx, query) {
+  const creds = await resolveCredentials(ctx);
+  if (!creds) {
+    emitOutput(ctx, NOT_LOGGED_IN, "warning");
+    return;
+  }
+  const trimmedQuery = (query || "").trim();
+  if (!trimmedQuery) {
+    emitOutput(ctx, "Usage: /antigravity search <query>", "warning");
+    return;
+  }
+  try {
+    if (ctx.hasUI) ctx.ui.notify(`Searching: "${trimmedQuery}"\u2026`, "info");
+    const result = await performWebSearch(creds, {
+      query: trimmedQuery,
+      signal: ctx.signal
+    });
+    emitOutput(ctx, result.formattedOutput);
+  } catch (err) {
+    emitOutput(ctx, `Search failed: ${err.message}`, "error");
+  }
+}
 async function runAntigravitySubcommand(args, ctx, quotaStatus, catalog) {
-  const parts = (args || "").trim().split(/\s+/).filter(Boolean);
+  const raw = (args || "").trim();
+  const parts = raw.split(/\s+/).filter(Boolean);
   const sub = parseAntigravitySubcommand(parts[0] || "");
+  const subArgs = raw.slice(parts[0]?.length || 0).trim();
   const cmd = SUBCOMMANDS.find((c) => c.name === sub);
   if (cmd) {
-    await cmd.run({ ctx, quotaStatus, catalog });
+    await cmd.run({ ctx, quotaStatus, catalog, subArgs });
     return;
   }
   emitOutput(ctx, USAGE_TEXT);
@@ -2051,6 +2205,53 @@ function index_default(pi) {
       await runAntigravitySubcommand(args, ctx, quotaStatus, catalog);
     }
   });
+  if (typeof pi.registerTool === "function") {
+    pi.registerTool({
+      name: "search_web",
+      label: "Search Web",
+      description: "Performs a web search for a given query. Returns a summary of relevant information along with URL citations.",
+      promptSnippet: "Search the web using Google Search Grounding with source citations",
+      promptGuidelines: [
+        "Use search_web when you need real-time web search, latest documentation, or up-to-date facts with verified citations."
+      ],
+      parameters: Type.Object({
+        query: Type.String({ description: "The search query string to look up on the web." }),
+        domain: Type.Optional(
+          Type.String({
+            description: "Optional domain to recommend the search prioritize (e.g. 'github.com', 'developer.mozilla.org')."
+          })
+        )
+      }),
+      async execute(_toolCallId, params, signal, _onUpdate, ctx) {
+        const creds = await resolveCredentials(ctx);
+        if (!creds) {
+          return {
+            content: [{ type: "text", text: "Error: Not logged in. Run /login antigravity first." }],
+            details: { error: "not_logged_in" }
+          };
+        }
+        try {
+          const result = await performWebSearch(creds, {
+            query: params.query,
+            domain: params.domain,
+            signal
+          });
+          return {
+            content: [{ type: "text", text: result.formattedOutput }],
+            details: {
+              sources: result.sources,
+              queries: result.queries
+            }
+          };
+        } catch (err) {
+          return {
+            content: [{ type: "text", text: `Search failed: ${err.message}` }],
+            details: { error: err.message }
+          };
+        }
+      }
+    });
+  }
   pi.on("session_start", async (_event, ctx) => {
     await quotaStatus.refreshAndPaint(ctx);
   });