fix(elements): dedupe in-flight code block tokenization

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<typeof import("shiki")>(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<void>();
+  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(<CodeBlockContent code="const a = 1;" language="javascript" />);
+
+    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 () => {

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<string, TokenizedCode>();
 // Subscribers for async token updates
 const subscribers = new Map<string, Set<(result: TokenizedCode) => void>>();
 
+// Keys with an in-flight highlight job, to avoid duplicate tokenization
+const pendingHighlights = new Set<string>();
+
 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);
     });