oh-my-pi/packages/ai/test/google-gemini-cli-3x-thinking.test.ts

import { afterEach, describe, expect, it, vi } from "bun:test";
import { Effort } from "@oh-my-pi/pi-ai";
import { enrichModelThinking } from "@oh-my-pi/pi-ai/model-thinking";
import { hookFetch } from "@oh-my-pi/pi-utils";
import { getBundledModel } from "../src/models";
import { streamSimple } from "../src/stream";
import type { Context, Model } from "../src/types";

interface GeminiCliThinkingConfig {
	thinkingLevel?: string;
	thinkingBudget?: number;
}

interface CapturedRequestBody {
	request?: {
		generationConfig?: {
			thinkingConfig?: GeminiCliThinkingConfig;
		};
	};
}

function createModel(id: string): Model<"google-gemini-cli"> {
	return enrichModelThinking({
		id,
		name: id,
		api: "google-gemini-cli",
		provider: "google-gemini-cli",
		baseUrl: "https://cloudcode-pa.googleapis.com",
		reasoning: true,
		input: ["text", "image"],
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
		contextWindow: 1_000_000,
		maxTokens: 65_536,
	});
}

const context: Context = {
	messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
};

function extractThinking(bodyText: string | undefined): GeminiCliThinkingConfig | undefined {
	if (!bodyText) return undefined;
	const parsed = JSON.parse(bodyText) as CapturedRequestBody;
	return parsed.request?.generationConfig?.thinkingConfig;
}

describe("google-gemini-cli Gemini 3.x thinking mapping", () => {
	afterEach(() => {
		vi.restoreAllMocks();
	});

	it("includes gemini-3.1-pro-preview in bundled google-gemini-cli models", () => {
		expect(getBundledModel("google-gemini-cli", "gemini-3.1-pro-preview")?.id).toBe("gemini-3.1-pro-preview");
	});
	it("uses thinkingLevel for gemini-3.1-pro-preview when the effort is supported", async () => {
		let requestBody: string | undefined;
		using _hook = hookFetch((_input, init) => {
			requestBody = typeof init?.body === "string" ? init.body : undefined;
			return new Response('{"error":{"message":"bad request"}}', { status: 400 });
		});

		const stream = streamSimple(createModel("gemini-3.1-pro-preview"), context, {
			apiKey: JSON.stringify({ token: "token", projectId: "proj-123" }),
			reasoning: Effort.High,
		});
		await stream.result();

		const thinking = extractThinking(requestBody);
		expect(thinking?.thinkingLevel).toBe("HIGH");
		expect(thinking?.thinkingBudget).toBeUndefined();
	});

	it("rejects unsupported gemini-3.1-pro-preview efforts instead of promoting them", () => {
		let requestBody: string | undefined;
		using _hook = hookFetch((_input, init) => {
			requestBody = typeof init?.body === "string" ? init.body : undefined;
			return new Response('{"error":{"message":"bad request"}}', { status: 400 });
		});

		expect(() =>
			streamSimple(createModel("gemini-3.1-pro-preview"), context, {
				apiKey: JSON.stringify({ token: "token", projectId: "proj-123" }),
				reasoning: Effort.Medium,
			}),
		).toThrow(/Supported efforts: low, high/);
		expect(requestBody).toBeUndefined();
	});

	it("uses thinkingLevel for gemini-3.1-flash-preview", async () => {
		let requestBody: string | undefined;
		using _hook = hookFetch((_input, init) => {
			requestBody = typeof init?.body === "string" ? init.body : undefined;
			return new Response('{"error":{"message":"bad request"}}', { status: 400 });
		});

		const stream = streamSimple(createModel("gemini-3.1-flash-preview"), context, {
			apiKey: JSON.stringify({ token: "token", projectId: "proj-123" }),
			reasoning: Effort.Medium,
		});
		await stream.result();

		const thinking = extractThinking(requestBody);
		expect(thinking?.thinkingLevel).toBe("MEDIUM");
		expect(thinking?.thinkingBudget).toBeUndefined();
	});

	it("keeps thinkingBudget for gemini-2.5-pro", async () => {
		let requestBody: string | undefined;
		using _hook = hookFetch((_input, init) => {
			requestBody = typeof init?.body === "string" ? init.body : undefined;
			return new Response('{"error":{"message":"bad request"}}', { status: 400 });
		});

		const stream = streamSimple(createModel("gemini-2.5-pro"), context, {
			apiKey: JSON.stringify({ token: "token", projectId: "proj-123" }),
			reasoning: Effort.Medium,
		});
		await stream.result();

		const thinking = extractThinking(requestBody);
		expect(thinking?.thinkingLevel).toBeUndefined();
		expect(thinking?.thinkingBudget).toBeDefined();
	});
});