refactor: replace Piscina with native worker_threads

Removes the piscina dependency in favor of a custom worker pool
implementation using Node.js built-in worker_threads. This reduces
external dependencies while maintaining the same async tokenization
functionality.

Changes:
- Added workerPool.ts to manage worker thread lifecycle and message
  passing
- Updated tokenizer.worker.ts to handle parentPort messages directly
- Modified tokenizer.ts to use new run() function instead of Piscina
- Migrated unit tests from Jest to bun test for consistency
- Updated CI and Makefile to use bun test for coverage

The custom pool creates a single persistent worker at module load
time and handles request/response matching via message IDs.

Author: ThomasK33

.github/workflows/ci.yml

@@ -75,7 +75,7 @@ jobs:
       - uses: ./.github/actions/setup-cmux
 
       - name: Run tests with coverage
-        run: bun x jest --coverage --coverageReporters=lcov ${{ github.event.inputs.test_filter || 'src' }}
+        run: bun test --coverage --coverage-reporter=lcov ${{ github.event.inputs.test_filter || 'src' }}
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5

Makefile

@@ -191,11 +191,11 @@ check-deadcode: node_modules/.installed ## Check for potential dead code (manual
 
 ## Testing
 test-integration: node_modules/.installed ## Run all tests (unit + integration)
-	@bun x jest src
+	@bun test src
 	@TEST_INTEGRATION=1 bun x jest tests
 
-test-unit: node_modules/.installed build-main ## Run unit tests (requires build-main for worker compilation)
-	@bun x jest src
+test-unit: node_modules/.installed ## Run unit tests
+	@bun test src
 
 test: test-unit ## Alias for test-unit
 

bun.lock

@@ -28,7 +28,6 @@
         "lru-cache": "^11.2.2",
         "markdown-it": "^14.1.0",
         "minimist": "^1.2.8",
-        "piscina": "^5.1.3",
         "rehype-harden": "^1.1.5",
         "shescape": "^2.1.6",
         "source-map-support": "^0.5.21",
@@ -395,42 +394,6 @@
 
     "@mermaid-js/parser": ["@mermaid-js/[email protected]", "", { "dependencies": { "langium": "3.3.1" } }, "sha512-lnjOhe7zyHjc+If7yT4zoedx2vo4sHaTmtkl1+or8BRTnCtDmcTpAjpzDSfCZrshM5bCoz0GyidzadJAH1xobA=="],
 
