workspaceStateMachine.ts

import { createWorkspaceIdentifier, extractAgents } from "../api/api-helper";
import {
	LazyStream,
	startWorkspaceIfStoppedOrFailed,
	streamAgentLogs,
	streamBuildLogs,
} from "../api/workspace";
import { maybeAskAgent } from "../promptUtils";
import { type AuthorityParts } from "../util";
import { vscodeProposed } from "../vscodeProposed";

import { TerminalSession } from "./terminalSession";

import type {
	ProvisionerJobLog,
	Workspace,
	WorkspaceAgentLog,
} from "coder/site/src/api/typesGenerated";
import type * as vscode from "vscode";

import type { CoderApi } from "../api/coderApi";
import type { FeatureSet } from "../featureSet";
import type { Logger } from "../logging/logger";
import type { CliAuth } from "../settings/cli";

/**
 * Manages workspace and agent state transitions until ready for SSH connection.
 * Streams build and agent logs, and handles socket lifecycle.
 */
export class WorkspaceStateMachine implements vscode.Disposable {
	private readonly terminal: TerminalSession;
	private readonly buildLogStream = new LazyStream<ProvisionerJobLog>();
	private readonly agentLogStream = new LazyStream<WorkspaceAgentLog[]>();

	private agent: { id: string; name: string } | undefined;

	constructor(
		private readonly parts: AuthorityParts,
		private readonly workspaceClient: CoderApi,
		private readonly firstConnect: boolean,
		private readonly binaryPath: string,
		private readonly featureSet: FeatureSet,
		private readonly logger: Logger,
		private readonly cliAuth: CliAuth,
	) {
		this.terminal = new TerminalSession("Workspace Build");
	}

	/**
	 * Process workspace state and determine if agent is ready.
	 * Reports progress updates and returns true if ready to connect, false if should wait for next event.
	 */
	async processWorkspace(
		workspace: Workspace,
		progress: vscode.Progress<{ message?: string }>,
	): Promise<boolean> {
		const workspaceName = createWorkspaceIdentifier(workspace);

		switch (workspace.latest_build.status) {
			case "running":
				this.buildLogStream.close();
				break;

			case "stopped":
			case "failed": {
				this.buildLogStream.close();

				if (!this.firstConnect && !(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`);
				return false;
			}

			case "pending":
			case "starting":
			case "stopping": {
				// Clear the agent since it's ID could change after a restart
				this.agent = undefined;
				this.agentLogStream.close();
				progress.report({
					message: `building ${workspaceName} (${workspace.latest_build.status})...`,
				});
				this.logger.info(`Waiting for ${workspaceName}`);

				const write = (line: string) =>
					this.terminal.writeEmitter.fire(line + "\r\n");
				await this.buildLogStream.open(() =>
					streamBuildLogs(
						this.workspaceClient,
						write,
						workspace.latest_build.id,
					),
				);
				return false;
			}

			case "deleted":
			case "deleting":
			case "canceled":
			case "canceling":
				this.buildLogStream.close();
				throw new Error(`${workspaceName} is ${workspace.latest_build.status}`);
		}

		const agents = extractAgents(workspace.latest_build.resources);
		if (this.agent === undefined) {
			this.logger.info(`Finding agent for ${workspaceName}`);
			const gotAgent = await maybeAskAgent(agents, this.parts.agent);
			if (!gotAgent) {
				// User declined to pick an agent.
				throw new Error("Agent selection cancelled");
			}
			this.agent = { id: gotAgent.id, name: gotAgent.name };
			this.logger.info(
				`Found agent ${gotAgent.name} with status`,
				gotAgent.status,
			);
		}
		const agent = agents.find((a) => a.id === this.agent?.id);
		if (!agent) {
			throw new Error(
				`Agent ${this.agent.name} not found in ${workspaceName} resources`,
			);
		}

		switch (agent.status) {
			case "connecting":
				progress.report({
					message: `connecting to agent ${agent.name}...`,
				});
				this.logger.debug(`Connecting to agent ${agent.name}`);
				return false;

			case "disconnected":
				throw new Error(`Agent ${workspaceName}/${agent.name} disconnected`);

			case "timeout":
				progress.report({
					message: `agent ${agent.name} timed out, retrying...`,
				});
				this.logger.debug(`Agent ${agent.name} timed out, retrying`);
				return false;

			case "connected":
				break;
		}

		switch (agent.lifecycle_state) {
			case "ready":
				this.agentLogStream.close();
				return true;

			case "starting": {
				const isBlocking = agent.scripts.some(
					(script) => script.start_blocks_login,
				);
				if (!isBlocking) {
					return true;
				}

				progress.report({
					message: `running agent ${agent.name} startup scripts...`,
				});
				this.logger.debug(`Running agent ${agent.name} startup scripts`);

				const writeAgent = (line: string) =>
					this.terminal.writeEmitter.fire(line + "\r\n");
				await this.agentLogStream.open(() =>
					streamAgentLogs(this.workspaceClient, writeAgent, agent.id),
				);
				return false;
			}

			case "created":
				progress.report({
					message: `starting agent ${agent.name}...`,
				});
				this.logger.debug(`Starting agent ${agent.name}`);
				return false;

			case "start_error":
				this.agentLogStream.close();
				this.logger.info(
					`Agent ${agent.name} startup scripts failed, but continuing`,
				);
				return true;

			case "start_timeout":
				this.agentLogStream.close();
				this.logger.info(
					`Agent ${agent.name} startup scripts timed out, but continuing`,
				);
				return true;

			case "shutting_down":
			case "off":
			case "shutdown_error":
			case "shutdown_timeout":
				this.agentLogStream.close();
				throw new Error(
					`Invalid lifecycle state '${agent.lifecycle_state}' for ${workspaceName}/${agent.name}`,
				);
		}
	}

	private async confirmStart(workspaceName: string): Promise<boolean> {
		const action = await vscodeProposed.window.showInformationMessage(
			`Unable to connect to the workspace ${workspaceName} because it is not running. Start the workspace?`,
			{
				useCustom: true,
				modal: true,
			},
			"Start",
		);
		return action === "Start";
	}

	public getAgentId(): string | undefined {
		return this.agent?.id;
	}

	dispose(): void {
		this.buildLogStream.close();
		this.agentLogStream.close();
		this.terminal.dispose();
	}
}