Merge pull request #1860 from microsoft/roblou/revert-chat-activation (#1878)

Restore activationBlockers on auth

Co-authored-by: Logan Ramos <[email protected]>

Author: roblourens

src/extension/common/contributions.ts

@@ -5,6 +5,7 @@
 
 import { ILogService } from '../../platform/log/common/logService';
 import { Disposable, isDisposable } from '../../util/vs/base/common/lifecycle';
+import { StopWatch } from '../../util/vs/base/common/stopwatch';
 import { IInstantiationService, ServicesAccessor } from '../../util/vs/platform/instantiation/common/instantiation';
 
 export interface IExtensionContribution {
@@ -15,6 +16,12 @@ export interface IExtensionContribution {
 	 * Dispose of the contribution.
 	 */
 	dispose?(): void;
+
+	/**
+	 * A promise that the extension `activate` method will wait on before completing.
+	 * USE this carefully as it will delay startup of our extension.
+	 */
+	activationBlocker?: Promise<any>;
 }
 
 export interface IExtensionContributionFactory {
@@ -31,6 +38,8 @@ export function asContributionFactory(ctor: { new(...args: any[]): any }): IExte
 }
 
 export class ContributionCollection extends Disposable {
+	private readonly allActivationBlockers: Promise<any>[] = [];
+
 	constructor(
 		contribs: IExtensionContributionFactory[],
 		@ILogService logService: ILogService,
@@ -46,9 +55,22 @@ export class ContributionCollection extends Disposable {
 				if (isDisposable(instance)) {
 					this._register(instance);
 				}
+
+				if (instance?.activationBlocker) {
+					const sw = StopWatch.create();
+					const id = instance.id || 'UNKNOWN';
+					this.allActivationBlockers.push(instance.activationBlocker.finally(() => {
+						logService.info(`activationBlocker from '${id}' took for ${Math.round(sw.elapsed())}ms`);
+					}));
+				}
 			} catch (error) {
 				logService.error(error, `Error while loading contribution`);
 			}
 		}
 	}
+
+	async waitForActivationBlockers(): Promise<void> {
+		// WAIT for all activation blockers to complete
+		await Promise.allSettled(this.allActivationBlockers);
+	}
 }

src/extension/conversation/vscode-node/chatParticipants.ts

@@ -10,9 +10,7 @@ import { IInteractionService } from '../../../platform/chat/common/interactionSe
 import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
 import { IEndpointProvider } from '../../../platform/endpoint/common/endpointProvider';
 import { IOctoKitService } from '../../../platform/github/common/githubService';
-import { ILogService } from '../../../platform/log/common/logService';
 import { IExperimentationService } from '../../../platform/telemetry/common/nullExperimentationService';
-import { DeferredPromise } from '../../../util/vs/base/common/async';
 import { Event, Relay } from '../../../util/vs/base/common/event';
 import { DisposableStore, IDisposable } from '../../../util/vs/base/common/lifecycle';
 import { autorun } from '../../../util/vs/base/common/observableInternal';
