🤖 perf: speed up CI with approx tokenizer and parallel static-check (#962)

Speeds up CI by:

1. **Approximate tokenizer for tests** - Defaults Jest to `length/4`
token counting instead of loading the 10s WASM tokenizer
2. **Parallel static-check** - Removes `-j3` limit, allowing all 6
subtasks to run in parallel
3. **Shallow checkout for tests** - Test jobs don't need full git
history

## Changes

### Tokenizer (`src/node/utils/main/tokenizer.ts`)
- Add `shouldUseApproxTokenizer()` that checks `MUX_APPROX_TOKENIZER=1`
- Short-circuit all tokenizer functions when in approx mode
- Override with `MUX_FORCE_REAL_TOKENIZER=1` for tests needing real
tokenization

### Test setup (`tests/setup.ts`)
- Default `MUX_APPROX_TOKENIZER=1` in Jest

### CI (`.github/workflows/pr.yml`)
- `make -j static-check` (was `-j3`)
- Remove `fetch-depth: 0` from test-unit, test-integration,
test-storybook, test-e2e

### Test fix (`src/cli/run.test.ts`)
- Update version test to accept both `vX.Y.Z` and hash-only formats
(shallow clones don't have tags)

_Generated with `mux`_

Author: ammar-agent

.github/workflows/pr.yml

@@ -77,7 +77,7 @@ jobs:
             experimental-features = nix-command flakes
       - run: curl -LsSf https://astral.sh/uv/install.sh | sh
       - run: echo "$HOME/.local/bin" >> $GITHUB_PATH
-      - run: make -j3 static-check
+      - run: make -j static-check
 
   test-unit:
     name: Test / Unit
@@ -86,8 +86,6 @@ jobs:
     runs-on: ${{ github.repository_owner == 'coder' && 'depot-ubuntu-22.04-16' || 'ubuntu-latest' }}
     steps:
       - uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
       - uses: ./.github/actions/setup-mux
       - run: make build-main
       - run: bun test --coverage --coverage-reporter=lcov ${{ github.event.inputs.test_filter || 'src' }}
@@ -106,8 +104,6 @@ jobs:
     runs-on: ${{ github.repository_owner == 'coder' && 'depot-ubuntu-22.04-16' || 'ubuntu-latest' }}
     steps:
       - uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
       - uses: ./.github/actions/setup-mux
       - run: make build-main
       - name: Run integration tests
@@ -146,8 +142,6 @@ jobs:
     runs-on: ${{ github.repository_owner == 'coder' && 'depot-ubuntu-22.04-16' || 'ubuntu-latest' }}
     steps:
       - uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
       - uses: ./.github/actions/setup-mux
       - uses: ./.github/actions/setup-playwright
       - run: make storybook-build
@@ -171,8 +165,6 @@ jobs:
     runs-on: ${{ matrix.runner }}
     steps:
       - uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
       - uses: ./.github/actions/setup-mux
       - name: Install xvfb
         if: matrix.os == 'linux'

src/cli/run.test.ts

@@ -102,8 +102,8 @@ describe("mux CLI", () => {
     test("--version shows version info", async () => {
       const result = await runCli(["--version"]);
       expect(result.exitCode).toBe(0);
-      // Version format: vX.Y.Z-N-gHASH (HASH)
-      expect(result.stdout).toMatch(/v\d+\.\d+\.\d+/);
+      // Version format: vX.Y.Z-N-gHASH (HASH) or just HASH (HASH) in shallow clones
+      expect(result.stdout).toMatch(/v\d+\.\d+\.\d+|^[0-9a-f]{7,}/);
     });
 
     test("unknown command shows error", async () => {

src/node/utils/main/tokenizer.ts

@@ -18,6 +18,31 @@ export interface Tokenizer {
   countTokens: (text: string) => Promise<number>;
 }
 
+const APPROX_ENCODING = "approx-4";
+
+function shouldUseApproxTokenizer(): boolean {
+  // MUX_FORCE_REAL_TOKENIZER=1 overrides approx mode (for tests that need real tokenization)
+  // MUX_APPROX_TOKENIZER=1 enables fast approximate mode (default in Jest)
+  if (process.env.MUX_FORCE_REAL_TOKENIZER === "1") {
+    return false;
+  }
+  return process.env.MUX_APPROX_TOKENIZER === "1";
+}
+
+function approximateCount(text: string): number {
+  if (typeof text !== "string" || text.length === 0) {
+    return 0;
+  }
+  return Math.ceil(text.length / 4);
+}
+
+function getApproxTokenizer(): Tokenizer {
+  return {
+    encoding: APPROX_ENCODING,
+    countTokens: (input: string) => Promise.resolve(approximateCount(input)),
+  };
+}
+
 const encodingPromises = new Map<ModelName, Promise<string>>();
 const inFlightCounts = new Map<string, Promise<number>>();
 const tokenCountCache = new LRUCache<string, number>({
@@ -138,6 +163,14 @@ async function countTokensInternal(modelName: ModelName, text: string): Promise<
 export function loadTokenizerModules(
   modelsToWarm: string[] = Array.from(DEFAULT_WARM_MODELS)
 ): Promise<Array<PromiseSettledResult<string>>> {
+  if (shouldUseApproxTokenizer()) {
+    const fulfilled: Array<PromiseFulfilledResult<string>> = modelsToWarm.map(() => ({
+      status: "fulfilled",
+      value: APPROX_ENCODING,
+    }));
+    return Promise.resolve(fulfilled);
+  }
+
   return Promise.allSettled(
     modelsToWarm.map((modelString) => {
       const modelName = normalizeModelKey(modelString);
@@ -151,6 +184,10 @@ export function loadTokenizerModules(
 }
 
 export async function getTokenizerForModel(modelString: string): Promise<Tokenizer> {
+  if (shouldUseApproxTokenizer()) {
+    return getApproxTokenizer();
+  }
+
   const modelName = resolveModelName(modelString);
   const encodingName = await resolveEncoding(modelName);
 
@@ -161,12 +198,21 @@ export async function getTokenizerForModel(modelString: string): Promise<Tokeniz
 }
 
 export function countTokens(modelString: string, text: string): Promise<number> {
+  if (shouldUseApproxTokenizer()) {
+    return Promise.resolve(approximateCount(text));
+  }
+
   const modelName = resolveModelName(modelString);
   return countTokensInternal(modelName, text);
 }
 
 export function countTokensBatch(modelString: string, texts: string[]): Promise<number[]> {
   assert(Array.isArray(texts), "Batch token counting expects an array of strings");
+
+  if (shouldUseApproxTokenizer()) {
+    return Promise.resolve(texts.map((text) => approximateCount(text)));
+  }
+
   const modelName = resolveModelName(modelString);
   return Promise.all(texts.map((text) => countTokensInternal(modelName, text)));
 }

tests/setup.ts

@@ -7,6 +7,12 @@ import assert from "assert";
 import "disposablestack/auto";
 
 assert.equal(typeof Symbol.dispose, "symbol");
+// Use fast approximate token counting in Jest to avoid 10s WASM cold starts
+// Individual tests can override with MUX_FORCE_REAL_TOKENIZER=1
+if (process.env.MUX_FORCE_REAL_TOKENIZER !== "1") {
+  process.env.MUX_APPROX_TOKENIZER ??= "1";
+}
+
 assert.equal(typeof Symbol.asyncDispose, "symbol");
 
 // Polyfill File for undici in jest environment