fix(ai): preserve Google simple completion sanitation

This commit is contained in:
Vincent Koc
2026-08-04 13:56:10 +08:00
parent 48a759f6a0
commit 9c854cb1a0
3 changed files with 57 additions and 5 deletions
@@ -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,
@@ -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,
@@ -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) => {