refactor: make confirm dialog update-aware and improve test quality

Make the stopped-workspace dialog offer "Update and Start" when
the workspace is outdated, so users can update without going
through a separate command. Inline the wantsUpdate getter since
it was just hiding a simple equality check on startupMode.

Rework the test file to use MockUserInteraction (instead of
mocking vscodeProposed directly), MockTerminalSession with
content capture, proper agent/resource factories from the shared
workspace mock, and a setup() function that returns all test
fixtures. Add getMessageCalls() to MockUserInteraction so tests
can inspect which buttons were offered in dialogs.

Add edge case coverage for agent timeout, created lifecycle,
start_timeout, all shutdown states, and getAgentId.

Author: EhabY

src/remote/workspaceStateMachine.ts

@@ -61,7 +61,7 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		switch (workspace.latest_build.status) {
 			case "running":
 				this.buildLogStream.close();
-				if (this.wantsUpdate) {
+				if (this.startupMode === "update") {
 					await this.triggerUpdate(workspace, workspaceName, progress);
 					// Agent IDs may have changed after an update.
 					this.agent = undefined;
@@ -72,14 +72,18 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 			case "failed": {
 				this.buildLogStream.close();
 
-				if (
-					this.startupMode === "none" &&
-					!(await this.confirmStart(workspaceName))
-				) {
-					throw new Error(`Workspace start cancelled`);
+				if (this.startupMode === "none") {
+					const choice = await this.confirmStartOrUpdate(
+						workspaceName,
+						workspace.outdated,
+					);
+					if (!choice) {
+						throw new Error(`Workspace start cancelled`);
+					}
+					this.startupMode = choice;
 				}
 
-				if (this.wantsUpdate) {
+				if (this.startupMode === "update") {
 					await this.triggerUpdate(workspace, workspaceName, progress);
 				} else {
 					await this.triggerStart(workspace, workspaceName, progress);
@@ -90,7 +94,7 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 			case "pending":
 			case "starting":
 			case "stopping": {
-				// Clear the agent since it's ID could change after a restart
+				// Clear the agent since its ID could change after a restart
 				this.agent = undefined;
 				this.agentLogStream.close();
 				progress.report({
@@ -219,10 +223,6 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		}
 	}
 
-	private get wantsUpdate(): boolean {
-		return this.startupMode === "update";
-	}
-
 	private buildCliContext(workspace: Workspace) {
 		return {
 			restClient: this.workspaceClient,
@@ -264,16 +264,22 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 		this.logger.info(`${workspaceName} update initiated`);
 	}
 
-	private async confirmStart(workspaceName: string): Promise<boolean> {
+	private async confirmStartOrUpdate(
+		workspaceName: string,
+		outdated: boolean,
+	): Promise<"start" | "update" | undefined> {
+		const buttons = outdated ? ["Start", "Update and Start"] : ["Start"];
 		const action = await vscodeProposed.window.showInformationMessage(
-			`Unable to connect to the workspace ${workspaceName} because it is not running. Start the workspace?`,
+			`The workspace ${workspaceName} is not running. How would you like to proceed?`,
 			{
 				useCustom: true,
 				modal: true,
 			},
-			"Start",
+			...buttons,
 		);
-		return action === "Start";
+		if (action === "Start") return "start";
+		if (action === "Update and Start") return "update";
+		return undefined;
 	}
 
 	public getAgentId(): string | undefined {

test/mocks/testHelpers.ts

@@ -173,12 +173,20 @@ export class MockProgressReporter {
 	}
 }
 
+/** A recorded call to one of the vscode.window.show*Message methods. */
+export interface MessageCall {
+	level: "information" | "warning" | "error";
+	message: string;
+	items: string[];
+}
+
 /**
  * Mock user interaction that integrates with vscode.window message dialogs and input boxes.
- * Use this to control user responses in tests.
+ * Use this to control user responses and inspect dialog calls in tests.
  */
 export class MockUserInteraction {
 	private readonly responses = new Map<string, string | undefined>();
+	private readonly _messageCalls: MessageCall[] = [];
 	private inputBoxValue: string | undefined;
 	private inputBoxValidateInput: ((value: string) => Promise<void>) | undefined;
 	private externalUrls: string[] = [];
@@ -194,6 +202,13 @@ export class MockUserInteraction {
 		this.responses.set(message, response);
 	}
 
+	/**
+	 * Get all message dialog calls that were made (across all levels).
+	 */
+	getMessageCalls(): readonly MessageCall[] {
+		return this._messageCalls;
+	}
+
 	/**
 	 * Set the value to return from showInputBox.
 	 * Pass undefined to simulate user cancelling.
@@ -229,6 +244,7 @@ export class MockUserInteraction {
 	 */
 	clear(): void {
 		this.responses.clear();
+		this._messageCalls.length = 0;
 		this.inputBoxValue = undefined;
 		this.inputBoxValidateInput = undefined;
 		this.externalUrls = [];
@@ -242,20 +258,27 @@ export class MockUserInteraction {
 			return this.responses.get(message);
 		};
 
-		// eslint-disable-next-line @typescript-eslint/no-explicit-any
-		const handleMessage = (message: string): Thenable<any> => {
-			const response = getResponse(message);
-			return Promise.resolve(response);
-		};
-
-		vi.mocked(vscode.window.showErrorMessage).mockImplementation(handleMessage);
+		const handleMessage =
+			(level: MessageCall["level"]) =>
+			// eslint-disable-next-line @typescript-eslint/no-explicit-any -- serves all show*Message overloads
+			(message: string, ...rest: unknown[]): Thenable<any> => {
+				const items = rest.filter(
+					(arg): arg is string => typeof arg === "string",
+				);
+				this._messageCalls.push({ level, message, items });
+				return Promise.resolve(getResponse(message));
+			};
+
+		vi.mocked(vscode.window.showErrorMessage).mockImplementation(
+			handleMessage("error"),
+		);
 
 		vi.mocked(vscode.window.showWarningMessage).mockImplementation(
-			handleMessage,
+			handleMessage("warning"),
 		);
 
 		vi.mocked(vscode.window.showInformationMessage).mockImplementation(
-			handleMessage,
+			handleMessage("information"),
 		);
 
 		vi.mocked(vscode.env.openExternal).mockImplementation(
@@ -922,3 +945,44 @@ export class MockContextManager {
 
 	readonly dispose = vi.fn();
 }
+
+/**
+ * Mock TerminalSession that captures all content written to the terminal.
+ * Use `lastInstance` to get the most recently created instance (set in the constructor),
+ * which is useful when the real TerminalSession is created inside the class under test.
+ */
+export class MockTerminalSession {
+	static lastInstance: MockTerminalSession | undefined;
+
+	private readonly _lines: string[] = [];
+
+	readonly writeEmitter = {
+		fire: vi.fn((data: string) => {
+			this._lines.push(data);
+		}),
+		event: vi.fn(),
+		dispose: vi.fn(),
+	};
+	readonly terminal = { show: vi.fn(), dispose: vi.fn() };
+	readonly dispose = vi.fn();
+
+	constructor(_name?: string) {
+		MockTerminalSession.lastInstance = this;
+	}
+
+	/** All lines written via writeEmitter.fire(). */
+	get lines(): readonly string[] {
+		return this._lines;
+	}
+
+	/** Concatenated terminal content. */
+	get content(): string {
+		return this._lines.join("");
+	}
+
+	/** Reset captured content and mock call history. */
+	clear(): void {
+		this._lines.length = 0;
+		this.writeEmitter.fire.mockClear();
+	}
+}

test/mocks/workspace.ts

@@ -1,10 +1,12 @@
 /**
- * Test factory for Coder SDK Workspace type.
+ * Test factories for Coder SDK workspace types.
  */
 
 import type {
 	Workspace,
+	WorkspaceAgent,
 	WorkspaceBuild,
+	WorkspaceResource,
 } from "coder/site/src/api/typesGenerated";
 
 const defaultBuild: WorkspaceBuild = {
@@ -92,3 +94,54 @@ export function workspace(
 		...rest,
 	};
 }
+
+/** Create a WorkspaceAgent with sensible defaults for a connected, ready agent. */
+export function agent(overrides: Partial<WorkspaceAgent> = {}): WorkspaceAgent {
+	return {
+		id: "agent-1",
+		parent_id: null,
+		created_at: "2024-01-01T00:00:00Z",
+		updated_at: "2024-01-01T00:00:00Z",
+		status: "connected",
+		lifecycle_state: "ready",
+		name: "main",
+		resource_id: "resource-1",
+		architecture: "amd64",
+		environment_variables: {},
+		operating_system: "linux",
+		logs_length: 0,
+		logs_overflowed: false,
+		version: "2.25.0",
+		api_version: "1.0",
+		apps: [],
+		connection_timeout_seconds: 120,
+		troubleshooting_url: "",
+		subsystems: [],
+		health: { healthy: true },
+		display_apps: [],
+		log_sources: [],
+		scripts: [],
+		startup_script_behavior: "non-blocking",
+		...overrides,
+	};
+}
+
+/** Create a WorkspaceResource with sensible defaults. */
+export function resource(
+	overrides: Partial<WorkspaceResource> = {},
+): WorkspaceResource {
+	return {
+		id: "resource-1",
+		created_at: "2024-01-01T00:00:00Z",
+		job_id: "job-1",
+		workspace_transition: "start",
+		type: "docker_container",
+		name: "main",
+		hide: false,
+		icon: "",
+		agents: [],
+		metadata: [],
+		daily_cost: 0,
+		...overrides,
+	};
+}