Introduce ContextManager to hold context state globally + Rely on secrets to propagate authentication events between windows

Author: EhabY

src/commands.ts

@@ -12,6 +12,7 @@ import { CoderApi } from "./api/coderApi";
 import { needToken } from "./api/utils";
 import { type CliManager } from "./core/cliManager";
 import { type ServiceContainer } from "./core/container";
+import { type ContextManager } from "./core/contextManager";
 import { type MementoManager } from "./core/mementoManager";
 import { type PathResolver } from "./core/pathResolver";
 import { type SecretsManager } from "./core/secretsManager";
@@ -32,6 +33,7 @@ export class Commands {
 	private readonly mementoManager: MementoManager;
 	private readonly secretsManager: SecretsManager;
 	private readonly cliManager: CliManager;
+	private readonly contextManager: ContextManager;
 	// These will only be populated when actively connected to a workspace and are
 	// used in commands.  Because commands can be executed by the user, it is not
 	// possible to pass in arguments, so we have to store the current workspace
@@ -53,6 +55,7 @@ export class Commands {
 		this.mementoManager = serviceContainer.getMementoManager();
 		this.secretsManager = serviceContainer.getSecretsManager();
 		this.cliManager = serviceContainer.getCliManager();
+		this.contextManager = serviceContainer.getContextManager();
 	}
 
 	/**
@@ -189,6 +192,11 @@ export class Commands {
 		label?: string;
 		autoLogin?: boolean;
 	}): Promise<void> {
+		if (this.contextManager.get("coder.authenticated")) {
+			return;
+		}
+		this.logger.info("Logging in");
+
 		const url = await this.maybeAskUrl(args?.url);
 		if (!url) {
 			return; // The user aborted.
@@ -219,13 +227,9 @@ export class Commands {
 		await this.secretsManager.setSessionToken(res.token);
 
 		// These contexts control various menu items and the sidebar.
-		await vscode.commands.executeCommand(
-			"setContext",
-			"coder.authenticated",
-			true,
-		);
+		this.contextManager.set("coder.authenticated", true);
 		if (res.user.roles.find((role) => role.name === "owner")) {
-			await vscode.commands.executeCommand("setContext", "coder.isOwner", true);
+			this.contextManager.set("coder.isOwner", true);
 		}
 
 		vscode.window
@@ -245,6 +249,7 @@ export class Commands {
 
 		// Fetch workspaces for the new deployment.
 		vscode.commands.executeCommand("coder.refreshWorkspaces");
+		this.secretsManager.triggerLoginStateChange("login");
 	}
 
 	/**
@@ -376,6 +381,10 @@ export class Commands {
 	}
 
 	public async forceLogout(): Promise<void> {
+		if (!this.contextManager.get("coder.authenticated")) {
+			return;
+		}
+		this.logger.info("Logging out");
 		// Clear from the REST client.  An empty url will indicate to other parts of
 		// the code that we are logged out.
 		this.restClient.setHost("");
@@ -385,11 +394,7 @@ export class Commands {
 		await this.mementoManager.setUrl(undefined);
 		await this.secretsManager.setSessionToken(undefined);
 
-		await vscode.commands.executeCommand(
-			"setContext",
-			"coder.authenticated",
-			false,
-		);
+		this.contextManager.set("coder.authenticated", false);
 		vscode.window
 			.showInformationMessage("You've been logged out of Coder!", "Login")
 			.then((action) => {
@@ -400,6 +405,7 @@ export class Commands {
 
 		// This will result in clearing the workspace list.
 		vscode.commands.executeCommand("coder.refreshWorkspaces");
+		this.secretsManager.triggerLoginStateChange("logout");
 	}
 
 	/**

src/core/container.ts

@@ -3,6 +3,7 @@ import * as vscode from "vscode";
 import { type Logger } from "../logging/logger";
 
 import { CliManager } from "./cliManager";
+import { ContextManager } from "./contextManager";
 import { MementoManager } from "./mementoManager";
 import { PathResolver } from "./pathResolver";
 import { SecretsManager } from "./secretsManager";
@@ -17,6 +18,7 @@ export class ServiceContainer implements vscode.Disposable {
 	private readonly mementoManager: MementoManager;
 	private readonly secretsManager: SecretsManager;
 	private readonly cliManager: CliManager;
+	private readonly contextManager: ContextManager;
 
 	constructor(
 		context: vscode.ExtensionContext,
@@ -34,6 +36,7 @@ export class ServiceContainer implements vscode.Disposable {
 			this.logger,
 			this.pathResolver,
 		);
+		this.contextManager = new ContextManager();
 	}
 
 	getVsCodeProposed(): typeof vscode {
@@ -60,10 +63,15 @@ export class ServiceContainer implements vscode.Disposable {
 		return this.cliManager;
 	}
 
+	getContextManager(): ContextManager {
+		return this.contextManager;
+	}
+
 	/**
 	 * Dispose of all services and clean up resources.
 	 */
 	dispose(): void {
+		this.contextManager.dispose();
 		this.logger.dispose();
 	}
 }

