Support code block // [!code -- or ++] notation (#4248)

viktorrenkema · web-flow · commit d994875bea2f · 2026-05-12T15:39:21.000+02:00
diff --git a/packages/gitbook/src/components/DocumentView/CodeBlock/CodeBlockRenderer.tsx b/packages/gitbook/src/components/DocumentView/CodeBlock/CodeBlockRenderer.tsx
@@ -114,27 +114,30 @@ function CodeHighlightLine(props: {
     withLineNumbers: boolean;
 }) {
     const { line, isLast, withLineNumbers, bg, fg } = props;
+    const lineStyle = {
+        color: fg?.color,
+        ...fg?.vars,
+        backgroundColor: bg?.color,
+        ...bg?.vars,
+    };
     return (
         <span
-            className={tcls('highlight-line', line.highlighted && 'highlighted')}
-            style={{
-                color: fg?.color,
-                ...fg?.vars,
-                backgroundColor: bg?.color,
-                ...bg?.vars,
-            }}
-        >
-            {withLineNumbers && (
-                <span
-                    className="highlight-line-number"
-                    style={{
-                        color: fg?.color,
-                        ...fg?.vars,
-                        backgroundColor: bg?.color,
-                        ...bg?.vars,
-                    }}
-                />
+            className={tcls(
+                'highlight-line',
+                line.diff === 'added' && 'diff-added',
+                line.diff === 'deleted' && 'diff-deleted',
+                line.highlighted && 'highlighted'
             )}
+            aria-label={
+                line.diff === 'added'
+                    ? 'Added line'
+                    : line.diff === 'deleted'
+                      ? 'Removed line'
+                      : undefined
+            }
+            style={lineStyle}
+        >
+            {withLineNumbers && <span className="highlight-line-number" style={lineStyle} />}
             <span className="highlight-line-content">
                 <CodeHighlightTokens tokens={line.tokens} />
                 {!isLast && '\n'}
diff --git a/packages/gitbook/src/components/DocumentView/CodeBlock/highlight.test.ts b/packages/gitbook/src/components/DocumentView/CodeBlock/highlight.test.ts
@@ -1,7 +1,13 @@
 import { expect, it } from 'bun:test';
 import type { DocumentBlockCode } from '@gitbook/api';
 
-import { type RenderedInline, getInlines, highlight } from './highlight';
+import {
+    type HighlightLine,
+    type HighlightToken,
+    type RenderedInline,
+    getInlines,
+    highlight,
+} from './highlight';
 
 async function highlightWithInlines(block: DocumentBlockCode) {
     const inlines: RenderedInline[] = getInlines(block).map((inline) => ({
@@ -690,6 +696,138 @@ it('should support multiple code tokens in an annotation', async () => {
     ]);
 });
 
+function joinLineContent(line: HighlightLine): string {
+    const visit = (tokens: HighlightToken[]): string =>
+        tokens
+            .map((t) => {
+                if (t.type === 'plain') return t.content;
+                if (t.type === 'shiki') return t.token.content;
+                return visit(t.children);
+            })
+            .join('');
+    return visit(line.tokens);
+}
+
+function singleLineBlock(syntax: string | undefined, text: string): DocumentBlockCode {
+    return {
+        object: 'block',
+        type: 'code',
+        data: syntax ? { syntax } : {},
+        nodes: [
+            {
+                object: 'block',
+                type: 'code-line',
+                data: {},
+                nodes: [
+                    {
+                        object: 'text',
+                        leaves: [{ object: 'leaf', marks: [], text }],
+                    },
+                ],
+            },
+        ],
+    };
+}
+
+it('classifies and strips trailing // [!code ++] in JS', async () => {
+    const lines = await highlightWithInlines(
+        singleLineBlock('javascript', 'const a = 1 // [!code ++]')
+    );
+    expect(lines).toHaveLength(1);
+    expect(lines[0]!.diff).toBe('added');
+    expect(joinLineContent(lines[0]!)).toBe('const a = 1');
+});
+
+it('classifies and strips trailing # [!code --] in Python', async () => {
+    const lines = await highlightWithInlines(singleLineBlock('python', 'x = 1 # [!code --]'));
+    expect(lines[0]!.diff).toBe('deleted');
+    expect(joinLineContent(lines[0]!)).toBe('x = 1');
+});
+
+it('classifies and strips <!-- [!code ++] --> in HTML', async () => {
+    const lines = await highlightWithInlines(
+        singleLineBlock('html', '<div></div> <!-- [!code ++] -->')
+    );
+    expect(lines[0]!.diff).toBe('added');
+    expect(joinLineContent(lines[0]!)).toBe('<div></div>');
+});
+
+it('classifies and strips /* [!code --] */ in CSS', async () => {
+    const lines = await highlightWithInlines(
+        singleLineBlock('css', '.a { color: red; } /* [!code --] */')
+    );
+    expect(lines[0]!.diff).toBe('deleted');
+    expect(joinLineContent(lines[0]!)).toBe('.a { color: red; }');
+});
+
+it('does not classify when marker is not at end of line', async () => {
+    const lines = await highlightWithInlines(
+        singleLineBlock('javascript', 'const x = 1 // [!code ++] trailing')
+    );
+    expect(lines[0]!.diff).toBeNull();
+    expect(joinLineContent(lines[0]!)).toBe('const x = 1 // [!code ++] trailing');
+});
+
+it('returns diff: null for lines without a marker', async () => {
+    const lines = await highlightWithInlines(singleLineBlock('javascript', 'console.log("hi")'));
+    expect(lines[0]!.diff).toBeNull();
+});
+
+it('classifies and strips marker via plainHighlighting fallback', async () => {
+    const lines = await highlightWithInlines(
+        singleLineBlock(undefined, 'plain text // [!code ++]')
+    );
+    expect(lines[0]!.diff).toBe('added');
+    expect(joinLineContent(lines[0]!)).toBe('plain text');
+});
+
+it('preserves inline annotation when marker is stripped', async () => {
+    const tokens = await highlightWithInlines({
+        object: 'block',
+        type: 'code',
+        data: { syntax: 'javascript' },
+        nodes: [
+            {
+                object: 'block',
+                type: 'code-line',
+                data: {},
+                nodes: [
+                    {
+                        object: 'text',
+                        leaves: [{ object: 'leaf', marks: [], text: 'console.' }],
+                    },
+                    {
+                        object: 'inline',
+                        type: 'annotation',
+                        nodes: [
+                            {
+                                object: 'text',
+                                leaves: [{ object: 'leaf', marks: [], text: 'log' }],
+                            },
+                        ],
+                        isVoid: false,
+                        fragments: [],
+                    },
+                    {
+                        object: 'text',
+                        leaves: [{ object: 'leaf', marks: [], text: '("Hi") // [!code ++]' }],
+                    },
+                ],
+            },
+        ],
+    });
+
+    expect(tokens[0]!.diff).toBe('added');
+    expect(joinLineContent(tokens[0]!)).toBe('console.log("Hi")');
+    // inline annotation around "log" must be preserved
+    const hasAnnotation = tokens[0]!.tokens.some(
+        (t) =>
+            t.type === 'annotation' &&
+            t.children.some((c) => c.type === 'shiki' && c.token.content === 'log')
+    );
+    expect(hasAnnotation).toBe(true);
+});
+
 it('should handle \\r', async () => {
     const tokens = await highlightWithInlines({
         object: 'block',
diff --git a/packages/gitbook/src/components/DocumentView/CodeBlock/highlight.ts b/packages/gitbook/src/components/DocumentView/CodeBlock/highlight.ts
@@ -34,11 +34,114 @@ export type HighlightTheme = {
     lines: HighlightLine[];
 };
 
+export type LineDiffNotation = 'added' | 'deleted';
+
 export type HighlightLine = {
     highlighted: boolean;
+    diff: LineDiffNotation | null;
     tokens: HighlightToken[];
 };
 
+/**
+ * Detects an in-source diff notation marker at the end of a line, e.g.
+ * `// [!code ++]`, `# [!code --]`, `<!-- [!code ++] -->`, `/* [!code --] *​/`.
+ * Mirror of the gitbook-x parser; keep regex byte-for-byte identical.
+ */
+const NOTATION_PATTERN =
+    /[ \t]*(?:(?:\/\/|#|--|;)\s*\[!code\s+(\+\+|--)\]|<!--\s*\[!code\s+(\+\+|--)\]\s*-->|\/\*\s*\[!code\s+(\+\+|--)\]\s*\*\/)\s*$/;
+
+export function parseDiffNotation(
+    line: string
+): { diff: LineDiffNotation; markerStart: number } | null {
+    const match = NOTATION_PATTERN.exec(line);
+    if (!match) {
+        return null;
+    }
+    const variant = match[1] ?? match[2] ?? match[3];
+    return {
+        diff: variant === '++' ? 'added' : 'deleted',
+        markerStart: match.index,
+    };
+}
+
+/**
+ * Truncate a sequence of HighlightTokens so only the first `maxLen` characters
+ * (counted across all tokens, recursing into annotations) remain. Used to strip
+ * trailing diff-notation markers from rendered output.
+ */
+export function truncateHighlightTokens(
+    tokens: HighlightToken[],
+    maxLen: number
+): HighlightToken[] {
+    const out: HighlightToken[] = [];
+    let remaining = maxLen;
+    for (const token of tokens) {
+        if (remaining <= 0) {
+            break;
+        }
+        const len = highlightTokenLength(token);
+        if (len <= remaining) {
+            out.push(token);
+            remaining -= len;
+            continue;
+        }
+        out.push(sliceHighlightToken(token, remaining));
+        remaining = 0;
+    }
+    return out;
+}
+
+function highlightTokenLength(token: HighlightToken): number {
+    switch (token.type) {
+        case 'plain':
+            return token.content.length;
+        case 'shiki':
+            return token.token.content.length;
+        case 'annotation':
+            return token.children.reduce((acc, child) => acc + highlightTokenLength(child), 0);
+    }
+}
+
+/**
+ * Concatenate the text content of a sequence of HighlightTokens, recursing
+ * into annotation children. Mirror of {@link highlightTokenLength}.
+ */
+export function getHighlightTokensText(tokens: HighlightToken[]): string {
+    return tokens
+        .map((token) => {
+            switch (token.type) {
+                case 'plain':
+                    return token.content;
+                case 'shiki':
+                    return token.token.content;
+                case 'annotation':
+                    return getHighlightTokensText(token.children);
+            }
+        })
+        .join('');
+}
+
+function sliceHighlightToken(token: HighlightToken, maxLen: number): HighlightToken {
+    switch (token.type) {
+        case 'plain':
+            return { type: 'plain', content: token.content.slice(0, maxLen) };
+        case 'shiki': {
+            const inner = token.token as ThemedToken & { start?: number; end?: number };
+            const newContent = inner.content.slice(0, maxLen);
+            const newToken: ThemedToken & { start?: number; end?: number } = {
+                ...inner,
+                content: newContent,
+            };
+            if (typeof inner.start === 'number') {
+                newToken.end = inner.start + newContent.length;
+            }
+            return { type: 'shiki', token: newToken };
+        }
+        case 'annotation':
+            return { ...token, children: truncateHighlightTokens(token.children, maxLen) };
+    }
+}
+
 export type HighlightToken =
     | { type: 'plain'; content: string }
     | { type: 'shiki'; token: ThemedToken }
@@ -136,6 +239,9 @@ export async function highlight(
             const lineBlock = block.nodes[index];
             const result: HighlightToken[] = [];
 
+            const lineText = tokens.map((token) => token.content).join('');
+            const notation = parseDiffNotation(lineText);
+
             const eatToken = (): PositionedToken | null => {
                 const token = tokens.shift();
                 if (token) {
@@ -152,9 +258,14 @@ export async function highlight(
 
             currentIndex += 1; // for the \n
 
+            const finalTokens = notation
+                ? truncateHighlightTokens(result, notation.markerStart)
+                : result;
+
             return {
                 highlighted: Boolean(lineBlock?.data.highlighted),
-                tokens: result,
+                diff: notation?.diff ?? null,
+                tokens: finalTokens,
             };
         }),
     };
diff --git a/packages/gitbook/src/components/DocumentView/CodeBlock/plain-highlight.ts b/packages/gitbook/src/components/DocumentView/CodeBlock/plain-highlight.ts
@@ -3,7 +3,14 @@ import type { CustomizationThemedCodeTheme, DocumentBlockCode } from '@gitbook/a
 import { getNodeText } from '@/lib/document';
 import { bundledThemesInfo } from 'shiki/themes';
 import { customThemes } from './customThemes';
-import type { HighlightTheme, HighlightToken, RenderedInline } from './highlight';
+import {
+    type HighlightTheme,
+    type HighlightToken,
+    type RenderedInline,
+    getHighlightTokensText,
+    parseDiffNotation,
+    truncateHighlightTokens,
+} from './highlight';
 
 /**
  * Parse a code block without highlighting it.
@@ -62,9 +69,14 @@ export function plainHighlight(
                 };
             });
 
+            // Detect diff notation against the built tokens (not the raw nodes)
+            // so any evaluated inline expressions are included in the offset math.
+            const notation = parseDiffNotation(getHighlightTokensText(tokens));
+
             return {
                 highlighted: Boolean(lineBlock.data.highlighted),
-                tokens,
+                diff: notation?.diff ?? null,
+                tokens: notation ? truncateHighlightTokens(tokens, notation.markerStart) : tokens,
             };
         }),
     };
diff --git a/packages/gitbook/src/components/RootLayout/globals.css b/packages/gitbook/src/components/RootLayout/globals.css
