diff --git a/src/api/workspace.ts b/src/api/workspace.ts index 698d212a..7d392836 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -45,66 +45,46 @@ export class LazyStream { } } +interface CliContext { + restClient: Api; + auth: CliAuth; + binPath: string; + workspace: Workspace; + writeEmitter: vscode.EventEmitter; + 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, - 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) => { + ctx.writeEmitter.fire(data.toString().replace(/\r?\n/g, "\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) => { + const text = data.toString(); + ctx.writeEmitter.fire(text.replace(/\r?\n/g, "\r\n")); + capturedStderr += text; }); - 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 +94,51 @@ export async function startWorkspaceIfStoppedOrFailed( }); } +/** + * Start a stopped or failed workspace using `coder start`. + * No-ops if the workspace is already running. + */ +export async function startWorkspace(ctx: CliContext): Promise { + 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.24). + * Falls back to the REST API: stop, wait, then updateWorkspaceVersion. + */ +export async function updateWorkspace(ctx: CliContext): Promise { + 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. 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..a1a1e701 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. + * - "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 { 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 "none" (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 ?? "none"; } /** 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..8cc17b8c 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` with stop transition (stops before updating) + cliUpdate: versionAtLeast(version, "2.24.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..20ef8f66 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,34 +61,40 @@ export class WorkspaceStateMachine implements vscode.Disposable { switch (workspace.latest_build.status) { case "running": this.buildLogStream.close(); + if (this.startupMode === "update") { + 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))) { - 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; } - 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.startupMode === "update") { + await this.triggerUpdate(workspace, workspaceName, progress); + } else { + await this.triggerStart(workspace, workspaceName, progress); + } return false; } 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({ @@ -215,16 +223,63 @@ export class WorkspaceStateMachine implements vscode.Disposable { } } - private async confirmStart(workspaceName: string): Promise { + private buildCliContext(workspace: Workspace) { + return { + restClient: this.workspaceClient, + auth: this.cliAuth, + binPath: this.binaryPath, + workspace, + writeEmitter: this.terminal.writeEmitter, + featureSet: this.featureSet, + }; + } + + 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.buildCliContext(workspace)); + 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.buildCliContext(workspace)); + // Downgrade so subsequent transitions don't re-trigger the update. + this.startupMode = "start"; + this.logger.info(`${workspaceName} update initiated`); + } + + 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/core/mementoManager.test.ts b/test/unit/core/mementoManager.test.ts index 16c3efbe..f7c19c5d 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("none"); + }); + + it("should return 'none' when nothing is set", async () => { + expect(await mementoManager.getAndClearStartupMode()).toBe("none"); + }); - 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("none"); }); 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("none"); }); }); diff --git a/test/unit/remote/workspaceStateMachine.test.ts b/test/unit/remote/workspaceStateMachine.test.ts new file mode 100644 index 00000000..e45c21c9 --- /dev/null +++ b/test/unit/remote/workspaceStateMachine.test.ts @@ -0,0 +1,395 @@ +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, + MockTerminalSession, + MockUserInteraction, +} from "../../mocks/testHelpers"; +import { + agent as createAgent, + resource as createResource, + workspace as createWorkspace, +} from "../../mocks/workspace"; + +import type { + Workspace, + WorkspaceAgent, +} 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 { 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("@/remote/terminalSession", async () => { + const helpers = await import("../../mocks/testHelpers"); + return { TerminalSession: helpers.MockTerminalSession }; +}); + +const DEFAULT_PARTS: Readonly = { + agent: "main", + sshHost: "coder-vscode--testuser--test-workspace.main", + safeHostname: "test.coder.com", + username: "testuser", + workspace: "test-workspace", +} as const; + +// 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 = {}, +): Workspace { + return createWorkspace({ + latest_build: { + status: "running", + resources: [createResource({ agents: [createAgent(agentOverrides)] })], + }, + }); +} + +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" }, + ); + return { sm, progress, userInteraction }; +} + +describe("WorkspaceStateMachine", () => { + beforeEach(() => { + vi.clearAllMocks(); + MockTerminalSession.lastInstance = undefined; + 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, progress } = setup(); + expect(await sm.processWorkspace(runningWorkspace(), progress)).toBe( + true, + ); + }); + + it("returns false when agent is connecting", async () => { + 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, progress } = setup(); + 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, progress } = setup("update"); + const ws = runningWorkspace(); + + expect(await sm.processWorkspace(ws, progress)).toBe(true); + expect(updateWorkspace).toHaveBeenCalledOnce(); + }); + + it("re-resolves agent after update", async () => { + const { sm, progress } = setup("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 { sm: smUpdate, progress: p2 } = setup("update"); + vi.mocked(maybeAskAgent).mockClear(); + await smUpdate.processWorkspace(ws, p2); + expect(maybeAskAgent).toHaveBeenCalledOnce(); + }); + + it("downgrades to 'start' mode after update", async () => { + const { sm, progress } = setup("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, progress } = setup("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, 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 '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("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( + "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, progress } = setup(); + 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, progress } = setup(); + 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, 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, 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, + 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, progress } = setup(); + const ws = runningWorkspace({ lifecycle_state: "start_error" }); + expect(await sm.processWorkspace(ws, progress)).toBe(true); + }); + + 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, progress } = setup(); + await expect( + sm.processWorkspace(runningWorkspace(), progress), + ).rejects.toThrow("Agent selection cancelled"); + }); + + it("throws when selected agent disappears from resources", async () => { + const { sm, progress } = setup(); + 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, progress } = setup("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, progress } = setup("update"); + await sm.processWorkspace(runningWorkspace(), progress); + + expect(progress.report).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining("updating"), + }), + ); + }); + }); + + 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", () => { + const { sm } = setup(); + expect(() => sm.dispose()).not.toThrow(); + }); + }); +});