Improve UX

Author: EhabY

src/commands.ts

@@ -19,18 +19,15 @@ import { type SecretsManager } from "./core/secretsManager";
 import { CertificateError } from "./error";
 import { getGlobalFlags } from "./globalFlags";
 import { type Logger } from "./logging/logger";
-import { OAuthMetadataClient } from "./oauth/metadataClient";
 import { type OAuthSessionManager } from "./oauth/sessionManager";
-import { maybeAskAgent, maybeAskUrl } from "./promptUtils";
+import { maybeAskAgent, maybeAskUrl, maybeAskAuthMethod } from "./promptUtils";
 import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util";
 import {
 	AgentTreeItem,
 	type OpenableTreeItem,
 	WorkspaceTreeItem,
 } from "./workspace/workspacesProvider";
 
-type AuthMethod = "oauth" | "legacy";
-
 export class Commands {
 	private readonly vscodeProposed: typeof vscode;
 	private readonly logger: Logger;
@@ -92,7 +89,7 @@ export class Commands {
 		// Try to get a token from the user, if we need one, and their user.
 		const autoLogin = args?.autoLogin === true;
 
-		const res = await this.maybeAskToken(url, args?.token, autoLogin);
+		const res = await this.attemptLogin(url, args?.token, autoLogin);
 		if (!res) {
 			return; // The user aborted, or unable to auth.
 		}
@@ -136,12 +133,12 @@ export class Commands {
 	}
 
 	/**
-	 * If necessary, ask for a token, and keep asking until the token has been
-	 * validated.  Return the token and user that was fetched to validate the
-	 * token.  Null means the user aborted or we were unable to authenticate with
-	 * mTLS (in the latter case, an error notification will have been displayed).
+	 * Attempt to authenticate using OAuth, token, or mTLS. If necessary, prompts
+	 * for authentication method and credentials. Returns the token and user upon
+	 * successful authentication. Null means the user aborted or authentication
+	 * failed (in which case an error notification will have been displayed).
 	 */
-	private async maybeAskToken(
+	private async attemptLogin(
 		url: string,
 		token: string | undefined,
 		isAutoLogin: boolean,
@@ -174,58 +171,18 @@ export class Commands {
 			}
 		}
 
-		// Check if server supports OAuth
-		const supportsOAuth = await this.checkOAuthSupport(client);
-
-		let choice: AuthMethod | undefined = "legacy";
-		if (supportsOAuth) {
-			choice = await this.askAuthMethod();
-		}
-
-		if (choice === "oauth") {
-			return this.loginWithOAuth(client);
-		} else if (choice === "legacy") {
-			const initialToken =
-				token || (await this.secretsManager.getSessionToken());
-			return this.loginWithToken(client, initialToken);
+		const authMethod = await maybeAskAuthMethod(client);
+		switch (authMethod) {
+			case "oauth":
+				return this.loginWithOAuth(client);
+			case "legacy": {
+				const initialToken =
+					token || (await this.secretsManager.getSessionToken());
+				return this.loginWithToken(client, initialToken);
+			}
+			case undefined:
+				return null; // User aborted
 		}
-
-		return null; // User aborted.
-	}
-
-	private async checkOAuthSupport(client: CoderApi): Promise<boolean> {
-		const metadataClient = new OAuthMetadataClient(
-			client.getAxiosInstance(),
-			this.logger,
-		);
-		return metadataClient.checkOAuthSupport();
-	}
-
-	/**
-	 * Ask user to choose between OAuth and legacy API token authentication.
-	 */
-	private async askAuthMethod(): Promise<AuthMethod | undefined> {
-		const choice = await vscode.window.showQuickPick(
-			[
-				{
-					label: "$(key) OAuth (Recommended)",
-					detail: "Secure authentication with automatic token refresh",
-					value: "oauth" as const,
-				},
-				{
-					label: "$(lock) API Token",
-					detail: "Use a manually created API key",
-					value: "legacy" as const,
-				},
-			],
-			{
-				title: "Choose Authentication Method",
-				placeHolder: "How would you like to authenticate?",
-				ignoreFocusOut: true,
-			},
-		);
-
-		return choice?.value;
 	}
 
 	private async loginWithToken(
@@ -297,7 +254,15 @@ export class Commands {
 		try {
 			this.logger.info("Starting OAuth authentication");
 
-			const tokenResponse = await this.oauthSessionManager.login(client);
+			const tokenResponse = await vscode.window.withProgress(
+				{
+					location: vscode.ProgressLocation.Notification,
+					title: "Authenticating",
+					cancellable: false,
+				},
+				async (progress) =>
+					await this.oauthSessionManager.login(client, progress),
+			);
 
 			// Validate token by fetching user
 			client.setSessionToken(tokenResponse.access_token);

src/oauth/metadataClient.ts

@@ -26,13 +26,13 @@ export class OAuthMetadataClient {
 	/**
 	 * Check if a server supports OAuth by attempting to fetch the well-known endpoint.
 	 */
-	async checkOAuthSupport(): Promise<boolean> {
+	public static async checkOAuthSupport(
+		axiosInstance: AxiosInstance,
+	): Promise<boolean> {
 		try {
-			await this.axiosInstance.get(OAUTH_DISCOVERY_ENDPOINT);
-			this.logger.debug("Server supports OAuth");
+			await axiosInstance.get(OAUTH_DISCOVERY_ENDPOINT);
 			return true;
-		} catch (error) {
-			this.logger.debug("Server does not support OAuth:", error);
+		} catch {
 			return false;
 		}
 	}

src/oauth/sessionManager.ts

@@ -263,7 +263,10 @@ export class OAuthSessionManager implements vscode.Disposable {
 	 * @param client CoderApi instance for the deployment to authenticate against
 	 * @returns TokenResponse containing access token and optional refresh token
 	 */
-	async login(client: CoderApi): Promise<TokenResponse> {
+	async login(
+		client: CoderApi,
+		progress: vscode.Progress<{ message?: string; increment?: number }>,
+	): Promise<TokenResponse> {
 		const baseUrl = client.getAxiosInstance().defaults.baseURL;
 		if (!baseUrl) {
 			throw new Error("CoderApi instance has no base URL set");
@@ -282,13 +285,16 @@ export class OAuthSessionManager implements vscode.Disposable {
 		const metadata = await metadataClient.getMetadata();
 
 		// Only register the client on login
+		progress.report({ message: "registering client...", increment: 10 });
 		const registration = await this.registerClient(axiosInstance, metadata);
 
+		progress.report({ message: "waiting for authorization...", increment: 30 });
 		const { code, verifier } = await this.startAuthorization(
 			metadata,
 			registration,
 		);
 
+		progress.report({ message: "exchanging token...", increment: 30 });
 		const tokenResponse = await this.exchangeToken(
 			code,
 			verifier,
@@ -297,6 +303,7 @@ export class OAuthSessionManager implements vscode.Disposable {
 			registration,
 		);
 
+		progress.report({ increment: 30 });
 		this.logger.info("OAuth login flow completed successfully");
 
 		return tokenResponse;

src/promptUtils.ts

@@ -1,7 +1,11 @@
 import { type WorkspaceAgent } from "coder/site/src/api/typesGenerated";
 import * as vscode from "vscode";
 
+import { type CoderApi } from "./api/coderApi";
 import { type MementoManager } from "./core/mementoManager";
+import { OAuthMetadataClient } from "./oauth/metadataClient";
+
+type AuthMethod = "oauth" | "legacy";
 
 /**
  * Find the requested agent if specified, otherwise return the agent if there
@@ -129,3 +133,54 @@ export async function maybeAskUrl(
 	}
 	return url;
 }
+
+export async function maybeAskAuthMethod(
+	client: CoderApi,
+): Promise<AuthMethod | undefined> {
+	// Check if server supports OAuth with progress indication
+	const supportsOAuth = await vscode.window.withProgress(
+		{
+			location: vscode.ProgressLocation.Notification,
+			title: "Checking authentication methods",
+			cancellable: false,
+		},
+		async () => {
+			return await OAuthMetadataClient.checkOAuthSupport(
+				client.getAxiosInstance(),
+			);
+		},
+	);
+
+	if (supportsOAuth) {
+		return await askAuthMethod();
+	} else {
+		return "legacy";
+	}
+}
+
+/**
+ * Ask user to choose between OAuth and legacy API token authentication.
+ */
+async function askAuthMethod(): Promise<AuthMethod | undefined> {
+	const choice = await vscode.window.showQuickPick(
+		[
+			{
+				label: "OAuth (Recommended)",
+				description: "Secure authentication with automatic token refresh",
+				value: "oauth" as const,
+			},
+			{
+				label: "Session Token (Legacy)",
+				description: "Generate and paste a session token manually",
+				value: "legacy" as const,
+			},
+		],
+		{
+			title: "Select authentication method",
+			placeHolder: "How would you like to authenticate?",
+			ignoreFocusOut: true,
+		},
+	);
+
+	return choice?.value;
+}