Hook up into the login/logout logic of the extension

Author: EhabY

package.json

@@ -255,16 +255,6 @@
 				"title": "Search",
 				"category": "Coder",
 				"icon": "$(search)"
-			},
-			{
-				"command": "coder.oauth.login",
-				"title": "OAuth Login",
-				"category": "Coder"
-			},
-			{
-				"command": "coder.oauth.logout",
-				"title": "OAuth Logout",
-				"category": "Coder"
 			}
 		],
 		"menus": {

src/commands.ts

@@ -19,6 +19,7 @@ import { type SecretsManager } from "./core/secretsManager";
 import { CertificateError } from "./error";
 import { getGlobalFlags } from "./globalFlags";
 import { type Logger } from "./logging/logger";
+import { type CoderOAuthHelper } from "./oauth/oauthHelper";
 import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util";
 import {
 	AgentTreeItem,
@@ -48,6 +49,7 @@ export class Commands {
 	public constructor(
 		serviceContainer: ServiceContainer,
 		private readonly restClient: Api,
+		private readonly oauthHelper: CoderOAuthHelper,
 	) {
 		this.vscodeProposed = serviceContainer.getVsCodeProposed();
 		this.logger = serviceContainer.getLogger();
@@ -182,59 +184,119 @@ export class Commands {
 	}
 
 	/**
-	 * Log into the provided deployment. If the deployment URL is not specified,
-	 * ask for it first with a menu showing recent URLs along with the default URL
-	 * and CODER_URL, if those are set.
+	 * Check if server supports OAuth by attempting to fetch the well-known endpoint.
 	 */
-	public async login(args?: {
-		url?: string;
-		token?: string;
-		label?: string;
-		autoLogin?: boolean;
-	}): Promise<void> {
-		if (this.contextManager.get("coder.authenticated")) {
-			return;
+	private async checkOAuthSupport(client: CoderApi): Promise<boolean> {
+		try {
+			await client
+				.getAxiosInstance()
+				.get("/.well-known/oauth-authorization-server");
+			this.logger.debug("Server supports OAuth");
+			return true;
+		} catch (error) {
+			this.logger.debug("Server does not support OAuth:", error);
+			return false;
 		}
-		this.logger.info("Logging in");
+	}
 
-		const url = await this.maybeAskUrl(args?.url);
-		if (!url) {
-			return; // The user aborted.
-		}
+	/**
+	 * Ask user to choose between OAuth and legacy API token authentication.
+	 */
+	private async askAuthMethod(): Promise<"oauth" | "legacy" | undefined> {
+		const choice = await vscode.window.showQuickPick(
+			[
+				{
+					label: "$(key) OAuth (Recommended)",
+					detail: "Secure authentication with automatic token refresh",
+					value: "oauth",
+				},
+				{
+					label: "$(lock) API Token",
+					detail: "Use a manually created API key",
+					value: "legacy",
+				},
+			],
+			{
+				title: "Choose Authentication Method",
+				placeHolder: "How would you like to authenticate?",
+				ignoreFocusOut: true,
+			},
+		);
 
-		// It is possible that we are trying to log into an old-style host, in which
-		// case we want to write with the provided blank label instead of generating
-		// a host label.
-		const label = args?.label === undefined ? toSafeHost(url) : args.label;
+		return choice?.value as "oauth" | "legacy" | undefined;
+	}
 
-		// 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);
-		if (!res) {
-			return; // The user aborted, or unable to auth.
+	/**
+	 * Authenticate using OAuth flow.
+	 * Returns the access token and authenticated user, or null if failed/cancelled.
+	 */
+	private async loginWithOAuth(
+		url: string,
+	): Promise<{ user: User; token: string } | null> {
+		try {
+			this.logger.info("Starting OAuth authentication");
+
+			// Start OAuth authorization flow
+			const { code, verifier } = await this.oauthHelper.startAuthorization(url);
+
+			// Exchange authorization code for tokens
+			const tokenResponse = await this.oauthHelper.exchangeToken(
+				code,
+				verifier,
+			);
+
+			// Validate token by fetching user
+			const client = CoderApi.create(
+				url,
+				tokenResponse.access_token,
+				this.logger,
+			);
+			const user = await client.getAuthenticatedUser();
+
+			this.logger.info("OAuth authentication successful");
+
+			return {
+				token: tokenResponse.access_token,
+				user,
+			};
+		} catch (error) {
+			this.logger.error("OAuth authentication failed:", error);
+			vscode.window.showErrorMessage(
+				`OAuth authentication failed: ${getErrorMessage(error, "Unknown error")}`,
+			);
+			return null;
 		}
+	}
 
-		// The URL is good and the token is either good or not required; authorize
-		// the global client.
+	/**
+	 * Complete the login process by storing credentials and updating context.
+	 */
+	private async completeLogin(
+		url: string,
+		label: string,
+		token: string,
+		user: User,
+	): Promise<void> {
+		// Authorize the global client
 		this.restClient.setHost(url);
-		this.restClient.setSessionToken(res.token);
+		this.restClient.setSessionToken(token);
 
-		// Store these to be used in later sessions.
+		// Store for later sessions
 		await this.mementoManager.setUrl(url);
-		await this.secretsManager.setSessionToken(res.token);
+		await this.secretsManager.setSessionToken(token);
 
-		// Store on disk to be used by the cli.
-		await this.cliManager.configure(label, url, res.token);
+		// Store on disk for CLI
+		await this.cliManager.configure(label, url, token);
 
-		// These contexts control various menu items and the sidebar.
+		// Update contexts
 		this.contextManager.set("coder.authenticated", true);
-		if (res.user.roles.find((role) => role.name === "owner")) {
+		if (user.roles.find((role) => role.name === "owner")) {
 			this.contextManager.set("coder.isOwner", true);
 		}
 
 		vscode.window
 			.showInformationMessage(
-				`Welcome to Coder, ${res.user.username}!`,
+				`Welcome to Coder, ${user.username}!`,
 				{
 					detail:
 						"You can now use the Coder extension to manage your Coder instance.",
@@ -252,6 +314,73 @@ export class Commands {
 		vscode.commands.executeCommand("coder.refreshWorkspaces");
 	}
 
+	/**
+	 * Log into the provided deployment. If the deployment URL is not specified,
+	 * ask for it first with a menu showing recent URLs along with the default URL
+	 * and CODER_URL, if those are set.
+	 */
+	public async login(args?: {
+		url?: string;
+		token?: string;
+		label?: string;
+		autoLogin?: boolean;
+	}): Promise<void> {
+		if (this.contextManager.get("coder.authenticated")) {
+			return;
+		}
+		this.logger.info("Logging in");
+
+		const url = await this.maybeAskUrl(args?.url);
+		if (!url) {
+			return; // The user aborted.
+		}
+
+		const label = args?.label ?? toSafeHost(url);
+		const autoLogin = args?.autoLogin === true;
+
+		// Check if we have an existing valid legacy token
+		const existingToken = await this.secretsManager.getSessionToken();
+		const client = CoderApi.create(url, existingToken, this.logger);
+		if (existingToken && !args?.token) {
+			try {
+				const user = await client.getAuthenticatedUser();
+				this.logger.info("Using existing valid session token");
+				await this.completeLogin(url, label, existingToken, user);
+				return;
+			} catch {
+				this.logger.debug("Existing token invalid, clearing it");
+				await this.secretsManager.setSessionToken();
+			}
+		}
+
+		// Check if server supports OAuth
+		const supportsOAuth = await this.checkOAuthSupport(client);
+
+		if (supportsOAuth && !autoLogin) {
+			const choice = await this.askAuthMethod();
+			if (!choice) {
+				return;
+			}
+
+			if (choice === "oauth") {
+				const res = await this.loginWithOAuth(url);
+				if (!res) {
+					return;
+				}
+				await this.completeLogin(url, label, res.token, res.user);
+				return;
+			}
+		}
+
+		// Use legacy token flow (existing behavior)
+		const res = await this.maybeAskToken(url, args?.token, autoLogin);
+		if (!res) {
+			return;
+		}
+
+		await this.completeLogin(url, label, res.token, res.user);
+	}
+
 	/**
 	 * 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
@@ -377,6 +506,22 @@ export class Commands {
 			// Sanity check; command should not be available if no url.
 			throw new Error("You are not logged in");
 		}
+
+		// Check if using OAuth
+		const hasOAuthTokens = await this.secretsManager.getOAuthTokens();
+		if (hasOAuthTokens) {
+			this.logger.info("Logging out via OAuth");
+			try {
+				await this.oauthHelper.logout();
+			} catch (error) {
+				this.logger.warn(
+					"OAuth logout failed, continuing with cleanup:",
+					error,
+				);
+			}
+		}
+
+		// Continue with standard logout (clears sessionToken, contexts, etc)
 		await this.forceLogout();
 	}
 

src/core/secretsManager.ts

@@ -84,6 +84,20 @@ export class SecretsManager {
 		});
 	}
 
+	/**
+	 * Listens for session token changes.
+	 */
+	public onDidChangeSessionToken(
+		listener: (token: string | undefined) => Promise<void>,
+	): Disposable {
+		return this.secrets.onDidChange(async (e) => {
+			if (e.key === SESSION_TOKEN_KEY) {
+				const token = await this.getSessionToken();
+				await listener(token);
+			}
+		});
+	}
+
 	/**
 	 * Store OAuth client registration data.
 	 */

src/extension.ts

@@ -124,11 +124,34 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	);
 
 	const oauthHelper = await activateCoderOAuth(
-		client,
+		url || "",
 		secretsManager,
 		output,
 		ctx,
 	);
+	ctx.subscriptions.push(oauthHelper);
+
+	// Listen for session token changes and sync state across all components
+	ctx.subscriptions.push(
+		secretsManager.onDidChangeSessionToken(async (token) => {
+			if (!token) {
+				output.debug("Session token cleared");
+				client.setSessionToken("");
+				return;
+			}
+
+			output.debug("Session token changed, syncing state");
+
+			client.setSessionToken(token);
+			const url = mementoManager.getUrl();
+			if (url) {
+				const cliManager = serviceContainer.getCliManager();
+				// TODO label might not match?
+				await cliManager.configure(toSafeHost(url), url, token);
+				output.debug("Updated CLI config with new token");
+			}
+		}),
+	);
 
 	// Handle vscode:// URIs.
 	const uriHandler = vscode.window.registerUriHandler({
@@ -293,7 +316,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 
 	// Register globally available commands.  Many of these have visibility
 	// controlled by contexts, see `when` in the package.json.
-	const commands = new Commands(serviceContainer, client);
+	const commands = new Commands(serviceContainer, client, oauthHelper);
 	ctx.subscriptions.push(
 		vscode.commands.registerCommand(
 			"coder.login",