Skip to content
← Developer Diaries

Writing a unified Plugin for Code Blocks

11 min read Revised
  • Tooling
  • #typescript
  • #unified
  • #rehype
  • #markdown

How to wrap markdown code blocks in a language label and a copy button by rewriting the HTML syntax tree with a rehype plugin, and what the inline onclick actually costs you.

Every markdown renderer turns a fenced block into <pre><code class="language-ts"> and stops there. To get a language label and a copy button, you write a rehype plugin: walk the HAST with unist-util-visit, match every element whose tagName is pre, and replace it in its parent's children array with a wrapper element that holds a header and the original <pre>. It is about forty lines, it runs at build time, and it never touches a string of HTML.

Update, September 2026: The plugin has since been corrected on every count below: className as an array rather than a raw class, the <code> child found by tag name, [SKIP, index + 1] returned after replacement, and a data-copy-code attribute in place of the inline onclick. The code quoted here is the version being examined.

The alternative is a regex over the serialised output. Don't. The right place to intervene is while the document is still a tree.

Where does a markdown pipeline let you intervene?#

unified processes documents as a pipeline of syntax trees. Markdown text becomes MDAST (a markdown tree), which becomes HAST (an HTML tree), which becomes a string. Each arrow is a plugin, and you can insert your own at any of them.

import { unified } from 'unified';          // 11.0.5
import remarkParse from 'remark-parse';     // 11.0.0
import remarkRehype from 'remark-rehype';   // 11.1.2
import customCodeBlocks from './customCodeblocks';
import rehypeHighlight from 'rehype-highlight'; // 7.0.2
import rehypeStringify from 'rehype-stringify'; // 10.0.1

const html = unified()
  .use(remarkParse)
  .use(remarkRehype)
  .use(customCodeBlocks)
  .use(rehypeHighlight)
  .use(rehypeStringify)
  .processSync(md)
  .toString();

Those are the current major versions, and they matter: the repo behind this post pins unified 10, remark-parse 10, rehype-highlight 6 and rehype-stringify 9. The whole ecosystem went to ESM-only and bumped majors together, so upgrading is one coordinated jump rather than a package at a time. The plugin code below is unchanged by it.

Working on HAST — after markdown has become HTML structure, before it is serialised — means you are manipulating real nodes with real properties, not pattern-matching against text. Order inside the rehype half is yours to pick. I put my plugin before rehype-highlight, and it makes no difference either way: highlighting looks for a <code> whose parent is <pre>, and wrapping the <pre> in a div leaves that relationship intact. I confirmed this rather than assumed it, because it is the sort of thing that quietly stops being true.

Finding every code block with unist-util-visit#

unist-util-visit (5.1.0) walks the tree and calls you for matching nodes. The signature is visit(tree, test, visitor), and the visitor receives (node, index, parent):

visit(tree, 'element', (node, index, parent) => {
  if (node.tagName !== 'pre' || parent === undefined || index === undefined) return;
  parent.children[index] = wrapper(node, language);
  return [SKIP, index + 1];
});

index and parent are what let you replace a node rather than merely mutate it. Mutating gives you a differently-dressed <pre>; replacing gives you a new subtree with the <pre> inside it.

The return [SKIP, index + 1] is the part the original version of this post got wrong by omission. The unist-util-visit readme is blunt about it: "Replacing node itself, if SKIP is not returned, still causes its descendants to be walked (which is a bug)." I instrumented both forms against a two-fence document and logged every element the visitor saw:

without SKIP:  pre code pre code
with SKIP:     pre pre

Without it, the traversal descends into the <pre> you just detached and re-parented. Here that is harmless — the <code> inside doesn't match pre, so nothing loops. But it is harmless by accident. The moment your wrapper contains anything the same visitor would match, you have an infinite loop, and you will find it at 2am. Returning [SKIP, index + 1] says: I have handled this subtree, carry on with the next sibling.

How do you get the language out of a code fence?#

remark-rehype puts it on the <code> child as a class token. Here is exactly what it produces, from an actual run:

pre children:  ["element:code"]
code props:    {"className":["language-javascript"]}

Note className, and note that it is an array. That is the HAST convention for space-separated attributes. So the original extraction works:

let language = '';
try {
  language = node.children[0].properties.className.find((className: string) =>
    className.startsWith('language-')
  );
  language = language.substring(9);
} catch (error) {
  language = 'plaintext';
}

The try/catch is doing real work, and I still defend the shape of it. Every step of that chain can legitimately fail on a valid document. A fence with no language produces properties: {} — I checked, it is an empty object, not a missing one — so .className is undefined and .find throws. Guarding each condition individually is four nested checks to arrive at the same place. Catching and defaulting says the thing you mean: if I can't tell what language this is, it's plain text.

