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;
};