refactor: migrate OAuth callback to WindowBroadcast

Author: EhabY

src/core/secretsManager.ts

@@ -2,6 +2,7 @@ import { z } from "zod";
 
 import { DeploymentSchema, type Deployment } from "../deployment/types";
 import { toSafeHost } from "../util";
+import { WindowBroadcast } from "../windowBroadcast";
 
 import type { OAuth2ClientRegistrationResponse } from "coder/site/src/api/typesGenerated";
 import type { Memento, SecretStorage, Disposable } from "vscode";
@@ -57,14 +58,24 @@ const OAuthCallbackDataSchema = z.object({
 	error: z.string().nullable(),
 });
 
-type OAuthCallbackData = z.infer<typeof OAuthCallbackDataSchema>;
+export type OAuthCallbackData = z.infer<typeof OAuthCallbackDataSchema>;
 
 export class SecretsManager {
+	public readonly oauthCallback: WindowBroadcast<OAuthCallbackData>;
+
 	constructor(
 		private readonly secrets: SecretStorage,
 		private readonly memento: Memento,
 		private readonly logger: Logger,
-	) {}
+	) {
+		this.oauthCallback = new WindowBroadcast(
+			secrets,
+			OAUTH_CALLBACK_KEY,
+			(v: unknown): v is OAuthCallbackData =>
+				OAuthCallbackDataSchema.safeParse(v).success,
+			logger,
+		);
+	}
 
 	private buildKey(prefix: SecretKeyPrefix, safeHostname: string): string {
 		return `${prefix}${safeHostname || "<legacy>"}`;
@@ -306,54 +317,6 @@ 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> {
-		const parsed = OAuthCallbackDataSchema.parse(data);
-		await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(parsed));
-	}
-
-	/**
-	 * 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;
-			}
-
-			const raw = await this.secrets.get(OAUTH_CALLBACK_KEY);
-			if (!raw) {
-				return;
-			}
-
-			let parsed: unknown;
-			try {
-				parsed = JSON.parse(raw);
-			} catch (err) {
-				this.logger.error("Failed to parse OAuth callback JSON", err);
-				return;
-			}
-
-			const result = OAuthCallbackDataSchema.safeParse(parsed);
-			if (!result.success) {
-				this.logger.error("Invalid OAuth callback data shape", result.error);
-				return;
-			}
-
-			try {
-				listener(result.data);
-			} catch (err) {
-				this.logger.error("Error in onDidChangeOAuthCallback listener", err);
-			}
-		});
-	}
-
 	public getOAuthClientRegistration(
 		safeHostname: string,
 	): Promise<OAuth2ClientRegistrationResponse | undefined> {

src/oauth/authorizer.ts

@@ -247,7 +247,7 @@ export class OAuthAuthorizer implements vscode.Disposable {
 					timeoutMins * 60 * 1000,
 				);
 
-				const listener = this.secretsManager.onDidChangeOAuthCallback(
+				const listener = this.secretsManager.oauthCallback.onReceive(
 					({ state: callbackState, code, error }) => {
 						if (callbackState !== state) {
 							this.logger.warn(

src/uri/uriHandler.ts

@@ -211,7 +211,7 @@ async function handleOAuthCallback(ctx: UriRouteContext): Promise<void> {
 	}
 
 	try {
-		await secretsManager.setOAuthCallback({ state, code, error });
+		await secretsManager.oauthCallback.send({ state, code, error });
 		logger.debug("OAuth callback processed successfully");
 	} catch (err) {
 		logger.error("Failed to process OAuth callback:", err);

test/unit/oauth/authorizer.test.ts

@@ -83,7 +83,7 @@ function createTestContext() {
 
 	/** Completes login by sending successful OAuth callback */
 	const completeLogin = async (state: string) => {
-		await base.secretsManager.setOAuthCallback({
+		await base.secretsManager.oauthCallback.send({
 			state,
 			code: "code",
 			error: null,
@@ -138,7 +138,7 @@ describe("OAuthAuthorizer", () => {
 			const { state } = await waitForBrowserToOpen();
 
 			// Set the callback with the correct state (simulate user clicking authorize)
-			await secretsManager.setOAuthCallback({
+			await secretsManager.oauthCallback.send({
 				state,
 				code: "auth-code-123",
 				error: null,
@@ -185,7 +185,7 @@ describe("OAuthAuthorizer", () => {
 			const { authUrl, state } = await waitForBrowserToOpen();
 			expect(authUrl.searchParams.get("client_id")).toBe("existing-client-id");
 
-			await secretsManager.setOAuthCallback({
+			await secretsManager.oauthCallback.send({
 				state,
 				code: "code",
 				error: null,
@@ -225,7 +225,7 @@ describe("OAuthAuthorizer", () => {
 			const { authUrl, state } = await waitForBrowserToOpen();
 			expect(authUrl.searchParams.get("client_id")).toBe("new-client-id");
 
-			await secretsManager.setOAuthCallback({
+			await secretsManager.oauthCallback.send({
 				state,
 				code: "code",
 				error: null,
@@ -267,7 +267,7 @@ describe("OAuthAuthorizer", () => {
 			const { loginPromise, state } = await startLogin();
 
 			// Send callback with wrong state - should be ignored
-			await secretsManager.setOAuthCallback({
+			await secretsManager.oauthCallback.send({
 				state: "wrong-state",
 				code: "code",
 				error: null,
@@ -292,7 +292,7 @@ describe("OAuthAuthorizer", () => {
 			setupOAuthRoutes();
 
 			const { loginPromise, state } = await startLogin();
-			await secretsManager.setOAuthCallback({
+			await secretsManager.oauthCallback.send({
 				state,
 				code: null,
 				error: "access_denied",
@@ -307,7 +307,11 @@ describe("OAuthAuthorizer", () => {
 			setupOAuthRoutes();
 
 			const { loginPromise, state } = await startLogin();
-			await secretsManager.setOAuthCallback({ state, code: null, error: null });
+			await secretsManager.oauthCallback.send({
+				state,
+				code: null,
+				error: null,
+			});
 
 			await expect(loginPromise).rejects.toThrow(
 				"No authorization code received",

test/unit/uri/uriHandler.test.ts

@@ -361,7 +361,7 @@ describe("uriHandler", () => {
 			const { handleUri, secretsManager } = createTestContext();
 
 			const callbackPromise = new Promise<CallbackData>((resolve) => {
-				secretsManager.onDidChangeOAuthCallback(resolve);
+				secretsManager.oauthCallback.onReceive(resolve);
 			});
 
 			await handleUri(
@@ -380,7 +380,7 @@ describe("uriHandler", () => {
 			const { handleUri, secretsManager } = createTestContext();
 
 			const callbackPromise = new Promise<CallbackData>((resolve) => {
-				secretsManager.onDidChangeOAuthCallback(resolve);
+				secretsManager.oauthCallback.onReceive(resolve);
 			});
 
 			await handleUri(
@@ -399,7 +399,7 @@ describe("uriHandler", () => {
 			const { handleUri, secretsManager } = createTestContext();
 
 			let callbackReceived = false;
-			secretsManager.onDidChangeOAuthCallback(() => {
+			secretsManager.oauthCallback.onReceive(() => {
 				callbackReceived = true;
 			});