Only set session auth if the token is valid

* Always return User if the login is successful

Author: EhabY

src/commands.ts

@@ -111,9 +111,6 @@ export class Commands {
 			return;
 		}
 
-		// Login might have happened in another process/window so we do not have the user yet.
-		result.user ??= await this.extensionClient.getAuthenticatedUser();
-
 		await this.deploymentManager.setDeployment({
 			url,
 			safeHostname,

src/deployment/deploymentManager.ts

@@ -83,13 +83,8 @@ export class DeploymentManager implements vscode.Disposable {
 	 * Returns true if deployment was changed, false otherwise.
 	 */
 	public async setDeploymentIfValid(
-		deployment: DeploymentWithAuth,
+		deployment: Deployment & { token?: string },
 	): Promise<boolean> {
-		if (deployment.user) {
-			await this.setDeployment({ ...deployment, user: deployment.user });
-			return true;
-		}
-
 		const auth = await this.secretsManager.getSessionAuth(
 			deployment.safeHostname,
 		);

src/login/loginCoordinator.ts

@@ -16,12 +16,13 @@ import type { Logger } from "../logging/logger";
 
 type LoginResult =
 	| { success: false }
-	| { success: true; user?: User; token: string };
+	| { success: true; user: User; token: string };
 
 interface LoginOptions {
 	safeHostname: string;
 	url: string | undefined;
 	autoLogin?: boolean;
+	token?: string;
 }
 
 /**
@@ -49,6 +50,7 @@ export class LoginCoordinator {
 			const result = await this.attemptLogin(
 				{ safeHostname, url },
 				options.autoLogin ?? false,
+				options.token,
 			);
 
 			await this.persistSessionAuth(result, safeHostname, url);
@@ -95,6 +97,7 @@ export class LoginCoordinator {
 						const result = await this.attemptLogin(
 							{ url: newUrl, safeHostname },
 							false,
+							options.token,
 						);
 
 						await this.persistSessionAuth(result, safeHostname, newUrl);
@@ -168,10 +171,17 @@ export class LoginCoordinator {
 		const promise = new Promise<LoginResult>((resolve) => {
 			disposable = this.secretsManager.onDidChangeSessionAuth(
 				safeHostname,
-				(auth) => {
+				async (auth) => {
 					if (auth?.token) {
 						disposable?.dispose();
-						resolve({ success: true, token: auth.token });
+						const client = CoderApi.create(auth.url, auth.token, this.logger);
+						try {
+							const user = await client.getAuthenticatedUser();
+							resolve({ success: true, token: auth.token, user });
+						} catch {
+							// Token from other window was invalid, ignore and keep waiting
+							// (or user can click Login/Cancel in the dialog)
+						}
 					}
 				},
 			);
@@ -191,54 +201,93 @@ export class LoginCoordinator {
 	private async attemptLogin(
 		deployment: Deployment,
 		isAutoLogin: boolean,
+		providedToken?: string,
 	): Promise<LoginResult> {
-		const needsToken = needToken(vscode.workspace.getConfiguration());
 		const client = CoderApi.create(deployment.url, "", this.logger);
 
-		let storedToken: string | undefined;
-		if (needsToken) {
-			const auth = await this.secretsManager.getSessionAuth(
-				deployment.safeHostname,
+		// mTLS authentication (no token needed)
+		if (!needToken(vscode.workspace.getConfiguration())) {
+			return this.tryMtlsAuth(client, isAutoLogin);
+		}
+
+		// Try provided token first
+		if (providedToken) {
+			const result = await this.tryTokenAuth(
+				client,
+				providedToken,
+				isAutoLogin,
 			);
-			storedToken = auth?.token;
-			if (storedToken) {
-				client.setSessionToken(storedToken);
+			if (result !== "retry") {
+				return result;
 			}
 		}
 
-		// Attempt authentication with current credentials (token or mTLS)
-		try {
-			if (!needsToken || storedToken) {
-				const user = await client.getAuthenticatedUser();
-				// Return the token that was used (empty string for mTLS since
-				// the `vscodessh` command currently always requires a token file)
-				return { success: true, token: storedToken ?? "", user };
+		// Try stored token (skip if same as provided)
+		const auth = await this.secretsManager.getSessionAuth(
+			deployment.safeHostname,
+		);
+		if (auth?.token && auth.token !== providedToken) {
+			const result = await this.tryTokenAuth(client, auth.token, isAutoLogin);
+			if (result !== "retry") {
+				return result;
 			}
+		}
+
+		// Prompt user for token
+		return this.loginWithToken(client);
+	}
+
+	private async tryMtlsAuth(
+		client: CoderApi,
+		isAutoLogin: boolean,
+	): Promise<LoginResult> {
+		try {
+			const user = await client.getAuthenticatedUser();
+			return { success: true, token: "", user };
 		} catch (err) {
-			const is401 = isAxiosError(err) && err.response?.status === 401;
-			if (needsToken && is401) {
-				// For token auth with 401: silently continue to prompt for new credentials
-			} else {
-				// For mTLS or non-401 errors: show error and abort
-				const message = getErrorMessage(err, "no response from the server");
-				if (isAutoLogin) {
-					this.logger.warn("Failed to log in to Coder server:", message);
-				} else {
-					this.vscodeProposed.window.showErrorMessage(
-						"Failed to log in to Coder server",
-						{
-							detail: message,
-							modal: true,
-							useCustom: true,
-						},
-					);
-				}
-				return { success: false };
+			this.showAuthError(err, isAutoLogin);
+			return { success: false };
+		}
+	}
+
+	/**
+	 * Returns 'retry' on 401 to signal trying next token.
+	 */
+	private async tryTokenAuth(
+		client: CoderApi,
+		token: string,
+		isAutoLogin: boolean,
+	): Promise<LoginResult | "retry"> {
+		client.setSessionToken(token);
+		try {
+			const user = await client.getAuthenticatedUser();
+			return { success: true, token, user };
+		} catch (err) {
+			if (isAxiosError(err) && err.response?.status === 401) {
+				return "retry";
 			}
+			this.showAuthError(err, isAutoLogin);
+			return { success: false };
 		}
+	}
 
-		const result = await this.loginWithToken(client);
-		return result;
+	/**
+	 * Shows auth error via dialog or logs it for autoLogin.
+	 */
+	private showAuthError(err: unknown, isAutoLogin: boolean): void {
+		const message = getErrorMessage(err, "no response from the server");
+		if (isAutoLogin) {
+			this.logger.warn("Failed to log in to Coder server:", message);
+		} else {
+			this.vscodeProposed.window.showErrorMessage(
+				"Failed to log in to Coder server",
+				{
+					detail: message,
+					modal: true,
+					useCustom: true,
+				},
+			);
+		}
 	}
 
 	/**

src/uri/uriHandler.ts

@@ -159,21 +159,18 @@ async function setupDeployment(
 
 	const safeHostname = toSafeHost(url);
 
-	const token: string | null = params.get("token");
-	if (token !== null) {
-		await secretsManager.setSessionAuth(safeHostname, { url, token });
-	}
-
+	const token: string | undefined = params.get("token") ?? undefined;
 	const result = await loginCoordinator.ensureLoggedIn({
 		safeHostname,
 		url,
+		token,
 	});
 
 	if (!result.success) {
 		throw new Error("Failed to login to deployment from URI");
 	}
 
-	await deploymentManager.setDeploymentIfValid({
+	await deploymentManager.setDeployment({
 		safeHostname,
 		url,
 		token: result.token,

test/unit/login/loginCoordinator.test.ts

@@ -331,4 +331,142 @@ describe("LoginCoordinator", () => {
 			expect(result.success).toBe(false);
 		});
 	});
+
+	describe("token fallback order", () => {
+		it("uses provided token first when valid", async () => {
+			const { secretsManager, coordinator, mockSuccessfulAuth } =
+				createTestContext();
+			const user = mockSuccessfulAuth();
+
+			// Store a different token
+			await secretsManager.setSessionAuth(TEST_HOSTNAME, {
+				url: TEST_URL,
+				token: "stored-token",
+			});
+
+			const result = await coordinator.ensureLoggedIn({
+				url: TEST_URL,
+				safeHostname: TEST_HOSTNAME,
+				token: "provided-token",
+			});
+
+			expect(result).toEqual({ success: true, user, token: "provided-token" });
+		});
+
+		it("falls back to stored token when provided token is invalid", async () => {
+			const { mockAdapter, secretsManager, coordinator } = createTestContext();
+			const user = createMockUser();
+
+			mockAdapter
+				.mockRejectedValueOnce({
+					isAxiosError: true,
+					response: { status: 401 }, // Fail the provided token with 401
+					message: "Unauthorized",
+				})
+				.mockResolvedValueOnce({
+					data: user,
+					status: 200, // Succeed the stored token
+					headers: {},
+					config: {},
+				});
+
+			await secretsManager.setSessionAuth(TEST_HOSTNAME, {
+				url: TEST_URL,
+				token: "stored-token",
+			});
+
+			const result = await coordinator.ensureLoggedIn({
+				url: TEST_URL,
+				safeHostname: TEST_HOSTNAME,
+				token: "invalid-provided-token",
+			});
+
+			expect(result).toEqual({ success: true, user, token: "stored-token" });
+		});
+
+		it("prompts user when both provided and stored tokens are invalid", async () => {
+			const { mockAdapter, userInteraction, secretsManager, coordinator } =
+				createTestContext();
+			const user = createMockUser();
+
+			mockAdapter
+				.mockRejectedValueOnce({
+					isAxiosError: true,
+					response: { status: 401 }, // provided token
+					message: "Unauthorized",
+				})
+				.mockRejectedValueOnce({
+					isAxiosError: true,
+					response: { status: 401 }, // stored token
+					message: "Unauthorized",
+				})
+				.mockResolvedValueOnce({
+					data: user,
+					status: 200, // user-entered token
+					headers: {},
+					config: {},
+				});
+
+			await secretsManager.setSessionAuth(TEST_HOSTNAME, {
+				url: TEST_URL,
+				token: "stored-token",
+			});
+
+			userInteraction.setInputBoxValue("user-entered-token");
+
+			const result = await coordinator.ensureLoggedIn({
+				url: TEST_URL,
+				safeHostname: TEST_HOSTNAME,
+				token: "invalid-provided-token",
+			});
+
+			expect(result).toEqual({
+				success: true,
+				user,
+				token: "user-entered-token",
+			});
+			expect(vscode.window.showInputBox).toHaveBeenCalled();
+		});
+
+		it("skips stored token check when same as provided token", async () => {
+			const { mockAdapter, userInteraction, secretsManager, coordinator } =
+				createTestContext();
+			const user = createMockUser();
+
+			mockAdapter
+				.mockRejectedValueOnce({
+					isAxiosError: true,
+					response: { status: 401 }, // provided token
+					message: "Unauthorized",
+				})
+				.mockResolvedValueOnce({
+					data: user,
+					status: 200, // user-entered token
+					headers: {},
+					config: {},
+				});
+
+			// Store the SAME token as will be provided
+			await secretsManager.setSessionAuth(TEST_HOSTNAME, {
+				url: TEST_URL,
+				token: "same-token",
+			});
+
+			userInteraction.setInputBoxValue("user-entered-token");
+
+			const result = await coordinator.ensureLoggedIn({
+				url: TEST_URL,
+				safeHostname: TEST_HOSTNAME,
+				token: "same-token",
+			});
+
+			expect(result).toEqual({
+				success: true,
+				user,
+				token: "user-entered-token",
+			});
+			// Provided/stored token check only called once + user prompt
+			expect(mockAdapter).toHaveBeenCalledTimes(2);
+		});
+	});
 });