perf(core): Share pre-compiled tiktoken WASM module across worker threads

[yamadashy]

Each worker thread was independently loading the ~5.3MB tiktoken WASM binary via synchronous readFileSync + new WebAssembly.Module() + new WebAssembly.Instance() at import time. This change compiles the WASM once in the main thread and passes the WebAssembly.Module to workers via workerData (structured clone), so workers only need to instantiate it.

Approach

Uses the tiktoken/init API (an official export in the tiktoken package) for explicit WASM initialization control:

1. Main thread: Reads and compiles the WASM binary once via WebAssembly.compile(), cached in wasmModuleCache.ts
2. Worker creation: The compiled WebAssembly.Module is passed to workers through extraWorkerData (new option on WorkerOptions)
3. Worker init: Each worker calls WebAssembly.instantiate(module, imports) — instantiation only, no recompilation

Fallback: If workerData.tiktokenWasmModule is absent or precompilation fails, workers fall back to reading and compiling from disk (backwards compatible).

Benchmark Results

Worker initialization (3 workers concurrent init, 10 runs, 2 warmup):

	Before	After	Improvement
Median	353ms	252ms	-29%
Avg	363ms	259ms	-29%
Min	312ms	221ms	-29%
P90	529ms	355ms	-33%

Changed Files

File	Change
src/core/metrics/wasmModuleCache.ts	New — Compile and cache WASM module
src/types/webassembly.d.ts	New — WebAssembly type declarations for es2022 target
src/core/metrics/TokenCounter.ts	Switch to tiktoken/init, add initTiktokenWasm()
src/core/metrics/workers/calculateMetricsWorker.ts	Init WASM from pre-compiled module via workerData[1]
src/core/metrics/calculateMetrics.ts	Pre-compile WASM and pass to workers via extraWorkerData
src/shared/processConcurrency.ts	Add extraWorkerData option to WorkerOptions
4 files	tiktoken to tiktoken/init type import change
tests/core/metrics/TokenCounter.test.ts	Update mock target to tiktoken/init

Checklist

• Run npm run test
• Run npm run lint

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 7cbabb6d-fc6e-4a7a-b374-8e45526ce4a3

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR refactors tiktoken imports to use tiktoken/init instead of tiktoken, introduces WebAssembly module precompilation and caching infrastructure, and updates worker thread initialization to accept shared precompiled WASM modules via extended WorkerOptions.

Changes

Cohort / File(s)	Summary
Tiktoken Type Import Updates src/config/configSchema.ts, src/core/metrics/calculateOutputMetrics.ts, src/core/metrics/calculateSelectiveFileMetrics.ts, src/core/metrics/tokenCounterFactory.ts	Updated TiktokenEncoding type import source from tiktoken to tiktoken/init across multiple files; no logic changes.
WASM Infrastructure src/core/metrics/TokenCounter.ts, src/core/metrics/wasmModuleCache.ts, src/types/webassembly.d.ts	Added initTiktokenWasm() helper for WASM initialization with optional precompiled module; introduced WASM caching with lazy compilation and timing; added WebAssembly TypeScript declarations for Module, Instance, and async APIs.
Worker Thread Integration src/core/metrics/calculateMetrics.ts, src/core/metrics/workers/calculateMetricsWorker.ts, src/shared/processConcurrency.ts	Extended WorkerOptions with optional extraWorkerData field; updated calculateMetrics to precompile WASM and pass module to workers; updated worker to initialize tiktoken WASM using shared precompiled module before token counting.
Test Updates tests/core/metrics/TokenCounter.test.ts	Updated test mocks to import from tiktoken/init instead of tiktoken; added mock for init function.

Sequence Diagram(s)

