From 0f54e4f4303311ec2d866388dbcb421463b22812 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Mon, 1 Sep 2025 13:32:09 +0000 Subject: [PATCH] fix: use local token counting for VS Code LM to prevent quota exhaustion - Replace VS Code LM API countTokens calls with local tiktoken-based counting - This prevents consuming GitHub Copilot quota for token counting operations - Fixes issue where users exhausted their entire quota in one day Fixes #7575 --- src/api/providers/__tests__/vscode-lm.spec.ts | 20 +++- src/api/providers/vscode-lm.ts | 108 +++++------------- 2 files changed, 47 insertions(+), 81 deletions(-) diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index afb349e5e0..93a07d6733 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -16,6 +16,14 @@ vi.mock("vscode", () => { ) {} } + class MockLanguageModelToolResultPart { + type = "tool_result" + constructor( + public callId: string, + public content: MockLanguageModelTextPart[], + ) {} + } + return { workspace: { onDidChangeConfiguration: vi.fn((_callback) => ({ @@ -48,6 +56,7 @@ vi.mock("vscode", () => { }, LanguageModelTextPart: MockLanguageModelTextPart, LanguageModelToolCallPart: MockLanguageModelToolCallPart, + LanguageModelToolResultPart: MockLanguageModelToolResultPart, lm: { selectChatModels: vi.fn(), }, @@ -135,7 +144,8 @@ describe("VsCodeLmHandler", () => { beforeEach(() => { const mockModel = { ...mockLanguageModelChat } ;(vscode.lm.selectChatModels as Mock).mockResolvedValueOnce([mockModel]) - mockLanguageModelChat.countTokens.mockResolvedValue(10) + // Note: We no longer mock countTokens as we're using local token counting + // to avoid consuming VS Code LM API quota // Override the default client with our test client handler["client"] = mockLanguageModelChat @@ -178,6 +188,10 @@ describe("VsCodeLmHandler", () => { inputTokens: expect.any(Number), outputTokens: expect.any(Number), }) + + // Verify that VS Code LM API's countTokens was NOT called + // to avoid consuming user quota + expect(mockLanguageModelChat.countTokens).not.toHaveBeenCalled() }) it("should handle tool calls", async () => { @@ -221,6 +235,10 @@ describe("VsCodeLmHandler", () => { type: "text", text: JSON.stringify({ type: "tool_call", ...toolCallData }), }) + + // Verify that VS Code LM API's countTokens was NOT called + // to avoid consuming user quota + expect(mockLanguageModelChat.countTokens).not.toHaveBeenCalled() }) it("should handle errors", async () => { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index d8a492f772..2d7e3fdd5e 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -5,6 +5,7 @@ import { type ModelInfo, openAiModelInfoSaneDefaults } from "@roo-code/types" import type { ApiHandlerOptions } from "../../shared/api" import { SELECTOR_SEPARATOR, stringifyVsCodeLmModelSelector } from "../../shared/vsCodeSelectorUtils" +import { countTokens } from "../../utils/countTokens" import { ApiStream } from "../transform/stream" import { convertToVsCodeLmMessages, extractTextCountFromMessage } from "../transform/vscode-lm-format" @@ -177,100 +178,46 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan /** * Implements the ApiHandler countTokens interface method - * Provides token counting for Anthropic content blocks + * Uses local token counting to avoid consuming VS Code LM API quota * * @param content The content blocks to count tokens for * @returns A promise resolving to the token count */ override async countTokens(content: Array): Promise { - // Convert Anthropic content blocks to a string for VSCode LM token counting - let textContent = "" - - for (const block of content) { - if (block.type === "text") { - textContent += block.text || "" - } else if (block.type === "image") { - // VSCode LM doesn't support images directly, so we'll just use a placeholder - textContent += "[IMAGE]" - } - } - - return this.internalCountTokens(textContent) + // Use the base provider's local token counting implementation + // This avoids calling the VS Code LM API which consumes user quota + return super.countTokens(content) } /** - * Private implementation of token counting used internally by VsCodeLmHandler + * Private implementation of local token counting for VS Code LM messages + * This replaces the previous implementation that used the VS Code LM API */ - private async internalCountTokens(text: string | vscode.LanguageModelChatMessage): Promise { - // Check for required dependencies - if (!this.client) { - console.warn("Roo Code : No client available for token counting") + private async localCountTokens(text: string | vscode.LanguageModelChatMessage): Promise { + // Convert to string if it's a message + let textContent = "" + + if (typeof text === "string") { + textContent = text + } else if (text && typeof text === "object") { + // Extract text content from the message object + textContent = extractTextCountFromMessage(text) + } + + // Use local tiktoken-based counting + if (!textContent) { return 0 } - if (!this.currentRequestCancellation) { - console.warn("Roo Code : No cancellation token available for token counting") - return 0 - } + // Convert to Anthropic content blocks for counting + const contentBlocks: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: textContent }] - // Validate input - if (!text) { - console.debug("Roo Code : Empty text provided for token counting") - return 0 - } - - try { - // Handle different input types - let tokenCount: number - - if (typeof text === "string") { - tokenCount = await this.client.countTokens(text, this.currentRequestCancellation.token) - } else if (text instanceof vscode.LanguageModelChatMessage) { - // For chat messages, ensure we have content - if (!text.content || (Array.isArray(text.content) && text.content.length === 0)) { - console.debug("Roo Code : Empty chat message content") - return 0 - } - const countMessage = extractTextCountFromMessage(text) - tokenCount = await this.client.countTokens(countMessage, this.currentRequestCancellation.token) - } else { - console.warn("Roo Code : Invalid input type for token counting") - return 0 - } - - // Validate the result - if (typeof tokenCount !== "number") { - console.warn("Roo Code : Non-numeric token count received:", tokenCount) - return 0 - } - - if (tokenCount < 0) { - console.warn("Roo Code : Negative token count received:", tokenCount) - return 0 - } - - return tokenCount - } catch (error) { - // Handle specific error types - if (error instanceof vscode.CancellationError) { - console.debug("Roo Code : Token counting cancelled by user") - return 0 - } - - const errorMessage = error instanceof Error ? error.message : "Unknown error" - console.warn("Roo Code : Token counting failed:", errorMessage) - - // Log additional error details if available - if (error instanceof Error && error.stack) { - console.debug("Token counting error stack:", error.stack) - } - - return 0 // Fallback to prevent stream interruption - } + return countTokens(contentBlocks, { useWorker: true }) } private async calculateTotalInputTokens(vsCodeLmMessages: vscode.LanguageModelChatMessage[]): Promise { - const messageTokens: number[] = await Promise.all(vsCodeLmMessages.map((msg) => this.internalCountTokens(msg))) + // Use local token counting instead of VS Code LM API + const messageTokens: number[] = await Promise.all(vsCodeLmMessages.map((msg) => this.localCountTokens(msg))) return messageTokens.reduce((sum: number, tokens: number): number => sum + tokens, 0) } @@ -440,8 +387,9 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan } } - // Count tokens in the accumulated text after stream completion - const totalOutputTokens: number = await this.internalCountTokens(accumulatedText) + // Use local token counting for the accumulated output text + // This avoids consuming VS Code LM API quota + const totalOutputTokens: number = await this.localCountTokens(accumulatedText) // Report final usage after stream completion yield {