From c1652f6b22575a16f2bdd2689d7bb44c002c7ad1 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 14 Apr 2026 13:00:36 +0300 Subject: [PATCH 1/3] feat: use startup mode for workspace updates instead of direct REST calls Replace the boolean firstConnect flag with a typed StartupMode ("prompt" | "start" | "update") so the extension can tell the workspace state machine exactly what to do on reconnect. The "Update Workspace" command now sets the mode to "update" and reloads the window, letting the state machine handle the update via `coder update` (or the REST fallback for older CLIs). Also refactors workspace.ts to extract a shared runCliCommand helper used by both startWorkspace and updateWorkspace. --- src/api/workspace.ts | 135 +++++--- src/commands.ts | 14 +- src/core/mementoManager.ts | 31 +- src/extension.ts | 6 +- src/featureSet.ts | 3 + src/remote/remote.ts | 11 +- src/remote/workspaceStateMachine.ts | 84 ++++- test/unit/core/mementoManager.test.ts | 26 +- .../unit/remote/workspaceStateMachine.test.ts | 327 ++++++++++++++++++ 9 files changed, 540 insertions(+), 97 deletions(-) create mode 100644 test/unit/remote/workspaceStateMachine.test.ts diff --git a/src/api/workspace.ts b/src/api/workspace.ts index 698d212a..677a539c 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -45,66 +45,47 @@ export class LazyStream { } } +interface CliContext { + auth: CliAuth; + binPath: string; + workspace: Workspace; + writeEmitter: vscode.EventEmitter; +} + /** - * 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, - featureSet: FeatureSet, -): Promise { - // 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 { 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 +95,68 @@ 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( + restClient: Api, + ctx: CliContext, + featureSet: FeatureSet, +): Promise { + const current = await restClient.getWorkspace(ctx.workspace.id); + if (!["stopped", "failed"].includes(current.latest_build.status)) { + return current; + } + + const args = ["start", "--yes"]; + if (featureSet.buildReason) { + args.push("--reason", "vscode_connection"); + } + + await runCliCommand(ctx, args); + return 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 → updateWorkspaceVersion. + */ +export async function updateWorkspace( + restClient: Api, + ctx: CliContext, + featureSet: FeatureSet, +): Promise { + if (featureSet.cliUpdate) { + await runCliCommand(ctx, ["update"]); + return restClient.getWorkspace(ctx.workspace.id); + } + + // REST API fallback for older CLIs. + const workspace = await restClient.getWorkspace(ctx.workspace.id); + if (workspace.latest_build.status === "running") { + ctx.writeEmitter.fire("Stopping workspace for update...\r\n"); + const stopBuild = await restClient.stopWorkspace(workspace.id); + const stoppedJob = await 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 restClient.updateWorkspaceVersion(workspace); + return restClient.getWorkspace(ctx.workspace.id); +} + /** * Streams build logs in real-time via a callback. * Returns the websocket for lifecycle management. diff --git a/src/commands.ts b/src/commands.ts index eef2835d..b8a08aa5 100644 --- a/src/commands.ts +++ b/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 { @@ -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", diff --git a/src/core/mementoManager.ts b/src/core/mementoManager.ts index 4a83cdb1..1a5c9169 100644 --- a/src/core/mementoManager.ts +++ b/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. + * - "prompt": Normal reconnection; 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 = "prompt" | "start" | "update"; + interface Stamped { 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 { - return this.setStamped("firstConnect", true); + /** Set the startup mode for the next workspace connection. */ + public async setStartupMode(mode: StartupMode): Promise { + 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 "prompt" (the default) when no mode was explicitly set. */ - public async getAndClearFirstConnect(): Promise { - const value = this.getStamped("firstConnect"); + public async getAndClearStartupMode(): Promise { + const value = this.getStamped("startupMode"); if (value !== undefined) { - await this.memento.update("firstConnect", undefined); + await this.memento.update("startupMode", undefined); } - return value === true; + return value ?? "prompt"; } /** Store a chat ID to open after a remote-authority reload. */ diff --git a/src/extension.ts b/src/extension.ts index 8531d929..3b54d1c0 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -70,8 +70,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // 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 { try { const details = await remote.setup( vscodeProposed.env.remoteAuthority, - isFirstConnect, + startupMode, remoteSshExtension.id, ); if (details) { diff --git a/src/featureSet.ts b/src/featureSet.ts index 09ff4574..a5b924a4 100644 --- a/src/featureSet.ts +++ b/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 diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 30d2df37..e8e9b057 100644 --- a/src/remote/remote.ts +++ b/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 { 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, diff --git a/src/remote/workspaceStateMachine.ts b/src/remote/workspaceStateMachine.ts index b874b969..d265a74c 100644 --- a/src/remote/workspaceStateMachine.ts +++ b/src/remote/workspaceStateMachine.ts @@ -1,12 +1,12 @@ import { createWorkspaceIdentifier, extractAgents } from "../api/api-helper"; import { LazyStream, - startWorkspaceIfStoppedOrFailed, + startWorkspace, + updateWorkspace, streamAgentLogs, streamBuildLogs, } from "../api/workspace"; import { maybeAskAgent } from "../promptUtils"; -import { type AuthorityParts } from "../util"; import { vscodeProposed } from "../vscodeProposed"; import { TerminalSession } from "./terminalSession"; @@ -19,9 +19,11 @@ import type { import type * as vscode from "vscode"; import type { CoderApi } from "../api/coderApi"; +import type { StartupMode } from "../core/mementoManager"; import type { FeatureSet } from "../featureSet"; import type { Logger } from "../logging/logger"; import type { CliAuth } from "../settings/cli"; +import type { AuthorityParts } from "../util"; /** * Manages workspace and agent state transitions until ready for SSH connection. @@ -37,7 +39,7 @@ export class WorkspaceStateMachine implements vscode.Disposable { constructor( private readonly parts: AuthorityParts, private readonly workspaceClient: CoderApi, - private readonly firstConnect: boolean, + private startupMode: StartupMode, private readonly binaryPath: string, private readonly featureSet: FeatureSet, private readonly logger: Logger, @@ -59,27 +61,29 @@ export class WorkspaceStateMachine implements vscode.Disposable { switch (workspace.latest_build.status) { case "running": this.buildLogStream.close(); + if (this.wantsUpdate) { + await this.triggerUpdate(workspace, workspaceName, progress); + // Agent IDs may have changed after an update. + this.agent = undefined; + } break; case "stopped": case "failed": { this.buildLogStream.close(); - if (!this.firstConnect && !(await this.confirmStart(workspaceName))) { + if ( + this.startupMode === "prompt" && + !(await this.confirmStart(workspaceName)) + ) { throw new Error(`Workspace start cancelled`); } - progress.report({ message: `starting ${workspaceName}...` }); - this.logger.info(`Starting ${workspaceName}`); - await startWorkspaceIfStoppedOrFailed( - this.workspaceClient, - this.cliAuth, - this.binaryPath, - workspace, - this.terminal.writeEmitter, - this.featureSet, - ); - this.logger.info(`${workspaceName} status is now running`); + if (this.wantsUpdate) { + await this.triggerUpdate(workspace, workspaceName, progress); + } else { + await this.triggerStart(workspace, workspaceName, progress); + } return false; } @@ -215,6 +219,56 @@ export class WorkspaceStateMachine implements vscode.Disposable { } } + private get wantsUpdate(): boolean { + return this.startupMode === "update"; + } + + private get cliContext() { + return { + auth: this.cliAuth, + binPath: this.binaryPath, + writeEmitter: this.terminal.writeEmitter, + }; + } + + private async triggerStart( + workspace: Workspace, + workspaceName: string, + progress: vscode.Progress<{ message?: string }>, + ): Promise { + progress.report({ message: `starting ${workspaceName}...` }); + this.logger.info(`Starting ${workspaceName}`, { + mode: this.startupMode, + status: workspace.latest_build.status, + }); + await startWorkspace( + this.workspaceClient, + { ...this.cliContext, workspace }, + this.featureSet, + ); + this.logger.info(`${workspaceName} start initiated`); + } + + private async triggerUpdate( + workspace: Workspace, + workspaceName: string, + progress: vscode.Progress<{ message?: string }>, + ): Promise { + progress.report({ message: `updating ${workspaceName}...` }); + this.logger.info(`Updating ${workspaceName}`, { + mode: this.startupMode, + status: workspace.latest_build.status, + }); + await updateWorkspace( + this.workspaceClient, + { ...this.cliContext, workspace }, + this.featureSet, + ); + // Downgrade so subsequent transitions don't re-trigger the update. + this.startupMode = "start"; + this.logger.info(`${workspaceName} update initiated`); + } + private async confirmStart(workspaceName: string): Promise { const action = await vscodeProposed.window.showInformationMessage( `Unable to connect to the workspace ${workspaceName} because it is not running. Start the workspace?`, diff --git a/test/unit/core/mementoManager.test.ts b/test/unit/core/mementoManager.test.ts index 16c3efbe..3275915a 100644 --- a/test/unit/core/mementoManager.test.ts +++ b/test/unit/core/mementoManager.test.ts @@ -66,23 +66,31 @@ describe("MementoManager", () => { }); }); - describe("firstConnect", () => { - it("should return true only once", async () => { - await mementoManager.setFirstConnect(); + describe("startupMode", () => { + it("should return the set mode and clear after read", async () => { + await mementoManager.setStartupMode("start"); + expect(await mementoManager.getAndClearStartupMode()).toBe("start"); + expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); + }); + + it("should return 'prompt' when nothing is set", async () => { + expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); + }); - expect(await mementoManager.getAndClearFirstConnect()).toBe(true); - expect(await mementoManager.getAndClearFirstConnect()).toBe(false); + it("should support 'update' mode", async () => { + await mementoManager.setStartupMode("update"); + expect(await mementoManager.getAndClearStartupMode()).toBe("update"); }); it("should treat legacy bare values as expired", async () => { - await memento.update("firstConnect", true); - expect(await mementoManager.getAndClearFirstConnect()).toBe(false); + await memento.update("startupMode", "start"); + expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); }); it("should expire after 5 minutes", async () => { - await mementoManager.setFirstConnect(); + await mementoManager.setStartupMode("update"); vi.advanceTimersByTime(5 * 60 * 1000 + 1); - expect(await mementoManager.getAndClearFirstConnect()).toBe(false); + expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); }); }); diff --git a/test/unit/remote/workspaceStateMachine.test.ts b/test/unit/remote/workspaceStateMachine.test.ts new file mode 100644 index 00000000..af753eec --- /dev/null +++ b/test/unit/remote/workspaceStateMachine.test.ts @@ -0,0 +1,327 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { WorkspaceStateMachine } from "@/remote/workspaceStateMachine"; + +import { createMockLogger, MockProgress } from "../../mocks/testHelpers"; +import { workspace as createWorkspace } from "../../mocks/workspace"; + +import type { + Workspace, + WorkspaceAgent, + WorkspaceResource, +} from "coder/site/src/api/typesGenerated"; + +import type { CoderApi } from "@/api/coderApi"; +import type { FeatureSet } from "@/featureSet"; +import type { CliAuth } from "@/settings/cli"; +import type { AuthorityParts } from "@/util"; + +vi.mock("@/api/workspace", async (importActual) => { + const { LazyStream } = await importActual(); + return { + LazyStream, + startWorkspace: vi.fn().mockResolvedValue({}), + updateWorkspace: vi.fn().mockResolvedValue({}), + streamBuildLogs: vi.fn().mockResolvedValue({}), + streamAgentLogs: vi.fn().mockResolvedValue({}), + }; +}); + +vi.mock("@/promptUtils", () => ({ + maybeAskAgent: vi.fn(), +})); + +vi.mock("@/vscodeProposed", () => ({ + vscodeProposed: { + window: { showInformationMessage: vi.fn() }, + }, +})); + +vi.mock("@/remote/terminalSession", () => ({ + TerminalSession: vi.fn().mockImplementation(function () { + return { + writeEmitter: { fire: vi.fn(), event: vi.fn(), dispose: vi.fn() }, + terminal: { show: vi.fn(), dispose: vi.fn() }, + dispose: vi.fn(), + }; + }), +})); + +const { startWorkspace, updateWorkspace, streamBuildLogs } = + await import("@/api/workspace"); +const { maybeAskAgent } = await import("@/promptUtils"); +const { vscodeProposed } = await import("@/vscodeProposed"); + +function createAgent(overrides: Partial = {}): WorkspaceAgent { + return { + id: "agent-1", + name: "main", + status: "connected", + lifecycle_state: "ready", + scripts: [], + ...overrides, + } as unknown as WorkspaceAgent; +} + +function runningWorkspace( + agentOverrides: Partial = {}, +): Workspace { + return createWorkspace({ + latest_build: { + status: "running", + resources: [ + { + agents: [createAgent(agentOverrides)], + } as unknown as WorkspaceResource, + ], + }, + }); +} + +function createStateMachine( + startupMode: "prompt" | "start" | "update" = "start", +) { + return new WorkspaceStateMachine( + { agent: "main" } as unknown as AuthorityParts, + {} as CoderApi, + startupMode, + "/usr/bin/coder", + {} as FeatureSet, + createMockLogger(), + { mode: "url", url: "https://test.coder.com" } as CliAuth, + ); +} + +describe("WorkspaceStateMachine", () => { + let progress: MockProgress<{ message?: string }>; + + beforeEach(() => { + vi.clearAllMocks(); + progress = new MockProgress(); + vi.mocked(maybeAskAgent).mockImplementation((agents) => + Promise.resolve(agents.length > 0 ? agents[0] : undefined), + ); + }); + + describe("running workspace", () => { + it("returns true when agent is connected and ready", async () => { + const sm = createStateMachine(); + expect(await sm.processWorkspace(runningWorkspace(), progress)).toBe( + true, + ); + }); + + it("returns false when agent is connecting", async () => { + const sm = createStateMachine(); + const ws = runningWorkspace({ status: "connecting" }); + expect(await sm.processWorkspace(ws, progress)).toBe(false); + }); + + it("throws when agent is disconnected", async () => { + const sm = createStateMachine(); + const ws = runningWorkspace({ status: "disconnected" }); + await expect(sm.processWorkspace(ws, progress)).rejects.toThrow( + "disconnected", + ); + }); + + it("triggers update and falls through to agent check", async () => { + const sm = createStateMachine("update"); + const ws = runningWorkspace(); + + expect(await sm.processWorkspace(ws, progress)).toBe(true); + expect(updateWorkspace).toHaveBeenCalledOnce(); + }); + + it("re-resolves agent after update", async () => { + const sm = createStateMachine("start"); + const ws = runningWorkspace(); + + // Resolve agent, then verify it's cached on the next call. + await sm.processWorkspace(ws, progress); + vi.mocked(maybeAskAgent).mockClear(); + + await sm.processWorkspace(ws, progress); + expect(maybeAskAgent).not.toHaveBeenCalled(); + + // With update mode, the agent is cleared so it gets re-resolved. + const smUpdate = createStateMachine("update"); + vi.mocked(maybeAskAgent).mockClear(); + await smUpdate.processWorkspace(ws, progress); + expect(maybeAskAgent).toHaveBeenCalledOnce(); + }); + + it("downgrades to 'start' mode after update", async () => { + const sm = createStateMachine("update"); + await sm.processWorkspace(runningWorkspace(), progress); + vi.mocked(updateWorkspace).mockClear(); + + await sm.processWorkspace(runningWorkspace(), progress); + expect(updateWorkspace).not.toHaveBeenCalled(); + }); + }); + + describe("stopped/failed workspace", () => { + for (const status of ["stopped", "failed"] as const) { + it(`auto-starts '${status}' workspace`, async () => { + const sm = createStateMachine("start"); + const ws = createWorkspace({ latest_build: { status } }); + + expect(await sm.processWorkspace(ws, progress)).toBe(false); + expect(startWorkspace).toHaveBeenCalledOnce(); + }); + } + + it("triggers update instead of start when mode is 'update'", async () => { + const sm = createStateMachine("update"); + const ws = createWorkspace({ latest_build: { status: "stopped" } }); + + expect(await sm.processWorkspace(ws, progress)).toBe(false); + expect(updateWorkspace).toHaveBeenCalledOnce(); + }); + + it("prompts user when mode is 'prompt' and user accepts", async () => { + vi.mocked(vscodeProposed.window.showInformationMessage).mockResolvedValue( + "Start" as never, + ); + const sm = createStateMachine("prompt"); + const ws = createWorkspace({ latest_build: { status: "stopped" } }); + + expect(await sm.processWorkspace(ws, progress)).toBe(false); + expect(startWorkspace).toHaveBeenCalledOnce(); + }); + + it("throws when user declines start prompt", async () => { + vi.mocked(vscodeProposed.window.showInformationMessage).mockResolvedValue( + undefined as never, + ); + const sm = createStateMachine("prompt"); + const ws = createWorkspace({ latest_build: { status: "stopped" } }); + + await expect(sm.processWorkspace(ws, progress)).rejects.toThrow( + "Workspace start cancelled", + ); + }); + }); + + describe("building workspace", () => { + for (const status of ["pending", "starting", "stopping"] as const) { + it(`returns false and streams build logs for '${status}'`, async () => { + const sm = createStateMachine(); + const ws = createWorkspace({ latest_build: { status } }); + + expect(await sm.processWorkspace(ws, progress)).toBe(false); + expect(streamBuildLogs).toHaveBeenCalledOnce(); + }); + } + }); + + describe("terminal states", () => { + for (const status of [ + "deleted", + "deleting", + "canceled", + "canceling", + ] as const) { + it(`throws for '${status}'`, async () => { + const sm = createStateMachine(); + const ws = createWorkspace({ latest_build: { status } }); + await expect(sm.processWorkspace(ws, progress)).rejects.toThrow(status); + }); + } + }); + + describe("agent lifecycle", () => { + it("returns true for non-blocking 'starting' agent", async () => { + const sm = createStateMachine(); + const ws = runningWorkspace({ lifecycle_state: "starting", scripts: [] }); + expect(await sm.processWorkspace(ws, progress)).toBe(true); + }); + + it("returns false for 'starting' agent with blocking scripts", async () => { + const sm = createStateMachine(); + const ws = runningWorkspace({ + lifecycle_state: "starting", + scripts: [ + { + start_blocks_login: true, + } as unknown as WorkspaceAgent["scripts"][0], + ], + }); + expect(await sm.processWorkspace(ws, progress)).toBe(false); + }); + + it("returns true for 'start_error' (continues anyway)", async () => { + const sm = createStateMachine(); + expect( + await sm.processWorkspace( + runningWorkspace({ lifecycle_state: "start_error" }), + progress, + ), + ).toBe(true); + }); + + it("throws for 'off' lifecycle state", async () => { + const sm = createStateMachine(); + await expect( + sm.processWorkspace( + runningWorkspace({ lifecycle_state: "off" }), + progress, + ), + ).rejects.toThrow("Invalid lifecycle state"); + }); + }); + + describe("agent selection", () => { + it("throws when user declines agent selection", async () => { + vi.mocked(maybeAskAgent).mockResolvedValue(undefined); + const sm = createStateMachine(); + await expect( + sm.processWorkspace(runningWorkspace(), progress), + ).rejects.toThrow("Agent selection cancelled"); + }); + + it("throws when selected agent disappears from resources", async () => { + const sm = createStateMachine(); + await sm.processWorkspace(runningWorkspace(), progress); + + const wsNoAgents = createWorkspace({ + latest_build: { status: "running", resources: [] }, + }); + await expect(sm.processWorkspace(wsNoAgents, progress)).rejects.toThrow( + "not found", + ); + }); + }); + + describe("progress reporting", () => { + it("reports starting for stopped workspace", async () => { + const sm = createStateMachine("start"); + const ws = createWorkspace({ latest_build: { status: "stopped" } }); + await sm.processWorkspace(ws, progress); + + expect(progress.report).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining("starting"), + }), + ); + }); + + it("reports updating for update mode", async () => { + const sm = createStateMachine("update"); + await sm.processWorkspace(runningWorkspace(), progress); + + expect(progress.report).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining("updating"), + }), + ); + }); + }); + + describe("dispose", () => { + it("can be disposed without errors", () => { + expect(() => createStateMachine().dispose()).not.toThrow(); + }); + }); +}); From 519d514cc79f130415621667ade035d39e2729df Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 14 Apr 2026 13:03:09 +0300 Subject: [PATCH 2/3] rename StartupMode "prompt" to "none" for clarity "prompt" was ambiguous since it could imply prompting for both start and update. "none" makes it clear that no explicit intent was set, and the default behavior (asking) applies. Also bundles restClient and featureSet into CliContext so startWorkspace and updateWorkspace each take a single argument, removes redundant getWorkspace() API calls in favor of using the workspace already passed via context, replaces the cliContext getter with buildCliContext(workspace), and fixes a stale test description. --- src/api/workspace.ts | 40 ++++++++----------- src/core/mementoManager.ts | 8 ++-- src/remote/workspaceStateMachine.ts | 19 ++++----- test/unit/core/mementoManager.test.ts | 10 ++--- .../unit/remote/workspaceStateMachine.test.ts | 6 +-- 5 files changed, 35 insertions(+), 48 deletions(-) diff --git a/src/api/workspace.ts b/src/api/workspace.ts index 677a539c..babcc80d 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -46,10 +46,12 @@ export class LazyStream { } interface CliContext { + restClient: Api; auth: CliAuth; binPath: string; workspace: Workspace; writeEmitter: vscode.EventEmitter; + featureSet: FeatureSet; } /** @@ -106,55 +108,45 @@ function splitLines(data: Buffer): string[] { * Start a stopped or failed workspace using `coder start`. * No-ops if the workspace is already running. */ -export async function startWorkspace( - restClient: Api, - ctx: CliContext, - featureSet: FeatureSet, -): Promise { - const current = await restClient.getWorkspace(ctx.workspace.id); - if (!["stopped", "failed"].includes(current.latest_build.status)) { - return current; +export async function startWorkspace(ctx: CliContext): Promise { + if (!["stopped", "failed"].includes(ctx.workspace.latest_build.status)) { + return ctx.workspace; } const args = ["start", "--yes"]; - if (featureSet.buildReason) { + if (ctx.featureSet.buildReason) { args.push("--reason", "vscode_connection"); } await runCliCommand(ctx, args); - return restClient.getWorkspace(ctx.workspace.id); + 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 → updateWorkspaceVersion. + * Falls back to the REST API: stop, wait, then updateWorkspaceVersion. */ -export async function updateWorkspace( - restClient: Api, - ctx: CliContext, - featureSet: FeatureSet, -): Promise { - if (featureSet.cliUpdate) { +export async function updateWorkspace(ctx: CliContext): Promise { + if (ctx.featureSet.cliUpdate) { await runCliCommand(ctx, ["update"]); - return restClient.getWorkspace(ctx.workspace.id); + return ctx.restClient.getWorkspace(ctx.workspace.id); } // REST API fallback for older CLIs. - const workspace = await restClient.getWorkspace(ctx.workspace.id); - if (workspace.latest_build.status === "running") { + if (ctx.workspace.latest_build.status === "running") { ctx.writeEmitter.fire("Stopping workspace for update...\r\n"); - const stopBuild = await restClient.stopWorkspace(workspace.id); - const stoppedJob = await restClient.waitForBuild(stopBuild); + 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 restClient.updateWorkspaceVersion(workspace); - return restClient.getWorkspace(ctx.workspace.id); + await ctx.restClient.updateWorkspaceVersion(ctx.workspace); + return ctx.restClient.getWorkspace(ctx.workspace.id); } /** diff --git a/src/core/mementoManager.ts b/src/core/mementoManager.ts index 1a5c9169..a1a1e701 100644 --- a/src/core/mementoManager.ts +++ b/src/core/mementoManager.ts @@ -9,12 +9,12 @@ const PENDING_TTL_MS = 5 * 60 * 1000; /** * Describes the startup intent when the extension connects to a workspace. - * - "prompt": Normal reconnection; ask before starting a stopped 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 = "prompt" | "start" | "update"; +export type StartupMode = "none" | "start" | "update"; interface Stamped { value: T; @@ -62,14 +62,14 @@ export class MementoManager { /** * Read and clear the startup mode. - * Returns "prompt" (the default) when no mode was explicitly set. + * Returns "none" (the default) when no mode was explicitly set. */ public async getAndClearStartupMode(): Promise { const value = this.getStamped("startupMode"); if (value !== undefined) { await this.memento.update("startupMode", undefined); } - return value ?? "prompt"; + return value ?? "none"; } /** Store a chat ID to open after a remote-authority reload. */ diff --git a/src/remote/workspaceStateMachine.ts b/src/remote/workspaceStateMachine.ts index d265a74c..3e946af0 100644 --- a/src/remote/workspaceStateMachine.ts +++ b/src/remote/workspaceStateMachine.ts @@ -73,7 +73,7 @@ export class WorkspaceStateMachine implements vscode.Disposable { this.buildLogStream.close(); if ( - this.startupMode === "prompt" && + this.startupMode === "none" && !(await this.confirmStart(workspaceName)) ) { throw new Error(`Workspace start cancelled`); @@ -223,11 +223,14 @@ export class WorkspaceStateMachine implements vscode.Disposable { return this.startupMode === "update"; } - private get cliContext() { + private buildCliContext(workspace: Workspace) { return { + restClient: this.workspaceClient, auth: this.cliAuth, binPath: this.binaryPath, + workspace, writeEmitter: this.terminal.writeEmitter, + featureSet: this.featureSet, }; } @@ -241,11 +244,7 @@ export class WorkspaceStateMachine implements vscode.Disposable { mode: this.startupMode, status: workspace.latest_build.status, }); - await startWorkspace( - this.workspaceClient, - { ...this.cliContext, workspace }, - this.featureSet, - ); + await startWorkspace(this.buildCliContext(workspace)); this.logger.info(`${workspaceName} start initiated`); } @@ -259,11 +258,7 @@ export class WorkspaceStateMachine implements vscode.Disposable { mode: this.startupMode, status: workspace.latest_build.status, }); - await updateWorkspace( - this.workspaceClient, - { ...this.cliContext, workspace }, - this.featureSet, - ); + await updateWorkspace(this.buildCliContext(workspace)); // Downgrade so subsequent transitions don't re-trigger the update. this.startupMode = "start"; this.logger.info(`${workspaceName} update initiated`); diff --git a/test/unit/core/mementoManager.test.ts b/test/unit/core/mementoManager.test.ts index 3275915a..f7c19c5d 100644 --- a/test/unit/core/mementoManager.test.ts +++ b/test/unit/core/mementoManager.test.ts @@ -70,11 +70,11 @@ describe("MementoManager", () => { it("should return the set mode and clear after read", async () => { await mementoManager.setStartupMode("start"); expect(await mementoManager.getAndClearStartupMode()).toBe("start"); - expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); + expect(await mementoManager.getAndClearStartupMode()).toBe("none"); }); - it("should return 'prompt' when nothing is set", async () => { - expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); + it("should return 'none' when nothing is set", async () => { + expect(await mementoManager.getAndClearStartupMode()).toBe("none"); }); it("should support 'update' mode", async () => { @@ -84,13 +84,13 @@ describe("MementoManager", () => { it("should treat legacy bare values as expired", async () => { await memento.update("startupMode", "start"); - expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); + expect(await mementoManager.getAndClearStartupMode()).toBe("none"); }); it("should expire after 5 minutes", async () => { await mementoManager.setStartupMode("update"); vi.advanceTimersByTime(5 * 60 * 1000 + 1); - expect(await mementoManager.getAndClearStartupMode()).toBe("prompt"); + expect(await mementoManager.getAndClearStartupMode()).toBe("none"); }); }); diff --git a/test/unit/remote/workspaceStateMachine.test.ts b/test/unit/remote/workspaceStateMachine.test.ts index af753eec..0898df85 100644 --- a/test/unit/remote/workspaceStateMachine.test.ts +++ b/test/unit/remote/workspaceStateMachine.test.ts @@ -79,7 +79,7 @@ function runningWorkspace( } function createStateMachine( - startupMode: "prompt" | "start" | "update" = "start", + startupMode: "none" | "start" | "update" = "start", ) { return new WorkspaceStateMachine( { agent: "main" } as unknown as AuthorityParts, @@ -184,7 +184,7 @@ describe("WorkspaceStateMachine", () => { vi.mocked(vscodeProposed.window.showInformationMessage).mockResolvedValue( "Start" as never, ); - const sm = createStateMachine("prompt"); + const sm = createStateMachine("none"); const ws = createWorkspace({ latest_build: { status: "stopped" } }); expect(await sm.processWorkspace(ws, progress)).toBe(false); @@ -195,7 +195,7 @@ describe("WorkspaceStateMachine", () => { vi.mocked(vscodeProposed.window.showInformationMessage).mockResolvedValue( undefined as never, ); - const sm = createStateMachine("prompt"); + const sm = createStateMachine("none"); const ws = createWorkspace({ latest_build: { status: "stopped" } }); await expect(sm.processWorkspace(ws, progress)).rejects.toThrow( From fc263e7da34b31bce249c6916c70c1293bab4314 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 14 Apr 2026 14:37:53 +0300 Subject: [PATCH 3/3] 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. --- src/remote/workspaceStateMachine.ts | 38 +-- test/mocks/testHelpers.ts | 84 +++++- test/mocks/workspace.ts | 55 +++- .../unit/remote/workspaceStateMachine.test.ts | 252 +++++++++++------- 4 files changed, 310 insertions(+), 119 deletions(-) diff --git a/src/remote/workspaceStateMachine.ts b/src/remote/workspaceStateMachine.ts index 3e946af0..20ef8f66 100644 --- a/src/remote/workspaceStateMachine.ts +++ b/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 { + 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 { diff --git a/test/mocks/testHelpers.ts b/test/mocks/testHelpers.ts index d92b96f7..052de4b4 100644 --- a/test/mocks/testHelpers.ts +++ b/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(); + private readonly _messageCalls: MessageCall[] = []; private inputBoxValue: string | undefined; private inputBoxValidateInput: ((value: string) => Promise) | 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 => { - 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 => { + 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(); + } +} diff --git a/test/mocks/workspace.ts b/test/mocks/workspace.ts index 315f377d..13378b19 100644 --- a/test/mocks/workspace.ts +++ b/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 { + 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 { + 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, + }; +} diff --git a/test/unit/remote/workspaceStateMachine.test.ts b/test/unit/remote/workspaceStateMachine.test.ts index 0898df85..e45c21c9 100644 --- a/test/unit/remote/workspaceStateMachine.test.ts +++ b/test/unit/remote/workspaceStateMachine.test.ts @@ -1,19 +1,34 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; +import { + startWorkspace, + updateWorkspace, + streamBuildLogs, + streamAgentLogs, +} from "@/api/workspace"; +import { maybeAskAgent } from "@/promptUtils"; import { WorkspaceStateMachine } from "@/remote/workspaceStateMachine"; -import { createMockLogger, MockProgress } from "../../mocks/testHelpers"; -import { workspace as createWorkspace } from "../../mocks/workspace"; +import { + createMockLogger, + MockProgress, + MockTerminalSession, + MockUserInteraction, +} from "../../mocks/testHelpers"; +import { + agent as createAgent, + resource as createResource, + workspace as createWorkspace, +} from "../../mocks/workspace"; import type { Workspace, WorkspaceAgent, - WorkspaceResource, } from "coder/site/src/api/typesGenerated"; import type { CoderApi } from "@/api/coderApi"; +import type { StartupMode } from "@/core/mementoManager"; import type { FeatureSet } from "@/featureSet"; -import type { CliAuth } from "@/settings/cli"; import type { AuthorityParts } from "@/util"; vi.mock("@/api/workspace", async (importActual) => { @@ -31,37 +46,22 @@ vi.mock("@/promptUtils", () => ({ maybeAskAgent: vi.fn(), })); -vi.mock("@/vscodeProposed", () => ({ - vscodeProposed: { - window: { showInformationMessage: vi.fn() }, - }, -})); - -vi.mock("@/remote/terminalSession", () => ({ - TerminalSession: vi.fn().mockImplementation(function () { - return { - writeEmitter: { fire: vi.fn(), event: vi.fn(), dispose: vi.fn() }, - terminal: { show: vi.fn(), dispose: vi.fn() }, - dispose: vi.fn(), - }; - }), -})); +vi.mock("@/remote/terminalSession", async () => { + const helpers = await import("../../mocks/testHelpers"); + return { TerminalSession: helpers.MockTerminalSession }; +}); -const { startWorkspace, updateWorkspace, streamBuildLogs } = - await import("@/api/workspace"); -const { maybeAskAgent } = await import("@/promptUtils"); -const { vscodeProposed } = await import("@/vscodeProposed"); +const DEFAULT_PARTS: Readonly = { + agent: "main", + sshHost: "coder-vscode--testuser--test-workspace.main", + safeHostname: "test.coder.com", + username: "testuser", + workspace: "test-workspace", +} as const; -function createAgent(overrides: Partial = {}): WorkspaceAgent { - return { - id: "agent-1", - name: "main", - status: "connected", - lifecycle_state: "ready", - scripts: [], - ...overrides, - } as unknown as WorkspaceAgent; -} +// The message shown by confirmStartOrUpdate for our test workspace. +const CONFIRM_MESSAGE = + "The workspace testuser/test-workspace is not running. How would you like to proceed?"; function runningWorkspace( agentOverrides: Partial = {}, @@ -69,35 +69,30 @@ function runningWorkspace( return createWorkspace({ latest_build: { status: "running", - resources: [ - { - agents: [createAgent(agentOverrides)], - } as unknown as WorkspaceResource, - ], + resources: [createResource({ agents: [createAgent(agentOverrides)] })], }, }); } -function createStateMachine( - startupMode: "none" | "start" | "update" = "start", -) { - return new WorkspaceStateMachine( - { agent: "main" } as unknown as AuthorityParts, +function setup(startupMode: StartupMode = "start") { + const progress = new MockProgress<{ message?: string }>(); + const userInteraction = new MockUserInteraction(); + const sm = new WorkspaceStateMachine( + DEFAULT_PARTS, {} as CoderApi, startupMode, "/usr/bin/coder", {} as FeatureSet, createMockLogger(), - { mode: "url", url: "https://test.coder.com" } as CliAuth, + { mode: "url", url: "https://test.coder.com" }, ); + return { sm, progress, userInteraction }; } describe("WorkspaceStateMachine", () => { - let progress: MockProgress<{ message?: string }>; - beforeEach(() => { vi.clearAllMocks(); - progress = new MockProgress(); + MockTerminalSession.lastInstance = undefined; vi.mocked(maybeAskAgent).mockImplementation((agents) => Promise.resolve(agents.length > 0 ? agents[0] : undefined), ); @@ -105,20 +100,26 @@ describe("WorkspaceStateMachine", () => { describe("running workspace", () => { it("returns true when agent is connected and ready", async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); expect(await sm.processWorkspace(runningWorkspace(), progress)).toBe( true, ); }); it("returns false when agent is connecting", async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); const ws = runningWorkspace({ status: "connecting" }); expect(await sm.processWorkspace(ws, progress)).toBe(false); }); + it("returns false when agent times out", async () => { + const { sm, progress } = setup(); + const ws = runningWorkspace({ status: "timeout" }); + expect(await sm.processWorkspace(ws, progress)).toBe(false); + }); + it("throws when agent is disconnected", async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); const ws = runningWorkspace({ status: "disconnected" }); await expect(sm.processWorkspace(ws, progress)).rejects.toThrow( "disconnected", @@ -126,7 +127,7 @@ describe("WorkspaceStateMachine", () => { }); it("triggers update and falls through to agent check", async () => { - const sm = createStateMachine("update"); + const { sm, progress } = setup("update"); const ws = runningWorkspace(); expect(await sm.processWorkspace(ws, progress)).toBe(true); @@ -134,7 +135,7 @@ describe("WorkspaceStateMachine", () => { }); it("re-resolves agent after update", async () => { - const sm = createStateMachine("start"); + const { sm, progress } = setup("start"); const ws = runningWorkspace(); // Resolve agent, then verify it's cached on the next call. @@ -145,14 +146,14 @@ describe("WorkspaceStateMachine", () => { expect(maybeAskAgent).not.toHaveBeenCalled(); // With update mode, the agent is cleared so it gets re-resolved. - const smUpdate = createStateMachine("update"); + const { sm: smUpdate, progress: p2 } = setup("update"); vi.mocked(maybeAskAgent).mockClear(); - await smUpdate.processWorkspace(ws, progress); + await smUpdate.processWorkspace(ws, p2); expect(maybeAskAgent).toHaveBeenCalledOnce(); }); it("downgrades to 'start' mode after update", async () => { - const sm = createStateMachine("update"); + const { sm, progress } = setup("update"); await sm.processWorkspace(runningWorkspace(), progress); vi.mocked(updateWorkspace).mockClear(); @@ -164,7 +165,7 @@ describe("WorkspaceStateMachine", () => { describe("stopped/failed workspace", () => { for (const status of ["stopped", "failed"] as const) { it(`auto-starts '${status}' workspace`, async () => { - const sm = createStateMachine("start"); + const { sm, progress } = setup("start"); const ws = createWorkspace({ latest_build: { status } }); expect(await sm.processWorkspace(ws, progress)).toBe(false); @@ -173,29 +174,59 @@ describe("WorkspaceStateMachine", () => { } it("triggers update instead of start when mode is 'update'", async () => { - const sm = createStateMachine("update"); + const { sm, progress } = setup("update"); const ws = createWorkspace({ latest_build: { status: "stopped" } }); expect(await sm.processWorkspace(ws, progress)).toBe(false); expect(updateWorkspace).toHaveBeenCalledOnce(); }); - it("prompts user when mode is 'prompt' and user accepts", async () => { - vi.mocked(vscodeProposed.window.showInformationMessage).mockResolvedValue( - "Start" as never, - ); - const sm = createStateMachine("none"); + it("prompts user when mode is 'none' and user picks 'Start'", async () => { + const { sm, progress, userInteraction } = setup("none"); + userInteraction.setResponse(CONFIRM_MESSAGE, "Start"); const ws = createWorkspace({ latest_build: { status: "stopped" } }); expect(await sm.processWorkspace(ws, progress)).toBe(false); expect(startWorkspace).toHaveBeenCalledOnce(); + expect(updateWorkspace).not.toHaveBeenCalled(); }); - it("throws when user declines start prompt", async () => { - vi.mocked(vscodeProposed.window.showInformationMessage).mockResolvedValue( - undefined as never, - ); - const sm = createStateMachine("none"); + it("offers 'Update and Start' for outdated workspace and triggers update", async () => { + const { sm, progress, userInteraction } = setup("none"); + userInteraction.setResponse(CONFIRM_MESSAGE, "Update and Start"); + const ws = createWorkspace({ + outdated: true, + latest_build: { status: "stopped" }, + }); + + expect(await sm.processWorkspace(ws, progress)).toBe(false); + + const calls = userInteraction.getMessageCalls(); + expect(calls).toHaveLength(1); + expect(calls[0].items).toEqual(["Start", "Update and Start"]); + + expect(updateWorkspace).toHaveBeenCalledOnce(); + expect(startWorkspace).not.toHaveBeenCalled(); + }); + + it("does not offer 'Update and Start' when workspace is not outdated", async () => { + const { sm, progress, userInteraction } = setup("none"); + userInteraction.setResponse(CONFIRM_MESSAGE, "Start"); + const ws = createWorkspace({ + outdated: false, + latest_build: { status: "stopped" }, + }); + + await sm.processWorkspace(ws, progress); + + const calls = userInteraction.getMessageCalls(); + expect(calls).toHaveLength(1); + expect(calls[0].items).toEqual(["Start"]); + }); + + it("throws when user declines the prompt", async () => { + const { sm, progress, userInteraction } = setup("none"); + userInteraction.setResponse(CONFIRM_MESSAGE, undefined); const ws = createWorkspace({ latest_build: { status: "stopped" } }); await expect(sm.processWorkspace(ws, progress)).rejects.toThrow( @@ -207,7 +238,7 @@ describe("WorkspaceStateMachine", () => { describe("building workspace", () => { for (const status of ["pending", "starting", "stopping"] as const) { it(`returns false and streams build logs for '${status}'`, async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); const ws = createWorkspace({ latest_build: { status } }); expect(await sm.processWorkspace(ws, progress)).toBe(false); @@ -224,7 +255,7 @@ describe("WorkspaceStateMachine", () => { "canceling", ] as const) { it(`throws for '${status}'`, async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); const ws = createWorkspace({ latest_build: { status } }); await expect(sm.processWorkspace(ws, progress)).rejects.toThrow(status); }); @@ -233,56 +264,79 @@ describe("WorkspaceStateMachine", () => { describe("agent lifecycle", () => { it("returns true for non-blocking 'starting' agent", async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); const ws = runningWorkspace({ lifecycle_state: "starting", scripts: [] }); expect(await sm.processWorkspace(ws, progress)).toBe(true); }); it("returns false for 'starting' agent with blocking scripts", async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); const ws = runningWorkspace({ lifecycle_state: "starting", scripts: [ { + id: "script-1", + log_source_id: "log-1", + log_path: "", + script: "#!/bin/bash", + cron: "", + run_on_start: true, + run_on_stop: false, start_blocks_login: true, - } as unknown as WorkspaceAgent["scripts"][0], + timeout: 0, + display_name: "Startup", + }, ], }); expect(await sm.processWorkspace(ws, progress)).toBe(false); + expect(streamAgentLogs).toHaveBeenCalledOnce(); + }); + + it("returns false for 'created' agent", async () => { + const { sm, progress } = setup(); + const ws = runningWorkspace({ lifecycle_state: "created" }); + expect(await sm.processWorkspace(ws, progress)).toBe(false); }); it("returns true for 'start_error' (continues anyway)", async () => { - const sm = createStateMachine(); - expect( - await sm.processWorkspace( - runningWorkspace({ lifecycle_state: "start_error" }), - progress, - ), - ).toBe(true); + const { sm, progress } = setup(); + const ws = runningWorkspace({ lifecycle_state: "start_error" }); + expect(await sm.processWorkspace(ws, progress)).toBe(true); }); - it("throws for 'off' lifecycle state", async () => { - const sm = createStateMachine(); - await expect( - sm.processWorkspace( - runningWorkspace({ lifecycle_state: "off" }), - progress, - ), - ).rejects.toThrow("Invalid lifecycle state"); + it("returns true for 'start_timeout' (continues anyway)", async () => { + const { sm, progress } = setup(); + const ws = runningWorkspace({ lifecycle_state: "start_timeout" }); + expect(await sm.processWorkspace(ws, progress)).toBe(true); }); + + for (const lifecycle_state of [ + "shutting_down", + "off", + "shutdown_error", + "shutdown_timeout", + ] as const) { + it(`throws for '${lifecycle_state}' lifecycle state`, async () => { + const { sm, progress } = setup(); + const ws = runningWorkspace({ lifecycle_state }); + await expect(sm.processWorkspace(ws, progress)).rejects.toThrow( + "Invalid lifecycle state", + ); + }); + } }); describe("agent selection", () => { it("throws when user declines agent selection", async () => { vi.mocked(maybeAskAgent).mockResolvedValue(undefined); - const sm = createStateMachine(); + const { sm, progress } = setup(); await expect( sm.processWorkspace(runningWorkspace(), progress), ).rejects.toThrow("Agent selection cancelled"); }); it("throws when selected agent disappears from resources", async () => { - const sm = createStateMachine(); + const { sm, progress } = setup(); await sm.processWorkspace(runningWorkspace(), progress); const wsNoAgents = createWorkspace({ @@ -296,7 +350,7 @@ describe("WorkspaceStateMachine", () => { describe("progress reporting", () => { it("reports starting for stopped workspace", async () => { - const sm = createStateMachine("start"); + const { sm, progress } = setup("start"); const ws = createWorkspace({ latest_build: { status: "stopped" } }); await sm.processWorkspace(ws, progress); @@ -308,7 +362,7 @@ describe("WorkspaceStateMachine", () => { }); it("reports updating for update mode", async () => { - const sm = createStateMachine("update"); + const { sm, progress } = setup("update"); await sm.processWorkspace(runningWorkspace(), progress); expect(progress.report).toHaveBeenCalledWith( @@ -319,9 +373,23 @@ describe("WorkspaceStateMachine", () => { }); }); + describe("getAgentId", () => { + it("returns undefined before agent is resolved", () => { + const { sm } = setup(); + expect(sm.getAgentId()).toBeUndefined(); + }); + + it("returns agent ID after processing a running workspace", async () => { + const { sm, progress } = setup(); + await sm.processWorkspace(runningWorkspace(), progress); + expect(sm.getAgentId()).toBe("agent-1"); + }); + }); + describe("dispose", () => { it("can be disposed without errors", () => { - expect(() => createStateMachine().dispose()).not.toThrow(); + const { sm } = setup(); + expect(() => sm.dispose()).not.toThrow(); }); }); });