What I would no longer defend is children[0]. In this pipeline it is reliably the <code> element, because mdast-util-to-hast emits exactly one child with no surrounding whitespace. It stops being reliable the moment anything else is in the pipeline — rehype-raw parsing an inline <pre> from the source, or a tree that came from rehype-parse rather than markdown. Then children[0] is a text node, and the code throws:

TypeError: Cannot read properties of undefined (reading 'className')

Which the catch swallows, and every block on the page silently labels itself plaintext. That is the failure mode of a broad catch: it converts a structural mismatch into a plausible-looking wrong answer. The fix is to find the child rather than index it:

const code = node.children.find(
  (child): child is Element => child.type === 'element' && child.tagName === 'code'
);

And substring(9) should be slice('language-'.length). Nine is correct; it just doesn't say why.

Two things remark-rehype will not hand you. Fence metadata is dropped. A fence opened with ts title="a.ts" still arrives as {"className":["language-ts"]} and nothing else, so a filename header needs a remark-stage plugin reading node.meta on the MDAST code node. And the language string is an identifier, not a label. The repo capitalises it, which gives you "Javascript", "Css" and "Sql" in your header. A small Record<string, string> of display names is uglier and correct.

class or className? The bug that works anyway#

Here is the thing I got wrong and shipped. The extraction above reads properties.className. The replacement node writes:

properties: {
  class: 'rounded-md bg-blue-500 prose-pre:m-0 ...',
}

One file, one convention each way. In HAST the property is className, and it takes an array of tokens. class as a raw string is not the HAST shape.

It also works, and understanding why is the useful part. hast-util-to-html resolves every property key through property-information, which normalises names case-insensitively:

class     -> {"property":"className","attribute":"class","spaceSeparated":true}
className -> {"property":"className","attribute":"class","spaceSeparated":true}

Both keys land on the same info object, and a string value is emitted as-is. So class: 'a b c', className: ['a','b','c'] and className: 'a b c' all serialise identically to class="a b c". I ran all three. The serialiser is forgiving, and my mistake never reached the output.

It reaches everything else. Any downstream plugin that inspects the tree — rehype-class-names, hast-util-select, your own second pass — reads properties.className and gets undefined, because the class list is filed under a key nothing looks for. And if a later plugin does the correct thing and adds a className alongside my class, you get this:

<div class="from-class" class="from-className"></div>

A duplicate attribute. Not a crash, not a warning — browsers keep the first and drop the second, so the styling that breaks is whichever one you weren't thinking about. Write className: ['rounded-md', 'bg-blue-500'] and the problem does not exist.

Is hastscript worth the dependency?#

The original post argued that hand-built node objects were verbose on purpose and that plain objects keep the dependency count at zero. Having lived with the sixty-line object literal in that repo, I've changed my mind. hastscript (9.0.1) is one small package and it removes an entire category of mistake:

import { h } from 'hastscript';

parent.children[index] = h('div.code-block', { 'data-language': language }, [
  h('header.code-block__bar', [
    h('span.code-block__lang', displayName(language)),
    h('button.code-block__copy', { type: 'button', 'data-copy-code': '' }, 'Copy')
  ]),
  node
]);

h('div.code-block') produces className: ['code-block'] — the correct array form — without you having to remember it. That is the point: the helper encodes the convention, so the class-versus-className bug above is unwritable. The argument for verbosity was never about safety, it was about avoiding a dependency, and here the dependency is the safety. That block plus the visitor above typechecks clean under tsc --strict, which the original could not: it needed two @ts-ignore comments and a tree: any.

What does an inline onclick actually cost?#

The repo emits this, and it is still the honest weak point:

properties: {
  class: 'btn btn-sm variant-soft-surface my-1',
  onclick: 'copyCode(this)',
}

The reasoning was sound as far as it went: the plugin emits static HTML with no way to attach a listener, so a named global that consumers define once is the smallest contract. What I did not account for is Content-Security-Policy.

Inline event handler attributes are inline script. A policy of script-src 'self' blocks them outright — the button renders and does nothing, with an error only in the console. Allowing it means either 'unsafe-inline', which defeats most of the point of having a CSP, or CSP Level 3's 'unsafe-hashes' plus a hash of the attribute's exact source text:

script-src 'self' 'unsafe-hashes' 'sha256-pUnTendwAdzhlWjppGh5zVws6PVloNc0pY9rdzjYdxc='

That is the real SHA-256 of the string copyCode(this). It works, and it is a maintenance trap: change the handler text and the hash silently stops matching. Nonces do not help — a nonce is an attribute on a <script> element, and there is no element here. SvelteKit's kit.csp will hash the inline scripts it generates, but content injected with {@html} is not its to see.

