Add OAuth 2.1 authentication support

src/commands.ts

@@ -19,6 +19,7 @@ import { type DeploymentManager } from "./deployment/deploymentManager";
 import { CertificateError } from "./error";
 import { type Logger } from "./logging/logger";
 import { type LoginCoordinator } from "./login/loginCoordinator";
+import { type OAuthSessionManager } from "./oauth/sessionManager";
 import { maybeAskAgent, maybeAskUrl } from "./promptUtils";
 import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util";
 import {
@@ -51,6 +52,7 @@ export class Commands {
 	public constructor(
 		serviceContainer: ServiceContainer,
 		private readonly extensionClient: CoderApi,
+		private readonly oauthSessionManager: OAuthSessionManager,
 		private readonly deploymentManager: DeploymentManager,
 	) {
 		this.vscodeProposed = serviceContainer.getVsCodeProposed();
@@ -105,6 +107,7 @@ export class Commands {
 			safeHostname,
 			url,
 			autoLogin: args?.autoLogin,
+			oauthSessionManager: this.oauthSessionManager,
 		});
 
 		if (!result.success) {

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";
@@ -7,7 +8,12 @@ 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 SESSION_KEY_PREFIX = "coder.session." as const;
+const OAUTH_CLIENT_PREFIX = "coder.oauth.client." as const;
+
+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" | "DPoP";
+	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,56 @@ 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;
+			}
+
+			try {
+				const data = await this.secrets.get(OAUTH_CALLBACK_KEY);
+				if (data) {
+					const parsed = JSON.parse(data) as OAuthCallbackData;
+					listener(parsed);
+				}
+			} catch {
+				// Ignore parse errors
+			}
+		});
+	}
+
+	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,17 @@
 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 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 +23,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 +42,7 @@ export class DeploymentManager implements vscode.Disposable {
 	private constructor(
 		serviceContainer: ServiceContainer,
 		private readonly client: CoderApi,
+		private readonly oauthSessionManager: OAuthSessionManager,
 		private readonly workspaceProviders: WorkspaceProvider[],
 	) {
 		this.secretsManager = serviceContainer.getSecretsManager();
@@ -52,11 +54,13 @@ export class DeploymentManager implements vscode.Disposable {
 	public static create(
 		serviceContainer: ServiceContainer,
 		client: CoderApi,
+		oauthSessionManager: OAuthSessionManager,
 		workspaceProviders: WorkspaceProvider[],
 	): DeploymentManager {
 		const manager = new DeploymentManager(
 			serviceContainer,
 			client,
+			oauthSessionManager,
 			workspaceProviders,
 		);
 		manager.subscribeToCrossWindowChanges();
@@ -125,9 +129,13 @@ 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.persistDeployment(deployment);
 	}
 
@@ -140,6 +148,7 @@ export class DeploymentManager implements vscode.Disposable {
 		this.#deployment = null;
 
 		this.client.setCredentials(undefined, undefined);
+		this.oauthSessionManager.clearDeployment();
 		this.updateAuthContexts();
 		this.refreshWorkspaces();