feat(perf): move token counting to Web Worker architecture

Improve performance by offloading token counting to a dedicated Web Worker.
This prevents blocking the main thread during token calculations.

Key changes:

Add token-counter worker with optimized Tiktoken usage
Implement robust WorkerManager with error handling and retries
Update base provider to use worker for token counting
Configure esbuild to properly bundle worker code

Author: samhvw8

esbuild.js

@@ -187,12 +187,28 @@ const extensionConfig = {
 	external: ["vscode"],
 }
 
+const workerConfig = {
+	bundle: true,
+	minify: production,
+	sourcemap: !production,
+	logLevel: "silent",
+	entryPoints: ["src/api/providers/workers/token-counter.worker.ts"],
+	format: "cjs",
+	platform: "node",
+	outfile: "dist/workers/token-counter.worker.js",
+	plugins: [esbuildProblemMatcherPlugin],
+}
+
 async function main() {
-	const extensionCtx = await esbuild.context(extensionConfig)
+	// Create contexts for both worker and extension
+	const [workerCtx, extensionCtx] = await Promise.all([
+		esbuild.context(workerConfig),
+		esbuild.context(extensionConfig),
+	])
 
 	if (watch) {
-		// Start the esbuild watcher
-		await extensionCtx.watch()
+		// Start the esbuild watchers
+		await Promise.all([workerCtx.watch(), extensionCtx.watch()])
 
 		// Copy and watch locale files
 		console.log("Copying locale files initially...")
@@ -201,8 +217,9 @@ async function main() {
 		// Set up the watcher for locale files
 		setupLocaleWatcher()
 	} else {
-		await extensionCtx.rebuild()
-		await extensionCtx.dispose()
+		// Build once and dispose
+		await Promise.all([workerCtx.rebuild(), extensionCtx.rebuild()])
+		await Promise.all([workerCtx.dispose(), extensionCtx.dispose()])
 	}
 }
 

knip.json

@@ -1,6 +1,6 @@
 {
 	"$schema": "https://unpkg.com/knip@latest/schema.json",
-	"entry": ["src/extension.ts", "src/activate/index.ts", "webview-ui/src/index.tsx"],
+	"entry": ["src/extension.ts", "src/activate/index.ts", "webview-ui/src/index.tsx", "src/api/providers/workers/token-counter.worker.ts"],
 	"project": ["src/**/*.ts", "webview-ui/src/**/*.{ts,tsx}"],
 	"ignore": [
 		"**/__mocks__/**",

src/api/providers/base-provider.ts

@@ -2,18 +2,34 @@ import { Anthropic } from "@anthropic-ai/sdk"
 import { ApiHandler } from ".."
 import { ModelInfo } from "../../shared/api"
 import { ApiStream } from "../transform/stream"
-import { Tiktoken } from "js-tiktoken/lite"
-import o200kBase from "js-tiktoken/ranks/o200k_base"
+import { workerManager } from "../../services/workers/WorkerManager"
 
-// Reuse the fudge factor used in the original code
-const TOKEN_FUDGE_FACTOR = 1.5
+const TOKEN_WORKER_ID = "token-counter"
+const TOKEN_WORKER_PATH = "workers/token-counter.worker.js"
 
 /**
  * Base class for API providers that implements common functionality
  */
 export abstract class BaseProvider implements ApiHandler {
-	// Cache the Tiktoken encoder instance since it's stateless
-	private encoder: Tiktoken | null = null
+	private isDestroyed: boolean = false
+
+	/**
+	 * Class destructor to ensure cleanup
+	 */
+	public async [Symbol.asyncDispose](): Promise<void> {
+		if (!this.isDestroyed) {
+			this.isDestroyed = true
+			await this.cleanup()
+		}
+	}
+
+	/**
+	 * Cleanup resources used by the provider
+	 * This method can be called explicitly or will be called automatically on destruction
+	 */
+	public async cleanup(): Promise<void> {
+		this.isDestroyed = true
+	}
 	abstract createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
 	abstract getModel(): { id: string; info: ModelInfo }
 
@@ -30,35 +46,33 @@ export abstract class BaseProvider implements ApiHandler {
 	async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
 		if (!content || content.length === 0) return 0
 
-		let totalTokens = 0
+		const worker = await workerManager.initializeWorker(TOKEN_WORKER_ID, TOKEN_WORKER_PATH)
 
-		// Lazily create and cache the encoder if it doesn't exist
-		if (!this.encoder) {
-			this.encoder = new Tiktoken(o200kBase)
-		}
+		return new Promise((resolve, reject) => {
+			// Handle worker messages
+			const messageHandler = (result: number | { error: string }) => {
+				worker.removeListener("message", messageHandler)
+				worker.removeListener("error", errorHandler)
 
-		// Process each content block using the cached encoder
-		for (const block of content) {
-			if (block.type === "text") {
-				// Use tiktoken for text token counting
-				const text = block.text || ""
-				if (text.length > 0) {
-					const tokens = this.encoder.encode(text)
-					totalTokens += tokens.length
-				}
-			} else if (block.type === "image") {
-				// For images, calculate based on data size
-				const imageSource = block.source
-				if (imageSource && typeof imageSource === "object" && "data" in imageSource) {
-					const base64Data = imageSource.data as string
-					totalTokens += Math.ceil(Math.sqrt(base64Data.length))
+				if (typeof result === "number") {
+					resolve(result)
 				} else {
-					totalTokens += 300 // Conservative estimate for unknown images
+					reject(new Error(result.error))
 				}
 			}
-		}
 
-		// Add a fudge factor to account for the fact that tiktoken is not always accurate
-		return Math.ceil(totalTokens * TOKEN_FUDGE_FACTOR)
+			// Handle worker errors
+			const errorHandler = (error: Error) => {
+				worker.removeListener("message", messageHandler)
+				worker.removeListener("error", errorHandler)
+				reject(error)
+			}
+
+			worker.once("message", messageHandler)
+			worker.once("error", errorHandler)
+
+			// Send content to worker
+			worker.postMessage(content)
+		})
 	}
 }

src/api/providers/workers/token-counter.worker.ts

@@ -0,0 +1,52 @@
+import { parentPort } from "worker_threads"
+import { Tiktoken } from "js-tiktoken/lite"
+import o200kBase from "js-tiktoken/ranks/o200k_base"
+import { Anthropic } from "@anthropic-ai/sdk"
+
+// Reuse the fudge factor used in the original code
+const TOKEN_FUDGE_FACTOR = 1.5
+
+type ContentBlock = Anthropic.Messages.ContentBlockParam
+
+// Initialize encoder once for reuse
+let encoder: Tiktoken | null = null
+
+parentPort?.on("message", async (content: Array<ContentBlock>) => {
+	try {
+		if (!content || content.length === 0) {
+			parentPort?.postMessage(0)
+			return
+		}
+
+		let totalTokens = 0
+
+		// Lazily create encoder if it doesn't exist
+		if (!encoder) {
+			encoder = new Tiktoken(o200kBase)
+		}
+
+		// Process each content block
+		for (const block of content) {
+			if (block.type === "text") {
+				const text = block.text || ""
+				if (text.length > 0) {
+					const tokens = encoder.encode(text)
+					totalTokens += tokens.length
+				}
+			} else if (block.type === "image") {
+				const imageSource = block.source
+				if (imageSource && typeof imageSource === "object" && "data" in imageSource) {
+					const base64Data = imageSource.data as string
+					totalTokens += Math.ceil(Math.sqrt(base64Data.length))
+				} else {
+					totalTokens += 300 // Conservative estimate for unknown images
+				}
+			}
+		}
+
+		// Apply fudge factor and send result
+		parentPort?.postMessage(Math.ceil(totalTokens * TOKEN_FUDGE_FACTOR))
+	} catch (error) {
+		parentPort?.postMessage({ error: error instanceof Error ? error.message : "Unknown error" })
+	}
+})

src/services/workers/WorkerManager.ts

@@ -0,0 +1,118 @@
+import { Worker } from "worker_threads"
+import path from "path"
+
+interface WorkerConfig {
+	retryAttempts: number
+	maxRetries: number
+	retryDelay: number
+}
+
+export class WorkerManager {
+	private static instance: WorkerManager
+	private workers: Map<string, Worker>
+	private configs: Map<string, WorkerConfig>
+	private isDestroyed: boolean
+
+	private constructor() {
+		this.workers = new Map()
+		this.configs = new Map()
+		this.isDestroyed = false
+	}
+
+	public static getInstance(): WorkerManager {
+		if (!WorkerManager.instance) {
+			WorkerManager.instance = new WorkerManager()
+		}
+		return WorkerManager.instance
+	}
+
+	public async initializeWorker(
+		workerId: string,
+		workerPath: string,
+		config: Partial<WorkerConfig> = {},
+	): Promise<Worker> {
+		if (this.isDestroyed) {
+			throw new Error("WorkerManager has been destroyed")
+		}
+
+		// If worker already exists, return it
+		if (this.workers.has(workerId)) {
+			return this.workers.get(workerId)!
+		}
+
+		const workerConfig: WorkerConfig = {
+			retryAttempts: 0,
+			maxRetries: config.maxRetries ?? 3,
+			retryDelay: config.retryDelay ?? 1000,
+		}
+
+		const absolutePath = path.resolve(__dirname, workerPath)
+		const worker = new Worker(absolutePath)
+
+		worker.on("error", async (error) => {
+			console.error(`Worker ${workerId} error:`, error)
+			await this.handleWorkerError(workerId, workerConfig)
+		})
+
+		worker.on("exit", async (code) => {
+			if (code !== 0) {
+				console.error(`Worker ${workerId} exited with code ${code}`)
+				await this.handleWorkerError(workerId, workerConfig)
+			}
+		})
+
+		this.workers.set(workerId, worker)
+		this.configs.set(workerId, workerConfig)
+
+		return worker
+	}
+
+	private async handleWorkerError(workerId: string, config: WorkerConfig): Promise<void> {
+		if (config.retryAttempts >= config.maxRetries) {
+			console.error(`Max retry attempts reached for worker ${workerId}, worker will not restart`)
+			return
+		}
+
+		config.retryAttempts++
+		const delay = config.retryDelay * Math.pow(2, config.retryAttempts - 1)
+
+		console.log(
+			`Attempting to restart worker ${workerId} (attempt ${config.retryAttempts} of ${config.maxRetries}) after ${delay}ms`,
+		)
+
+		await new Promise((resolve) => setTimeout(resolve, delay))
+
+		const worker = this.workers.get(workerId)
+		if (worker) {
+			await worker.terminate()
+			this.workers.delete(workerId)
+			this.configs.delete(workerId)
+		}
+	}
+
+	public getWorker(workerId: string): Worker | undefined {
+		return this.workers.get(workerId)
+	}
+
+	public async cleanup(): Promise<void> {
+		if (this.isDestroyed) {
+			return
+		}
+
+		this.isDestroyed = true
+
+		const terminationPromises = Array.from(this.workers.values()).map((worker) => worker.terminate())
+
+		await Promise.all(terminationPromises)
+
+		this.workers.clear()
+		this.configs.clear()
+	}
+
+	public async [Symbol.asyncDispose](): Promise<void> {
+		await this.cleanup()
+	}
+}
+
+// Export singleton instance
+export const workerManager = WorkerManager.getInstance()