Another fix for "element with id not registered" error (#281269)

Fixes microsoft/vscode-pull-request-github#8073

Author: alexr00

extensions/vscode-api-tests/src/singlefolder-tests/tree.test.ts

@@ -103,4 +103,143 @@ suite('vscode API - tree', () => {
 			assert.fail(error.message);
 		}
 	});
+
+	test('TreeView - element already registered after refresh', async function () {
+		this.timeout(60_000);
+
+		type ParentElement = { readonly kind: 'parent' };
+		type ChildElement = { readonly kind: 'leaf'; readonly version: number };
+		type TreeElement = ParentElement | ChildElement;
+
+		class ParentRefreshTreeDataProvider implements vscode.TreeDataProvider<TreeElement> {
+			private readonly changeEmitter = new vscode.EventEmitter<TreeElement | undefined>();
+			private readonly rootRequestEmitter = new vscode.EventEmitter<number>();
+			private readonly childRequestEmitter = new vscode.EventEmitter<number>();
+			private readonly rootRequests: DeferredPromise<TreeElement[]>[] = [];
+			private readonly childRequests: DeferredPromise<TreeElement[]>[] = [];
+			private readonly parentElement: ParentElement = { kind: 'parent' };
+			private childVersion = 0;
+			private currentChild: ChildElement = { kind: 'leaf', version: 0 };
+
+			readonly onDidChangeTreeData = this.changeEmitter.event;
+
+			getChildren(element?: TreeElement): Thenable<TreeElement[]> {
+				if (!element) {
+					const deferred = new DeferredPromise<TreeElement[]>();
+					this.rootRequests.push(deferred);
+					this.rootRequestEmitter.fire(this.rootRequests.length);
+					return deferred.p;
+				}
+				if (element.kind === 'parent') {
+					const deferred = new DeferredPromise<TreeElement[]>();
+					this.childRequests.push(deferred);
+					this.childRequestEmitter.fire(this.childRequests.length);
+					return deferred.p;
+				}
+				return Promise.resolve([]);
+			}
+
+			getTreeItem(element: TreeElement): vscode.TreeItem {
+				if (element.kind === 'parent') {
+					const item = new vscode.TreeItem('parent', vscode.TreeItemCollapsibleState.Collapsed);
+					item.id = 'parent';
+					return item;
+				}
+				const item = new vscode.TreeItem('duplicate', vscode.TreeItemCollapsibleState.None);
+				item.id = 'dup';
+				return item;
+			}
+
+			getParent(element: TreeElement): TreeElement | undefined {
+				if (element.kind === 'leaf') {
+					return this.parentElement;
+				}
+				return undefined;
+			}
+
+			getCurrentChild(): ChildElement {
+				return this.currentChild;
+			}
+
+			replaceChild(): ChildElement {
+				this.childVersion++;
+				this.currentChild = { kind: 'leaf', version: this.childVersion };
+				return this.currentChild;
+			}
+
+			async waitForRootRequestCount(count: number): Promise<void> {
+				while (this.rootRequests.length < count) {
+					await asPromise(this.rootRequestEmitter.event);
+				}
+			}
+
+			async waitForChildRequestCount(count: number): Promise<void> {
+				while (this.childRequests.length < count) {
+					await asPromise(this.childRequestEmitter.event);
+				}
+			}
+
+			async resolveNextRootRequest(elements?: TreeElement[]): Promise<void> {
+				const next = this.rootRequests.shift();
+				if (!next) {
+					return;
+				}
+				await next.complete(elements ?? [this.parentElement]);
+			}
+
+			async resolveChildRequestAt(index: number, elements?: TreeElement[]): Promise<void> {
+				const request = this.childRequests[index];
+				if (!request) {
+					return;
+				}
+				this.childRequests.splice(index, 1);
+				await request.complete(elements ?? [this.currentChild]);
+			}
+
+			dispose(): void {
+				this.changeEmitter.dispose();
+				this.rootRequestEmitter.dispose();
+				this.childRequestEmitter.dispose();
+				while (this.rootRequests.length) {
+					this.rootRequests.shift()!.complete([]);
+				}
+				while (this.childRequests.length) {
+					this.childRequests.shift()!.complete([]);
+				}
+			}
+		}
+
+		const provider = new ParentRefreshTreeDataProvider();
+		disposables.push(provider);
+
+		const treeView = vscode.window.createTreeView('test.treeRefresh', { treeDataProvider: provider });
+		disposables.push(treeView);
+
+		const initialChild = provider.getCurrentChild();
+		const firstReveal = (treeView.reveal(initialChild, { expand: true })
+			.then(() => ({ error: undefined as Error | undefined })) as Promise<{ error: Error | undefined }>)
+			.catch(error => ({ error }));
+
+		await provider.waitForRootRequestCount(1);
+		await provider.resolveNextRootRequest();
+
+		await provider.waitForChildRequestCount(1);
+		const staleChild = provider.getCurrentChild();
+		const refreshedChild = provider.replaceChild();
+		const secondReveal = (treeView.reveal(refreshedChild, { expand: true })
+			.then(() => ({ error: undefined as Error | undefined })) as Promise<{ error: Error | undefined }>)
+			.catch(error => ({ error }));
+
+		await provider.waitForChildRequestCount(2);
+
+		await provider.resolveChildRequestAt(1, [refreshedChild]);
+		await delay(0);
+		await provider.resolveChildRequestAt(0, [staleChild]);
+
+		const [firstResult, secondResult] = await Promise.all([firstReveal, secondReveal]);
+		const error = firstResult.error ?? secondResult.error;
+		if (error && /Element with id .+ is already registered/.test(error.message)) {
+			assert.fail(error.message);
+		}
+	});
 });

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

