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
This commit is contained in:
Roo Code 2025-07-17 17:57:56 +00:00
parent 4b1c9021d6
commit 56224459f5
3 changed files with 44 additions and 3 deletions

View file

@ -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)
}
}

View file

@ -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)
}
}

View file

@ -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)
}
}