fix: harden mock tokenization fallback

ThomasK33 · ThomasK33 · commit c92b63393449 · 2025-10-30T14:27:05.000+01:00
Introduce bounded tokenizer fallback for mock streams to avoid hangs.
Add assertions and debug logging to surface invalid tokenizer results.
Switch mock scenarios to openai:gpt-5 to align with tokenizer choice.
Expose PLAYWRIGHT_ARGS in test-e2e make target to ease overrides.
diff --git a/Makefile b/Makefile
@@ -220,7 +220,7 @@ test-coverage: ## Run tests with coverage
 
 test-e2e: ## Run end-to-end tests
 	@$(MAKE) build
-	@CMUX_E2E_LOAD_DIST=1 CMUX_E2E_SKIP_BUILD=1 PLAYWRIGHT_SKIP_BROWSER_DOWNLOAD=1 bun x playwright test --project=electron
+	@CMUX_E2E_LOAD_DIST=1 CMUX_E2E_SKIP_BUILD=1 PLAYWRIGHT_SKIP_BROWSER_DOWNLOAD=1 bun x playwright test --project=electron $(PLAYWRIGHT_ARGS)
 
 ## Distribution
 dist: build ## Build distributable packages
diff --git a/src/main-desktop.ts b/src/main-desktop.ts
@@ -334,7 +334,7 @@ async function loadServices(): Promise<void> {
     );
   } else {
     console.log(
-      `[${timestamp()}] Updater service disabled in dev mode (et DEBUG_UPDATER=1 or DEBUG_UPDATER=<version> to enable)`
+      `[${timestamp()}] Updater service disabled in dev mode (set DEBUG_UPDATER=1 or DEBUG_UPDATER=<version> to enable)`
     );
   }
 
diff --git a/src/services/mock/mockScenarioPlayer.ts b/src/services/mock/mockScenarioPlayer.ts
@@ -1,10 +1,12 @@
+import assert from "@/utils/assert";
 import type { CmuxMessage } from "@/types/message";
 import { createCmuxMessage } from "@/types/message";
 import type { HistoryService } from "@/services/historyService";
 import type { Result } from "@/types/result";
 import { Ok, Err } from "@/types/result";
 import type { SendMessageError } from "@/types/errors";
 import type { AIService } from "@/services/aiService";