@@ -314,13 +314,17 @@ class ExtHostTreeView<T> extends Disposable {
 
 	private static readonly LABEL_HANDLE_PREFIX = '0';
 	private static readonly ID_HANDLE_PREFIX = '1';
+	private static readonly ROOT_FETCH_KEY = Symbol('extHostTreeViewRoot');
 
 	private readonly _dataProvider: vscode.TreeDataProvider<T>;
 	private readonly _dndController: vscode.TreeDragAndDropController<T> | undefined;
 
 	private _roots: TreeNode[] | undefined = undefined;
 	private _elements: Map<TreeItemHandle, T> = new Map<TreeItemHandle, T>();
 	private _nodes: Map<T, TreeNode> = new Map<T, TreeNode>();
+	// Track the latest child-fetch per element so that refresh-triggered cache clears ignore stale results.
+	// Without these tokens, an earlier getChildren promise resolving after refresh would re-register handles and hit the duplicate-id guard.
+	private readonly _childrenFetchTokens = new Map<T | typeof ExtHostTreeView.ROOT_FETCH_KEY, number>();
 
 	private _visible: boolean = false;
 	get visible(): boolean { return this._visible; }
@@ -725,14 +729,25 @@ class ExtHostTreeView<T> extends Disposable {
 		return this._roots;
 	}
 
+	private _getFetchKey(parentElement?: T): T | typeof ExtHostTreeView.ROOT_FETCH_KEY {
+		return parentElement ?? ExtHostTreeView.ROOT_FETCH_KEY;
+	}
+
 	private async _fetchChildrenNodes(parentElement?: T): Promise<TreeNode[] | undefined> {
 		// clear children cache
 		this._addChildrenToClear(parentElement);
+		const fetchKey = this._getFetchKey(parentElement);
+		let requestId = this._childrenFetchTokens.get(fetchKey) ?? 0;
+		requestId++;
+		this._childrenFetchTokens.set(fetchKey, requestId);
 
 		const cts = new CancellationTokenSource(this._refreshCancellationSource.token);
 
 		try {
 			const elements = await this._dataProvider.getChildren(parentElement);
+			if (this._childrenFetchTokens.get(fetchKey) !== requestId) {
+				return undefined;
+			}
 			const parentNode = parentElement ? this._nodes.get(parentElement) : undefined;
 
 			if (cts.token.isCancellationRequested) {
@@ -743,12 +758,18 @@ class ExtHostTreeView<T> extends Disposable {
 			const treeItems = await Promise.all(coalesce(coalescedElements).map(element => {
 				return this._dataProvider.getTreeItem(element);
 			}));
+			if (this._childrenFetchTokens.get(fetchKey) !== requestId) {
+				return undefined;
+			}
 			if (cts.token.isCancellationRequested) {
 				return undefined;
 			}
 
 			// createAndRegisterTreeNodes adds the nodes to a cache. This must be done sync so that they get added in the correct order.
 			const items = treeItems.map((item, index) => item ? this._createAndRegisterTreeNode(coalescedElements[index], item, parentNode) : null);
+			if (this._childrenFetchTokens.get(fetchKey) !== requestId) {
+				return undefined;
+			}
 
 			return coalesce(items);
 		} finally {
@@ -1062,6 +1083,7 @@ class ExtHostTreeView<T> extends Disposable {
 		});
 		this._nodes.clear();
 		this._elements.clear();
+		this._childrenFetchTokens.clear();
 	}
 
 	private _clearNodes(nodes: TreeNode[]): void {
@@ -1075,6 +1097,7 @@ class ExtHostTreeView<T> extends Disposable {
 		this._nodes.clear();
 		dispose(this._nodesToClear);
 		this._nodesToClear.clear();
+		this._childrenFetchTokens.clear();
 	}
 
 	override dispose() {