fix(code-block): key the token cache on the full code and highlight unknown languages as text

.changeset/code-block-cache-key-unknown-lang.md

@@ -0,0 +1,5 @@
+---
+"ai-elements": patch
+---
+
+Fix CodeBlock showing the tokens of other code with the same length and ends, and highlight languages shiki does not bundle as plain text

packages/elements/__tests__/code-block.test.tsx

@@ -1,7 +1,13 @@
+import type { BundledLanguage } from "shiki";
+
 import { render, screen, waitFor } from "@testing-library/react";
 import { userEvent } from "@testing-library/user-event";
 
-import { CodeBlock, CodeBlockCopyButton } from "../src/code-block";
+import {
+  CodeBlock,
+  CodeBlockCopyButton,
+  highlightCode,
+} from "../src/code-block";
 
 describe("codeBlock", () => {
   it("renders code content", async () => {
@@ -43,6 +49,36 @@ describe("codeBlock", () => {
     expect(container.firstChild).toHaveClass("group");
     expect(container.firstChild).toHaveClass("relative");
   });
+
+  it("does not show the tokens of other code with the same length and ends", async () => {
+    const edge = "x".repeat(100);
+    const first = `${edge}\nconst first = 1;\n${edge}`;
+    const other = `${edge}\nconst other = 2;\n${edge}`;
+    const onHighlighted = vi.fn();
+
+    highlightCode(first, "javascript", onHighlighted);
+    await waitFor(() => {
+      expect(onHighlighted).toHaveBeenCalled();
+    });
+
+    const { container } = render(
+      <CodeBlock code={other} language="javascript" />
+    );
+    expect(container.textContent).toContain("const other = 2;");
+  });
+
+  it("highlights a language shiki does not bundle as plain text", async () => {
+    const onHighlighted = vi.fn();
+
+    highlightCode(
+      "IF ready THEN start",
+      "pseudocode" as BundledLanguage,
+      onHighlighted
+    );
+    await waitFor(() => {
+      expect(onHighlighted).toHaveBeenCalled();
+    });
+  });
 });
 
 const setupCopyButtonTests = () => {

packages/elements/src/code-block.tsx

@@ -27,7 +27,7 @@ import type {
   HighlighterGeneric,
   ThemedToken,
 } from "shiki";
-import { createHighlighter } from "shiki";
+import { bundledLanguages, createHighlighter, isSpecialLang } from "shiki";
 
 // Shiki uses bitflags for font styles: 1=italic, 2=bold, 4=underline
 // oxlint-disable-next-line eslint(no-bitwise)
@@ -141,26 +141,30 @@ const tokensCache = new Map<string, TokenizedCode>();
 // Subscribers for async token updates
 const subscribers = new Map<string, Set<(result: TokenizedCode) => void>>();
 
-const getTokensCacheKey = (code: string, language: BundledLanguage) => {
-  const start = code.slice(0, 100);
-  const end = code.length > 100 ? code.slice(-100) : "";
-  return `${language}:${code.length}:${start}:${end}`;
-};
+// Keyed on the full code: a sample of it lets different code share tokens
+const getTokensCacheKey = (code: string, language: BundledLanguage) =>
+  `${language}:${code}`;
 
 const getHighlighter = (
   language: BundledLanguage
 ): Promise<HighlighterGeneric<BundledLanguage, BundledTheme>> => {
-  const cached = highlighterCache.get(language);
+  // createHighlighter rejects a language shiki does not bundle
+  const lang =
+    Object.hasOwn(bundledLanguages, language) || isSpecialLang(language)
+      ? language
+      : "text";
+
+  const cached = highlighterCache.get(lang);
   if (cached) {
     return cached;
   }
 
   const highlighterPromise = createHighlighter({
-    langs: [language],
+    langs: [lang],
     themes: ["github-light", "github-dark"],
   });
 
-  highlighterCache.set(language, highlighterPromise);
+  highlighterCache.set(lang, highlighterPromise);
   return highlighterPromise;
 };