New TAS assignments endpoint updates (lifecycle, readiness, disposal) (#329736)

* New TAS assignments endpoint updates (lifecycle, readiness, disposal)

* feedback updates

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

@@ -132,6 +132,7 @@ export class BaseExperimentationService extends Disposable implements IExperimen
 
 	private readonly _delegateFn: TASClientDelegateFn;
 	private readonly _globalState: vscode.Memento;
+	private _delegateGeneration = 0;
 
 	protected _onDidTreatmentsChange = this._register(new Emitter<TreatmentsChangeEvent>());
 	readonly onDidTreatmentsChange = this._onDidTreatmentsChange.event;
@@ -167,10 +168,17 @@ 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 as unknown as { dispose?(): void }).dispose?.());
+		this._delegateDisposable.value = toDisposable(() => delegate.dispose());
 		delegate.initialFetch.then(() => {
+			if (generation !== this._delegateGeneration || this._store.isDisposed) {
+				return; // superseded by a newer delegate, or the service was disposed
+			}
 			this._logService.trace(`[BaseExperimentationService] Initial fetch completed`);
+			// A replacement delegate (e.g. once the assignments endpoint arrives) may carry
+			// different assignments; announce any changes so experiment-based config updates.
+			this._signalTreatmentsChangeEvent();
 		});
 		return delegate;
 	}

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

@@ -639,3 +639,138 @@ describe('ExP Service Tests', () => {
 		await extensionContext.globalState.update(UserInfoStore.IS_SN_STORAGE_KEY, undefined);
 	});
 });
+
+/**
+ * A delegate whose initial fetch is completed manually, so tests can control ordering, and
+ * that records when it is disposed. Used to exercise endpoint-driven delegate recreation.
+ */
+class ControllableMockTAS implements ITASExperimentationService {
+	public disposed = false;
+	private _resolveInitialFetch!: () => void;
+	private readonly _treatments = new Map<string, boolean | number | string>();
+
+	readonly initializePromise = Promise.resolve();
+	readonly initialFetch: Promise<void>;
+
+	constructor() {
+		this.initialFetch = new Promise<void>(resolve => { this._resolveInitialFetch = resolve; });
+	}
+
+	completeInitialFetch(): void {
+		this._resolveInitialFetch();
+	}
+
+	setTreatment(name: string, value: boolean | number | string): void {
+		this._treatments.set(name, value);
+	}
+
+	isFlightEnabled(): boolean { throw new Error('Method not implemented.'); }
+	isCachedFlightEnabled(): Promise<boolean> { throw new Error('Method not implemented.'); }
+	isFlightEnabledAsync(): Promise<boolean> { throw new Error('Method not implemented.'); }
+	getTreatmentVariable<T extends boolean | number | string>(_configId: string, name: string): T | undefined {
+		return this._treatments.get(name) as T | undefined;
+	}
+	getTreatmentVariableAsync<T extends boolean | number | string>(configId: string, name: string): Promise<T | undefined> {
+		return Promise.resolve(this.getTreatmentVariable<T>(configId, name));
+	}
+	dispose(): void { this.disposed = true; }
+}
+
+class RecreatableExperimentationService extends BaseExperimentationService {
+	public readonly delegates: ControllableMockTAS[];
+
+	constructor(
+		@IVSCodeExtensionContext extensionContext: IVSCodeExtensionContext,
+		@ICopilotTokenStore tokenStore: ICopilotTokenStore,
+		@IConfigurationService configurationService: IConfigurationService,
+		@ILogService logService: ILogService
+	) {
+		const delegates: ControllableMockTAS[] = [];
+		const delegateFn: TASClientDelegateFn = () => {
+			const delegate = new ControllableMockTAS();
+			delegates.push(delegate);
+			return delegate;
+		};
+
+		super(delegateFn, extensionContext, tokenStore, configurationService, logService);
+		this.delegates = delegates;
+	}
+
+	triggerRecreate(): void {
+		this.recreateDelegate();
+	}
+}
+
+describe('ExP Service delegate recreation', () => {
+	let accessor: ITestingServicesAccessor;
+
+	beforeAll(() => {
+		const testingServiceCollection = createPlatformServices();
+		accessor = testingServiceCollection.createTestingAccessor();
+	});
+
+	const create = () => accessor.get(IInstantiationService).createInstance(RecreatableExperimentationService);
+
+	it('disposes the previous delegate and switches to the new one on recreation', () => {
+		const service = create();
+		expect(service.delegates.length).toBe(1);
+		const first = service.delegates[0];
+
+		service.triggerRecreate();
+
+		expect(service.delegates.length).toBe(2);
+		const second = service.delegates[1];
+		expect(first.disposed).toBe(true);
+		expect(second.disposed).toBe(false);
+
+		// The final delegate is disposed together with the service.
+		service.dispose();
+		expect(second.disposed).toBe(true);
+	});
+
+	it('ignores a superseded delegate\'s initial fetch completion', async () => {
+		const service = create();
+		const first = service.delegates[0];
+
+		// Seed a previously-read treatment so a change would be detectable if (wrongly) signaled.
+		first.setTreatment('x', 'v0');
+		expect(service.getTreatmentVariable<string>('x')).toBe('v0');
+
+		service.triggerRecreate();
+		service.delegates[1].setTreatment('x', 'v1');
+
+		let fired = false;
+		service.onDidTreatmentsChange(() => { fired = true; });
+
+		// The superseded (first) delegate completing its fetch must not signal changes.
+		first.completeInitialFetch();
+		await new Promise(resolve => setTimeout(resolve, 0));
+		expect(fired).toBe(false);
+
+		service.dispose();
+	});
+
+	it('signals treatment changes after the replacement delegate completes its initial fetch', async () => {
+		const service = create();
+		const first = service.delegates[0];
+
+		first.setTreatment('x', 'v0');
+		expect(service.getTreatmentVariable<string>('x')).toBe('v0');
+
+		service.triggerRecreate();
+		const second = service.delegates[1];
+		second.setTreatment('x', 'v1');
+
+		const changePromise = new Promise<TreatmentsChangeEvent>(resolve => {
+			service.onDidTreatmentsChange(resolve);
+		});
+
+		second.completeInitialFetch();
+		const event = await changePromise;
+
+		expect(event.affectedTreatmentVariables).toContain('x');
+		expect(service.getTreatmentVariable<string>('x')).toBe('v1');
+
+		service.dispose();
+	});
+});

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

