Add OAuth 2.1 authentication support

src/core/container.ts

@@ -48,6 +48,7 @@ export class ServiceContainer implements vscode.Disposable {
 			this.mementoManager,
 			this.vscodeProposed,
 			this.logger,
+			context.extension.id,
 		);
 	}
 
@@ -89,5 +90,6 @@ export class ServiceContainer implements vscode.Disposable {
 	dispose(): void {
 		this.contextManager.dispose();
 		this.logger.dispose();
+		this.loginCoordinator.dispose();
 	}
 }

src/core/secretsManager.ts

@@ -1,4 +1,5 @@
 import { type Logger } from "../logging/logger";
+import { type ClientRegistrationResponse } from "../oauth/types";
 import { toSafeHost } from "../util";
 
 import type { Memento, SecretStorage, Disposable } from "vscode";
@@ -8,6 +9,11 @@ import type { Deployment } from "../deployment/types";
 // Each deployment has its own key to ensure atomic operations (multiple windows
 // writing to a shared key could drop data) and to receive proper VS Code events.
 const SESSION_KEY_PREFIX = "coder.session.";
+const OAUTH_CLIENT_PREFIX = "coder.oauth.client.";
+
+type SecretKeyPrefix = typeof SESSION_KEY_PREFIX | typeof OAUTH_CLIENT_PREFIX;
+
+const OAUTH_CALLBACK_KEY = "coder.oauthCallback";
 
 const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment";
 
@@ -20,9 +26,22 @@ export interface CurrentDeploymentState {
 	deployment: Deployment | null;
 }
 
+/**
+ * OAuth token data stored alongside session auth.
+ * When present, indicates the session is authenticated via OAuth.
+ */
+export interface OAuthTokenData {
+	token_type: "Bearer";
+	refresh_token?: string;
+	scope?: string;
+	expiry_timestamp: number;
+}
+
 export interface SessionAuth {
 	url: string;
 	token: string;
+	/** If present, this session uses OAuth authentication */
+	oauth?: OAuthTokenData;
 }
 
 // Tracks when a deployment was last accessed for LRU pruning.
@@ -31,13 +50,57 @@ interface DeploymentUsage {
 	lastAccessedAt: string;
 }
 
