Unify login experience

Author: EhabY

src/commands.ts

@@ -9,6 +9,7 @@ import { createWorkspaceIdentifier, extractAgents } from "./api/api-helper";
 import { type CliManager } from "./core/cliManager";
 import { type ServiceContainer } from "./core/container";
 import { type ContextManager } from "./core/contextManager";
+import { type Deployment } from "./core/deployment";
 import { type MementoManager } from "./core/mementoManager";
 import { type PathResolver } from "./core/pathResolver";
 import { type SecretsManager } from "./core/secretsManager";
@@ -61,6 +62,17 @@ export class Commands {
 		this.loginCoordinator = serviceContainer.getLoginCoordinator();
 	}
 
+	/**
+	 * Get the current deployment, throwing if not logged in.
+	 */
+	private async requireDeployment(): Promise<Deployment> {
+		const deployment = await this.secretsManager.getCurrentDeployment();
+		if (!deployment) {
+			throw new Error("You are not logged in");
+		}
+		return deployment;
+	}
+
 	/**
 	 * Log into the provided deployment. If the deployment URL is not specified,
 	 * ask for it first with a menu showing recent URLs along with the default URL
@@ -76,7 +88,12 @@ export class Commands {
 		}
 		this.logger.info("Logging in");
 
-		const url = await maybeAskUrl(this.mementoManager, args?.url);
+		const currentDeployment = await this.secretsManager.getCurrentDeployment();
+		const url = await maybeAskUrl(
+			this.mementoManager,
+			args?.url,
+			currentDeployment?.url,
+		);
 		if (!url) {
 			return;
 		}
@@ -98,16 +115,13 @@ export class Commands {
 			return;
 		}
 
-		// Authorize the global client
+		// Set client immediately so subsequent operations in this function have the correct host/token.
+		// The cross-window listener will also update the client, but that's async.
 		this.restClient.setHost(url);
 		this.restClient.setSessionToken(result.token);
 
-		// Store for later sessions
-		await this.mementoManager.setUrl(url);
-		await this.secretsManager.setSessionAuth(label, {
-			url,
-			token: result.token,
-		});
+		// Set as current deployment (triggers cross-window sync).
+		await this.secretsManager.setCurrentDeployment({ url, label });
 
 		// Update contexts
 		this.contextManager.set("coder.authenticated", true);
@@ -130,7 +144,6 @@ export class Commands {
 				}
 			});
 
-		await this.secretsManager.triggerLoginStateChange(label, "login");
 		vscode.commands.executeCommand("coder.refreshWorkspaces");
 	}
 
@@ -163,13 +176,8 @@ export class Commands {
 	 * Log out from the currently logged-in deployment.
 	 */
 	public async logout(): Promise<void> {
-		const url = this.mementoManager.getUrl();
-		if (!url) {
-			// Sanity check; command should not be available if no url.
-			throw new Error("You are not logged in");
-		}
-
-		await this.forceLogout(toSafeHost(url));
+		const deployment = await this.requireDeployment();
+		await this.forceLogout(deployment.label);
 	}
 
 	public async forceLogout(label: string): Promise<void> {
@@ -178,8 +186,7 @@ export class Commands {
 		}
 		this.logger.info(`Logging out of deployment: ${label}`);
 
-		// Only clear REST client and UI context if logging out of current deployment
-		// Fire and forget
+		// Fire and forget OAuth logout
 		this.oauthSessionManager.logout().catch((error) => {
 			this.logger.warn("OAuth logout failed, continuing with cleanup:", error);
 		});
@@ -189,8 +196,10 @@ export class Commands {
 		this.restClient.setHost("");
 		this.restClient.setSessionToken("");
 
-		// Clear from memory.
-		await this.mementoManager.setUrl(undefined);
+		// Clear current deployment (triggers cross-window sync)
+		await this.secretsManager.setCurrentDeployment(undefined);
+
+		// Clear all auth data for this deployment
 		await this.secretsManager.clearAllAuthData(label);
 
 		this.contextManager.set("coder.authenticated", false);
@@ -204,8 +213,6 @@ export class Commands {
 
 		// This will result in clearing the workspace list.
 		vscode.commands.executeCommand("coder.refreshWorkspaces");
-
-		await this.secretsManager.triggerLoginStateChange(label, "logout");
 	}
 
 	/**
@@ -214,7 +221,8 @@ export class Commands {
 	 * Must only be called if currently logged in.
 	 */
 	public async createWorkspace(): Promise<void> {
-		const uri = this.mementoManager.getUrl() + "/templates";
+		const deployment = await this.requireDeployment();
+		const uri = deployment.url + "/templates";
 		await vscode.commands.executeCommand("vscode.open", uri);
 	}
 
@@ -228,8 +236,9 @@ export class Commands {
 	 */
 	public async navigateToWorkspace(item: OpenableTreeItem) {
 		if (item) {
+			const deployment = await this.requireDeployment();
 			const workspaceId = createWorkspaceIdentifier(item.workspace);
-			const uri = this.mementoManager.getUrl() + `/@${workspaceId}`;
+			const uri = deployment.url + `/@${workspaceId}`;
 			await vscode.commands.executeCommand("vscode.open", uri);
 		} else if (this.workspace && this.workspaceRestClient) {
 			const baseUrl =
@@ -251,8 +260,9 @@ export class Commands {
 	 */
 	public async navigateToWorkspaceSettings(item: OpenableTreeItem) {
 		if (item) {
+			const deployment = await this.requireDeployment();
 			const workspaceId = createWorkspaceIdentifier(item.workspace);
-			const uri = this.mementoManager.getUrl() + `/@${workspaceId}/settings`;
+			const uri = deployment.url + `/@${workspaceId}/settings`;
 			await vscode.commands.executeCommand("vscode.open", uri);
 		} else if (this.workspace && this.workspaceRestClient) {
 			const baseUrl =
@@ -329,18 +339,14 @@ export class Commands {
 					const terminal = vscode.window.createTerminal(app.name);
 
 					// If workspace_name is provided, run coder ssh before the command
-
-					const url = this.mementoManager.getUrl();
-					if (!url) {
-						throw new Error("No coder url found for sidebar");
-					}
+					const deployment = await this.requireDeployment();
 					const binary = await this.cliManager.fetchBinary(
 						this.restClient,
-						toSafeHost(url),
+						deployment.label,
 					);
 
 					const configDir = this.pathResolver.getGlobalConfigDir(
-						toSafeHost(url),
+						deployment.label,
 					);
 					const globalFlags = getGlobalFlags(
 						vscode.workspace.getConfiguration(),

src/core/mementoManager.ts

@@ -7,27 +7,16 @@ export class MementoManager {
 	constructor(private readonly memento: Memento) {}
 
 	/**
-	 * Add the URL to the list of recently accessed URLs in global storage, then
-	 * set it as the last used URL.
-	 *
-	 * If the URL is falsey, then remove it as the last used URL and do not touch
-	 * the history.
+	 * Add a URL to the history of recently accessed URLs.
+	 * Used by the URL picker to show recent deployments.
 	 */
-	public async setUrl(url: string | undefined): Promise<void> {
-		await this.memento.update("url", url);
+	public async addToUrlHistory(url: string): Promise<void> {
 		if (url) {
 			const history = this.withUrlHistory(url);
 			await this.memento.update("urlHistory", history);
 		}
 	}
 
-	/**
-	 * Get the last used URL.
-	 */
-	public getUrl(): string | undefined {
-		return this.memento.get("url");
-	}
-
 	/**
 	 * Get the most recently accessed URLs (oldest to newest) with the provided
 	 * values appended. Duplicates will be removed.

src/core/secretsManager.ts

@@ -1,15 +1,16 @@
-import {
-	type TokenResponse,
-	type ClientRegistrationResponse,
-} from "../oauth/types";
+import { toSafeHost } from "../util";
 
 import type { Memento, SecretStorage, Disposable } from "vscode";
 
+import type { TokenResponse, ClientRegistrationResponse } from "../oauth/types";
+
+import type { Deployment } from "./deployment";
+
 const SESSION_KEY_PREFIX = "coder.session.";
 const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens.";
 const OAUTH_CLIENT_PREFIX = "coder.oauth.client.";
 
-const LOGIN_STATE_KEY = "coder.loginState";
+const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment";
 const OAUTH_CALLBACK_KEY = "coder.oauthCallback";
 
 const KNOWN_LABELS_KEY = "coder.knownLabels";
@@ -32,10 +33,8 @@ interface OAuthCallbackData {
 	error: string | null;
 }
 
-export enum AuthAction {
-	LOGIN,
-	LOGOUT,
-	INVALID,
+export interface CurrentDeploymentState {
+	deployment: Deployment | null;
 }
 
 export class SecretsManager {
@@ -45,54 +44,57 @@ export class SecretsManager {
 	) {}
 
 	/**
-	 * Triggers a login/logout event that propagates across all VS Code windows.
+	 * Sets the current deployment and triggers a cross-window sync event.
+	 * This is the single source of truth for which deployment is currently active.
 	 */
-	public async triggerLoginStateChange(
-		label: string,
-		action: "login" | "logout",
+	public async setCurrentDeployment(
+		deployment: Deployment | undefined,
 	): Promise<void> {
-		const loginState = {
-			action,
-			label,
+		const state = {
+			deployment: deployment ?? null,
 			timestamp: new Date().toISOString(),
 		};
-		await this.secrets.store(LOGIN_STATE_KEY, JSON.stringify(loginState));
+		await this.secrets.store(CURRENT_DEPLOYMENT_KEY, JSON.stringify(state));
 	}
 
 	/**
-	 * Listens for login/logout events from any VS Code window.
+	 * Gets the current deployment from storage.
 	 */
-	public onDidChangeLoginState(
-		listener: (state: AuthAction, label: string) => Promise<void>,
-	): Disposable {
-		return this.secrets.onDidChange(async (e) => {
-			if (e.key !== LOGIN_STATE_KEY) {
-				return;
+	public async getCurrentDeployment(): Promise<Deployment | undefined> {
+		try {
+			const data = await this.secrets.get(CURRENT_DEPLOYMENT_KEY);
+			if (!data) {
+				return undefined;
 			}
+			const parsed = JSON.parse(data) as { deployment: Deployment | null };
+			return parsed.deployment ?? undefined;
+		} catch {
+			return undefined;
+		}
+	}
 
-			const stateStr = await this.secrets.get(LOGIN_STATE_KEY);
-			if (!stateStr) {
-				await listener(AuthAction.INVALID, "");
+	/**
+	 * Listens for deployment changes from any VS Code window.
+	 * Fires when login, logout, or deployment switch occurs.
+	 */
+	public onDidChangeCurrentDeployment(
+		listener: (state: CurrentDeploymentState) => void | Promise<void>,
+	): Disposable {
+		return this.secrets.onDidChange(async (e) => {
+			if (e.key !== CURRENT_DEPLOYMENT_KEY) {
 				return;
 			}
 
 			try {
-				const parsed = JSON.parse(stateStr) as {
-					action: string;
-					label: string;
-					timestamp: string;
-				};
-
-				if (parsed.action === "login") {
-					await listener(AuthAction.LOGIN, parsed.label);
-				} else if (parsed.action === "logout") {
-					await listener(AuthAction.LOGOUT, parsed.label);
-				} else {
-					await listener(AuthAction.INVALID, parsed.label);
+				const data = await this.secrets.get(CURRENT_DEPLOYMENT_KEY);
+				if (data) {
+					const parsed = JSON.parse(data) as {
+						deployment: Deployment | null;
+					};
+					await listener({ deployment: parsed.deployment });
 				}
 			} catch {
-				// Invalid JSON, treat as invalid state
-				await listener(AuthAction.INVALID, "");
+				// Ignore parse errors
 			}
 		});
 	}
@@ -132,7 +134,7 @@ export class SecretsManager {
 	/**
 	 * Listen for changes to a specific deployment's session auth.
 	 */
-	public onDidChangeDeploymentAuth(
+	public onDidChangeSessionAuth(
 		label: string,
 		listener: (auth: SessionAuth | undefined) => void | Promise<void>,
 	): Disposable {
@@ -147,6 +149,10 @@ export class SecretsManager {
 	}
 
 	public async getSessionAuth(label: string): Promise<SessionAuth | undefined> {
+		if (!label) {
+			return undefined;
+		}
+
 		try {
 			const data = await this.secrets.get(`${SESSION_KEY_PREFIX}${label}`);
 			if (!data) {
@@ -286,24 +292,35 @@ export class SecretsManager {
 
 	/**
 	 * Migrate from legacy flat sessionToken storage to new format.
+	 * Also sets the current deployment if none exists.
 	 */
-	public async migrateFromLegacyStorage(
-		url: string,
-		label: string,
-	): Promise<boolean> {
+	public async migrateFromLegacyStorage(): Promise<string | undefined> {
+		const legacyUrl = this.memento.get<string>("url");
+		if (!legacyUrl) {
+			return undefined;
+		}
+
+		const label = toSafeHost(legacyUrl);
+
 		const existing = await this.getSessionAuth(label);
 		if (existing) {
-			return false;
+			return undefined;
 		}
 
 		const oldToken = await this.secrets.get(LEGACY_SESSION_TOKEN_KEY);
 		if (!oldToken) {
-			return false;
+			return undefined;
 		}
 
-		await this.setSessionAuth(label, { url, token: oldToken });
+		await this.setSessionAuth(label, { url: legacyUrl, token: oldToken });
 		await this.secrets.delete(LEGACY_SESSION_TOKEN_KEY);
 
-		return true;
+		// Also set as current deployment if none exists
+		const currentDeployment = await this.getCurrentDeployment();
+		if (!currentDeployment) {
+			await this.setCurrentDeployment({ url: legacyUrl, label });
+		}
+
+		return label;
 	}
 }