@@ -33,7 +33,7 @@ function trimVersionSuffix(version: string): string {
 
 /**
  * Formats an ISO date into the `yyyymmddHH` form the experimentation backend expects
- * (10 digits, fits within int32). Returns an empty string when unavailable.
+ * (10 digits: yyyymmddHH). Returns an empty string when unavailable.
  */
 function formatReleaseDate(iso: string): string {
 	if (!iso) {

src/vs/platform/assignment/common/assignment.ts

@@ -174,7 +174,7 @@ function trimVersionSuffix(version: string): string {
 
 /**
  * Formats an ISO release date into the `yyyymmddHH` form the experimentation backend
- * expects (10 digits, fits within int32). Returns an empty string when unavailable.
+ * expects (10 digits: yyyymmddHH). Returns an empty string when unavailable.
  */
 function formatReleaseDate(iso: string): string {
 	if (!iso) {

src/vs/workbench/services/assignment/common/assignmentService.ts

@@ -152,12 +152,14 @@ export class WorkbenchAssignmentService extends Disposable implements IAssignmen
 	declare readonly _serviceBrand: undefined;
 
 	private tasClient: Promise<TASClient> | undefined;
-	private readonly tasSetupDisposables = new DisposableStore();
+	private readonly tasSetupDisposables = this._register(new DisposableStore());
 
 	private assignmentsEndpoint: string | undefined;
 
 	private networkInitialized = false;
 	private setupGeneration = 0;
+	/** Revokes the current setup's storage/telemetry/fetch wrappers, neutralizing a superseded in-flight client. */
+	private revokeCurrentSetup: (() => void) | undefined;
 	private readonly overrideInitDelay: Promise<void>;
 
 	private readonly contextFilter: AssignmentContextFilter;
@@ -187,16 +189,17 @@ export class WorkbenchAssignmentService extends Disposable implements IAssignmen
 			this.tasClient = this.setupTASClient();
 
 			// The assignments endpoint is sourced from account entitlements, which load
-			// asynchronously. Re-setup the client when it first appears or changes.
-			this._register(this.defaultAccountService.onDidChangeDefaultAccount(() => {
-				const next = this.getAssignmentsEndpoint();
-				if (next !== this.assignmentsEndpoint) {
-					this.tasClient = this.setupTASClient();
-				}
-			}));
+			// asynchronously. The initial account load resolves the readiness barrier without
+			// firing onDidChangeDefaultAccount, so proactively re-check once it is ready, and
+			// again whenever the account changes later.
+			this.defaultAccountService.getDefaultAccount().then(() => this.recreateTasClientIfEndpointChanged());
+			this._register(this.defaultAccountService.onDidChangeDefaultAccount(() => this.recreateTasClientIfEndpointChanged()));
 
-			// Ensure the final client's auto-polling is stopped when the service is disposed.
-			this._register(toDisposable(() => WorkbenchAssignmentService.disposeTasClient(this.tasClient)));
+			// Stop the final client's auto-polling and revoke its wrappers when the service is disposed.
+			this._register(toDisposable(() => {
+				this.revokeCurrentSetup?.();
+				WorkbenchAssignmentService.disposeTasClient(this.tasClient);
+			}));
 		}
 
 		this.contextFilter = this._register(new AssignmentContextFilter(storageService));
@@ -286,6 +289,17 @@ export class WorkbenchAssignmentService extends Disposable implements IAssignmen
 		return `${exp.replace(/\/+$/, '')}/api/v1/assignments`;
 	}
 
+	/** Recreates the TAS client when the resolved assignments endpoint has changed. */
+	private recreateTasClientIfEndpointChanged(): void {
+		if (this._store.isDisposed) {
+			return; // the service was disposed before the (async) account load resolved
+		}
+		const next = this.getAssignmentsEndpoint();
+		if (next !== this.assignmentsEndpoint) {
+			this.tasClient = this.setupTASClient();
+		}
+	}
+
 	/**
 	 * Transport for the new assignments endpoint, backed by the main-process request service
 	 * (avoids renderer CORS). Shape matches tas-client's injectable `assignmentsFetch`.
@@ -313,9 +327,44 @@ export class WorkbenchAssignmentService extends Disposable implements IAssignmen
 		const generation = ++this.setupGeneration;
 		this.networkInitialized = false;
 
-		// Dispose the previously created client so it stops auto-polling the (legacy) endpoint.
+		// Revoke the previous setup's wrappers, then dispose its client. Revoking neutralizes a
+		// superseded, still-in-flight client: after replacement it can no longer write the shared
+		// memento, emit telemetry, or hit the assignments endpoint. This is needed because the
+		// tas-client's dispose() only stops its polling timer, not an already-running fetch.
+		this.revokeCurrentSetup?.();
 		WorkbenchAssignmentService.disposeTasClient(this.tasClient);
 
+		let revoked = false;
+		this.revokeCurrentSetup = () => { revoked = true; };
+
+		// Reference the shared memento/telemetry/fetch lazily (at call time): they are assigned in
+		// the constructor body after the initial setupTASClient() call has already started.
+		const service = this;
+
+		const keyValueStorage: IKeyValueStorage = {
+			getValue<T>(key: string, defaultValue?: T): Promise<T | undefined> {
+				return service.keyValueStorage.getValue<T>(key, defaultValue);
+			},
+			setValue<T>(key: string, value: T): void {
+				if (!revoked) {
+					service.keyValueStorage.setValue<T>(key, value);
+				}
+			},
+		};
+
+		const telemetry: IExperimentationTelemetry = {
+			setSharedProperty(name: string, value: string): void {
+				if (!revoked) {
+					service.telemetry.setSharedProperty(name, value);
+				}
+			},
+			postEvent(eventName: string, props: Map<string, string>): void {
+				if (!revoked) {
+					service.telemetry.postEvent(eventName, props);
+				}
+			},
+		};
+
 		const targetPopulation = this.productService.quality === 'stable' ?
 			TargetPopulation.Public : (this.productService.quality === 'exploration' ?
 				TargetPopulation.Exploration : TargetPopulation.Insiders);
@@ -369,9 +418,9 @@ export class WorkbenchAssignmentService extends Disposable implements IAssignmen
 		const fetchStopWatch = StopWatch.create();
 		const tasClient = new tasClientModule.ExperimentationService({
 			filterProviders: [filterProvider, extensionsFilterProvider],
-			telemetry: this.telemetry,
+			telemetry,
 			storageKey: ASSIGNMENT_STORAGE_KEY,
-			keyValueStorage: this.keyValueStorage,
+			keyValueStorage,
 			assignmentContextTelemetryPropertyName: tasConfig.assignmentContextTelemetryPropertyName,
 			telemetryEventName: tasConfig.telemetryEventName,
 			endpoint: tasConfig.endpoint,
@@ -379,7 +428,9 @@ export class WorkbenchAssignmentService extends Disposable implements IAssignmen
 			assignmentsFilterProviders,
 			// Route the assignments request through the main-process request service so it is
 			// not subject to renderer CORS (parity with how core reaches api.github.com).
-			assignmentsFetch: assignmentsEndpoint ? this.assignmentsFetch : undefined,
+			assignmentsFetch: assignmentsEndpoint
+				? (url, init) => (revoked ? Promise.resolve({ status: 0, json: async () => ({}) }) : service.assignmentsFetch(url, init))
+				: undefined,
 			refetchInterval: ASSIGNMENT_REFETCH_INTERVAL,
 		});
 
@@ -449,7 +500,7 @@ export class WorkbenchAssignmentService extends Disposable implements IAssignmen
 
 	/** Stops a TAS client's auto-polling once it resolves. Safe to call with `undefined`. */
 	private static disposeTasClient(client: Promise<TASClient> | undefined): void {
-		client?.then(c => (c as unknown as { dispose?(): void }).dispose?.()).catch(() => undefined);
+		client?.then(c => c.dispose()).catch(() => undefined);
 	}
 }