fix(aikit): google family fixes

sanchitrk · Sep 5, 2026 · Sep 4, 2026 · Sep 5, 2026 · Sep 5, 2026 · Sep 5, 2026
commit df50a15230482ba22b9e0c814e849bfc1d6b22da
diff --git a/packages/aikit/src/llm/thinking.ts b/packages/aikit/src/llm/thinking.ts
@@ -25,6 +25,7 @@ import {
  */
 export interface Plan {
 	readonly level: Model.ThinkingLevel;
+	/** Fitted tokens, or Google's -1 sentinel for provider-managed dynamic thinking. */
 	readonly budget: number;
 	readonly maxTokens: number | undefined;
 	/** The caller configured a budget for this level, rather than inheriting a default. */
@@ -56,21 +57,28 @@ export function resolvePlan(model: Model.Info, context: Message.Context, options
 
 	const budgets = options.thinkingBudgets;
 	const explicit = budgets?.[level] !== undefined;
+	const key = Model.optionsKey(model);
+	const usesGoogleBudget = (key === "google" || key === "google-vertex") && !usesGoogleThinkingLevel(model);
+	const requestedBudget = usesGoogleBudget
+		? (budgets?.[level] ?? googleThinkingBudget(model, googleThinkingLevel(model.thinkingLevelMap?.[level] ?? level)))
+		: thinkingBudgetForLevel(level, budgets);
 
 	// A catalog entry without a ceiling gives the fit nothing to work against.
 	if (model.maxTokens <= 0) {
-		const budget = thinkingBudgetForLevel(level, budgets);
-		return { level, budget, maxTokens: clampToContext(options.maxTokens), explicit };
+		return { level, budget: requestedBudget, maxTokens: clampToContext(options.maxTokens), explicit };
 	}
 
-	const adjusted = adjustMaxTokensForThinking(options.maxTokens, model.maxTokens, level, budgets);
+	// Google's -1 is a dynamic-thinking sentinel, not a negative token count.
+	const dynamic = usesGoogleBudget && requestedBudget === -1;
+	const adjusted = adjustMaxTokensForThinking(options.maxTokens, model.maxTokens, level, {
+		[level]: dynamic ? 0 : requestedBudget,
+	});
 	const maxTokens = clampMaxTokensToContext(model, context, adjusted.maxTokens);
 
 	/*
-	 * Re-fit only when the context clamp actually took room away. The budget was
-	 * already fitted under `adjusted.maxTokens` with the answer floor reserved, so
-	 * clamping again against an unchanged ceiling would reserve that floor twice
-	 * and can push the budget under a provider's own minimum.
+	 * Preserve an explicit small answer allowance when it fits, and apply the
+	 * 1024-token reserve when context pressure reduces it. Pi's Anthropic wrapper
+	 * applies this reserve unconditionally; aikit deliberately permits small answers.
 	 */
 	const budget =
 		maxTokens < adjusted.maxTokens
@@ -86,7 +94,7 @@ export function resolvePlan(model: Model.Info, context: Message.Context, options
 		return { level: "off", budget: 0, maxTokens, explicit };
 	}
 
-	return { level, budget, maxTokens, explicit };
+	return { level, budget: dynamic ? -1 : budget, maxTokens, explicit };
 }
 
 type GoogleThinkingLevel = "minimal" | "low" | "medium" | "high";
@@ -113,8 +121,7 @@ const GOOGLE_THINKING_BUDGETS: ReadonlyArray<[RegExp, Record<GoogleThinkingLevel
 	[/2\.5-flash/, { minimal: 128, low: 2048, medium: 8192, high: 24576 }],
 ];
 
-function googleThinkingBudget(model: Model.Info, level: GoogleThinkingLevel, plan: Plan): number {
-	if (plan.explicit) return plan.budget;
+function googleThinkingBudget(model: Model.Info, level: GoogleThinkingLevel): number {
 	for (const [pattern, budgets] of GOOGLE_THINKING_BUDGETS) {
 		if (pattern.test(model.id)) return budgets[level];
 	}
@@ -231,9 +238,8 @@ export function reasoningProviderOptions(model: Model.Info, plan: Plan): Provide
 		return {
 			[key]: {
 				thinkingConfig: {
-					thinkingLevel,
 					includeThoughts: true,
-					thinkingBudget: googleThinkingBudget(model, thinkingLevel, plan),
+					thinkingBudget: budget,
 				},
 			},
 		};

diff --git a/packages/aikit/src/utils/estimate.ts b/packages/aikit/src/utils/estimate.ts
@@ -57,9 +57,11 @@ export function estimateContentTokens(content: ReadonlyArray<Message.TextContent
  * results in their own messages -- the result it already holds.
  */
 function toolCallChars(part: Message.ToolCall): number {
-	let chars = part.name.length + safeJsonStringify(part.arguments).length;
-	if ("result" in part && part.result) chars += contentChars(part.result.content);
-	return chars;
+	return part.name.length + safeJsonStringify(part.arguments).length + toolResultChars(part);
+}
+
+function toolResultChars(part: Message.ToolCall): number {
+	return "result" in part && part.result ? contentChars(part.result.content) : 0;
 }
 
 export function estimateMessageTokens(message: Message.Message): number {
@@ -118,7 +120,14 @@ function estimateMessages(messages: ReadonlyArray<Message.Message>): ContextUsag
 	const usage = lastAssistantUsage(messages);
 	if (usage) {
 		const usageTokens = calculateContextTokens(usage.usage);
-		let trailingTokens = 0;
+		// Unlike Pi's separate tool-result messages, our results are attached to
+		// the assistant turn after its usage was reported. Only its output is counted.
+		const latest = messages[usage.index]!;
+		let resultChars = 0;
+		for (const part of latest.parts) {
+			if (part.type === "toolCall") resultChars += toolResultChars(part);
+		}
+		let trailingTokens = Math.ceil(resultChars / CHARS_PER_TOKEN);
 		for (let i = usage.index + 1; i < messages.length; i++) trailingTokens += estimateMessageTokens(messages[i]!);
 		return { tokens: usageTokens + trailingTokens, usageTokens, trailingTokens, lastUsageIndex: usage.index };
 	}
@@ -133,10 +142,10 @@ function estimateToolsTokens(tools: ReadonlyArray<Message.Tool> | undefined): nu
 	return estimateTextTokens(safeJsonStringify(tools));
 }
 
-/** Tool names that only became available after the given message index. */
+/** Tools discovered by execution of the given assistant turn or later turns. */
 function toolNamesAddedAfter(messages: ReadonlyArray<Message.Message>, index: number): Set<string> {
 	const names = new Set<string>();
-	for (let i = index + 1; i < messages.length; i++) {
+	for (let i = index; i < messages.length; i++) {
 		const message = messages[i]!;
 		if (message.role !== "assistant") continue;
 		for (const part of message.parts) {

diff --git a/packages/aikit/test/estimate.test.ts b/packages/aikit/test/estimate.test.ts
@@ -2,7 +2,7 @@ import { describe, expect, it } from "vite-plus/test";
 import { clampMaxTokensToContext } from "../src/llm/shared.ts";
 import * as Message from "../src/message/message.ts";
 import { estimateContextTokens, estimateMessageTokens, estimateTextTokens } from "../src/utils/estimate.ts";
-import { makeModel, makeUsage } from "./utils/fixtures.ts";
+import { makeCompletedToolCall, makeModel, makeUsage } from "./utils/fixtures.ts";
 
 const model = makeModel({ contextWindow: 10_000, maxTokens: 8_000 });
 
@@ -70,6 +70,22 @@ describe("estimateMessageTokens", () => {
 });
 
 describe("estimateContextTokens", () => {
+	it("adds results and discovered tools on the usage-bearing turn only once", () => {
+		const latest = assistant(100, 1_000);
+		latest.parts = [makeCompletedToolCall("read", "read", [{ type: "text", text: "r".repeat(40_000) }])];
+		const call = latest.parts[0]!;
+		if (call.type === "toolCall") call.addedToolNames = ["new_tool"];
+		const tool: Message.Tool = { name: "new_tool", description: "new", parameters: { type: "object" } };
+		const context: Message.Context = { messages: [latest], tools: [tool] };
+		const added = 10_000 + estimateTextTokens(JSON.stringify([tool]));
+		expect(estimateContextTokens(context)).toMatchObject({ tokens: 1_000 + added, trailingTokens: added });
+		expect(clampMaxTokensToContext(makeModel({ contextWindow: 20_000 }), context, 16_000)).toBe(
+			20_000 - 1_000 - added - 4_096,
+		);
+		context.messages.push(assistant(200, 12_000));
+		expect(estimateContextTokens(context)).toMatchObject({ tokens: 12_000, trailingTokens: 0 });
+	});
+
 	it("ignores stale assistant usage after a newer message is inserted before it", () => {
 		const context: Message.Context = {
 			systemPrompt: "system",

diff --git a/packages/aikit/test/thinking.provideroptions.test.ts b/packages/aikit/test/thinking.provideroptions.test.ts
@@ -102,7 +102,7 @@ describe("Thinking.reasoningProviderOptions", () => {
 		for (const [id, expected] of cases) {
 			const model = google(id);
 			const options = Thinking.reasoningProviderOptions(model, planFor(model, { reasoning: "high" }));
-			expect(options.google?.thinkingConfig).toMatchObject({ thinkingLevel: "high", thinkingBudget: expected });
+			expect(options.google?.thinkingConfig).toEqual({ includeThoughts: true, thinkingBudget: expected });
 		}
 	});
 

diff --git a/packages/aikit/test/thinking.wire.test.ts b/packages/aikit/test/thinking.wire.test.ts
@@ -0,0 +1,92 @@
+import { describe, expect, it } from "vite-plus/test";
+import { stream } from "../src/llm/stream.ts";
+import type { RuntimeOptions } from "../src/llm/options.ts";
+import { makeModel, makeUserMessage } from "./utils/fixtures.ts";
+
+async function capture(model: ReturnType<typeof makeModel>, options: RuntimeOptions) {
+	let body: unknown;
+	const fetch: typeof globalThis.fetch = async (_input, init) => {
+		if (typeof init?.body !== "string") throw new Error("Expected a JSON request body");
+		body = JSON.parse(init.body);
+		return new Response(JSON.stringify({ error: { message: "Captured request", type: "invalid_request_error" } }), {
+			status: 400,
+			headers: { "content-type": "application/json" },
+		});
+	};
+	await stream(
+		model,
+		{ messages: [makeUserMessage("hello")] },
+		{
+			...options,
+			apiKey: "test-key",
+			factoryOptions: { fetch },
+		},
+	).result();
+	expect(body).toBeDefined();
+	return body;
+}
+
+const google = (overrides: Parameters<typeof makeModel>[0] = {}) =>
+	makeModel({
+		id: "gemini-2.5-pro",
+		protocol: "google",
+		npm: "@ai-sdk/google",
+		reasoning: true,
+		maxTokens: 64_000,
+		...overrides,
+	});
+
+describe("thinking wire budgets", () => {
+	it("adds Google's family budget to the requested answer allowance", async () => {
+		expect(await capture(google(), { reasoning: "high", maxTokens: 4_096 })).toMatchObject({
+			generationConfig: {
+				maxOutputTokens: 36_864,
+				thinkingConfig: { thinkingBudget: 32_768, includeThoughts: true },
+			},
+		});
+	});
+
+	it.each([undefined, { high: 32_768 }])(
+		"keeps Google's fitted budget under context pressure (%j)",
+		async (thinkingBudgets) => {
+			// hello = 2 estimated tokens; 10,000 - 2 - 4,096 = 5,902.
+			expect(
+				await capture(google({ contextWindow: 10_000 }), {
+					reasoning: "high",
+					maxTokens: 4_096,
+					...(thinkingBudgets ? { thinkingBudgets } : {}),
+				}),
+			).toMatchObject({
+				generationConfig: {
+					maxOutputTokens: 5_902,
+					thinkingConfig: { thinkingBudget: 4_878 },
+				},
+			});
+		},
+	);
+
+	it("fits Google's budget under the model ceiling", async () => {
+		expect(await capture(google({ maxTokens: 8_192 }), { reasoning: "high" })).toMatchObject({
+			generationConfig: { maxOutputTokens: 8_192, thinkingConfig: { thinkingBudget: 7_168 } },
+		});
+	});
+
+	it("preserves Google's dynamic sentinel without subtracting it from the ceiling", async () => {
+		expect(await capture(google({ id: "gemini-unknown" }), { reasoning: "high", maxTokens: 4_096 })).toMatchObject({
+			generationConfig: { maxOutputTokens: 4_096, thinkingConfig: { thinkingBudget: -1 } },
+		});
+	});
+
+	it("lets the Anthropic adapter add the budget exactly once", async () => {
+		const model = makeModel({
+			id: "claude-sonnet-4-5",
+			npm: "@ai-sdk/anthropic",
+			reasoning: true,
+			maxTokens: 64_000,
+		});
+		expect(await capture(model, { reasoning: "high", maxTokens: 4_096 })).toMatchObject({
+			max_tokens: 20_480,
+			thinking: { type: "enabled", budget_tokens: 16_384 },
+		});
+	});
+});