Add a race between login event and dialog promise when connecting and logged out

Author: EhabY

src/extension.ts

@@ -309,6 +309,15 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		commands.viewLogs.bind(commands),
 	);
 
+	const remote = new Remote(
+		vscodeProposed,
+		output,
+		commands,
+		ctx.extensionMode,
+		pathResolver,
+		cliManager,
+	);
+
 	ctx.subscriptions.push(
 		secretsManager.onDidChangeSessionToken(async () => {
 			const token = await secretsManager.getSessionToken();
@@ -320,6 +329,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 				output.info("Logging in");
 				// Should login the user directly if the URL+Token are valid
 				await commands.login({ url, token });
+				// Resolve any pending login detection promises
+				remote.resolveLoginDetected();
 			}
 		}),
 	);
@@ -334,14 +345,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	// (this would require the user to uninstall the Coder extension and
 	// reinstall after installing the remote SSH extension, which is annoying)
 	if (remoteSSHExtension && vscodeProposed.env.remoteAuthority) {
-		const remote = new Remote(
-			vscodeProposed,
-			output,
-			commands,
-			ctx.extensionMode,
-			pathResolver,
-			cliManager,
-		);
 		try {
 			const details = await remote.setup(
 				vscodeProposed.env.remoteAuthority,

src/remote.ts

@@ -44,6 +44,9 @@ export interface RemoteDetails extends vscode.Disposable {
 }
 
 export class Remote {
+	private loginDetectedResolver: (() => void) | undefined;
+	private loginDetectedPromise: Promise<void> = Promise.resolve();
+
 	public constructor(
 		// We use the proposed API to get access to useCustom in dialogs.
 		private readonly vscodeProposed: typeof vscode,
@@ -54,6 +57,27 @@ export class Remote {
 		private readonly cliManager: CliManager,
 	) {}
 
+	/**
+	 * Creates a new promise that will be resolved when login is detected in another window.
+	 * This should be called when starting a setup operation that might need login.
+	 */
+	private createLoginDetectionPromise(): void {
+		this.loginDetectedPromise = new Promise<void>((resolve) => {
+			this.loginDetectedResolver = resolve;
+		});
+	}
+
+	/**
+	 * Resolves the current login detection promise if one exists.
+	 * This should be called from the extension when login is detected.
+	 */
+	public resolveLoginDetected(): void {
+		if (this.loginDetectedResolver) {
+			this.loginDetectedResolver();
+			this.loginDetectedResolver = undefined;
+		}
+	}
+
 	private async confirmStart(workspaceName: string): Promise<boolean> {
 		const action = await this.vscodeProposed.window.showInformationMessage(
 			`Unable to connect to the workspace ${workspaceName} because it is not running. Start the workspace?`,
@@ -206,6 +230,8 @@ export class Remote {
 		remoteAuthority: string,
 		firstConnect: boolean,
 	): Promise<RemoteDetails | undefined> {
+		this.createLoginDetectionPromise();
+
 		const parts = parseRemoteAuthority(remoteAuthority);
 		if (!parts) {
 			// Not a Coder host.
@@ -218,7 +244,7 @@ export class Remote {
 		await this.migrateSessionToken(parts.label);
 
 		// Get the URL and token belonging to this host.
-		const { url: baseUrlRaw, token } = await this.cliManager.readConfig(
+		let { url: baseUrlRaw, token } = await this.cliManager.readConfig(
 			parts.label,
 		);
 
@@ -227,24 +253,38 @@ export class Remote {
 			!baseUrlRaw ||
 			(!token && needToken(vscode.workspace.getConfiguration()))
 		) {
-			const result = await this.vscodeProposed.window.showInformationMessage(
+			const dialogPromise = this.vscodeProposed.window.showInformationMessage(
 				"You are not logged in...",
 				{
 					useCustom: true,
 					modal: true,
-					detail: `You must log in to access ${workspaceName}.`,
+					detail: `You must log in to access ${workspaceName}. If you've already logged in, you may close this dialog.`,
 				},
 				"Log In",
 			);
-			if (!result) {
-				// User declined to log in.
-				await this.closeRemote();
+
+			// Race between dialog and login detection
+			const result = await Promise.race([
+				this.loginDetectedPromise.then(() => ({ type: "login" as const })),
+				dialogPromise.then((userChoice) => ({
+					type: "dialog" as const,
+					userChoice,
+				})),
+			]);
+
+			if (result.type === "login") {
+				return this.setup(remoteAuthority, firstConnect);
 			} else {
-				// Log in then try again.
-				await this.commands.login({ url: baseUrlRaw, label: parts.label });
-				await this.setup(remoteAuthority, firstConnect);
+				if (!result.userChoice) {
+					// User declined to log in.
+					await this.closeRemote();
+					return;
+				} else {
+					// Log in then try again.
+					await this.commands.login({ url: baseUrlRaw, label: parts.label });
+					return this.setup(remoteAuthority, firstConnect);
+				}
 			}
-			return;
 		}
 
 		this.logger.info("Using deployment URL", baseUrlRaw);
@@ -360,11 +400,11 @@ export class Remote {
 						);
 					if (!result) {
 						await this.closeRemote();
+						return;
 					} else {
 						await this.commands.login({ url: baseUrlRaw, label: parts.label });
-						await this.setup(remoteAuthority, firstConnect);
+						return this.setup(remoteAuthority, firstConnect);
 					}
-					return;
 				}
 				default:
 					throw error;