feat: use startup mode for workspace updates

src/api/workspace.ts

@@ -45,66 +45,49 @@ export class LazyStream<T> {
 	}
 }
 
+interface CliContext {
+	restClient: Api;
+	auth: CliAuth;
+	binPath: string;
+	workspace: Workspace;
+	writeEmitter: vscode.EventEmitter<string>;
+	featureSet: FeatureSet;
+}
+
 /**
- * Start or update a workspace and return the updated workspace.
+ * Spawn a Coder CLI subcommand and stream its output.
+ * Resolves when the process exits successfully; rejects on non-zero exit.
  */
-export async function startWorkspaceIfStoppedOrFailed(
-	restClient: Api,
-	auth: CliAuth,
-	binPath: string,
-	workspace: Workspace,
-	writeEmitter: vscode.EventEmitter<string>,
-	featureSet: FeatureSet,
-): Promise<Workspace> {
-	// Before we start a workspace, we make an initial request to check it's not already started
-	const updatedWorkspace = await restClient.getWorkspace(workspace.id);
-
-	if (!["stopped", "failed"].includes(updatedWorkspace.latest_build.status)) {
-		return updatedWorkspace;
-	}
-
+function runCliCommand(ctx: CliContext, args: string[]): Promise<void> {
 	return new Promise((resolve, reject) => {
-		const startArgs = [
-			...getGlobalShellFlags(vscode.workspace.getConfiguration(), auth),
-			"start",
-			"--yes",
-			createWorkspaceIdentifier(workspace),
+		const fullArgs = [
+			...getGlobalShellFlags(vscode.workspace.getConfiguration(), ctx.auth),
+			...args,
+			createWorkspaceIdentifier(ctx.workspace),
 		];
-		if (featureSet.buildReason) {
-			startArgs.push("--reason", "vscode_connection");
-		}
 
-		// { shell: true } requires one shell-safe command string, otherwise we lose all escaping
-		const cmd = `${escapeCommandArg(binPath)} ${startArgs.join(" ")}`;
-		const startProcess = spawn(cmd, { shell: true });
-
-		startProcess.stdout.on("data", (data: Buffer) => {
-			const lines = data
-				.toString()
-				.split(/\r*\n/)
-				.filter((line) => line !== "");
-			for (const line of lines) {
-				writeEmitter.fire(line.toString() + "\r\n");
+		const cmd = `${escapeCommandArg(ctx.binPath)} ${fullArgs.join(" ")}`;
+		const proc = spawn(cmd, { shell: true });
+
+		proc.stdout.on("data", (data: Buffer) => {
+			for (const line of splitLines(data)) {
+				ctx.writeEmitter.fire(line + "\r\n");
 			}
 		});
 
 		let capturedStderr = "";
-		startProcess.stderr.on("data", (data: Buffer) => {
-			const lines = data
-				.toString()
-				.split(/\r*\n/)
-				.filter((line) => line !== "");
-			for (const line of lines) {
-				writeEmitter.fire(line.toString() + "\r\n");
-				capturedStderr += line.toString() + "\n";
+		proc.stderr.on("data", (data: Buffer) => {
+			for (const line of splitLines(data)) {
+				ctx.writeEmitter.fire(line + "\r\n");
+				capturedStderr += line + "\n";
 			}
 		});
 
-		startProcess.on("close", (code: number) => {
+		proc.on("close", (code: number) => {
 			if (code === 0) {
-				resolve(restClient.getWorkspace(workspace.id));
+				resolve();
 			} else {
-				let errorText = `"${startArgs.join(" ")}" exited with code ${code}`;
+				let errorText = `"${fullArgs.join(" ")}" exited with code ${code}`;
 				if (capturedStderr !== "") {
 					errorText += `: ${capturedStderr}`;
 				}
@@ -114,6 +97,58 @@ export async function startWorkspaceIfStoppedOrFailed(
 	});
 }
 
+function splitLines(data: Buffer): string[] {
+	return data
+		.toString()
+		.split(/\r*\n/)
+		.filter((line) => line !== "");
+}
+
+/**
+ * Start a stopped or failed workspace using `coder start`.
+ * No-ops if the workspace is already running.
+ */
+export async function startWorkspace(ctx: CliContext): Promise<Workspace> {
+	if (!["stopped", "failed"].includes(ctx.workspace.latest_build.status)) {
+		return ctx.workspace;
+	}
+
+	const args = ["start", "--yes"];
+	if (ctx.featureSet.buildReason) {
+		args.push("--reason", "vscode_connection");
+	}
+
+	await runCliCommand(ctx, args);
+	return ctx.restClient.getWorkspace(ctx.workspace.id);
+}
+
+/**
+ * Update a workspace to the latest template version.
+ *
+ * Uses `coder update` when the CLI supports it (>= 2.25).
+ * Falls back to the REST API: stop, wait, then updateWorkspaceVersion.
+ */
+export async function updateWorkspace(ctx: CliContext): Promise<Workspace> {
+	if (ctx.featureSet.cliUpdate) {
+		await runCliCommand(ctx, ["update"]);
+		return ctx.restClient.getWorkspace(ctx.workspace.id);
+	}
+
+	// REST API fallback for older CLIs.
+	if (ctx.workspace.latest_build.status === "running") {
+		ctx.writeEmitter.fire("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 canceled during stop");
+		}
+	}
+
+	ctx.writeEmitter.fire("Starting workspace with updated template...\r\n");
+	await ctx.restClient.updateWorkspaceVersion(ctx.workspace);
+	return ctx.restClient.getWorkspace(ctx.workspace.id);
+}
+
 /**
  * Streams build logs in real-time via a callback.
  * Returns the websocket for lifecycle management.

src/commands.ts

@@ -741,7 +741,7 @@ export class Commands {
 		}
 
 		// Only set the memento when opening a new folder
-		await this.mementoManager.setFirstConnect();
+		await this.mementoManager.setStartupMode("start");
 		await vscode.commands.executeCommand(
 			"vscode.openFolder",
 			vscode.Uri.from({
@@ -770,9 +770,15 @@ export class Commands {
 			},
 			"Update and Restart",
 		);
-		if (action === "Update and Restart") {
-			await this.remoteWorkspaceClient.updateWorkspaceVersion(this.workspace);
+		if (action !== "Update and Restart") {
+			return;
 		}
+
+		this.logger.info(
+			`Updating workspace ${createWorkspaceIdentifier(this.workspace)}`,
+		);
+		await this.mementoManager.setStartupMode("update");
+		await vscode.commands.executeCommand("workbench.action.reloadWindow");
 	}
 
 	public async pingWorkspace(item?: OpenableTreeItem): Promise<void> {
@@ -1041,7 +1047,7 @@ export class Commands {
 		}
 
 		// Only set the memento when opening a new folder/window
-		await this.mementoManager.setFirstConnect();
+		await this.mementoManager.setStartupMode("start");
 		if (folderPath) {
 			await vscode.commands.executeCommand(
 				"vscode.openFolder",

src/core/mementoManager.ts

@@ -7,6 +7,15 @@ const MAX_URLS = 10;
 // state from crashes or interrupted reloads.
 const PENDING_TTL_MS = 5 * 60 * 1000;
 
+/**
+ * Describes the startup intent when the extension connects to a workspace.
+ * - "none":    No explicit intent; ask before starting a stopped workspace.
+ * - "start":   User-initiated open/restart; auto-start without prompting.
+ * - "update":  User-initiated restart + update; use `coder update` to apply
+ *              the latest template version, auto-starting without prompting.
+ */
+export type StartupMode = "none" | "start" | "update";
+
 interface Stamped<T> {
 	value: T;
 	setAt: number;
@@ -46,25 +55,21 @@ export class MementoManager {
 			: Array.from(urls);
 	}
 
-	/**
-	 * Mark this as the first connection to a workspace, which influences whether
-	 * the workspace startup confirmation is shown to the user.
-	 */
-	public async setFirstConnect(): Promise<void> {
-		return this.setStamped("firstConnect", true);
+	/** Set the startup mode for the next workspace connection. */
+	public async setStartupMode(mode: StartupMode): Promise<void> {
+		await this.setStamped("startupMode", mode);
 	}
 
 	/**
-	 * Check if this is the first connection to a workspace and clear the flag.
-	 * Used to determine whether to automatically start workspaces without
-	 * prompting the user for confirmation.
+	 * Read and clear the startup mode.
+	 * Returns "none" (the default) when no mode was explicitly set.
 	 */
-	public async getAndClearFirstConnect(): Promise<boolean> {
-		const value = this.getStamped<boolean>("firstConnect");
+	public async getAndClearStartupMode(): Promise<StartupMode> {
+		const value = this.getStamped<StartupMode>("startupMode");
 		if (value !== undefined) {
-			await this.memento.update("firstConnect", undefined);
+			await this.memento.update("startupMode", undefined);
 		}
-		return value === true;
+		return value ?? "none";
 	}
 
 	/** Store a chat ID to open after a remote-authority reload. */

src/extension.ts

@@ -70,8 +70,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	// Migrate auth storage from old flat format to new label-based format
 	await migrateAuthStorage(serviceContainer);
 
-	// Try to clear this flag ASAP
-	const isFirstConnect = await mementoManager.getAndClearFirstConnect();
+	// Clear and capture the startup mode before anything else.
+	const startupMode = await mementoManager.getAndClearStartupMode();
 
 	const deployment = await secretsManager.getCurrentDeployment();
 
@@ -348,7 +348,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		try {
 			const details = await remote.setup(
 				vscodeProposed.env.remoteAuthority,
-				isFirstConnect,
+				startupMode,
 				remoteSshExtension.id,
 			);
 			if (details) {

src/featureSet.ts

@@ -5,6 +5,7 @@ export interface FeatureSet {
 	proxyLogDirectory: boolean;
 	wildcardSSH: boolean;
 	buildReason: boolean;
+	cliUpdate: boolean;
 	keyringAuth: boolean;
 	keyringTokenRead: boolean;
 	supportBundle: boolean;
@@ -44,6 +45,8 @@ export function featureSetForVersion(
 		wildcardSSH: versionAtLeast(version, "2.19.0"),
 		// --reason flag for `coder start`
 		buildReason: versionAtLeast(version, "2.25.0"),
+		// `coder update` subcommand
+		cliUpdate: versionAtLeast(version, "2.25.0"),
 		// Keyring-backed token storage via `coder login`
 		keyringAuth: versionAtLeast(version, "2.29.0"),
 		// `coder login token` for reading tokens from the keyring

src/remote/remote.ts

@@ -26,6 +26,7 @@ import { version as cliVersion } from "../core/cliExec";
 import { type CliManager } from "../core/cliManager";
 import { type ServiceContainer } from "../core/container";
 import { type ContextManager } from "../core/contextManager";
+import { type StartupMode } from "../core/mementoManager";
 import { type PathResolver } from "../core/pathResolver";
 import { type SecretsManager } from "../core/secretsManager";
 import { toError } from "../error/errorUtils";
@@ -97,7 +98,7 @@ export class Remote {
 	 */
 	public async setup(
 		remoteAuthority: string,
-		firstConnect: boolean,
+		startupMode: StartupMode,
 		remoteSshExtensionId: string,
 	): Promise<RemoteDetails | undefined> {
 		const parts = parseRemoteAuthority(remoteAuthority);
@@ -165,11 +166,7 @@ export class Remote {
 				});
 				if (result.success) {
 					// Login successful, retry setup
-					return this.setup(
-						remoteAuthority,
-						firstConnect,
-						remoteSshExtensionId,
-					);
+					return this.setup(remoteAuthority, startupMode, remoteSshExtensionId);
 				} else {
 					// User cancelled or login failed
 					await this.closeRemote();
@@ -372,7 +369,7 @@ export class Remote {
 			const stateMachine = new WorkspaceStateMachine(
 				parts,
 				workspaceClient,
-				firstConnect,
+				startupMode,
 				binaryPath,
 				featureSet,
 				this.logger,