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 b05fb472..8e667807 100644 --- a/packages/elements/__tests__/code-block.test.tsx +++ b/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( + + ); + 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 820142d2..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) @@ -141,26 +141,30 @@ 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 ): 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; };