fix(ui): preserve Markdown code selections (#2449)

* fix(ui): preserve markdown code selections

* test(ui): clarify markdown selection fixtures

* fix(ui): preserve highlighted code languages

* fix(ui): preserve markdown block boundaries

packages/ui/src/components/chat/message/TextSelectionMenu.tsx

@@ -15,6 +15,7 @@ import { resolveProjectForSessionDirectory } from '@/lib/projectResolution';
 import { useEffectiveDirectory } from '@/hooks/useEffectiveDirectory';
 import { isVSCodeRuntime } from '@/lib/desktop';
 import { useI18n } from '@/lib/i18n';
+import { rangeToMarkdown, trimSelectionValue, wrapMarkdownSelectionForChat } from './selectionMarkdown';
 
 interface TextSelectionMenuProps {
   containerRef: React.RefObject<HTMLElement | null>;
@@ -44,167 +45,6 @@ const appendDistilledInsightToNotes = (existingNotes: string, insight: string):
 
 const DESKTOP_MENU_SIDE_MARGIN_PX = 8;
 const DESKTOP_MENU_FALLBACK_WIDTH_PX = 280;
-const BLOCK_TAGS = new Set([
-  'address', 'article', 'aside', 'blockquote', 'dd', 'div', 'dl', 'dt',
-  'fieldset', 'figcaption', 'figure', 'footer', 'form', 'h1', 'h2', 'h3',
-  'h4', 'h5', 'h6', 'header', 'hr', 'li', 'main', 'nav', 'ol', 'p', 'pre',
-  'section', 'table', 'ul',
-]);
-
-const normalizeLineBreaks = (value: string): string => value.replace(/\r\n?/g, '\n');
-
-const trimSelectionValue = (value: string): string => normalizeLineBreaks(value).trim();
-
-const textToMarkdownInline = (value: string): string => value.replace(/\s+/g, ' ').trim();
-
-const renderInlineMarkdownNode = (node: Node): string => {
-  if (node.nodeType === Node.TEXT_NODE) {
-    return textToMarkdownInline(node.textContent || '');
-  }
-
-  if (node.nodeType !== Node.ELEMENT_NODE) {
-    return '';
-  }
-
-  const element = node as HTMLElement;
-  const tag = element.tagName.toLowerCase();
-  const childText = Array.from(element.childNodes)
-    .map((child) => renderInlineMarkdownNode(child))
-    .join('')
-    .replace(/\s+/g, ' ')
-    .trim();
-
-  if (!childText && tag !== 'br') {
-    return '';
-  }
-
-  if (tag === 'br') return '\n';
-  if (tag === 'strong' || tag === 'b') return `**${childText}**`;
-  if (tag === 'em' || tag === 'i') return `*${childText}*`;
-  if (tag === 'code') return `\`${childText.replace(/`/g, '\\`')}\``;
-  if (tag === 'a') {
-    const href = element.getAttribute('href');
-    return href ? `[${childText}](${href})` : childText;
-  }
-
-  return childText;
-};
-
-const renderListMarkdown = (list: HTMLElement, ordered: boolean): string => {
-  const items = Array.from(list.children).filter(
-    (child): child is HTMLElement => child instanceof HTMLElement && child.tagName.toLowerCase() === 'li'
-  );
-
-  return items
-    .map((item, index) => {
-      const prefix = ordered ? `${index + 1}. ` : '- ';
-      const body = Array.from(item.childNodes)
-        .map((child) => renderInlineMarkdownNode(child))
-        .join('')
-        .replace(/\s+/g, ' ')
-        .trim();
-      return body ? `${prefix}${body}` : '';
-    })
-    .filter(Boolean)
-    .join('\n');
-};
-
-const renderBlockMarkdownNode = (node: Node): string => {
-  if (node.nodeType === Node.TEXT_NODE) {
-    return trimSelectionValue(node.textContent || '');
-  }
-
-  if (node.nodeType !== Node.ELEMENT_NODE) {
-    return '';
-  }
-
-  const element = node as HTMLElement;
-  const tag = element.tagName.toLowerCase();
-
-  if (tag === 'pre') {
-    const codeElement = element.querySelector('code');
-    const languageClass = codeElement?.className || '';
-    const language = (languageClass.match(/language-([\w-]+)/)?.[1] || '').trim();
-    const code = normalizeLineBreaks(codeElement?.textContent || element.textContent || '').replace(/\n$/, '');
-    return `\`\`\`${language}\n${code}\n\`\`\``;
-  }
-
-  if (tag === 'code') {
-    const code = normalizeLineBreaks(element.textContent || '').trim();
-    return code ? `\`${code.replace(/`/g, '\\`')}\`` : '';
-  }
-
-  if (tag === 'ul') return renderListMarkdown(element, false);
-  if (tag === 'ol') return renderListMarkdown(element, true);
-
-  if (tag === 'blockquote') {
-    const content = trimSelectionValue(
-      Array.from(element.childNodes).map((child) => renderBlockMarkdownNode(child)).join('\n')
-    );
-    return content
-      .split('\n')
-      .filter((line) => line.length > 0)
-      .map((line) => `> ${line}`)
-      .join('\n');
-  }
-
-  if (/^h[1-6]$/.test(tag)) {
-    const level = Number.parseInt(tag[1], 10);
-    const text = trimSelectionValue(Array.from(element.childNodes).map((child) => renderInlineMarkdownNode(child)).join(''));
-    return text ? `${'#'.repeat(level)} ${text}` : '';
-  }
-
-  if (tag === 'p' || tag === 'div' || tag === 'li') {
-    return trimSelectionValue(Array.from(element.childNodes).map((child) => renderInlineMarkdownNode(child)).join(''));
-  }
-
-  const blockChildren = Array.from(element.childNodes)
-    .map((child) => renderBlockMarkdownNode(child))
-    .filter((child) => child.length > 0);
-  if (blockChildren.length > 0) {
-    return blockChildren.join('\n\n');
-  }
-
-  return trimSelectionValue(Array.from(element.childNodes).map((child) => renderInlineMarkdownNode(child)).join(''));
-};
-
-const isInlineSelectionFragment = (fragment: DocumentFragment): boolean => {
-  return Array.from(fragment.childNodes).every((node) => {
-    if (node.nodeType === Node.TEXT_NODE) {
-      return true;
-    }
-    if (node.nodeType !== Node.ELEMENT_NODE) {
-      return true;
-    }
-
-    const element = node as HTMLElement;
-    return !BLOCK_TAGS.has(element.tagName.toLowerCase());
-  });
-};
-
-const rangeToMarkdown = (range: Range, plainText: string): string => {
-  const fragment = range.cloneContents();
-
-  if (isInlineSelectionFragment(fragment)) {
-    const inlineMarkdown = trimSelectionValue(
-      Array.from(fragment.childNodes)
-        .map((node) => renderInlineMarkdownNode(node))
-        .join('')
-    );
-    if (inlineMarkdown) {
-      return inlineMarkdown;
-    }
-  }
-
-  const markdown = Array.from(fragment.childNodes)
-    .map((node) => renderBlockMarkdownNode(node))
-    .filter((value) => value.length > 0)
-    .join('\n\n')
-    .trim();
-
-  return markdown || trimSelectionValue(plainText);
-};
-
 export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerRef }) => {
   const { t } = useI18n();
   const [position, setPosition] = React.useState<MenuPosition>({ x: 0, y: 0, show: false });
@@ -462,7 +302,7 @@ export const TextSelectionMenu: React.FC<TextSelectionMenuProps> = ({ containerR
   const handleAddToChat = React.useCallback(() => {
     if (!selectedTextMarkdown) return;
 
-    const markdownBlock = `\`\`\`md\n${selectedTextMarkdown}\n\`\`\``;
+    const markdownBlock = wrapMarkdownSelectionForChat(selectedTextMarkdown);
     setPendingInputText(markdownBlock, 'append');
     
     hideMenu();

packages/ui/src/components/chat/message/selectionMarkdown.test.ts

@@ -0,0 +1,183 @@
+import { describe, expect, test } from 'bun:test';
+
+import {
+  formatCodeSelectionMarkdown,
+  selectionNodesToMarkdown,
+  trimSelectionValue,
+  wrapMarkdownSelectionForChat,
+} from './selectionMarkdown';
+
+type TestNode =
+  | { type: 'text'; value: string }
+  | {
+      type: 'element';
+      tag: string;
+      className: string;
+      href: string;
+      component: string;
+      markdownLanguage: string;
+      isMarkdownBlock: boolean;
+      isCodeLines: boolean;
+      isCodeLineNumber: boolean;
+      children: TestNode[];
+    };
+
+const multiline = (...lines: string[]): string => lines.join('\n');
+const text = (value: string): TestNode => ({ type: 'text', value });
+const element = (
+  tag: string,
+  children: TestNode[],
+  options: Partial<Omit<Extract<TestNode, { type: 'element' }>, 'type' | 'tag' | 'children'>> = {},
+): TestNode => ({
+  type: 'element',
+  tag,
+  className: '',
+  href: '',
+  component: '',
+  markdownLanguage: '',
+  isMarkdownBlock: false,
+  isCodeLines: false,
+  isCodeLineNumber: false,
+  children,
+  ...options,
+});
+
+const codeLine = (number: number, content: string): TestNode => element('span', [
+  element('span', [text(String(number))], { isCodeLineNumber: true }),
+  element('span', [text(content)]),
+]);
+
+const markdownBlock = (children: TestNode[]): TestNode => element('div', children, { isMarkdownBlock: true });
+
+const codeWrapper = (lines: string[], language = 'ts'): TestNode => element('div', [
+  element('div', [text(language)]),
+  element('div', [
+    element('pre', [
+      element('code', lines.flatMap((line, index) => [
+        codeLine(index + 12, line),
+        ...(index < lines.length - 1 ? [element('span', [text('\n')])] : []),
+      ]), { isCodeLines: true }),
+    ], { markdownLanguage: language }),
+  ]),
+], { component: 'markdown-code' });
+
+describe('selectionNodesToMarkdown', () => {
+  test('serializes a complete grid code block without its header or line numbers', () => {
+    expect(selectionNodesToMarkdown([codeWrapper(['range.cloneContents()', 'next()'])], '')).toBe(multiline(
+      '```ts',
+      'range.cloneContents()',
+      'next()',
+      '```',
+    ));
+  });
+
+  test('preserves a code block nested between production Markdown block wrappers', () => {
+    const nodes = [
+      markdownBlock([element('p', [
+        text('Method:'),
+        element('code', [text('Selection.toString()')]),
+        text('. Add to Chat uses a cloned range.'),
+      ])]),
+      markdownBlock([codeWrapper(['range.cloneContents()'])]),
+      markdownBlock([element('p', [text('Following explanation')])]),
+    ];
+
+    expect(selectionNodesToMarkdown(nodes, '')).toBe(multiline(
+      'Method:`Selection.toString()`. Add to Chat uses a cloned range.',
+      '',
+      '```ts',
+      'range.cloneContents()',
+      '```',
+      '',
+      'Following explanation',
+    ));
+  });
+
+  test('preserves a partial code block selected before prose', () => {
+    expect(selectionNodesToMarkdown([
+      codeWrapper(['range.cloneContents()']),
+      element('p', [text('Following explanation')]),
+    ], '')).toBe(multiline(
+      '```ts',
+      'range.cloneContents()',
+      '```',
+      '',
+      'Following explanation',
+    ));
+  });
+});
+
+describe('formatCodeSelectionMarkdown', () => {
+  test('preserves indentation and blank lines', () => {
+    expect(formatCodeSelectionMarkdown(multiline(
+      'if (ready) {',
+      '  run();',
+      '',
+      '  stop();',
+      '}',
+    ), 'ts')).toBe(multiline(
+      '```ts',
+      'if (ready) {',
+      '  run();',
+      '',
+      '  stop();',
+      '}',
+      '```',
+    ));
+  });
+
+  test('normalizes line endings without duplicating a trailing newline', () => {
+    expect(formatCodeSelectionMarkdown('first\r\nsecond\r\n', 'text')).toBe(multiline(
+      '```text',
+      'first',
+      'second',
+      '```',
+    ));
+  });
+
+  test('uses a longer fence when selected code contains backtick fences', () => {
+    expect(formatCodeSelectionMarkdown(multiline(
+      'before',
+      '```',
+      'after',
+    ), 'md')).toBe(multiline(
+      '````md',
+      'before',
+      '```',
+      'after',
+      '````',
+    ));
+  });
+
+  test('preserves punctuation in language identifiers', () => {
+    expect(selectionNodesToMarkdown([codeWrapper(['std::vector<int> values;'], 'c++')], '')).toBe(multiline(
+      '```c++',
+      'std::vector<int> values;',
+      '```',
+    ));
+  });
+});
+
+describe('trimSelectionValue', () => {
+  test('normalizes line endings before trimming the selection', () => {
+    expect(trimSelectionValue('  first\r\nsecond  ')).toBe(multiline('first', 'second'));
+  });
+});
+
+describe('wrapMarkdownSelectionForChat', () => {
+  test('uses a longer outer fence when the selection contains fenced code', () => {
+    const selectedMarkdown = multiline(
+      '```ts',
+      'run();',
+      '```',
+    );
+
+    expect(wrapMarkdownSelectionForChat(selectedMarkdown)).toBe(multiline(
+      '````md',
+      '```ts',
+      'run();',
+      '```',
+      '````',
+    ));
+  });
+});

packages/ui/src/components/chat/message/selectionMarkdown.ts

@@ -0,0 +1,239 @@
+type SelectionNode =
+  | { type: 'text'; value: string }
+  | {
+      type: 'element';
+      tag: string;
+      className: string;
+      href: string;
+      component: string;
+      markdownLanguage: string;
+      isMarkdownBlock: boolean;
+      isCodeLines: boolean;
+      isCodeLineNumber: boolean;
+      children: SelectionNode[];
+    };
+
+const BLOCK_TAGS = new Set([
+  'address', 'article', 'aside', 'blockquote', 'dd', 'div', 'dl', 'dt',
+  'fieldset', 'figcaption', 'figure', 'footer', 'form', 'h1', 'h2', 'h3',
+  'h4', 'h5', 'h6', 'header', 'hr', 'li', 'main', 'nav', 'ol', 'p', 'pre',
+  'section', 'table', 'ul',
+]);
+
+const normalizeLineBreaks = (value: string): string => value.replace(/\r\n?/g, '\n');
+export const trimSelectionValue = (value: string): string => normalizeLineBreaks(value).trim();
+const textToMarkdownInline = (value: string): string => value.replace(/\s+/g, ' ').trim();
+
+const getCodeLanguageFromClassName = (className: string): string => {
+  return (className.match(/language-([\w+#.-]+)/)?.[1] || '').trim();
+};
+
+const getBlockCodeLanguage = (code: HTMLElement): string => {
+  return code.closest('pre')?.getAttribute('data-md-lang')
+    || getCodeLanguageFromClassName(code.className);
+};
+
+const toSelectionNode = (node: Node): SelectionNode | null => {
+  if (node.nodeType === 3) {
+    return { type: 'text', value: node.textContent || '' };
+  }
+  if (node.nodeType !== 1) {
+    return null;
+  }
+
+  const element = node as Element;
+  return {
+    type: 'element',
+    tag: element.tagName.toLowerCase(),
+    className: element.getAttribute('class') || '',
+    href: element.getAttribute('href') || '',
+    component: element.getAttribute('data-component') || '',
+    markdownLanguage: element.getAttribute('data-md-lang') || '',
+    isMarkdownBlock: element.hasAttribute('data-md-block'),
+    isCodeLines: element.hasAttribute('data-md-code-lines'),
+    isCodeLineNumber: element.hasAttribute('data-md-code-line-number'),
+    children: Array.from(element.childNodes)
+      .map((child) => toSelectionNode(child))
+      .filter((child): child is SelectionNode => child !== null),
+  };
+};
+
+export const trimSelectionNodes = (nodes: SelectionNode[]): SelectionNode[] => {
+  return nodes
+    .filter((node) => node.type === 'text' || !node.isCodeLineNumber)
+    .map((node) => node.type === 'text'
+      ? node
+      : { ...node, children: trimSelectionNodes(node.children) });
+};
+
+const toSelectionNodes = (root: ParentNode): SelectionNode[] => {
+  return Array.from(root.childNodes)
+    .map((child) => toSelectionNode(child))
+    .filter((child): child is SelectionNode => child !== null);
+};
+
+const getSelectionText = (node: SelectionNode): string => {
+  return node.type === 'text'
+    ? node.value
+    : node.children.map((child) => getSelectionText(child)).join('');
+};
+
+const findElement = (
+  node: SelectionNode,
+  predicate: (element: Extract<SelectionNode, { type: 'element' }>) => boolean,
+): Extract<SelectionNode, { type: 'element' }> | null => {
+  if (node.type === 'text') return null;
+  if (predicate(node)) return node;
+  for (const child of node.children) {
+    const match = findElement(child, predicate);
+    if (match) return match;
+  }
+  return null;
+};
+
+export const formatCodeSelectionMarkdown = (code: string, language = ''): string => {
+  const normalizedCode = normalizeLineBreaks(code).replace(/\n$/, '');
+  const longestBacktickRun = Math.max(0, ...Array.from(normalizedCode.matchAll(/`+/g), (match) => match[0].length));
+  const fence = '`'.repeat(Math.max(3, longestBacktickRun + 1));
+  return `${fence}${language}\n${normalizedCode}\n${fence}`;
+};
+
+const renderInlineMarkdownNode = (node: SelectionNode): string => {
+  if (node.type === 'text') {
+    return textToMarkdownInline(node.value);
+  }
+
+  const childText = node.children
+    .map((child) => renderInlineMarkdownNode(child))
+    .join('')
+    .replace(/\s+/g, ' ')
+    .trim();
+
+  if (!childText && node.tag !== 'br') return '';
+  if (node.tag === 'br') return '\n';
+  if (node.tag === 'strong' || node.tag === 'b') return `**${childText}**`;
+  if (node.tag === 'em' || node.tag === 'i') return `*${childText}*`;
+  if (node.tag === 'code') return `\`${childText.replace(/`/g, '\\`')}\``;
+  if (node.tag === 'a') return node.href ? `[${childText}](${node.href})` : childText;
+  return childText;
+};
+
+const renderListMarkdown = (list: Extract<SelectionNode, { type: 'element' }>, ordered: boolean): string => {
+  return list.children
+    .filter((child): child is Extract<SelectionNode, { type: 'element' }> => child.type === 'element' && child.tag === 'li')
+    .map((item, index) => {
+      const body = item.children
+        .map((child) => renderInlineMarkdownNode(child))
+        .join('')
+        .replace(/\s+/g, ' ')
+        .trim();
+      return body ? `${ordered ? `${index + 1}.` : '-'} ${body}` : '';
+    })
+    .filter(Boolean)
+    .join('\n');
+};
+
+const renderBlockMarkdownNode = (node: SelectionNode): string => {
+  if (node.type === 'text') return trimSelectionValue(node.value);
+
+  if (node.isMarkdownBlock) {
+    return node.children
+      .map((child) => renderBlockMarkdownNode(child))
+      .filter((child) => child.length > 0)
+      .join('\n\n');
+  }
+
+  if (node.component === 'markdown-code') {
+    const pre = findElement(node, (element) => element.tag === 'pre');
+    return pre ? renderBlockMarkdownNode(pre) : '';
+  }
+
+  if (node.tag === 'pre' || (node.tag === 'code' && node.isCodeLines)) {
+    const code = node.tag === 'code'
+      ? node
+      : findElement(node, (element) => element.tag === 'code');
+    return formatCodeSelectionMarkdown(
+      code ? getSelectionText(code) : getSelectionText(node),
+      node.markdownLanguage || getCodeLanguageFromClassName(code?.className || ''),
+    );
+  }
+
+  if (node.tag === 'code') {
+    const code = normalizeLineBreaks(getSelectionText(node)).trim();
+    return code ? `\`${code.replace(/`/g, '\\`')}\`` : '';
+  }
+  if (node.tag === 'ul') return renderListMarkdown(node, false);
+  if (node.tag === 'ol') return renderListMarkdown(node, true);
+
+  if (node.tag === 'blockquote') {
+    const content = trimSelectionValue(node.children.map((child) => renderBlockMarkdownNode(child)).join('\n'));
+    return content
+      .split('\n')
+      .filter((line) => line.length > 0)
+      .map((line) => `> ${line}`)
+      .join('\n');
+  }
+
+  if (/^h[1-6]$/.test(node.tag)) {
+    const level = Number.parseInt(node.tag[1], 10);
+    const text = trimSelectionValue(node.children.map((child) => renderInlineMarkdownNode(child)).join(''));
+    return text ? `${'#'.repeat(level)} ${text}` : '';
+  }
+
+  if (node.tag === 'p' || node.tag === 'div' || node.tag === 'li') {
+    return trimSelectionValue(node.children.map((child) => renderInlineMarkdownNode(child)).join(''));
+  }
+
+  const blockChildren = node.children
+    .map((child) => renderBlockMarkdownNode(child))
+    .filter((child) => child.length > 0);
+  return blockChildren.length > 0
+    ? blockChildren.join('\n\n')
+    : trimSelectionValue(node.children.map((child) => renderInlineMarkdownNode(child)).join(''));
+};
+
+const isInlineSelectionNode = (node: SelectionNode): boolean => {
+  if (node.type === 'text') return true;
+  return !node.isMarkdownBlock && !node.isCodeLines && node.component !== 'markdown-code' && !BLOCK_TAGS.has(node.tag);
+};
+
+export const selectionNodesToMarkdown = (nodes: SelectionNode[], plainText: string): string => {
+  const trimmedNodes = trimSelectionNodes(nodes);
+  if (trimmedNodes.every((node) => isInlineSelectionNode(node))) {
+    const inlineMarkdown = trimSelectionValue(trimmedNodes.map((node) => renderInlineMarkdownNode(node)).join(''));
+    if (inlineMarkdown) return inlineMarkdown;
+  }
+
+  const markdown = trimmedNodes
+    .map((node) => renderBlockMarkdownNode(node))
+    .filter((value) => value.length > 0)
+    .join('\n\n')
+    .trim();
+  return markdown || trimSelectionValue(plainText);
+};
+
+const getContainingBlockCode = (node: Node): HTMLElement | null => {
+  const element = node.nodeType === 1 ? node as Element : node.parentElement;
+  return element?.closest<HTMLElement>('pre code') ?? null;
+};
+
+export const rangeToMarkdown = (range: Range, plainText: string): string => {
+  const startCode = getContainingBlockCode(range.startContainer);
+  const endCode = getContainingBlockCode(range.endContainer);
+  const nodes = trimSelectionNodes(toSelectionNodes(range.cloneContents()));
+
+  if (startCode && startCode === endCode) {
+    return formatCodeSelectionMarkdown(
+      nodes.map((node) => getSelectionText(node)).join(''),
+      getBlockCodeLanguage(startCode),
+    );
+  }
+
+  return selectionNodesToMarkdown(nodes, plainText);
+};
+
+export const wrapMarkdownSelectionForChat = (markdown: string): string => {
+  const longestBacktickRun = Math.max(0, ...Array.from(markdown.matchAll(/`+/g), (match) => match[0].length));
+  const fence = '`'.repeat(Math.max(3, longestBacktickRun + 1));
+  return `${fence}md\n${markdown}\n${fence}`;
+};