Skip to content

Commit c383ef8

Browse files
committed
fix: allow tokenizer fallback when Piscina unavailable
Detect Bun runtime and bypass Piscina so token counting works in tests. Assert encoding responses and reuse worker logic when pooling is off. Disable Git signing in tests and align fixtures with current adapters.
1 parent 4250fd6 commit c383ef8

File tree

8 files changed

+52
-43
lines changed

8 files changed

+52
-43
lines changed

src/git.test.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ describe("createWorktree", () => {
2121
await execAsync(`git init`, { cwd: tempGitRepo });
2222
await execAsync(`git config user.email "test@example.com"`, { cwd: tempGitRepo });
2323
await execAsync(`git config user.name "Test User"`, { cwd: tempGitRepo });
24+
await execAsync(`git config commit.gpgsign false`, { cwd: tempGitRepo });
2425
await execAsync(`echo "test" > README.md`, { cwd: tempGitRepo });
2526
await execAsync(`git add .`, { cwd: tempGitRepo });
2627
await execAsync(`git commit -m "Initial commit"`, { cwd: tempGitRepo });

src/runtime/sshConnectionPool.test.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import * as os from "os";
2+
import * as path from "path";
23
import { getControlPath } from "./sshConnectionPool";
34
import type { SSHRuntimeConfig } from "./SSHRuntime";
45

@@ -127,6 +128,9 @@ describe("username isolation", () => {
127128

128129
// The path should be deterministic for this user
129130
expect(controlPath).toBe(getControlPath(config));
130-
expect(controlPath).toMatch(/^\/tmp\/cmux-ssh-[a-f0-9]{12}$/);
131+
132+
const expectedPrefix = path.join(os.tmpdir(), "cmux-ssh-");
133+
expect(controlPath.startsWith(expectedPrefix)).toBe(true);
134+
expect(controlPath).toMatch(/cmux-ssh-[a-f0-9]{12}$/);
131135
});
132136
});