src/core/contextManager.ts

@@ -0,0 +1,33 @@
+import * as vscode from "vscode";
+
+const CONTEXT_DEFAULTS = {
+	"coder.authenticated": false,
+	"coder.isOwner": false,
+	"coder.loaded": false,
+	"coder.workspace.updatable": false,
+} as const;
+
+type CoderContext = keyof typeof CONTEXT_DEFAULTS;
+
+export class ContextManager implements vscode.Disposable {
+	private readonly context = new Map<CoderContext, boolean>();
+
+	public constructor() {
+		(Object.keys(CONTEXT_DEFAULTS) as CoderContext[]).forEach((key) => {
+			this.set(key, CONTEXT_DEFAULTS[key]);
+		});
+	}
+
+	public set(key: CoderContext, value: boolean): void {
+		this.context.set(key, value);
+		vscode.commands.executeCommand("setContext", key, value);
+	}
+
+	public get(key: CoderContext): boolean {
+		return this.context.get(key) ?? CONTEXT_DEFAULTS[key];
+	}
+
+	public dispose() {
+		this.context.clear();
+	}
+}

src/core/secretsManager.ts

@@ -2,8 +2,13 @@ import type { SecretStorage, Disposable } from "vscode";
 
 const SESSION_TOKEN_KEY = "sessionToken";
 
