Fix Gemini thought signature validation and token counting errors (#9380)

Co-authored-by: Matt Rubens <[email protected]>

Author: hannesrudolph

src/api/providers/__tests__/gemini.spec.ts

@@ -206,23 +206,6 @@ describe("GeminiHandler", () => {
 			expect(handler.calculateCost({ info: mockInfo, inputTokens: 0, outputTokens })).toBeCloseTo(expectedCost)
 		})
 
-		it("should calculate cost with cache write tokens", () => {
-			const inputTokens = 10000
-			const outputTokens = 20000
-			const cacheWriteTokens = 5000
-			const CACHE_TTL = 5 // Match the constant in gemini.ts
-
-			// Added non-null assertions (!)
-			const expectedInputCost = (inputTokens / 1_000_000) * mockInfo.inputPrice!
-			const expectedOutputCost = (outputTokens / 1_000_000) * mockInfo.outputPrice!
-			const expectedCacheWriteCost =
-				mockInfo.cacheWritesPrice! * (cacheWriteTokens / 1_000_000) * (CACHE_TTL / 60)
-			const expectedCost = expectedInputCost + expectedOutputCost + expectedCacheWriteCost
-
-			const cost = handler.calculateCost({ info: mockInfo, inputTokens, outputTokens })
-			expect(cost).toBeCloseTo(expectedCost)
-		})
-
 		it("should calculate cost with cache read tokens", () => {
 			const inputTokens = 10000 // Total logical input
 			const outputTokens = 20000

src/api/providers/gemini.ts

@@ -82,9 +82,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 			: (maxTokens ?? undefined)
 
 		// Only forward encrypted reasoning continuations (thoughtSignature) when we are
-		// using effort-based reasoning (thinkingLevel). Budget-only configs should NOT
-		// send thoughtSignature parts back to Gemini.
-		const includeThoughtSignatures = Boolean(thinkingConfig?.thinkingLevel)
+		// using reasoning (thinkingConfig is present). Both effort-based (thinkingLevel)
+		// and budget-based (thinkingBudget) models require this for active loops.
+		const includeThoughtSignatures = Boolean(thinkingConfig)
 
 		// The message list can include provider-specific meta entries such as
 		// `{ type: "reasoning", ... }` that are intended only for providers like
@@ -162,10 +162,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 						}>) {
 							// Capture thought signatures so they can be persisted into API history.
 							const thoughtSignature = part.thoughtSignature
-							// Only persist encrypted reasoning when an effort-based thinking level is set
-							// (i.e. thinkingConfig.thinkingLevel is present). Budget-based configs that only
-							// set thinkingBudget should NOT trigger encrypted continuation.
-							if (thinkingConfig?.thinkingLevel && thoughtSignature) {
+							// Persist encrypted reasoning when using reasoning. Both effort-based
+							// and budget-based models require this for active loops.
+							if (thinkingConfig && thoughtSignature) {
 								this.lastThoughtSignature = thoughtSignature
 							}
 
@@ -351,7 +350,12 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 			const countTokensRequest = {
 				model,
 				// Token counting does not need encrypted continuation; always drop thoughtSignature.
-				contents: convertAnthropicContentToGemini(content, { includeThoughtSignatures: false }),
+				contents: [
+					{
+						role: "user",
+						parts: convertAnthropicContentToGemini(content, { includeThoughtSignatures: false }),
+					},
+				],
 			}
 
 			const response = await this.client.models.countTokens(countTokensRequest)

src/api/transform/__tests__/gemini-format.spec.ts

@@ -124,6 +124,7 @@ describe("convertAnthropicMessageToGemini", () => {
 						name: "calculator",
 						args: { operation: "add", numbers: [2, 3] },
 					},
+					thoughtSignature: "skip_thought_signature_validator",
 				},
 			],
 		})

src/api/transform/gemini-format.ts

@@ -19,15 +19,30 @@ export function convertAnthropicContentToGemini(
 ): Part[] {
 	const includeThoughtSignatures = options?.includeThoughtSignatures ?? true
 
+	// First pass: find thoughtSignature if it exists in the content blocks
+	let activeThoughtSignature: string | undefined
+	if (Array.isArray(content)) {
+		const sigBlock = content.find((block) => isThoughtSignatureContentBlock(block)) as ThoughtSignatureContentBlock
+		if (sigBlock?.thoughtSignature) {
+			activeThoughtSignature = sigBlock.thoughtSignature
+		}
+	}
+
+	// Determine the signature to attach to function calls.
+	// If we're in a mode that expects signatures (includeThoughtSignatures is true):
+	// 1. Use the actual signature if we found one in the history/content.
+	// 2. Fallback to "skip_thought_signature_validator" if missing (e.g. cross-model history).
+	let functionCallSignature: string | undefined
+	if (includeThoughtSignatures) {
+		functionCallSignature = activeThoughtSignature || "skip_thought_signature_validator"
+	}
+
 	if (typeof content === "string") {
 		return [{ text: content }]
 	}
 
 	return content.flatMap((block): Part | Part[] => {
-		// Handle thoughtSignature blocks first so that the main switch can continue
-		// to operate on the standard Anthropic content union. This preserves strong
-		// typing for known block types while still allowing provider-specific
-		// extensions when needed.
+		// Handle thoughtSignature blocks first
 		if (isThoughtSignatureContentBlock(block)) {
 			if (includeThoughtSignatures && typeof block.thoughtSignature === "string") {
 				// The Google GenAI SDK currently exposes thoughtSignature as an
@@ -54,7 +69,10 @@ export function convertAnthropicContentToGemini(
 						name: block.name,
 						args: block.input as Record<string, unknown>,
 					},
-				}
+					// Inject the thoughtSignature into the functionCall part if required.
+					// This is necessary for Gemini 2.5/3+ thinking models to validate the tool call.
+					...(functionCallSignature ? { thoughtSignature: functionCallSignature } : {}),
+				} as Part
 			case "tool_result": {
 				if (!block.content) {
 					return []
@@ -108,6 +126,10 @@ export function convertAnthropicMessageToGemini(
 ): Content {
 	return {
 		role: message.role === "assistant" ? "model" : "user",
-		parts: convertAnthropicContentToGemini(message.content, options),
+		parts: convertAnthropicContentToGemini(message.content, {
+			...options,
+			includeThoughtSignatures:
+				message.role === "assistant" ? (options?.includeThoughtSignatures ?? true) : false,
+		}),
 	}
 }