Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 1 addition & 9 deletions .github/workflows/pr.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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' }}
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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'
Expand Down
4 changes: 2 additions & 2 deletions src/cli/run.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down
46 changes: 46 additions & 0 deletions src/node/utils/main/tokenizer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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>({
Expand Down Expand Up @@ -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);
Expand All @@ -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);

Expand All @@ -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)));
}
Expand Down
6 changes: 6 additions & 0 deletions tests/setup.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down