Update src/api/providers/vertex.ts

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

.changeset/lemon-pears-talk.md

@@ -0,0 +1,5 @@
+---
+"claude-dev": minor
+---
+
+enabled pricing calculation for gemini and vertex + more robust caching & cache tracking for gemini & vertex

package.json

@@ -389,6 +389,7 @@
 		"jschardet": "^3.1.4",
 		"mammoth": "^1.8.0",
 		"monaco-vscode-textmate-theme-converter": "^0.1.7",
+		"node-cache": "^5.1.2",
 		"ollama": "^0.5.13",
 		"open-graph-scraper": "^6.9.0",
 		"openai": "^4.83.0",

src/api/providers/gemini.ts

@@ -1,6 +1,7 @@
 import type { Anthropic } from "@anthropic-ai/sdk"
 // Restore GenerateContentConfig import and add GenerateContentResponseUsageMetadata
 import { GoogleGenAI, type Content, type GenerateContentConfig, type GenerateContentResponseUsageMetadata } from "@google/genai"
+import NodeCache from "node-cache"
 import { withRetry } from "../retry"
 import { ApiHandler } from "../"
 import { ApiHandlerOptions, geminiDefaultModelId, GeminiModelId, geminiModels, ModelInfo } from "@shared/api"
@@ -10,100 +11,183 @@ import { ApiStream } from "../transform/stream"
 // Define a default TTL for the cache (e.g., 15 minutes in seconds)
 const DEFAULT_CACHE_TTL_SECONDS = 900
 
