Adds authentication migration

Adds better log out support

package.json

@@ -3625,7 +3625,7 @@
 			},
 			{
 				"command": "gitlens.plus.logout",
-				"title": "Disconnect from GitLens+",
+				"title": "Sign out of GitLens+",
 				"category": "GitLens+"
 			},
 			{

src/container.ts

@@ -160,7 +160,9 @@ export class Container {
 
 		const server = new ServerConnection(this);
 		context.subscriptions.push(server);
-		context.subscriptions.push(new SubscriptionAuthenticationProvider(this, server));
+		context.subscriptions.push(
+			(this._subscriptionAuthentication = new SubscriptionAuthenticationProvider(this, server)),
+		);
 		context.subscriptions.push((this._subscription = new SubscriptionService(this)));
 
 		context.subscriptions.push((this._git = new GitProviderService(this)));
@@ -450,15 +452,16 @@ export class Container {
 		return this._searchAndCompareView;
 	}
 
-	private _subscription: SubscriptionService | undefined;
+	private _subscription: SubscriptionService;
 	get subscription() {
-		if (this._subscription == null) {
-			this._subscription = new SubscriptionService(this);
-		}
-
 		return this._subscription;
 	}
 
+	private _subscriptionAuthentication: SubscriptionAuthenticationProvider;
+	get subscriptionAuthentication() {
+		return this._subscriptionAuthentication;
+	}
+
 	private _settingsWebview: SettingsWebview;
 	get settingsWebview() {
 		return this._settingsWebview;

src/plus/subscription/authenticationProvider.ts

@@ -6,10 +6,12 @@ import {
 	AuthenticationSession,
 	Disposable,
 	EventEmitter,
+	extensions,
 	window,
 } from 'vscode';
 import type { Container } from '../../container';
 import { Logger } from '../../logger';
+import { StorageKeys } from '../../storage';
 import { debug } from '../../system/decorators/log';
 import { ServerConnection } from './serverConnection';
 
@@ -26,7 +28,6 @@ interface StoredSession {
 
 const authenticationId = 'gitlens+';
 const authenticationLabel = 'GitLens+';
-const authenticationSecretKey = `gitlens.plus.auth`;
 
 export class SubscriptionAuthenticationProvider implements AuthenticationProvider, Disposable {
 	private _onDidChangeSessions = new EventEmitter<AuthenticationProviderAuthenticationSessionsChangeEvent>();
@@ -53,6 +54,10 @@ export class SubscriptionAuthenticationProvider implements AuthenticationProvide
 		this._disposable.dispose();
 	}
 
+	private get secretStorageKey(): string {
+		return `gitlens.plus.auth:${this.container.env}`;
+	}
+
 	@debug()
 	public async createSession(scopes: string[]): Promise<AuthenticationSession> {
 		const cc = Logger.getCorrelationContext();
@@ -124,11 +129,55 @@ export class SubscriptionAuthenticationProvider implements AuthenticationProvide
 			this._onDidChangeSessions.fire({ added: [], removed: [session], changed: [] });
 		} catch (ex) {
 			Logger.error(ex, cc);
-			void window.showErrorMessage(`Unable to sign out: ${ex}`);
+			void window.showErrorMessage(`Unable to sign out of GitLens+: ${ex}`);
 			throw ex;
 		}
 	}
 
+	private _migrated: boolean | undefined;
+	async tryMigrateSession(): Promise<AuthenticationSession | undefined> {
+		if (this._migrated == null) {
+			this._migrated = this.container.storage.get<boolean>(StorageKeys.MigratedAuthentication, false);
+		}
+		if (this._migrated) return undefined;
+
+		let session: AuthenticationSession | undefined;
+		try {
+			if (extensions.getExtension('gitkraken.gitkraken-authentication') == null) return;
+
+			session = await authentication.getSession('gitkraken', ['gitlens'], {
+				createIfNone: false,
+			});
+			if (session == null) return;
+
+			session = {
+				id: uuid(),
+				accessToken: session.accessToken,
+				account: { ...session.account },
+				scopes: session.scopes,
+			};
+
+			const sessions = await this._sessionsPromise;
+			const scopesKey = getScopesKey(session.scopes);
+			const sessionIndex = sessions.findIndex(s => s.id === session!.id || getScopesKey(s.scopes) === scopesKey);
+			if (sessionIndex > -1) {
+				sessions.splice(sessionIndex, 1, session);
+			} else {
+				sessions.push(session);
+			}
+
+			await this.storeSessions(sessions);
+
+			this._onDidChangeSessions.fire({ added: [session], removed: [], changed: [] });
+		} catch (ex) {
+			Logger.error(ex, 'Unable to migrate authentication');
+		} finally {
+			this._migrated = true;
+			void this.container.storage.store<boolean>(StorageKeys.MigratedAuthentication, true);
+		}
+		return session;
+	}
+
 	private async checkForUpdates() {
 		const previousSessions = await this._sessionsPromise;
 		this._sessionsPromise = this.getSessionsFromStorage();
@@ -171,14 +220,14 @@ export class SubscriptionAuthenticationProvider implements AuthenticationProvide
 		let storedSessions: StoredSession[];
 
 		try {
-			const sessionsJSON = await this.container.storage.getSecret(authenticationSecretKey);
+			const sessionsJSON = await this.container.storage.getSecret(this.secretStorageKey);
 			if (!sessionsJSON || sessionsJSON === '[]') return [];
 
 			try {
 				storedSessions = JSON.parse(sessionsJSON);
 			} catch (ex) {
 				try {
-					await this.container.storage.deleteSecret(authenticationSecretKey);
+					await this.container.storage.deleteSecret(this.secretStorageKey);
 				} catch {}
 
 				throw ex;
@@ -233,7 +282,7 @@ export class SubscriptionAuthenticationProvider implements AuthenticationProvide
 	private async storeSessions(sessions: AuthenticationSession[]): Promise<void> {
 		try {
 			this._sessionsPromise = Promise.resolve(sessions);
-			await this.container.storage.storeSecret(authenticationSecretKey, JSON.stringify(sessions));
+			await this.container.storage.storeSecret(this.secretStorageKey, JSON.stringify(sessions));
 		} catch (ex) {
 			Logger.error(ex, `Unable to store ${sessions.length} sessions`);
 		}

src/plus/subscription/serverConnection.ts

@@ -151,12 +151,9 @@ export class ServerConnection implements Disposable {
 
 	private updateStatusBarItem(signingIn?: boolean) {
 		if (signingIn && this._statusBarItem == null) {
-			this._statusBarItem = window.createStatusBarItem(
-				'gitkraken-authentication.signIn',
-				StatusBarAlignment.Left,
-			);
-			this._statusBarItem.name = 'GitKraken Sign-in';
-			this._statusBarItem.text = 'Signing into gitkraken.com...';
+			this._statusBarItem = window.createStatusBarItem('gitlens.plus.signIn', StatusBarAlignment.Left);
+			this._statusBarItem.name = 'GitLens+ Sign in';
+			this._statusBarItem.text = 'Signing into GitLens+...';
 			this._statusBarItem.show();
 		}
 

src/plus/subscription/subscriptionService.ts

@@ -1,7 +1,7 @@
 import {
 	authentication,
+	AuthenticationProviderAuthenticationSessionsChangeEvent,
 	AuthenticationSession,
-	AuthenticationSessionsChangeEvent,
 	version as codeVersion,
 	commands,
 	Disposable,
@@ -25,10 +25,11 @@ import { setContext } from '../../context';
 import { AccountValidationError } from '../../errors';
 import { RepositoriesChangeEvent } from '../../git/gitProviderService';
 import { Logger } from '../../logger';
-import { StorageKeys, WorkspaceStorageKeys } from '../../storage';
+import { StorageKeys } from '../../storage';
 import {
 	computeSubscriptionState,
 	getSubscriptionPlan,
+	getSubscriptionPlanName,
 	getSubscriptionPlanPriority,
 	getSubscriptionTimeRemaining,
 	getTimeRemaining,
@@ -74,7 +75,10 @@ export class SubscriptionService implements Disposable {
 	constructor(private readonly container: Container) {
 		this._disposable = Disposable.from(
 			once(container.onReady)(this.onReady, this),
-			authentication.onDidChangeSessions(this.onAuthenticationChanged, this),
+			this.container.subscriptionAuthentication.onDidChangeSessions(
+				e => setTimeout(() => this.onAuthenticationChanged(e), 0),
+				this,
+			),
 		);
 
 		this.changeSubscription(this.getStoredSubscription(), true);
@@ -87,10 +91,28 @@ export class SubscriptionService implements Disposable {
 		this._disposable.dispose();
 	}
 
-	private onAuthenticationChanged(e: AuthenticationSessionsChangeEvent): void {
-		if (e.provider.id !== SubscriptionService.authenticationProviderId) return;
+	private async onAuthenticationChanged(e: AuthenticationProviderAuthenticationSessionsChangeEvent) {
+		let session = this._session;
+		if (session == null && this._sessionPromise != null) {
+			session = await this._sessionPromise;
+		}
+
+		if (session != null && e.removed?.some(s => s.id === session!.id)) {
+			this._session = undefined;
+			this._sessionPromise = undefined;
+			void this.logout();
+			return;
+		}
+
+		const updated = e.added?.[0] ?? e.changed?.[0];
+		if (updated == null) return;
 
-		void this.ensureSession(false, true);
+		if (updated.id === session?.id && updated.accessToken === session?.accessToken) {
+			return;
+		}
+
+		this._session = session;
+		void this.validate();
 	}
 
 	@memoize()
@@ -135,10 +157,6 @@ export class SubscriptionService implements Disposable {
 		return Uri.parse('https://gitkraken.com');
 	}
 
-	private get connectedKey(): `${WorkspaceStorageKeys.ConnectedPrefix}${string}` {
-		return `${WorkspaceStorageKeys.ConnectedPrefix}gitkraken`;
-	}
-
 	private _etag: number = 0;
 	get etag(): number {
 		return this._etag;
@@ -202,8 +220,6 @@ export class SubscriptionService implements Disposable {
 
 		void this.showHomeView();
 
-		await this.container.storage.deleteWorkspace(this.connectedKey);
-
 		const session = await this.ensureSession(true);
 		const loggedIn = Boolean(session);
 		if (loggedIn) {
@@ -254,8 +270,10 @@ export class SubscriptionService implements Disposable {
 	@log()
 	logout(reset: boolean = false): void {
 		this._sessionPromise = undefined;
-		this._session = undefined;
-		void this.container.storage.storeWorkspace(this.connectedKey, false);
+		if (this._session != null) {
+			void this.container.subscriptionAuthentication.removeSession(this._session.id);
+			this._session = undefined;
+		}
 
 		if (reset && this.container.debugging) {
 			this.changeSubscription(undefined);
@@ -575,26 +593,29 @@ export class SubscriptionService implements Disposable {
 	@debug()
 	private async ensureSession(createIfNeeded: boolean, force?: boolean): Promise<AuthenticationSession | undefined> {
 		if (this._sessionPromise != null && this._session === undefined) {
-			this._session = await this._sessionPromise;
-			this._sessionPromise = undefined;
+			void (await this._sessionPromise);
 		}
 
 		if (!force && this._session != null) return this._session;
 		if (this._session === null && !createIfNeeded) return undefined;
 
-		if (createIfNeeded) {
-			await this.container.storage.deleteWorkspace(this.connectedKey);
-		} else if (this.container.storage.getWorkspace<boolean>(this.connectedKey) === false) {
-			return undefined;
-		}
-
 		if (this._sessionPromise === undefined) {
-			this._sessionPromise = this.getOrCreateSession(createIfNeeded);
+			this._sessionPromise = this.getOrCreateSession(createIfNeeded).then(
+				s => {
+					this._session = s;
+					this._sessionPromise = undefined;
+					return this._session;
+				},
+				() => {
+					this._session = null;
+					this._sessionPromise = undefined;
+					return this._session;
+				},
+			);
 		}
 
-		this._session = await this._sessionPromise;
-		this._sessionPromise = undefined;
-		return this._session ?? undefined;
+		const session = await this._sessionPromise;
+		return session ?? undefined;
 	}
 
 	@debug()
@@ -623,6 +644,11 @@ export class SubscriptionService implements Disposable {
 			Logger.error(ex, cc);
 		}
 
+		// If we didn't find a session, check if we could migrate one from the GK auth provider
+		if (session === undefined) {
+			session = await this.container.subscriptionAuthentication.tryMigrateSession();
+		}
+
 		if (session == null) {
 			this.logout();
 			return session ?? null;
@@ -702,7 +728,19 @@ export class SubscriptionService implements Disposable {
 
 	private getStoredSubscription(): Subscription | undefined {
 		const storedSubscription = this.container.storage.get<Stored<Subscription>>(StorageKeys.Subscription);
-		return storedSubscription?.data;
+
+		const subscription = storedSubscription?.data;
+		if (subscription != null) {
+			// Migrate the plan names to the latest names
+			(subscription.plan.actual as Mutable<Subscription['plan']['actual']>).name = getSubscriptionPlanName(
+				subscription.plan.actual.id,
+			);
+			(subscription.plan.effective as Mutable<Subscription['plan']['effective']>).name = getSubscriptionPlanName(
+				subscription.plan.effective.id,
+			);
+		}
+
+		return subscription;
 	}
 
 	private async storeSubscription(subscription: Subscription): Promise<void> {
@@ -756,13 +794,13 @@ export class SubscriptionService implements Disposable {
 
 		if (this._statusBarSubscription == null) {
 			this._statusBarSubscription = window.createStatusBarItem(
-				'gitlens.subscription',
+				'gitlens.plus.subscription',
 				StatusBarAlignment.Left,
 				1,
 			);
 		}
 
-		this._statusBarSubscription.name = 'GitLens Subscription';
+		this._statusBarSubscription.name = 'GitLens+ Subscription';
 		this._statusBarSubscription.command = Commands.ShowHomeView;
 
 		if (account?.verified === false) {

src/storage.ts

@@ -74,6 +74,7 @@ export const enum StorageKeys {
 	PendingWhatsNewOnFocus = 'gitlens:pendingWhatsNewOnFocus',
 	Version = 'gitlens:version',
 
+	MigratedAuthentication = 'gitlens:plus:migratedAuthentication',
 	Subscription = 'gitlens:premium:subscription', // Don't change this key name as its the stored subscription
 
 	Deprecated_Version = 'gitlensVersion',