Optimistically show paste with imports if TS server takes to long when computing imports to add

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 = TsPasteEdit.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.