sequenceDiagram
    participant Main as Main Thread
    participant Cache as WASM Cache
    participant Worker as Worker Thread
    participant Tiktoken as Tiktoken

    Main->>Cache: getCompiledTiktokenWasmModule()
    activate Cache
    Cache->>Cache: Read tiktoken_bg.wasm binary
    Cache->>Cache: WebAssembly.compile(binary)
    Cache-->>Main: Compiled Module (cached)
    deactivate Cache

    Main->>Main: createWorkerPool(extraWorkerData: {tiktokenWasmModule})
    Main->>Worker: Initialize with workerData
    
    activate Worker
    Worker->>Worker: initTiktokenWasm(workerData.tiktokenWasmModule)
    Worker->>Tiktoken: init(precompiledModule)
    Tiktoken-->>Worker: WASM initialized
    deactivate Worker

    Main->>Worker: countTokens(text, encoding)
    activate Worker
    Worker->>Worker: Await wasmInitPromise
    Worker->>Tiktoken: get_encoding(encoding)
    Tiktoken-->>Worker: Encoder instance
    Worker-->>Main: Token count
    deactivate Worker

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• repomix#816: Modifies src/shared/processConcurrency.ts for worker data handling; directly related to extraWorkerData propagation changes in this PR.
• repomix#817: Updates metrics worker initialization and worker-pool integration with token-counter cleanup; shares WASM precompilation and worker initialization patterns with this PR.
• repomix#851: Modifies calculateMetrics.ts and processConcurrency.ts worker initialization; directly overlaps with worker data/initialization changes.

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'perf(core): Share pre-compiled tiktoken WASM module across worker threads' is concise, specific, and directly summarizes the main optimization change in the PR.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description check	✅ Passed	The pull request description is comprehensive and well-structured, covering problem statement, approach, benchmark results, file changes, and completed checklist items.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/tiktoken-wasm-module-sharing

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant performance enhancement by optimizing how the tiktoken WebAssembly module is handled across worker threads. Instead of each worker independently compiling the large WASM binary, the main thread now pre-compiles it once and shares the resulting module with workers. This change drastically reduces startup overhead for workers, leading to a substantial speedup in token counting operations by allowing workers to only instantiate the module.

Highlights

• Performance Optimization: Significantly reduced tiktoken WebAssembly (WASM) compilation overhead in worker threads, leading to up to 65% faster average task execution.
• Shared WASM Module: Implemented a mechanism to compile the tiktoken WASM module once in the main thread and share the pre-compiled module with worker threads via workerData.
• Explicit WASM Initialization: Utilized the tiktoken/init API for controlled WASM initialization, allowing workers to instantiate the module directly without recompilation.
• Worker Data Integration: Extended worker thread options to pass the pre-compiled WASM module as extraWorkerData to new workers, enabling efficient module transfer.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a significant performance improvement by pre-compiling the tiktoken WASM module in the main thread and sharing it with worker threads. The implementation is solid, leveraging the tiktoken/init API and a new extraWorkerData option to pass the compiled module. The inclusion of a fallback mechanism and the clear benchmark results are commendable. My review includes one suggestion to enhance type safety in the worker thread by validating data received via workerData instead of using a type assertion.

[gemini-code-assist]

For improved type safety and clarity, it's better to validate the type of workerData.tiktokenWasmModule before using it, instead of relying on a type assertion. workerData can contain arbitrary data, so runtime checks are crucial. An instanceof check ensures the object is a WebAssembly.Module. This change also introduces a temporary variable for readability.

// Initialize tiktoken WASM. If a pre-compiled module is passed from the main thread,
// use it to avoid re-compiling. This is much faster (~6ms vs ~250ms).
const wasmModule = workerData?.tiktokenWasmModule;
const wasmInitPromise = initTiktokenWasm(
  wasmModule instanceof WebAssembly.Module ? wasmModule : undefined,
);

[claude]

Code Review — perf(core): Share pre-compiled tiktoken WASM module across worker threads

Overall: Excellent performance optimization with a clean, well-structured implementation. The 2.9x speedup is impressive and the approach (compile once, share via structured clone) is the right pattern for WebAssembly in worker threads.

Positive Highlights

• Clean separation: wasmModuleCache.ts handles compilation, initTiktokenWasm handles initialization, workers consume via workerData
• Good fallback path — workers still work even without the pre-compiled module
• extraWorkerData is a nice generic extension to WorkerOptions that keeps the concurrency module decoupled from tiktoken specifics

Issues & Suggestions

1. Race condition in WASM module cache (low risk)

In wasmModuleCache.ts, if getCompiledTiktokenWasmModule() is called concurrently before the first call completes, multiple compilations will run in parallel since the compiledModule check and assignment aren't atomic. Currently this is safe because calculateMetrics only calls it once, but this is a latent bug if usage changes.

A simple fix would be to cache the promise instead:

let compilationPromise: Promise<WebAssembly.Module> | null = null;

export const getCompiledTiktokenWasmModule = (): Promise<WebAssembly.Module> => {
  if (!compilationPromise) {
    compilationPromise = compileModule();
  }
  return compilationPromise;
};

2. Duplicated WASM resolution logic

The fallback path in initTiktokenWasm (TokenCounter.ts lines 14-20) duplicates the createRequire + require.resolve pattern from wasmModuleCache.ts. Consider extracting the WASM path resolution into a shared utility, or having the fallback call getCompiledTiktokenWasmModule directly (though that would add a dependency from worker to main-thread code, so the duplication may be intentional).

3. Consider adding "WebWorker" lib to tsconfig instead of custom type declarations

