From 56224459f5ccf23155bcf20159bbc53618df0dec Mon Sep 17 00:00:00 2001
From: Roo Code <roomote@roocode.com>
Date: Thu, 17 Jul 2025 17:57:56 +0000
Subject: [PATCH] fix: improve Vertex AI context tracking and token counting

- Enhanced GeminiHandler.countTokens() to handle zero token counts and provide better fallback
- Added countTokens override in AnthropicVertexHandler to use reliable tiktoken for Claude models
- Added robust token counting with validation in VertexHandler for Gemini models
- Improved error handling and logging for token counting failures
- Ensures consistent context tracking for both Vertex AI Claude and Gemini models

Fixes #5830
---
 src/api/providers/anthropic-vertex.ts |  9 ++++++++
 src/api/providers/gemini.ts           |  6 ++---
 src/api/providers/vertex.ts           | 32 +++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/api/providers/anthropic-vertex.ts b/src/api/providers/anthropic-vertex.ts
index c70a15926d..0609d0bbd9 100644
--- a/src/api/providers/anthropic-vertex.ts
+++ b/src/api/providers/anthropic-vertex.ts
@@ -217,4 +217,13 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
 			throw error
 		}
 	}
+
+	/**
+		* Override token counting to use tiktoken fallback for Vertex AI Claude models.
+		* The Vertex AI token counting API can be unreliable, so we use the base provider's
+		* tiktoken implementation which is more consistent for context tracking.
+		*/
+	override async countTokens(content: Anthropic.Messages.ContentBlockParam[]): Promise<number> {
+		return super.countTokens(content)
+	}
 }
diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts
index 6765c8676d..20e022e1bb 100644
--- a/src/api/providers/gemini.ts
+++ b/src/api/providers/gemini.ts
@@ -176,14 +176,14 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 				contents: convertAnthropicContentToGemini(content),
 			})
 
-			if (response.totalTokens === undefined) {
-				console.warn("Gemini token counting returned undefined, using fallback")
+			if (response.totalTokens === undefined || response.totalTokens === 0) {
+				console.warn("Gemini token counting returned undefined or zero, using fallback")
 				return super.countTokens(content)
 			}
 
 			return response.totalTokens
 		} catch (error) {
-			console.warn("Gemini token counting failed, using fallback", error)
+			console.warn("Gemini token counting failed, using fallback:", error)
 			return super.countTokens(content)
 		}
 	}
diff --git a/src/api/providers/vertex.ts b/src/api/providers/vertex.ts
index 2c077d97b7..0fbab4c828 100644
--- a/src/api/providers/vertex.ts
+++ b/src/api/providers/vertex.ts
@@ -1,4 +1,5 @@
 import { type ModelInfo, type VertexModelId, vertexDefaultModelId, vertexModels } from "@roo-code/types"
+import type { Anthropic } from "@anthropic-ai/sdk"
 
 import type { ApiHandlerOptions } from "../../shared/api"
 
@@ -24,4 +25,35 @@ export class VertexHandler extends GeminiHandler implements SingleCompletionHand
 		// suffix.
 		return { id: id.endsWith(":thinking") ? id.replace(":thinking", "") : id, info, ...params }
 	}
+
+	/**
+	 * Override token counting to add additional error handling for Vertex AI.
+	 * Falls back to tiktoken if the Gemini API token counting fails or returns unreliable results.
+	 */
+	override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
+		try {
+			// Try the parent GeminiHandler's token counting first
+			const tokenCount = await super.countTokens(content)
+			
+			// Additional validation: if token count seems unreasonably low for non-empty content,
+			// fall back to tiktoken
+			if (content.length > 0 && tokenCount === 0) {
+				console.warn("Vertex AI token counting returned 0 for non-empty content, using tiktoken fallback")
+				return this.fallbackTokenCount(content)
+			}
+			
+			return tokenCount
+		} catch (error) {
+			console.warn("Vertex AI token counting failed completely, using tiktoken fallback:", error)
+			return this.fallbackTokenCount(content)
+		}
+	}
+
+	/**
+	 * Fallback token counting using the base provider's tiktoken implementation
+	 */
+	private async fallbackTokenCount(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
+		// Call the base provider's countTokens method (which uses tiktoken)
+		return super.countTokens(content)
+	}
 }