From 9c854cb1a0514101d01505cc784a839f4e015bef Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Tue, 4 Aug 2026 13:56:10 +0800
Subject: [PATCH] fix(ai): preserve Google simple completion sanitation

---
 .../simple-completion-transport.test.ts       | 19 +++++++---
 .../transports/simple-completion-transport.ts |  5 +++
 .../google-simple-completion-stream.test.ts   | 38 +++++++++++++++++++
 3 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/packages/ai/src/transports/simple-completion-transport.test.ts b/packages/ai/src/transports/simple-completion-transport.test.ts
index fc9b6622d304..c10bbf1bb735 100644
--- a/packages/ai/src/transports/simple-completion-transport.test.ts
+++ b/packages/ai/src/transports/simple-completion-transport.test.ts
@@ -354,10 +354,10 @@ describe("prepareModelForSimpleCompletion", () => {
     });
   });
 
-  it("uses the Google simple-completion sanitizer alias after transport checks pass through", () => {
+  it("keeps registered Google models on the sanitizer path when the provider owns a stream", () => {
     const model: Model<"google-generative-ai"> = {
-      id: "gemini-flash-latest",
-      name: "Gemini Flash Latest",
+      id: "gemma-4-26b-a4b-it",
+      name: "Gemma 4 26B",
       api: "google-generative-ai",
       provider: "google",
       baseUrl: "https://generativelanguage.googleapis.com",
@@ -368,14 +368,23 @@ describe("prepareModelForSimpleCompletion", () => {
       maxTokens: 8192,
       headers: {},
     };
+    const googleStream = vi.fn(() => createAssistantMessageEventStream());
+    apiRegistry.registerApiProvider(
+      {
+        api: "google-generative-ai",
+        stream: googleStream,
+        streamSimple: googleStream,
+      },
+      SIMPLE_COMPLETION_SOURCE_ID,
+    );
     prepareGoogleSimpleCompletionModel.mockImplementationOnce((_registry: unknown, m: unknown) => ({
       ...(m as Model<"google-generative-ai">),
       api: "openclaw-google-generative-ai-simple",
     }));
-    resolveProviderStreamFn.mockReturnValueOnce(undefined);
 
     const result = prepareModelForSimpleCompletion({ model });
 
+    expect(resolveProviderStreamFn).not.toHaveBeenCalled();
     expect(prepareTransportAwareSimpleModel).toHaveBeenCalledWith(model, { cfg: undefined });
     expect(prepareGoogleSimpleCompletionModel).toHaveBeenCalledWith(apiRegistry, model);
     expect(buildTransportAwareSimpleStreamFn).not.toHaveBeenCalled();
@@ -404,12 +413,12 @@ describe("prepareModelForSimpleCompletion", () => {
       ...model,
       api: "openclaw-google-generative-ai-transport",
     };
-    resolveProviderStreamFn.mockReturnValueOnce(undefined);
     buildTransportAwareSimpleStreamFn.mockReturnValueOnce("google-transport-stream");
     prepareTransportAwareSimpleModel.mockReturnValueOnce(transportModel);
 
     const result = prepareModelForSimpleCompletion({ model });
 
+    expect(resolveProviderStreamFn).not.toHaveBeenCalled();
     expect(buildTransportAwareSimpleStreamFn).toHaveBeenCalledWith(model, { cfg: undefined });
     expect(ensureCustomApiRegistered).toHaveBeenCalledWith(
       apiRegistry,
diff --git a/packages/ai/src/transports/simple-completion-transport.ts b/packages/ai/src/transports/simple-completion-transport.ts
index 7b58fbe25a0b..c61b11ad08cb 100644
--- a/packages/ai/src/transports/simple-completion-transport.ts
+++ b/packages/ai/src/transports/simple-completion-transport.ts
@@ -177,6 +177,11 @@ function prepareProviderStreamModel<TApi extends Api>(params: {
   cfg?: unknown;
   apiRegistry: ApiRegistry;
 }): Model | undefined {
+  // Google simple completions have managed transport and sanitizer paths below.
+  // A plugin-native stream here would bypass both and emit unsupported payloads.
+  if (params.model.api === "google-generative-ai") {
+    return undefined;
+  }
   const pluginModel = resolveModelHeaderSentinels(params.model);
   const providerStreamFn = getAiTransportHost().plugin.resolveProviderStream({
     provider: params.model.provider,
diff --git a/src/agents/google-simple-completion-stream.test.ts b/src/agents/google-simple-completion-stream.test.ts
index c6f98aa21081..c8540f3aa035 100644
--- a/src/agents/google-simple-completion-stream.test.ts
+++ b/src/agents/google-simple-completion-stream.test.ts
@@ -156,6 +156,44 @@ describe("prepareGoogleSimpleCompletionModel", () => {
     ).not.toHaveProperty("thinkingBudget");
   });
 
+  it("removes disabled thinking budget for Gemma 4 when reasoning is omitted", async () => {
+    const actual = await vi.importActual<
+      typeof import("../llm/providers/stream-wrappers/google-thinking-payload.js")
+    >("../llm/providers/stream-wrappers/google-thinking-payload.js");
+    sanitizeGoogleThinkingPayload.mockImplementationOnce(actual.sanitizeGoogleThinkingPayload);
+    streamSimple.mockImplementationOnce((_model, _context, options) => {
+      const payload = {
+        generationConfig: {
+          thinkingConfig: { thinkingBudget: 0 },
+        },
+      };
+      options?.onPayload?.(payload, _model);
+      return { content: [{ type: "text", text: "ok" }], payload };
+    });
+    const model = makeGoogleModel("gemma-4-26b-a4b-it");
+    const wrapped = prepareGoogleSimpleCompletionModel(apiRegistry, model);
+    const streamFn = ensureCustomApiRegistered.mock.calls[0]?.[2] as (
+      ...args: unknown[]
+    ) => unknown;
+
+    const result = await streamFn(wrapped, { messages: [] }, { apiKey: "key" });
+
+    expect(sanitizeGoogleThinkingPayload).toHaveBeenCalledWith({
+      payload: {
+        generationConfig: {},
+      },
+      modelId: "gemma-4-26b-a4b-it",
+      thinkingLevel: undefined,
+    });
+    expect(
+      (
+        result as {
+          payload: { generationConfig: Record<string, unknown> };
+        }
+      ).payload.generationConfig,
+    ).not.toHaveProperty("thinkingConfig");
+  });
+
   it.each(["xhigh", "max"] as const)(
     "preserves clamped-off intent in the final Gemini 3 payload for reasoning=%s",
     async (reasoning) => {