Optimistically show paste with imports if TS server takes to long w…

…hen computing imports to add

Fixes microsoft#235959

When pasting across files, sometimes users can see a spinner while we wait to see if TS has a `paste with imports` edit. This needs to be done even if we insert plain text by default so that we can then show the paste widget and user can select `paste with imports`

This change makes it so that if TS takes over 200ms, we assume that a paste operation may be available and return an unresolved edit for it. Only when the edit then actually needs to be applied, we wait for TS to respond

In most cases we won't need to wait because plain text will be pasted by default. Users will then only see the spinner after selecting `paste with imports`

extensions/typescript-language-features/src/languageFeatures/copyPaste.ts

@@ -9,39 +9,61 @@ import { LanguageDescription } from '../configuration/languageDescription';
 import { API } from '../tsServer/api';
 import protocol from '../tsServer/protocol/protocol';
 import * as typeConverters from '../typeConverters';
-import { ClientCapability, ITypeScriptServiceClient } from '../typescriptService';
+import { ClientCapability, ITypeScriptServiceClient, ServerResponse } from '../typescriptService';
+import { raceTimeout } from '../utils/async';
 import FileConfigurationManager from './fileConfigurationManager';
 import { conditionalRegistration, requireGlobalConfiguration, requireMinVersion, requireSomeCapability } from './util/dependentRegistration';
 
 class CopyMetadata {
 	constructor(
-		readonly resource: vscode.Uri,
-		readonly ranges: readonly vscode.Range[],
+		public readonly resource: vscode.Uri,
+		public readonly ranges: readonly vscode.Range[],
+		public readonly copyOperation: Promise<ServerResponse.Response<protocol.PreparePasteEditsResponse>> | undefined
 	) { }
+}
 
