import type * as GoogleGenAi from "@google/genai";
import type { GenerateContentParameters } from "@google/genai";
import { describe, expect, it, vi } from "vitest";
import { getModel } from "../src/models.js";
import {
convertMessages,
convertTools,
isThinkingPart,
retainThoughtSignature,
} from "../src/providers/google-shared.js";
import { streamSimpleGoogleVertex } from "../src/providers/google-vertex.js";
import type { Context, Model, Tool } from "../src/types.js";
vi.mock("@google/genai", async (importOriginal) => {
const actual = await importOriginal<typeof GoogleGenAi>();
class GoogleGenAI {
models = {
generateContentStream: async function* () {
yield {
candidates: [{ content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }],
usageMetadata: { promptTokenCount: 1, candidatesTokenCount: 1, totalTokenCount: 2 },
};
},
};
}
return {
...actual,
GoogleGenAI,
ResourceScope: { COLLECTION: "COLLECTION" },
ThinkingLevel: {
THINKING_LEVEL_UNSPECIFIED: "THINKING_LEVEL_UNSPECIFIED",
MINIMAL: "MINIMAL",
LOW: "LOW",
MEDIUM: "MEDIUM",
HIGH: "HIGH",
},
};
});
function makeTool(parameters: Record<string, unknown>): Tool {
return {
name: "test_tool",
description: "A test tool",
parameters: parameters as Tool["parameters"],
};
}
describe("google-shared convertTools", () => {
it("strips JSON Schema meta keys from parameters when useParameters=true", () => {
const tools = [
makeTool({
$schema: "http://json-schema.org/draft-07/schema#",
$id: "urn:bash-tool",
$comment: "A bash tool for demonstration",
$defs: {
commandDef: { type: "string" },
},
definitions: {
legacyDef: { type: "number" },
},
type: "object",
properties: {
command: { type: "string" },
},
required: ["command"],
}),
];
const result = convertTools(tools, true);
const decl = result?.[0]?.functionDeclarations?.[0];
expect(decl).toBeDefined();
expect(decl?.parameters).toEqual({
type: "object",
properties: {
command: { type: "string" },
},
required: ["command"],
});
expect(decl?.parameters).not.toHaveProperty("$schema");
expect(decl?.parameters).not.toHaveProperty("$id");
expect(decl?.parameters).not.toHaveProperty("$comment");
expect(decl?.parameters).not.toHaveProperty("$defs");
expect(decl?.parameters).not.toHaveProperty("definitions");
});
it("recursively strips nested JSON Schema meta keys", () => {
const tools = [
makeTool({
$schema: "http://json-schema.org/draft-07/schema#",
type: "object",
properties: {
deep: {
$schema: "http://json-schema.org/draft-07/schema#",
$id: "urn:nested",
type: "string",
},
},
}),
];
const result = convertTools(tools, true);
const decl = result?.[0]?.functionDeclarations?.[0];
expect(decl).toBeDefined();
expect(decl?.parameters).toEqual({
type: "object",
properties: {
deep: {
type: "string",
},
},
});
});
it("preserves $ref while stripping meta keys", () => {
const tools = [
makeTool({
$schema: "http://json-schema.org/draft-07/schema#",
type: "object",
properties: {
refProp: {
$ref: "#/$defs/someDef",
type: "string",
},
},
}),
];
const result = convertTools(tools, true);
const decl = result?.[0]?.functionDeclarations?.[0];
expect(decl).toBeDefined();
expect(decl?.parameters).toEqual({
type: "object",
properties: {
refProp: {
$ref: "#/$defs/someDef",
type: "string",
},
},
});
});
it("does not mutate the original Tool.parameters object", () => {
const originalParameters = {
$schema: "http://json-schema.org/draft-07/schema#",
type: "object",
properties: {
command: { type: "string" },
},
required: ["command"],
};
const tools = [makeTool(originalParameters)];
convertTools(tools, true);
expect(originalParameters).toEqual({
$schema: "http://json-schema.org/draft-07/schema#",
type: "object",
properties: {
command: { type: "string" },
},
required: ["command"],
});
});
it("preserves $schema in parametersJsonSchema when useParameters=false", () => {
const tools = [
makeTool({
$schema: "http://json-schema.org/draft-07/schema#",
type: "object",
properties: {
command: { type: "string" },
},
required: ["command"],
}),
];
const result = convertTools(tools, false);
const decl = result?.[0]?.functionDeclarations?.[0];
expect(decl).toBeDefined();
expect(decl?.parametersJsonSchema).toEqual({
$schema: "http://json-schema.org/draft-07/schema#",
type: "object",
properties: {
command: { type: "string" },
},
required: ["command"],
});
});
it("handles tools without $schema gracefully", () => {
const tools = [
makeTool({
type: "object",
properties: {
path: { type: "string" },
},
required: ["path"],
}),
];
const result = convertTools(tools, true);
const decl = result?.[0]?.functionDeclarations?.[0];
expect(decl).toBeDefined();
expect(decl?.parameters).toEqual({
type: "object",
properties: {
path: { type: "string" },
},
required: ["path"],
});
});
it("returns undefined for empty tool list", () => {
expect(convertTools([])).toBeUndefined();
expect(convertTools([], true)).toBeUndefined();
});
});
function makeGoogleModel(
api: "google-generative-ai" | "google-vertex",
provider: string,
id: string,
): Model<"google-generative-ai" | "google-vertex"> {
return {
id,
name: id,
api,
provider,
baseUrl: "https://example.com",
reasoning: true,
input: ["text", "image"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 128000,
maxTokens: 8192,
} as Model<"google-generative-ai" | "google-vertex">;
}
const emptyUsage = {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
describe("google-shared convertMessages — Gemini 3 tool call signatures", () => {
function toolCallContext(source: { api: string; provider: string; id: string }, thoughtSignature?: string): Context {
return {
messages: [
{ role: "user", content: "Hi", timestamp: 1 },
{
role: "assistant",
content: [
{
type: "toolCall",
id: "call_1",
name: "bash",
arguments: { command: "echo hi" },
...(thoughtSignature && { thoughtSignature }),
},
{ type: "toolCall", id: "call_2", name: "bash", arguments: { command: "ls -la" } },
],
api: source.api,
provider: source.provider,
model: source.id,
usage: emptyUsage,
stopReason: "toolUse",
timestamp: 2,
},
],
} as Context;
}
it.each([
{
name: "unsigned Google Gen AI tool calls from another model",
model: makeGoogleModel("google-generative-ai", "google", "gemini-3-pro-preview"),
sourceId: "other-model",
signature: undefined,
expected: [undefined, undefined],
},
{
name: "unsigned Vertex tool calls",
model: makeGoogleModel("google-vertex", "google-vertex", "gemini-3-pro-preview"),
sourceId: undefined,
signature: undefined,
expected: [undefined, undefined],
},
{
name: "a valid signature from the same provider and model",
model: makeGoogleModel("google-generative-ai", "google", "gemini-3-pro-preview"),
sourceId: undefined,
signature: "AAAAAAAAAAAAAAAAAAAAAA==",
expected: ["AAAAAAAAAAAAAAAAAAAAAA==", undefined],
},
{
name: "non-Gemini-3 models",
model: makeGoogleModel("google-generative-ai", "google", "gemini-2.5-flash"),
sourceId: "other-model",
signature: undefined,
expected: [undefined, undefined],
},
])("never fakes a thought signature for $name", ({ model, sourceId, signature, expected }) => {
const source = { api: model.api, provider: model.provider, id: sourceId ?? model.id };
const contents = convertMessages(model, toolCallContext(source, signature));
const modelTurn = contents.find((content) => content.role === "model");
const functionCallParts = modelTurn?.parts?.filter((part) => part.functionCall !== undefined) ?? [];
expect(functionCallParts).toHaveLength(2);
expect(functionCallParts.map((part) => part.thoughtSignature)).toEqual(expected);
expect(JSON.stringify(modelTurn)).not.toContain("skip_thought_signature_validator");
});
});
describe("google-shared image tool result routing", () => {
function imageToolResultContext(model: Model<"google-generative-ai" | "google-vertex">): Context {
const toolResult = (toolCallId: string, content: unknown) => ({
role: "toolResult" as const,
toolCallId,
toolName: "read",
content,
isError: false,
timestamp: 3,
});
return {
messages: [
{ role: "user", content: "read the files", timestamp: 1 },
{
role: "assistant",
content: [
{ type: "toolCall", id: "call_a", name: "read", arguments: { path: "a.txt" } },
{ type: "toolCall", id: "call_img", name: "read", arguments: { path: "image.png" } },
{ type: "toolCall", id: "call_b", name: "read", arguments: { path: "b.txt" } },
],
api: model.api,
provider: model.provider,
model: model.id,
usage: emptyUsage,
stopReason: "toolUse",
timestamp: 2,
},
toolResult("call_a", [{ type: "text", text: "alpha text" }]),
toolResult("call_img", [{ type: "image", data: "abc", mimeType: "image/png" }]),
toolResult("call_b", [{ type: "text", text: "beta text" }]),
],
} as Context;
}
it("keeps a separate synthetic image turn for Gemini 2.x models", () => {
const model = makeGoogleModel("google-generative-ai", "google", "gemini-2.5-flash");
const contents = convertMessages(model, imageToolResultContext(model));
expect(contents).toHaveLength(5);
expect(contents[2].parts?.every((part) => part.functionResponse)).toBe(true);
expect(contents[3].parts?.[1]?.inlineData).toBeTruthy();
expect(contents[4].parts?.[0]?.functionResponse).toBeTruthy();
});
it("nests image tool results for Gemini 3 models", () => {
const model = makeGoogleModel("google-generative-ai", "google", "gemini-3-pro-preview");
const contents = convertMessages(model, imageToolResultContext(model));
expect(contents).toHaveLength(3);
expect(contents[2].parts).toHaveLength(3);
const imageResponse = contents[2].parts?.[1]?.functionResponse;
expect(imageResponse?.parts).toHaveLength(1);
expect(imageResponse?.parts?.[0]?.inlineData).toBeTruthy();
});
});
describe("Google Vertex thinking budget payload", () => {
const stableFlashLite = getModel("google-vertex", "gemini-2.5-flash-lite");
it.each([stableFlashLite, { ...stableFlashLite, id: "gemini-2.5-flash-lite-preview" }])(
"uses the supported minimal budget for $id",
async (model) => {
let capturedPayload: GenerateContentParameters | undefined;
await streamSimpleGoogleVertex(
model,
{ messages: [{ role: "user", content: "hello", timestamp: 1 }] },
{
apiKey: "fake-key",
reasoning: "minimal",
onPayload: (payload) => {
capturedPayload = payload as GenerateContentParameters;
return payload;
},
},
).result();
expect(capturedPayload?.config?.thinkingConfig).toEqual({ includeThoughts: true, thinkingBudget: 512 });
},
);
});
describe("Google thinking detection (thoughtSignature)", () => {
it.each([
{ part: { thought: true, thoughtSignature: undefined }, thinking: true },
{ part: { thought: true, thoughtSignature: "opaque-signature" }, thinking: true },
{ part: { thought: undefined, thoughtSignature: "opaque-signature" }, thinking: false },
{ part: { thought: false, thoughtSignature: "opaque-signature" }, thinking: false },
{ part: { thought: undefined, thoughtSignature: undefined }, thinking: false },
{ part: { thought: false, thoughtSignature: "" }, thinking: false },
])("isThinkingPart($part) === $thinking", ({ part, thinking }) => {
expect(isThinkingPart(part)).toBe(thinking);
});
it.each([
{ previous: undefined, next: "sig-1", expected: "sig-1" },
{ previous: "sig-1", next: undefined, expected: "sig-1" },
{ previous: "sig-1", next: "", expected: "sig-1" },
{ previous: "sig-1", next: "sig-2", expected: "sig-2" },
])("retainThoughtSignature($previous, $next) === $expected", ({ previous, next, expected }) => {
expect(retainThoughtSignature(previous, next)).toBe(expected);
});
});