Move memento and secrets managment from storage.ts

EhabY · EhabY · commit 1896c541b04c · 2025-09-17T12:42:01.000+03:00
diff --git a/src/commands.ts b/src/commands.ts
@@ -10,7 +10,9 @@ import { createWorkspaceIdentifier, extractAgents } from "./api/api-helper";
 import { CoderApi } from "./api/coderApi";
 import { needToken } from "./api/utils";
 import { CliConfigManager } from "./core/cliConfig";
+import { MementoManager } from "./core/mementoManager";
 import { PathResolver } from "./core/pathResolver";
+import { SecretsManager } from "./core/secretsManager";
 import { CertificateError } from "./error";
 import { getGlobalFlags } from "./globalFlags";
 import { Storage } from "./storage";
@@ -39,6 +41,8 @@ export class Commands {
 		private readonly restClient: Api,
 		private readonly storage: Storage,
 		private readonly pathResolver: PathResolver,
+		private readonly mementoManager: MementoManager,
+		private readonly secretsManager: SecretsManager,
 	) {
 		this.cliConfigManager = new CliConfigManager(pathResolver);
 	}
@@ -108,7 +112,7 @@ export class Commands {
 		quickPick.title = "Enter the URL of your Coder deployment.";
 
 		// Initial items.
-		quickPick.items = this.storage
+		quickPick.items = this.mementoManager
 			.withUrlHistory(defaultURL, process.env.CODER_URL)
 			.map((url) => ({
 				alwaysShow: true,
@@ -119,7 +123,7 @@ export class Commands {
 		// an option in case the user wants to connect to something that is not in
 		// the list.
 		quickPick.onDidChangeValue((value) => {
-			quickPick.items = this.storage
+			quickPick.items = this.mementoManager
 				.withUrlHistory(defaultURL, process.env.CODER_URL, value)
 				.map((url) => ({
 					alwaysShow: true,
@@ -199,8 +203,8 @@ export class Commands {
 		this.restClient.setSessionToken(res.token);
 
 		// Store these to be used in later sessions.
-		await this.storage.setUrl(url);
-		await this.storage.setSessionToken(res.token);
+		await this.mementoManager.setUrl(url);
+		await this.secretsManager.setSessionToken(res.token);
 
 		// Store on disk to be used by the cli.
 		await this.cliConfigManager.configure(label, url, res.token);
@@ -288,7 +292,7 @@ export class Commands {
 			title: "Coder API Key",
 			password: true,
 			placeHolder: "Paste your API key.",
-			value: token || (await this.storage.getSessionToken()),
+			value: token || (await this.secretsManager.getSessionToken()),
 			ignoreFocusOut: true,
 			validateInput: async (value) => {
 				client.setSessionToken(value);
@@ -354,7 +358,7 @@ export class Commands {
 	 * Log out from the currently logged-in deployment.
 	 */
 	public async logout(): Promise<void> {
-		const url = this.storage.getUrl();
+		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");
@@ -366,8 +370,8 @@ export class Commands {
 		this.restClient.setSessionToken("");
 
 		// Clear from memory.
-		await this.storage.setUrl(undefined);
-		await this.storage.setSessionToken(undefined);
+		await this.mementoManager.setUrl(undefined);
+		await this.secretsManager.setSessionToken(undefined);
 
 		await vscode.commands.executeCommand(
 			"setContext",
@@ -392,7 +396,7 @@ export class Commands {
 	 * Must only be called if currently logged in.
 	 */
 	public async createWorkspace(): Promise<void> {
-		const uri = this.storage.getUrl() + "/templates";
+		const uri = this.mementoManager.getUrl() + "/templates";
 		await vscode.commands.executeCommand("vscode.open", uri);
 	}
 
@@ -407,7 +411,7 @@ export class Commands {
 	public async navigateToWorkspace(item: OpenableTreeItem) {
 		if (item) {
 			const workspaceId = createWorkspaceIdentifier(item.workspace);
-			const uri = this.storage.getUrl() + `/@${workspaceId}`;
+			const uri = this.mementoManager.getUrl() + `/@${workspaceId}`;
 			await vscode.commands.executeCommand("vscode.open", uri);
 		} else if (this.workspace && this.workspaceRestClient) {
 			const baseUrl =
@@ -430,7 +434,7 @@ export class Commands {
 	public async navigateToWorkspaceSettings(item: OpenableTreeItem) {
 		if (item) {
 			const workspaceId = createWorkspaceIdentifier(item.workspace);
-			const uri = this.storage.getUrl() + `/@${workspaceId}/settings`;
+			const uri = this.mementoManager.getUrl() + `/@${workspaceId}/settings`;
 			await vscode.commands.executeCommand("vscode.open", uri);
 		} else if (this.workspace && this.workspaceRestClient) {
 			const baseUrl =
@@ -508,7 +512,7 @@ export class Commands {
 
 					// If workspace_name is provided, run coder ssh before the command
 
-					const url = this.storage.getUrl();
+					const url = this.mementoManager.getUrl();
 					if (!url) {
 						throw new Error("No coder url found for sidebar");
 					}
@@ -650,8 +654,8 @@ export class Commands {
 			newWindow = false;
 		}
 
-		// Only set the memento if when opening a new folder
-		await this.storage.setFirstConnect();
+		// Only set the memento when opening a new folder
+		await this.mementoManager.setFirstConnect();
 		await vscode.commands.executeCommand(
 			"vscode.openFolder",
 			vscode.Uri.from({
@@ -831,8 +835,8 @@ export class Commands {
 			}
 		}
 
-		// Only set the memento if when opening a new folder/window
-		await this.storage.setFirstConnect();
+		// Only set the memento when opening a new folder/window
+		await this.mementoManager.setFirstConnect();
 		if (folderPath) {
 			await vscode.commands.executeCommand(
 				"vscode.openFolder",
diff --git a/src/core/mementoManager.ts b/src/core/mementoManager.ts
@@ -0,0 +1,71 @@
+import type { Memento } from "vscode";
+
+// Maximum number of recent URLs to store.
+const MAX_URLS = 10;
+
+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.
+	 */
+	public async setUrl(url?: string): Promise<void> {
+		await this.memento.update("url", url);
+		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.
+	 */
+	public withUrlHistory(...append: (string | undefined)[]): string[] {
+		const val = this.memento.get("urlHistory");
+		const urls = Array.isArray(val) ? new Set(val) : new Set();
+		for (const url of append) {
+			if (url) {
+				// It might exist; delete first so it gets appended.
+				urls.delete(url);
+				urls.add(url);
+			}
+		}
+		// Slice off the head if the list is too large.
+		return urls.size > MAX_URLS
+			? Array.from(urls).slice(urls.size - MAX_URLS, urls.size)
+			: Array.from(urls);
+	}
+
+	/**
+	 * Mark this as the first connection to a workspace, which influences whether
+	 * the workspace startup confirmation is shown to the user.
+	 */
+	public async setFirstConnect(): Promise<void> {
+		return this.memento.update("firstConnect", true);
+	}
+
+	/**
+	 * Check if this is the first connection to a workspace and clear the flag.
+	 * Used to determine whether to automatically start workspaces without
+	 * prompting the user for confirmation.
+	 */
+	public async getAndClearFirstConnect(): Promise<boolean> {
+		const isFirst = this.memento.get<boolean>("firstConnect");
+		if (isFirst !== undefined) {
+			await this.memento.update("firstConnect", undefined);
+		}
+		return isFirst === true;
+	}
+}
diff --git a/src/core/pathResolver.ts b/src/core/pathResolver.ts
@@ -4,6 +4,7 @@ import type { WorkspaceConfiguration } from "vscode";
 export class PathResolver {
 	constructor(
 		private readonly basePath: string,
+		private readonly codeLogPath: string,
 		private readonly configurations: WorkspaceConfiguration,
 	) {}
 
@@ -100,4 +101,16 @@ export class PathResolver {
 	public getUrlPath(label: string): string {
 		return path.join(this.getGlobalConfigDir(label), "url");
 	}
+
+	/**
+	 * The uri of a directory in which the extension can create log files.
+	 *
+	 * The directory might not exist on disk and creation is up to the extension.
+	 * However, the parent directory is guaranteed to be existent.
+	 *
+	 * This directory is provided by VS Code and may not be the same as the directory where the Coder CLI writes its log files.
+	 */
+	public getCodeLogDir(): string {
+		return this.codeLogPath;
+	}
 }
diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts
@@ -0,0 +1,29 @@
+import type { SecretStorage } from "vscode";
+
+export class SecretsManager {
+	constructor(private readonly secrets: SecretStorage) {}
+
+	/**
+	 * Set or unset the last used token.
+	 */
+	public async setSessionToken(sessionToken?: string): Promise<void> {
+		if (!sessionToken) {
+			await this.secrets.delete("sessionToken");
+		} else {
+			await this.secrets.store("sessionToken", sessionToken);
+		}
+	}
+
+	/**
+	 * Get the last used token.
+	 */
+	public async getSessionToken(): Promise<string | undefined> {
+		try {
+			return await this.secrets.get("sessionToken");
+		} catch (ex) {
+			// The VS Code session store has become corrupt before, and
+			// will fail to get the session token...
+			return undefined;
+		}
+	}
+}
diff --git a/src/extension.ts b/src/extension.ts
@@ -8,7 +8,9 @@ import { CoderApi } from "./api/coderApi";
 import { needToken } from "./api/utils";
 import { Commands } from "./commands";
 import { CliConfigManager } from "./core/cliConfig";
+import { MementoManager } from "./core/mementoManager";
 import { PathResolver } from "./core/pathResolver";
+import { SecretsManager } from "./core/secretsManager";
 import { CertificateError, getErrorDetail } from "./error";
 import { Remote } from "./remote";
 import { Storage } from "./storage";
@@ -52,28 +54,26 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 
 	const pathResolver = new PathResolver(
 		ctx.globalStorageUri.fsPath,
+		ctx.logUri.fsPath,
 		vscode.workspace.getConfiguration(),
 	);
+	const cliConfigManager = new CliConfigManager(pathResolver);
+	const mementoManager = new MementoManager(ctx.globalState);
+	const secretsManager = new SecretsManager(ctx.secrets);
+
 	const output = vscode.window.createOutputChannel("Coder", { log: true });
-	const storage = new Storage(
-		vscodeProposed,
-		output,
-		ctx.globalState,
-		ctx.secrets,
-		ctx.logUri,
-		pathResolver,
-	);
+	const storage = new Storage(output, pathResolver);
 
-	// Try to clear this flag ASAP then pass it around if needed
-	const isFirstConnect = await storage.getAndClearFirstConnect();
+	// Try to clear this flag ASAP
+	const isFirstConnect = await mementoManager.getAndClearFirstConnect();
 
 	// This client tracks the current login and will be used through the life of
 	// the plugin to poll workspaces for the current login, as well as being used
 	// in commands that operate on the current login.
-	const url = storage.getUrl();
+	const url = mementoManager.getUrl();
 	const client = CoderApi.create(
 		url || "",
-		await storage.getSessionToken(),
+		await secretsManager.getSessionToken(),
 		storage.output,
 		() => vscode.workspace.getConfiguration(),
 	);
@@ -108,8 +108,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		allWorkspacesProvider.setVisibility(event.visible);
 	});
 
-	const cliConfigManager = new CliConfigManager(pathResolver);
-
 	// Handle vscode:// URIs.
 	vscode.window.registerUriHandler({
 		handleUri: async (uri) => {
@@ -137,11 +135,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 				// hit enter and move on.
 				const url = await commands.maybeAskUrl(
 					params.get("url"),
-					storage.getUrl(),
+					mementoManager.getUrl(),
 				);
 				if (url) {
 					client.setHost(url);
-					await storage.setUrl(url);
+					await mementoManager.setUrl(url);
 				} else {
 					throw new Error(
 						"url must be provided or specified as a query parameter",
@@ -159,7 +157,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 					: (params.get("token") ?? "");
 				if (token) {
 					client.setSessionToken(token);
-					await storage.setSessionToken(token);
+					await secretsManager.setSessionToken(token);
 				}
 
 				// Store on disk to be used by the cli.
@@ -219,11 +217,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 				// hit enter and move on.
 				const url = await commands.maybeAskUrl(
 					params.get("url"),
-					storage.getUrl(),
+					mementoManager.getUrl(),
 				);
 				if (url) {
 					client.setHost(url);
-					await storage.setUrl(url);
+					await mementoManager.setUrl(url);
 				} else {
 					throw new Error(
 						"url must be provided or specified as a query parameter",
@@ -261,7 +259,14 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 
 	// Register globally available commands.  Many of these have visibility
 	// controlled by contexts, see `when` in the package.json.
-	const commands = new Commands(vscodeProposed, client, storage, pathResolver);
+	const commands = new Commands(
+		vscodeProposed,
+		client,
+		storage,
+		pathResolver,
+		mementoManager,
+		secretsManager,
+	);
 	vscode.commands.registerCommand("coder.login", commands.login.bind(commands));
 	vscode.commands.registerCommand(
 		"coder.logout",
diff --git a/src/remote.ts b/src/remote.ts
diff --git a/src/storage.ts b/src/storage.ts
