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

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(