fix: address shared workspaces review findings

verifyAndUpdateSession keeps verify-before-apply: it verifies the rotated
token before touching the live client, rotates credentials in place when the
user is unchanged (no revision bump or tree rebuild), and recovers a session
that was suspended mid-verify instead of getting stuck signed out (CRF-1, CRF-2).

Also:
- bound fetch() retries with a per-attempt disposed check, and dispose a
  metadata watcher created while dispose() races the fetch (CRF-3, CRF-8)
- dispose the tree-data EventEmitter on dispose() (CRF-15)
- rename clearSideEffects to clearAuthState (CRF-12)
- document the session snapshot revision contract (CRF-13)
- drop the dead getAxiosInstance mock helper and restore an unrelated blank
  line (CRF-14, CRF-16)
- trim comments that restated the code (CRF-9)
- add tests for dispose lifecycle, verify failure, signed-out startup, and
  stale-fetch recovery (CRF-4 through CRF-7)

Author: EhabY

src/deployment/deploymentManager.ts

@@ -200,7 +200,7 @@ export class DeploymentManager
 		this.#authListenerDisposable?.dispose();
 		this.#authListenerDisposable = undefined;
 		this.#sessionStore.signOut(null);
-		this.clearSideEffects();
+		this.clearAuthState();
 		this.telemetryService.setDeploymentUrl("");
 
 		await this.secretsManager.setCurrentDeployment(undefined);
@@ -212,10 +212,10 @@ export class DeploymentManager
 	 */
 	public suspendSession(): void {
 		this.#sessionStore.signOut(this.#sessionStore.current.deployment);
-		this.clearSideEffects();
+		this.clearAuthState();
 	}
 
-	private clearSideEffects(): void {
+	private clearAuthState(): void {
 		this.oauthSessionManager.clearDeployment();
 		this.client.setCredentials(undefined, undefined);
 		this.updateAuthContexts(undefined);
@@ -278,6 +278,10 @@ export class DeploymentManager
 		);
 	}
 
+	/**
+	 * Handle a token change while already signed in: verify the new token before
+	 * applying it, so an unverified token never reaches the live client.
+	 */
 	private async verifyAndUpdateSession(
 		deployment: Deployment & { token: string },
 	): Promise<void> {
@@ -287,10 +291,30 @@ export class DeploymentManager
 				deployment.url,
 				deployment.token,
 			);
-			if (this.#disposed || this.#sessionStore.current !== sessionBefore) {
+			if (this.#disposed) {
 				return;
 			}
-			await this.setDeployment({ ...deployment, user });
+			const current = this.#sessionStore.current;
+
+			if (current === sessionBefore) {
+				// Same-user rotation only needs fresh credentials; skipping
+				// setDeployment avoids a revision bump and tree rebuild.
+				if (current.kind === "signedIn" && current.user.id === user.id) {
+					this.client.setCredentials(deployment.url, deployment.token);
+				} else {
+					await this.setDeployment({ ...deployment, user });
+				}
+				return;
+			}
+
+			// Session changed under us: recover only a same-deployment suspension
+			// (e.g. a concurrent 401); a logout or switch wins.
+			if (
+				current.kind === "signedOut" &&
+				current.deployment?.safeHostname === deployment.safeHostname
+			) {
+				await this.setDeployment({ ...deployment, user });
+			}
 		} catch (e) {
 			this.logger.warn("Failed to authenticate updated session:", e);
 		}

src/deployment/sessionStore.ts