The webassembly.d.ts file (62 lines) provides manual WebAssembly type declarations because the project targets es2022. An alternative is adding "WebWorker" to the lib array in tsconfig.json, which includes the standard WebAssembly namespace types. This would eliminate the maintenance burden of keeping custom declarations in sync. If there's a reason WebWorker lib was avoided (e.g., it pulls in unwanted browser globals), a comment explaining the choice would help.

4. No test coverage for initTiktokenWasm

The initTiktokenWasm function is a new public export but has no unit test. The existing TokenCounter.test.ts mocks init as vi.fn() but doesn't verify the initialization logic (pre-compiled module path vs fallback path). Consider adding tests that verify:

• When a WebAssembly.Module is provided, init is called with an instantiation callback
• When no module is provided, it falls back to reading from disk

Premortem Analysis

Scenario	Risk	Mitigation
tiktoken/init API changes in a future version	Medium	The tiktoken/init export is an official API; pin tiktoken version in lockfile
WebAssembly.Module structured clone fails in edge Node.js versions	Low	Supported since Node 12; the fallback path handles this gracefully
child_process runtime used with extraWorkerData containing non-serializable data	Low	child_process doesn't use workerData at all in this codebase — but extraWorkerData silently does nothing for child_process runtime. Consider documenting this limitation.
WASM compilation fails (corrupted binary, OOM)	Low	Error propagates to calculateMetrics caller — appropriate behavior, no silent failure

Verdict: This is a solid, well-thought-out optimization. The suggestions above are minor improvements — none are blockers.

---

Reviewed by Claude

[codecov]

Codecov Report

❌ Patch coverage is 51.42857% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.89%. Comparing base (f0393f7) to head (6a7a4a4).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
src/core/metrics/wasmModuleCache.ts	20.00%	12 Missing ⚠️
src/core/metrics/calculateMetrics.ts	20.00%	4 Missing ⚠️
src/core/metrics/TokenCounter.ts	90.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1243      +/-   ##
==========================================
- Coverage   87.18%   86.89%   -0.30%     
==========================================
  Files         115      116       +1     
  Lines        4324     4358      +34     
  Branches     1002     1009       +7     
==========================================
+ Hits         3770     3787      +17     
- Misses        554      571      +17     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[cloudflare-workers-and-pages]

Deploying repomix with    Cloudflare Pages

Latest commit:	6a7a4a4
Status:	✅  Deploy successful!
Preview URL:	https://8249ea27.repomix.pages.dev
Branch Preview URL:	https://perf-tiktoken-wasm-module-sh.repomix.pages.dev

View logs

[claude]

Code Review (Updated) — perf(core): Share pre-compiled tiktoken WASM module across worker threads

Overall: This is a well-executed performance optimization. The second commit (c8dc1e6) properly addresses the earlier feedback from Devin (workerData wrapping), Gemini (instanceof check), and CodeRabbit (conditional precompile + try/catch). The implementation is clean and the ~29% median speedup is a meaningful improvement.

What looks good

• The !deps.taskRunner guard correctly avoids unnecessary WASM compilation when a test double is injected
• try/catch around getCompiledTiktokenWasmModule() ensures graceful fallback — workers will still function
• instanceof WebAssembly.Module runtime check in the worker is defensive and appropriate
• Tinypool workerData wrapping (workerData[1]) is correctly handled with an Array.isArray guard

Remaining items from previous review (still applicable)

1. Race condition in wasmModuleCache (low priority — nitpick)

The getCompiledTiktokenWasmModule function still has a theoretical race if called concurrently — multiple compilations could run before compiledModule is assigned. Both CodeRabbit and the previous Claude review flagged this. Currently safe since there's only one call site, but caching the promise instead of the result would be a trivial fix:

let compilationPromise: Promise<WebAssembly.Module> | null = null;

export const getCompiledTiktokenWasmModule = (): Promise<WebAssembly.Module> => {
  if (!compilationPromise) {
    compilationPromise = (async () => {
      const startTime = process.hrtime.bigint();
      const wasmPath = getTiktokenWasmPath();
      const wasmBinary = await fs.readFile(wasmPath);
      const module = await WebAssembly.compile(wasmBinary);
      const compileTime = Number(process.hrtime.bigint() - startTime) / 1e6;
      logger.debug(`Tiktoken WASM compilation took ${compileTime.toFixed(2)}ms`);
      return module;
    })();
  }
  return compilationPromise;
};

2. extraWorkerData can override reserved workerData keys

CodeRabbit flagged that ...extraWorkerData in processConcurrency.ts:88 is spread after workerType and logLevel, so a caller could accidentally shadow them. A simple reorder (spread extraWorkerData first, then the reserved keys) would prevent this without any API change:

workerData: {
  ...extraWorkerData,
  workerType,
  logLevel: logger.getLogLevel(),
},

