Always attach the auth interceptor and unify 401 handling

Author: EhabY

src/api/authInterceptor.ts

@@ -0,0 +1,123 @@
+import { type AxiosError, isAxiosError } from "axios";
+
+import { toSafeHost } from "../util";
+
+import type * as vscode from "vscode";
+
+import type { SecretsManager } from "../core/secretsManager";
+import type { Logger } from "../logging/logger";
+import type { RequestConfigWithMeta } from "../logging/types";
+import type { OAuthSessionManager } from "../oauth/sessionManager";
+
+import type { CoderApi } from "./coderApi";
+
+const coderSessionTokenHeader = "Coder-Session-Token";
+
+/**
+ * Callback invoked when authentication is required.
+ * Returns true if user successfully re-authenticated.
+ */
+export type AuthRequiredHandler = (hostname: string) => Promise<boolean>;
+
+/**
+ * Intercepts 401 responses and handles re-authentication.
+ *
+ * Always attached to the axios instance. Handles both OAuth (automatic refresh)
+ * and non-OAuth (interactive re-auth via callback) authentication failures.
+ */
+export class AuthInterceptor implements vscode.Disposable {
+	private readonly interceptorId: number;
+
+	constructor(
+		private readonly client: CoderApi,
+		private readonly logger: Logger,
+		private readonly oauthSessionManager: OAuthSessionManager,
+		private readonly secretsManager: SecretsManager,
+		private readonly onAuthRequired?: AuthRequiredHandler,
+	) {
+		this.interceptorId = this.client
+			.getAxiosInstance()
+			.interceptors.response.use(
+				(r) => r,
+				(error: unknown) => this.handleError(error),
+			);
+		this.logger.debug("Auth interceptor attached");
+	}
+
+	private async handleError(error: unknown): Promise<unknown> {
+		if (!isAxiosError(error)) {
+			throw error;
+		}
+
+		if (error.config) {
+			const config = error.config as { _retryAttempted?: boolean };
+			if (config._retryAttempted) {
+				throw error;
+			}
+		}
+
+		if (error.response?.status !== 401) {
+			throw error;
+		}
+
+		const baseUrl = this.client.getHost();
+		if (!baseUrl) {
+			throw error;
+		}
+		const hostname = toSafeHost(baseUrl);
+
+		return this.handle401Error(error, hostname);
+	}
+
+	private async handle401Error(
+		error: AxiosError,
+		hostname: string,
+	): Promise<unknown> {
+		this.logger.debug("Received 401 response, attempting recovery");
+
+		if (await this.oauthSessionManager.isLoggedInWithOAuth(hostname)) {
+			try {
+				const newTokens = await this.oauthSessionManager.refreshToken();
+				this.client.setSessionToken(newTokens.access_token);
+				this.logger.debug("Token refresh successful, retrying request");
+				return this.retryRequest(error, newTokens.access_token);
+			} catch (refreshError) {
+				this.logger.error("OAuth refresh failed:", refreshError);
+			}
+		}
+
+		if (this.onAuthRequired) {
+			this.logger.debug("Triggering interactive re-authentication");
+			const success = await this.onAuthRequired(hostname);
+			if (success) {
+				const auth = await this.secretsManager.getSessionAuth(hostname);
+				if (auth) {
+					this.logger.debug("Re-authentication successful, retrying request");
+					return this.retryRequest(error, auth.token);
+				}
+			}
+		}
+
+		throw error;
+	}
+
+	private retryRequest(error: AxiosError, token: string): Promise<unknown> {
+		if (!error.config) {
+			throw error;
+		}
+
+		const config = error.config as RequestConfigWithMeta & {
+			_retryAttempted?: boolean;
+		};
+		config._retryAttempted = true;
+		config.headers[coderSessionTokenHeader] = token;
+		return this.client.getAxiosInstance().request(config);
+	}
+
+	public dispose(): void {
+		this.client
+			.getAxiosInstance()
+			.interceptors.response.eject(this.interceptorId);
+		this.logger.debug("Auth interceptor detached");
+	}
+}

src/deployment/deploymentManager.ts

@@ -4,7 +4,6 @@ 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";
 
@@ -44,7 +43,6 @@ export class DeploymentManager implements vscode.Disposable {
 		serviceContainer: ServiceContainer,
 		private readonly client: CoderApi,
 		private readonly oauthSessionManager: OAuthSessionManager,
-		private readonly oauthInterceptor: OAuthInterceptor,
 		private readonly workspaceProviders: WorkspaceProvider[],
 	) {
 		this.secretsManager = serviceContainer.getSecretsManager();
@@ -57,14 +55,12 @@ export class DeploymentManager implements vscode.Disposable {
 		serviceContainer: ServiceContainer,
 		client: CoderApi,
 		oauthSessionManager: OAuthSessionManager,
-		oauthInterceptor: OAuthInterceptor,
 		workspaceProviders: WorkspaceProvider[],
 	): DeploymentManager {
 		const manager = new DeploymentManager(
 			serviceContainer,
 			client,
 			oauthSessionManager,
-			oauthInterceptor,
 			workspaceProviders,
 		);
 		manager.subscribeToCrossWindowChanges();
@@ -140,7 +136,6 @@ export class DeploymentManager implements vscode.Disposable {
 		this.refreshWorkspaces();
 
 		await this.oauthSessionManager.setDeployment(deployment);
-		await this.oauthInterceptor.setDeployment(deployment);
 		await this.persistDeployment(deployment);
 	}
 
@@ -154,7 +149,6 @@ export class DeploymentManager implements vscode.Disposable {
 
 		this.client.setCredentials(undefined, undefined);
 		this.oauthSessionManager.clearDeployment();
-		this.oauthInterceptor.clearDeployment();
 		this.updateAuthContexts();
 		this.refreshWorkspaces();
 

src/extension.ts

@@ -7,14 +7,14 @@ import * as path from "node:path";
 import * as vscode from "vscode";
 
 import { errToStr } from "./api/api-helper";
+import { AuthInterceptor } from "./api/authInterceptor";
 import { CoderApi } from "./api/coderApi";
 import { Commands } from "./commands";
 import { ServiceContainer } from "./core/container";
 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";
@@ -88,15 +88,20 @@ 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(
+	// Handles 401 responses (OAuth and otherwise)
+	const authInterceptor = new AuthInterceptor(
 		client,
 		output,
 		oauthSessionManager,
 		secretsManager,
-		deployment?.safeHostname ?? "",
+		() => {
+			void vscode.window.showWarningMessage(
+				"Session expired. Please log in again using the Coder sidebar.",
+			);
+			return Promise.resolve(false);
+		},
 	);
-	ctx.subscriptions.push(oauthInterceptor);
+	ctx.subscriptions.push(authInterceptor);
 
 	const myWorkspacesProvider = new WorkspaceProvider(
 		WorkspaceQuery.Mine,
@@ -146,7 +151,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		serviceContainer,
 		client,
 		oauthSessionManager,
-		oauthInterceptor,
 		[myWorkspacesProvider, allWorkspacesProvider],
 	);
 	ctx.subscriptions.push(deploymentManager);