Allow passing in a DataTransferFile to workspace edit (#175809)

* Allow passing in a `DataTransferItem` to workspace edit

Fixes #175800

Allows you to pass a file `DataTransferItem` to `WorkspaceEdit.createFile`. This lets us avoid transferring the data back and forth to the extension host, and also avoid having to base64 encode and decode it, significantly improving performance for large files

* Take data transfer file instead of data transfer item

extensions/markdown-language-features/src/languageFeatures/copyFiles/copyPaste.ts

@@ -31,8 +31,9 @@ class PasteEditProvider implements vscode.DocumentPasteEditProvider {
 		}
 
 		for (const imageMime of supportedImageMimes) {
-			const file = dataTransfer.get(imageMime)?.asFile();
-			if (file) {
+			const item = dataTransfer.get(imageMime);
+			const file = item?.asFile();
+			if (item && file) {
 				const edit = await this._makeCreateImagePasteEdit(document, file, token);
 				if (token.isCancellationRequested) {
 					return;
@@ -70,7 +71,7 @@ class PasteEditProvider implements vscode.DocumentPasteEditProvider {
 
 		// Note that there is currently no way to undo the file creation :/
 		const workspaceEdit = new vscode.WorkspaceEdit();
-		workspaceEdit.createFile(uri, { contents: await file.data() });
+		workspaceEdit.createFile(uri, { contents: file });
 
 		const pasteEdit = new vscode.DocumentPasteEdit(snippet);
 		pasteEdit.additionalEdit = workspaceEdit;

src/vs/editor/common/languages.ts

@@ -3,14 +3,15 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { VSBuffer } from 'vs/base/common/buffer';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { Codicon } from 'vs/base/common/codicons';
-import { ThemeIcon } from 'vs/base/common/themables';
 import { Color } from 'vs/base/common/color';
 import { VSDataTransfer } from 'vs/base/common/dataTransfer';
 import { Event } from 'vs/base/common/event';
 import { IMarkdownString } from 'vs/base/common/htmlContent';
 import { IDisposable } from 'vs/base/common/lifecycle';
+import { ThemeIcon } from 'vs/base/common/themables';
 import { URI, UriComponents } from 'vs/base/common/uri';
 import { ISingleEditOperation } from 'vs/editor/common/core/editOperation';
 import { IPosition, Position } from 'vs/editor/common/core/position';
@@ -1478,7 +1479,11 @@ export interface WorkspaceFileEditOptions {
 	folder?: boolean;
 	skipTrashBin?: boolean;
 	maxSize?: number;
-	contentsBase64?: string;
+
+	/**
+	 * @internal
+	 */
+	contents?: Promise<VSBuffer>;
 }
 
 export interface IWorkspaceFileEdit {

src/vs/monaco.d.ts

@@ -7457,7 +7457,6 @@ declare namespace monaco.languages {
 		folder?: boolean;
 		skipTrashBin?: boolean;
 		maxSize?: number;
-		contentsBase64?: string;
 	}
 
 	export interface IWorkspaceFileEdit {

src/vs/workbench/api/browser/mainThreadBulkEdits.ts

@@ -3,14 +3,15 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { VSBuffer, decodeBase64 } from 'vs/base/common/buffer';
 import { revive } from 'vs/base/common/marshalling';
 import { IBulkEditService, ResourceFileEdit, ResourceTextEdit } from 'vs/editor/browser/services/bulkEditService';
 import { WorkspaceEdit } from 'vs/editor/common/languages';
 import { ILogService } from 'vs/platform/log/common/log';
 import { IUriIdentityService } from 'vs/platform/uriIdentity/common/uriIdentity';
-import { IWorkspaceEditDto, MainContext, MainThreadBulkEditsShape } from 'vs/workbench/api/common/extHost.protocol';
+import { IWorkspaceEditDto, IWorkspaceFileEditDto, MainContext, MainThreadBulkEditsShape } from 'vs/workbench/api/common/extHost.protocol';
 import { ResourceNotebookCellEdit } from 'vs/workbench/contrib/bulkEdit/browser/bulkCellEdits';
-import { extHostNamedCustomer, IExtHostContext } from 'vs/workbench/services/extensions/common/extHostCustomers';
+import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers';
 
 
 @extHostNamedCustomer(MainContext.MainThreadBulkEdits)
@@ -34,9 +35,9 @@ export class MainThreadBulkEdits implements MainThreadBulkEditsShape {
 	}
 }
 
-export function reviveWorkspaceEditDto(data: IWorkspaceEditDto, uriIdentityService: IUriIdentityService): WorkspaceEdit;
-export function reviveWorkspaceEditDto(data: IWorkspaceEditDto | undefined, uriIdentityService: IUriIdentityService): WorkspaceEdit | undefined;
-export function reviveWorkspaceEditDto(data: IWorkspaceEditDto | undefined, uriIdentityService: IUriIdentityService): WorkspaceEdit | undefined {
+export function reviveWorkspaceEditDto(data: IWorkspaceEditDto, uriIdentityService: IUriIdentityService, resolveDataTransferFile?: (id: string) => Promise<VSBuffer>): WorkspaceEdit;
+export function reviveWorkspaceEditDto(data: IWorkspaceEditDto | undefined, uriIdentityService: IUriIdentityService, resolveDataTransferFile?: (id: string) => Promise<VSBuffer>): WorkspaceEdit | undefined;
+export function reviveWorkspaceEditDto(data: IWorkspaceEditDto | undefined, uriIdentityService: IUriIdentityService, resolveDataTransferFile?: (id: string) => Promise<VSBuffer>): WorkspaceEdit | undefined {
 	if (!data || !data.edits) {
 		return <WorkspaceEdit>data;
 	}
@@ -46,6 +47,20 @@ export function reviveWorkspaceEditDto(data: IWorkspaceEditDto | undefined, uriI
 			edit.resource = uriIdentityService.asCanonicalUri(edit.resource);
 		}
 		if (ResourceFileEdit.is(edit)) {
+			if (edit.options) {
+				const inContents = (edit as IWorkspaceFileEditDto).options?.contents;
+				if (inContents) {
+					if (inContents.type === 'base64') {
+						edit.options.contents = Promise.resolve(decodeBase64(inContents.value));
+					} else {
+						if (resolveDataTransferFile) {
+							edit.options.contents = resolveDataTransferFile(inContents.id);
+						} else {
+							throw new Error('Could not revive data transfer file');
+						}
+					}
+				}
+			}
 			edit.newResource = edit.newResource && uriIdentityService.asCanonicalUri(edit.newResource);
 			edit.oldResource = edit.oldResource && uriIdentityService.asCanonicalUri(edit.oldResource);
 		}

src/vs/workbench/api/browser/mainThreadInteractiveEditor.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { DisposableMap } from 'vs/base/common/lifecycle';
-import { IInteractiveEditorResponse, IInteractiveEditorService } from 'vs/workbench/contrib/interactiveEditor/common/interactiveEditor';
+import { IInteractiveEditorBulkEditResponse, IInteractiveEditorResponse, IInteractiveEditorService } from 'vs/workbench/contrib/interactiveEditor/common/interactiveEditor';
 import { IUriIdentityService } from 'vs/platform/uriIdentity/common/uriIdentity';
 import { reviveWorkspaceEditDto } from 'vs/workbench/api/browser/mainThreadBulkEdits';
 import { ExtHostContext, ExtHostInteractiveEditorShape, MainContext, MainThreadInteractiveEditorShape } from 'vs/workbench/api/common/extHost.protocol';
@@ -55,7 +55,7 @@ export class MainThreadInteractiveEditor implements MainThreadInteractiveEditorS
 			provideResponse: async (item, request, token) => {
 				const result = await this._proxy.$provideResponse(handle, item, request, token);
 				if (result?.type === 'bulkEdit') {
-					result.edits = reviveWorkspaceEditDto(result.edits, this._uriIdentService);
+					(<IInteractiveEditorBulkEditResponse>result).edits = reviveWorkspaceEditDto(result.edits, this._uriIdentService);
 				}
 				return <IInteractiveEditorResponse | undefined>result;
 			},

src/vs/workbench/api/browser/mainThreadLanguageFeatures.ts

@@ -974,7 +974,7 @@ class MainThreadPasteEditProvider implements languages.DocumentPasteEditProvider
 
 			return {
 				insertText: result.insertText,
-				additionalEdit: result.additionalEdit ? reviveWorkspaceEditDto(result.additionalEdit, this._uriIdentService) : undefined,
+				additionalEdit: result.additionalEdit ? reviveWorkspaceEditDto(result.additionalEdit, this._uriIdentService, dataId => this.resolveFileData(request.id, dataId)) : undefined,
 			};
 		} finally {
 			request.dispose();
@@ -982,7 +982,7 @@ class MainThreadPasteEditProvider implements languages.DocumentPasteEditProvider
 	}
 
 	resolveFileData(requestId: number, dataId: string): Promise<VSBuffer> {
-		return this.dataTransfers.resolveDropFileData(requestId, dataId);
+		return this.dataTransfers.resolveFileData(requestId, dataId);
 	}
 }
 
@@ -1014,7 +1014,7 @@ class MainThreadDocumentOnDropEditProvider implements languages.DocumentOnDropEd
 			return {
 				label: edit.label,
 				insertText: edit.insertText,
-				additionalEdit: reviveWorkspaceEditDto(edit.additionalEdit, this._uriIdentService),
+				additionalEdit: reviveWorkspaceEditDto(edit.additionalEdit, this._uriIdentService, dataId => this.resolveDocumentOnDropFileData(request.id, dataId)),
 			};
 		} finally {
 			request.dispose();
@@ -1022,7 +1022,7 @@ class MainThreadDocumentOnDropEditProvider implements languages.DocumentOnDropEd
 	}
 
 	public resolveDocumentOnDropFileData(requestId: number, dataId: string): Promise<VSBuffer> {
-		return this.dataTransfers.resolveDropFileData(requestId, dataId);
+		return this.dataTransfers.resolveFileData(requestId, dataId);
 	}
 }
 

src/vs/workbench/api/browser/mainThreadTreeViews.ts

@@ -245,7 +245,7 @@ class TreeViewDragAndDropController implements ITreeViewDragAndDropController {
 	}
 
 	public resolveDropFileData(requestId: number, dataItemId: string): Promise<VSBuffer> {
-		return this.dataTransfersCache.resolveDropFileData(requestId, dataItemId);
+		return this.dataTransfersCache.resolveFileData(requestId, dataItemId);
 	}
 }
 

src/vs/workbench/api/common/extHost.protocol.ts

@@ -1738,7 +1738,10 @@ export type ICellEditOperationDto =
 
 export type IWorkspaceCellEditDto = Dto<Omit<notebookCommon.IWorkspaceNotebookCellEdit, 'cellEdit'>> & { cellEdit: ICellEditOperationDto };
 
-export type IWorkspaceFileEditDto = Dto<languages.IWorkspaceFileEdit>;
+export type IWorkspaceFileEditDto = Dto<
+	Omit<languages.IWorkspaceFileEdit, 'options'> & {
+		options?: Omit<languages.WorkspaceFileEditOptions, 'contents'> & { contents?: { type: 'base64'; value: string } | { type: 'dataTransferItem'; id: string } };
+	}>;
 
 export type IWorkspaceTextEditDto = Dto<languages.IWorkspaceTextEdit>;
 

src/vs/workbench/api/common/extHostTypeConverters.ts

@@ -590,11 +590,20 @@ export namespace WorkspaceEdit {
 			for (const entry of value._allEntries()) {
 
 				if (entry._type === types.FileEditType.File) {
+					let contents: { type: 'base64'; value: string } | { type: 'dataTransferItem'; id: string } | undefined;
+					if (entry.options?.contents) {
+						if (ArrayBuffer.isView(entry.options.contents)) {
+							contents = { type: 'base64', value: encodeBase64(VSBuffer.wrap(entry.options.contents)) };
+						} else {
+							contents = { type: 'dataTransferItem', id: (entry.options.contents as types.DataTransferFile)._itemId };
+						}
+					}
+
 					// file operation
-					result.edits.push(<languages.IWorkspaceFileEdit>{
+					result.edits.push(<extHostProtocol.IWorkspaceFileEditDto>{
 						oldResource: entry.from,
 						newResource: entry.to,
-						options: { ...entry.options, contentsBase64: entry.options?.contents && encodeBase64(VSBuffer.wrap(entry.options.contents)) },
+						options: { ...entry.options, contents },
 						metadata: entry.metadata
 					});
 
@@ -2027,12 +2036,8 @@ export namespace DataTransferItem {
 		const file = item.fileData;
 		if (file) {
 			return new class extends types.DataTransferItem {
-				override asFile(): vscode.DataTransferFile {
-					return {
-						name: file.name,
-						uri: URI.revive(file.uri),
-						data: once(() => resolveFileData()),
-					};
+				override asFile() {
+					return new types.DataTransferFile(file.name, URI.revive(file.uri), item.id, once(() => resolveFileData()));
 				}
 			}('', item.id);
 		}

src/vs/workbench/api/common/extHostTypes.ts

@@ -708,7 +708,7 @@ export interface IFileOperationOptions {
 	readonly ignoreIfExists?: boolean;
 	readonly ignoreIfNotExists?: boolean;
 	readonly recursive?: boolean;
-	readonly contents?: Uint8Array;
+	readonly contents?: Uint8Array | vscode.DataTransferFile;
 }
 
 export const enum FileEditType {
@@ -778,7 +778,7 @@ export class WorkspaceEdit implements vscode.WorkspaceEdit {
 		this._edits.push({ _type: FileEditType.File, from, to, options, metadata });
 	}
 
-	createFile(uri: vscode.Uri, options?: { readonly overwrite?: boolean; readonly ignoreIfExists?: boolean; readonly contents?: Uint8Array }, metadata?: vscode.WorkspaceEditEntryMetadata): void {
+	createFile(uri: vscode.Uri, options?: { readonly overwrite?: boolean; readonly ignoreIfExists?: boolean; readonly contents?: Uint8Array | vscode.DataTransferFile }, metadata?: vscode.WorkspaceEditEntryMetadata): void {
 		this._edits.push({ _type: FileEditType.File, from: undefined, to: uri, options, metadata });
 	}
 
@@ -2582,7 +2582,7 @@ export enum TreeItemCheckboxState {
 }
 
 @es5ClassCompat
-export class DataTransferItem {
+export class DataTransferItem implements vscode.DataTransferItem {
 
 	async asString(): Promise<string> {
 		return typeof this.value === 'string' ? this.value : JSON.stringify(this.value);
@@ -2602,6 +2602,29 @@ export class DataTransferItem {
 	}
 }
 
+/**
+ * Intentionally not exported to extensions
+ */
+export class DataTransferFile implements vscode.DataTransferFile {
+
+	public readonly name: string;
+	public readonly uri: vscode.Uri | undefined;
+
+	public readonly _itemId: string;
+	private readonly _getData: () => Promise<Uint8Array>;
+
+	constructor(name: string, uri: vscode.Uri | undefined, itemId: string, getData: () => Promise<Uint8Array>) {
+		this.name = name;
+		this.uri = uri;
+		this._itemId = itemId;
+		this._getData = getData;
+	}
+
+	data(): Promise<Uint8Array> {
+		return this._getData();
+	}
+}
+
 @es5ClassCompat
 export class DataTransfer implements vscode.DataTransfer {
 	#items = new Map<string, DataTransferItem[]>();

src/vs/workbench/api/common/shared/dataTransferCache.ts

@@ -22,7 +22,7 @@ export class DataTransferCache {
 		};
 	}
 
-	async resolveDropFileData(requestId: number, dataItemId: string): Promise<VSBuffer> {
+	async resolveFileData(requestId: number, dataItemId: string): Promise<VSBuffer> {
 		const entry = this.dataTransfers.get(requestId);
 		if (!entry) {
 			throw new Error('No data transfer found');

src/vs/workbench/contrib/bulkEdit/browser/bulkFileEdits.ts

@@ -13,7 +13,7 @@ import { IWorkspaceUndoRedoElement, UndoRedoElementType, IUndoRedoService, UndoR
 import { URI } from 'vs/base/common/uri';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ILogService } from 'vs/platform/log/common/log';
-import { decodeBase64, VSBuffer } from 'vs/base/common/buffer';
+import { VSBuffer } from 'vs/base/common/buffer';
 import { ResourceFileEdit } from 'vs/editor/browser/services/bulkEditService';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { tail } from 'vs/base/common/arrays';
@@ -345,7 +345,7 @@ export class BulkFileEdits {
 			} else if (!edit.newResource && edit.oldResource) {
 				edits.push(new DeleteEdit(edit.oldResource, edit.options ?? {}, false));
 			} else if (edit.newResource && !edit.oldResource) {
-				edits.push(new CreateEdit(edit.newResource, edit.options ?? {}, edit.options.contentsBase64 ? decodeBase64(edit.options.contentsBase64) : undefined));
+				edits.push(new CreateEdit(edit.newResource, edit.options ?? {}, await edit.options.contents));
 			}
 		}
 

src/vs/workbench/contrib/interactiveEditor/browser/interactiveEditorController.ts

@@ -42,7 +42,6 @@ import { InteractiveEditorZoneWidget } from 'vs/workbench/contrib/interactiveEdi
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { URI } from 'vs/base/common/uri';
 import { isEqual } from 'vs/base/common/resources';
-import { decodeBase64 } from 'vs/base/common/buffer';
 
 
 type Exchange = { req: IInteractiveEditorRequest; res: IInteractiveEditorResponse };
@@ -157,7 +156,7 @@ class InlineDiffDecorations {
 export class EditResponse {
 
 	readonly localEdits: TextEdit[] = [];
-	readonly singleCreateFileEdit: { uri: URI; edits: TextEdit[] } | undefined;
+	readonly singleCreateFileEdit: { uri: URI; edits: Promise<TextEdit>[] } | undefined;
 	readonly workspaceEdits: ResourceEdit[] | undefined;
 	readonly workspaceEditsIncludeLocalEdits: boolean = false;
 
@@ -184,9 +183,8 @@ export class EditResponse {
 							this.singleCreateFileEdit = undefined;
 						} else {
 							this.singleCreateFileEdit = { uri: edit.newResource, edits: [] };
-							if (edit.options.contentsBase64) {
-								const newText = decodeBase64(edit.options.contentsBase64).toString();
-								this.singleCreateFileEdit.edits.push({ range: new Range(1, 1, 1, 1), text: newText });
+							if (edit.options.contents) {
+								this.singleCreateFileEdit.edits.push(edit.options.contents.then(x => ({ range: new Range(1, 1, 1, 1), text: x.toString() })));
 							}
 						}
 					}
@@ -197,7 +195,7 @@ export class EditResponse {
 						this.workspaceEditsIncludeLocalEdits = true;
 
 					} else if (isEqual(this.singleCreateFileEdit?.uri, edit.resource)) {
-						this.singleCreateFileEdit!.edits.push(edit.textEdit);
+						this.singleCreateFileEdit!.edits.push(Promise.resolve(edit.textEdit));
 					} else {
 						isComplexEdit = true;
 					}
@@ -604,7 +602,7 @@ export class InteractiveEditorController implements IEditorContribution {
 			this._zone.widget.updateToolbar(true);
 
 			if (editResponse.singleCreateFileEdit) {
-				this._zone.widget.showCreatePreview(editResponse.singleCreateFileEdit.uri, editResponse.singleCreateFileEdit.edits);
+				this._zone.widget.showCreatePreview(editResponse.singleCreateFileEdit.uri, await Promise.all(editResponse.singleCreateFileEdit.edits));
 			} else {
 				this._zone.widget.hideCreatePreview();
 			}

src/vscode-dts/vscode.d.ts

@@ -3710,7 +3710,18 @@ declare module 'vscode' {
 		 * the file is being created with.
 		 * @param metadata Optional metadata for the entry.
 		 */
-		createFile(uri: Uri, options?: { readonly overwrite?: boolean; readonly ignoreIfExists?: boolean; readonly contents?: Uint8Array }, metadata?: WorkspaceEditEntryMetadata): void;
+		createFile(uri: Uri, options?: {
+			readonly overwrite?: boolean;
+			readonly ignoreIfExists?: boolean;
+
+			/**
+			 * The initial contents of the new file.
+			 *
+			 * If creating a file from a {@link DocumentDropEditProvider drop operation}, you can
+			 * pass in a {@link DataTransferFile} to improve performance by avoiding extra data copying.
+			 */
+			readonly contents?: Uint8Array | DataTransferFile;
+		}, metadata?: WorkspaceEditEntryMetadata): void;
 
 		/**
 		 * Delete a file or folder.
@@ -10414,6 +10425,8 @@ declare module 'vscode' {
 
 	/**
 	 * A file associated with a {@linkcode DataTransferItem}.
+	 *
+	 * Instances of this type can only be created by the editor and not by extensions.
 	 */
 	export interface DataTransferFile {
 		/**