+interface OAuthCallbackData {
+	state: string;
+	code: string | null;
+	error: string | null;
+}
+
 export class SecretsManager {
 	constructor(
 		private readonly secrets: SecretStorage,
 		private readonly memento: Memento,
 		private readonly logger: Logger,
 	) {}
 
+	private buildKey(prefix: SecretKeyPrefix, safeHostname: string): string {
+		return `${prefix}${safeHostname || "<legacy>"}`;
+	}
+
+	private async getSecret<T>(
+		prefix: SecretKeyPrefix,
+		safeHostname: string,
+	): Promise<T | undefined> {
+		try {
+			const data = await this.secrets.get(this.buildKey(prefix, safeHostname));
+			if (!data) {
+				return undefined;
+			}
+			return JSON.parse(data) as T;
+		} catch {
+			return undefined;
+		}
+	}
+
+	private async setSecret<T>(
+		prefix: SecretKeyPrefix,
+		safeHostname: string,
+		value: T,
+	): Promise<void> {
+		await this.secrets.store(
+			this.buildKey(prefix, safeHostname),
+			JSON.stringify(value),
+		);
+		await this.recordDeploymentAccess(safeHostname);
+	}
+
+	private async clearSecret(
+		prefix: SecretKeyPrefix,
+		safeHostname: string,
+	): Promise<void> {
+		await this.secrets.delete(this.buildKey(prefix, safeHostname));
+	}
+
 	/**
 	 * Sets the current deployment and triggers a cross-window sync event.
 	 */
@@ -104,7 +167,7 @@ export class SecretsManager {
 		safeHostname: string,
 		listener: (auth: SessionAuth | undefined) => void | Promise<void>,
 	): Disposable {
-		const sessionKey = this.getSessionKey(safeHostname);
+		const sessionKey = this.buildKey(SESSION_KEY_PREFIX, safeHostname);
 		return this.secrets.onDidChange(async (e) => {
 			if (e.key !== sessionKey) {
 				return;
@@ -118,39 +181,27 @@ export class SecretsManager {
 		});
 	}
 
-	public async getSessionAuth(
+	public getSessionAuth(
 		safeHostname: string,
 	): Promise<SessionAuth | undefined> {
-		const sessionKey = this.getSessionKey(safeHostname);
-		try {
-			const data = await this.secrets.get(sessionKey);
-			if (!data) {
-				return undefined;
-			}
-			return JSON.parse(data) as SessionAuth;
-		} catch {
-			return undefined;
-		}
+		return this.getSecret<SessionAuth>(SESSION_KEY_PREFIX, safeHostname);
 	}
 
 	public async setSessionAuth(
 		safeHostname: string,
 		auth: SessionAuth,
 	): Promise<void> {
-		const sessionKey = this.getSessionKey(safeHostname);
-		// Extract only url and token before serializing
-		const state: SessionAuth = { url: auth.url, token: auth.token };
-		await this.secrets.store(sessionKey, JSON.stringify(state));
-		await this.recordDeploymentAccess(safeHostname);
-	}
-
-	private async clearSessionAuth(safeHostname: string): Promise<void> {
-		const sessionKey = this.getSessionKey(safeHostname);
-		await this.secrets.delete(sessionKey);
+		// Extract relevant fields before serializing
+		const state: SessionAuth = {
+			url: auth.url,
+			token: auth.token,
+			...(auth.oauth && { oauth: auth.oauth }),
+		};
+		await this.setSecret(SESSION_KEY_PREFIX, safeHostname, state);
 	}
 
-	private getSessionKey(safeHostname: string): string {
-		return `${SESSION_KEY_PREFIX}${safeHostname || "<legacy>"}`;
+	private clearSessionAuth(safeHostname: string): Promise<void> {
+		return this.clearSecret(SESSION_KEY_PREFIX, safeHostname);
 	}
 
 	/**
@@ -181,7 +232,10 @@ export class SecretsManager {
 	 * Clear all auth data for a deployment and remove it from the usage list.
 	 */
 	public async clearAllAuthData(safeHostname: string): Promise<void> {
-		await this.clearSessionAuth(safeHostname);
+		await Promise.all([
+			this.clearSessionAuth(safeHostname),
+			this.clearOAuthClientRegistration(safeHostname),
+		]);
 		const usage = this.getDeploymentUsage().filter(
 			(u) => u.safeHostname !== safeHostname,
 		);
@@ -234,4 +288,64 @@ export class SecretsManager {
 
 		return safeHostname;
 	}
+
+	/**
+	 * Write an OAuth callback result to secrets storage.
+	 * Used for cross-window communication when OAuth callback arrives in a different window.
+	 */
+	public async setOAuthCallback(data: OAuthCallbackData): Promise<void> {
+		await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(data));
+	}
+
+	/**
+	 * Listen for OAuth callback results from any VS Code window.
+	 * The listener receives the state parameter, code (if success), and error (if failed).
+	 */
+	public onDidChangeOAuthCallback(
+		listener: (data: OAuthCallbackData) => void,
+	): Disposable {
+		return this.secrets.onDidChange(async (e) => {
+			if (e.key !== OAUTH_CALLBACK_KEY) {
+				return;
+			}
+
+			let parsed: OAuthCallbackData;
+			try {
+				const data = await this.secrets.get(OAUTH_CALLBACK_KEY);
+				if (!data) {
+					return;
+				}
+				parsed = JSON.parse(data) as OAuthCallbackData;
+			} catch (err) {
+				this.logger.error("Failed to parse OAuth callback data", err);
+				return;
+			}
+
+			try {
+				listener(parsed);
+			} catch (err) {
+				this.logger.error("Error in onDidChangeOAuthCallback listener", err);
+			}
+		});
+	}
+
+	public getOAuthClientRegistration(
+		safeHostname: string,
+	): Promise<ClientRegistrationResponse | undefined> {
+		return this.getSecret<ClientRegistrationResponse>(
+			OAUTH_CLIENT_PREFIX,
+			safeHostname,
+		);
+	}
+
+	public setOAuthClientRegistration(
+		safeHostname: string,
+		registration: ClientRegistrationResponse,
+	): Promise<void> {
+		return this.setSecret(OAUTH_CLIENT_PREFIX, safeHostname, registration);
+	}
+
+	public clearOAuthClientRegistration(safeHostname: string): Promise<void> {
+		return this.clearSecret(OAUTH_CLIENT_PREFIX, safeHostname);
+	}
 }

src/deployment/deploymentManager.ts

@@ -1,17 +1,18 @@
 import { CoderApi } from "../api/coderApi";
+import { type ServiceContainer } from "../core/container";
+import { type ContextManager } from "../core/contextManager";
+import { type MementoManager } from "../core/mementoManager";
+import { type SecretsManager } from "../core/secretsManager";
+import { type Logger } from "../logging/logger";
+import { type OAuthInterceptor } from "../oauth/axiosInterceptor";
+import { type OAuthSessionManager } from "../oauth/sessionManager";
+import { type WorkspaceProvider } from "../workspace/workspacesProvider";
+
+import { type Deployment, type DeploymentWithAuth } from "./types";
 
 import type { User } from "coder/site/src/api/typesGenerated";
 import type * as vscode from "vscode";
 
-import type { ServiceContainer } from "../core/container";
-import type { ContextManager } from "../core/contextManager";
-import type { MementoManager } from "../core/mementoManager";
-import type { SecretsManager } from "../core/secretsManager";
-import type { Logger } from "../logging/logger";
-import type { WorkspaceProvider } from "../workspace/workspacesProvider";
-
-import type { Deployment, DeploymentWithAuth } from "./types";
-
 /**
  * Internal state type that allows mutation of user property.
  */
@@ -23,6 +24,7 @@ type DeploymentWithUser = Deployment & { user: User };
  * Centralizes:
  * - In-memory deployment state (url, label, token, user)
  * - Client credential updates
+ * - OAuth session management
  * - Auth listener registration
  * - Context updates (coder.authenticated, coder.isOwner)
  * - Workspace provider refresh
@@ -41,6 +43,8 @@ export class DeploymentManager implements vscode.Disposable {
 	private constructor(
 		serviceContainer: ServiceContainer,
 		private readonly client: CoderApi,
+		private readonly oauthSessionManager: OAuthSessionManager,
+		private readonly oauthInterceptor: OAuthInterceptor,
 		private readonly workspaceProviders: WorkspaceProvider[],
 	) {
 		this.secretsManager = serviceContainer.getSecretsManager();
@@ -52,11 +56,15 @@ export class DeploymentManager implements vscode.Disposable {
 	public static create(
 		serviceContainer: ServiceContainer,
 		client: CoderApi,
+		oauthSessionManager: OAuthSessionManager,
+		oauthInterceptor: OAuthInterceptor,
 		workspaceProviders: WorkspaceProvider[],
 	): DeploymentManager {
 		const manager = new DeploymentManager(
 			serviceContainer,
 			client,
+			oauthSessionManager,
+			oauthInterceptor,
 			workspaceProviders,
 		);
 		manager.subscribeToCrossWindowChanges();
@@ -125,9 +133,14 @@ export class DeploymentManager implements vscode.Disposable {
 			this.client.setCredentials(deployment.url, deployment.token);
 		}
 
+		// Register auth listener before setDeployment so background token refresh
+		// can update client credentials via the listener
 		this.registerAuthListener();
 		this.updateAuthContexts();
 		this.refreshWorkspaces();
+
+		await this.oauthSessionManager.setDeployment(deployment);
+		await this.oauthInterceptor.setDeployment(deployment);
 		await this.persistDeployment(deployment);
 	}
 
@@ -140,6 +153,8 @@ export class DeploymentManager implements vscode.Disposable {
 		this.#deployment = null;
 
 		this.client.setCredentials(undefined, undefined);
+		this.oauthSessionManager.clearDeployment();
+		this.oauthInterceptor.clearDeployment();
 		this.updateAuthContexts();
 		this.refreshWorkspaces();
 

src/extension.ts

@@ -14,6 +14,8 @@ import { type SecretsManager } from "./core/secretsManager";
 import { DeploymentManager } from "./deployment/deploymentManager";
 import { CertificateError } from "./error/certificateError";
 import { getErrorDetail, toError } from "./error/errorUtils";
+import { OAuthInterceptor } from "./oauth/axiosInterceptor";
+import { OAuthSessionManager } from "./oauth/sessionManager";
 import { Remote } from "./remote/remote";
 import { getRemoteSshExtension } from "./remote/sshExtension";
 import { registerUriHandler } from "./uri/uriHandler";
@@ -68,6 +70,13 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 
 	const deployment = await secretsManager.getCurrentDeployment();
 
+	// Create OAuth session manager with login coordinator
+	const oauthSessionManager = OAuthSessionManager.create(
+		deployment,
+		serviceContainer,
+	);
+	ctx.subscriptions.push(oauthSessionManager);
+
 	// This client tracks the current login and will be used through the life of
 	// the plugin to poll workspaces for the current login, as well as being used
 	// in commands that operate on the current login.
@@ -79,6 +88,16 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	);
 	ctx.subscriptions.push(client);
 
+	// Create OAuth interceptor - auto attaches/detaches based on token state
+	const oauthInterceptor = await OAuthInterceptor.create(
+		client,
+		output,
+		oauthSessionManager,
+		secretsManager,
+		deployment?.safeHostname ?? "",
+	);
+	ctx.subscriptions.push(oauthInterceptor);
+
 	const myWorkspacesProvider = new WorkspaceProvider(
 		WorkspaceQuery.Mine,
 		client,
@@ -123,10 +142,13 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	);
 
 	// Create deployment manager to centralize deployment state management
-	const deploymentManager = DeploymentManager.create(serviceContainer, client, [
-		myWorkspacesProvider,
-		allWorkspacesProvider,
-	]);
+	const deploymentManager = DeploymentManager.create(
+		serviceContainer,
+		client,
+		oauthSessionManager,
+		oauthInterceptor,
+		[myWorkspacesProvider, allWorkspacesProvider],
+	);
 	ctx.subscriptions.push(deploymentManager);
 
 	// Register globally available commands.  Many of these have visibility