Add OAuth 2.1 authentication support

src/api/oauthInterceptors.ts

@@ -0,0 +1,116 @@
+import { type AxiosError, isAxiosError } from "axios";
+
+import { type Logger } from "../logging/logger";
+import { type RequestConfigWithMeta } from "../logging/types";
+import { parseOAuthError, requiresReAuthentication } from "../oauth/errors";
+import { type OAuthSessionManager } from "../oauth/sessionManager";
+
+import { type CoderApi } from "./coderApi";
+
+const coderSessionTokenHeader = "Coder-Session-Token";
+
+/**
+ * Attach OAuth token refresh interceptors to a CoderApi instance.
+ * This should be called after creating the CoderApi when OAuth authentication is being used.
+ *
+ * Success interceptor: proactively refreshes token when approaching expiry.
+ * Error interceptor: reactively refreshes token on 401 responses.
+ */
+export function attachOAuthInterceptors(
+	client: CoderApi,
+	logger: Logger,
+	oauthSessionManager: OAuthSessionManager,
+): void {
+	client.getAxiosInstance().interceptors.response.use(
+		// Success response interceptor: proactive token refresh
+		(response) => {
+			// Fire-and-forget: don't await, don't block response
+			oauthSessionManager.refreshIfAlmostExpired().catch((error) => {
+				logger.warn("Proactive background token refresh failed:", error);
+			});
+
+			return response;
+		},
+		// Error response interceptor: reactive token refresh on 401
+		async (error: unknown) => {
+			if (!isAxiosError(error)) {
+				throw error;
+			}
+
+			if (error.config) {
+				const config = error.config as {
+					_oauthRetryAttempted?: boolean;
+				};
+				if (config._oauthRetryAttempted) {
+					throw error;
+				}
+			}
+
+			const status = error.response?.status;
+
+			// These could indicate permanent auth failures that won't be fixed by token refresh
+			if (status === 400 || status === 403) {
+				handlePossibleOAuthError(error, logger, oauthSessionManager);
+				throw error;
+			} else if (status === 401) {
+				return handle401Error(error, client, logger, oauthSessionManager);
+			}
+
+			throw error;
+		},
+	);
+}
+
+function handlePossibleOAuthError(
+	error: unknown,
+	logger: Logger,
+	oauthSessionManager: OAuthSessionManager,
+): void {
+	const oauthError = parseOAuthError(error);
+	if (oauthError && requiresReAuthentication(oauthError)) {
+		logger.error(
+			`OAuth error requires re-authentication: ${oauthError.errorCode}`,
+		);
+
+		oauthSessionManager.showReAuthenticationModal(oauthError).catch((err) => {
+			logger.error("Failed to show re-auth modal:", err);
+		});
+	}
+}
+
+async function handle401Error(
+	error: AxiosError,
+	client: CoderApi,
+	logger: Logger,
+	oauthSessionManager: OAuthSessionManager,
+): Promise<void> {
+	if (!oauthSessionManager.isLoggedInWithOAuth()) {
+		throw error;
+	}
+
+	logger.info("Received 401 response, attempting token refresh");
+
+	try {
+		const newTokens = await oauthSessionManager.refreshToken();
+		client.setSessionToken(newTokens.access_token);
+
+		logger.info("Token refresh successful, retrying request");
+
+		// Retry the original request with the new token
+		if (error.config) {
+			const config = error.config as RequestConfigWithMeta & {
+				_oauthRetryAttempted?: boolean;
+			};
+			config._oauthRetryAttempted = true;
+			config.headers[coderSessionTokenHeader] = newTokens.access_token;
+			return client.getAxiosInstance().request(config);
+		}
+
+		throw error;
+	} catch (refreshError) {
+		logger.error("Token refresh failed:", refreshError);
+
+		handlePossibleOAuthError(refreshError, logger, oauthSessionManager);
+		throw error;
+	}
+}

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,8 @@
 import { type Logger } from "../logging/logger";
+import {
+	type ClientRegistrationResponse,
+	type TokenResponse,
+} from "../oauth/types";
 import { toSafeHost } from "../util";
 
 import type { Memento, SecretStorage, Disposable } from "vscode";
@@ -8,6 +12,10 @@ 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_TOKENS_PREFIX = "coder.oauth.tokens.";
+const OAUTH_CLIENT_PREFIX = "coder.oauth.client.";
+
+const OAUTH_CALLBACK_KEY = "coder.oauthCallback";
 
 const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment";
 
@@ -31,6 +39,17 @@ interface DeploymentUsage {
 	lastAccessedAt: string;
 }
 
+export type StoredOAuthTokens = Omit<TokenResponse, "expires_in"> & {
+	expiry_timestamp: number;
+	deployment_url: string;
+};
+
+interface OAuthCallbackData {
+	state: string;
+	code: string | null;
+	error: string | null;
+}
+
 export class SecretsManager {
 	constructor(
 		private readonly secrets: SecretStorage,
@@ -181,6 +200,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 Promise.all([
+			this.clearSessionAuth(safeHostname),
+			this.clearOAuthData(safeHostname),
+		]);
 		await this.clearSessionAuth(safeHostname);
 		const usage = this.getDeploymentUsage().filter(
 			(u) => u.safeHostname !== safeHostname,
@@ -234,4 +257,107 @@ 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 async getOAuthTokens(
+		safeHostname: string,
+	): Promise<StoredOAuthTokens | undefined> {
+		try {
+			const data = await this.secrets.get(
+				`${OAUTH_TOKENS_PREFIX}${safeHostname}`,
+			);
+			if (!data) {
+				return undefined;
+			}
+			return JSON.parse(data) as StoredOAuthTokens;
+		} catch {
+			return undefined;
+		}
+	}
+
+	public async setOAuthTokens(
+		safeHostname: string,
+		tokens: StoredOAuthTokens,
+	): Promise<void> {
+		await this.secrets.store(
+			`${OAUTH_TOKENS_PREFIX}${safeHostname}`,
+			JSON.stringify(tokens),
+		);
+		await this.recordDeploymentAccess(safeHostname);
+	}
+
+	public async clearOAuthTokens(safeHostname: string): Promise<void> {
+		await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${safeHostname}`);
+	}
+
+	public async getOAuthClientRegistration(
+		safeHostname: string,
+	): Promise<ClientRegistrationResponse | undefined> {
+		try {
+			const data = await this.secrets.get(
+				`${OAUTH_CLIENT_PREFIX}${safeHostname}`,
+			);
+			if (!data) {
+				return undefined;
+			}
+			return JSON.parse(data) as ClientRegistrationResponse;
+		} catch {
+			return undefined;
+		}
+	}
+
+	public async setOAuthClientRegistration(
+		safeHostname: string,
+		registration: ClientRegistrationResponse,
+	): Promise<void> {
+		await this.secrets.store(
+			`${OAUTH_CLIENT_PREFIX}${safeHostname}`,
+			JSON.stringify(registration),
+		);
+		await this.recordDeploymentAccess(safeHostname);
+	}
+
+	public async clearOAuthClientRegistration(
+		safeHostname: string,
+	): Promise<void> {
+		await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${safeHostname}`);
+	}
+
+	public async clearOAuthData(safeHostname: string): Promise<void> {
+		await Promise.all([
+			this.clearOAuthTokens(safeHostname),
+			this.clearOAuthClientRegistration(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();