-    "@napi-rs/nice": ["@napi-rs/[email protected]", "", { "optionalDependencies": { "@napi-rs/nice-android-arm-eabi": "1.1.1", "@napi-rs/nice-android-arm64": "1.1.1", "@napi-rs/nice-darwin-arm64": "1.1.1", "@napi-rs/nice-darwin-x64": "1.1.1", "@napi-rs/nice-freebsd-x64": "1.1.1", "@napi-rs/nice-linux-arm-gnueabihf": "1.1.1", "@napi-rs/nice-linux-arm64-gnu": "1.1.1", "@napi-rs/nice-linux-arm64-musl": "1.1.1", "@napi-rs/nice-linux-ppc64-gnu": "1.1.1", "@napi-rs/nice-linux-riscv64-gnu": "1.1.1", "@napi-rs/nice-linux-s390x-gnu": "1.1.1", "@napi-rs/nice-linux-x64-gnu": "1.1.1", "@napi-rs/nice-linux-x64-musl": "1.1.1", "@napi-rs/nice-openharmony-arm64": "1.1.1", "@napi-rs/nice-win32-arm64-msvc": "1.1.1", "@napi-rs/nice-win32-ia32-msvc": "1.1.1", "@napi-rs/nice-win32-x64-msvc": "1.1.1" } }, "sha512-xJIPs+bYuc9ASBl+cvGsKbGrJmS6fAKaSZCnT0lhahT5rhA2VVy9/EcIgd2JhtEuFOJNx7UHNn/qiTPTY4nrQw=="],
-
-    "@napi-rs/nice-android-arm-eabi": ["@napi-rs/[email protected]", "", { "os": "android", "cpu": "arm" }, "sha512-kjirL3N6TnRPv5iuHw36wnucNqXAO46dzK9oPb0wj076R5Xm8PfUVA9nAFB5ZNMmfJQJVKACAPd/Z2KYMppthw=="],
-
-    "@napi-rs/nice-android-arm64": ["@napi-rs/[email protected]", "", { "os": "android", "cpu": "arm64" }, "sha512-blG0i7dXgbInN5urONoUCNf+DUEAavRffrO7fZSeoRMJc5qD+BJeNcpr54msPF6qfDD6kzs9AQJogZvT2KD5nw=="],
-
-    "@napi-rs/nice-darwin-arm64": ["@napi-rs/[email protected]", "", { "os": "darwin", "cpu": "arm64" }, "sha512-s/E7w45NaLqTGuOjC2p96pct4jRfo61xb9bU1unM/MJ/RFkKlJyJDx7OJI/O0ll/hrfpqKopuAFDV8yo0hfT7A=="],
-
-    "@napi-rs/nice-darwin-x64": ["@napi-rs/[email protected]", "", { "os": "darwin", "cpu": "x64" }, "sha512-dGoEBnVpsdcC+oHHmW1LRK5eiyzLwdgNQq3BmZIav+9/5WTZwBYX7r5ZkQC07Nxd3KHOCkgbHSh4wPkH1N1LiQ=="],
-
-    "@napi-rs/nice-freebsd-x64": ["@napi-rs/[email protected]", "", { "os": "freebsd", "cpu": "x64" }, "sha512-kHv4kEHAylMYmlNwcQcDtXjklYp4FCf0b05E+0h6nDHsZ+F0bDe04U/tXNOqrx5CmIAth4vwfkjjUmp4c4JktQ=="],
-
-    "@napi-rs/nice-linux-arm-gnueabihf": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "arm" }, "sha512-E1t7K0efyKXZDoZg1LzCOLxgolxV58HCkaEkEvIYQx12ht2pa8hoBo+4OB3qh7e+QiBlp1SRf+voWUZFxyhyqg=="],
-
-    "@napi-rs/nice-linux-arm64-gnu": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "arm64" }, "sha512-CIKLA12DTIZlmTaaKhQP88R3Xao+gyJxNWEn04wZwC2wmRapNnxCUZkVwggInMJvtVElA+D4ZzOU5sX4jV+SmQ=="],
-
-    "@napi-rs/nice-linux-arm64-musl": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "arm64" }, "sha512-+2Rzdb3nTIYZ0YJF43qf2twhqOCkiSrHx2Pg6DJaCPYhhaxbLcdlV8hCRMHghQ+EtZQWGNcS2xF4KxBhSGeutg=="],
-
-    "@napi-rs/nice-linux-ppc64-gnu": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "ppc64" }, "sha512-4FS8oc0GeHpwvv4tKciKkw3Y4jKsL7FRhaOeiPei0X9T4Jd619wHNe4xCLmN2EMgZoeGg+Q7GY7BsvwKpL22Tg=="],
-
-    "@napi-rs/nice-linux-riscv64-gnu": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "none" }, "sha512-HU0nw9uD4FO/oGCCk409tCi5IzIZpH2agE6nN4fqpwVlCn5BOq0MS1dXGjXaG17JaAvrlpV5ZeyZwSon10XOXw=="],
-
-    "@napi-rs/nice-linux-s390x-gnu": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "s390x" }, "sha512-2YqKJWWl24EwrX0DzCQgPLKQBxYDdBxOHot1KWEq7aY2uYeX+Uvtv4I8xFVVygJDgf6/92h9N3Y43WPx8+PAgQ=="],
-
-    "@napi-rs/nice-linux-x64-gnu": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "x64" }, "sha512-/gaNz3R92t+dcrfCw/96pDopcmec7oCcAQ3l/M+Zxr82KT4DljD37CpgrnXV+pJC263JkW572pdbP3hP+KjcIg=="],
-
-    "@napi-rs/nice-linux-x64-musl": ["@napi-rs/[email protected]", "", { "os": "linux", "cpu": "x64" }, "sha512-xScCGnyj/oppsNPMnevsBe3pvNaoK7FGvMjT35riz9YdhB2WtTG47ZlbxtOLpjeO9SqqQ2J2igCmz6IJOD5JYw=="],
-
-    "@napi-rs/nice-openharmony-arm64": ["@napi-rs/[email protected]", "", { "os": "none", "cpu": "arm64" }, "sha512-6uJPRVwVCLDeoOaNyeiW0gp2kFIM4r7PL2MczdZQHkFi9gVlgm+Vn+V6nTWRcu856mJ2WjYJiumEajfSm7arPQ=="],
-
-    "@napi-rs/nice-win32-arm64-msvc": ["@napi-rs/[email protected]", "", { "os": "win32", "cpu": "arm64" }, "sha512-uoTb4eAvM5B2aj/z8j+Nv8OttPf2m+HVx3UjA5jcFxASvNhQriyCQF1OB1lHL43ZhW+VwZlgvjmP5qF3+59atA=="],
-
-    "@napi-rs/nice-win32-ia32-msvc": ["@napi-rs/[email protected]", "", { "os": "win32", "cpu": "ia32" }, "sha512-CNQqlQT9MwuCsg1Vd/oKXiuH+TcsSPJmlAFc5frFyX/KkOh0UpBLEj7aoY656d5UKZQMQFP7vJNa1DNUNORvug=="],
-
-    "@napi-rs/nice-win32-x64-msvc": ["@napi-rs/[email protected]", "", { "os": "win32", "cpu": "x64" }, "sha512-vB+4G/jBQCAh0jelMTY3+kgFy00Hlx2f2/1zjMoH821IbplbWZOkLiTYXQkygNTzQJTq5cvwBDgn2ppHD+bglQ=="],
-
     "@napi-rs/wasm-runtime": ["@napi-rs/[email protected]", "", { "dependencies": { "@emnapi/core": "^1.4.3", "@emnapi/runtime": "^1.4.3", "@tybys/wasm-util": "^0.10.0" } }, "sha512-ZVWUcfwY4E/yPitQJl481FjFo3K22D6qF0DuFH6Y/nbnE11GY5uguDxZMGXPQ8WQ0128MXQD7TnfHyK4oWoIJQ=="],
 
     "@nodelib/fs.scandir": ["@nodelib/[email protected]", "", { "dependencies": { "@nodelib/fs.stat": "2.0.5", "run-parallel": "^1.1.9" } }, "sha512-vq24Bq3ym5HEQm2NKCr3yXDwjc7vTsEThRDnkp2DK9p1uqLR+DHurm/NOTo0KG7HYHU7eppKZj3MyqYuMBf62g=="],