This is low risk since the only current caller passes { tiktokenWasmModule }, but it's a good defensive practice for a generic extraWorkerData API.

3. initTiktokenWasm fallback duplicates WASM resolution

The fallback path in TokenCounter.ts:14-20 duplicates createRequire + require.resolve from wasmModuleCache.ts. This is minor since the fallback is a safety net, but extracting a shared getTiktokenWasmPath() utility would reduce duplication. Understandable if kept as-is to avoid cross-module coupling in the worker.

New observations on the fix commit

4. wasmInitPromise error handling in worker

In calculateMetricsWorker.ts:27, wasmInitPromise is created at module load time. If initTiktokenWasm rejects (e.g., corrupted WASM), the rejection is stored in the promise and re-thrown on every countTokens call via await wasmInitPromise (line 37). This is correct behavior — each task will fail with a clear error. However, since initTiktokenWasm does file I/O in the fallback path, consider whether an unhandled rejection warning could fire before the first countTokens call. In Node.js, a promise rejection is "unhandled" if no handler is attached by the end of the microtask queue. Since wasmInitPromise is awaited inside countTokens (which is called shortly after), this should be fine in practice.

Premortem (updated)

Scenario	Risk	Status
Concurrent WASM compilation race	Low	Not fixed, but safe with current single call site
child_process runtime ignores extraWorkerData silently	Low	calculateMetrics hardcodes worker_threads — no issue in practice
tiktoken/init API breakage	Medium	Standard tiktoken export; mitigated by lockfile
Worker WASM init failure	Low	Graceful — falls back to per-worker compilation

Verdict: The fix commit properly addressed the critical issues. The remaining suggestions are minor hardening. This PR is ready to merge.

---

Reviewed by Claude

[yamadashy]

• refactor(metrics): Replace tiktoken with gpt-tokenizer #1245

[claude]

Code Review (Third Pass) — perf(core): Share pre-compiled tiktoken WASM module across worker threads

Overall: Excellent performance optimization. The code at 6a7a4a4 correctly addresses all critical feedback from Devin, Gemini, and CodeRabbit. Two prior Claude reviews covered this thoroughly — no new issues found. The remaining suggestions from those reviews are minor hardening, not blockers.

Summary of outstanding suggestions (all low priority, non-blocking)

#	Suggestion	Priority	Status
1	Cache the promise in wasmModuleCache.ts instead of the result to prevent theoretical race condition	Low	Not addressed — safe with single call site
2	Spread extraWorkerData before reserved keys in processConcurrency.ts:85-88 to prevent accidental override	Low	Not addressed — safe with current callers
3	Extract shared getTiktokenWasmPath() to reduce duplication between wasmModuleCache.ts and TokenCounter.ts	Low	Intentional — avoids cross-module coupling in worker
4	Add unit tests for initTiktokenWasm (both paths) and wasmModuleCache — patch coverage is 51%	Recommended	Not addressed

Verdict

No new issues beyond what was previously flagged. The PR is well-structured with proper fallbacks, defensive instanceof checks, and correct Tinypool workerData handling. Ready to merge — the outstanding items can be addressed in follow-up if desired.

---

Reviewed by Claude

[github-actions]

⚡ Performance Benchmark

Packing the repomix repository with node bin/repomix.cjs

Runner	PR	main	Diff
Ubuntu	2.46s (±0.04s)	2.47s (±0.01s)	-0.01s (-0.3%)
macOS	1.89s (±0.37s)	1.84s (±0.36s)	+0.05s (+2.7%)
Windows	3.16s (±0.14s)	3.07s (±0.04s)	+0.09s (+3.0%)

Details

• Warmup: 2 runs (discarded)
• Measurement: 10 runs (median)
• ±: IQR (Interquartile Range) — middle 50% of measurements spread
• Workflow run

[yamadashy]

Closing this PR

CI ベンチマーク結果で改善が確認できませんでした：

Runner	Diff
Ubuntu	-0.3% (誤差)
macOS	+2.7% (微劣化)
Windows	+3.0% (微劣化)

WASM モジュール共有による Worker 初期化の高速化は理論的には有効ですが、E2E では改善が見られず、むしろ微劣化している環境もあります。現状の tiktoken の動作で問題がないため、この最適化は見送ります。

[yamadashy]

(Updated comment in English)

Closing this PR

CI benchmark results showed no improvement:

Runner	Diff
Ubuntu	-0.3% (within noise)
macOS	+2.7% (slight regression)
Windows	+3.0% (slight regression)

While sharing the pre-compiled WASM module across workers is theoretically sound for reducing initialization time, the E2E benchmarks show no measurable improvement — and slight regressions on macOS/Windows. Since the current tiktoken setup works well without issues, this optimization is not worth the added complexity.