From e9a0981b89c0a4ca7027a5b87459c1cdafb93e57 Mon Sep 17 00:00:00 2001 From: HardMax71 Date: Fri, 2 Oct 2026 16:34:55 +0200 Subject: [PATCH 1/2] fix(code-block): key the token cache on the full code The token cache key was the language, the code length and the first and last 100 characters, so two blocks with the same length and ends shared one entry and the second block showed the first block's code. The key already held the length, so every distinct code string is already its own entry; using the full code only makes the key exact. --- .../elements/__tests__/code-block.test.tsx | 23 ++++++++++++++++++- packages/elements/src/code-block.tsx | 8 +++---- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/packages/elements/__tests__/code-block.test.tsx b/packages/elements/__tests__/code-block.test.tsx index b05fb472..b6e8cc30 100644 --- a/packages/elements/__tests__/code-block.test.tsx +++ b/packages/elements/__tests__/code-block.test.tsx @@ -1,7 +1,11 @@ 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 +47,23 @@ 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( + + ); + expect(container.textContent).toContain("const other = 2;"); + }); }); const setupCopyButtonTests = () => { diff --git a/packages/elements/src/code-block.tsx b/packages/elements/src/code-block.tsx index 820142d2..51ba5507 100644 --- a/packages/elements/src/code-block.tsx +++ b/packages/elements/src/code-block.tsx @@ -141,11 +141,9 @@ const tokensCache = new Map(); // Subscribers for async token updates const subscribers = new Map 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 From 4dbe27d47146c760819f5ae9044747e3e28d40c4 Mon Sep 17 00:00:00 2001 From: HardMax71 Date: Fri, 2 Oct 2026 16:34:56 +0200 Subject: [PATCH 2/2] fix(code-block): highlight languages shiki does not bundle as text createHighlighter rejects a language that is not in the shiki bundle, so the "text" fallback in highlightCode never ran: the rejected promise was cached for that language, and every block in it stayed unhighlighted and logged "Failed to highlight code". The language is now resolved to "text" before the highlighter is created. --- .changeset/code-block-cache-key-unknown-lang.md | 5 +++++ packages/elements/__tests__/code-block.test.tsx | 15 +++++++++++++++ packages/elements/src/code-block.tsx | 14 ++++++++++---- 3 files changed, 30 insertions(+), 4 deletions(-) create mode 100644 .changeset/code-block-cache-key-unknown-lang.md diff --git a/.changeset/code-block-cache-key-unknown-lang.md b/.changeset/code-block-cache-key-unknown-lang.md new file mode 100644 index 00000000..b2c28bf7 --- /dev/null +++ b/.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 diff --git a/packages/elements/__tests__/code-block.test.tsx b/packages/elements/__tests__/code-block.test.tsx index b6e8cc30..8e667807 100644 --- a/packages/elements/__tests__/code-block.test.tsx +++ b/packages/elements/__tests__/code-block.test.tsx @@ -1,3 +1,5 @@ +import type { BundledLanguage } from "shiki"; + import { render, screen, waitFor } from "@testing-library/react"; import { userEvent } from "@testing-library/user-event"; @@ -64,6 +66,19 @@ describe("codeBlock", () => { ); 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 = () => { diff --git a/packages/elements/src/code-block.tsx b/packages/elements/src/code-block.tsx index 51ba5507..d5e24d99 100644 --- a/packages/elements/src/code-block.tsx +++ b/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) @@ -148,17 +148,23 @@ const getTokensCacheKey = (code: string, language: BundledLanguage) => const getHighlighter = ( language: BundledLanguage ): Promise> => { - 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; };