From e069295d6ac6f836d0c6ecc233546117deca5d36 Mon Sep 17 00:00:00 2001 From: ephraimduncan Date: Thu, 30 Jul 2026 17:54:28 +0000 Subject: [PATCH 1/2] fix(elements): dedupe in-flight code block tokenization An uncached code/language pair launched two Shiki jobs on first render: CodeBlockContent calls highlightCode from both its render-time memo and its passive effect before the token cache is warm. Track in-flight keys so a cache miss with a pending job only subscribes instead of starting duplicate tokenization. Also evict rejected createHighlighter promises from the highlighter cache so a transient failure no longer makes a language permanently unhighlightable. --- .../code-block-highlight-dedupe.test.tsx | 90 +++++++++++++++++++ packages/elements/src/code-block.tsx | 35 ++++++-- 2 files changed, 117 insertions(+), 8 deletions(-) create mode 100644 packages/elements/__tests__/code-block-highlight-dedupe.test.tsx diff --git a/packages/elements/__tests__/code-block-highlight-dedupe.test.tsx b/packages/elements/__tests__/code-block-highlight-dedupe.test.tsx new file mode 100644 index 00000000..03b9834c --- /dev/null +++ b/packages/elements/__tests__/code-block-highlight-dedupe.test.tsx @@ -0,0 +1,90 @@ +import { render, waitFor } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import { CodeBlockContent, highlightCode } from "../src/code-block"; + +const mocks = vi.hoisted(() => ({ + codeToTokensSpy: vi.fn((code: string) => ({ + bg: "#fff", + fg: "#000", + tokens: code.split("\n").map((line) => [{ color: "#000", content: line }]), + })), + failNext: { value: false }, +})); + +// oxlint-disable-next-line typescript-eslint(consistent-type-imports) +vi.mock(import("shiki"), () => ({ + createHighlighter: vi.fn(() => { + if (mocks.failNext.value) { + mocks.failNext.value = false; + return Promise.reject(new Error("boom")); + } + return Promise.resolve({ + codeToTokens: mocks.codeToTokensSpy, + getLoadedLanguages: () => ["javascript"], + }); + }), +})); + +const flushPending = () => { + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 20); + return promise; +}; + +describe("highlightCode in-flight dedupe", () => { + it("tokenizes an uncached block exactly once on first mount", async () => { + // CodeBlockContent calls highlightCode from both its render-time memo + // and its passive effect; only one tokenization job may result. + render(); + + await waitFor(() => { + expect(mocks.codeToTokensSpy).toHaveBeenCalled(); + }); + // Let any duplicate in-flight promise settle before counting + await flushPending(); + + expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); + }); + + it("delivers one in-flight result to all subscribers", async () => { + mocks.codeToTokensSpy.mockClear(); + const results: unknown[] = []; + + highlightCode("const b = 2;", "javascript", (r) => results.push(r)); + highlightCode("const b = 2;", "javascript", (r) => results.push(r)); + + await waitFor(() => { + expect(results).toHaveLength(2); + }); + expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); + expect(results[0]).toBe(results[1]); + }); + + it("permits a retry after a rejected highlight", async () => { + mocks.codeToTokensSpy.mockClear(); + mocks.failNext.value = true; + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => { + // suppress expected failure log + }); + + // The highlighter promise is cached per language; use a fresh language + // so the rejecting createHighlighter call is actually exercised. + highlightCode("const c = 3;", "python", () => { + // never called - highlight fails + }); + await waitFor(() => { + expect(consoleSpy).toHaveBeenCalled(); + }); + expect(mocks.codeToTokensSpy).not.toHaveBeenCalled(); + + // Pending-key state must be released so a later call can retry + const results: unknown[] = []; + highlightCode("const c = 3;", "python", (r) => results.push(r)); + await waitFor(() => { + expect(results).toHaveLength(1); + }); + expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); + consoleSpy.mockRestore(); + }); +}); diff --git a/packages/elements/src/code-block.tsx b/packages/elements/src/code-block.tsx index 820142d2..35f33fe8 100644 --- a/packages/elements/src/code-block.tsx +++ b/packages/elements/src/code-block.tsx @@ -1,5 +1,13 @@ "use client"; +import type { ComponentProps, CSSProperties, HTMLAttributes } from "react"; +import type { + BundledLanguage, + BundledTheme, + HighlighterGeneric, + ThemedToken, +} from "shiki"; + import { Button } from "@repo/shadcn-ui/components/ui/button"; import { Select, @@ -10,7 +18,6 @@ import { } from "@repo/shadcn-ui/components/ui/select"; import { cn } from "@repo/shadcn-ui/lib/utils"; import { CheckIcon, CopyIcon } from "lucide-react"; -import type { ComponentProps, CSSProperties, HTMLAttributes } from "react"; import { createContext, memo, @@ -21,12 +28,6 @@ import { useRef, useState, } from "react"; -import type { - BundledLanguage, - BundledTheme, - HighlighterGeneric, - ThemedToken, -} from "shiki"; import { createHighlighter } from "shiki"; // Shiki uses bitflags for font styles: 1=italic, 2=bold, 4=underline @@ -141,6 +142,9 @@ const tokensCache = new Map(); // Subscribers for async token updates const subscribers = new Map void>>(); +// Keys with an in-flight highlight job, to avoid duplicate tokenization +const pendingHighlights = new Set(); + const getTokensCacheKey = (code: string, language: BundledLanguage) => { const start = code.slice(0, 100); const end = code.length > 100 ? code.slice(-100) : ""; @@ -161,6 +165,12 @@ const getHighlighter = ( }); highlighterCache.set(language, highlighterPromise); + // Evict on failure so a later attempt can retry instead of reusing + // a permanently rejected promise + // oxlint-disable-next-line eslint-plugin-promise(prefer-await-to-then) + highlighterPromise.catch(() => { + highlighterCache.delete(language); + }); return highlighterPromise; }; @@ -203,7 +213,14 @@ export const highlightCode = ( subscribers.get(tokensCacheKey)?.add(callback); } - // Start highlighting in background - fire-and-forget async pattern + // Start highlighting in background - fire-and-forget async pattern. + // Skip if a job for this key is already in flight; its completion will + // populate the cache and notify all subscribers, including ours. + if (pendingHighlights.has(tokensCacheKey)) { + return null; + } + pendingHighlights.add(tokensCacheKey); + getHighlighter(language) // oxlint-disable-next-line eslint-plugin-promise(prefer-await-to-then) .then((highlighter) => { @@ -226,6 +243,7 @@ export const highlightCode = ( // Cache the result tokensCache.set(tokensCacheKey, tokenized); + pendingHighlights.delete(tokensCacheKey); // Notify all subscribers const subs = subscribers.get(tokensCacheKey); @@ -239,6 +257,7 @@ export const highlightCode = ( // oxlint-disable-next-line eslint-plugin-promise(prefer-await-to-then), eslint-plugin-promise(prefer-await-to-callbacks) .catch((error) => { console.error("Failed to highlight code:", error); + pendingHighlights.delete(tokensCacheKey); subscribers.delete(tokensCacheKey); }); From 40d12117b2b986bf93dc051422fa6f2400d6684e Mon Sep 17 00:00:00 2001 From: ephraimduncan Date: Thu, 30 Jul 2026 17:59:24 +0000 Subject: [PATCH 2/2] test(elements): fold highlight dedupe tests into code-block suite --- .../code-block-highlight-dedupe.test.tsx | 90 ------------------ .../elements/__tests__/code-block.test.tsx | 93 ++++++++++++++++++- 2 files changed, 92 insertions(+), 91 deletions(-) delete mode 100644 packages/elements/__tests__/code-block-highlight-dedupe.test.tsx diff --git a/packages/elements/__tests__/code-block-highlight-dedupe.test.tsx b/packages/elements/__tests__/code-block-highlight-dedupe.test.tsx deleted file mode 100644 index 03b9834c..00000000 --- a/packages/elements/__tests__/code-block-highlight-dedupe.test.tsx +++ /dev/null @@ -1,90 +0,0 @@ -import { render, waitFor } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; - -import { CodeBlockContent, highlightCode } from "../src/code-block"; - -const mocks = vi.hoisted(() => ({ - codeToTokensSpy: vi.fn((code: string) => ({ - bg: "#fff", - fg: "#000", - tokens: code.split("\n").map((line) => [{ color: "#000", content: line }]), - })), - failNext: { value: false }, -})); - -// oxlint-disable-next-line typescript-eslint(consistent-type-imports) -vi.mock(import("shiki"), () => ({ - createHighlighter: vi.fn(() => { - if (mocks.failNext.value) { - mocks.failNext.value = false; - return Promise.reject(new Error("boom")); - } - return Promise.resolve({ - codeToTokens: mocks.codeToTokensSpy, - getLoadedLanguages: () => ["javascript"], - }); - }), -})); - -const flushPending = () => { - const { promise, resolve } = Promise.withResolvers(); - setTimeout(resolve, 20); - return promise; -}; - -describe("highlightCode in-flight dedupe", () => { - it("tokenizes an uncached block exactly once on first mount", async () => { - // CodeBlockContent calls highlightCode from both its render-time memo - // and its passive effect; only one tokenization job may result. - render(); - - await waitFor(() => { - expect(mocks.codeToTokensSpy).toHaveBeenCalled(); - }); - // Let any duplicate in-flight promise settle before counting - await flushPending(); - - expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); - }); - - it("delivers one in-flight result to all subscribers", async () => { - mocks.codeToTokensSpy.mockClear(); - const results: unknown[] = []; - - highlightCode("const b = 2;", "javascript", (r) => results.push(r)); - highlightCode("const b = 2;", "javascript", (r) => results.push(r)); - - await waitFor(() => { - expect(results).toHaveLength(2); - }); - expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); - expect(results[0]).toBe(results[1]); - }); - - it("permits a retry after a rejected highlight", async () => { - mocks.codeToTokensSpy.mockClear(); - mocks.failNext.value = true; - const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => { - // suppress expected failure log - }); - - // The highlighter promise is cached per language; use a fresh language - // so the rejecting createHighlighter call is actually exercised. - highlightCode("const c = 3;", "python", () => { - // never called - highlight fails - }); - await waitFor(() => { - expect(consoleSpy).toHaveBeenCalled(); - }); - expect(mocks.codeToTokensSpy).not.toHaveBeenCalled(); - - // Pending-key state must be released so a later call can retry - const results: unknown[] = []; - highlightCode("const c = 3;", "python", (r) => results.push(r)); - await waitFor(() => { - expect(results).toHaveLength(1); - }); - expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); - consoleSpy.mockRestore(); - }); -}); diff --git a/packages/elements/__tests__/code-block.test.tsx b/packages/elements/__tests__/code-block.test.tsx index b05fb472..951553f4 100644 --- a/packages/elements/__tests__/code-block.test.tsx +++ b/packages/elements/__tests__/code-block.test.tsx @@ -1,7 +1,98 @@ import { render, screen, waitFor } from "@testing-library/react"; import { userEvent } from "@testing-library/user-event"; -import { CodeBlock, CodeBlockCopyButton } from "../src/code-block"; +import { + CodeBlock, + CodeBlockContent, + CodeBlockCopyButton, + highlightCode, +} from "../src/code-block"; + +const mocks = vi.hoisted(() => ({ + codeToTokensSpy: vi.fn((code: string) => ({ + bg: "#fff", + fg: "#000", + tokens: code.split("\n").map((line) => [{ color: "#000", content: line }]), + })), + failNext: { value: false }, +})); + +// oxlint-disable-next-line typescript-eslint(consistent-type-imports) +vi.mock(import("shiki"), () => ({ + createHighlighter: vi.fn(() => { + if (mocks.failNext.value) { + mocks.failNext.value = false; + return Promise.reject(new Error("boom")); + } + return Promise.resolve({ + codeToTokens: mocks.codeToTokensSpy, + getLoadedLanguages: () => ["javascript"], + }); + }), +})); + +const flushPending = () => { + const { promise, resolve } = Promise.withResolvers(); + setTimeout(resolve, 20); + return promise; +}; + +describe(highlightCode, () => { + it("tokenizes an uncached block exactly once on first mount", async () => { + // CodeBlockContent calls highlightCode from both its render-time memo + // and its passive effect; only one tokenization job may result. + render(); + + await waitFor(() => { + expect(mocks.codeToTokensSpy).toHaveBeenCalled(); + }); + // Let any duplicate in-flight promise settle before counting + await flushPending(); + + expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); + }); + + it("delivers one in-flight result to all subscribers", async () => { + mocks.codeToTokensSpy.mockClear(); + const results: unknown[] = []; + + highlightCode("const b = 2;", "javascript", (r) => results.push(r)); + highlightCode("const b = 2;", "javascript", (r) => results.push(r)); + + await waitFor(() => { + expect(results).toHaveLength(2); + }); + expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); + expect(results[0]).toBe(results[1]); + }); + + it("permits a retry after a rejected highlight", async () => { + mocks.codeToTokensSpy.mockClear(); + mocks.failNext.value = true; + const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => { + // suppress expected failure log + }); + + // The highlighter promise is cached per language; use a fresh language + // so the rejecting createHighlighter call is actually exercised. + highlightCode("const c = 3;", "python", () => { + // never called - highlight fails + }); + await waitFor(() => { + expect(consoleSpy).toHaveBeenCalled(); + }); + expect(mocks.codeToTokensSpy).not.toHaveBeenCalled(); + + // Pending-key state must be released so a later call can retry + const results: unknown[] = []; + highlightCode("const c = 3;", "python", (r) => results.push(r)); + await waitFor(() => { + expect(results).toHaveLength(1); + }); + expect(mocks.codeToTokensSpy).toHaveBeenCalledOnce(); + consoleSpy.mockRestore(); + }); +}); describe("codeBlock", () => { it("renders code content", async () => {