+const LOGIN_STATE_KEY = "loginState";
+
+type AuthAction = "login" | "logout";
 export class SecretsManager {
-	constructor(private readonly secrets: SecretStorage) {}
+	constructor(private readonly secrets: SecretStorage) {
+		void this.secrets.delete(LOGIN_STATE_KEY);
+	}
 
 	/**
 	 * Set or unset the last used token.
@@ -29,13 +34,17 @@ export class SecretsManager {
 		}
 	}
 
-	/**
-	 * Subscribe to changes to the session token which can be used to indicate user login status.
-	 */
-	public onDidChangeSessionToken(listener: () => Promise<void>): Disposable {
-		return this.secrets.onDidChange((e) => {
-			if (e.key === SESSION_TOKEN_KEY) {
-				listener();
+	public triggerLoginStateChange(action: AuthAction): void {
+		this.secrets.store(LOGIN_STATE_KEY, action);
+	}
+
+	public onDidChangeLoginState(
+		listener: (state?: AuthAction) => Promise<void>,
+	): Disposable {
+		return this.secrets.onDidChange(async (e) => {
+			if (e.key === LOGIN_STATE_KEY) {
+				const state = await this.secrets.get(LOGIN_STATE_KEY);
+				listener(state as AuthAction | undefined);
 			}
 		});
 	}

src/extension.ts

@@ -62,6 +62,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	const output = serviceContainer.getLogger();
 	const mementoManager = serviceContainer.getMementoManager();
 	const secretsManager = serviceContainer.getSecretsManager();
+	const contextManager = serviceContainer.getContextManager();
 
 	// Try to clear this flag ASAP
 	const isFirstConnect = await mementoManager.getAndClearFirstConnect();
@@ -331,18 +332,21 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	const remote = new Remote(serviceContainer, commands, ctx.extensionMode);
 
 	ctx.subscriptions.push(
-		secretsManager.onDidChangeSessionToken(async () => {
-			const token = await secretsManager.getSessionToken();
-			const url = mementoManager.getUrl();
-			if (!token) {
-				output.info("Logging out");
-				await commands.forceLogout();
-			} else if (url) {
-				output.info("Logging in");
+		secretsManager.onDidChangeLoginState(async (state) => {
+			if (state === undefined) {
+				// Initalization - Ignore those events
+				return;
+			}
+
+			if (state === "login") {
+				const token = await secretsManager.getSessionToken();
+				const url = mementoManager.getUrl();
 				// Should login the user directly if the URL+Token are valid
 				await commands.login({ url, token });
 				// Resolve any pending login detection promises
 				remote.resolveLoginDetected();
+			} else {
+				await commands.forceLogout();
 			}
 		}),
 	);
@@ -413,20 +417,12 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		output.info(`Logged in to ${baseUrl}; checking credentials`);
 		client
 			.getAuthenticatedUser()
-			.then(async (user) => {
+			.then((user) => {
 				if (user && user.roles) {
 					output.info("Credentials are valid");
-					vscode.commands.executeCommand(
-						"setContext",
-						"coder.authenticated",
-						true,
-					);
+					contextManager.set("coder.authenticated", true);
 					if (user.roles.find((role) => role.name === "owner")) {
-						await vscode.commands.executeCommand(
-							"setContext",
-							"coder.isOwner",
-							true,
-						);
+						contextManager.set("coder.isOwner", true);
 					}
 
 					// Fetch and monitor workspaces, now that we know the client is good.
@@ -445,11 +441,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 				);
 			})
 			.finally(() => {
-				vscode.commands.executeCommand("setContext", "coder.loaded", true);
+				contextManager.set("coder.loaded", true);
 			});
 	} else {
 		output.info("Not currently logged in");
-		vscode.commands.executeCommand("setContext", "coder.loaded", true);
+		contextManager.set("coder.loaded", true);
 
 		// Handle autologin, if not already logged in.
 		const cfg = vscode.workspace.getConfiguration();

src/remote/remote.ts

@@ -30,6 +30,7 @@ import { type Commands } from "../commands";
 import { type CliManager } from "../core/cliManager";
 import * as cliUtils from "../core/cliUtils";
 import { type ServiceContainer } from "../core/container";
+import { type ContextManager } from "../core/contextManager";
 import { type PathResolver } from "../core/pathResolver";
 import { featureSetForVersion, type FeatureSet } from "../featureSet";
 import { getGlobalFlags } from "../globalFlags";
@@ -58,6 +59,7 @@ export class Remote {
 	private readonly logger: Logger;
 	private readonly pathResolver: PathResolver;
 	private readonly cliManager: CliManager;
+	private readonly contextManager: ContextManager;
 
 	// Used to race between the login dialog and the logging in from a different window
 	private loginDetectedResolver: (() => void) | undefined;
@@ -73,6 +75,7 @@ export class Remote {
 		this.logger = serviceContainer.getLogger();
 		this.pathResolver = serviceContainer.getPathResolver();
 		this.cliManager = serviceContainer.getCliManager();
+		this.contextManager = serviceContainer.getContextManager();
 	}
 
 	/**
@@ -545,6 +548,7 @@ export class Remote {
 				workspaceClient,
 				this.logger,
 				this.vscodeProposed,
+				this.contextManager,
 			);
 			disposables.push(monitor);
 			disposables.push(

src/workspace/workspaceMonitor.ts

@@ -7,6 +7,7 @@ import * as vscode from "vscode";
 
 import { createWorkspaceIdentifier, errToStr } from "../api/api-helper";
 import { type CoderApi } from "../api/coderApi";
+import { type ContextManager } from "../core/contextManager";
 import { type Logger } from "../logging/logger";
 import { type OneWayWebSocket } from "../websocket/oneWayWebSocket";
 
@@ -41,6 +42,7 @@ export class WorkspaceMonitor implements vscode.Disposable {
 		private readonly logger: Logger,
 		// We use the proposed API to get access to useCustom in dialogs.
 		private readonly vscodeProposed: typeof vscode,
+		private readonly contextManager: ContextManager,
 	) {
 		this.name = createWorkspaceIdentifier(workspace);
 		const socket = this.client.watchWorkspace(workspace);
@@ -217,11 +219,7 @@ export class WorkspaceMonitor implements vscode.Disposable {
 	}
 
 	private updateContext(workspace: Workspace) {
-		vscode.commands.executeCommand(
-			"setContext",
-			"coder.workspace.updatable",
-			workspace.outdated,
-		);
+		this.contextManager.set("coder.workspace.updatable", workspace.outdated);
 	}
 
 	private updateStatusBar(workspace: Workspace) {