fix: run workspace updates in interactive task

src/api/workspace.ts

@@ -1,8 +1,8 @@
 import { type Api } from "coder/site/src/api/api";
 import {
-	type WorkspaceAgentLog,
 	type ProvisionerJobLog,
 	type Workspace,
+	type WorkspaceAgentLog,
 } from "coder/site/src/api/typesGenerated";
 import { spawn } from "node:child_process";
 import * as vscode from "vscode";
@@ -115,16 +115,18 @@ export async function startWorkspace(ctx: CliContext): Promise<Workspace> {
 /**
  * Update a workspace to the latest template version.
  *
- * Uses `coder update` when the CLI supports it (>= 2.24).
- * Falls back to the REST API: stop, wait, then updateWorkspaceVersion.
+ * Uses an interactive VS Code task so the CLI can handle parameter prompts.
  */
 export async function updateWorkspace(ctx: CliContext): Promise<Workspace> {
-	if (ctx.featureSet.cliUpdate) {
-		await runCliCommand(ctx, ["update"]);
-		return ctx.restClient.getWorkspace(ctx.workspace.id);
+	if (!ctx.featureSet.cliUpdate) {
+		return updateWorkspaceVersion(ctx);
 	}
 
-	// REST API fallback for older CLIs.
+	await runCliCommandInTask(ctx, ["update"]);
+	return ctx.restClient.getWorkspace(ctx.workspace.id);
+}
+
+async function updateWorkspaceVersion(ctx: CliContext): Promise<Workspace> {
 	if (ctx.workspace.latest_build.status === "running") {
 		ctx.write("Stopping workspace for update...\r\n");
 		const stopBuild = await ctx.restClient.stopWorkspace(ctx.workspace.id);
@@ -139,6 +141,71 @@ export async function updateWorkspace(ctx: CliContext): Promise<Workspace> {
 	return ctx.restClient.getWorkspace(ctx.workspace.id);
 }
 
+function runCliCommandInTask(ctx: CliContext, args: string[]): Promise<void> {
+	return new Promise((resolve, reject) => {
+		const workspaceIdentifier = createWorkspaceIdentifier(ctx.workspace);
+		const fullArgs = [
+			...getGlobalShellFlags(vscode.workspace.getConfiguration(), ctx.auth),
+			...args,
+			workspaceIdentifier,
+		];
+		const cmd = `${escapeCommandArg(ctx.binPath)} ${fullArgs.join(" ")}`;
+		const task = new vscode.Task(
+			{ type: "coder", command: args[0] },
+			vscode.TaskScope.Workspace,
+			`Coder: ${args[0]} ${workspaceIdentifier}`,
+			"coder",
+			new vscode.ShellExecution(cmd),
+			[],
+		);
+		task.presentationOptions = {
+			reveal: vscode.TaskRevealKind.Always,
+			panel: vscode.TaskPanelKind.Dedicated,
+			focus: true,
+			clear: true,
+		};
+
+		let execution: vscode.TaskExecution | undefined;
+		let settled = false;
+
+		function finish(exitCode: number | undefined) {
+			if (settled) return;
+			settled = true;
+			disposable.dispose();
+			if (exitCode === 0) {
+				resolve();
+			} else {
+				reject(
+					new Error(`"${fullArgs.join(" ")}" exited with code ${exitCode}`),
+				);
+			}
+		}
+
+		function fail(error: unknown) {
+			if (settled) return;
+			settled = true;
+			disposable.dispose();
+			reject(error instanceof Error ? error : new Error(String(error)));
+		}
+
+		const disposable = vscode.tasks.onDidEndTaskProcess((event) => {
+			if (execution) {
+				if (event.execution !== execution) {
+					return;
+				}
+			} else if (event.execution.task !== task) {
+				return;
+			}
+
+			finish(event.exitCode);
+		});
+
+		vscode.tasks.executeTask(task).then((taskExecution) => {
+			execution = taskExecution;
+		}, fail);
+	});
+}
+
 /**
  * Streams build logs in real-time via a callback.
  * Returns the websocket for lifecycle management.

src/remote/remote.ts

@@ -406,7 +406,7 @@ export class Remote {
 									);
 									if (isReady) {
 										subscription.dispose();
-										resolve(w);
+										resolve(stateMachine.getWorkspace() ?? w);
 										return;
 									}
 								} catch (error: unknown) {

src/remote/workspaceStateMachine.ts

@@ -35,6 +35,7 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 	private readonly agentLogStream = new LazyStream<WorkspaceAgentLog[]>();
 
 	private agent: { id: string; name: string } | undefined;
+	private workspace: Workspace | undefined;
 
 	constructor(
 		private readonly parts: AuthorityParts,
@@ -56,13 +57,19 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		workspace: Workspace,
 		progress: vscode.Progress<{ message?: string }>,
 	): Promise<boolean> {
+		this.workspace = workspace;
 		const workspaceName = createWorkspaceIdentifier(workspace);
 
 		switch (workspace.latest_build.status) {
 			case "running":
 				this.buildLogStream.close();
 				if (this.startupMode === "update") {
-					await this.triggerUpdate(workspace, workspaceName, progress);
+					workspace = await this.triggerUpdate(
+						workspace,
+						workspaceName,
+						progress,
+					);
+					this.workspace = workspace;
 					// Agent IDs may have changed after an update.
 					this.agent = undefined;
 				}
@@ -84,7 +91,15 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 				}
 
 				if (this.startupMode === "update") {
-					await this.triggerUpdate(workspace, workspaceName, progress);
+					workspace = await this.triggerUpdate(
+						workspace,
+						workspaceName,
+						progress,
+					);
+					this.workspace = workspace;
+					if (workspace.latest_build.status === "running") {
+						break;
+					}
 				} else {
 					await this.triggerStart(workspace, workspaceName, progress);
 				}
@@ -252,16 +267,19 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		workspace: Workspace,
 		workspaceName: string,
 		progress: vscode.Progress<{ message?: string }>,
-	): Promise<void> {
+	): Promise<Workspace> {
 		progress.report({ message: `updating ${workspaceName}...` });
 		this.logger.info(`Updating ${workspaceName}`, {
 			mode: this.startupMode,
 			status: workspace.latest_build.status,
 		});
-		await updateWorkspace(this.buildCliContext(workspace));
+		const updatedWorkspace = await updateWorkspace(
+			this.buildCliContext(workspace),
+		);
 		// Downgrade so subsequent transitions don't re-trigger the update.
 		this.startupMode = "start";
 		this.logger.info(`${workspaceName} update initiated`);
+		return updatedWorkspace;
 	}
 
 	private async confirmStartOrUpdate(
@@ -286,6 +304,10 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		return this.agent?.id;
 	}
 
+	public getWorkspace(): Workspace | undefined {
+		return this.workspace;
+	}
+
 	dispose(): void {
 		this.buildLogStream.close();
 		this.agentLogStream.close();

test/mocks/vscode.runtime.ts

@@ -47,6 +47,9 @@ export const InputBoxValidationSeverity = E({
 	Warning: 2,
 	Error: 3,
 });
+export const TaskScope = E({ Global: 1, Workspace: 2 });
+export const TaskRevealKind = E({ Always: 1, Silent: 2, Never: 3 });
+export const TaskPanelKind = E({ Shared: 1, Dedicated: 2, New: 3 });
 
 export class MarkdownString {
 	value: string;
@@ -65,6 +68,25 @@ export class ThemeColor {
 	}
 }
 
+export class ShellExecution {
+	commandLine: string | undefined;
+	constructor(commandLine: string) {
+		this.commandLine = commandLine;
+	}
+}
+
+export class Task {
+	presentationOptions: unknown;
+	constructor(
+		public definition: unknown,
+		public scope: unknown,
+		public name: string,
+		public source: string,
+		public execution?: unknown,
+		public problemMatchers: string[] = [],
+	) {}
+}
+
 export class Uri {
 	constructor(
 		public scheme: string,
@@ -114,6 +136,7 @@ export class EventEmitter<T> {
 const onDidChangeConfiguration = new EventEmitter<unknown>();
 const onDidChangeWorkspaceFolders = new EventEmitter<unknown>();
 const onDidChangeActiveColorTheme = new EventEmitter<unknown>();
+const onDidEndTaskProcess = new EventEmitter<unknown>();
 
 export const window = {
 	activeColorTheme: { kind: ColorThemeKind.Dark },
@@ -148,6 +171,12 @@ export const commands = {
 	executeCommand: vi.fn(),
 };
 
+export const tasks = {
+	executeTask: vi.fn(),
+	onDidEndTaskProcess: onDidEndTaskProcess.event,
+	__fireDidEndTaskProcess: (e: unknown) => onDidEndTaskProcess.fire(e),
+};
+
 export const workspace = {
 	getConfiguration: vi.fn(), // your helpers override this
 	workspaceFolders: [] as unknown[],
@@ -197,12 +226,18 @@ const vscode = {
 	ExtensionMode,
 	UIKind,
 	InputBoxValidationSeverity,
+	TaskScope,
+	TaskRevealKind,
+	TaskPanelKind,
+	ShellExecution,
+	Task,
 	Uri,
 	EventEmitter,
 	MarkdownString,
 	ThemeColor,
 	window,
 	commands,
+	tasks,
 	workspace,
 	env,
 	extensions,