+interface GeminiHandlerOptions extends ApiHandlerOptions {
+	isVertex?: boolean
+}
+
+/**
+ * Handler for Google's Gemini API with optimized caching strategy and accurate cost accounting.
+ *
+ * Key features:
+ * - One cache per task: Creates a single cache per task and reuses it for subsequent turns
+ * - Stable cache keys: Uses taskId as a stable identifier for caches
+ * - Efficient cache updates: Only updates caches when there's new content to add
+ * - Split cost accounting: Separates immediate costs from ongoing cache storage costs
+ *
+ * Cost accounting approach:
+ * - Immediate costs (per message): Input tokens, output tokens, and cache read costs
+ * - Ongoing costs (per task): Cache storage costs for the TTL period
+ *
+ * Gemini's caching system is unique in that it charges for holding tokens in cache by the hour.
+ * This implementation optimizes for both performance and cost by:
+ * 1. Minimizing redundant cache creations
+ * 2. Properly accounting for cache costs in the billing calculations
+ * 3. Using a stable cache key to ensure cache reuse across turns
+ * 4. Separating immediate costs from ongoing costs to avoid double-counting
+ */
 export class GeminiHandler implements ApiHandler {
 	private options: ApiHandlerOptions
-	private client: GoogleGenAI // Updated client type
+	private client: GoogleGenAI
 
-	// Internal state for caching
-	private cacheName: string | null = null
-	private cacheExpireTime: number | null = null
-	private isFirstApiCall = true
+	// Enhanced caching system
+	private contentCaches: NodeCache // Stores cache details (key, count, etc.)
+	private isCacheBusy = false
+	private taskCacheNames: Map<string, string> = new Map() // Maps taskId to cache name for stable lookup
+	private taskCacheTokens: Map<string, number> = new Map() // Maps taskId to total tokens in cache
 
-	constructor(options: ApiHandlerOptions) {
-		if (!options.geminiApiKey) {
-			throw new Error("API key is required for Google Gemini")
-		}
+	constructor(options: GeminiHandlerOptions) {
+		// Store the options
 		this.options = options
-		// Updated client initialization
-		this.client = new GoogleGenAI({ apiKey: options.geminiApiKey })
+
+		if (options.isVertex) {
+			// Initialize with Vertex AI configuration
+			const project = this.options.vertexProjectId ?? "not-provided"
+			const location = this.options.vertexRegion ?? "not-provided"
+
+			this.client = new GoogleGenAI({
+				vertexai: true,
+				project,
+				location,
+			})
+		} else {
+			// Initialize with standard API key
+			if (!options.geminiApiKey) {
+				throw new Error("API key is required for Google Gemini when not using Vertex AI")
+			}
+
+			this.client = new GoogleGenAI({ apiKey: options.geminiApiKey })
+		}
+
+		// Initialize cache with TTL and check period
+		this.contentCaches = new NodeCache({
+			stdTTL: DEFAULT_CACHE_TTL_SECONDS,
+			checkperiod: DEFAULT_CACHE_TTL_SECONDS,
+		})
 	}
 
+	/**
+	 * Creates a message using the Gemini API with optimized caching and split cost accounting.
+	 *
+	 * This method implements a task-based caching strategy:
+	 * 1. Each task gets its own cache, identified by taskId
+	 * 2. On first call for a task, a new cache is created
+	 * 3. On subsequent calls, the existing cache is reused and only new messages are sent
+	 * 4. Cache operations are tracked for accurate cost accounting
+	 *
+	 * Cost accounting:
+	 * - Immediate costs (returned in the usage object): Input tokens, output tokens, cache read costs
+	 * - Ongoing costs (tracked at task level): Cache storage costs for the TTL period
+	 *
+	 * @param systemPrompt The system prompt to use for the message
+	 * @param messages The conversation history to include in the message
+	 * @returns An async generator that yields chunks of the response with accurate immediate costs
+	 */
 	@withRetry()
 	async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
-		const { id: modelId, info: modelInfo } = this.getModel()
+		const { id: model, info } = this.getModel()
+		const contents = messages.map(convertAnthropicMessageToGemini)
 
-		// --- Cache Handling Logic ---
-		const isCacheValid = this.cacheName && this.cacheExpireTime && Date.now() < this.cacheExpireTime
-		let useCache = !this.isFirstApiCall && isCacheValid
+		// Ensure we have a stable cache key (taskId)
+		if (!this.options.taskId) {
+			console.warn("[GeminiHandler] No taskId provided, caching will be disabled")
+		}
 
-		if (this.isFirstApiCall && !isCacheValid && systemPrompt) {
-			// It's the first call, no valid cache exists, and we have a system prompt. Attempt cache creation.
-			this.isFirstApiCall = false
+		const taskId = this.options.taskId
 
-			// Minimum token check heuristic (simple length check for now, could be improved)
-			// Gemini requires minimum 4096 tokens. A simple length check isn't accurate but avoids complex token counting here.
-			// Let's assume a generous average of 4 chars/token. 4096 tokens * 4 chars/token = 16384 chars.
-			const MIN_SYSTEM_PROMPT_LENGTH_FOR_CACHE = 16384
-			if (systemPrompt.length >= MIN_SYSTEM_PROMPT_LENGTH_FOR_CACHE) {
-				// Start cache creation asynchronously, don't block the main request
-				this.createCacheInBackground(modelId, systemPrompt)
+		// Calculate total content length for cache eligibility check
+		const contentsLength = systemPrompt.length + this.getMessagesLength(contents)
+
+		// Minimum token threshold for caching (approx 4096 tokens)
+		const CONTEXT_CACHE_TOKEN_MINIMUM = 4096
+
+		let uncachedContent: Content[] | undefined = undefined
+		let cachedContent: string | undefined = undefined
+
+		// Check if caching is available and content is large enough to benefit from caching
+		// We only enable caching for conversations above a certain size to avoid overhead for small requests
+		const isCacheAvailable = info.supportsPromptCache && contentsLength > 4 * CONTEXT_CACHE_TOKEN_MINIMUM && taskId
+
+		// This flag tracks whether this operation involves a cache write/update
+		// It's used to track task-level ongoing costs, not immediate costs
+		let cacheWrite = false
+
+		if (isCacheAvailable) {
+			// Check if we already have a cache for this task
+			const existingCacheName = this.taskCacheNames.get(taskId)
+			const cacheEntry = existingCacheName ? this.contentCaches.get<{ key: string; count: number }>(taskId) : undefined
+
+			if (cacheEntry) {
+				// Use existing cache
+				uncachedContent = contents.slice(cacheEntry.count, contents.length)
+				cachedContent = cacheEntry.key
+				console.log(
+					`[GeminiHandler] using existing cache for task ${taskId}: ${cacheEntry.count} cached messages (${cacheEntry.key}) and ${uncachedContent.length} uncached messages`,
+				)
 			}
-			// Proceed with the first request *without* using the cache, as it's being created.
-			useCache = false
-		} else if (!isCacheValid && this.cacheName) {
-			// Cache exists but has expired
-			this.cacheName = null
-			this.cacheExpireTime = null
-			useCache = false
-		}
-		// --- End Cache Handling Logic ---
 
-		// Re-implement thinking budget logic based on new SDK structure
+			// Create or update cache only if there's new content to add
+			const shouldUpdateCache = !existingCacheName || (cacheEntry && uncachedContent && uncachedContent.length > 0)
+
+			if (shouldUpdateCache) {
+				// If we should update the cache, then there will be a cache write
+				cacheWrite = true
+			}
+		}
+		const isCacheUsed = !!cachedContent
+
+		// Configure thinking budget if supported
 		const thinkingBudget = this.options.thinkingBudgetTokens ?? 0
-		const maxBudget = modelInfo.thinkingConfig?.maxBudget ?? 0
+		const maxBudget = info.thinkingConfig?.maxBudget ?? 0
 
-		// port add baseUrl configuration for gemini api requests (#2843)
-		const httpOptions = this.options.geminiBaseUrl ? { baseUrl: this.options.geminiBaseUrl } : undefined
-
-		// Base generation config - Conditionally include systemInstruction based on cache usage
-		const generationConfig: GenerateContentConfig = {
-			httpOptions,
-			temperature: 0, // Default temperature
-			// Only include systemInstruction if NOT using the cache
-			...(useCache ? {} : { systemInstruction: systemPrompt }),
-		}
-
-		// Convert messages to the format expected by @google/genai
-		// Note: convertAnthropicMessageToGemini might need adjustments
-		const contents: Content[] = messages.map(convertAnthropicMessageToGemini)
-
-		// Construct the main request config - Type as GenerateContentConfig
+		// Set up base generation config
 		const requestConfig: GenerateContentConfig = {
-			...generationConfig,
+			// Add base URL if configured
+			httpOptions: this.options.geminiBaseUrl ? { baseUrl: this.options.geminiBaseUrl } : undefined,
+
+			// Only include systemInstruction if NOT using the cache
+			...(isCacheUsed ? {} : { systemInstruction: systemPrompt }),
+
+			// Set temperature (default to 0)
+			temperature: 0,
 		}
 
 		// Add thinking config if the model supports it
-		if (modelInfo.thinkingConfig?.outputPrice !== undefined && maxBudget > 0) {
+		if (info.thinkingConfig?.outputPrice !== undefined && maxBudget > 0) {
 			requestConfig.thinkingConfig = {
 				thinkingBudget: thinkingBudget,
 			}
 		}
 
-		// Generate content using the new SDK structure via client.models
+		// Generate content using the configured parameters
 		const result = await this.client.models.generateContentStream({
-			model: modelId, // Pass model ID directly
-			contents,
-			// Add cachedContent if using the cache
+			model,
+			contents: uncachedContent ?? contents,
 			config: {
 				...requestConfig,
-				...(useCache ? { cachedContent: this.cacheName! } : {}),
+				...(isCacheUsed ? { cachedContent } : {}),
 			},
 		})
 
-		// Declare variable to hold the last usage metadata found
+		// Update the cache after the LLM request is already sent to avoid blocking
+		// We only update the cache if we have a taskId and the cache write flag is set
+		// This is a non-blocking operation and will not affect the response time
+		if (cacheWrite && taskId) {
+			this.updateCacheContent(taskId, model, contents, systemPrompt)
+		}
+		// Track usage metadata
 		let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined
 
-		// Iterate directly over the stream
+		// Process the stream
 		for await (const chunk of result) {
 			if (chunk.text) {
 				yield {
@@ -111,47 +195,337 @@ export class GeminiHandler implements ApiHandler {
 					text: chunk.text,
 				}
 			}
+
 			if (chunk.usageMetadata) {
 				lastUsageMetadata = chunk.usageMetadata
 			}
 		}
 
+		// Yield usage information at the end
 		if (lastUsageMetadata) {
+			const inputTokens = lastUsageMetadata.promptTokenCount ?? 0
+			const outputTokens = lastUsageMetadata.candidatesTokenCount ?? 0
+			const cacheReadTokens = lastUsageMetadata.cachedContentTokenCount
+
+			// Calculate immediate costs only (excluding cache write/storage costs)
+			const totalCost = this.calculateCost({
+				info,
+				inputTokens,
+				outputTokens,
+				cacheReadTokens,
+			})
+
+			// Store the token count for task-level ongoing cost tracking
+			// This is not included in the immediate costs returned to the user
+			const cacheWriteTokens = cacheWrite ? inputTokens : undefined
+
+			// If this is a cache write operation, update the task's ongoing costs
+			if (cacheWrite && this.options.taskId && inputTokens > 0) {
+				// Log the ongoing costs for debugging
+				const ongoingCosts = this.getTaskOngoingCosts(this.options.taskId)
+				console.log(
+					`[GeminiHandler] Task ${this.options.taskId} ongoing costs: $${ongoingCosts?.toFixed(6) ?? "unknown"}`,
+				)
+			}
+
 			yield {
 				type: "usage",
-				inputTokens: lastUsageMetadata.promptTokenCount ?? 0,
-				outputTokens: lastUsageMetadata.candidatesTokenCount ?? 0,
-				cacheWriteTokens: lastUsageMetadata.cachedContentTokenCount ?? 0,
-				cacheReadTokens: useCache ? (lastUsageMetadata.promptTokenCount ?? 0) : 0, // If cache used, prompt tokens are read from cache
+				inputTokens,
+				outputTokens,
+				cacheReadTokens,
+				cacheWriteTokens,
+				totalCost,
 			}
 		}
 	}
 
-	private async createCacheInBackground(modelId: string, systemInstruction: string): Promise<void> {
+	/**
+	 * Lists all caches for the current API key.
+	 *
+	 * According to the Gemini API documentation, you can retrieve metadata for all uploaded caches
+	 * using the caches.list() method. This is useful for monitoring cache usage and cleanup.
+	 *
+	 * @param pageSize Optional number of caches to return per page (default: 10)
+	 * @returns A promise that resolves to an array of cache metadata objects
+	 */
+	public async listCaches(pageSize: number = 10): Promise<any[]> {
 		try {
-			const cache = await this.client.caches.create({
-				model: modelId,
+			const caches: any[] = []
+			const pager = await this.client.caches.list({ config: { pageSize } })
+
+			let page = pager.page
+			while (true) {
+				for (const cache of page) {
+					caches.push(cache)
+				}
+
+				if (!pager.hasNextPage()) {
+					break
+				}
+				page = await pager.nextPage()
+			}
+
+			return caches
+		} catch (error) {
+			console.error(`[GeminiHandler] Failed to list caches:`, error)
+			return []
+		}
+	}
+
+	/**
+	 * Updates the content of a cache for a specific task.
+	 *
+	 * Since the Gemini API doesn't support incremental updates to cache content,
+	 * this method:
+	 * 1. Creates a new cache with the full content (old + new)
+	 * 2. Deletes the old cache if it exists
+	 * 3. Updates our local tracking to point to the new cache
+	 *
+	 * @param taskId The ID of the task whose cache should be updated
+	 * @param model The model to use for the cache
+	 * @param contents The full content to cache (including both old and new messages)
+	 * @param systemInstruction The system instruction to include in the cache
+	 */
+	private async updateCacheContent(
+		taskId: string,
+		model: string,
+		contents: Content[],
+		systemInstruction: string,
+	): Promise<void> {
+		if (this.isCacheBusy) {
+			console.log(`[GeminiHandler] Cache is busy, skipping update for task ${taskId}`)
+			return
+		}
+
+		this.isCacheBusy = true
+		const timestamp = Date.now()
+		const existingCacheName = this.taskCacheNames.get(taskId)
+
+		try {
+			// 1. Create a new cache with the full content
+			const result = await this.client.caches.create({
+				model,
 				config: {
-					systemInstruction: systemInstruction,
+					contents,
+					systemInstruction,
 					ttl: `${DEFAULT_CACHE_TTL_SECONDS}s`,
+					httpOptions: { timeout: 120_000 },
 				},
 			})
 
-			if (cache?.name) {
-				this.cacheName = cache.name
-				// Calculate expiry timestamp using the default TTL, as the response object might not contain it directly.
-				this.cacheExpireTime = Date.now() + DEFAULT_CACHE_TTL_SECONDS * 1000
-			} else {
-				console.warn("Gemini cache creation call succeeded but returned no cache name.")
+			const { name, usageMetadata } = result
+
+			if (name) {
+				// 2. Delete the old cache if it exists (non-blocking)
+				// We don't await this operation to avoid blocking the main flow if deletion fails
+				if (existingCacheName) {
+					// Schedule cache deletion in the background
+					setTimeout(() => {
+						this.client.caches
+							.delete({ name: existingCacheName })
+							.then(() => {
+								console.log(`[GeminiHandler] Deleted old cache ${existingCacheName} for task ${taskId}`)
+							})
+							.catch((error) => {
+								console.error(`[GeminiHandler] Failed to delete old cache ${existingCacheName}:`, error)
+								console.log(`[GeminiHandler] Continuing without deleting old cache. It will expire after TTL.`)
+							})
+					}, 1000)
+				}
+
+				// 3. Update our local tracking
+				this.contentCaches.set<{ key: string; count: number }>(taskId, {
+					key: name,
+					count: contents.length,
+				})
+				this.taskCacheNames.set(taskId, name)
+
+				// Track total tokens in cache for ongoing cost calculation
+				const totalTokens = usageMetadata?.totalTokenCount ?? 0
+				this.taskCacheTokens.set(taskId, totalTokens)
+
+				const operation = existingCacheName ? "Updated" : "Created new"
+				console.log(
+					`[GeminiHandler] ${operation} cache for task ${taskId}: ${contents.length} messages (${totalTokens} tokens) in ${Date.now() - timestamp}ms`,
+				)
+
+				return // Indicate that a cache write occurred
 			}
+
+			return
 		} catch (error) {
-			console.error("Failed to create Gemini cache in background:", error)
-			// Reset state if creation failed definitively
-			this.cacheName = null
-			this.cacheExpireTime = null
+			console.error(`[GeminiHandler] Failed to update cache for task ${taskId}:`, error)
+			return
+		} finally {
+			this.isCacheBusy = false
 		}
 	}
 
+	/**
+	 * Updates the TTL of an existing cache.
+	 *
+	 * According to the Gemini API documentation, you can update the TTL of a cache
+	 * using the caches.update() method. This is useful for extending the lifetime
+	 * of a cache that's still being used.
+	 *
+	 * @param taskId The ID of the task whose cache TTL should be updated
+	 * @param ttlSeconds The new TTL in seconds
+	 * @returns A promise that resolves to the updated cache, or undefined if the update fails
+	 */
+	public async updateCacheTTL(taskId: string, ttlSeconds: number = DEFAULT_CACHE_TTL_SECONDS): Promise<any> {
+		const cacheName = this.taskCacheNames.get(taskId)
+		if (!cacheName) {
+			console.warn(`[GeminiHandler] No cache found for task ${taskId}, cannot update TTL`)
+			return
+		}
+
+		try {
+			const updatedCache = await this.client.caches.update({
+				name: cacheName,
+				config: { ttl: `${ttlSeconds}s` },
+			})
+
+			console.log(`[GeminiHandler] Updated TTL for cache ${cacheName} to ${ttlSeconds}s`)
+			return updatedCache
+		} catch (error) {
+			console.error(`[GeminiHandler] Failed to update TTL for cache ${cacheName}:`, error)
+		}
+	}
+
+	/**
+	 * Calculate the ongoing costs for a task based on cache storage.
+	 *
+	 * This method calculates the cost of holding tokens in cache for the TTL period.
+	 * These costs are separate from the immediate costs of API calls and should be
+	 * tracked at the task level rather than the message level.
+	 *
+	 * TODO: Surface these ongoing costs to the user in the UI, possibly in:
+	 * - The task header/summary
+	 * - A dedicated "costs" panel or tooltip
+	 * - As part of the total cost calculation for the task
+	 *
+	 * @param taskId The ID of the task to calculate ongoing costs for
+	 * @returns The ongoing cost in dollars, or undefined if no cache exists for the task
+	 */
+	public getTaskOngoingCosts(taskId: string): number | undefined {
+		const tokens = this.taskCacheTokens.get(taskId)
+		if (!tokens) {
+			return undefined
+		}
+
+		const { info } = this.getModel()
+		if (!info.cacheWritesPrice) {
+			return undefined
+		}
+
+		// Calculate the cost of holding tokens in cache for the TTL period
+		// (tokens / 1M) * (price per 1M tokens) * (cache TTL in hours)
+		return info.cacheWritesPrice * (tokens / 1_000_000) * (DEFAULT_CACHE_TTL_SECONDS / 3600)
+	}
+
+	/**
+	 * Calculate the immediate dollar cost of the API call based on token usage and model pricing.
+	 *
+	 * This method accounts for the immediate costs of the API call:
+	 * - Input token costs (for uncached tokens)
+	 * - Output token costs
+	 * - Cache read costs
+	 *
+	 * It does NOT include ongoing costs like cache storage, which are tracked separately
+	 * at the task level through getTaskOngoingCosts().
+	 */
+	public calculateCost({
+		info,
+		inputTokens,
+		outputTokens,
+		cacheWriteTokens = 0,
+		cacheReadTokens = 0,
+	}: {
+		info: ModelInfo
+		inputTokens: number
+		outputTokens: number
+		cacheWriteTokens?: number
+		cacheReadTokens?: number
+	}) {
+		// Exit early if any required pricing information is missing
+		if (!info.inputPrice || !info.outputPrice) {
+			return undefined
+		}
+
+		let inputPrice = info.inputPrice
+		let outputPrice = info.outputPrice
+		let cacheWritesPrice = info.cacheWritesPrice ?? 0
+		// Right now, we only show the immediate costs of caching and not the ongoing costs of storing the cache
+		cacheWritesPrice = 0
+		let cacheReadsPrice = info.cacheReadsPrice ?? 0
+
+		// If there's tiered pricing then adjust prices based on the input tokens used
+		if (info.tiers) {
+			const tier = info.tiers.find((tier) => inputTokens <= tier.contextWindow)
+			if (tier) {
+				inputPrice = tier.inputPrice ?? inputPrice
+				outputPrice = tier.outputPrice ?? outputPrice
+				cacheWritesPrice = tier.cacheWritesPrice ?? cacheWritesPrice
+				cacheReadsPrice = tier.cacheReadsPrice ?? cacheReadsPrice
+			}
+		}
+
+		// Subtract the cached input tokens from the total input tokens
+		const uncachedInputTokens = inputTokens - (cacheReadTokens ?? 0)
+
+		// Calculate immediate costs only
+
+		// 1. Input token costs (for uncached tokens)
+		const inputTokensCost = inputPrice * (uncachedInputTokens / 1_000_000)
+
+		// 2. Output token costs
+		const outputTokensCost = outputPrice * (outputTokens / 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
+
+		// 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 },
+		}
+
+		// Only include cache read costs in the trace (cache write costs are tracked separately)
+		if ((cacheReadTokens ?? 0) > 0) {
+			trace.cacheRead = { price: cacheReadsPrice, tokens: cacheReadTokens ?? 0, cost: cacheReadCost }
+		}
+
+		// console.log(`[GeminiHandler] calculateCost -> ${totalCost}`, trace)
+		return totalCost
+	}
+
+	/**
+	 * Calculate the total length of all messages for cache eligibility check
+	 */
+	private getMessagesLength(contents: Content[]): number {
+		return contents.reduce((total, content) => {
+			if (!content.parts) {
+				return total
+			}
+
+			return (
+				total +
+				content.parts.reduce((partTotal, part) => {
+					if (typeof part.text === "string") {
+						return partTotal + part.text.length
+					}
+					return partTotal
+				}, 0)
+			)
+		}, 0)
+	}
+
+	/**
+	 * Get the model ID and info for the current configuration
+	 */
 	getModel(): { id: GeminiModelId; info: ModelInfo } {
 		const modelId = this.options.apiModelId
 		if (modelId && modelId in geminiModels) {
@@ -163,4 +537,61 @@ export class GeminiHandler implements ApiHandler {
 			info: geminiModels[geminiDefaultModelId],
 		}
 	}
+
+	/**
+	 * Count tokens in content using the Gemini API
+	 */
+	async countTokens(content: Array<any>): Promise<number> {
+		try {
+			const { id: model } = this.getModel()
+
+			// Convert content to Gemini format
+			const geminiContent = content.map((block) => {
+				if (typeof block === "string") {
+					return { text: block }
+				}
+				return { text: JSON.stringify(block) }
+			})
+
+			// Use Gemini's token counting API
+			const response = await this.client.models.countTokens({
+				model,
+				contents: [{ parts: geminiContent }],
+			})
+
+			if (response.totalTokens === undefined) {
+				console.warn("Gemini token counting returned undefined, using fallback")
+				return this.estimateTokens(content)
+			}
+
+			return response.totalTokens
+		} catch (error) {
+			console.warn("Gemini token counting failed, using fallback", error)
+			return this.estimateTokens(content)
+		}
+	}
+
+	/**
+	 * Fallback token estimation method
+	 */
+	private estimateTokens(content: Array<any>): number {
+		// Simple estimation: ~4 characters per token
+		const totalChars = content.reduce((total, block) => {
+			if (typeof block === "string") {
+				return total + block.length
+			} else if (block && typeof block === "object") {
+				// Safely stringify the object
+				try {
+					const jsonStr = JSON.stringify(block)
+					return total + jsonStr.length
+				} catch (e) {
+					console.warn("Failed to stringify block for token estimation", e)
+					return total
+				}
+			}
+			return total
+		}, 0)
+
+		return Math.ceil(totalChars / 4)
+	}
 }

src/api/providers/vertex.ts

@@ -4,26 +4,28 @@ import { withRetry } from "../retry"
 import { ApiHandler } from "../"
 import { ApiHandlerOptions, ModelInfo, vertexDefaultModelId, VertexModelId, vertexModels } from "@shared/api"
 import { ApiStream } from "@api/transform/stream"
-import { VertexAI } from "@google-cloud/vertexai"
-import { calculateApiCostOpenAI } from "@utils/cost"
+import { GeminiHandler } from "./gemini"
 
-// https://docs.anthropic.com/en/api/claude-on-vertex-ai
 export class VertexHandler implements ApiHandler {
-	private options: ApiHandlerOptions
+	private geminiHandler: GeminiHandler
 	private clientAnthropic: AnthropicVertex
-	private clientVertex: VertexAI
+	private options: ApiHandlerOptions
 
 	constructor(options: ApiHandlerOptions) {
 		this.options = options
+
+		// Create a GeminiHandler with isVertex flag for Gemini models
+		this.geminiHandler = new GeminiHandler({
+			...options,
+			isVertex: true,
+		})
+
+		// Initialize Anthropic client for Claude models
 		this.clientAnthropic = new AnthropicVertex({
 			projectId: this.options.vertexProjectId,
 			// https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#regions
 			region: this.options.vertexRegion,
 		})
-		this.clientVertex = new VertexAI({
-			project: this.options.vertexProjectId,
-			location: this.options.vertexRegion,
-		})
 	}
 
 	@withRetry()
@@ -31,66 +33,46 @@ export class VertexHandler implements ApiHandler {
 		const model = this.getModel()
 		const modelId = model.id
 
-		if (modelId.includes("claude")) {
-			let budget_tokens = this.options.thinkingBudgetTokens || 0
-			const reasoningOn = modelId.includes("3-7") && budget_tokens !== 0 ? true : false
+		// For Gemini models, use the GeminiHandler
+		if (!modelId.includes("claude")) {
+			yield* this.geminiHandler.createMessage(systemPrompt, messages)
+			return
+		}
 
-			let stream
-			switch (modelId) {
-				case "claude-3-7-sonnet@20250219":
-				case "claude-3-5-sonnet-v2@20241022":
-				case "claude-3-5-sonnet@20240620":
-				case "claude-3-5-haiku@20241022":
-				case "claude-3-opus@20240229":
-				case "claude-3-haiku@20240307": {
-					// Find indices of user messages for cache control
-					const userMsgIndices = messages.reduce(
-						(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
-						[] as number[],
-					)
-					const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
-					const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
+		// Claude implementation
+		let budget_tokens = this.options.thinkingBudgetTokens || 0
+		const reasoningOn = modelId.includes("3-7") && budget_tokens !== 0 ? true : false
+		let stream
 
-					stream = await this.clientAnthropic.beta.messages.create(
-						{
-							model: modelId,
-							max_tokens: model.info.maxTokens || 8192,
-							thinking: reasoningOn ? { type: "enabled", budget_tokens: budget_tokens } : undefined,
-							temperature: reasoningOn ? undefined : 0,
-							system: [
-								{
-									text: systemPrompt,
-									type: "text",
-									cache_control: { type: "ephemeral" },
-								},
-							],
-							messages: messages.map((message, index) => {
-								if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) {
-									return {
-										...message,
-										content:
-											typeof message.content === "string"
-												? [
-														{
-															type: "text",
-															text: message.content,
-															cache_control: {
-																type: "ephemeral",
-															},
-														},
-													]
-												: message.content.map((content, contentIndex) =>
-														contentIndex === message.content.length - 1
-															? {
-																	...content,
-																	cache_control: {
-																		type: "ephemeral",
-																	},
-																}
-															: content,
-													),
-									}
-								}
+		switch (modelId) {
+			case "claude-3-7-sonnet@20250219":
+			case "claude-3-5-sonnet-v2@20241022":
+			case "claude-3-5-sonnet@20240620":
+			case "claude-3-5-haiku@20241022":
+			case "claude-3-opus@20240229":
+			case "claude-3-haiku@20240307": {
+				// Find indices of user messages for cache control
+				const userMsgIndices = messages.reduce(
+					(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
+					[] as number[],
+				)
+				const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
+				const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
+				stream = await this.clientAnthropic.beta.messages.create(
+					{
+						model: modelId,
+						max_tokens: model.info.maxTokens || 8192,
+						thinking: reasoningOn ? { type: "enabled", budget_tokens: budget_tokens } : undefined,
+						temperature: reasoningOn ? undefined : 0,
+						system: [
+							{
+								text: systemPrompt,
+								type: "text",
+								cache_control: { type: "ephemeral" },
+							},
+						],
+						messages: messages.map((message, index) => {
+							if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) {
 								return {
 									...message,
 									content:
@@ -99,184 +81,142 @@ export class VertexHandler implements ApiHandler {
 													{
 														type: "text",
 														text: message.content,
+														cache_control: {
+															type: "ephemeral",
+														},
 													},
 												]
-											: message.content,
+											: message.content.map((content, contentIndex) =>
+													contentIndex === message.content.length - 1
+														? {
+																...content,
+																cache_control: {
+																	type: "ephemeral",
+																},
+															}
+														: content,
+												),
 								}
-							}),
-							stream: true,
-						},
-						{
-							headers: {},
-						},
-					)
-					break
-				}
-				default: {
-					stream = await this.clientAnthropic.beta.messages.create({
-						model: modelId,
-						max_tokens: model.info.maxTokens || 8192,
-						temperature: 0,
-						system: [
-							{
-								text: systemPrompt,
-								type: "text",
-							},
-						],
-						messages: messages.map((message) => ({
-							...message,
-							content:
-								typeof message.content === "string"
-									? [
-											{
-												type: "text",
-												text: message.content,
-											},
-										]
-									: message.content,
-						})),
-						stream: true,
-					})
-					break
-				}
-			}
-			for await (const chunk of stream) {
-				switch (chunk.type) {
-					case "message_start":
-						const usage = chunk.message.usage
-						yield {
-							type: "usage",
-							inputTokens: usage.input_tokens || 0,
-							outputTokens: usage.output_tokens || 0,
-							cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
-							cacheReadTokens: usage.cache_read_input_tokens || undefined,
-						}
-						break
-					case "message_delta":
-						yield {
-							type: "usage",
-							inputTokens: 0,
-							outputTokens: chunk.usage.output_tokens || 0,
-						}
-						break
-					case "message_stop":
-						break
-					case "content_block_start":
-						switch (chunk.content_block.type) {
-							case "thinking":
-								yield {
-									type: "reasoning",
-									reasoning: chunk.content_block.thinking || "",
-								}
-								break
-							case "redacted_thinking":
-								// Handle redacted thinking blocks - we still mark it as reasoning
-								// but note that the content is encrypted
-								yield {
-									type: "reasoning",
-									reasoning: "[Redacted thinking block]",
-								}
-								break
-
-							case "text":
-								if (chunk.index > 0) {
-									yield {
-										type: "text",
-										text: "\n",
-									}
-								}
-								yield {
-									type: "text",
-									text: chunk.content_block.text,
-								}
-								break
-						}
-						break
-					case "content_block_delta":
-						switch (chunk.delta.type) {
-							case "thinking_delta":
-								yield {
-									type: "reasoning",
-									reasoning: chunk.delta.thinking,
-								}
-								break
-							case "text_delta":
-								yield {
-									type: "text",
-									text: chunk.delta.text,
-								}
-								break
-						}
-						break
-					case "content_block_stop":
-						break
-				}
-			}
-		} else {
-			// gemini
-			const generativeModel = this.clientVertex.getGenerativeModel({
-				model: this.getModel().id,
-				systemInstruction: {
-					role: "system",
-					parts: [{ text: systemPrompt }],
-				},
-			})
-			const request = {
-				contents: [
-					{
-						role: "user",
-						parts: messages.map((m) => {
-							if (typeof m.content === "string") {
-								return { text: m.content }
-							} else if (Array.isArray(m.content)) {
-								return {
-									text: m.content
-										.map((block) => {
-											if (typeof block === "string") {
-												return block
-											} else if (block.type === "text") {
-												return block.text
-											} else {
-												console.log("Unsupported block type", block)
-												return ""
-											}
-										})
-										.join(" "),
-								}
-							} else {
-								return { text: "" }
+							}
+							return {
+								...message,
+								content:
+									typeof message.content === "string"
+										? [
+												{
+													type: "text",
+													text: message.content,
+												},
+											]
+										: message.content,
 							}
 						}),
+						stream: true,
 					},
-				],
+					{
+						headers: {},
+					},
+				)
+				break
 			}
-			const streamingResult = await generativeModel.generateContentStream(request)
-			for await (const chunk of streamingResult.stream) {
-				// If usage data is available, yield it similarly:
-				// yield { type: "usage", inputTokens: 0, outputTokens: 0 }
-				// Otherwise, just yield text:
-				const candidates = chunk.candidates || []
-				for (const candidate of candidates) {
-					for (const part of candidate.content?.parts || []) {
-						if (part.text) {
+			default: {
+				stream = await this.clientAnthropic.beta.messages.create({
+					model: modelId,
+					max_tokens: model.info.maxTokens || 8192,
+					temperature: 0,
+					system: [
+						{
+							text: systemPrompt,
+							type: "text",
+						},
+					],
+					messages: messages.map((message) => ({
+						...message,
+						content:
+							typeof message.content === "string"
+								? [
+										{
+											type: "text",
+											text: message.content,
+										},
+									]
+								: message.content,
+					})),
+					stream: true,
+				})
+				break
+			}
+		}
+
+		for await (const chunk of stream) {
+			switch (chunk.type) {
+				case "message_start":
+					const usage = chunk.message.usage
+					yield {
+						type: "usage",
+						inputTokens: usage.input_tokens || 0,
+						outputTokens: usage.output_tokens || 0,
+						cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
+						cacheReadTokens: usage.cache_read_input_tokens || undefined,
+					}
+					break
+				case "message_delta":
+					yield {
+						type: "usage",
+						inputTokens: 0,
+						outputTokens: chunk.usage?.output_tokens || 0,
+					}
+					break
+				case "message_stop":
+					break
+				case "content_block_start":
+					switch (chunk.content_block.type) {
+						case "thinking":
+							yield {
+								type: "reasoning",
+								reasoning: chunk.content_block.thinking || "",
+							}
+							break
+						case "redacted_thinking":
+							// Handle redacted thinking blocks - we still mark it as reasoning
+							// but note that the content is encrypted
+							yield {
+								type: "reasoning",
+								reasoning: "[Redacted thinking block]",
+							}
+							break
+						case "text":
+							if (chunk.index > 0) {
+								yield {
+									type: "text",
+									text: "\n",
+								}
+							}
 							yield {
 								type: "text",
-								text: part.text,
+								text: chunk.content_block.text,
 							}
-						}
+							break
 					}
-				}
-			}
-			// Handle token usage metadata
-			const { usageMetadata } = await streamingResult.response
-			if (usageMetadata) {
-				const { promptTokenCount = 0, candidatesTokenCount = 0 } = usageMetadata
-				yield {
-					type: "usage",
-					inputTokens: promptTokenCount,
-					outputTokens: candidatesTokenCount,
-					totalCost: calculateApiCostOpenAI(model.info, promptTokenCount, candidatesTokenCount, 0, 0),
-				}
+					break
+				case "content_block_delta":
+					switch (chunk.delta.type) {
+						case "thinking_delta":
+							yield {
+								type: "reasoning",
+								reasoning: chunk.delta.thinking,
+							}
+							break
+						case "text_delta":
+							yield {
+								type: "text",
+								text: chunk.delta.text,
+							}
+							break
+					}
+					break
+				case "content_block_stop":
+					break
 			}
 		}
 	}

src/shared/api.ts

@@ -103,9 +103,7 @@ export interface ModelInfo {
 	supportsImages?: boolean
 	supportsPromptCache: boolean // this value is hardcoded for now
 	inputPrice?: number // Keep for non-tiered input models
-	inputPriceTiers?: PriceTier[] // Add for tiered input pricing
 	outputPrice?: number // Keep for non-tiered output models
-	outputPriceTiers?: PriceTier[] // Add for tiered output pricing
 	thinkingConfig?: {
 		maxBudget?: number // Max allowed thinking budget tokens
 		outputPrice?: number // Output price per million tokens when budget > 0
@@ -114,6 +112,13 @@ export interface ModelInfo {
 	cacheWritesPrice?: number
 	cacheReadsPrice?: number
 	description?: string
+	tiers?: {
+		contextWindow: number
+		inputPrice?: number
+		outputPrice?: number
+		cacheWritesPrice?: number
+		cacheReadsPrice?: number
+	}[]
 }
 
 export interface OpenAiCompatibleModelInfo extends ModelInfo {
@@ -321,10 +326,15 @@ export const vertexModels = {
 		maxTokens: 8192,
 		contextWindow: 200_000,
 		supportsImages: true,
-
 		supportsPromptCache: true,
 		inputPrice: 3.0,
 		outputPrice: 15.0,
+		cacheWritesPrice: 3.75,
+		cacheReadsPrice: 0.3,
+		thinkingConfig: {
+			maxBudget: 64000,
+			outputPrice: 15.0,
+		},
 	},
 	"claude-3-5-sonnet-v2@20241022": {
 		maxTokens: 8192,
@@ -378,12 +388,22 @@ export const vertexModels = {
 		cacheReadsPrice: 0.03,
 	},
 	"gemini-2.0-flash-001": {
+		maxTokens: 8192,
+		contextWindow: 1_048_576,
+		supportsImages: true,
+		supportsPromptCache: true,
+		inputPrice: 0.15,
+		outputPrice: 0.6,
+		cacheWritesPrice: 1.0,
+		cacheReadsPrice: 0.025,
+	},
+	"gemini-2.0-flash-lite-001": {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
 		supportsPromptCache: false,
-		inputPrice: 0.1,
-		outputPrice: 0.4,
+		inputPrice: 0.075,
+		outputPrice: 0.3,
 	},
 	"gemini-2.0-flash-thinking-exp-1219": {
 		maxTokens: 8192,
@@ -413,16 +433,20 @@ export const vertexModels = {
 		maxTokens: 65536,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
-		// inputPrice: 1.25, // Removed
-		// outputPrice: 10, // Removed
-		inputPriceTiers: [
-			{ tokenLimit: 200000, price: 1.25 }, // Input price for <= 200k input tokens
-			{ tokenLimit: Infinity, price: 2.5 }, // Input price for > 200k input tokens
-		],
-		outputPriceTiers: [
-			{ tokenLimit: 200000, price: 10.0 }, // Output price for <= 200k input tokens
-			{ tokenLimit: Infinity, price: 15.0 }, // Output price for > 200k input tokens
+		supportsPromptCache: true,
+		inputPrice: 2.5,
+		outputPrice: 15,
+		tiers: [
+			{
+				contextWindow: 200000,
+				inputPrice: 1.25,
+				outputPrice: 10,
+			},
+			{
+				contextWindow: Infinity,
+				inputPrice: 2.5,
+				outputPrice: 15,
+			},
 		],
 	},
 	"gemini-2.5-flash-preview-04-17": {
@@ -457,9 +481,25 @@ export const vertexModels = {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
-		inputPrice: 0,
-		outputPrice: 0,
+		supportsPromptCache: true,
+		inputPrice: 0.15,
+		outputPrice: 0.6,
+		cacheWritesPrice: 1.0,
+		cacheReadsPrice: 0.0375,
+		tiers: [
+			{
+				contextWindow: 128000,
+				inputPrice: 0.075,
+				outputPrice: 0.3,
+				cacheReadsPrice: 0.01875,
+			},
+			{
+				contextWindow: Infinity,
+				inputPrice: 0.15,
+				outputPrice: 0.6,
+				cacheReadsPrice: 0.0375,
+			},
+		],
 	},
 	"gemini-1.5-flash-exp-0827": {
 		maxTokens: 8192,
@@ -482,8 +522,8 @@ export const vertexModels = {
 		contextWindow: 2_097_152,
 		supportsImages: true,
 		supportsPromptCache: false,
-		inputPrice: 0,
-		outputPrice: 0,
+		inputPrice: 1.25,
+		outputPrice: 5,
 	},
 	"gemini-1.5-pro-exp-0827": {
 		maxTokens: 8192,
@@ -523,14 +563,24 @@ export const geminiModels = {
 		maxTokens: 65536,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
-		inputPriceTiers: [
-			{ tokenLimit: 200000, price: 1.25 }, // Input price for <= 200k input tokens
-			{ tokenLimit: Infinity, price: 2.5 }, // Input price for > 200k input tokens
-		],
-		outputPriceTiers: [
-			{ tokenLimit: 200000, price: 10.0 }, // Output price for <= 200k input tokens
-			{ tokenLimit: Infinity, price: 15.0 }, // Output price for > 200k input tokens
+		supportsPromptCache: true,
+		inputPrice: 2.5, // Default price (highest tier)
+		outputPrice: 15, // Default price (highest tier)
+		cacheReadsPrice: 0.625,
+		cacheWritesPrice: 4.5,
+		tiers: [
+			{
+				contextWindow: 200000,
+				inputPrice: 1.25,
+				outputPrice: 10,
+				cacheReadsPrice: 0.31,
+			},
+			{
+				contextWindow: Infinity,
+				inputPrice: 2.5,
+				outputPrice: 15,
+				cacheReadsPrice: 0.625,
+			},
 		],
 	},
 	"gemini-2.5-flash-preview-04-17": {
@@ -549,9 +599,11 @@ export const geminiModels = {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
-		inputPrice: 0,
-		outputPrice: 0,
+		supportsPromptCache: true,
+		inputPrice: 0.1,
+		outputPrice: 0.4,
+		cacheReadsPrice: 0.025,
+		cacheWritesPrice: 1.0,
 	},
 	"gemini-2.0-flash-lite-preview-02-05": {
 		maxTokens: 8192,
@@ -597,9 +649,25 @@ export const geminiModels = {
 		maxTokens: 8192,
 		contextWindow: 1_048_576,
 		supportsImages: true,
-		supportsPromptCache: false,
-		inputPrice: 0,
-		outputPrice: 0,
+		supportsPromptCache: true,
+		inputPrice: 0.15, // Default price (highest tier)
+		outputPrice: 0.6, // Default price (highest tier)
+		cacheReadsPrice: 0.0375,
+		cacheWritesPrice: 1.0,
+		tiers: [
+			{
+				contextWindow: 128000,
+				inputPrice: 0.075,
+				outputPrice: 0.3,
+				cacheReadsPrice: 0.01875,
+			},
+			{
+				contextWindow: Infinity,
+				inputPrice: 0.15,
+				outputPrice: 0.6,
+				cacheReadsPrice: 0.0375,
+			},
+		],
 	},
 	"gemini-1.5-flash-exp-0827": {
 		maxTokens: 8192,

src/utils/cost.ts

@@ -11,51 +11,56 @@ function calculateApiCostInternal(
 ): number {
 	const usedThinkingBudget = thinkingBudgetTokens && thinkingBudgetTokens > 0
 
-	// Determine effective input price
+	// Default prices
 	let effectiveInputPrice = modelInfo.inputPrice || 0
-	if (modelInfo.inputPriceTiers && modelInfo.inputPriceTiers.length > 0 && totalInputTokensForPricing !== undefined) {
-		// Ensure tiers are sorted by tokenLimit ascending before finding
-		const sortedInputTiers = [...modelInfo.inputPriceTiers].sort((a, b) => a.tokenLimit - b.tokenLimit)
+	let effectiveOutputPrice = modelInfo.outputPrice || 0
+	let effectiveCacheReadsPrice = modelInfo.cacheReadsPrice || 0
+	let effectiveCacheWritesPrice = modelInfo.cacheWritesPrice || 0
+
+	// Handle tiered pricing if available
+	if (modelInfo.tiers && modelInfo.tiers.length > 0 && totalInputTokensForPricing !== undefined) {
+		// Ensure tiers are sorted by contextWindow ascending before finding
+		const sortedTiers = [...modelInfo.tiers].sort((a, b) => a.contextWindow - b.contextWindow)
+
 		// Find the first tier where the total input tokens are less than or equal to the limit
-		const tier = sortedInputTiers.find((t) => totalInputTokensForPricing! <= t.tokenLimit)
+		const tier = sortedTiers.find((t) => totalInputTokensForPricing <= t.contextWindow)
+
 		if (tier) {
-			effectiveInputPrice = tier.price
+			// Apply all tiered price values if they exist
+			effectiveInputPrice = tier.inputPrice ?? effectiveInputPrice
+			effectiveOutputPrice = tier.outputPrice ?? effectiveOutputPrice
+			effectiveCacheReadsPrice = tier.cacheReadsPrice ?? effectiveCacheReadsPrice
+			effectiveCacheWritesPrice = tier.cacheWritesPrice ?? effectiveCacheWritesPrice
 		} else {
 			// Should ideally not happen if Infinity is used for the last tier, but fallback just in case
-			effectiveInputPrice = sortedInputTiers[sortedInputTiers.length - 1]?.price || 0
+			const lastTier = sortedTiers[sortedTiers.length - 1]
+			if (lastTier) {
+				effectiveInputPrice = lastTier.inputPrice ?? effectiveInputPrice
+				effectiveOutputPrice = lastTier.outputPrice ?? effectiveOutputPrice
+				effectiveCacheReadsPrice = lastTier.cacheReadsPrice ?? effectiveCacheReadsPrice
+				effectiveCacheWritesPrice = lastTier.cacheWritesPrice ?? effectiveCacheWritesPrice
+			}
 		}
 	}
 
-	// Determine effective output price
-	let effectiveOutputPrice = modelInfo.outputPrice || 0
-	// Check if thinking budget was used and has a specific price
+	// Override output price for thinking mode if applicable
 	if (usedThinkingBudget && modelInfo.thinkingConfig?.outputPrice !== undefined) {
 		effectiveOutputPrice = modelInfo.thinkingConfig.outputPrice
 		// TODO: Add support for tiered thinking budget output pricing if needed in the future
-		// } else if (usedThinkingBudget && modelInfo.thinkingConfig?.outputPriceTiers) { ... }
-	} else if (modelInfo.outputPriceTiers && modelInfo.outputPriceTiers.length > 0 && totalInputTokensForPricing !== undefined) {
-		// Use standard tiered output pricing (based on total *input* tokens for pricing)
-		const sortedOutputTiers = [...modelInfo.outputPriceTiers].sort((a, b) => a.tokenLimit - b.tokenLimit)
-		const tier = sortedOutputTiers.find((t) => totalInputTokensForPricing! <= t.tokenLimit)
-		if (tier) {
-			effectiveOutputPrice = tier.price
-		} else {
-			// Should ideally not happen if Infinity is used for the last tier, but fallback just in case
-			effectiveOutputPrice = sortedOutputTiers[sortedOutputTiers.length - 1]?.price || 0
-		}
 	}
 
-	const cacheWritesCost = ((modelInfo.cacheWritesPrice || 0) / 1_000_000) * cacheCreationInputTokens
-	const cacheReadsCost = ((modelInfo.cacheReadsPrice || 0) / 1_000_000) * cacheReadInputTokens
+	const cacheWritesCost = (effectiveCacheWritesPrice / 1_000_000) * cacheCreationInputTokens
+	const cacheReadsCost = (effectiveCacheReadsPrice / 1_000_000) * cacheReadInputTokens
+
 	// Use effectiveInputPrice for baseInputCost. Note: 'inputTokens' here is the potentially adjusted count (e.g., non-cached for OpenAI)
 	const baseInputCost = (effectiveInputPrice / 1_000_000) * inputTokens
+
 	// Use effectiveOutputPrice for outputCost
 	const outputCost = (effectiveOutputPrice / 1_000_000) * outputTokens
 
 	const totalCost = cacheWritesCost + cacheReadsCost + baseInputCost + outputCost
 	return totalCost
 }
-
 // For Anthropic compliant usage, the input tokens count does NOT include the cached tokens
 export function calculateApiCostAnthropic(
 	modelInfo: ModelInfo,

webview-ui/src/components/chat/TaskHeader.tsx

@@ -130,8 +130,7 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({
 		return (
 			apiConfiguration?.apiProvider !== "vscode-lm" &&
 			apiConfiguration?.apiProvider !== "ollama" &&
-			apiConfiguration?.apiProvider !== "lmstudio" &&
-			apiConfiguration?.apiProvider !== "gemini"
+			apiConfiguration?.apiProvider !== "lmstudio"
 		)
 	}, [apiConfiguration?.apiProvider, apiConfiguration?.openAiModelInfo])
 

webview-ui/src/components/settings/ApiOptions.tsx

@@ -1920,31 +1920,39 @@ export const formatPrice = (price: number) => {
 }
 
 // Returns an array of formatted tier strings
-const formatTiers = (tiers: ModelInfo["inputPriceTiers"]): JSX.Element[] => {
+const formatTiers = (
+	tiers: ModelInfo["tiers"],
+	priceType: "inputPrice" | "outputPrice" | "cacheReadsPrice" | "cacheWritesPrice",
+): JSX.Element[] => {
 	if (!tiers || tiers.length === 0) {
 		return []
 	}
-	return tiers.map((tier, index, arr) => {
-		const prevLimit = index > 0 ? arr[index - 1].tokenLimit : 0
-		return (
-			<span style={{ paddingLeft: "15px" }} key={index}>
-				{formatPrice(tier.price)}/million tokens (
-				{tier.tokenLimit === Number.POSITIVE_INFINITY ? (
-					<span>
-						{"> "}
-						{prevLimit.toLocaleString()}
-					</span>
-				) : (
-					<span>
-						{"<= "}
-						{tier.tokenLimit.toLocaleString()}
-					</span>
-				)}
-				{" tokens)"}
-				{index < arr.length - 1 && <br />}
-			</span>
-		)
-	})
+
+	return tiers
+		.map((tier, index, arr) => {
+			const prevLimit = index > 0 ? arr[index - 1].contextWindow : 0
+			const price = tier[priceType]
+
+			if (price === undefined) return null
+
+			return (
+				<span style={{ paddingLeft: "15px" }} key={index}>
+					{formatPrice(price)}/million tokens (
+					{tier.contextWindow === Number.POSITIVE_INFINITY ? (
+						<span>
+							{">"} {prevLimit.toLocaleString()}
+						</span>
+					) : (
+						<span>
+							{"<="} {tier.contextWindow.toLocaleString()}
+						</span>
+					)}
+					{" tokens)"}
+					{index < arr.length - 1 && <br />}
+				</span>
+			)
+		})
+		.filter((element): element is JSX.Element => element !== null)
 }
 
 export const ModelInfoView = ({
@@ -1962,13 +1970,14 @@ export const ModelInfoView = ({
 }) => {
 	const isGemini = Object.keys(geminiModels).includes(selectedModelId)
 	const hasThinkingConfig = !!modelInfo.thinkingConfig
+	const hasTiers = !!modelInfo.tiers && modelInfo.tiers.length > 0
 
-	// Create elements for tiered pricing separately
-	const inputPriceElement = modelInfo.inputPriceTiers ? (
+	// Create elements for input pricing
+	const inputPriceElement = hasTiers ? (
 		<Fragment key="inputPriceTiers">
 			<span style={{ fontWeight: 500 }}>Input price:</span>
 			<br />
-			{formatTiers(modelInfo.inputPriceTiers)}
+			{formatTiers(modelInfo.tiers, "inputPrice")}
 		</Fragment>
 	) : modelInfo.inputPrice !== undefined && modelInfo.inputPrice > 0 ? (
 		<span key="inputPrice">
@@ -1989,14 +1998,14 @@ export const ModelInfoView = ({
 				{formatPrice(modelInfo.thinkingConfig.outputPrice)}/million tokens
 			</Fragment>
 		)
-	} else if (modelInfo.outputPriceTiers) {
+	} else if (hasTiers) {
 		// Display tiered output pricing
 		outputPriceElement = (
 			<Fragment key="outputPriceTiers">
 				<span style={{ fontWeight: 500 }}>Output price:</span>
 				<span style={{ fontStyle: "italic" }}> (based on input tokens)</span>
 				<br />
-				{formatTiers(modelInfo.outputPriceTiers)}
+				{formatTiers(modelInfo.tiers, "outputPrice")}
 			</Fragment>
 		)
 	} else if (modelInfo.outputPrice !== undefined && modelInfo.outputPrice > 0) {