testing: improve decoration syncing (microsoft#159705)

So, two problems. One, we didn't actually fully re-sync testing
decorations when explicitly updated by the extension, since we attempted
to get the test URI from the _diff_ which would never actually be
present on single updates (since the test item URI cannot be changed).
I think this was the main problem people saw.

So, this fixes that. It also applies a change so that we only sync the
ranges to what the extension gives us if the document version is up to
date with what was in the extension host. This should avoid syncing
decorations to the wrong place--instead just use VS Code's own decoration
location tracking until we get a newer update from the extension.

Fixes microsoft#158475
Fixes microsoft#153304

Author: connor4312

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

@@ -183,7 +183,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 	const extHostWebviewPanels = rpcProtocol.set(ExtHostContext.ExtHostWebviewPanels, new ExtHostWebviewPanels(rpcProtocol, extHostWebviews, extHostWorkspace));
 	const extHostCustomEditors = rpcProtocol.set(ExtHostContext.ExtHostCustomEditors, new ExtHostCustomEditors(rpcProtocol, extHostDocuments, extensionStoragePaths, extHostWebviews, extHostWebviewPanels));
 	const extHostWebviewViews = rpcProtocol.set(ExtHostContext.ExtHostWebviewViews, new ExtHostWebviewViews(rpcProtocol, extHostWebviews));
-	const extHostTesting = rpcProtocol.set(ExtHostContext.ExtHostTesting, new ExtHostTesting(rpcProtocol, extHostCommands));
+	const extHostTesting = rpcProtocol.set(ExtHostContext.ExtHostTesting, new ExtHostTesting(rpcProtocol, extHostCommands, extHostDocumentsAndEditors));
 	const extHostUriOpeners = rpcProtocol.set(ExtHostContext.ExtHostUriOpeners, new ExtHostUriOpeners(rpcProtocol));
 	rpcProtocol.set(ExtHostContext.ExtHostInteractive, new ExtHostInteractive(rpcProtocol, extHostNotebook, extHostDocumentsAndEditors, extHostCommands, extHostLogService));
 

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

@@ -10,6 +10,7 @@ import { denamespaceTestTag, ITestItem, ITestItemContext } from 'vs/workbench/co
 import type * as vscode from 'vscode';
 import * as Convert from 'vs/workbench/api/common/extHostTypeConverters';
 import { URI } from 'vs/base/common/uri';
+import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 
 const testItemPropAccessor = <K extends keyof vscode.TestItem>(
 	api: IExtHostTestItemApi,
@@ -163,9 +164,10 @@ export class TestItemRootImpl extends TestItemImpl {
 }
 
 export class ExtHostTestItemCollection extends TestItemCollection<TestItemImpl> {
-	constructor(controllerId: string, controllerLabel: string) {
+	constructor(controllerId: string, controllerLabel: string, editors: ExtHostDocumentsAndEditors) {
 		super({
 			controllerId,
+			getDocumentVersion: (uri: URI) => editors.getDocument(uri)?.version,
 			getApiFor: getPrivateApiFor as (impl: TestItemImpl) => ITestItemApi<TestItemImpl>,
 			getChildren: (item) => item.children as ITestChildrenLike<TestItemImpl>,
 			root: new TestItemRootImpl(controllerId, controllerLabel),

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

@@ -16,6 +16,7 @@ import { isDefined } from 'vs/base/common/types';
 import { generateUuid } from 'vs/base/common/uuid';
 import { ExtHostTestingShape, ILocationDto, MainContext, MainThreadTestingShape } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
+import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 import { IExtHostRpcService } from 'vs/workbench/api/common/extHostRpcService';
 import { ExtHostTestItemCollection, TestItemImpl, TestItemRootImpl, toItemFromContext } from 'vs/workbench/api/common/extHostTestItem';
 import * as Convert from 'vs/workbench/api/common/extHostTypeConverters';
@@ -41,7 +42,11 @@ export class ExtHostTesting implements ExtHostTestingShape {
 	public onResultsChanged = this.resultsChangedEmitter.event;
 	public results: ReadonlyArray<vscode.TestRunResult> = [];
 
-	constructor(@IExtHostRpcService rpc: IExtHostRpcService, commands: ExtHostCommands) {
+	constructor(
+		@IExtHostRpcService rpc: IExtHostRpcService,
+		commands: ExtHostCommands,
+		private readonly editors: ExtHostDocumentsAndEditors,
+	) {
 		this.proxy = rpc.getProxy(MainContext.MainThreadTesting);
 		this.observer = new TestObservers(this.proxy);
 		this.runTracker = new TestRunCoordinator(this.proxy);
@@ -61,7 +66,7 @@ export class ExtHostTesting implements ExtHostTestingShape {
 		}
 
 		const disposable = new DisposableStore();
-		const collection = disposable.add(new ExtHostTestItemCollection(controllerId, label));
+		const collection = disposable.add(new ExtHostTestItemCollection(controllerId, label, this.editors));
 		collection.root.label = label;
 
 		const profiles = new Map<number, vscode.TestRunProfile>();

src/vs/workbench/api/test/browser/extHostTesting.test.ts

@@ -17,6 +17,7 @@ import { Location, Position, Range, TestMessage, TestResultState, TestRunProfile
 import { TestDiffOpType, TestItemExpandState, TestMessageType, TestsDiff } from 'vs/workbench/contrib/testing/common/testTypes';
 import { TestId } from 'vs/workbench/contrib/testing/common/testId';
 import type { TestItem, TestRunRequest } from 'vscode';
+import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
 
 const simplify = (item: TestItem) => ({
 	id: item.id,
@@ -69,7 +70,9 @@ suite('ExtHost Testing', () => {
 
 	let single: TestExtHostTestItemCollection;
 	setup(() => {
-		single = new TestExtHostTestItemCollection('ctrlId', 'root');
+		single = new TestExtHostTestItemCollection('ctrlId', 'root', {
+			getDocument: () => undefined,
+		} as Partial<ExtHostDocumentsAndEditors> as ExtHostDocumentsAndEditors);
 		single.resolveHandler = item => {
 			if (item === undefined) {
 				const a = new TestItemImpl('ctrlId', 'id-a', 'a', URI.file('/'));
@@ -150,7 +153,7 @@ suite('ExtHost Testing', () => {
 			assert.deepStrictEqual(single.collectDiff(), [
 				{
 					op: TestDiffOpType.Update,
-					item: { extId: new TestId(['ctrlId', 'id-a']).toString(), item: { description: 'Hello world' } },
+					item: { extId: new TestId(['ctrlId', 'id-a']).toString(), docv: undefined, item: { description: 'Hello world' } },
 				}
 			]);
 		});
@@ -251,15 +254,15 @@ suite('ExtHost Testing', () => {
 			assert.deepStrictEqual(single.collectDiff(), [
 				{
 					op: TestDiffOpType.Update,
-					item: { extId: new TestId(['ctrlId', 'id-a']).toString(), expand: TestItemExpandState.Expanded, item: { label: 'Hello world' } },
+					item: { extId: new TestId(['ctrlId', 'id-a']).toString(), expand: TestItemExpandState.Expanded, docv: undefined, item: { label: 'Hello world' } },
 				},
 			]);
 
 			newA.label = 'still connected';
 			assert.deepStrictEqual(single.collectDiff(), [
 				{
 					op: TestDiffOpType.Update,
-					item: { extId: new TestId(['ctrlId', 'id-a']).toString(), item: { label: 'still connected' } }
+					item: { extId: new TestId(['ctrlId', 'id-a']).toString(), docv: undefined, item: { label: 'still connected' } }
 				},
 			]);
 
@@ -286,7 +289,7 @@ suite('ExtHost Testing', () => {
 				},
 				{
 					op: TestDiffOpType.Update,
-					item: { extId: TestId.fromExtHostTestItem(oldAB, 'ctrlId').toString(), item: { label: 'Hello world' } },
+					item: { extId: TestId.fromExtHostTestItem(oldAB, 'ctrlId').toString(), docv: undefined, item: { label: 'Hello world' } },
 				},
 			]);
 
@@ -296,11 +299,11 @@ suite('ExtHost Testing', () => {
 			assert.deepStrictEqual(single.collectDiff(), [
 				{
 					op: TestDiffOpType.Update,
-					item: { extId: new TestId(['ctrlId', 'id-a', 'id-aa']).toString(), item: { label: 'still connected1' } }
+					item: { extId: new TestId(['ctrlId', 'id-a', 'id-aa']).toString(), docv: undefined, item: { label: 'still connected1' } }
 				},
 				{
 					op: TestDiffOpType.Update,
-					item: { extId: new TestId(['ctrlId', 'id-a', 'id-ab']).toString(), item: { label: 'still connected2' } }
+					item: { extId: new TestId(['ctrlId', 'id-a', 'id-ab']).toString(), docv: undefined, item: { label: 'still connected2' } }
 				},
 			]);
 
@@ -330,7 +333,7 @@ suite('ExtHost Testing', () => {
 			assert.deepStrictEqual(single.collectDiff(), [
 				{
 					op: TestDiffOpType.Update,
-					item: { extId: new TestId(['ctrlId', 'id-a', 'id-b']).toString(), item: { label: 'still connected' } }
+					item: { extId: new TestId(['ctrlId', 'id-a', 'id-b']).toString(), docv: undefined, item: { label: 'still connected' } }
 				},
 			]);
 

src/vs/workbench/contrib/testing/browser/testingDecorations.ts

@@ -75,8 +75,8 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 	private generation = 0;
 	private readonly changeEmitter = new Emitter<void>();
 	private readonly decorationCache = new ResourceMap<{
-		/** Whether tests in the resource have been updated, requiring rerendering */
-		testRangesUpdated: boolean;
+		/** The document version at which ranges have been updated, requiring rerendering */
+		rangeUpdateVersionId?: number;
 		/** Counter for the results rendered in the document */
 		generation: number;
 		value: TestDecorations<ITestDecoration>;
@@ -115,16 +115,14 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 		// is up to date. This prevents issues, as in #138632, #138835, #138922.
 		this._register(this.testService.onWillProcessDiff(diff => {
 			for (const entry of diff) {
-				let uri: URI | undefined | null;
-				if (entry.op === TestDiffOpType.Add || entry.op === TestDiffOpType.Update) {
-					uri = entry.item.item?.uri;
-				} else if (entry.op === TestDiffOpType.Remove) {
-					uri = this.testService.collection.getNodeById(entry.itemId)?.item.uri;
+				if (entry.op !== TestDiffOpType.Update || entry.item.docv === undefined || entry.item.item?.range === undefined) {
+					continue;
 				}
 
+				const uri = this.testService.collection.getNodeById(entry.item.extId)?.item.uri;
 				const rec = uri && this.decorationCache.get(uri);
 				if (rec) {
-					rec.testRangesUpdated = true;
+					rec.rangeUpdateVersionId = entry.item.docv;
 				}
 			}
 
@@ -160,7 +158,7 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 		}
 
 		const cached = this.decorationCache.get(resource);
-		if (cached && cached.generation === this.generation && !cached.testRangesUpdated) {
+		if (cached && cached.generation === this.generation && (cached.rangeUpdateVersionId === undefined || cached.rangeUpdateVersionId !== model.getVersionId())) {
 			return cached.value;
 		}
 
@@ -194,7 +192,7 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 		const gutterEnabled = getTestingConfiguration(this.configurationService, TestingConfigKeys.GutterEnabled);
 		const uriStr = model.uri.toString();
 		const cached = this.decorationCache.get(model.uri);
-		const testRangesUpdated = cached?.testRangesUpdated;
+		const testRangesUpdated = cached?.rangeUpdateVersionId === model.getVersionId();
 		const lastDecorations = cached?.value ?? new TestDecorations();
 		const newDecorations = new TestDecorations<ITestDecoration>();
 
@@ -298,7 +296,7 @@ export class TestingDecorationService extends Disposable implements ITestingDeco
 
 			this.decorationCache.set(model.uri, {
 				generation: this.generation,
-				testRangesUpdated: false,
+				rangeUpdateVersionId: cached?.rangeUpdateVersionId,
 				value: newDecorations,
 			});
 		});

src/vs/workbench/contrib/testing/common/testItemCollection.ts

@@ -9,6 +9,7 @@ import { Disposable } from 'vs/base/common/lifecycle';
 import { assertNever } from 'vs/base/common/assert';
 import { applyTestItemUpdate, ITestItem, ITestTag, namespaceTestTag, TestDiffOpType, TestItemExpandState, TestsDiff, TestsDiffOp } from 'vs/workbench/contrib/testing/common/testTypes';
 import { TestId } from 'vs/workbench/contrib/testing/common/testId';
+import { URI } from 'vs/base/common/uri';
 
 /**
  * @private
@@ -82,6 +83,9 @@ export interface ITestItemCollectionOptions<T> {
 	/** Controller ID to use to prefix these test items. */
 	controllerId: string;
 
+	/** Gets the document version at the given URI, if it's open */
+	getDocumentVersion(uri: URI | undefined): number | undefined;
+
 	/** Gets API for the given test item, used to listen for events and set parents. */
 	getApiFor(item: T): ITestItemApi<T>;
 
@@ -142,6 +146,7 @@ export interface ITestChildrenLike<T> extends Iterable<[string, T]> {
 export interface ITestItemLike {
 	id: string;
 	tags: readonly ITestTag[];
+	uri?: URI;
 	canResolveChildren: boolean;
 }
 
@@ -283,7 +288,11 @@ export class TestItemCollection<T extends ITestItemLike> extends Disposable {
 			case TestItemEventOp.SetProp:
 				this.pushDiff({
 					op: TestDiffOpType.Update,
-					item: { extId: internal.fullId.toString(), item: evt.update }
+					item: {
+						extId: internal.fullId.toString(),
+						item: evt.update,
+						docv: this.options.getDocumentVersion(internal.actual.uri),
+					}
 				});
 				break;
 			default:

src/vs/workbench/contrib/testing/common/testTypes.ts

@@ -369,13 +369,20 @@ export interface ITestItemUpdate {
 	extId: string;
 	expand?: TestItemExpandState;
 	item?: Partial<ITestItem>;
+
+	/**
+	 * The document version at the time the operation was made, if the test has
+	 * a URI and the document was open in the extension host.
+	 */
+	docv?: number;
 }
 
 export namespace ITestItemUpdate {
 	export interface Serialized {
 		extId: string;
 		expand?: TestItemExpandState;
 		item?: Partial<ITestItem.Serialized>;
+		docv?: number;
 	}
 
 	export const serialize = (u: ITestItemUpdate): Serialized => {
@@ -392,7 +399,7 @@ export namespace ITestItemUpdate {
 			if (u.item.sortText !== undefined) { item.sortText = u.item.sortText; }
 		}
 
-		return { extId: u.extId, expand: u.expand, item };
+		return { extId: u.extId, expand: u.expand, item, docv: u.docv };
 	};
 
 	export const deserialize = (u: Serialized): ITestItemUpdate => {
@@ -408,7 +415,7 @@ export namespace ITestItemUpdate {
 			if (u.item.sortText !== undefined) { item.sortText = u.item.sortText; }
 		}
 
-		return { extId: u.extId, expand: u.expand, item };
+		return { extId: u.extId, expand: u.expand, item, docv: u.docv };
 	};
 
 }

src/vs/workbench/contrib/testing/test/common/testStubs.ts

@@ -81,6 +81,7 @@ export class TestTestCollection extends TestItemCollection<TestTestItem> {
 			getApiFor: t => t.api,
 			toITestItem: t => t.toTestItem(),
 			getChildren: t => t.children,
+			getDocumentVersion: () => undefined,
 			root: new TestTestItem(controllerId, controllerId, 'root'),
 		});
 	}