@@ -38,12 +38,12 @@ export class SessionStore implements WorkspaceSessionState {
 
 	public readonly onDidChange = this.#onDidChange.event;
 
-	/** The full current session, including deployment and user. */
+	/** Full session state, including deployment and user. */
 	public get current(): SessionData {
 		return this.#data;
 	}
 
-	/** The lean projection exposed to session-state consumers. */
+	/** Lean projection for consumers that only track auth status and revision. */
 	public getSnapshot(): WorkspaceSessionSnapshot {
 		if (this.#data.kind === "signedIn") {
 			return {

src/extension.ts

@@ -101,6 +101,7 @@ async function doActivate(
 		? await secretsManager.getSessionAuth(deployment.safeHostname)
 		: undefined;
 	tracer.setAuthState(deploymentSessionAuth ? "stored" : "none");
+
 	// Shared handler for auth failures (used by interceptor + session manager)
 	const handleAuthFailure = (): Promise<void> => {
 		deploymentManager.suspendSession();

src/workspace/session.ts

@@ -1,5 +1,12 @@
 import type * as vscode from "vscode";
 
+/**
+ * A point-in-time view of the deployment session for workspace providers.
+ *
+ * `revision` increments on every sign-in or sign-out. Consumers snapshot the
+ * revision before an async call and compare afterward to detect that the
+ * session changed while the call was in flight.
+ */
 export type WorkspaceSessionSnapshot =
 	| { readonly kind: "signedOut"; readonly revision: number }
 	| {

src/workspace/workspacesProvider.ts

@@ -60,6 +60,9 @@ export interface WorkspaceProviderOptions {
 	readonly refreshIntervalMs?: number;
 }
 
+// Bounds fetch() retries when the session keeps changing mid-request.
+const MAX_FETCH_ATTEMPTS = 3;
+
 /**
  * Polls workspaces using the provided REST client and renders them in a tree.
  *
@@ -119,7 +122,6 @@ export class WorkspaceProvider
 			this.fetching = false;
 		}
 
-		// Keep polling only after a clean fetch while still signed in and visible.
 		if (
 			!hadError &&
 			!this.disposed &&
@@ -143,65 +145,75 @@ export class WorkspaceProvider
 	 * signed out, and throws if the query fails.
 	 */
 	private async fetch(): Promise<WorkspaceTreeItem[]> {
-		const session = this.sessionState.getSnapshot();
-		if (session.kind !== "signedIn") {
-			return [];
-		}
+		for (let attempt = 0; attempt < MAX_FETCH_ATTEMPTS; attempt++) {
+			if (this.disposed) {
+				return [];
+			}
+			const session = this.sessionState.getSnapshot();
+			if (session.kind !== "signedIn") {
+				return [];
+			}
 
-		const resp = await this.client.getWorkspaces({
-			q: this.getWorkspacesQuery,
-		});
+			const resp = await this.client.getWorkspaces({
+				q: this.getWorkspacesQuery,
+			});
 
-		// If the session changed while the request was in flight, this result is
-		// stale. Drop it and fetch again for the current session.
-		const latest = this.sessionState.getSnapshot();
-		if (latest.kind !== "signedIn" || latest.revision !== session.revision) {
-			return this.fetch();
-		}
+			// Session changed mid-request; this result is stale, so retry.
+			const latest = this.sessionState.getSnapshot();
+			if (latest.kind !== "signedIn" || latest.revision !== session.revision) {
+				continue;
+			}
 
-		const workspaces = this.filterWorkspaces(resp.workspaces, session);
-		const oldWatcherIds = [...this.agentWatchers.keys()];
-		const reusedWatcherIds: string[] = [];
-
-		// TODO: I think it might make more sense for the tree items to contain
-		// their own watchers, rather than recreate the tree items every time and
-		// have this separate map held outside the tree.
-		if (this.config.showMetadata) {
-			const agents = extractAllAgents(workspaces);
-			for (const agent of agents) {
-				// If we have an existing watcher, re-use it.
-				const oldWatcher = this.agentWatchers.get(agent.id);
-				if (oldWatcher) {
-					reusedWatcherIds.push(agent.id);
-				} else {
-					// Otherwise create a new watcher.
+			const workspaces = this.filterWorkspaces(resp.workspaces, session);
+			const oldWatcherIds = [...this.agentWatchers.keys()];
+			const reusedWatcherIds: string[] = [];
+
+			// TODO: I think it might make more sense for the tree items to contain
+			// their own watchers, rather than recreate the tree items every time
+			// and have this separate map held outside the tree.
+			if (this.config.showMetadata) {
+				const agents = extractAllAgents(workspaces);
+				for (const agent of agents) {
+					// If we have an existing watcher, re-use it.
+					const oldWatcher = this.agentWatchers.get(agent.id);
+					if (oldWatcher) {
+						reusedWatcherIds.push(agent.id);
+						continue;
+					}
 					const watcher = await createAgentMetadataWatcher(
 						agent.id,
 						this.client,
 					);
+					// dispose() may have cleared the map mid-create; don't leak
+					// this watcher.
+					if (this.disposed) {
+						watcher.dispose();
+						return [];
+					}
 					watcher.onChange(() => this.refreshTree());
 					this.agentWatchers.set(agent.id, watcher);
 				}
 			}
-		}
 
-		// Dispose of watchers we ended up not reusing.
-		for (const id of oldWatcherIds) {
-			if (!reusedWatcherIds.includes(id)) {
-				this.agentWatchers.get(id)?.dispose();
-				this.agentWatchers.delete(id);
+			// Dispose of watchers we ended up not reusing.
+			for (const id of oldWatcherIds) {
+				if (!reusedWatcherIds.includes(id)) {
+					this.agentWatchers.get(id)?.dispose();
+					this.agentWatchers.delete(id);
+				}
 			}
-		}
 
-		// Create tree items for each workspace
-		return workspaces.map(
-			(workspace: Workspace) =>
-				new WorkspaceTreeItem(
-					workspace,
-					this.config.showOwner,
-					this.config.showMetadata,
-				),
-		);
+			return workspaces.map(
+				(workspace: Workspace) =>
+					new WorkspaceTreeItem(
+						workspace,
+						this.config.showOwner,
+						this.config.showMetadata,
+					),
+			);
+		}
+		// Session changed on every attempt; the next refresh will catch up.
+		return [];
 	}
 
 	private filterWorkspaces(
@@ -244,10 +256,7 @@ export class WorkspaceProvider
 		}
 	}
 
-	/**
-	 * Schedule a refresh if one is not already scheduled or underway and a
-	 * timeout length was provided.
-	 */
+	/** Schedule the next poll, unless one is pending or no interval is set. */
 	private maybeScheduleRefresh() {
 		if (this.options.refreshIntervalMs && !this.timeout) {
 			this.timeout = setTimeout(() => {
@@ -370,6 +379,7 @@ export class WorkspaceProvider
 		this.disposed = true;
 		this.clearState();
 		this.sessionChangeDisposable.dispose();
+		this._onDidChangeTreeData.dispose();
 	}
 }
 

test/mocks/testHelpers.ts

@@ -1254,7 +1254,6 @@ interface WorkspacesResponse {
  * tests drive the production watcher rather than a stub.
  */
 export class MockWorkspacesClient {
-	baseURL: string | undefined = "https://coder.example.com";
 	readonly metadataStreams = new Map<
 		string,
 		MockEventStream<{ data: AgentMetadataEvent[] }>
@@ -1271,10 +1270,6 @@ export class MockWorkspacesClient {
 		return Promise.resolve(stream);
 	}
 
-	getAxiosInstance() {
-		return { defaults: { baseURL: this.baseURL } };
-	}
-
 	/** Resolve the next getWorkspaces call with these workspaces. */
 	respondOnce(workspaces: readonly Workspace[]): void {
 		this.getWorkspaces.mockResolvedValueOnce({

test/unit/deployment/deploymentManager.test.ts

@@ -478,6 +478,94 @@ describe("DeploymentManager", () => {
 			expect(manager.getCurrentDeployment()).toBeNull();
 			expect(currentUserId(manager)).toBeUndefined();
 		});
+
+		it("rotates the token in place without a revision bump when the user is unchanged", async () => {
+			const { mockClient, validationMockClient, secretsManager, manager } =
+				createTestContext();
+			const user = createMockUser({ id: "stable-user" });
+			validationMockClient.setAuthenticatedUserResponse(user);
+
+			await manager.setDeployment({
+				url: TEST_URL,
+				safeHostname: TEST_HOSTNAME,
+				token: "initial-token",
+				user,
+			});
+			const revisionBefore = manager.getSnapshot().revision;
+
+			await secretsManager.setSessionAuth(TEST_HOSTNAME, {
+				url: TEST_URL,
+				token: "rotated-token",
+			});
+			await flush();
+
+			expect(mockClient.token).toBe("rotated-token");
+			expect(currentUserId(manager)).toBe("stable-user");
+			// No sign-in fired, so the workspace trees are not rebuilt.
+			expect(manager.getSnapshot().revision).toBe(revisionBefore);
+		});
+
+		it("keeps the existing session when verifying a rotated token fails", async () => {
+			const { mockClient, validationMockClient, secretsManager, manager } =
+				createTestContext();
+			const user = createMockUser({ id: "stable-user" });
+
+			await manager.setDeployment({
+				url: TEST_URL,
+				safeHostname: TEST_HOSTNAME,
+				token: "initial-token",
+				user,
+			});
+
+			// Verifying the rotated token fails (e.g. a transient network blip).
+			validationMockClient.getAuthenticatedUser.mockRejectedValueOnce(
+				new Error("network down"),
+			);
+			await secretsManager.setSessionAuth(TEST_HOSTNAME, {
+				url: TEST_URL,
+				token: "rotated-token",
+			});
+			await flush();
+
+			// Verify-before-apply: the unverified token never reaches the client.
+			expect(mockClient.token).toBe("initial-token");
+			expect(manager.isAuthenticated()).toBe(true);
+			expect(currentUserId(manager)).toBe("stable-user");
+		});
+
+		it("recovers a session suspended while a rotated token is verified", async () => {
+			const { mockClient, validationMockClient, secretsManager, manager } =
+				createTestContext();
+			const user = createMockUser({ id: "stable-user" });
+			const validation = Promise.withResolvers<typeof user>();
+			validationMockClient.getAuthenticatedUser.mockReturnValueOnce(
+				validation.promise,
+			);
+
+			await manager.setDeployment({
+				url: TEST_URL,
+				safeHostname: TEST_HOSTNAME,
+				token: "initial-token",
+				user,
+			});
+			await secretsManager.setSessionAuth(TEST_HOSTNAME, {
+				url: TEST_URL,
+				token: "rotated-token",
+			});
+			await flush();
+
+			// A concurrent 401 on the old token suspends the session mid-verify.
+			manager.suspendSession();
+			expect(manager.isAuthenticated()).toBe(false);
+
+			validation.resolve(user);
+			await flush();
+
+			// The verified token recovers the session instead of staying suspended.
+			expect(manager.isAuthenticated()).toBe(true);
+			expect(mockClient.token).toBe("rotated-token");
+			expect(currentUserId(manager)).toBe("stable-user");
+		});
 	});
 
 	describe("logout", () => {