@@ -2341,8 +2304,6 @@
 
     "pirates": ["[email protected]", "", {}, "sha512-TfySrs/5nm8fQJDcBDuUng3VOUKsd7S+zqvbOTiGXHfxX4wK31ard+hoNuvkicM/2YFzlpDgABOevKSsB4G/FA=="],
 
-    "piscina": ["[email protected]", "", { "optionalDependencies": { "@napi-rs/nice": "^1.0.4" } }, "sha512-0u3N7H4+hbr40KjuVn2uNhOcthu/9usKhnw5vT3J7ply79v3D3M8naI00el9Klcy16x557VsEkkUQaHCWFXC/g=="],
-
     "pkg-dir": ["[email protected]", "", { "dependencies": { "find-up": "^4.0.0" } }, "sha512-HRDzbaKjC+AOWVXxAU/x54COGeIv9eb+6CkDSQoNTt4XyWoIJvuPsXizxu/Fr23EiekbtZwmh1IcIG/l/a10GQ=="],
 
     "pkg-types": ["[email protected]", "", { "dependencies": { "confbox": "^0.2.2", "exsolve": "^1.0.7", "pathe": "^2.0.3" } }, "sha512-SIqCzDRg0s9npO5XQ3tNZioRY1uK06lA41ynBC1YmFTmnY6FjUjVt6s4LoADmwoig1qqD0oK8h1p/8mlMx8Oig=="],

package.json

@@ -69,7 +69,6 @@
     "lru-cache": "^11.2.2",
     "markdown-it": "^14.1.0",
     "minimist": "^1.2.8",
-    "piscina": "^5.1.3",
     "rehype-harden": "^1.1.5",
     "shescape": "^2.1.6",
     "source-map-support": "^0.5.21",

src/utils/main/tokenizer.ts

@@ -2,14 +2,13 @@ import assert from "@/utils/assert";
 import CRC32 from "crc-32";
 import { LRUCache } from "lru-cache";
 import { getAvailableTools, getToolSchemas } from "@/utils/tools/toolDefinitions";
-import Piscina from "piscina";
-import { createRequire } from "module";
 import type { CountTokensInput } from "./tokenizer.worker";
 import { models, type ModelName } from "ai-tokenizer";
