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 () => { 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); });