Improve Katex copy robustness

src/vs/workbench/contrib/chat/browser/actions/chatCopyActions.ts

@@ -12,6 +12,7 @@ import { CHAT_CATEGORY, stringifyItem } from './chatActions.js';
 import { ChatTreeItem, IChatWidgetService } from '../chat.js';
 import { ChatContextKeys } from '../../common/chatContextKeys.js';
 import { IChatRequestViewModel, IChatResponseViewModel, isChatTreeItem, isRequestVM, isResponseVM } from '../../common/chatViewModel.js';
+import { katexContainerClassName, katexContainerLatexAttributeName } from '../../../markdown/common/markedKatexExtension.js';
 
 export function registerChatCopyActions() {
 	registerAction2(class CopyAllAction extends Action2 {
@@ -135,15 +136,13 @@ export function registerChatCopyActions() {
 			// Otherwise, fallback to querying from the active element
 			if (!selectedElement) {
 				// eslint-disable-next-line no-restricted-syntax
-				selectedElement = activeElement?.querySelector('.katex') ?? null;
+				selectedElement = activeElement?.querySelector(`.${katexContainerClassName}`) ?? null;
 			}
 
 			// Extract the LaTeX source from the annotation element
-			const katexElement = dom.isHTMLElement(selectedElement) ? selectedElement.closest('.katex') : null;
-			// eslint-disable-next-line no-restricted-syntax
-			const annotation = katexElement?.querySelector('annotation[encoding="application/x-tex"]');
-			if (annotation) {
-				const latexSource = annotation.textContent || '';
+			const katexElement = dom.isHTMLElement(selectedElement) ? selectedElement.closest(`.${katexContainerClassName}`) : null;
+			const latexSource = katexElement?.getAttribute(katexContainerLatexAttributeName) || '';
+			if (latexSource) {
 				await clipboardService.writeText(latexSource);
 			}
 		}

src/vs/workbench/contrib/chat/browser/chatWidget.ts

@@ -95,6 +95,7 @@ import { ChatViewPane } from './chatViewPane.js';
 import { ChatViewWelcomePart, IChatSuggestedPrompts, IChatViewWelcomeContent } from './viewsWelcome/chatViewWelcomeController.js';
 import { IWorkspaceContextService, WorkbenchState } from '../../../../platform/workspace/common/workspace.js';
 import { ILifecycleService } from '../../../services/lifecycle/common/lifecycle.js';
+import { katexContainerClassName } from '../../markdown/common/markedKatexExtension.js';
 
 const $ = dom.$;
 
@@ -2015,7 +2016,7 @@ export class ChatWidget extends Disposable implements IChatWidget {
 
 		// Check if the context menu was opened on a KaTeX element
 		const target = e.browserEvent.target as HTMLElement;
-		const isKatexElement = target.closest('.katex') !== null;
+		const isKatexElement = target.closest(`.${katexContainerClassName}`) !== null;
 
 		const scopedContextKeyService = this.contextKeyService.createOverlay([
 			[ChatContextKeys.responseIsFiltered.key, isResponseVM(selected) && !!selected.errorDetails?.responseIsFiltered],

src/vs/workbench/contrib/markdown/browser/markedKatexSupport.ts

@@ -9,7 +9,7 @@ import { MarkdownSanitizerConfig } from '../../../../base/browser/markdownRender
 import { CodeWindow } from '../../../../base/browser/window.js';
 import { Lazy } from '../../../../base/common/lazy.js';
 import type * as marked from '../../../../base/common/marked/marked.js';
-import { MarkedKatexExtension } from '../common/markedKatexExtension.js';
+import { katexContainerLatexAttributeName, MarkedKatexExtension } from '../common/markedKatexExtension.js';
 
 export class MarkedKatexSupport {
 
@@ -32,6 +32,7 @@ export class MarkedKatexSupport {
 					'stretchy',
 					'encoding',
 					'accent',
+					katexContainerLatexAttributeName,
 
 					// SVG
 					'd',

src/vs/workbench/contrib/markdown/common/markedKatexExtension.ts

@@ -3,9 +3,11 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 import type * as marked from '../../../../base/common/marked/marked.js';
+import { htmlAttributeEncodeValue } from '../../../../base/common/strings.js';
 
 export const mathInlineRegExp = /(?<![a-zA-Z0-9])(?<dollars>\${1,2})(?!\$)((?:\\.|[^\\\n])*?(?:\\.|[^\\\n\$]))\k<dollars>(?![a-zA-Z0-9])/; // Non-standard, but ensure opening $ is not preceded and closing $ is not followed by word/number characters
-
+export const katexContainerClassName = 'vscode-katex-container';
+export const katexContainerLatexAttributeName = 'data-latex';
 
 const inlineRule = new RegExp('^' + mathInlineRegExp.source);
 
@@ -32,11 +34,15 @@ export namespace MarkedKatexExtension {
 		return (token: marked.Tokens.Generic) => {
 			let out: string;
 			try {
-				out = katex.renderToString(token.text, {
+				const html = katex.renderToString(token.text, {
 					...options,
 					throwOnError: true,
 					displayMode: token.displayMode,
 				});
+
+				// Wrap in a container with attribute as a fallback for extracting the original LaTeX source
+				// This ensures we can always retrieve the source even if the annotation element is not present
+				out = `<span class="${katexContainerClassName}" ${katexContainerLatexAttributeName}="${htmlAttributeEncodeValue(token.text)}">${html}</span>`;
 			} catch {
 				// On failure, just use the original text including the wrapping $ or $$
 				out = token.raw;

src/vs/workbench/contrib/markdown/test/browser/__snapshots__/Markdown_Katex_Support_Test_Basic_inline_equation.0.snap

@@ -1 +1 @@
-<p>Hello <span class="katex"><span class="katex-mathml"><math><semantics><mrow><mfrac><mn>1</mn><mn>2</mn></mfrac></mrow><annotation encoding="application/x-tex">\frac{1}{2}</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 1.1901em; vertical-align: -0.345em"></span><span class="mord"><span class="mopen nulldelimiter"></span><span class="mfrac"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height: 0.8451em"><span style="top: -2.655em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">2</span></span></span></span><span style="top: -3.23em"><span class="pstrut" style="height: 3em"></span><span class="frac-line" style="border-bottom-width: 0.04em"></span></span><span style="top: -3.394em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">1</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height: 0.345em"><span></span></span></span></span></span><span class="mclose nulldelimiter"></span></span></span></span></span> World!</p>
+<p>Hello <span class="vscode-katex-container" data-latex="\frac{1}{2}"><span class="katex"><span class="katex-mathml"><math><semantics><mrow><mfrac><mn>1</mn><mn>2</mn></mfrac></mrow><annotation encoding="application/x-tex">\frac{1}{2}</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 1.1901em; vertical-align: -0.345em"></span><span class="mord"><span class="mopen nulldelimiter"></span><span class="mfrac"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height: 0.8451em"><span style="top: -2.655em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">2</span></span></span></span><span style="top: -3.23em"><span class="pstrut" style="height: 3em"></span><span class="frac-line" style="border-bottom-width: 0.04em"></span></span><span style="top: -3.394em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">1</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height: 0.345em"><span></span></span></span></span></span><span class="mclose nulldelimiter"></span></span></span></span></span></span> World!</p>

src/vs/workbench/contrib/markdown/test/browser/__snapshots__/Markdown_Katex_Support_Test_Should_still_render_math_at_start_and_end_of_line.0.snap

@@ -1 +1 @@
-<p><span class="katex"><span class="katex-mathml"><math><semantics><mrow><mfrac><mn>1</mn><mn>2</mn></mfrac></mrow><annotation encoding="application/x-tex">\frac{1}{2}</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 1.1901em; vertical-align: -0.345em"></span><span class="mord"><span class="mopen nulldelimiter"></span><span class="mfrac"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height: 0.8451em"><span style="top: -2.655em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">2</span></span></span></span><span style="top: -3.23em"><span class="pstrut" style="height: 3em"></span><span class="frac-line" style="border-bottom-width: 0.04em"></span></span><span style="top: -3.394em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">1</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height: 0.345em"><span></span></span></span></span></span><span class="mclose nulldelimiter"></span></span></span></span></span> at start, and at end <span class="katex"><span class="katex-mathml"><math><semantics><mrow><msup><mi>x</mi><mn>2</mn></msup></mrow><annotation encoding="application/x-tex">x^2</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 0.8141em"></span><span class="mord"><span class="mord mathnormal">x</span><span class="msupsub"><span class="vlist-t"><span class="vlist-r"><span class="vlist" style="height: 0.8141em"><span style="top: -3.063em; margin-right: 0.05em"><span class="pstrut" style="height: 2.7em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight">2</span></span></span></span></span></span></span></span></span></span></span></p>
+<p><span class="vscode-katex-container" data-latex="\frac{1}{2}"><span class="katex"><span class="katex-mathml"><math><semantics><mrow><mfrac><mn>1</mn><mn>2</mn></mfrac></mrow><annotation encoding="application/x-tex">\frac{1}{2}</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 1.1901em; vertical-align: -0.345em"></span><span class="mord"><span class="mopen nulldelimiter"></span><span class="mfrac"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height: 0.8451em"><span style="top: -2.655em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">2</span></span></span></span><span style="top: -3.23em"><span class="pstrut" style="height: 3em"></span><span class="frac-line" style="border-bottom-width: 0.04em"></span></span><span style="top: -3.394em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">1</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height: 0.345em"><span></span></span></span></span></span><span class="mclose nulldelimiter"></span></span></span></span></span></span> at start, and at end <span class="vscode-katex-container" data-latex="x^2"><span class="katex"><span class="katex-mathml"><math><semantics><mrow><msup><mi>x</mi><mn>2</mn></msup></mrow><annotation encoding="application/x-tex">x^2</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 0.8141em"></span><span class="mord"><span class="mord mathnormal">x</span><span class="msupsub"><span class="vlist-t"><span class="vlist-r"><span class="vlist" style="height: 0.8141em"><span style="top: -3.063em; margin-right: 0.05em"><span class="pstrut" style="height: 2.7em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight">2</span></span></span></span></span></span></span></span></span></span></span></span></p>

src/vs/workbench/contrib/markdown/test/browser/__snapshots__/Markdown_Katex_Support_Test_Should_still_render_math_with_special_characters_around_dollars.0.snap

@@ -1 +1 @@
-<p>Hello (<span class="katex"><span class="katex-mathml"><math><semantics><mrow><mfrac><mn>1</mn><mn>2</mn></mfrac></mrow><annotation encoding="application/x-tex">\frac{1}{2}</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 1.1901em; vertical-align: -0.345em"></span><span class="mord"><span class="mopen nulldelimiter"></span><span class="mfrac"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height: 0.8451em"><span style="top: -2.655em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">2</span></span></span></span><span style="top: -3.23em"><span class="pstrut" style="height: 3em"></span><span class="frac-line" style="border-bottom-width: 0.04em"></span></span><span style="top: -3.394em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">1</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height: 0.345em"><span></span></span></span></span></span><span class="mclose nulldelimiter"></span></span></span></span></span>) and [<span class="katex"><span class="katex-mathml"><math><semantics><mrow><msup><mi>x</mi><mn>2</mn></msup></mrow><annotation encoding="application/x-tex">x^2</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 0.8141em"></span><span class="mord"><span class="mord mathnormal">x</span><span class="msupsub"><span class="vlist-t"><span class="vlist-r"><span class="vlist" style="height: 0.8141em"><span style="top: -3.063em; margin-right: 0.05em"><span class="pstrut" style="height: 2.7em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight">2</span></span></span></span></span></span></span></span></span></span></span>] work fine</p>
+<p>Hello (<span class="vscode-katex-container" data-latex="\frac{1}{2}"><span class="katex"><span class="katex-mathml"><math><semantics><mrow><mfrac><mn>1</mn><mn>2</mn></mfrac></mrow><annotation encoding="application/x-tex">\frac{1}{2}</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 1.1901em; vertical-align: -0.345em"></span><span class="mord"><span class="mopen nulldelimiter"></span><span class="mfrac"><span class="vlist-t vlist-t2"><span class="vlist-r"><span class="vlist" style="height: 0.8451em"><span style="top: -2.655em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">2</span></span></span></span><span style="top: -3.23em"><span class="pstrut" style="height: 3em"></span><span class="frac-line" style="border-bottom-width: 0.04em"></span></span><span style="top: -3.394em"><span class="pstrut" style="height: 3em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mtight">1</span></span></span></span></span><span class="vlist-s">​</span></span><span class="vlist-r"><span class="vlist" style="height: 0.345em"><span></span></span></span></span></span><span class="mclose nulldelimiter"></span></span></span></span></span></span>) and [<span class="vscode-katex-container" data-latex="x^2"><span class="katex"><span class="katex-mathml"><math><semantics><mrow><msup><mi>x</mi><mn>2</mn></msup></mrow><annotation encoding="application/x-tex">x^2</annotation></semantics></math></span><span class="katex-html"><span class="base"><span class="strut" style="height: 0.8141em"></span><span class="mord"><span class="mord mathnormal">x</span><span class="msupsub"><span class="vlist-t"><span class="vlist-r"><span class="vlist" style="height: 0.8141em"><span style="top: -3.063em; margin-right: 0.05em"><span class="pstrut" style="height: 2.7em"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight">2</span></span></span></span></span></span></span></span></span></span></span></span>] work fine</p>