Better error handling on 400/401/403 due to auth

Author: EhabY

src/api/coderApi.ts

@@ -4,6 +4,7 @@ import {
 	type AxiosHeaders,
 	type AxiosResponseTransformer,
 	isAxiosError,
+	type AxiosError,
 } from "axios";
 import { Api } from "coder/site/src/api/api";
 import {
@@ -32,11 +33,7 @@ import {
 	HttpClientLogLevel,
 } from "../logging/types";
 import { sizeOf } from "../logging/utils";
-import {
-	parseOAuthError,
-	requiresReAuthentication,
-	isNetworkError,
-} from "../oauth/errors";
+import { parseOAuthError, requiresReAuthentication } from "../oauth/errors";
 import { type OAuthSessionManager } from "../oauth/sessionManager";
 import { type UnidirectionalStream } from "../websocket/eventStreamConnection";
 import {
@@ -421,7 +418,7 @@ function addLoggingInterceptors(client: AxiosInstance, logger: Logger) {
 /**
  * Add OAuth token refresh interceptors.
  * Success interceptor: proactively refreshes token when approaching expiry.
- * Error interceptor: reactively refreshes token on 401/403 responses.
+ * Error interceptor: reactively refreshes token on 401 responses.
  */
 function addOAuthInterceptors(
 	client: CoderApi,
@@ -444,53 +441,90 @@ function addOAuthInterceptors(
 
 			return response;
 		},
-		// Error response interceptor: reactive token refresh on 401/403
+		// Error response interceptor: reactive token refresh on 401
 		async (error: unknown) => {
 			if (!isAxiosError(error)) {
 				throw error;
 			}
 
-			const status = error.response?.status;
-			if (status !== 401 && status !== 403) {
-				throw error;
+			if (error.config) {
+				const config = error.config as {
+					_oauthRetryAttempted?: boolean;
+				};
+				if (config._oauthRetryAttempted) {
+					throw error;
+				}
 			}
 
-			if (!oauthSessionManager.isLoggedInWithOAuth()) {
+			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);
 			}
 
-			logger.info(`Received ${status} response, attempting token refresh`);
-
-			try {
-				const newTokens = await oauthSessionManager.refreshToken();
-				client.setSessionToken(newTokens.access_token);
-
-				logger.info("Token refresh successful, updated session token");
-			} catch (refreshError) {
-				logger.error("Token refresh failed:", refreshError);
-
-				const oauthError = parseOAuthError(refreshError);
-				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);
-						});
-				} else if (isNetworkError(refreshError)) {
-					logger.warn(
-						"Token refresh failed due to network error, will retry later",
-					);
-				}
-			}
 			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;
+	}
+}
+
 function wrapRequestTransform(
 	transformer: AxiosResponseTransformer | AxiosResponseTransformer[],
 	config: RequestConfigWithMeta,

src/oauth/errors.ts

@@ -164,10 +164,3 @@ export function requiresReAuthentication(error: OAuthError): boolean {
 		error instanceof InvalidGrantError || error instanceof InvalidClientError
 	);
 }
-
-/**
- * Checks if an error is a network/connectivity error
- */
-export function isNetworkError(error: unknown): boolean {
-	return isAxiosError(error) && !error.response && Boolean(error.request);
-}

src/oauth/sessionManager.ts

