Review comments

Author: EhabY

src/remote/remote.ts

@@ -291,8 +291,8 @@ export class Remote {
 				this.vscodeProposed,
 				this.contextManager,
 			);
-			disposables.push(monitor);
 			disposables.push(
+				monitor,
 				monitor.onChange.event((w) => (this.commands.workspace = w)),
 			);
 
@@ -310,7 +310,7 @@ export class Remote {
 			disposables.push(stateMachine);
 
 			try {
-				await this.vscodeProposed.window.withProgress(
+				workspace = await this.vscodeProposed.window.withProgress(
 					{
 						location: vscode.ProgressLocation.Notification,
 						cancellable: false,
@@ -320,10 +320,8 @@ export class Remote {
 						let inProgress = false;
 						let pendingWorkspace: Workspace | null = null;
 
-						await new Promise<void>((resolve, reject) => {
+						return new Promise<Workspace>((resolve, reject) => {
 							const processWorkspace = async (w: Workspace) => {
-								workspace = w;
-
 								if (inProgress) {
 									// Process one workspace at a time, keeping only the last
 									pendingWorkspace = w;
@@ -340,26 +338,19 @@ export class Remote {
 									);
 									if (isReady) {
 										subscription.dispose();
-										resolve();
+										resolve(w);
 										return;
 									}
-
-									if (pendingWorkspace) {
-										const isReadyAfter = await stateMachine.processWorkspace(
-											pendingWorkspace,
-											progress,
-										);
-										if (isReadyAfter) {
-											subscription.dispose();
-											resolve();
-										}
-									}
 								} catch (error) {
 									subscription.dispose();
 									reject(error);
 								} finally {
 									inProgress = false;
 								}
+
+								if (pendingWorkspace) {
+									processWorkspace(pendingWorkspace);
+								}
 							};
 
 							processWorkspace(workspace);
@@ -373,6 +364,9 @@ export class Remote {
 				stateMachine.dispose();
 			}
 
+			// Mark initial setup as complete so the monitor can start notifying about state changes
+			monitor.markInitialSetupComplete();
+
 			const agents = extractAgents(workspace.latest_build.resources);
 			const agent = agents.find(
 				(agent) => agent.id === stateMachine.getAgentId(),

src/remote/workspaceStateMachine.ts

@@ -95,6 +95,8 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 			case "pending":
 			case "starting":
 			case "stopping":
+				// Clear the agent ID since it could change after a restart
+				this.agentId = undefined;
 				progress?.report({ message: "Waiting for workspace build..." });
 				this.logger.info(`Waiting for ${workspaceName}...`);
 

src/workspace/workspaceMonitor.ts

@@ -29,7 +29,7 @@ export class WorkspaceMonitor implements vscode.Disposable {
 	private notifiedDeletion = false;
 	private notifiedOutdated = false;
 	private notifiedNotRunning = false;
-	private isReady = false;
+	private completedInitialSetup = false;
 
 	readonly onChange = new vscode.EventEmitter<Workspace>();
 	private readonly statusBarItem: vscode.StatusBarItem;
@@ -111,6 +111,10 @@ export class WorkspaceMonitor implements vscode.Disposable {
 		return monitor;
 	}
 
+	public markInitialSetupComplete(): void {
+		this.completedInitialSetup = true;
+	}
+
 	/**
 	 * Permanently close the websocket.
 	 */
@@ -124,7 +128,6 @@ export class WorkspaceMonitor implements vscode.Disposable {
 	}
 
 	private update(workspace: Workspace) {
-		this.updateReadyState(workspace);
 		this.updateContext(workspace);
 		this.updateStatusBar(workspace);
 	}
@@ -133,7 +136,7 @@ export class WorkspaceMonitor implements vscode.Disposable {
 		this.maybeNotifyOutdated(workspace);
 		this.maybeNotifyAutostop(workspace);
 		this.maybeNotifyDeletion(workspace);
-		if (this.isReady) {
+		if (this.completedInitialSetup) {
 			// This instance might be created before the workspace is running
 			this.maybeNotifyNotRunning(workspace);
 		}
@@ -249,12 +252,6 @@ export class WorkspaceMonitor implements vscode.Disposable {
 		this.logger.error(message);
 	}
 
-	private updateReadyState(workspace: Workspace): void {
-		if (workspace.latest_build.status === "running") {
-			this.isReady = true;
-		}
-	}
-
 	private updateContext(workspace: Workspace) {
 		this.contextManager.set("coder.workspace.updatable", workspace.outdated);
 	}