+import { run } from "./workerPool";
 
 /**
  * Public tokenizer interface exposed to callers.
- * countTokens is async because the heavy lifting happens in a Piscina worker.
+ * countTokens is async because the heavy lifting happens in a worker thread.
  */
 export interface Tokenizer {
   encoding: string;
@@ -26,13 +25,6 @@ const DEFAULT_WARM_MODELS = [
   "anthropic:claude-sonnet-4-5",
 ] as const;
 
-const requireForResolve = createRequire(__filename);
-const workerPath = requireForResolve.resolve("./tokenizer.worker");
-const tokenizerPool = new Piscina({
-  filename: workerPath,
-  idleTimeout: Infinity,
-});
-
 const encodingPromises = new Map<ModelName, Promise<string>>();
 const inFlightCounts = new Map<string, Promise<number>>();
 const tokenCountCache = new LRUCache<string, number>({
@@ -59,8 +51,7 @@ function normalizeModelKey(modelName: string): ModelName {
 function resolveEncoding(modelName: ModelName): Promise<string> {
   let promise = encodingPromises.get(modelName);
   if (!promise) {
-    promise = tokenizerPool
-      .run(modelName, { name: "encodingName" })
+    promise = run<string>("encodingName", modelName)
       .then((result: unknown) => {
         assert(
           typeof result === "string" && result.length > 0,
@@ -97,8 +88,7 @@ async function countTokensInternal(modelName: ModelName, text: string): Promise<
   let pending = inFlightCounts.get(key);
   if (!pending) {
     const payload: CountTokensInput = { modelName, input: text };
-    pending = tokenizerPool
-      .run(payload, { name: "countTokens" })
+    pending = run<number>("countTokens", payload)
       .then((value: unknown) => {
         assert(
           typeof value === "number" && Number.isFinite(value) && value >= 0,

src/utils/main/tokenizer.worker.ts

@@ -1,4 +1,5 @@
 import assert from "node:assert";
+import { parentPort } from "node:worker_threads";
 import { Tokenizer, models } from "ai-tokenizer";
 import type { ModelName } from "ai-tokenizer";
 import * as encoding from "ai-tokenizer/encoding";
@@ -38,3 +39,36 @@ export function encodingName(modelName: ModelName): string {
   assert(model, `Unknown tokenizer model '${modelName}'`);
   return model.encoding;
 }
+
+// Handle messages from main thread
+if (parentPort) {
+  parentPort.on("message", (message: { messageId: number; taskName: string; data: unknown }) => {
+    try {
+      let result: unknown;
+
+      switch (message.taskName) {
+        case "countTokens":
+          result = countTokens(message.data as CountTokensInput);
+          break;
+        case "encodingName":
+          result = encodingName(message.data as ModelName);
+          break;
+        default:
+          throw new Error(`Unknown task: ${message.taskName}`);
+      }
+
+      parentPort!.postMessage({
+        messageId: message.messageId,
+        result,
+      });
+    } catch (error) {
+      parentPort!.postMessage({
+        messageId: message.messageId,
+        error: {
+          message: error instanceof Error ? error.message : String(error),
+          stack: error instanceof Error ? error.stack : undefined,
+        },
+      });
+    }
+  });
+}

src/utils/main/workerPool.ts

@@ -0,0 +1,95 @@
+import { Worker } from "node:worker_threads";
+import { createRequire } from "node:module";
+
+interface WorkerRequest {
+  messageId: number;
+  taskName: string;
+  data: unknown;
+}
+
+interface WorkerSuccessResponse {
+  messageId: number;
+  result: unknown;
+}
+
+interface WorkerErrorResponse {
+  messageId: number;
+  error: {
+    message: string;
+    stack?: string;
+  };
+}
+
+type WorkerResponse = WorkerSuccessResponse | WorkerErrorResponse;
+
+let messageIdCounter = 0;
+const pendingPromises = new Map<
+  number,
+  { resolve: (value: unknown) => void; reject: (error: Error) => void }
+>();
+
+// Create the worker at module load time
+// This is safe because the heavy ai-tokenizer imports happen in the worker thread, not the main thread
+const requireForResolve = createRequire(__filename);
+const workerPath = requireForResolve.resolve("./tokenizer.worker");
+const worker = new Worker(workerPath);
+
+// Handle messages from worker
+worker.on("message", (response: WorkerResponse) => {
+  const pending = pendingPromises.get(response.messageId);
+  if (!pending) {
+    console.error(`[workerPool] No pending promise for messageId ${response.messageId}`);
+    return;
+  }
+
+  pendingPromises.delete(response.messageId);
+
+  if ("error" in response) {
+    const error = new Error(response.error.message);
+    error.stack = response.error.stack;
+    pending.reject(error);
+  } else {
+    pending.resolve(response.result);
+  }
+});
+
+// Handle worker errors
+worker.on("error", (error) => {
+  console.error("[workerPool] Worker error:", error);
+  // Reject all pending promises
+  for (const pending of pendingPromises.values()) {
+    pending.reject(error);
+  }
+  pendingPromises.clear();
+});
+
+// Handle worker exit
+worker.on("exit", (code) => {
+  if (code !== 0) {
+    console.error(`[workerPool] Worker stopped with exit code ${code}`);
+    const error = new Error(`Worker stopped with exit code ${code}`);
+    for (const pending of pendingPromises.values()) {
+      pending.reject(error);
+    }
+    pendingPromises.clear();
+  }
+});
+
+/**
+ * Run a task on the worker thread
+ * @param taskName The name of the task to run (e.g., "countTokens", "encodingName")
+ * @param data The data to pass to the task
+ * @returns A promise that resolves with the task result
+ */
+export function run<T>(taskName: string, data: unknown): Promise<T> {
+  const messageId = messageIdCounter++;
+  const request: WorkerRequest = { messageId, taskName, data };
+
+  return new Promise<T>((resolve, reject) => {
+    pendingPromises.set(messageId, {
+      resolve: resolve as (value: unknown) => void,
+      reject,
+    });
+    worker.postMessage(request);
+  });
+}