Adding Thinking UX for Gemini (RooCodeInc#4137)

Author: arafatkatze

.changeset/wet-rice-hunt.md

@@ -0,0 +1,5 @@
+---
+"claude-dev": patch
+---
+
+Adding Thinking UX for Gemini

package-lock.json

@@ -404,7 +404,7 @@
 		"@bufbuild/protobuf": "^2.2.5",
 		"@cerebras/cerebras_cloud_sdk": "^1.35.0",
 		"@google-cloud/vertexai": "^1.9.3",
-		"@google/genai": "^0.13.0",
+		"@google/genai": "1.0.0",
 		"@grpc/grpc-js": "^1.9.15",
 		"@grpc/reflection": "^1.0.4",
 		"@mistralai/mistralai": "^1.5.0",

package.json

@@ -2,6 +2,7 @@ import type { Anthropic } from "@anthropic-ai/sdk"
 // Restore GenerateContentConfig import and add GenerateContentResponseUsageMetadata
 import { GoogleGenAI, type GenerateContentConfig, type GenerateContentResponseUsageMetadata } from "@google/genai"
 import { withRetry } from "../retry"
+import { Part } from "@google/genai"
 import { ApiHandler } from "../"
 import { ApiHandlerOptions, geminiDefaultModelId, GeminiModelId, geminiModels, ModelInfo } from "@shared/api"
 import { convertAnthropicMessageToGemini } from "../transform/gemini-format"
@@ -96,9 +97,10 @@ export class GeminiHandler implements ApiHandler {
 		}
 
 		// Add thinking config if the model supports it
-		if (info.thinkingConfig?.outputPrice !== undefined && maxBudget > 0) {
+		if (thinkingBudget > 0) {
 			requestConfig.thinkingConfig = {
 				thinkingBudget: thinkingBudget,
+				includeThoughts: true,
 			}
 		}
 
@@ -111,6 +113,7 @@ export class GeminiHandler implements ApiHandler {
 		let promptTokens = 0
 		let outputTokens = 0
 		let cacheReadTokens = 0
+		let thoughtsTokenCount = 0 // Initialize thought token counts
 		let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined
 
 		try {
@@ -130,6 +133,31 @@ export class GeminiHandler implements ApiHandler {
 					isFirstSdkChunk = false
 				}
 
+				// Handle thinking content from Gemini's response
+				const candidateForThoughts = chunk?.candidates?.[0]
+				const partsForThoughts = candidateForThoughts?.content?.parts
+				let thoughts = "" // Initialize as empty string
+
+				if (partsForThoughts) {
+					// This ensures partsForThoughts is a Part[] array
+					for (const part of partsForThoughts) {
+						const { thought, text } = part as Part
+						if (thought && text) {
+							// Ensure part.text exists
+							// Handle the thought part
+							thoughts += text + "\n" // Append thought and a newline
+						}
+					}
+				}
+
+				if (thoughts.trim() !== "") {
+					yield {
+						type: "reasoning",
+						reasoning: thoughts.trim(),
+					}
+					thoughts = "" // Reset thoughts after yielding
+				}
+
 				if (chunk.text) {
 					yield {
 						type: "text",
@@ -141,6 +169,7 @@ export class GeminiHandler implements ApiHandler {
 					lastUsageMetadata = chunk.usageMetadata
 					promptTokens = lastUsageMetadata.promptTokenCount ?? promptTokens
 					outputTokens = lastUsageMetadata.candidatesTokenCount ?? outputTokens
+					thoughtsTokenCount = lastUsageMetadata.thoughtsTokenCount ?? thoughtsTokenCount
 					cacheReadTokens = lastUsageMetadata.cachedContentTokenCount ?? cacheReadTokens
 				}
 			}
@@ -151,12 +180,14 @@ export class GeminiHandler implements ApiHandler {
 					info,
 					inputTokens: promptTokens,
 					outputTokens,
+					thoughtsTokenCount,
 					cacheReadTokens,
 				})
 				yield {
 					type: "usage",
 					inputTokens: promptTokens,
 					outputTokens,
+					thoughtsTokenCount,
 					cacheReadTokens,
 					cacheWriteTokens: 0,
 					totalCost,
@@ -239,11 +270,13 @@ export class GeminiHandler implements ApiHandler {
 		info,
 		inputTokens,
 		outputTokens,
+		thoughtsTokenCount = 0,
 		cacheReadTokens = 0,
 	}: {
 		info: ModelInfo
 		inputTokens: number
 		outputTokens: number
+		thoughtsTokenCount: number
 		cacheReadTokens?: number
 	}) {
 		// Exit early if any required pricing information is missing
@@ -275,18 +308,18 @@ export class GeminiHandler implements ApiHandler {
 		const inputTokensCost = inputPrice * (uncachedInputTokens / 1_000_000)
 
 		// 2. Output token costs
-		const outputTokensCost = outputPrice * (outputTokens / 1_000_000)
+		const responseTokensCost = outputPrice * ((outputTokens + thoughtsTokenCount) / 1_000_000)
 
 		// 3. Cache read costs (immediate)
 		const cacheReadCost = (cacheReadTokens ?? 0) > 0 ? cacheReadsPrice * ((cacheReadTokens ?? 0) / 1_000_000) : 0
 
 		// Calculate total immediate cost (excluding cache write/storage costs)
-		const totalCost = inputTokensCost + outputTokensCost + cacheReadCost
+		const totalCost = inputTokensCost + responseTokensCost + cacheReadCost
 
 		// Create the trace object for debugging
 		const trace: Record<string, { price: number; tokens: number; cost: number }> = {
 			input: { price: inputPrice, tokens: uncachedInputTokens, cost: inputTokensCost },
-			output: { price: outputPrice, tokens: outputTokens, cost: outputTokensCost },
+			output: { price: outputPrice, tokens: outputTokens, cost: responseTokensCost },
 		}
 
 		// Only include cache read costs in the trace (cache write costs are tracked separately)

src/api/providers/gemini.ts

@@ -17,5 +17,6 @@ export interface ApiStreamUsageChunk {
 	outputTokens: number
 	cacheWriteTokens?: number
 	cacheReadTokens?: number
+	thoughtsTokenCount?: number // openrouter
 	totalCost?: number // openrouter
 }

src/api/transform/stream.ts

@@ -598,6 +598,9 @@ export const vertexModels = {
 				cacheReadsPrice: 0.625,
 			},
 		],
+		thinkingConfig: {
+			maxBudget: 32768,
+		},
 	},
 	"gemini-2.5-flash-preview-04-17": {
 		maxTokens: 65536,
@@ -766,6 +769,9 @@ export const geminiModels = {
 				cacheReadsPrice: 0.625,
 			},
 		],
+		thinkingConfig: {
+			maxBudget: 32768,
+		},
 	},
 	"gemini-2.5-flash-preview-05-20": {
 		maxTokens: 65536,

src/shared/api.ts

@@ -119,7 +119,14 @@ const OpenRouterBalanceDisplay = ({ apiKey }: { apiKey: string }) => {
 
 const SUPPORTED_THINKING_MODELS: Record<string, string[]> = {
 	anthropic: ["claude-3-7-sonnet-20250219", "claude-sonnet-4-20250514", "claude-opus-4-20250514"],
-	vertex: ["claude-3-7-sonnet@20250219", "claude-sonnet-4@20250514", "claude-opus-4@20250514"],
+	vertex: [
+		"claude-3-7-sonnet@20250219",
+		"claude-sonnet-4@20250514",
+		"claude-opus-4@20250514",
+		"gemini-2.5-flash-preview-05-20",
+		"gemini-2.5-flash-preview-04-17",
+		"gemini-2.5-pro-preview-06-05",
+	],
 	qwen: [
 		"qwen3-235b-a22b",
 		"qwen3-32b",
@@ -132,6 +139,7 @@ const SUPPORTED_THINKING_MODELS: Record<string, string[]> = {
 		"qwen-plus-latest",
 		"qwen-turbo-latest",
 	],
+	gemini: ["gemini-2.5-flash-preview-05-20", "gemini-2.5-flash-preview-04-17", "gemini-2.5-pro-preview-06-05"],
 }
 
 // This is necessary to ensure dropdown opens downward, important for when this is used in popup
@@ -1057,15 +1065,6 @@ const ApiOptions = ({
 							</VSCodeLink>
 						)}
 					</p>
-
-					{/* Add Thinking Budget Slider specifically for gemini-2.5-flash-preview-04-17 */}
-					{selectedProvider === "gemini" && selectedModelId === "gemini-2.5-flash-preview-04-17" && (
-						<ThinkingBudgetSlider
-							apiConfiguration={apiConfiguration}
-							setApiConfiguration={setApiConfiguration}
-							maxBudget={selectedModelInfo.thinkingConfig?.maxBudget}
-						/>
-					)}
 				</div>
 			)}