+import { log } from "@/services/log";
 import type {
   MockAssistantEvent,
   MockStreamErrorEvent,
@@ -17,6 +19,91 @@ import type { ToolCallStartEvent, ToolCallEndEvent } from "@/types/stream";
 import type { ReasoningDeltaEvent } from "@/types/stream";
 import { getTokenizerForModel } from "@/utils/main/tokenizer";
 
+const MOCK_TOKENIZER_MODEL = "openai:gpt-5";
+const TOKENIZE_TIMEOUT_MS = 150;
+let tokenizerFallbackLogged = false;
+
+function approximateTokenCount(text: string): number {
+  const normalizedLength = text.trim().length;
+  if (normalizedLength === 0) {
+    return 0;
+  }
+  return Math.max(1, Math.ceil(normalizedLength / 4));
+}
+
+async function tokenizeWithMockModel(text: string, context: string): Promise<number> {
+  assert(typeof text === "string", `Mock scenario ${context} expects string input`);
+  const approximateTokens = approximateTokenCount(text);
+  let fallbackUsed = false;
+  let timeoutId: NodeJS.Timeout | undefined;
+
+  const fallbackPromise = new Promise<number>((resolve) => {
+    timeoutId = setTimeout(() => {
+      fallbackUsed = true;
+      resolve(approximateTokens);
+    }, TOKENIZE_TIMEOUT_MS);
+  });
+
+  const actualPromise = (async () => {
+    const tokenizer = await getTokenizerForModel(MOCK_TOKENIZER_MODEL);
+    assert(
+      typeof tokenizer.encoding === "string" && tokenizer.encoding.length > 0,
+      `Tokenizer for ${MOCK_TOKENIZER_MODEL} must expose a non-empty encoding`
+    );
+    const tokens = await tokenizer.countTokens(text);
+    assert(
+      Number.isFinite(tokens) && tokens >= 0,
+      `Tokenizer for ${MOCK_TOKENIZER_MODEL} returned invalid token count`
+    );
+    return tokens;
+  })();
+
+  let tokens: number;
+  try {
+    tokens = await Promise.race([actualPromise, fallbackPromise]);
+  } catch (error) {
+    if (timeoutId !== undefined) {
+      clearTimeout(timeoutId);
+    }
+    const errorMessage = error instanceof Error ? error.message : String(error);
+    throw new Error(
+      `[MockScenarioPlayer] Failed to tokenize ${context} with ${MOCK_TOKENIZER_MODEL}: ${errorMessage}`
+    );
+  }
+
+  if (!fallbackUsed && timeoutId !== undefined) {
+    clearTimeout(timeoutId);
+  }
+
+  actualPromise
+    .then((resolvedTokens) => {
+      if (fallbackUsed && !tokenizerFallbackLogged) {
+        tokenizerFallbackLogged = true;
+        log.debug(
+          `[MockScenarioPlayer] Tokenizer fallback used for ${context}; emitted ${approximateTokens}, background tokenizer returned ${resolvedTokens}`
+        );
+      }
+    })
+    .catch((error) => {
+      if (fallbackUsed && !tokenizerFallbackLogged) {
+        tokenizerFallbackLogged = true;
+        const errorMessage = error instanceof Error ? error.message : String(error);
+        log.debug(
+          `[MockScenarioPlayer] Tokenizer fallback used for ${context}; background error: ${errorMessage}`
+        );
+      }
+    });
+
+  if (fallbackUsed) {
+    assert(
+      Number.isFinite(tokens) && tokens >= 0,
+      `Token fallback produced invalid count for ${context}`
+    );
+  }
+
+  return tokens;
+}
+
 interface MockPlayerDeps {
   aiService: AIService;
   historyService: HistoryService;
@@ -159,8 +246,7 @@ export class MockScenarioPlayer {
       }
       case "reasoning-delta": {
         // Mock scenarios use the same tokenization logic as real streams for consistency
-        const tokenizer = await getTokenizerForModel("gpt-4"); // Mock uses GPT-4 tokenizer
-        const tokens = await tokenizer.countTokens(event.text);
+        const tokens = await tokenizeWithMockModel(event.text, "reasoning-delta text");
         const payload: ReasoningDeltaEvent = {
           type: "reasoning-delta",
           workspaceId,
@@ -175,8 +261,7 @@ export class MockScenarioPlayer {
       case "tool-start": {
         // Mock scenarios use the same tokenization logic as real streams for consistency
         const inputText = JSON.stringify(event.args);
-        const tokenizer = await getTokenizerForModel("gpt-4"); // Mock uses GPT-4 tokenizer
-        const tokens = await tokenizer.countTokens(inputText);
+        const tokens = await tokenizeWithMockModel(inputText, "tool-call args");
         const payload: ToolCallStartEvent = {
           type: "tool-call-start",
           workspaceId,
@@ -204,8 +289,13 @@ export class MockScenarioPlayer {
       }
       case "stream-delta": {
         // Mock scenarios use the same tokenization logic as real streams for consistency
-        const tokenizer = await getTokenizerForModel("gpt-4"); // Mock uses GPT-4 tokenizer
-        const tokens = await tokenizer.countTokens(event.text);
+        let tokens: number;
+        try {
+          tokens = await tokenizeWithMockModel(event.text, "stream-delta text");
+        } catch (error) {
+          console.error("[MockScenarioPlayer] tokenize failed for stream-delta", error);
+          throw error;
+        }
         const payload: StreamDeltaEvent = {
           type: "stream-delta",
           workspaceId,
diff --git a/src/services/mock/scenarios/basicChat.ts b/src/services/mock/scenarios/basicChat.ts
@@ -12,7 +12,7 @@ const listProgrammingLanguagesTurn: ScenarioTurn = {
   assistant: {
     messageId: "msg-basic-1",
     events: [
-      { kind: "stream-start", delay: 0, messageId: "msg-basic-1", model: "mock:planner" },
+      { kind: "stream-start", delay: 0, messageId: "msg-basic-1", model: "openai:gpt-5" },
       {
         kind: "stream-delta",
         delay: STREAM_BASE_DELAY,
@@ -37,7 +37,7 @@ const listProgrammingLanguagesTurn: ScenarioTurn = {
         kind: "stream-end",
         delay: STREAM_BASE_DELAY * 5,
         metadata: {
-          model: "mock:planner",
+          model: "openai:gpt-5",
           inputTokens: 64,
           outputTokens: 48,
           systemMessageTokens: 12,
diff --git a/src/services/mock/scenarios/permissionModes.ts b/src/services/mock/scenarios/permissionModes.ts
@@ -19,7 +19,7 @@ const planRefactorTurn: ScenarioTurn = {
         kind: "stream-start",
         delay: 0,
         messageId: "msg-plan-refactor",
-        model: "mock:planner",
+        model: "openai:gpt-5",
       },
       {
         kind: "stream-delta",
@@ -45,7 +45,7 @@ const planRefactorTurn: ScenarioTurn = {
         kind: "stream-end",
         delay: STREAM_BASE_DELAY * 5,
         metadata: {
-          model: "mock:planner",
+          model: "openai:gpt-5",
           inputTokens: 180,
           outputTokens: 130,
           systemMessageTokens: 24,
diff --git a/src/services/mock/scenarios/review.ts b/src/services/mock/scenarios/review.ts
@@ -16,7 +16,7 @@ const summarizeBranchesTurn: ScenarioTurn = {
   assistant: {
     messageId: "msg-plan-1",
     events: [
-      { kind: "stream-start", delay: 0, messageId: "msg-plan-1", model: "mock:planner" },
+      { kind: "stream-start", delay: 0, messageId: "msg-plan-1", model: "openai:gpt-5" },
       {
         kind: "reasoning-delta",
         delay: STREAM_BASE_DELAY,
@@ -61,7 +61,7 @@ const summarizeBranchesTurn: ScenarioTurn = {
         kind: "stream-end",
         delay: STREAM_BASE_DELAY * 6,
         metadata: {
-          model: "mock:planner",
+          model: "openai:gpt-5",
           inputTokens: 128,
           outputTokens: 85,
           systemMessageTokens: 32,
diff --git a/src/services/mock/scenarios/slashCommands.ts b/src/services/mock/scenarios/slashCommands.ts
@@ -24,7 +24,7 @@ const compactConversationTurn: ScenarioTurn = {
         kind: "stream-start",
         delay: 0,
         messageId: "msg-slash-compact-1",
-        model: "anthropic:claude-sonnet-4-5",
+        model: "openai:gpt-5",
       },
       {
         kind: "stream-delta",
@@ -35,7 +35,7 @@ const compactConversationTurn: ScenarioTurn = {
         kind: "stream-end",
         delay: STREAM_BASE_DELAY * 2,
         metadata: {
-          model: "anthropic:claude-sonnet-4-5",
+          model: "openai:gpt-5",
           inputTokens: 220,
           outputTokens: 96,
           systemMessageTokens: 18,
diff --git a/src/services/mock/scenarios/toolFlows.ts b/src/services/mock/scenarios/toolFlows.ts
@@ -269,7 +269,7 @@ const recallTestFileTurn: ScenarioTurn = {
         kind: "stream-start",
         delay: 0,
         messageId: "msg-tool-recall-test-file",
-        model: "mock:planner",
+        model: "openai:gpt-5",
       },
       {
         kind: "stream-delta",
@@ -280,7 +280,7 @@ const recallTestFileTurn: ScenarioTurn = {
         kind: "stream-end",
         delay: STREAM_BASE_DELAY * 2,
         metadata: {
-          model: "mock:planner",
+          model: "openai:gpt-5",
           inputTokens: 60,
           outputTokens: 34,
           systemMessageTokens: 10,

Original file line number	Diff line number	Diff line change
`@@ -334,7 +334,7 @@ async function loadServices(): Promise<void> {`
`334`	`334`	`);`
`335`	`335`	`} else {`
`336`	`336`	`console.log(`
`337`		- `[${timestamp()}] Updater service disabled in dev mode (et DEBUG_UPDATER=1 or DEBUG_UPDATER=<version> to enable)`
	`337`	+ `[${timestamp()}] Updater service disabled in dev mode (set DEBUG_UPDATER=1 or DEBUG_UPDATER=<version> to enable)`
`338`	`338`	`);`
`339`	`339`	`}`
`340`	`340`