Merge pull request microsoft#191846 from microsoft/aamunger/flakyImageCopy

look for re-used output id containing the image

Author: amunger

src/vs/workbench/contrib/notebook/browser/controller/cellOutputActions.ts

@@ -55,7 +55,7 @@ registerAction2(class CopyCellOutputAction extends Action2 {
 		const mimeType = outputViewModel.pickedMimeType?.mimeType;
 
 		if (mimeType?.startsWith('image/')) {
-			const focusOptions = { skipReveal: true, outputId: outputViewModel.model.outputId };
+			const focusOptions = { skipReveal: true, outputId: outputViewModel.model.outputId, altOutputId: outputViewModel.model.alternativeOutputId };
 			await notebookEditor.focusNotebookCell(outputViewModel.cellViewModel as ICellViewModel, 'output', focusOptions);
 			notebookEditor.copyOutputImage(outputViewModel);
 		} else {

src/vs/workbench/contrib/notebook/browser/notebookBrowser.ts

@@ -152,6 +152,7 @@ export interface IFocusNotebookCellOptions {
 	readonly focusEditorLine?: number;
 	readonly minimalScrolling?: boolean;
 	readonly outputId?: string;
+	readonly altOutputId?: string;
 }
 
 //#endregion

src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts

@@ -2350,7 +2350,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 			}
 
 			const focusElementId = options?.outputId ?? cell.id;
-			this._webview.focusOutput(focusElementId, this._webviewFocused);
+			this._webview.focusOutput(focusElementId, options?.altOutputId, this._webviewFocused);
 
 			cell.updateEditState(CellEditState.Preview, 'focusNotebookCell');
 			cell.focusMode = CellFocusMode.Output;

src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

@@ -1548,7 +1548,8 @@ export class BackLayerWebView<T extends ICommonCellInfo> extends Themable {
 	async copyImage(output: ICellOutputViewModel): Promise<void> {
 		this._sendMessageToWebview({
 			type: 'copyImage',
-			outputId: output.model.outputId
+			outputId: output.model.outputId,
+			altOutputId: output.model.alternativeOutputId
 		});
 	}
 
@@ -1608,7 +1609,7 @@ export class BackLayerWebView<T extends ICommonCellInfo> extends Themable {
 		this.webview?.focus();
 	}
 
-	focusOutput(cellOrOutputId: string, viewFocused: boolean) {
+	focusOutput(cellOrOutputId: string, alternateId: string | undefined, viewFocused: boolean) {
 		if (this._disposed) {
 			return;
 		}
@@ -1620,6 +1621,7 @@ export class BackLayerWebView<T extends ICommonCellInfo> extends Themable {
 		this._sendMessageToWebview({
 			type: 'focus-output',
 			cellOrOutputId: cellOrOutputId,
+			alternateId: alternateId
 		});
 	}
 

src/vs/workbench/contrib/notebook/browser/view/renderers/webviewMessages.ts

@@ -269,11 +269,13 @@ export interface IShowOutputMessage {
 export interface ICopyImageMessage {
 	readonly type: 'copyImage';
 	readonly outputId: string;
+	readonly altOutputId: string;
 }
 
 export interface IFocusOutputMessage {
 	readonly type: 'focus-output';
 	readonly cellOrOutputId: string;
+	readonly alternateId?: string;
 }
 
 export interface IAckOutputHeight {

src/vs/workbench/contrib/notebook/browser/view/renderers/webviewPreloads.ts

@@ -474,8 +474,9 @@ async function webviewPreloads(ctx: PreloadContext) {
 		});
 	};
 
-	function focusFirstFocusableOrContainerInOutput(cellOrOutputId: string) {
-		const cellOutputContainer = document.getElementById(cellOrOutputId);
+	function focusFirstFocusableOrContainerInOutput(cellOrOutputId: string, alternateId?: string) {
+		const cellOutputContainer = document.getElementById(cellOrOutputId) ??
+			alternateId ? document.getElementById(alternateId!) : undefined;
 		if (cellOutputContainer) {
 			if (cellOutputContainer.contains(document.activeElement)) {
 				return;
@@ -1362,17 +1363,18 @@ async function webviewPreloads(ctx: PreloadContext) {
 		});
 	};
 
-	const copyOutputImage = async (outputId: string, retries = 5) => {
+	const copyOutputImage = async (outputId: string, altOutputId: string, retries = 5) => {
 		if (!document.hasFocus() && retries > 0) {
 			// copyImage can be called from outside of the webview, which means this function may be running whilst the webview is gaining focus.
 			// Since navigator.clipboard.write requires the document to be focused, we need to wait for focus.
 			// We cannot use a listener, as there is a high chance the focus is gained during the setup of the listener resulting in us missing it.
-			setTimeout(() => { copyOutputImage(outputId, retries - 1); }, 20);
+			setTimeout(() => { copyOutputImage(outputId, altOutputId, retries - 1); }, 20);
 			return;
 		}
 
 		try {
-			const image = document.getElementById(outputId)?.querySelector('img');
+			const image = document.getElementById(outputId)?.querySelector('img')
+				?? document.getElementById(altOutputId)?.querySelector('img');
 			if (image) {
 				await navigator.clipboard.write([new ClipboardItem({
 					'image/png': new Promise((resolve) => {
@@ -1500,9 +1502,7 @@ async function webviewPreloads(ctx: PreloadContext) {
 				break;
 			}
 			case 'copyImage': {
-
-				await copyOutputImage(event.data.outputId);
-
+				await copyOutputImage(event.data.outputId, event.data.altOutputId);
 				break;
 			}
 			case 'ack-dimension': {
@@ -1524,7 +1524,7 @@ async function webviewPreloads(ctx: PreloadContext) {
 				break;
 			}
 			case 'focus-output':
-				focusFirstFocusableOrContainerInOutput(event.data.cellOrOutputId);
+				focusFirstFocusableOrContainerInOutput(event.data.cellOrOutputId, event.data.alternateId);
 				break;
 			case 'decorations': {
 				let outputContainer = document.getElementById(event.data.cellId);