@@ -30,18 +28,18 @@ import { getAdditionalWelcomeMessage } from './welcomeMessageProvider';
 export class ChatAgentService implements IChatAgentService {
 	declare readonly _serviceBrand: undefined;
 
-	private _lastChatAgents: ChatParticipants | undefined; // will be cleared when disposed
+	private _lastChatAgents: ChatAgents | undefined; // will be cleared when disposed
 
 	constructor(
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 	) { }
 
-	public debugGetCurrentChatAgents(): ChatParticipants | undefined {
+	public debugGetCurrentChatAgents(): ChatAgents | undefined {
 		return this._lastChatAgents;
 	}
 
 	register(): IDisposable {
-		const chatAgents = this.instantiationService.createInstance(ChatParticipants);
+		const chatAgents = this.instantiationService.createInstance(ChatAgents);
 		chatAgents.register();
 		this._lastChatAgents = chatAgents;
 		return {
@@ -53,13 +51,11 @@ export class ChatAgentService implements IChatAgentService {
 	}
 }
 
-class ChatParticipants implements IDisposable {
+class ChatAgents implements IDisposable {
 	private readonly _disposables = new DisposableStore();
 
 	private additionalWelcomeMessage: vscode.MarkdownString | undefined;
 
-	private requestBlocker: Promise<void>;
-
 	constructor(
 		@IOctoKitService private readonly octoKitService: IOctoKitService,
 		@IAuthenticationService private readonly authenticationService: IAuthenticationService,
@@ -71,24 +67,7 @@ class ChatParticipants implements IDisposable {
 		@IChatQuotaService private readonly _chatQuotaService: IChatQuotaService,
 		@IConfigurationService private readonly configurationService: IConfigurationService,
 		@IExperimentationService private readonly experimentationService: IExperimentationService,
-		@ILogService private readonly logService: ILogService,
-	) {
-		// TODO@roblourens This may not be necessary - the point for now is to match the old behavior of blocking chat until auth completes
-		const requestBlockerDeferred = new DeferredPromise<void>();
-		this.requestBlocker = requestBlockerDeferred.p;
-		if (authenticationService.copilotToken) {
-			this.logService.debug(`ChatParticipants: Copilot token already available`);
-			requestBlockerDeferred.complete();
-		} else {
-			this._disposables.add(authenticationService.onDidAuthenticationChange(async () => {
-				const hasSession = !!authenticationService.copilotToken;
-				this.logService.debug(`ChatParticipants: onDidAuthenticationChange has token: ${hasSession}`);
-				if (hasSession) {
-					requestBlockerDeferred.complete();
-				}
-			}));
-		}
-	}
+	) { }
 
 	dispose() {
 		this._disposables.dispose();
@@ -280,7 +259,6 @@ Learn more about [GitHub Copilot](https://docs.github.com/copilot/using-github-c
 
 	private getChatParticipantHandler(id: string, name: string, defaultIntentIdOrGetter: IntentOrGetter, onRequestPaused: Event<vscode.ChatParticipantPauseStateEvent>): vscode.ChatExtendedRequestHandler {
 		return async (request, context, stream, token): Promise<vscode.ChatResult> => {
-			await this.requestBlocker;
 
 			// If we need privacy confirmation, i.e with 3rd party models. We will return a confirmation response and return early
 			const privacyConfirmation = await this.requestPolicyConfirmation(request, stream);

src/extension/conversation/vscode-node/conversationFeature.ts

@@ -18,6 +18,7 @@ import { ILogService } from '../../../platform/log/common/logService';
 import { ISettingsEditorSearchService } from '../../../platform/settingsEditor/common/settingsEditorSearchService';
 import { ITelemetryService } from '../../../platform/telemetry/common/telemetry';
 import { isUri } from '../../../util/common/types';
+import { DeferredPromise } from '../../../util/vs/base/common/async';
 import { CancellationToken } from '../../../util/vs/base/common/cancellation';
 import { DisposableStore, IDisposable, combinedDisposable } from '../../../util/vs/base/common/lifecycle';
 import { URI } from '../../../util/vs/base/common/uri';
@@ -61,6 +62,7 @@ export class ConversationFeature implements IExtensionContribution {
 	private _settingsSearchProviderRegistered = false;
 
 	readonly id = 'conversationFeature';
+	readonly activationBlocker?: Promise<any>;
 
 	constructor(
 		@IInstantiationService private instantiationService: IInstantiationService,
@@ -85,7 +87,25 @@ export class ConversationFeature implements IExtensionContribution {
 		// Register Copilot token listener
 		this.registerCopilotTokenListener();
 
-		this.activated = true;
+		const activationBlockerDeferred = new DeferredPromise<void>();
+		this.activationBlocker = activationBlockerDeferred.p;
+		if (authenticationService.copilotToken) {
+			this.logService.debug(`ConversationFeature: Copilot token already available`);
+			this.activated = true;
+			activationBlockerDeferred.complete();
+		}
+
+		this._disposables.add(authenticationService.onDidAuthenticationChange(async () => {
+			const hasSession = !!authenticationService.copilotToken;
+			this.logService.debug(`ConversationFeature: onDidAuthenticationChange has token: ${hasSession}`);
+			if (hasSession) {
+				this.activated = true;
+			} else {
+				this.activated = false;
+			}
+
+			activationBlockerDeferred.complete();
+		}));
 	}
 
 	get enabled() {

src/extension/conversation/vscode-node/languageModelAccess.ts

@@ -29,7 +29,7 @@ import { isEncryptedThinkingDelta } from '../../../platform/thinking/common/thin
 import { BaseTokensPerCompletion } from '../../../platform/tokenizer/node/tokenizer';
 import { TelemetryCorrelationId } from '../../../util/common/telemetryCorrelationId';
 import { Emitter } from '../../../util/vs/base/common/event';
-import { Disposable } from '../../../util/vs/base/common/lifecycle';
+import { Disposable, MutableDisposable } from '../../../util/vs/base/common/lifecycle';
 import { isDefined, isNumber, isString, isStringArray } from '../../../util/vs/base/common/types';
 import { localize } from '../../../util/vs/nls';
 import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
@@ -44,6 +44,8 @@ export class LanguageModelAccess extends Disposable implements IExtensionContrib
 
 	readonly id = 'languageModelAccess';
 
+	readonly activationBlocker?: Promise<any>;
+
 	private readonly _onDidChange = this._register(new Emitter<void>());
 	private _currentModels: vscode.LanguageModelChatInformation[] = []; // Store current models for reference
 	private _chatEndpoints: IChatEndpoint[] = [];
@@ -71,8 +73,10 @@ export class LanguageModelAccess extends Disposable implements IExtensionContrib
 		}
 
 		// initial
-		this._registerChatProvider();
-		this._registerEmbeddings();
+		this.activationBlocker = Promise.all([
+			this._registerChatProvider(),
+			this._registerEmbeddings(),
+		]);
 	}
 
 	override dispose(): void {
@@ -83,7 +87,7 @@ export class LanguageModelAccess extends Disposable implements IExtensionContrib
 		return this._currentModels;
 	}
 
-	private _registerChatProvider(): void {
+	private async _registerChatProvider(): Promise<void> {
 		const provider: vscode.LanguageModelChatProvider = {
 			onDidChangeLanguageModelChatInformation: this._onDidChange.event,
 			provideLanguageModelChatInformation: this._provideLanguageModelChatInfo.bind(this),
@@ -256,22 +260,34 @@ export class LanguageModelAccess extends Disposable implements IExtensionContrib
 
 	private async _registerEmbeddings(): Promise<void> {
 
-		const embeddingsComputer = this._embeddingsComputer;
-		const embeddingType = EmbeddingType.text3small_512;
-		const model = getWellKnownEmbeddingTypeInfo(embeddingType)?.model;
-		if (!model) {
-			throw new Error(`No model found for embedding type ${embeddingType.id}`);
-		}
+		const dispo = this._register(new MutableDisposable());
+
 
-		const that = this;
-		this._register(vscode.lm.registerEmbeddingsProvider(`copilot.${model}`, new class implements vscode.EmbeddingsProvider {
-			async provideEmbeddings(input: string[], token: vscode.CancellationToken): Promise<vscode.Embedding[]> {
-				await that._getToken();
+		const update = async () => {
 
-				const result = await embeddingsComputer.computeEmbeddings(embeddingType, input, {}, new TelemetryCorrelationId('EmbeddingsProvider::provideEmbeddings'), token);
-				return result.values.map(embedding => ({ values: embedding.value.slice(0) }));
+			if (!await this._getToken()) {
+				dispo.clear();
+				return;
 			}
-		}));
+
+			const embeddingsComputer = this._embeddingsComputer;
+			const embeddingType = EmbeddingType.text3small_512;
+			const model = getWellKnownEmbeddingTypeInfo(embeddingType)?.model;
+			if (!model) {
+				throw new Error(`No model found for embedding type ${embeddingType.id}`);
+			}
+
+			dispo.clear();
+			dispo.value = vscode.lm.registerEmbeddingsProvider(`copilot.${model}`, new class implements vscode.EmbeddingsProvider {
+				async provideEmbeddings(input: string[], token: vscode.CancellationToken): Promise<vscode.Embedding[]> {
+					const result = await embeddingsComputer.computeEmbeddings(embeddingType, input, {}, new TelemetryCorrelationId('EmbeddingsProvider::provideEmbeddings'), token);
+					return result.values.map(embedding => ({ values: embedding.value.slice(0) }));
+				}
+			});
+		};
+
+		this._register(this._authenticationService.onDidAuthenticationChange(() => update()));
+		await update();
 	}
 
 	private async _getToken(): Promise<CopilotToken | undefined> {

src/extension/conversation/vscode-node/test/conversationFeature.test.ts

@@ -70,6 +70,16 @@ suite('Conversation feature test suite', function () {
 		}
 	});
 
+	test("If the 'interactive' version does not match, the feature is not enabled and not activated", function () {
+		const conversationFeature = instaService.createInstance(ConversationFeature);
+		try {
+			assert.deepStrictEqual(conversationFeature.enabled, false);
+			assert.deepStrictEqual(conversationFeature.activated, false);
+		} finally {
+			conversationFeature.dispose();
+		}
+	});
+
 	test('The feature is enabled and activated in test mode', function () {
 		const conversationFeature = instaService.createInstance(ConversationFeature);
 		try {

src/extension/extension/vscode/extension.ts

@@ -9,7 +9,6 @@ import { isScenarioAutomation } from '../../../platform/env/common/envService';
 import { isProduction } from '../../../platform/env/common/packagejson';
 import { IHeatmapService } from '../../../platform/heatmap/common/heatmapService';
 import { IIgnoreService } from '../../../platform/ignore/common/ignoreService';
-import { ILogService } from '../../../platform/log/common/logService';
 import { IExperimentationService } from '../../../platform/telemetry/common/nullExperimentationService';
 import { IInstantiationServiceBuilder, InstantiationServiceBuilder } from '../../../util/common/services';
 import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
@@ -64,16 +63,16 @@ export async function baseActivate(configuration: IExtensionActivationConfigurat
 
 	await instantiationService.invokeFunction(async accessor => {
 		const expService = accessor.get(IExperimentationService);
-		const logService = accessor.get(ILogService);
 
 		// Await intialization of exp service. This ensure cache is fresh.
 		// It will then auto refresh every 30 minutes after that.
 		await expService.hasTreatments();
 
+		// THIS is awaited because some contributions can block activation
+		// via `IExtensionContribution#activationBlocker`
 		const contributions = instantiationService.createInstance(ContributionCollection, configuration.contributions);
 		context.subscriptions.push(contributions);
-
-		logService.trace('Copilot Chat extension activated');
+		await contributions.waitForActivationBlockers();
 	});
 
 	if (ExtensionMode.Test === context.extensionMode && !isScenarioAutomation) {