Add tests and feedback update

Author: vijayupadya

extensions/copilot/src/platform/telemetry/node/baseExperimentationService.ts

@@ -116,7 +116,49 @@ export class UserInfoStore extends Disposable {
 	}
 }
 
-export type TASClientDelegateFn = (globalState: vscode.Memento, userInfoStore: UserInfoStore) => ITASExperimentationService;
+/**
+ * A one-way switch used to neutralize a superseded delegate. `tas-client`'s `dispose()` only
+ * stops polling; an already in-flight fetch can still complete and write shared state. Revoking
+ * the gate makes that delegate's storage/telemetry writes no-ops so it cannot overwrite the
+ * memento (`VSCode.ABExp.FeatureData`) or `abexp.assignmentcontext` after being replaced.
+ */
+export class RevocationGate {
+	private _revoked = false;
+	get isRevoked(): boolean {
+		return this._revoked;
+	}
+	revoke(): void {
+		this._revoked = true;
+	}
+}
+
+/** Wraps a memento so writes are dropped once the gate is revoked (reads still pass through). */
+class RevocableMemento implements vscode.Memento {
+	constructor(
+		private readonly _actual: vscode.Memento,
+		private readonly _gate: RevocationGate,
+	) { }
+
+	keys(): readonly string[] {
+		return this._actual.keys();
+	}
+
+	get<T>(key: string): T | undefined;
+	get<T>(key: string, defaultValue: T): T;
+	get<T>(key: string, defaultValue?: T): T | undefined {
+		return defaultValue === undefined ? this._actual.get<T>(key) : this._actual.get<T>(key, defaultValue);
+	}
+
+	// eslint-disable-next-line @typescript-eslint/no-explicit-any
+	update(key: string, value: any): Thenable<void> {
+		if (this._gate.isRevoked) {
+			return Promise.resolve();
+		}
+		return this._actual.update(key, value);
+	}
+}
+
+export type TASClientDelegateFn = (memento: vscode.Memento, userInfoStore: UserInfoStore, gate: RevocationGate) => ITASExperimentationService;
 
 export class BaseExperimentationService extends Disposable implements IExperimentationService {
 
@@ -169,8 +211,16 @@ export class BaseExperimentationService extends Disposable implements IExperimen
 
 	private _createDelegate(): ITASExperimentationService {
 		const generation = ++this._delegateGeneration;
-		const delegate = this._delegateFn(this._globalState, this._userInfoStore);
-		this._delegateDisposable.value = toDisposable(() => delegate.dispose());
+		const gate = new RevocationGate();
+		const memento = new RevocableMemento(this._globalState, gate);
+		const delegate = this._delegateFn(memento, this._userInfoStore, gate);
+		// Revoke this generation's storage/telemetry writes and stop its polling when it is
+		// superseded (a newer delegate is assigned) or the service is disposed, so a still
+		// in-flight fetch cannot overwrite the shared memento / assignment context afterwards.
+		this._delegateDisposable.value = toDisposable(() => {
+			gate.revoke();
+			delegate.dispose();
+		});
 		delegate.initialFetch.then(() => {
 			if (generation !== this._delegateGeneration || this._store.isDisposed) {
 				return; // superseded by a newer delegate, or the service was disposed

extensions/copilot/src/platform/telemetry/test/node/experimentation.spec.ts

@@ -4,16 +4,21 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { beforeAll, beforeEach, describe, expect, it } from 'vitest';
+import type * as vscode from 'vscode';
 import { IExperimentationService as ITASExperimentationService } from 'vscode-tas-client';
+import { mock } from '../../../../util/common/test/simpleMock';
+import { Event } from '../../../../util/vs/base/common/event';
 import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
 import { CopilotToken, createTestExtendedTokenInfo } from '../../../authentication/common/copilotToken';
 import { ICopilotTokenStore } from '../../../authentication/common/copilotTokenStore';
 import { IConfigurationService } from '../../../configuration/common/configurationService';
 import { IVSCodeExtensionContext } from '../../../extContext/common/extensionContext';
 import { ILogService } from '../../../log/common/logService';
+import { FetchOptions, HeadersImpl, IFetcherService, Response } from '../../../networking/common/fetcherService';
 import { createPlatformServices, ITestingServicesAccessor } from '../../../test/node/services';
 import { TreatmentsChangeEvent } from '../../common/nullExperimentationService';
-import { BaseExperimentationService, TASClientDelegateFn, UserInfoStore } from '../../node/baseExperimentationService';
+import { createTasFetch } from '../../vscode-node/tasFetch';
+import { BaseExperimentationService, RevocationGate, TASClientDelegateFn, UserInfoStore } from '../../node/baseExperimentationService';
 
 
 function toExpectedTreatment(name: string, org: string | undefined, sku: string | undefined): string | undefined {
@@ -45,6 +50,31 @@ class TestExperimentationService extends BaseExperimentationService {
 	}
 }
 
+/** Captures the memento + revocation gate handed to each delegate generation. */
+class RevocationTestExperimentationService extends BaseExperimentationService {
+	public readonly captured: { memento: vscode.Memento; gate: RevocationGate }[];
+
+	constructor(
+		@IVSCodeExtensionContext extensionContext: IVSCodeExtensionContext,
+		@ICopilotTokenStore tokenStore: ICopilotTokenStore,
+		@IConfigurationService configurationService: IConfigurationService,
+		@ILogService logService: ILogService
+	) {
+		const captured: { memento: vscode.Memento; gate: RevocationGate }[] = [];
+		const delegateFn: TASClientDelegateFn = (memento, userInfoStore, gate) => {
+			captured.push({ memento, gate });
+			return new MockTASExperimentationService(userInfoStore);
+		};
+
+		super(delegateFn, extensionContext, tokenStore, configurationService, logService);
+		this.captured = captured;
+	}
+
+	recreate(): void {
+		this.recreateDelegate();
+	}
+}
+
 class MockTASExperimentationService implements ITASExperimentationService {
 	private _initializePromise: Promise<void> | undefined;
 	private _initialFetch: Promise<void> | undefined;
@@ -169,6 +199,27 @@ describe('ExP Service Tests', () => {
 		});
 	};
 
+	it('revokes a superseded delegate so its late writes are dropped', async () => {
+		const svc = accessor.get(IInstantiationService).createInstance(RevocationTestExperimentationService);
+		const globalState = accessor.get(IVSCodeExtensionContext).globalState;
+
+		expect(svc.captured.length).toBe(1);
+		svc.recreate();
+		expect(svc.captured.length).toBe(2);
+
+		const [gen1, gen2] = svc.captured;
+		expect(gen1.gate.isRevoked).toBe(true);
+		expect(gen2.gate.isRevoked).toBe(false);
+
+		// A superseded (revoked) generation's writes are dropped; the current one's land.
+		await gen1.memento.update('exp.revoke.test', 'stale');
+		expect(globalState.get('exp.revoke.test')).toBeUndefined();
+		await gen2.memento.update('exp.revoke.test', 'fresh');
+		expect(globalState.get('exp.revoke.test')).toBe('fresh');
+
+		svc.dispose();
+	});
+
 	it('should return treatments based on copilot token', async () => {
 		await expService.hasTreatments();
 		let expectedTreatment = toExpectedTreatment('a', undefined, undefined);
@@ -774,3 +825,39 @@ describe('ExP Service delegate recreation', () => {
 		service.dispose();
 	});
 });
+
+/**
+ * Records every request routed through the fetcher service so a test can assert that both TAS
+ * endpoints go through it (proxy-aware transport) with the expected method and call site.
+ */
+class RecordingFetcherService extends mock<IFetcherService>() {
+	public readonly calls: { url: string; method: string; callSite: string; body?: string }[] = [];
+
+	override readonly onDidFetch = Event.None;
+	override readonly onDidCompleteFetch = Event.None;
+
+	override getUserAgentLibrary(): string {
+		return 'test-fetcher';
+	}
+
+	override fetch(url: string, options: FetchOptions): Promise<Response> {
+		this.calls.push({ url, method: options.method ?? 'GET', callSite: options.callSite, body: options.body });
+		return Promise.resolve(Response.fromText(200, 'OK', new HeadersImpl({}), '{}', 'test-stub'));
+	}
+}
+
+describe('TAS proxy transport adapter', () => {
+
+	it('routes both the legacy GET and the assignments POST through the fetcher service with the expected call sites', async () => {
+		const fetcher = new RecordingFetcherService();
+		const tasFetch = createTasFetch(fetcher);
+
+		await tasFetch('https://default.exp-tas.com/vscode/ab', { method: 'GET', headers: { 'X-Legacy': '1' } });
+		await tasFetch('https://exp.example.test/vscode/api/v1/assignments', { method: 'POST', headers: { 'X-New': '1' }, body: '{"parameters":{}}' });
+
+		expect(fetcher.calls).toEqual([
+			{ url: 'https://default.exp-tas.com/vscode/ab', method: 'GET', callSite: 'exp.legacy', body: undefined },
+			{ url: 'https://exp.example.test/vscode/api/v1/assignments', method: 'POST', callSite: 'exp.assignments', body: '{"parameters":{}}' },
+		]);
+	});
+});

extensions/copilot/src/platform/telemetry/vscode-node/microsoftExperimentationService.ts

@@ -16,8 +16,9 @@ import { IVSCodeExtensionContext } from '../../extContext/common/extensionContex
 import { ILogService } from '../../log/common/logService';
 import { IFetcherService } from '../../networking/common/fetcherService';
 import { FetcherService } from '../../networking/vscode-node/fetcherServiceImpl';
-import { ITelemetryService } from '../common/telemetry';
-import { BaseExperimentationService, UserInfoStore } from '../node/baseExperimentationService';
+import { IExperimentationTelemetry, ITelemetryService } from '../common/telemetry';
+import { BaseExperimentationService, RevocationGate, UserInfoStore } from '../node/baseExperimentationService';
+import { createTasFetch } from './tasFetch';
 
 function getTargetPopulation(isPreRelease: boolean): TargetPopulation {
 	if (isPreRelease) {
@@ -277,24 +278,20 @@ export class MicrosoftExperimentationService extends BaseExperimentationService
 		const version = context.extension.packageJSON['version'];
 		const targetPopulation = getTargetPopulation(envService.isPreRelease());
 		let self: MicrosoftExperimentationService | undefined = undefined;
-		const delegateFn = (globalState: vscode.Memento, userInfoStore: UserInfoStore) => {
-			const wrappedMemento = new ExpMementoWrapper(globalState, envService);
+		const delegateFn = (memento: vscode.Memento, userInfoStore: UserInfoStore, gate: RevocationGate) => {
+			const wrappedMemento = new ExpMementoWrapper(memento, envService);
 			const exp = copilotTokenStore.copilotToken?.endpoints?.exp;
 			const assignmentsEndpoint = exp ? `${exp.replace(/\/+$/, '')}/api/v1/assignments` : undefined;
 			// Route both the legacy (GET) and assignments (POST) requests through the extension's
 			// fetcher service so they get proxy handling, retries/fallback, and the standard user-agent.
-			const tasFetch = (url: string, init: { method: 'GET' | 'POST'; headers: Record<string, string>; body?: string }) =>
-				fetcherService.fetch(url, {
-					method: init.method,
-					headers: init.headers,
-					body: init.body,
-					callSite: init.method === 'POST' ? 'exp.assignments' : 'exp.legacy',
-				});
+			const tasFetch = createTasFetch(fetcherService);
 			return getExperimentationServiceFromConfig({
 				extensionName: id,
 				extensionVersion: version,
 				targetPopulation,
-				telemetry: telemetryService,
+				// Wrapped per generation so a superseded delegate's in-flight fetch cannot write
+				// telemetry (e.g. overwrite `abexp.assignmentcontext`) after being replaced.
+				telemetry: new RevocableExpTelemetry(telemetryService, gate),
 				memento: wrappedMemento,
 				filterProviders: [
 					new GithubAccountFilterProvider(userInfoStore, logService),
@@ -334,6 +331,32 @@ export class MicrosoftExperimentationService extends BaseExperimentationService
 	}
 }
 
+/**
+ * Wraps the telemetry service so a superseded delegate's writes are dropped once its generation
+ * gate is revoked, preventing a stale in-flight fetch from overwriting shared telemetry
+ * properties (notably `abexp.assignmentcontext`) after a newer delegate has replaced it.
+ */
+class RevocableExpTelemetry implements IExperimentationTelemetry {
+	constructor(
+		private readonly _actual: IExperimentationTelemetry,
+		private readonly _gate: RevocationGate,
+	) { }
+
+	setSharedProperty(name: string, value: string): void {
+		if (this._gate.isRevoked) {
+			return;
+		}
+		this._actual.setSharedProperty(name, value);
+	}
+
+	postEvent(eventName: string, props: Map<string, string>): void {
+		if (this._gate.isRevoked) {
+			return;
+		}
+		this._actual.postEvent(eventName, props);
+	}
+}
+
 class ExpMementoWrapper implements vscode.Memento {
 
 	constructor(

extensions/copilot/src/platform/telemetry/vscode-node/tasFetch.ts

@@ -0,0 +1,23 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IFetcherService } from '../../networking/common/fetcherService';
+
+/**
+ * Builds the transport that TAS uses for both the legacy (GET) and the assignments (POST)
+ * endpoints, routing every request through the extension's fetcher service so they get proxy
+ * handling, retries/fallback, and the standard user-agent. Wiring both endpoints to this single
+ * adapter is what keeps the assignments call from silently bypassing proxy handling; the method
+ * determines the call site used for fetch telemetry.
+ */
+export function createTasFetch(fetcherService: IFetcherService) {
+	return (url: string, init: { method: 'GET' | 'POST'; headers: Record<string, string>; body?: string }) =>
+		fetcherService.fetch(url, {
+			method: init.method,
+			headers: init.headers,
+			body: init.body,
+			callSite: init.method === 'POST' ? 'exp.assignments' : 'exp.legacy',
+		});
+}