fix(penpal): render pending highlights via rehype instead of DOM mutation

apps/penpal/frontend/src/components/MarkdownViewer.test.tsx

@@ -76,4 +76,31 @@ describe('MarkdownViewer', () => {
     expect(mermaidContainer).toBeDefined();
     expect(mermaidContainer?.getAttribute('data-mermaid-source')).toContain('graph TD');
   });
+
+  it('renders comment highlights via rehype plugin', () => {
+    const md = 'Hello world';
+    const highlights = [
+      { threadId: 't1', selectedText: 'world', startLine: 1 },
+    ];
+    const { container } = render(
+      <MarkdownViewer content={md} rawMarkdown={md} highlights={highlights} />,
+    );
+    const mark = container.querySelector('mark.comment-highlight');
+    expect(mark).not.toBeNull();
+    expect(mark?.textContent).toBe('world');
+    expect(mark?.getAttribute('data-thread-id')).toBe('t1');
+  });
+
+  it('renders pending highlights with pending-highlight class', () => {
+    const md = 'Hello world';
+    const highlights = [
+      { threadId: 'pending', selectedText: 'world', startLine: 1, pending: true },
+    ];
+    const { container } = render(
+      <MarkdownViewer content={md} rawMarkdown={md} highlights={highlights} />,
+    );
+    const mark = container.querySelector('mark.pending-highlight');
+    expect(mark).not.toBeNull();
+    expect(mark?.classList.contains('comment-highlight')).toBe(true);
+  });
 });

apps/penpal/frontend/src/components/SelectionToolbar.tsx

@@ -73,6 +73,8 @@ function computeAnchor(
     }
   }
 