-	toJSON() {
-		return JSON.stringify({
-			resource: this.resource.toJSON(),
-			ranges: this.ranges,
-		});
-	}
+class TsPasteEdit extends vscode.DocumentPasteEdit {
 
-	static fromJSON(str: string): CopyMetadata | undefined {
-		try {
-			const parsed = JSON.parse(str);
-			return new CopyMetadata(
-				vscode.Uri.from(parsed.resource),
-				parsed.ranges.map((r: any) => new vscode.Range(r[0].line, r[0].character, r[1].line, r[1].character)));
-		} catch {
-			// ignore
+	static tryCreateFromResponse(
+		client: ITypeScriptServiceClient,
+		response: ServerResponse.Response<protocol.GetPasteEditsResponse>
+	): TsPasteEdit | undefined {
+		if (response.type !== 'response' || !response.body?.edits.length) {
+			return undefined;
 		}
-		return undefined;
+
+		const pasteEdit = new TsPasteEdit();
+
+		const additionalEdit = new vscode.WorkspaceEdit();
+		for (const edit of response.body.edits) {
+			additionalEdit.set(client.toResource(edit.fileName), edit.textChanges.map(typeConverters.TextEdit.fromCodeEdit));
+		}
+		pasteEdit.additionalEdit = additionalEdit;
+
+		return pasteEdit;
+	}
+
+	constructor() {
+		super('', vscode.l10n.t("Paste with imports"), DocumentPasteProvider.kind);
+		this.yieldTo = [
+			vscode.DocumentDropOrPasteEditKind.Text.append('plain')
+		];
+	}
+}
+
+class TsPendingPasteEdit extends TsPasteEdit {
+	constructor(
+		text: string,
+		public readonly operation: Promise<ServerResponse.Response<protocol.GetPasteEditsResponse>>
+	) {
+		super();
+		this.insertText = text;
 	}
 }
 
 const enabledSettingId = 'updateImportsOnPaste.enabled';
 
-class DocumentPasteProvider implements vscode.DocumentPasteEditProvider {
+class DocumentPasteProvider implements vscode.DocumentPasteEditProvider<TsPasteEdit> {
 
 	static readonly kind = vscode.DocumentDropOrPasteEditKind.TextUpdateImports.append('jsts');
 	static readonly metadataMimeType = 'application/vnd.code.jsts.metadata';
@@ -62,16 +84,32 @@ class DocumentPasteProvider implements vscode.DocumentPasteEditProvider {
 			return;
 		}
 
-		const response = await this._client.interruptGetErr(() => this._client.execute('preparePasteEdits', {
+		const copyRequest = this._client.interruptGetErr(() => this._client.execute('preparePasteEdits', {
 			file,
 			copiedTextSpan: ranges.map(typeConverters.Range.toTextSpan),
 		}, token));
-		if (token.isCancellationRequested || response.type !== 'response' || !response.body) {
+
+		const copyTimeout = 200;
+		const response = await raceTimeout(copyRequest, copyTimeout);
+		if (token.isCancellationRequested) {
 			return;
 		}
 
-		dataTransfer.set(DocumentPasteProvider.metadataMimeType,
-			new vscode.DataTransferItem(new CopyMetadata(document.uri, ranges).toJSON()));
+		if (response) {
+			if (response.type !== 'response' || !response.body) {
+				// We got a response which told us no to bother with the paste
+				// Don't store anything so that we don't trigger on paste
+				return;
+			}
+
+			dataTransfer.set(DocumentPasteProvider.metadataMimeType,
+				new vscode.DataTransferItem(new CopyMetadata(document.uri, ranges, undefined)));
+		} else {
+			// We are still waiting on the response. Store the pending request so that we can try checking it on paste
+			// when it has hopefully resolved
+			dataTransfer.set(DocumentPasteProvider.metadataMimeType,
+				new vscode.DataTransferItem(new CopyMetadata(document.uri, ranges, copyRequest)));
+		}
 	}
 
 	async provideDocumentPasteEdits(
@@ -80,7 +118,7 @@ class DocumentPasteProvider implements vscode.DocumentPasteEditProvider {
 		dataTransfer: vscode.DataTransfer,
 		_context: vscode.DocumentPasteEditContext,
 		token: vscode.CancellationToken,
-	): Promise<vscode.DocumentPasteEdit[] | undefined> {
+	): Promise<TsPasteEdit[] | undefined> {
 		if (!this.isEnabled(document)) {
 			return;
 		}
@@ -114,42 +152,68 @@ class DocumentPasteProvider implements vscode.DocumentPasteEditProvider {
 		}
 
 		if (copiedFrom?.file === file) {
+			// We are pasting in the same file we copied from. No need to do anything
 			return;
 		}
 
-		const response = await this._client.interruptGetErr(() => {
-			this.fileConfigurationManager.ensureConfigurationForDocument(document, token);
+		const pasteCts = new vscode.CancellationTokenSource();
+		token.onCancellationRequested(() => pasteCts.cancel());
 
-			return this._client.execute('getPasteEdits', {
-				file,
-				// TODO: only supports a single paste for now
-				pastedText: [text],
-				pasteLocations: ranges.map(typeConverters.Range.toTextSpan),
-				copiedFrom
-			}, token);
+		// If we have a copy operation, use that to potentially eagerly cancel the paste if it resolves to false
+		metadata?.copyOperation?.then(copyResponse => {
+			if (copyResponse.type !== 'response' || !copyResponse.body) {
+				pasteCts.cancel();
+			}
+		}, (_err) => {
+			// Expected. May have been cancelled.
 		});
-		if (response.type !== 'response' || !response.body?.edits.length || token.isCancellationRequested) {
-			return;
-		}
 
-		const edit = new vscode.DocumentPasteEdit('', vscode.l10n.t("Paste with imports"), DocumentPasteProvider.kind);
-		edit.yieldTo = [vscode.DocumentDropOrPasteEditKind.Text.append('plain')];
+		try {
+			const pasteOperation = this._client.interruptGetErr(() => {
+				this.fileConfigurationManager.ensureConfigurationForDocument(document, token);
+
+				return this._client.execute('getPasteEdits', {
+					file,
+					// TODO: only supports a single paste for now
+					pastedText: [text],
+					pasteLocations: ranges.map(typeConverters.Range.toTextSpan),
+					copiedFrom
+				}, pasteCts.token);
+			});
+
+			const pasteTimeout = 200;
+			const response = await raceTimeout(pasteOperation, pasteTimeout);
+			if (response) {
+				// Success, can return real paste edit.
+				const edit = TsPendingPasteEdit.tryCreateFromResponse(this._client, response);
+				return edit ? [edit] : undefined;
+			} else {
+				// Still waiting on the response. Eagerly return a paste edit that we will resolve when we
+				// really need to apply it
+				return [new TsPendingPasteEdit(text, pasteOperation)];
+			}
+		} finally {
+			pasteCts.dispose();
+		}
+	}
 
-		const additionalEdit = new vscode.WorkspaceEdit();
-		for (const edit of response.body.edits) {
-			additionalEdit.set(this._client.toResource(edit.fileName), edit.textChanges.map(typeConverters.TextEdit.fromCodeEdit));
+	async resolveDocumentPasteEdit(inEdit: TsPasteEdit, _token: vscode.CancellationToken): Promise<TsPasteEdit | undefined> {
+		if (!(inEdit instanceof TsPendingPasteEdit)) {
+			return;
 		}
-		edit.additionalEdit = additionalEdit;
-		return [edit];
+
+		const response = await inEdit.operation;
+		const pasteEdit = TsPendingPasteEdit.tryCreateFromResponse(this._client, response);
+		return pasteEdit ?? inEdit;
 	}
 
 	private async extractMetadata(dataTransfer: vscode.DataTransfer, token: vscode.CancellationToken): Promise<CopyMetadata | undefined> {
-		const metadata = await dataTransfer.get(DocumentPasteProvider.metadataMimeType)?.asString();
+		const metadata = await dataTransfer.get(DocumentPasteProvider.metadataMimeType)?.value;
 		if (token.isCancellationRequested) {
 			return undefined;
 		}
 
-		return metadata ? CopyMetadata.fromJSON(metadata) : undefined;
+		return metadata instanceof CopyMetadata ? metadata : undefined;
 	}
 
 	private isEnabled(document: vscode.TextDocument) {

extensions/typescript-language-features/src/utils/async.ts

@@ -161,3 +161,17 @@ export class Throttler {
 		this.isDisposed = true;
 	}
 }
+
+export function raceTimeout<T>(promise: Promise<T>, timeout: number, onTimeout?: () => void): Promise<T | undefined> {
+	let promiseResolve: ((value: T | undefined) => void) | undefined = undefined;
+
+	const timer = setTimeout(() => {
+		promiseResolve?.(undefined);
+		onTimeout?.();
+	}, timeout);
+
+	return Promise.race([
+		promise.finally(() => clearTimeout(timer)),
+		new Promise<T | undefined>(resolve => promiseResolve = resolve)
+	]);
+}

src/vs/editor/contrib/dropOrPasteInto/browser/copyPasteController.ts

@@ -9,7 +9,7 @@ import { coalesce } from '../../../../base/common/arrays.js';
 import { CancelablePromise, createCancelablePromise, DeferredPromise, raceCancellation } from '../../../../base/common/async.js';
 import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js';
 import { createStringDataTransferItem, IReadonlyVSDataTransfer, matchesMimeType, UriList, VSDataTransfer } from '../../../../base/common/dataTransfer.js';
-import { CancellationError, isCancellationError } from '../../../../base/common/errors.js';
+import { isCancellationError } from '../../../../base/common/errors.js';
 import { HierarchicalKind } from '../../../../base/common/hierarchicalKind.js';
 import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js';
 import { Mimes } from '../../../../base/common/mime.js';
@@ -379,27 +379,22 @@ export class CopyPasteController extends Disposable implements IEditorContributi
 
 				if (editSession.edits.length) {
 					const canShowWidget = editor.getOption(EditorOption.pasteAs).showPasteSelector === 'afterPaste';
-					return this._postPasteWidgetManager.applyEditAndShowIfNeeded(selections, { activeEditIndex: this.getInitialActiveEditIndex(model, editSession.edits), allEdits: editSession.edits }, canShowWidget, (edit, token) => {
-						return new Promise<PasteEditWithProvider>((resolve, reject) => {
-							(async () => {
-								try {
-									const resolveP = edit.provider.resolveDocumentPasteEdit?.(edit, token);
-									const showP = new DeferredPromise<void>();
-									const resolved = resolveP && await this._pasteProgressManager.showWhile(selections[0].getEndPosition(), localize('resolveProcess', "Resolving paste edit. Click to cancel"), Promise.race([showP.p, resolveP]), {
-										cancel: () => {
-											showP.cancel();
-											return reject(new CancellationError());
-										}
-									}, 0);
-									if (resolved) {
-										edit.additionalEdit = resolved.additionalEdit;
-									}
-									return resolve(edit);
-								} catch (err) {
-									return reject(err);
-								}
-							})();
-						});
+					return this._postPasteWidgetManager.applyEditAndShowIfNeeded(selections, { activeEditIndex: this.getInitialActiveEditIndex(model, editSession.edits), allEdits: editSession.edits }, canShowWidget, async (edit, resolveToken) => {
+						if (!edit.provider.resolveDocumentPasteEdit) {
+							return edit;
+						}
+
+						const resolveP = edit.provider.resolveDocumentPasteEdit(edit, resolveToken);
+						const showP = new DeferredPromise<void>();
+						const resolved = await this._pasteProgressManager.showWhile(selections[0].getEndPosition(), localize('resolveProcess', "Resolving paste edit for '{0}'. Click to cancel", edit.title), raceCancellation(Promise.race([showP.p, resolveP]), resolveToken), {
+							cancel: () => showP.cancel()
+						}, 0);
+
+						if (resolved) {
+							edit.insertText = resolved.insertText;
+							edit.additionalEdit = resolved.additionalEdit;
+						}
+						return edit;
 					}, token);
 				}
 

src/vs/editor/contrib/dropOrPasteInto/browser/postEditWidget.ts

@@ -6,6 +6,7 @@
 import * as dom from '../../../../base/browser/dom.js';
 import { Button } from '../../../../base/browser/ui/button/button.js';
 import { IAction } from '../../../../base/common/actions.js';
+import { raceCancellationError } from '../../../../base/common/async.js';
 import { CancellationToken } from '../../../../base/common/cancellation.js';
 import { Codicon } from '../../../../base/common/codicons.js';
 import { toErrorMessage } from '../../../../base/common/errorMessage.js';
@@ -25,6 +26,7 @@ import { IBulkEditResult, IBulkEditService } from '../../../browser/services/bul
 import { Range } from '../../../common/core/range.js';
 import { DocumentDropEdit, DocumentPasteEdit } from '../../../common/languages.js';
 import { TrackedRangeStickiness } from '../../../common/model.js';
+import { CodeEditorStateFlag, EditorStateCancellationTokenSource } from '../../editorState/browser/editorState.js';
 import { createCombinedWorkspaceEdit } from './edit.js';
 import './postEditWidget.css';
 
@@ -169,11 +171,11 @@ export class PostEditWidgetManager<T extends DocumentPasteEdit | DocumentDropEdi
 	}
 
 	public async applyEditAndShowIfNeeded(ranges: readonly Range[], edits: EditSet<T>, canShowWidget: boolean, resolve: (edit: T, token: CancellationToken) => Promise<T>, token: CancellationToken) {
-		const model = this._editor.getModel();
-		if (!model || !ranges.length) {
+		if (!ranges.length || !this._editor.hasModel()) {
 			return;
 		}
 
+		const model = this._editor.getModel();
 		const edit = edits.allEdits.at(edits.activeEditIndex);
 		if (!edit) {
 			return;
@@ -200,11 +202,14 @@ export class PostEditWidgetManager<T extends DocumentPasteEdit | DocumentDropEdi
 			}
 		};
 
+		const editorStateCts = new EditorStateCancellationTokenSource(this._editor, CodeEditorStateFlag.Value | CodeEditorStateFlag.Selection, undefined, token);
 		let resolvedEdit: T;
 		try {
-			resolvedEdit = await resolve(edit, token);
+			resolvedEdit = await raceCancellationError(resolve(edit, editorStateCts.token), editorStateCts.token);
 		} catch (e) {
 			return handleError(e, localize('resolveError', "Error resolving edit '{0}':\n{1}", edit.title, toErrorMessage(e)));
+		} finally {
+			editorStateCts.dispose();
 		}
 
 		if (token.isCancellationRequested) {

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

@@ -1088,6 +1088,10 @@ class MainThreadPasteEditProvider implements languages.DocumentPasteEditProvider
 		if (metadata.supportsResolve) {
 			this.resolveDocumentPasteEdit = async (edit: languages.DocumentPasteEdit, token: CancellationToken) => {
 				const resolved = await this._proxy.$resolvePasteEdit(this._handle, (<IPasteEditDto>edit)._cacheId!, token);
+				if (typeof resolved.insertText !== 'undefined') {
+					edit.insertText = resolved.insertText;
+				}
+
 				if (resolved.additionalEdit) {
 					edit.additionalEdit = reviveWorkspaceEditDto(resolved.additionalEdit, this._uriIdentService);
 				}

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

@@ -2291,7 +2291,7 @@ export interface ExtHostLanguageFeaturesShape {
 	$releaseCodeActions(handle: number, cacheId: number): void;
 	$prepareDocumentPaste(handle: number, uri: UriComponents, ranges: readonly IRange[], dataTransfer: DataTransferDTO, token: CancellationToken): Promise<DataTransferDTO | undefined>;
 	$providePasteEdits(handle: number, requestId: number, uri: UriComponents, ranges: IRange[], dataTransfer: DataTransferDTO, context: IDocumentPasteContextDto, token: CancellationToken): Promise<IPasteEditDto[] | undefined>;
-	$resolvePasteEdit(handle: number, id: ChainedCacheId, token: CancellationToken): Promise<{ additionalEdit?: IWorkspaceEditDto }>;
+	$resolvePasteEdit(handle: number, id: ChainedCacheId, token: CancellationToken): Promise<{ insertText?: string; additionalEdit?: IWorkspaceEditDto }>;
 	$releasePasteEdits(handle: number, cacheId: number): void;
 	$provideDocumentFormattingEdits(handle: number, resource: UriComponents, options: languages.FormattingOptions, token: CancellationToken): Promise<languages.TextEdit[] | undefined>;
 	$provideDocumentRangeFormattingEdits(handle: number, resource: UriComponents, range: IRange, options: languages.FormattingOptions, token: CancellationToken): Promise<languages.TextEdit[] | undefined>;

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

@@ -649,16 +649,18 @@ class DocumentPasteEditProvider {
 		}));
 	}
 
-	async resolvePasteEdit(id: extHostProtocol.ChainedCacheId, token: CancellationToken): Promise<{ additionalEdit?: extHostProtocol.IWorkspaceEditDto }> {
+	async resolvePasteEdit(id: extHostProtocol.ChainedCacheId, token: CancellationToken): Promise<{ insertText?: string | vscode.SnippetString; additionalEdit?: extHostProtocol.IWorkspaceEditDto }> {
 		const [sessionId, itemId] = id;
 		const item = this._cache.get(sessionId, itemId);
 		if (!item || !this._provider.resolveDocumentPasteEdit) {
 			return {}; // this should not happen...
 		}
 
 		const resolvedItem = (await this._provider.resolveDocumentPasteEdit(item, token)) ?? item;
-		const additionalEdit = resolvedItem.additionalEdit ? typeConvert.WorkspaceEdit.from(resolvedItem.additionalEdit, undefined) : undefined;
-		return { additionalEdit };
+		return {
+			insertText: resolvedItem.insertText,
+			additionalEdit: resolvedItem.additionalEdit ? typeConvert.WorkspaceEdit.from(resolvedItem.additionalEdit, undefined) : undefined
+		};
 	}
 
 	releasePasteEdits(id: number): any {

src/vscode-dts/vscode.d.ts

@@ -6324,7 +6324,7 @@ declare module 'vscode' {
 		 * Optional method which fills in the {@linkcode DocumentPasteEdit.additionalEdit} before the edit is applied.
 		 *
 		 * This is called once per edit and should be used if generating the complete edit may take a long time.
-		 * Resolve can only be used to change {@linkcode DocumentPasteEdit.additionalEdit}.
+		 * Resolve can only be used to change {@linkcode DocumentPasteEdit.insertText} or {@linkcode DocumentPasteEdit.additionalEdit}.
 		 *
 		 * @param pasteEdit The {@linkcode DocumentPasteEdit} to resolve.
 		 * @param token A cancellation token.