Allow callback handling in multiple VS Code windows

Author: EhabY

src/core/secretsManager.ts

@@ -13,11 +13,19 @@ const OAUTH_CLIENT_REGISTRATION_KEY = "oauthClientRegistration";
 
 const OAUTH_TOKENS_KEY = "oauthTokens";
 
+const OAUTH_CALLBACK_KEY = "coder.oauthCallback";
+
 export type StoredOAuthTokens = Omit<TokenResponse, "expires_in"> & {
 	expiry_timestamp: number;
 	deployment_url: string;
 };
 
+interface OAuthCallbackData {
+	state: string;
+	code: string | null;
+	error: string | null;
+}
+
 export enum AuthAction {
 	LOGIN,
 	LOGOUT,
@@ -163,4 +171,36 @@ export class SecretsManager {
 		}
 		return undefined;
 	}
+
+	/**
+	 * 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
+			}
+		});
+	}
 }

src/extension.ts

@@ -163,7 +163,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 				const code = params.get("code");
 				const state = params.get("state");
 				const error = params.get("error");
-				oauthSessionManager.handleCallback(code, state, error);
+				await oauthSessionManager.handleCallback(code, state, error);
 				return;
 			}
 

src/oauth/sessionManager.ts

@@ -64,13 +64,7 @@ export class OAuthSessionManager implements vscode.Disposable {
 	private refreshInProgress = false;
 	private lastRefreshAttempt = 0;
 
-	// Pending authorization flow state
-	private pendingAuthResolve:
-		| ((value: { code: string; verifier: string }) => void)
-		| undefined;
 	private pendingAuthReject: ((reason: Error) => void) | undefined;
-	private expectedState: string | undefined;
-	private pendingVerifier: string | undefined;
 
 	/**
 	 * Create and initialize a new OAuth session manager.
@@ -370,106 +364,80 @@ export class OAuthSessionManager implements vscode.Disposable {
 			challenge,
 		);
 
-		return new Promise<{ code: string; verifier: string }>(
+		const callbackPromise = new Promise<{ code: string; verifier: string }>(
 			(resolve, reject) => {
 				const timeoutMins = 5;
-				const timeout = setTimeout(
+				const timeoutHandle = setTimeout(
 					() => {
-						this.clearPendingAuth();
+						cleanup();
 						reject(
 							new Error(`OAuth flow timed out after ${timeoutMins} minutes`),
 						);
 					},
 					timeoutMins * 60 * 1000,
 				);
 
-				const clearPromise = () => {
-					clearTimeout(timeout);
-					this.clearPendingAuth();
-				};
-
-				this.pendingAuthResolve = (result) => {
-					clearPromise();
-					resolve(result);
-				};
-
-				this.pendingAuthReject = (error) => {
-					clearPromise();
-					reject(error);
-				};
+				const listener = this.secretsManager.onDidChangeOAuthCallback(
+					({ state: callbackState, code, error }) => {
+						if (callbackState !== state) {
+							return;
+						}
 
-				this.expectedState = state;
-				this.pendingVerifier = verifier;
+						cleanup();
 
-				vscode.env.openExternal(vscode.Uri.parse(authUrl)).then(
-					() => {},
-					(error) => {
-						if (error instanceof Error) {
-							this.pendingAuthReject?.(error);
+						if (error) {
+							reject(new Error(`OAuth error: ${error}`));
+						} else if (code) {
+							resolve({ code, verifier });
 						} else {
-							this.pendingAuthReject?.(new Error("Failed to open browser"));
+							reject(new Error("No authorization code received"));
 						}
 					},
 				);
+
+				const cleanup = () => {
+					clearTimeout(timeoutHandle);
+					listener.dispose();
+				};
+
+				this.pendingAuthReject = (error) => {
+					cleanup();
+					reject(error);
+				};
 			},
 		);
-	}
 
-	/**
-	 * Clear pending authorization flow state.
-	 */
-	private clearPendingAuth(): void {
-		this.pendingAuthResolve = undefined;
-		this.pendingAuthReject = undefined;
-		this.expectedState = undefined;
-		this.pendingVerifier = undefined;
+		try {
+			await vscode.env.openExternal(vscode.Uri.parse(authUrl));
+		} catch (error) {
+			throw error instanceof Error
+				? error
+				: new Error("Failed to open browser");
+		}
+
+		return callbackPromise;
 	}
 
 	/**
 	 * Handle OAuth callback from browser redirect.
-	 * Validates state and resolves pending authorization promise.
-	 *
-	 * // TODO this has to work across windows!
+	 * Writes the callback result to secrets storage, triggering the waiting window to proceed.
 	 */
-	handleCallback(
+	async handleCallback(
 		code: string | null,
 		state: string | null,
 		error: string | null,
-	): void {
-		if (!this.pendingAuthResolve || !this.pendingAuthReject) {
-			this.logger.warn("Received OAuth callback but no pending auth flow");
-			return;
-		}
-
-		if (error) {
-			this.pendingAuthReject(new Error(`OAuth error: ${error}`));
-			return;
-		}
-
-		if (!code) {
-			this.pendingAuthReject(new Error("No authorization code received"));
-			return;
-		}
-
+	): Promise<void> {
 		if (!state) {
-			this.pendingAuthReject(new Error("No state received"));
-			return;
-		}
-
-		if (state !== this.expectedState) {
-			this.pendingAuthReject(
-				new Error("State mismatch - possible CSRF attack"),
-			);
+			this.logger.warn("Received OAuth callback with no state parameter");
 			return;
 		}
 
-		const verifier = this.pendingVerifier;
-		if (!verifier) {
-			this.pendingAuthReject(new Error("No PKCE verifier found"));
-			return;
+		try {
+			await this.secretsManager.setOAuthCallback({ state, code, error });
+			this.logger.debug("OAuth callback processed successfully");
+		} catch (err) {
+			this.logger.error("Failed to process OAuth callback:", err);
 		}
-
-		this.pendingAuthResolve({ code, verifier });
 	}
 
 	/**
@@ -712,13 +680,13 @@ export class OAuthSessionManager implements vscode.Disposable {
 	}
 
 	/**
-	 * Clears all in-memory state and rejects any pending operations.
+	 * Clears all in-memory state.
 	 */
 	dispose(): void {
 		if (this.pendingAuthReject) {
 			this.pendingAuthReject(new Error("OAuth session manager disposed"));
 		}
-		this.clearPendingAuth();
+		this.pendingAuthReject = undefined;
 		this.storedTokens = undefined;
 		this.refreshInProgress = false;
 		this.lastRefreshAttempt = 0;