@kt3k/tku 1.0.9 → 1.0.11

package/dist/main.js

@@ -75,35 +75,61 @@ async function listTextFiles(repoPath, options = {}) {
 }
 
 // src/tokenize.ts
-import { readFile } from "node:fs/promises";
 import { resolve as resolve2 } from "node:path";
+import { availableParallelism } from "node:os";
+import { Worker } from "node:worker_threads";
 import { get_encoding } from "tiktoken";
+var workerExt = import.meta.url.endsWith(".ts") ? ".ts" : ".js";
+var workerUrl = new URL(`./tokenize-worker${workerExt}`, import.meta.url);
 async function tokenizeFiles(repoPath, files, encoding, options = {}) {
-  const enc = get_encoding(encoding);
-  try {
-    const results = [];
-    let totalTokens = 0;
-    for (let i = 0; i < files.length; i++) {
-      const file = files[i];
-      options.onProgress?.(file, i + 1, files.length);
-      const fullPath = resolve2(repoPath, file);
-      const content = await readFile(fullPath, "utf-8");
-      const tokens = enc.encode_ordinary(content).length;
-      results.push({
-        path: file,
-        tokens
-      });
-      totalTokens += tokens;
-    }
-    return {
-      encoding,
-      files: results,
-      totalTokens,
-      totalFiles: results.length
-    };
-  } finally {
-    enc.free();
+  const resolvedRepoPath = resolve2(repoPath);
+  const numWorkers = Math.min(availableParallelism(), files.length || 1);
+  const chunks = Array.from({
+    length: numWorkers
+  }, () => []);
+  for (let i = 0; i < files.length; i++) {
+    chunks[i % numWorkers].push(files[i]);
   }
+  let progressCount = 0;
+  const total = files.length;
+  const workerPromises = chunks.map((chunk) => {
+    if (chunk.length === 0) return Promise.resolve([]);
+    return new Promise((resolvePromise, reject) => {
+      const worker = new Worker(workerUrl);
+      worker.on("message", (msg) => {
+        if (msg.type === "progress") {
+          progressCount++;
+          options.onProgress?.(msg.file, progressCount, total);
+        } else if (msg.type === "done") {
+          resolvePromise(msg.results);
+          worker.terminate();
+        }
+      });
+      worker.on("error", (err) => {
+        reject(err);
+        worker.terminate();
+      });
+      worker.postMessage({
+        repoPath: resolvedRepoPath,
+        files: chunk,
+        encoding
+      });
+    });
+  });
+  const chunkResults = await Promise.all(workerPromises);
+  const results = chunkResults.flat();
+  const orderMap = new Map(files.map((f, i) => [
+    f,
+    i
+  ]));
+  results.sort((a, b) => orderMap.get(a.path) - orderMap.get(b.path));
+  const totalTokens = results.reduce((sum, f) => sum + f.tokens, 0);
+  return {
+    encoding,
+    files: results,
+    totalTokens,
+    totalFiles: results.length
+  };
 }
 
 // src/format.ts
@@ -214,7 +240,9 @@ async function main() {
   const isTTY = process.stderr.isTTY;
   function status(msg) {
     if (isTTY) {
-      process.stderr.write(`\r\x1B[K${msg}`);
+      const cols = process.stderr.columns || 80;
+      const truncated = msg.length > cols ? msg.slice(0, cols) : msg;
+      process.stderr.write(`\r\x1B[K${truncated}`);
     }
   }
   function clearStatus() {

package/dist/tokenize-worker.js

@@ -0,0 +1,30 @@
+// src/tokenize-worker.ts
+import { parentPort } from "node:worker_threads";
+import { readFile } from "node:fs/promises";
+import { resolve } from "node:path";
+import { get_encoding } from "tiktoken";
+parentPort.on("message", async (msg) => {
+  const enc = get_encoding(msg.encoding);
+  try {
+    const results = [];
+    for (const file of msg.files) {
+      const fullPath = resolve(msg.repoPath, file);
+      const content = await readFile(fullPath, "utf-8");
+      const tokens = enc.encode_ordinary(content).length;
+      results.push({
+        path: file,
+        tokens
+      });
+      parentPort.postMessage({
+        type: "progress",
+        file
+      });
+    }
+    parentPort.postMessage({
+      type: "done",
+      results
+    });
+  } finally {
+    enc.free();
+  }
+});

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@kt3k/tku",
-  "version": "1.0.9",
+  "version": "1.0.11",
   "description": "",
   "main": "dist/main.js",
   "bin": {
@@ -10,8 +10,8 @@
     "dist"
   ],
   "scripts": {
-    "build": "deno bundle --external tiktoken --external picomatch -o dist/main.js src/main.ts && chmod +x dist/main.js",
-    "prepublish": "npm run build",
+    "build": "deno bundle --external tiktoken --external picomatch -o dist/main.js src/main.ts && deno bundle --external tiktoken -o dist/tokenize-worker.js src/tokenize-worker.ts && chmod +x dist/main.js",
+    "prepublishOnly": "npm run build",
     "test": "vitest run"
   },
   "keywords": [],