So the plugin's output is quietly incompatible with a strict CSP, and nothing in the code says so. That is the correction I most wanted to make to this post.

The fix is to emit data, not behaviour. h('button', { type: 'button', 'data-copy-code': '' }, 'Copy') is inert markup. One delegated listener, in a real script file the CSP already trusts, handles every block on the page including ones added later:

document.addEventListener('click', async (event) => {
  const button = event.target.closest('[data-copy-code]');
  if (!button) return;
  const code = button.closest('.code-block')?.querySelector('code')?.textContent;
  if (!code) return;
  await navigator.clipboard.writeText(code);
  button.textContent = 'Copied';
  setTimeout(() => { button.textContent = 'Copy'; }, 1000);
});

One listener, no globals on window, no CSP exception. If you want the behaviour to travel with the markup instead of living in a page script, emit <copy-code> as a custom element and ship a small definition — the plugin's output then declares what it needs and the consumer supplies it once, which is the contract I was reaching for with the global and didn't get.

Worth saying plainly: the demo in that repo defines copyCode inside a <script lang="ts"> nested in the markup of +page.svelte. Only a component's top-level script goes through the TypeScript preprocessor, so the annotations in it are shipped to the browser verbatim. The plugin was the part I was thinking about; the contract it depends on got the least attention, which is exactly backwards.

rehype-highlight or Shiki?#

Both are current and the choice is genuine. rehype-highlight wraps highlight.js through lowlight: 37 common languages bundled by default, small, fast, and it emits hljs-* classes you style with a stylesheet. @shikijs/rehype (4.4.3) uses the same TextMate grammars as VS Code and emits inline styles or CSS variables, which is far more faithful on languages that defeat regex highlighters.

For a build-time pipeline where bytes do not reach the browser, Shiki wins on fidelity and costs you nothing at runtime. For anything highlighting in the browser, lowlight is the smaller thing to ship. rehype-pretty-code sits on top of Shiki and gives you line highlighting and titles without writing a plugin at all — if all you want is a nicer code block, start there and skip this post.

This blog takes the other route entirely: Shiki at build time with the wrapper added by the renderer directly, no tree transformation at all. Same problem, different seam. The plugin approach is more powerful and more machinery; which one is right depends on whether you need to transform anything else while you are in there.

Where this stops being the right tool#

Once you can rewrite the tree, a code block stops being a rendering detail and becomes a component you control. Line numbers, filename headers, diff markers, collapsible blocks — same technique, different node shape.

What the tree cannot give you is behaviour. Every interactive thing this plugin adds has to be met halfway by a script the consumer loads, because the output is a string and there is nothing left to bind to. That is the same lesson as building a shop with no component model: you learn precisely what components are for by doing without one. It is also the same shape as audio state in React — something outside your abstraction owns the state, and the abstraction has to be honest about that rather than pretend.

If your renderer already produces components — MDX, or a Svelte or React markdown pipeline — do this as a component and skip the plugin. The tree transformation earns its keep when the output really is a string, or when you have several transformations to make and want them in one pass over one tree.

Common questions#

Should I use class or className in a HAST node?#

className, as an array of tokens. Both serialise to the same class attribute, so a raw class string looks like it works, but only the serialiser normalises it — every plugin, selector and utility that inspects the tree reads properties.className and will not find it. Using hastscript's h('div.foo') gets it right without you thinking about it.

Why does my rehype plugin loop forever?#

Almost certainly because you replaced a node without returning SKIP. unist-util-visit still walks the replaced node's descendants, so if your new wrapper contains something the same visitor matches, it wraps it again. Return [SKIP, index + 1] from the visitor after any replacement.

How do I get the filename out of a fence info string like ts title="app.ts"?#

Not from HAST. mdast-util-to-hast drops fence metadata and gives you only className: ['language-ts']. Read node.meta on the MDAST code node in a remark-stage plugin and pass it through as a data attribute before remark-rehype runs.

Will a copy button with an inline onclick work under a Content-Security-Policy?#

Not under a strict one. Inline handler attributes count as inline script, so you need 'unsafe-inline' or 'unsafe-hashes' with a SHA-256 of the exact attribute text. Nonces cannot help, because there is no <script> element to put one on. Emit a data- attribute and use one delegated listener instead.

Does wrapping <pre> break rehype-highlight?#

No. Highlighting matches a <code> whose parent is <pre>, and wrapping the <pre> leaves that relationship untouched. Plugin order between the two does not matter. What would break it is wrapping the <code> rather than the <pre>.

Comments

Sign in to comment. No password required.

Loading comments…