src/services/ipcMain.ts

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import assert from "@/utils/assert";
2-
import crypto from "crypto";
32
import type { BrowserWindow, IpcMain as ElectronIpcMain } from "electron";
43
import { spawn, spawnSync } from "child_process";
54
import * as fs from "fs";
@@ -29,7 +28,6 @@ import { BashExecutionService } from "@/services/bashExecutionService";
2928
import { InitStateManager } from "@/services/initStateManager";
3029
import { createRuntime } from "@/runtime/runtimeFactory";
3130
import type { RuntimeConfig } from "@/types/runtime";
32-
import { uniqueSuffix } from "@/utils/hasher";
3331
/**
3432
* IpcMain - Manages all IPC handlers and service coordination
3533
*

src/services/streamManager.test.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -379,7 +379,7 @@ describe("StreamManager - Unavailable Tool Handling", () => {
379379
messageId: "test-message-1",
380380
token: "test-token",
381381
startTime: Date.now(),
382-
model: "test-model",
382+
model: "anthropic:claude-sonnet-4-5",
383383
historySequence: 1,
384384
parts: [],
385385
lastPartialWriteTime: 0,

src/stores/WorkspaceConsumerManager.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -250,8 +250,8 @@ export class WorkspaceConsumerManager {
250250
typeof queueMicrotask === "function"
251251
? queueMicrotask
252252
: (callback: () => void) => {
253-
void Promise.resolve().then(callback);
254-
};
253+
void Promise.resolve().then(callback);
254+
};
255255

256256
schedule(() => {
257257
this.pendingNotifications.delete(workspaceId);

src/utils/main/tokenizer.ts

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,10 @@ import { getAvailableTools, getToolSchemas } from "@/utils/tools/toolDefinitions
55
import Piscina from "piscina";
66
import { createRequire } from "module";
77
import type { CountTokensInput } from "./tokenizer.worker";
8+
import {
9+
countTokens as workerCountTokens,
10+
encodingName as workerEncodingName,
11+
} from "./tokenizer.worker";
812
import { models, type ModelName } from "ai-tokenizer";
913

1014
/**
@@ -27,11 +31,18 @@ const DEFAULT_WARM_MODELS = [
2731
] as const;
2832

2933
const requireForResolve = createRequire(__filename);
30-
const workerPath = requireForResolve.resolve("./tokenizer-pool.worker");
31-
const tokenizerPool = new Piscina({
32-
filename: workerPath,
33-
idleTimeout: Infinity,
34-
});
34+
const supportsPiscina =
35+
typeof process !== "undefined" &&
36+
!!process.versions?.node &&
37+
!("bun" in (process.versions ?? {}));
38+
39+
const workerPath = supportsPiscina ? requireForResolve.resolve("./tokenizer.worker") : null;
40+
const tokenizerPool = supportsPiscina
41+
? new Piscina({
42+
filename: workerPath!,
43+
idleTimeout: Infinity,
44+
})
45+
: null;
3546

3647
const encodingPromises = new Map<ModelName, Promise<string>>();
3748
const inFlightCounts = new Map<string, Promise<number>>();
@@ -57,6 +68,12 @@ function normalizeModelKey(modelName: string): ModelName {
5768
}
5869

5970
async function resolveEncoding(modelName: ModelName): Promise<string> {
71+
if (!tokenizerPool) {
72+
const encoding = workerEncodingName(modelName);
73+
assert(typeof encoding === "string" && encoding.length > 0);
74+
return encoding;
75+
}
76+
6077
let promise = encodingPromises.get(modelName);
6178
if (!promise) {
6279
promise = tokenizerPool
@@ -94,6 +111,16 @@ async function countTokensInternal(modelName: ModelName, text: string): Promise<
94111
return cached;
95112
}
96113

114+
if (!tokenizerPool) {
115+
const value = workerCountTokens({ modelName, input: text });
116+
assert(
117+
typeof value === "number" && Number.isFinite(value) && value >= 0,
118+
"Tokenizer must return a non-negative finite token count"
119+
);
120+
tokenCountCache.set(key, value);
121+
return value;
122+
}
123+
97124
let pending = inFlightCounts.get(key);
98125
if (!pending) {
99126
const payload: CountTokensInput = { modelName, input: text };

src/utils/main/tokenizer.worker.ts

Lines changed: 6 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,7 @@
11
import assert from "@/utils/assert";
22
import { Tokenizer, models } from "ai-tokenizer";
33
import type { ModelName } from "ai-tokenizer";
4-
// import * as claude from "ai-tokenizer/encoding/claude";
5-
// import * as o200k_base from "ai-tokenizer/encoding/o200k_base";
6-
//
7-
// import * as encoding from "ai-tokenizer/encoding";
4+
import * as encoding from "ai-tokenizer/encoding";
85

96
export interface CountTokensInput {
107
modelName: ModelName;
@@ -13,43 +10,25 @@ export interface CountTokensInput {
1310

1411
const tokenizerCache = new Map<ModelName, Tokenizer>();
1512

16-
type importedClaude = typeof import("ai-tokenizer/encoding/claude");
17-
type importedO200k_base = typeof import("ai-tokenizer/encoding/o200k_base");
18-
let claude: importedClaude | null = null;
19-
let o200k_base: importedO200k_base | null = null;
20-
21-
async function jitLoadEncoding(encoding: string): Promise<importedClaude | importedO200k_base> {
22-
if (encoding == "claude") {
23-
claude ??= await import("ai-tokenizer/encoding/claude");
24-
return claude;
25-
}
26-
27-
o200k_base ??= await import("ai-tokenizer/encoding/o200k_base");
28-
return o200k_base;
29-
}
30-
31-
async function getTokenizer(modelName: ModelName): Promise<Tokenizer> {
13+
function getTokenizer(modelName: ModelName): Tokenizer {
3214
const cached = tokenizerCache.get(modelName);
3315
if (cached) {
3416
return cached;
3517
}
3618

3719
const model = models[modelName];
3820
assert(model, `Unknown tokenizer model '${modelName}'`);
39-
assert(
40-
model.encoding == "claude" || model.encoding == "o200k_base",
41-
`Unknown encoding '${model.encoding}' for '${modelName}'`
42-
);
4321

44-
const encodingModule = await jitLoadEncoding(model.encoding);
22+
const encodingModule = encoding[model.encoding];
23+
assert(encodingModule, `Unknown tokenizer encoding '${model.encoding}'`);
4524

4625
const tokenizer = new Tokenizer(encodingModule);
4726
tokenizerCache.set(modelName, tokenizer);
4827
return tokenizer;
4928
}
5029

51-
export async function countTokens({ modelName, input }: CountTokensInput): Promise<number> {
52-
const tokenizer = await getTokenizer(modelName);
30+
export function countTokens({ modelName, input }: CountTokensInput): number {
31+
const tokenizer = getTokenizer(modelName);
5332
const count = tokenizer.count(input);
5433
return count;
5534
}

src/utils/tokens/tokenStatsCalculator.test.ts

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ describe("mergeResults", () => {
343343
{ consumer: "Assistant", promise: Promise.resolve(200) },
344344
];
345345
const results = [100, 200];
346-
const toolDefinitions = new Map();
346+
const toolDefinitions = new Map<string, number>();
347347
const systemMessageTokens = 0;
348348

349349
const consumerMap = mergeResults(jobs, results, toolDefinitions, systemMessageTokens);
@@ -358,7 +358,7 @@ describe("mergeResults", () => {
358358
{ consumer: "User", promise: Promise.resolve(50) },
359359
];
360360
const results = [100, 50];
361-
const toolDefinitions = new Map();
361+
const toolDefinitions = new Map<string, number>();
362362
const systemMessageTokens = 0;
363363

364364
const consumerMap = mergeResults(jobs, results, toolDefinitions, systemMessageTokens);
@@ -372,7 +372,7 @@ describe("mergeResults", () => {
372372
{ consumer: "Read", promise: Promise.resolve(50) },
373373
];
374374
const results = [100, 50];
375-
const toolDefinitions = new Map([["Read", 25]]);
375+
const toolDefinitions = new Map<string, number>([["Read", 25]]);
376376
const systemMessageTokens = 0;
377377

378378
const consumerMap = mergeResults(jobs, results, toolDefinitions, systemMessageTokens);
@@ -384,7 +384,7 @@ describe("mergeResults", () => {
384384
test("adds system message tokens", () => {
385385
const jobs: TokenCountJob[] = [];
386386
const results: number[] = [];
387-
const toolDefinitions = new Map();
387+
const toolDefinitions = new Map<string, number>();
388388
const systemMessageTokens = 300;
389389

390390
const consumerMap = mergeResults(jobs, results, toolDefinitions, systemMessageTokens);
@@ -398,7 +398,7 @@ describe("mergeResults", () => {
398398
{ consumer: "Assistant", promise: Promise.resolve(100) },
399399
];
400400
const results = [0, 100];
401-
const toolDefinitions = new Map();
401+
const toolDefinitions = new Map<string, number>();
402402
const systemMessageTokens = 0;
403403

404404
const consumerMap = mergeResults(jobs, results, toolDefinitions, systemMessageTokens);

0 commit comments

Comments
 (0)