Add support for copying math source from katex blocks (#271590)

5593ecd8 · Paul · GitHub · 52a89abf · 5593ecd8 · 5593ecd8
--- a/src/vs/workbench/contrib/chat/browser/actions/chatCopyActions.ts
+++ b/src/vs/workbench/contrib/chat/browser/actions/chatCopyActions.ts
@@ -87,4 +87,63 @@ export function registerChatCopyActions() {
 			await clipboardService.writeText(text);
 		}
 	});
+
+	registerAction2(class CopyKatexMathSourceAction extends Action2 {
+		constructor() {
+			super({
+				id: 'workbench.action.chat.copyKatexMathSource',
+				title: localize2('chat.copyKatexMathSource.label', "Copy Math Source"),
+				f1: false,
+				category: CHAT_CATEGORY,
+				menu: {
+					id: MenuId.ChatContext,
+					group: 'copy',
+					when: ChatContextKeys.isKatexMathElement,
+				}
+			});
+		}
+
+		async run(accessor: ServicesAccessor, ...args: unknown[]) {
+			const chatWidgetService = accessor.get(IChatWidgetService);
+			const clipboardService = accessor.get(IClipboardService);
+
+			const widget = chatWidgetService.lastFocusedWidget;
+			let item = args[0] as ChatTreeItem | undefined;
+			if (!isChatTreeItem(item)) {
+				item = widget?.getFocus();
+				if (!item) {
+					return;
+				}
+			}
+
+			// Try to find a KaTeX element from the selection or active element
+			let selectedElement: Node | null = null;
+
+			// If there is a selection, and focus is inside the widget, extract the inner KaTeX element.
+			const activeElement = dom.getActiveElement();
+			const nativeSelection = dom.getActiveWindow().getSelection();
+			if (widget && nativeSelection && nativeSelection.rangeCount > 0 && dom.isAncestor(activeElement, widget.domNode)) {
+				const range = nativeSelection.getRangeAt(0);
+				selectedElement = range.commonAncestorContainer;
+
+				// If it's a text node, get its parent element
+				if (selectedElement.nodeType === Node.TEXT_NODE) {
+					selectedElement = selectedElement.parentElement;
+				}
+			}
+
+			// Otherwise, fallback to querying from the active element
+			if (!selectedElement) {
+				selectedElement = activeElement?.querySelector('.katex') ?? null;
+			}
+
+			// Extract the LaTeX source from the annotation element
+			const katexElement = dom.isHTMLElement(selectedElement) ? selectedElement.closest('.katex') : null;
+			const annotation = katexElement?.querySelector('annotation[encoding="application/x-tex"]');
+			if (annotation) {
+				const latexSource = annotation.textContent || '';
+				await clipboardService.writeText(latexSource);
+			}
+		}
+	});
 }
--- a/src/vs/workbench/contrib/chat/browser/chatWidget.ts
+++ b/src/vs/workbench/contrib/chat/browser/chatWidget.ts
@@ -1963,8 +1963,14 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		e.browserEvent.stopPropagation();

 		const selected = e.element;
+
+		// 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 scopedContextKeyService = this.contextKeyService.createOverlay([
-			[ChatContextKeys.responseIsFiltered.key, isResponseVM(selected) && !!selected.errorDetails?.responseIsFiltered]
+			[ChatContextKeys.responseIsFiltered.key, isResponseVM(selected) && !!selected.errorDetails?.responseIsFiltered],
+			[ChatContextKeys.isKatexMathElement.key, isKatexElement]
 		]);
 		this.contextMenuService.showContextMenu({
 			menuId: MenuId.ChatContext,

--- a/src/vs/workbench/contrib/chat/common/chatContextKeys.ts
+++ b/src/vs/workbench/contrib/chat/common/chatContextKeys.ts
@@ -92,6 +92,7 @@ export namespace ChatContextKeys {
 	export const sessionType = new RawContextKey<string>('chatSessionType', '', { type: 'string', description: localize('chatSessionType', "The type of the current chat session item.") });
 	export const isHistoryItem = new RawContextKey<boolean>('chatIsHistoryItem', false, { type: 'boolean', description: localize('chatIsHistoryItem', "True when the chat session item is from history.") });
 	export const isActiveSession = new RawContextKey<boolean>('chatIsActiveSession', false, { type: 'boolean', description: localize('chatIsActiveSession', "True when the chat session is currently active (not deletable).") });
+	export const isKatexMathElement = new RawContextKey<boolean>('chatIsKatexMathElement', false, { type: 'boolean', description: localize('chatIsKatexMathElement', "True when focusing a KaTeX math element.") });
 }

 export namespace ChatContextKeyExprs {