+  anchor.occurrenceIndex = occurrenceIndex;
+
   // Find nth occurrence helper
   function findNthOccurrence(haystack: string, needle: string, n: number): number {
     let pos = 0;
@@ -128,54 +130,6 @@ function computeAnchor(
   return anchor;
 }
 
-/**
- * Applies pending highlight marks to the current selection.
- */
-export function applyPendingHighlight(sel: Selection, contentEl: HTMLElement) {
-  removePendingHighlight();
-  if (!sel.rangeCount) return;
-  const range = sel.getRangeAt(0);
-  if (!contentEl.contains(range.commonAncestorContainer)) return;
-
-  const treeWalker = document.createTreeWalker(contentEl, NodeFilter.SHOW_TEXT);
-  const textNodes: Text[] = [];
-  while (treeWalker.nextNode()) {
-    if (range.intersectsNode(treeWalker.currentNode)) {
-      textNodes.push(treeWalker.currentNode as Text);
-    }
-  }
-
-  for (let i = textNodes.length - 1; i >= 0; i--) {
-    const node = textNodes[i];
-    const start = node === range.startContainer ? range.startOffset : 0;
-    const end = node === range.endContainer ? range.endOffset : node.nodeValue!.length;
-    if (start >= end) continue;
-
-    const mark = document.createElement('mark');
-    mark.className = 'pending-highlight';
-
-    if (start === 0 && end === node.nodeValue!.length) {
-      node.parentNode!.insertBefore(mark, node);
-      mark.appendChild(node);
-    } else {
-      const subRange = document.createRange();
-      subRange.setStart(node, start);
-      subRange.setEnd(node, end);
-      subRange.surroundContents(mark);
-    }
-  }
-  sel.removeAllRanges();
-}
-
-export function removePendingHighlight() {
-  document.querySelectorAll('.pending-highlight').forEach((mark) => {
-    const parent = mark.parentNode!;
-    while (mark.firstChild) parent.insertBefore(mark.firstChild, mark);
-    parent.removeChild(mark);
-    parent.normalize();
-  });
-}
-
 /**
  * Finds which mermaid container index corresponds to a given startLine
  * by counting ```mermaid fences in the raw markdown.
@@ -322,7 +276,7 @@ export default function SelectionToolbar({
     if (!selectedText) return;
 
     const anchor = computeAnchor(sel, selectedText, rawMarkdown, contentRef.current);
-    applyPendingHighlight(sel, contentRef.current);
+    sel.removeAllRanges();
     setVisible(false);
     onComment(anchor, selectedText);
   };

apps/penpal/frontend/src/components/rehypeCommentHighlights.test.ts

@@ -0,0 +1,141 @@
+import { describe, it, expect } from 'vitest';
+import rehypeCommentHighlights from './rehypeCommentHighlights';
+import type { Root, Element, Text } from 'hast';
+
+/** Build a minimal HAST tree: <p data-source-line="N">...text...</p> */
+function makeTree(line: number, text: string): Root {
+  return {
+    type: 'root',
+    children: [
+      {
+        type: 'element',
+        tagName: 'p',
+        properties: {},
+        children: [{ type: 'text', value: text } as Text],
+        position: { start: { line, column: 1, offset: 0 }, end: { line, column: 1 + text.length, offset: text.length } },
+      } as Element,
+    ],
+  };
+}
+
+/** Find all <mark> elements in a tree */
+function findMarks(node: Root | Element): Element[] {
+  const marks: Element[] = [];
+  for (const child of ('children' in node ? node.children : [])) {
+    if (child.type === 'element') {
+      if (child.tagName === 'mark') marks.push(child);
+      marks.push(...findMarks(child));
+    }
+  }
+  return marks;
+}
+
+describe('rehypeCommentHighlights', () => {
+  it('wraps matching text in a <mark> with comment-highlight class', () => {
+    const tree = makeTree(1, 'Hello world');
+    const transform = rehypeCommentHighlights({
+      highlights: [{ threadId: 't1', selectedText: 'world', startLine: 1 }],
+    });
+    transform(tree);
+    const marks = findMarks(tree);
+    expect(marks).toHaveLength(1);
+    expect((marks[0].properties as Record<string, unknown>).className).toEqual(['comment-highlight']);
+    expect((marks[0].properties as Record<string, unknown>).dataThreadId).toBe('t1');
+    expect((marks[0].children[0] as Text).value).toBe('world');
+  });
+
+  it('does not wrap text when selectedText is not found', () => {
+    const tree = makeTree(1, 'Hello world');
+    const transform = rehypeCommentHighlights({
+      highlights: [{ threadId: 't1', selectedText: 'missing', startLine: 1 }],
+    });
+    transform(tree);
+    expect(findMarks(tree)).toHaveLength(0);
+  });
+
+  it('does not wrap text when startLine does not match', () => {
+    const tree = makeTree(1, 'Hello world');
+    const transform = rehypeCommentHighlights({
+      highlights: [{ threadId: 't1', selectedText: 'world', startLine: 99 }],
+    });
+    transform(tree);
+    expect(findMarks(tree)).toHaveLength(0);
+  });
+
+  it('adds pending-highlight class when pending is true', () => {
+    const tree = makeTree(1, 'Hello world');
+    const transform = rehypeCommentHighlights({
+      highlights: [{ threadId: 'pending', selectedText: 'world', startLine: 1, pending: true }],
+    });
+    transform(tree);
+    const marks = findMarks(tree);
+    expect(marks).toHaveLength(1);
+    expect((marks[0].properties as Record<string, unknown>).className).toEqual(['comment-highlight', 'pending-highlight']);
+  });
+
+  it('does not add pending-highlight class when pending is false/undefined', () => {
+    const tree = makeTree(1, 'Hello world');
+    const transform = rehypeCommentHighlights({
+      highlights: [{ threadId: 't1', selectedText: 'world', startLine: 1 }],
+    });
+    transform(tree);
+    const marks = findMarks(tree);
+    expect(marks).toHaveLength(1);
+    expect((marks[0].properties as Record<string, unknown>).className).toEqual(['comment-highlight']);
+  });
+
+  it('splits text node correctly around the match', () => {
+    const tree = makeTree(1, 'The quick brown fox');
+    const transform = rehypeCommentHighlights({
+      highlights: [{ threadId: 't1', selectedText: 'quick', startLine: 1 }],
+    });
+    transform(tree);
+    const p = tree.children[0] as Element;
+    // Should be: "The " + <mark>quick</mark> + " brown fox"
+    expect(p.children).toHaveLength(3);
+    expect((p.children[0] as Text).value).toBe('The ');
+    expect((p.children[1] as Element).tagName).toBe('mark');
+    expect((p.children[2] as Text).value).toBe(' brown fox');
+  });
+
+  it('returns tree unchanged when highlights array is empty', () => {
+    const tree = makeTree(1, 'Hello world');
+    const original = JSON.stringify(tree);
+    const transform = rehypeCommentHighlights({ highlights: [] });
+    transform(tree);
+    expect(JSON.stringify(tree)).toBe(original);
+  });
+
+  it('handles multiple highlights on different lines', () => {
+    const tree: Root = {
+      type: 'root',
+      children: [
+        {
+          type: 'element',
+          tagName: 'p',
+          properties: {},
+          children: [{ type: 'text', value: 'First line' } as Text],
+          position: { start: { line: 1, column: 1, offset: 0 }, end: { line: 1, column: 11, offset: 10 } },
+        } as Element,
+        {
+          type: 'element',
+          tagName: 'p',
+          properties: {},
+          children: [{ type: 'text', value: 'Second line' } as Text],
+          position: { start: { line: 3, column: 1, offset: 12 }, end: { line: 3, column: 12, offset: 23 } },
+        } as Element,
+      ],
+    };
+    const transform = rehypeCommentHighlights({
+      highlights: [
+        { threadId: 't1', selectedText: 'First', startLine: 1 },
+        { threadId: 't2', selectedText: 'Second', startLine: 3 },
+      ],
+    });
+    transform(tree);
+    const marks = findMarks(tree);
+    expect(marks).toHaveLength(2);
+    expect((marks[0].properties as Record<string, unknown>).dataThreadId).toBe('t1');
+    expect((marks[1].properties as Record<string, unknown>).dataThreadId).toBe('t2');
+  });
+});

apps/penpal/frontend/src/components/rehypeCommentHighlights.ts

@@ -5,6 +5,8 @@ export interface ThreadHighlight {
   threadId: string;
   selectedText: string;
   startLine: number;
+  occurrenceIndex?: number;
+  pending?: boolean;
 }
 
 interface Options {
@@ -87,7 +89,22 @@ function applyHighlight(element: Element, highlight: ThreadHighlight) {
   // Normalize whitespace for matching (mirrors DOM-based approach)
   const normalizedText = text.replace(/\s+/g, ' ');
   const normalizedSelected = highlight.selectedText.replace(/\s+/g, ' ');
-  const matchIndex = normalizedText.indexOf(normalizedSelected);
+
+  // Use occurrenceIndex to find the Nth match within the block,
+  // falling back to the first occurrence if the index is not set or not found.
+  let matchIndex = -1;
+  const targetOccurrence = highlight.occurrenceIndex ?? 0;
+  let pos = 0;
+  for (let i = 0; i <= targetOccurrence; i++) {
+    const found = normalizedText.indexOf(normalizedSelected, pos);
+    if (found === -1) break;
+    if (i === targetOccurrence) {
+      matchIndex = found;
+    }
+    pos = found + 1;
+  }
+  // Fall back to first occurrence if Nth not found
+  if (matchIndex === -1) matchIndex = normalizedText.indexOf(normalizedSelected);
   if (matchIndex === -1) return;
 
   const matchStart = matchIndex;
@@ -139,7 +156,7 @@ function applyHighlight(element: Element, highlight: ThreadHighlight) {
       type: 'element',
       tagName: 'mark',
       properties: {
-        className: ['comment-highlight'],
+        className: highlight.pending ? ['comment-highlight', 'pending-highlight'] : ['comment-highlight'],
         dataThreadId: highlight.threadId,
       },
       children: [