@@ -61,7 +61,7 @@ const DEFAULT_OAUTH_SCOPES = [
  */
 export class OAuthSessionManager implements vscode.Disposable {
 	private storedTokens: StoredOAuthTokens | undefined;
-	private refreshInProgress = false;
+	private refreshPromise: Promise<TokenResponse> | null = null;
 	private lastRefreshAttempt = 0;
 
 	private pendingAuthReject: ((reason: Error) => void) | undefined;
@@ -134,7 +134,7 @@ export class OAuthSessionManager implements vscode.Disposable {
 	 */
 	private async clearTokenState(): Promise<void> {
 		this.storedTokens = undefined;
-		this.refreshInProgress = false;
+		this.refreshPromise = null;
 		this.lastRefreshAttempt = 0;
 		await this.secretsManager.setOAuthTokens(undefined);
 		await this.secretsManager.setOAuthClientRegistration(undefined);
@@ -489,56 +489,64 @@ export class OAuthSessionManager implements vscode.Disposable {
 
 	/**
 	 * Refresh the access token using the stored refresh token.
-	 * Uses a mutex to prevent concurrent refresh attempts.
+	 * Uses a shared promise to handle concurrent refresh attempts.
 	 */
 	async refreshToken(): Promise<TokenResponse> {
-		if (this.refreshInProgress) {
-			throw new Error("Token refresh already in progress");
+		// If a refresh is already in progress, return the existing promise
+		if (this.refreshPromise) {
+			this.logger.debug(
+				"Token refresh already in progress, waiting for result",
+			);
+			return this.refreshPromise;
 		}
 
 		if (!this.storedTokens?.refresh_token) {
 			throw new Error("No refresh token available");
 		}
 
-		this.refreshInProgress = true;
+		const refreshToken = this.storedTokens.refresh_token;
+		const accessToken = this.storedTokens.access_token;
+
 		this.lastRefreshAttempt = Date.now();
 
-		try {
-			const { axiosInstance, metadata, registration } =
-				await this.prepareOAuthOperation(
-					this.deploymentUrl,
-					this.storedTokens.access_token,
-				);
+		// Create and store the refresh promise
+		this.refreshPromise = (async () => {
+			try {
+				const { axiosInstance, metadata, registration } =
+					await this.prepareOAuthOperation(this.deploymentUrl, accessToken);
 
-			this.logger.debug("Refreshing access token");
+				this.logger.debug("Refreshing access token");
 
-			const params: RefreshTokenRequestParams = {
-				grant_type: REFRESH_GRANT_TYPE,
-				refresh_token: this.storedTokens.refresh_token,
-				client_id: registration.client_id,
-				client_secret: registration.client_secret,
-			};
+				const params: RefreshTokenRequestParams = {
+					grant_type: REFRESH_GRANT_TYPE,
+					refresh_token: refreshToken,
+					client_id: registration.client_id,
+					client_secret: registration.client_secret,
+				};
 
-			const tokenRequest = toUrlSearchParams(params);
+				const tokenRequest = toUrlSearchParams(params);
 
-			const response = await axiosInstance.post<TokenResponse>(
-				metadata.token_endpoint,
-				tokenRequest,
-				{
-					headers: {
-						"Content-Type": "application/x-www-form-urlencoded",
+				const response = await axiosInstance.post<TokenResponse>(
+					metadata.token_endpoint,
+					tokenRequest,
+					{
+						headers: {
+							"Content-Type": "application/x-www-form-urlencoded",
+						},
 					},
-				},
-			);
+				);
 
-			this.logger.debug("Token refresh successful");
+				this.logger.debug("Token refresh successful");
 
-			await this.saveTokens(response.data);
+				await this.saveTokens(response.data);
 
-			return response.data;
-		} finally {
-			this.refreshInProgress = false;
-		}
+				return response.data;
+			} finally {
+				this.refreshPromise = null;
+			}
+		})();
+
+		return this.refreshPromise;
 	}
 
 	/**
@@ -579,7 +587,7 @@ export class OAuthSessionManager implements vscode.Disposable {
 		if (
 			!this.isLoggedInWithOAuth() ||
 			!this.storedTokens?.refresh_token ||
-			this.refreshInProgress
+			this.refreshPromise !== null
 		) {
 			return false;
 		}
@@ -695,7 +703,7 @@ export class OAuthSessionManager implements vscode.Disposable {
 		}
 		this.pendingAuthReject = undefined;
 		this.storedTokens = undefined;
-		this.refreshInProgress = false;
+		this.refreshPromise = null;
 		this.lastRefreshAttempt = 0;
 
 		this.logger.debug("OAuth session manager disposed");