fix: queue workspace updates as one build and stop autostart races

`coder update` runs a stop build and then a start build. Between the
two, the extension's own `coder ssh` reconnect autostarts the workspace
on the old template version, so the update's start build is rejected
with "A workspace build is already active" and the workspace comes back
unchanged.

On servers from 2.36, post a single build that carries the start in
`on_success`, so nothing can take the build slot in between. Leave
`template_version_id` unset on the follow-up build, because pinning it
requires template update permission and the update targets the active
version anyway. Servers before 2.36 keep the `coder update` path, and
CLIs before 2.24 keep the REST path.

Pass `--disable-autostart` as an internal SSH flag so a custom
`coder.sshFlags` value cannot drop it, and ask before connecting to the
existing version when an update fails instead of falling back silently.

Fixes #1095

Author: EhabY

package.json

@@ -183,7 +183,7 @@
 					]
 				},
 				"coder.sshFlags": {
-					"markdownDescription": "Additional flags to pass to the `coder ssh` command when establishing SSH connections. Enter each flag as a separate array item; values are passed verbatim and in order. See the [CLI ssh reference](https://coder.com/docs/reference/cli/ssh) for available flags.\n\nNote: `--network-info-dir` and `--ssh-host-prefix` are ignored (managed internally). Prefer `#coder.proxyLogDirectory#` over `--log-dir`/`-l` for full functionality.",
+					"markdownDescription": "Additional flags to pass to the `coder ssh` command when establishing SSH connections. Enter each flag as a separate array item; values are passed verbatim and in order. See the [CLI ssh reference](https://coder.com/docs/reference/cli/ssh) for available flags.\n\nNote: `--disable-autostart`, `--network-info-dir`, and `--ssh-host-prefix` are managed internally. Prefer `#coder.proxyLogDirectory#` over `--log-dir`/`-l` for full functionality.",
 					"type": "array",
 					"items": {
 						"type": "string"

src/api/workspace.ts

@@ -1,23 +1,30 @@
 import { spawn } from "node:child_process";
+import * as semver from "semver";
 import * as vscode from "vscode";
 
+import { versionAtLeast, type FeatureSet } from "../featureSet";
 import { getGlobalFlags, type CliAuth } from "../settings/cli";
 
 import { errToStr, createWorkspaceIdentifier } from "./api-helper";
 
 import type { Api } from "coder/site/src/api/api";
 import type {
+	CreateWorkspaceBuildOnSuccessRequest,
 	ProvisionerJobLog,
 	Workspace,
 	WorkspaceAgentLog,
 	WorkspaceBuildParameter,
 } from "coder/site/src/api/typesGenerated";
 
-import type { FeatureSet } from "../featureSet";
 import type { UnidirectionalStream } from "../websocket/eventStreamConnection";
 
 import type { CoderApi } from "./coderApi";
 
+/** Server version that stops and starts in one build via `on_success`. */
+const RESTART_BUILD_VERSION = "2.36.0";
+
+const UPDATE_REASON = "vscode_connection";
+
 /** Opens a stream once; subsequent open() calls are no-ops until closed. */
 export class LazyStream<T> {
 	private stream: UnidirectionalStream<T> | null = null;
@@ -108,7 +115,7 @@ export async function startWorkspace(ctx: CliContext): Promise<Workspace> {
 
 	const args = ["start", "--yes"];
 	if (ctx.featureSet.buildReason) {
-		args.push("--reason", "vscode_connection");
+		args.push("--reason", UPDATE_REASON);
 	}
 
 	await runCliCommand(ctx, args);
@@ -118,13 +125,18 @@ export async function startWorkspace(ctx: CliContext): Promise<Workspace> {
 /**
  * Update a workspace to the latest template version. Callers must collect
  * any newly-required parameters via `collectUpdateParameters` first; this
- * function does not prompt. Falls back to the REST API on CLIs older than
- * 2.24.
+ * function does not prompt. Servers older than 2.36 need two builds, via
+ * `coder update` or, on CLIs older than 2.24, the REST API.
  */
 export async function updateWorkspace(
 	ctx: CliContext,
 	parameters: WorkspaceBuildParameter[],
 ): Promise<Workspace> {
+	const { version } = await ctx.restClient.getBuildInfo();
+	if (versionAtLeast(semver.parse(version), RESTART_BUILD_VERSION)) {
+		return updateWorkspaceInOneBuild(ctx, parameters);
+	}
+
 	if (!ctx.featureSet.cliUpdate) {
 		return updateWorkspaceViaApi(ctx, parameters);
 	}
@@ -137,6 +149,43 @@ export async function updateWorkspace(
 	return ctx.restClient.getWorkspace(ctx.workspace.id);
 }
 
+/**
+ * Queues the stop and the start as a single build request, so no other build
+ * can take the slot in between and start the outdated version. Returns the
+ * accepted build; the server creates the start build once the stop succeeds.
+ */
+async function updateWorkspaceInOneBuild(
+	ctx: CliContext,
+	parameters: WorkspaceBuildParameter[],
+): Promise<Workspace> {
+	// Re-read so the transition matches the current build, which may have
+	// changed while parameters were being collected.
+	const workspace = await ctx.restClient.getWorkspace(ctx.workspace.id);
+	const start: CreateWorkspaceBuildOnSuccessRequest = {
+		transition: "start",
+		rich_parameter_values: parameters,
+	};
+	if (workspace.latest_build.status !== "running") {
+		ctx.write("Starting workspace with the updated template...\r\n");
+		const build = await ctx.restClient.postWorkspaceBuild(workspace.id, {
+			...start,
+			reason: UPDATE_REASON,
+			template_version_id: workspace.template_active_version_id,
+		});
+		return { ...workspace, latest_build: build };
+	}
+
+	ctx.write("Restarting workspace with the updated template...\r\n");
+	// A follow-up build can only pin a template version with template update
+	// permission, so leave it unset to take the active version at start time.
+	const build = await ctx.restClient.postWorkspaceBuild(workspace.id, {
+		transition: "stop",
+		reason: UPDATE_REASON,
+		on_success: start,
+	});
+	return { ...workspace, latest_build: build };
+}
+
 async function updateWorkspaceViaApi(
 	ctx: CliContext,
 	parameters: WorkspaceBuildParameter[],
@@ -145,8 +194,8 @@ async function updateWorkspaceViaApi(
 		ctx.write("Stopping workspace for update...\r\n");
 		const stopBuild = await ctx.restClient.stopWorkspace(ctx.workspace.id);
 		const stoppedJob = await ctx.restClient.waitForBuild(stopBuild);
-		if (stoppedJob?.status === "canceled") {
-			throw new Error("Workspace update cancelled during stop");
+		if (stoppedJob?.status !== "succeeded") {
+			throw new Error("Workspace update stop build did not succeed");
 		}
 	}
 

src/featureSet.ts

@@ -14,10 +14,10 @@ export interface FeatureSet {
 }
 
 /**
- * True when the CLI version is at least `minVersion`, or is a dev build.
+ * True when the version is at least `minVersion`, or is a dev build.
  * Returns false for null (unknown) versions.
  */
-function versionAtLeast(
+export function versionAtLeast(
 	version: semver.SemVer | null,
 	minVersion: string,
 ): boolean {

src/instrumentation/workspace.ts

@@ -11,7 +11,7 @@ import type { TelemetryReporter } from "../telemetry/reporter";
 import type { Span } from "../telemetry/span";
 
 export type WorkspacePromptAction = "start" | "update";
-export type WorkspaceUpdatePrompt = "parameters" | "confirmation";
+export type WorkspaceUpdatePrompt = "parameters" | "confirmation" | "failure";
 
 /**
  * Emits `workspace.state_transitioned` for a detected workspace transition.
@@ -152,6 +152,19 @@ export class WorkspaceOperationTelemetry {
 		});
 	}
 
+	/** Records whether the user connects to the existing version anyway. */
+	public traceFailurePrompt(fn: () => Promise<boolean>): Promise<boolean> {
+		return this.traceUpdatePrompt("failure", async (span) => {
+			const connect = await fn();
+			if (!connect) {
+				span.markAborted();
+				return false;
+			}
+			span.setProperty("action", "connect");
+			return true;
+		});
+	}
+
 	private traceUpdatePrompt<T>(
 		prompt: WorkspaceUpdatePrompt,
 		fn: (span: Span) => Promise<T>,

src/remote/remote.ts

@@ -855,6 +855,7 @@ export class Remote {
 			const userSshFlags = getSshFlags(vscodeConfig);
 			// Make sure to update the `coder.sshFlags` description if we add more internal flags here!
 			const internalFlags = [
+				"--disable-autostart",
 				"--stdio",
 				"--usage-app=vscode",
 				"--network-info-dir",

src/remote/workspaceStateMachine.ts

@@ -1,5 +1,3 @@
-import * as vscode from "vscode";
-
 import {
 	createWorkspaceIdentifier,
 	errToStr,
@@ -27,6 +25,7 @@ import type {
 	Workspace,
 	WorkspaceAgentLog,
 } from "coder/site/src/api/typesGenerated";
+import type * as vscode from "vscode";
 
 import type { CoderApi } from "../api/coderApi";
 import type { ServiceContainer } from "../core/container";
@@ -48,6 +47,8 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 
 	private agent: { id: string; name: string } | undefined;
 	private workspace: Workspace | undefined;
+	/** Build number of the update we queued, once one has been accepted. */
+	private updatedBuildNumber: number | undefined;
 
 	private readonly logger: Logger;
 
@@ -78,6 +79,18 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		workspace: Workspace,
 		progress: vscode.Progress<{ message?: string }>,
 	): Promise<boolean> {
+		if (this.updatedBuildNumber !== undefined) {
+			// Monitor events queued before the update must not resolve the
+			// connection or move log streaming back to the previous build.
+			if (workspace.latest_build.build_number < this.updatedBuildNumber) {
+				return false;
+			}
+			this.updatedBuildNumber = workspace.latest_build.build_number;
+		}
+		if (workspace.latest_build.id !== this.workspace?.latest_build.id) {
+			// Build logs stream from one build, so a new build needs a new stream.
+			this.buildLogStream.close();
+		}
 		this.workspace = workspace;
 		const workspaceName = createWorkspaceIdentifier(workspace);
 
@@ -90,17 +103,29 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 					progress,
 				);
 				if (updated) {
-					workspace = updated;
 					// Agent IDs may have changed after an update.
 					this.resetAgent();
-					if (workspace.latest_build.status !== "running") return false;
+					return this.processWorkspace(updated, progress);
 				}
 				break;
 			}
 
 			case "stopped":
 			case "failed": {
 				this.buildLogStream.close();
+				if (this.updatedBuildNumber !== undefined) {
+					if (workspace.latest_build.status === "failed") {
+						throw new Error(
+							`Update failed for ${workspaceName}. Check the workspace in the dashboard before retrying.`,
+						);
+					}
+					// The server starts the workspace once the stop build
+					// succeeds; starting it here would race that build.
+					progress.report({
+						message: `waiting for the server to start ${workspaceName}...`,
+					});
+					return false;
+				}
 
 				if (this.startupMode === "none") {
 					const choice = await this.confirmStartOrUpdate(
@@ -119,13 +144,11 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 					progress,
 				);
 				if (updated) {
-					workspace = updated;
 					// Agent IDs may have changed after an update.
 					this.resetAgent();
-					if (workspace.latest_build.status !== "running") return false;
-					break;
+					return this.processWorkspace(updated, progress);
 				}
-				// Either we weren't in update mode, or the update failed: start.
+				// Start only when no update was requested.
 				await this.triggerStart(workspace, workspaceName, progress);
 				return false;
 			}
@@ -289,7 +312,7 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		this.logger.info(`${workspaceName} start initiated`);
 	}
 
-	/** No-op if not in update mode. Falls through to start on failure. */
+	/** No-op outside update mode; asks before falling back to the old version. */
 	private async maybeUpdate(
 		workspace: Workspace,
 		workspaceName: string,
@@ -309,8 +332,6 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 			this.workspace = await this.operationTelemetry.traceUpdate(() =>
 				updateWorkspace(this.buildCliContext(workspace), parameters),
 			);
-			this.logger.info(`${workspaceName} update initiated`);
-			return this.workspace;
 		} catch (error) {
 			if (error instanceof WorkspaceUpdateCancelledError) {
 				this.logger.info(
@@ -320,11 +341,36 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 			}
 			const reason = errToStr(error);
 			this.logger.warn(`Update failed for ${workspaceName}: ${reason}`);
-			vscode.window.showWarningMessage(
-				`Workspace update failed: ${reason}. Continuing with the existing version.`,
+			const connect = await this.operationTelemetry.traceFailurePrompt(() =>
+				this.confirmConnectToExisting(workspaceName, reason),
 			);
+			if (!connect) {
+				throw error;
+			}
+			this.logger.info(`Connecting to the existing ${workspaceName} version`);
 			return undefined;
 		}
+		this.updatedBuildNumber = this.workspace.latest_build.build_number;
+		this.logger.info(`${workspaceName} update initiated`);
+		return this.workspace;
+	}
+
+	/** Offers the existing version when the update could not be queued. */
+	private async confirmConnectToExisting(
+		workspaceName: string,
+		reason: string,
+	): Promise<boolean> {
+		const action = "Connect Anyway";
+		const choice = await vscodeProposed.window.showWarningMessage(
+			`Failed to update ${workspaceName}`,
+			{
+				useCustom: true,
+				modal: true,
+				detail: `${reason}\n\nTo connect without updating, choose ${action}.`,
+			},
+			action,
+		);
+		return choice === action;
 	}
 
 	private async confirmStartOrUpdate(