refactor(google): move thinking policy into plugin

CHANGELOG.md

@@ -78,6 +78,7 @@ Docs: https://docs.openclaw.ai
 - Gateway/auth: reject gateway auth credentials that match published example placeholders at startup and secret reload, and keep cloud install snippets from publishing copy-paste gateway/keyring secrets. (#68404) Thanks @coygeek.
 - CLI/update: preserve macOS restart helper launchctl failures in the update restart log without letting log setup block the restart path. (#68492) Thanks @hclsys.
 - Slack/threads: keep file-only root messages as starter context so first thread replies can still hydrate starter media. (#68594) Thanks @martingarramon.
+- Google/Antigravity: resolve forward-compatible Gemini 3.1 Pro custom-tools and Flash variants from the bundled Google plugin templates, so `google-antigravity/gemini-3.1-pro-preview-customtools` no longer falls through to an unknown-model error. Fixes #35512.
 
 ## 2026.4.15
 

extensions/google/api.ts

@@ -10,6 +10,19 @@ import {
   normalizeGoogleGenerativeAiBaseUrl,
 } from "./provider-policy.js";
 export { normalizeAntigravityModelId, normalizeGoogleModelId } from "./model-id.js";
+export {
+  createGoogleThinkingPayloadWrapper,
+  createGoogleThinkingStreamWrapper,
+  isGoogleGemini3FlashModel,
+  isGoogleGemini3ProModel,
+  isGoogleGemini3ThinkingLevelModel,
+  isGoogleThinkingRequiredModel,
+  resolveGoogleGemini3ThinkingLevel,
+  sanitizeGoogleThinkingPayload,
+  stripInvalidGoogleThinkingBudget,
+  type GoogleThinkingInputLevel,
+  type GoogleThinkingLevel,
+} from "./thinking-api.js";
 export {
   DEFAULT_GOOGLE_API_BASE_URL,
   isGoogleGenerativeAiApi,

extensions/google/provider-hooks.ts

@@ -1,9 +1,9 @@
 import { buildProviderReplayFamilyHooks } from "openclaw/plugin-sdk/provider-model-shared";
-import { GOOGLE_THINKING_STREAM_HOOKS } from "openclaw/plugin-sdk/provider-stream-family";
+import { createGoogleThinkingStreamWrapper } from "./thinking-api.js";
 
 export const GOOGLE_GEMINI_PROVIDER_HOOKS = {
   ...buildProviderReplayFamilyHooks({
     family: "google-gemini",
   }),
-  ...GOOGLE_THINKING_STREAM_HOOKS,
+  wrapStreamFn: createGoogleThinkingStreamWrapper,
 };

extensions/google/provider-models.test.ts

@@ -228,6 +228,112 @@ describe("resolveGoogleGeminiForwardCompatModel", () => {
     });
   });
 
+  it("resolves Antigravity Gemini 3.1 pro customtools from the low template", () => {
+    const model = resolveGoogleGeminiForwardCompatModel({
+      providerId: "google-antigravity",
+      ctx: createContext({
+        provider: "google-antigravity",
+        modelId: "gemini-3.1-pro-preview-customtools",
+        models: [
+          createTemplateModel("google-antigravity", "gemini-3-pro-low", {
+            api: "openai-completions",
+            baseUrl: "https://antigravity.example/v1",
+            contextWindow: 1_048_576,
+            reasoning: true,
+          }),
+        ],
+      }),
+    });
+
+    expect(model).toMatchObject({
+      provider: "google-antigravity",
+      id: "gemini-3.1-pro-preview-customtools",
+      api: "openai-completions",
+      baseUrl: "https://antigravity.example/v1",
+      contextWindow: 1_048_576,
+      reasoning: true,
+    });
+  });
+
+  it("falls back to the Antigravity high template when the low template is unavailable", () => {
+    const model = resolveGoogleGeminiForwardCompatModel({
+      providerId: "google-antigravity",
+      ctx: createContext({
+        provider: "google-antigravity",
+        modelId: "gemini-3.1-pro-preview",
+        models: [
+          createTemplateModel("google-antigravity", "gemini-3-pro-high", {
+            api: "openai-completions",
+            maxTokens: 65_536,
+            reasoning: true,
+          }),
+        ],
+      }),
+    });
+
+    expect(model).toMatchObject({
+      provider: "google-antigravity",
+      id: "gemini-3.1-pro-preview",
+      api: "openai-completions",
+      maxTokens: 65_536,
+      reasoning: true,
+    });
+  });
+
+  it("resolves Antigravity Gemini 3.1 flash variants from the flash template", () => {
+    const models = [
+      createTemplateModel("google-antigravity", "gemini-3-flash", {
+        api: "openai-completions",
+        contextWindow: 1_048_576,
+      }),
+    ];
+
+    expect(
+      resolveGoogleGeminiForwardCompatModel({
+        providerId: "google-antigravity",
+        ctx: createContext({
+          provider: "google-antigravity",
+          modelId: "gemini-3.1-flash-preview",
+          models,
+        }),
+      }),
+    ).toMatchObject({
+      provider: "google-antigravity",
+      id: "gemini-3.1-flash-preview",
+      api: "openai-completions",
+      contextWindow: 1_048_576,
+    });
+
+    expect(
+      resolveGoogleGeminiForwardCompatModel({
+        providerId: "google-antigravity",
+        ctx: createContext({
+          provider: "google-antigravity",
+          modelId: "gemini-3.1-flash-lite-preview",
+          models,
+        }),
+      }),
+    ).toMatchObject({
+      provider: "google-antigravity",
+      id: "gemini-3.1-flash-lite-preview",
+      api: "openai-completions",
+      contextWindow: 1_048_576,
+    });
+  });
+
+  it("returns undefined for Antigravity Gemini 3.1 models without a matching template", () => {
+    const model = resolveGoogleGeminiForwardCompatModel({
+      providerId: "google-antigravity",
+      ctx: createContext({
+        provider: "google-antigravity",
+        modelId: "gemini-3.1-pro-preview-customtools",
+        models: [createTemplateModel("google-antigravity", "claude-opus-4-6-thinking")],
+      }),
+    });
+
+    expect(model).toBeUndefined();
+  });
+
   it("prefers the flash-lite template before the broader flash prefix", () => {
     const model = resolveGoogleGeminiForwardCompatModel({
       providerId: "google-vertex",

extensions/google/provider-models.ts

@@ -6,6 +6,7 @@ import { cloneFirstTemplateModel } from "openclaw/plugin-sdk/provider-model-shar
 import { normalizeOptionalLowercaseString } from "openclaw/plugin-sdk/text-runtime";
 
 const GOOGLE_GEMINI_CLI_PROVIDER_ID = "google-gemini-cli";
+const GOOGLE_ANTIGRAVITY_PROVIDER_ID = "google-antigravity";
 const GEMINI_2_5_PRO_PREFIX = "gemini-2.5-pro";
 const GEMINI_2_5_FLASH_LITE_PREFIX = "gemini-2.5-flash-lite";
 const GEMINI_2_5_FLASH_PREFIX = "gemini-2.5-flash";
@@ -22,6 +23,8 @@ const GEMINI_2_5_FLASH_TEMPLATE_IDS = ["gemini-2.5-flash"] as const;
 const GEMINI_3_1_PRO_TEMPLATE_IDS = ["gemini-3-pro-preview"] as const;
 const GEMINI_3_1_FLASH_LITE_TEMPLATE_IDS = ["gemini-3.1-flash-lite-preview"] as const;
 const GEMINI_3_1_FLASH_TEMPLATE_IDS = ["gemini-3-flash-preview"] as const;
+const GEMINI_3_PRO_ANTIGRAVITY_TEMPLATE_IDS = ["gemini-3-pro-low", "gemini-3-pro-high"] as const;
+const GEMINI_3_FLASH_ANTIGRAVITY_TEMPLATE_IDS = ["gemini-3-flash"] as const;
 // Gemma uses the Gemini flash template as a forward-compat approximation
 // until a dedicated Gemma template is registered in the catalog.
 const GEMMA_TEMPLATE_IDS = GEMINI_3_1_FLASH_TEMPLATE_IDS;
@@ -29,6 +32,7 @@ const GEMMA_TEMPLATE_IDS = GEMINI_3_1_FLASH_TEMPLATE_IDS;
 type GoogleForwardCompatFamily = {
   googleTemplateIds: readonly string[];
   cliTemplateIds: readonly string[];
+  antigravityTemplateIds?: readonly string[];
   preferExternalFirstForCli?: boolean;
 };
 
@@ -61,13 +65,21 @@ function isGoogleGeminiCliProvider(providerId: string): boolean {
   return normalizeOptionalLowercaseString(providerId) === GOOGLE_GEMINI_CLI_PROVIDER_ID;
 }
 
+function isGoogleAntigravityProvider(providerId: string): boolean {
+  return normalizeOptionalLowercaseString(providerId) === GOOGLE_ANTIGRAVITY_PROVIDER_ID;
+}
+
 function templateIdsForProvider(
   templateProviderId: string,
   family: GoogleForwardCompatFamily,
 ): readonly string[] {
-  return isGoogleGeminiCliProvider(templateProviderId)
-    ? family.cliTemplateIds
-    : family.googleTemplateIds;
+  if (isGoogleGeminiCliProvider(templateProviderId)) {
+    return family.cliTemplateIds;
+  }
+  if (isGoogleAntigravityProvider(templateProviderId)) {
+    return family.antigravityTemplateIds ?? family.googleTemplateIds;
+  }
+  return family.googleTemplateIds;
 }
 
 function buildGoogleTemplateSources(params: {
@@ -135,6 +147,7 @@ export function resolveGoogleGeminiForwardCompatModel(params: {
     family = {
       googleTemplateIds: GEMINI_3_1_PRO_TEMPLATE_IDS,
       cliTemplateIds: GEMINI_3_1_PRO_TEMPLATE_IDS,
+      antigravityTemplateIds: GEMINI_3_PRO_ANTIGRAVITY_TEMPLATE_IDS,
     };
     if (params.providerId === "google" || params.providerId === GOOGLE_GEMINI_CLI_PROVIDER_ID) {
       patch = { reasoning: true };
@@ -146,11 +159,13 @@ export function resolveGoogleGeminiForwardCompatModel(params: {
     family = {
       googleTemplateIds: GEMINI_3_1_FLASH_LITE_TEMPLATE_IDS,
       cliTemplateIds: GEMINI_3_1_FLASH_LITE_TEMPLATE_IDS,
+      antigravityTemplateIds: GEMINI_3_FLASH_ANTIGRAVITY_TEMPLATE_IDS,
     };
   } else if (lower.startsWith(GEMINI_3_1_FLASH_PREFIX) || lower === GEMINI_FLASH_LATEST_ID) {
     family = {
       googleTemplateIds: GEMINI_3_1_FLASH_TEMPLATE_IDS,
       cliTemplateIds: GEMINI_3_1_FLASH_TEMPLATE_IDS,
+      antigravityTemplateIds: GEMINI_3_FLASH_ANTIGRAVITY_TEMPLATE_IDS,
     };
   } else if (lower.startsWith(GEMMA_PREFIX)) {
     family = {

extensions/google/runtime-api.ts

@@ -1,7 +1,17 @@
 export {
   DEFAULT_GOOGLE_API_BASE_URL,
+  createGoogleThinkingPayloadWrapper,
+  createGoogleThinkingStreamWrapper,
+  isGoogleGemini3FlashModel,
+  isGoogleGemini3ProModel,
+  isGoogleGemini3ThinkingLevelModel,
+  isGoogleThinkingRequiredModel,
   normalizeGoogleApiBaseUrl,
   normalizeGoogleModelId,
   parseGeminiAuth,
+  resolveGoogleGemini3ThinkingLevel,
   resolveGoogleGenerativeAiHttpRequestConfig,
+  sanitizeGoogleThinkingPayload,
+  stripInvalidGoogleThinkingBudget,
 } from "./api.js";
+export type { GoogleThinkingInputLevel, GoogleThinkingLevel } from "./api.js";

extensions/google/thinking-api.ts

@@ -0,0 +1,13 @@
+export {
+  createGoogleThinkingPayloadWrapper,
+  createGoogleThinkingStreamWrapper,
+  isGoogleGemini3FlashModel,
+  isGoogleGemini3ProModel,
+  isGoogleGemini3ThinkingLevelModel,
+  isGoogleThinkingRequiredModel,
+  resolveGoogleGemini3ThinkingLevel,
+  sanitizeGoogleThinkingPayload,
+  stripInvalidGoogleThinkingBudget,
+  type GoogleThinkingInputLevel,
+  type GoogleThinkingLevel,
+} from "./thinking.js";

extensions/google/thinking.test.ts

@@ -0,0 +1,116 @@
+import { describe, expect, it } from "vitest";
+import {
+  resolveGoogleGemini3ThinkingLevel,
+  sanitizeGoogleThinkingPayload,
+} from "./thinking-api.js";
+
+describe("google thinking policy", () => {
+  it.each([
+    ["off", "LOW"],
+    ["minimal", "LOW"],
+    ["low", "LOW"],
+    ["medium", "HIGH"],
+    ["adaptive", "HIGH"],
+    ["high", "HIGH"],
+    ["xhigh", "HIGH"],
+  ] as const)("maps Gemini 3 Pro thinking level %s to %s", (thinkingLevel, expected) => {
+    expect(
+      resolveGoogleGemini3ThinkingLevel({
+        modelId: "gemini-3.1-pro-preview",
+        thinkingLevel,
+      }),
+    ).toBe(expected);
+  });
+
+  it.each([
+    [0, "LOW"],
+    [2048, "LOW"],
+    [2049, "HIGH"],
+  ] as const)("maps Gemini 3 Pro budget %s to %s", (thinkingBudget, expected) => {
+    expect(
+      resolveGoogleGemini3ThinkingLevel({
+        modelId: "gemini-pro-latest",
+        thinkingBudget,
+      }),
+    ).toBe(expected);
+  });
+
+  it.each([
+    ["off", "MINIMAL"],
+    ["minimal", "MINIMAL"],
+    ["low", "LOW"],
+    ["medium", "MEDIUM"],
+    ["adaptive", "MEDIUM"],
+    ["high", "HIGH"],
+    ["xhigh", "HIGH"],
+  ] as const)("maps Gemini 3 Flash thinking level %s to %s", (thinkingLevel, expected) => {
+    expect(
+      resolveGoogleGemini3ThinkingLevel({
+        modelId: "gemini-flash-latest",
+        thinkingLevel,
+      }),
+    ).toBe(expected);
+  });
+
+  it.each([
+    [-1, "MINIMAL"],
+    [0, "MINIMAL"],
+    [2048, "LOW"],
+    [8192, "MEDIUM"],
+    [8193, "HIGH"],
+  ] as const)("maps Gemini 3 Flash budget %s to %s", (thinkingBudget, expected) => {
+    expect(
+      resolveGoogleGemini3ThinkingLevel({
+        modelId: "gemini-3.1-flash-lite-preview",
+        thinkingBudget,
+      }),
+    ).toBe(expected);
+  });
+
+  it("removes thinkingBudget=0 for Gemini 2.5 Pro", () => {
+    const payload = {
+      config: {
+        thinkingConfig: { thinkingBudget: 0 },
+      },
+    };
+
+    sanitizeGoogleThinkingPayload({ payload, modelId: "google/gemini-2.5-pro-preview" });
+
+    expect(payload.config).not.toHaveProperty("thinkingConfig");
+  });
+
+  it("rewrites Gemini 3 thinking budgets to thinkingLevel", () => {
+    const payload = {
+      generationConfig: {
+        thinkingConfig: { thinkingBudget: 8193, includeThoughts: true },
+      },
+    };
+
+    sanitizeGoogleThinkingPayload({
+      payload,
+      modelId: "gemini-3-flash-preview",
+      thinkingLevel: "medium",
+    });
+
+    expect(payload.generationConfig.thinkingConfig).toEqual({
+      includeThoughts: true,
+      thinkingLevel: "MEDIUM",
+    });
+  });
+
+  it("maps Gemma 4 thinking mode without sending thinkingBudget", () => {
+    const payload = {
+      config: {
+        thinkingConfig: { thinkingBudget: 4096 },
+      },
+    };
+
+    sanitizeGoogleThinkingPayload({
+      payload,
+      modelId: "gemma-4-26b-a4b-it",
+      thinkingLevel: "high",
+    });
+
+    expect(payload.config.thinkingConfig).toEqual({ thinkingLevel: "HIGH" });
+  });
+});

extensions/google/thinking.ts

@@ -0,0 +1,263 @@
+import type { ProviderWrapStreamFnContext } from "openclaw/plugin-sdk/plugin-entry";
+import { createPayloadPatchStreamWrapper } from "openclaw/plugin-sdk/provider-stream-shared";
+import { normalizeLowercaseStringOrEmpty } from "openclaw/plugin-sdk/text-runtime";
+
+export type GoogleThinkingLevel = "MINIMAL" | "LOW" | "MEDIUM" | "HIGH";
+export type GoogleThinkingInputLevel =
+  | "off"
+  | "minimal"
+  | "low"
+  | "medium"
+  | "adaptive"
+  | "high"
+  | "xhigh";
+
+export function isGoogleThinkingRequiredModel(modelId: string): boolean {
+  return normalizeLowercaseStringOrEmpty(modelId).includes("gemini-2.5-pro");
+}
+
+export function isGoogleGemini3ProModel(modelId: string): boolean {
+  const normalized = normalizeLowercaseStringOrEmpty(modelId);
+  return /(?:^|\/)gemini-(?:3(?:\.\d+)?-pro|pro-latest)(?:-|$)/.test(normalized);
+}
+
+export function isGoogleGemini3FlashModel(modelId: string): boolean {
+  const normalized = normalizeLowercaseStringOrEmpty(modelId);
+  return /(?:^|\/)gemini-(?:3(?:\.\d+)?-flash|flash(?:-lite)?-latest)(?:-|$)/.test(normalized);
+}
+
+export function isGoogleGemini3ThinkingLevelModel(modelId: string): boolean {
+  return isGoogleGemini3ProModel(modelId) || isGoogleGemini3FlashModel(modelId);
+}
+
+export function resolveGoogleGemini3ThinkingLevel(params: {
+  modelId?: string;
+  thinkingLevel?: GoogleThinkingInputLevel;
+  thinkingBudget?: number;
+}): GoogleThinkingLevel | undefined {
+  if (typeof params.modelId !== "string") {
+    return undefined;
+  }
+  if (isGoogleGemini3ProModel(params.modelId)) {
+    switch (params.thinkingLevel) {
+      case "off":
+      case "minimal":
+      case "low":
+        return "LOW";
+      case "medium":
+      case "adaptive":
+      case "high":
+      case "xhigh":
+        return "HIGH";
+    }
+    if (typeof params.thinkingBudget === "number") {
+      return params.thinkingBudget <= 2048 ? "LOW" : "HIGH";
+    }
+    return undefined;
+  }
+  if (!isGoogleGemini3FlashModel(params.modelId)) {
+    return undefined;
+  }
+  switch (params.thinkingLevel) {
+    case "off":
+    case "minimal":
+      return "MINIMAL";
+    case "low":
+      return "LOW";
+    case "medium":
+    case "adaptive":
+      return "MEDIUM";
+    case "high":
+    case "xhigh":
+      return "HIGH";
+  }
+  if (typeof params.thinkingBudget !== "number") {
+    return undefined;
+  }
+  if (params.thinkingBudget <= 0) {
+    return "MINIMAL";
+  }
+  if (params.thinkingBudget <= 2048) {
+    return "LOW";
+  }
+  if (params.thinkingBudget <= 8192) {
+    return "MEDIUM";
+  }
+  return "HIGH";
+}
+
+export function stripInvalidGoogleThinkingBudget(params: {
+  thinkingConfig: Record<string, unknown>;
+  modelId?: string;
+}): boolean {
+  if (
+    params.thinkingConfig.thinkingBudget !== 0 ||
+    typeof params.modelId !== "string" ||
+    !isGoogleThinkingRequiredModel(params.modelId)
+  ) {
+    return false;
+  }
+  delete params.thinkingConfig.thinkingBudget;
+  return true;
+}
+
+function isGemma4Model(modelId: string): boolean {
+  return normalizeLowercaseStringOrEmpty(modelId).startsWith("gemma-4");
+}
+
+function mapThinkLevelToGemma4ThinkingLevel(
+  thinkingLevel?: GoogleThinkingInputLevel,
+): "MINIMAL" | "HIGH" | undefined {
+  switch (thinkingLevel) {
+    case "off":
+      return undefined;
+    case "minimal":
+    case "low":
+      return "MINIMAL";
+    case "medium":
+    case "adaptive":
+    case "high":
+    case "xhigh":
+      return "HIGH";
+    default:
+      return undefined;
+  }
+}
+
+function normalizeGemma4ThinkingLevel(value: unknown): "MINIMAL" | "HIGH" | undefined {
+  if (typeof value !== "string") {
+    return undefined;
+  }
+  switch (value.trim().toUpperCase()) {
+    case "MINIMAL":
+    case "LOW":
+      return "MINIMAL";
+    case "MEDIUM":
+    case "HIGH":
+      return "HIGH";
+    default:
+      return undefined;
+  }
+}
+
+export function sanitizeGoogleThinkingPayload(params: {
+  payload: unknown;
+  modelId?: string;
+  thinkingLevel?: GoogleThinkingInputLevel;
+}): void {
+  if (!params.payload || typeof params.payload !== "object") {
+    return;
+  }
+  const payloadObj = params.payload as Record<string, unknown>;
+  sanitizeGoogleThinkingConfigContainer({
+    container: payloadObj.config,
+    modelId: params.modelId,
+    thinkingLevel: params.thinkingLevel,
+  });
+  sanitizeGoogleThinkingConfigContainer({
+    container: payloadObj.generationConfig,
+    modelId: params.modelId,
+    thinkingLevel: params.thinkingLevel,
+  });
+}
+
+function sanitizeGoogleThinkingConfigContainer(params: {
+  container: unknown;
+  modelId?: string;
+  thinkingLevel?: GoogleThinkingInputLevel;
+}): void {
+  if (!params.container || typeof params.container !== "object") {
+    return;
+  }
+  const configObj = params.container as Record<string, unknown>;
+  const thinkingConfig = configObj.thinkingConfig;
+  if (!thinkingConfig || typeof thinkingConfig !== "object") {
+    return;
+  }
+  const thinkingConfigObj = thinkingConfig as Record<string, unknown>;
+
+  if (typeof params.modelId === "string" && isGemma4Model(params.modelId)) {
+    const normalizedThinkingLevel = normalizeGemma4ThinkingLevel(thinkingConfigObj.thinkingLevel);
+    const explicitMappedLevel = mapThinkLevelToGemma4ThinkingLevel(params.thinkingLevel);
+    const disabledViaBudget =
+      typeof thinkingConfigObj.thinkingBudget === "number" && thinkingConfigObj.thinkingBudget <= 0;
+    const hadThinkingBudget = thinkingConfigObj.thinkingBudget !== undefined;
+    delete thinkingConfigObj.thinkingBudget;
+
+    if (
+      params.thinkingLevel === "off" ||
+      (disabledViaBudget && explicitMappedLevel === undefined && !normalizedThinkingLevel)
+    ) {
+      delete thinkingConfigObj.thinkingLevel;
+      if (Object.keys(thinkingConfigObj).length === 0) {
+        delete configObj.thinkingConfig;
+      }
+      return;
+    }
+
+    const mappedLevel =
+      explicitMappedLevel ?? normalizedThinkingLevel ?? (hadThinkingBudget ? "MINIMAL" : undefined);
+
+    if (mappedLevel) {
+      thinkingConfigObj.thinkingLevel = mappedLevel;
+    }
+    return;
+  }
+
+  const thinkingBudget = thinkingConfigObj.thinkingBudget;
+
+  if (typeof params.modelId === "string" && isGoogleGemini3ThinkingLevelModel(params.modelId)) {
+    const mappedLevel = resolveGoogleGemini3ThinkingLevel({
+      modelId: params.modelId,
+      thinkingLevel: params.thinkingLevel,
+      thinkingBudget: typeof thinkingBudget === "number" ? thinkingBudget : undefined,
+    });
+    delete thinkingConfigObj.thinkingBudget;
+    if (mappedLevel) {
+      thinkingConfigObj.thinkingLevel = mappedLevel;
+    }
+    if (Object.keys(thinkingConfigObj).length === 0) {
+      delete configObj.thinkingConfig;
+    }
+    return;
+  }
+
+  if (
+    stripInvalidGoogleThinkingBudget({ thinkingConfig: thinkingConfigObj, modelId: params.modelId })
+  ) {
+    if (Object.keys(thinkingConfigObj).length === 0) {
+      delete configObj.thinkingConfig;
+    }
+    return;
+  }
+
+  if (typeof thinkingBudget !== "number" || thinkingBudget >= 0) {
+    return;
+  }
+
+  delete thinkingConfigObj.thinkingBudget;
+  if (Object.keys(thinkingConfigObj).length === 0) {
+    delete configObj.thinkingConfig;
+  }
+}
+
+export function createGoogleThinkingPayloadWrapper(
+  baseStreamFn: ProviderWrapStreamFnContext["streamFn"],
+  thinkingLevel?: GoogleThinkingInputLevel,
+): NonNullable<ProviderWrapStreamFnContext["streamFn"]> {
+  return createPayloadPatchStreamWrapper(baseStreamFn, ({ payload, model }) => {
+    if (model.api === "google-generative-ai") {
+      sanitizeGoogleThinkingPayload({
+        payload,
+        modelId: model.id,
+        thinkingLevel,
+      });
+    }
+  });
+}
+
+export function createGoogleThinkingStreamWrapper(
+  ctx: ProviderWrapStreamFnContext,
+): NonNullable<ProviderWrapStreamFnContext["streamFn"]> {
+  return createGoogleThinkingPayloadWrapper(ctx.streamFn, ctx.thinkingLevel);
+}

src/agents/google-thinking-compat.ts

@@ -1,102 +1,10 @@
-import { normalizeLowercaseStringOrEmpty } from "../shared/string-coerce.js";
-
-export type GoogleThinkingLevel = "MINIMAL" | "LOW" | "MEDIUM" | "HIGH";
-export type GoogleThinkingInputLevel =
-  | "off"
-  | "minimal"
-  | "low"
-  | "medium"
-  | "adaptive"
-  | "high"
-  | "xhigh";
-
-// Gemini 2.5 Pro only works in thinking mode and rejects thinkingBudget=0 with
-// "Budget 0 is invalid. This model only works in thinking mode."
-export function isGoogleThinkingRequiredModel(modelId: string): boolean {
-  return normalizeLowercaseStringOrEmpty(modelId).includes("gemini-2.5-pro");
-}
-
-export function isGoogleGemini3ProModel(modelId: string): boolean {
-  const normalized = normalizeLowercaseStringOrEmpty(modelId);
-  return /(?:^|\/)gemini-(?:3(?:\.\d+)?-pro|pro-latest)(?:-|$)/.test(normalized);
-}
-
-export function isGoogleGemini3FlashModel(modelId: string): boolean {
-  const normalized = normalizeLowercaseStringOrEmpty(modelId);
-  return /(?:^|\/)gemini-(?:3(?:\.\d+)?-flash|flash(?:-lite)?-latest)(?:-|$)/.test(normalized);
-}
-
-export function isGoogleGemini3ThinkingLevelModel(modelId: string): boolean {
-  return isGoogleGemini3ProModel(modelId) || isGoogleGemini3FlashModel(modelId);
-}
-
-export function resolveGoogleGemini3ThinkingLevel(params: {
-  modelId?: string;
-  thinkingLevel?: GoogleThinkingInputLevel;
-  thinkingBudget?: number;
-}): GoogleThinkingLevel | undefined {
-  if (typeof params.modelId !== "string") {
-    return undefined;
-  }
-  if (isGoogleGemini3ProModel(params.modelId)) {
-    switch (params.thinkingLevel) {
-      case "off":
-      case "minimal":
-      case "low":
-        return "LOW";
-      case "medium":
-      case "adaptive":
-      case "high":
-      case "xhigh":
-        return "HIGH";
-    }
-    if (typeof params.thinkingBudget === "number") {
-      return params.thinkingBudget <= 2048 ? "LOW" : "HIGH";
-    }
-    return undefined;
-  }
-  if (!isGoogleGemini3FlashModel(params.modelId)) {
-    return undefined;
-  }
-  switch (params.thinkingLevel) {
-    case "off":
-    case "minimal":
-      return "MINIMAL";
-    case "low":
-      return "LOW";
-    case "medium":
-    case "adaptive":
-      return "MEDIUM";
-    case "high":
-    case "xhigh":
-      return "HIGH";
-  }
-  if (typeof params.thinkingBudget !== "number") {
-    return undefined;
-  }
-  if (params.thinkingBudget <= 0) {
-    return "MINIMAL";
-  }
-  if (params.thinkingBudget <= 2048) {
-    return "LOW";
-  }
-  if (params.thinkingBudget <= 8192) {
-    return "MEDIUM";
-  }
-  return "HIGH";
-}
-
-export function stripInvalidGoogleThinkingBudget(params: {
-  thinkingConfig: Record<string, unknown>;
-  modelId?: string;
-}): boolean {
-  if (
-    params.thinkingConfig.thinkingBudget !== 0 ||
-    typeof params.modelId !== "string" ||
-    !isGoogleThinkingRequiredModel(params.modelId)
-  ) {
-    return false;
-  }
-  delete params.thinkingConfig.thinkingBudget;
-  return true;
-}
+export {
+  isGoogleGemini3FlashModel,
+  isGoogleGemini3ProModel,
+  isGoogleGemini3ThinkingLevelModel,
+  isGoogleThinkingRequiredModel,
+  resolveGoogleGemini3ThinkingLevel,
+  stripInvalidGoogleThinkingBudget,
+  type GoogleThinkingInputLevel,
+  type GoogleThinkingLevel,
+} from "../../extensions/google/thinking-api.js";

src/agents/pi-embedded-runner/google-stream-wrappers.ts

@@ -1,170 +1,4 @@
-import type { StreamFn } from "@mariozechner/pi-agent-core";
-import { streamSimple } from "@mariozechner/pi-ai";
-import type { ThinkLevel } from "../../auto-reply/thinking.js";
-import { normalizeLowercaseStringOrEmpty } from "../../shared/string-coerce.js";
-import {
-  isGoogleGemini3ThinkingLevelModel,
-  resolveGoogleGemini3ThinkingLevel,
-  stripInvalidGoogleThinkingBudget,
-} from "../google-thinking-compat.js";
-import { streamWithPayloadPatch } from "./stream-payload-utils.js";
-
-function isGemma4Model(modelId: string): boolean {
-  return normalizeLowercaseStringOrEmpty(modelId).startsWith("gemma-4");
-}
-
-function mapThinkLevelToGemma4ThinkingLevel(
-  thinkingLevel?: ThinkLevel,
-): "MINIMAL" | "HIGH" | undefined {
-  switch (thinkingLevel) {
-    case "off":
-      return undefined;
-    case "minimal":
-    case "low":
-      return "MINIMAL";
-    case "medium":
-    case "adaptive":
-    case "high":
-    case "xhigh":
-      return "HIGH";
-    default:
-      return undefined;
-  }
-}
-
-function normalizeGemma4ThinkingLevel(value: unknown): "MINIMAL" | "HIGH" | undefined {
-  if (typeof value !== "string") {
-    return undefined;
-  }
-  switch (value.trim().toUpperCase()) {
-    case "MINIMAL":
-    case "LOW":
-      return "MINIMAL";
-    case "MEDIUM":
-    case "HIGH":
-      return "HIGH";
-    default:
-      return undefined;
-  }
-}
-
-export function sanitizeGoogleThinkingPayload(params: {
-  payload: unknown;
-  modelId?: string;
-  thinkingLevel?: ThinkLevel;
-}): void {
-  if (!params.payload || typeof params.payload !== "object") {
-    return;
-  }
-  const payloadObj = params.payload as Record<string, unknown>;
-  sanitizeGoogleThinkingConfigContainer({
-    container: payloadObj.config,
-    modelId: params.modelId,
-    thinkingLevel: params.thinkingLevel,
-  });
-  sanitizeGoogleThinkingConfigContainer({
-    container: payloadObj.generationConfig,
-    modelId: params.modelId,
-    thinkingLevel: params.thinkingLevel,
-  });
-}
-
-function sanitizeGoogleThinkingConfigContainer(params: {
-  container: unknown;
-  modelId?: string;
-  thinkingLevel?: ThinkLevel;
-}): void {
-  if (!params.container || typeof params.container !== "object") {
-    return;
-  }
-  const configObj = params.container as Record<string, unknown>;
-  const thinkingConfig = configObj.thinkingConfig;
-  if (!thinkingConfig || typeof thinkingConfig !== "object") {
-    return;
-  }
-  const thinkingConfigObj = thinkingConfig as Record<string, unknown>;
-
-  if (typeof params.modelId === "string" && isGemma4Model(params.modelId)) {
-    const normalizedThinkingLevel = normalizeGemma4ThinkingLevel(thinkingConfigObj.thinkingLevel);
-    const explicitMappedLevel = mapThinkLevelToGemma4ThinkingLevel(params.thinkingLevel);
-    const disabledViaBudget =
-      typeof thinkingConfigObj.thinkingBudget === "number" && thinkingConfigObj.thinkingBudget <= 0;
-    const hadThinkingBudget = thinkingConfigObj.thinkingBudget !== undefined;
-    delete thinkingConfigObj.thinkingBudget;
-
-    if (
-      params.thinkingLevel === "off" ||
-      (disabledViaBudget && explicitMappedLevel === undefined && !normalizedThinkingLevel)
-    ) {
-      delete thinkingConfigObj.thinkingLevel;
-      if (Object.keys(thinkingConfigObj).length === 0) {
-        delete configObj.thinkingConfig;
-      }
-      return;
-    }
-
-    const mappedLevel =
-      explicitMappedLevel ?? normalizedThinkingLevel ?? (hadThinkingBudget ? "MINIMAL" : undefined);
-
-    if (mappedLevel) {
-      thinkingConfigObj.thinkingLevel = mappedLevel;
-    }
-    return;
-  }
-
-  const thinkingBudget = thinkingConfigObj.thinkingBudget;
-
-  if (typeof params.modelId === "string" && isGoogleGemini3ThinkingLevelModel(params.modelId)) {
-    const mappedLevel = resolveGoogleGemini3ThinkingLevel({
-      modelId: params.modelId,
-      thinkingLevel: params.thinkingLevel,
-      thinkingBudget: typeof thinkingBudget === "number" ? thinkingBudget : undefined,
-    });
-    delete thinkingConfigObj.thinkingBudget;
-    if (mappedLevel) {
-      thinkingConfigObj.thinkingLevel = mappedLevel;
-    }
-    if (Object.keys(thinkingConfigObj).length === 0) {
-      delete configObj.thinkingConfig;
-    }
-    return;
-  }
-
-  if (
-    stripInvalidGoogleThinkingBudget({ thinkingConfig: thinkingConfigObj, modelId: params.modelId })
-  ) {
-    if (Object.keys(thinkingConfigObj).length === 0) {
-      delete configObj.thinkingConfig;
-    }
-    return;
-  }
-
-  if (typeof thinkingBudget !== "number" || thinkingBudget >= 0) {
-    return;
-  }
-
-  // pi-ai can emit thinkingBudget=-1 for some Google model IDs; a negative budget
-  // is invalid for Google-compatible backends and can lead to malformed handling.
-  delete thinkingConfigObj.thinkingBudget;
-  if (Object.keys(thinkingConfigObj).length === 0) {
-    delete configObj.thinkingConfig;
-  }
-}
-
-export function createGoogleThinkingPayloadWrapper(
-  baseStreamFn: StreamFn | undefined,
-  thinkingLevel?: ThinkLevel,
-): StreamFn {
-  const underlying = baseStreamFn ?? streamSimple;
-  return (model, context, options) => {
-    return streamWithPayloadPatch(underlying, model, context, options, (payload) => {
-      if (model.api === "google-generative-ai") {
-        sanitizeGoogleThinkingPayload({
-          payload,
-          modelId: model.id,
-          thinkingLevel,
-        });
-      }
-    });
-  };
-}
+export {
+  createGoogleThinkingPayloadWrapper,
+  sanitizeGoogleThinkingPayload,
+} from "../../../extensions/google/thinking-api.js";

src/plugin-sdk/provider-stream-shared.ts

@@ -1,5 +1,6 @@
 import type { StreamFn } from "@mariozechner/pi-agent-core";
 import { streamSimple } from "@mariozechner/pi-ai";
+import { streamWithPayloadPatch } from "../agents/pi-embedded-runner/stream-payload-utils.js";
 
 export type ProviderStreamWrapperFactory =
   | ((streamFn: StreamFn | undefined) => StreamFn | undefined)
@@ -132,6 +133,20 @@ export function createHtmlEntityToolCallArgumentDecodingWrapper(
   };
 }
 
+export function createPayloadPatchStreamWrapper(
+  baseStreamFn: StreamFn | undefined,
+  patchPayload: (params: {
+    payload: Record<string, unknown>;
+    model: Parameters<StreamFn>[0];
+  }) => void,
+): StreamFn {
+  const underlying = baseStreamFn ?? streamSimple;
+  return (model, context, options) =>
+    streamWithPayloadPatch(underlying, model, context, options, (payload) =>
+      patchPayload({ payload, model }),
+    );
+}
+
 export {
   applyAnthropicPayloadPolicyToParams,
   resolveAnthropicPayloadPolicy,
@@ -149,7 +164,7 @@ export {
   createMoonshotThinkingWrapper,
   resolveMoonshotThinkingType,
 } from "../agents/pi-embedded-runner/moonshot-thinking-stream-wrappers.js";
-export { streamWithPayloadPatch } from "../agents/pi-embedded-runner/stream-payload-utils.js";
+export { streamWithPayloadPatch };
 export {
   createToolStreamWrapper,
   createZaiToolStreamWrapper,