diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index da624bad..ef120ce4 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -11,6 +11,7 @@ import { type ProvisionerJobLog, type Workspace, type WorkspaceAgent, + type WorkspaceAgentLog, } from "coder/site/src/api/typesGenerated"; import * as vscode from "vscode"; import { type ClientOptions, type CloseEvent, type ErrorEvent } from "ws"; @@ -109,18 +110,42 @@ export class CoderApi extends Api { logs: ProvisionerJobLog[], options?: ClientOptions, ) => { + return this.watchLogs( + `/api/v2/workspacebuilds/${buildId}/logs`, + logs, + options, + ); + }; + + watchWorkspaceAgentLogs = async ( + agentId: string, + logs: WorkspaceAgentLog[], + options?: ClientOptions, + ) => { + return this.watchLogs( + `/api/v2/workspaceagents/${agentId}/logs`, + logs, + options, + ); + }; + + private async watchLogs( + apiRoute: string, + logs: { id: number }[], + options?: ClientOptions, + ) { const searchParams = new URLSearchParams({ follow: "true" }); const lastLog = logs.at(-1); if (lastLog) { searchParams.append("after", lastLog.id.toString()); } - return this.createWebSocket({ - apiRoute: `/api/v2/workspacebuilds/${buildId}/logs`, + return this.createWebSocket({ + apiRoute, searchParams, options, }); - }; + } private async createWebSocket( configs: Omit, diff --git a/src/api/workspace.ts b/src/api/workspace.ts index cb03d9fc..a24d3a64 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -1,11 +1,17 @@ -import { spawn } from "child_process"; import { type Api } from "coder/site/src/api/api"; -import { type Workspace } from "coder/site/src/api/typesGenerated"; +import { + type WorkspaceAgentLog, + type ProvisionerJobLog, + type Workspace, + type WorkspaceAgent, +} from "coder/site/src/api/typesGenerated"; +import { spawn } from "node:child_process"; import * as vscode from "vscode"; import { type FeatureSet } from "../featureSet"; import { getGlobalFlags } from "../globalFlags"; import { escapeCommandArg } from "../util"; +import { type OneWayWebSocket } from "../websocket/oneWayWebSocket"; import { errToStr, createWorkspaceIdentifier } from "./api-helper"; import { type CoderApi } from "./coderApi"; @@ -36,7 +42,7 @@ export async function startWorkspaceIfStoppedOrFailed( createWorkspaceIdentifier(workspace), ]; if (featureSet.buildReason) { - startArgs.push(...["--reason", "vscode_connection"]); + startArgs.push("--reason", "vscode_connection"); } // { shell: true } requires one shell-safe command string, otherwise we lose all escaping @@ -44,27 +50,25 @@ export async function startWorkspaceIfStoppedOrFailed( const startProcess = spawn(cmd, { shell: true }); startProcess.stdout.on("data", (data: Buffer) => { - data + const lines = data .toString() .split(/\r*\n/) - .forEach((line: string) => { - if (line !== "") { - writeEmitter.fire(line.toString() + "\r\n"); - } - }); + .filter((line) => line !== ""); + for (const line of lines) { + writeEmitter.fire(line.toString() + "\r\n"); + } }); let capturedStderr = ""; startProcess.stderr.on("data", (data: Buffer) => { - data + const lines = data .toString() .split(/\r*\n/) - .forEach((line: string) => { - if (line !== "") { - writeEmitter.fire(line.toString() + "\r\n"); - capturedStderr += line.toString() + "\n"; - } - }); + .filter((line) => line !== ""); + for (const line of lines) { + writeEmitter.fire(line.toString() + "\r\n"); + capturedStderr += line.toString() + "\n"; + } }); startProcess.on("close", (code: number) => { @@ -82,51 +86,72 @@ export async function startWorkspaceIfStoppedOrFailed( } /** - * Wait for the latest build to finish while streaming logs to the emitter. - * - * Once completed, fetch the workspace again and return it. + * Streams build logs to the emitter in real-time. + * Returns the websocket for lifecycle management. */ -export async function waitForBuild( +export async function streamBuildLogs( client: CoderApi, writeEmitter: vscode.EventEmitter, workspace: Workspace, -): Promise { - // This fetches the initial bunch of logs. - const logs = await client.getWorkspaceBuildLogs(workspace.latest_build.id); - logs.forEach((log) => writeEmitter.fire(log.output + "\r\n")); - +): Promise> { const socket = await client.watchBuildLogsByBuildId( workspace.latest_build.id, - logs, + [], ); - await new Promise((resolve, reject) => { - socket.addEventListener("message", (data) => { - if (data.parseError) { - writeEmitter.fire( - errToStr(data.parseError, "Failed to parse message") + "\r\n", - ); - } else { - writeEmitter.fire(data.parsedMessage.output + "\r\n"); - } - }); + socket.addEventListener("message", (data) => { + if (data.parseError) { + writeEmitter.fire( + errToStr(data.parseError, "Failed to parse message") + "\r\n", + ); + } else { + writeEmitter.fire(data.parsedMessage.output + "\r\n"); + } + }); + + socket.addEventListener("error", (error) => { + const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; + writeEmitter.fire( + `Error watching workspace build logs on ${baseUrlRaw}: ${errToStr(error, "no further details")}\r\n`, + ); + }); + + socket.addEventListener("close", () => { + writeEmitter.fire("Build complete\r\n"); + }); + + return socket; +} - socket.addEventListener("error", (error) => { - const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; - return reject( - new Error( - `Failed to watch workspace build on ${baseUrlRaw}: ${errToStr(error, "no further details")}`, - ), +/** + * Streams agent logs to the emitter in real-time. + * Returns the websocket for lifecycle management. + */ +export async function streamAgentLogs( + client: CoderApi, + writeEmitter: vscode.EventEmitter, + agent: WorkspaceAgent, +): Promise> { + const socket = await client.watchWorkspaceAgentLogs(agent.id, []); + + socket.addEventListener("message", (data) => { + if (data.parseError) { + writeEmitter.fire( + errToStr(data.parseError, "Failed to parse message") + "\r\n", ); - }); + } else { + for (const log of data.parsedMessage) { + writeEmitter.fire(log.output + "\r\n"); + } + } + }); - socket.addEventListener("close", () => resolve()); + socket.addEventListener("error", (error) => { + const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; + writeEmitter.fire( + `Error watching agent logs on ${baseUrlRaw}: ${errToStr(error, "no further details")}\r\n`, + ); }); - writeEmitter.fire("Build complete\r\n"); - const updatedWorkspace = await client.getWorkspace(workspace.id); - writeEmitter.fire( - `Workspace is now ${updatedWorkspace.latest_build.status}\r\n`, - ); - return updatedWorkspace; + return socket; } diff --git a/src/commands.ts b/src/commands.ts index 5abeb026..682d745b 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -19,6 +19,7 @@ import { type SecretsManager } from "./core/secretsManager"; import { CertificateError } from "./error"; import { getGlobalFlags } from "./globalFlags"; import { type Logger } from "./logging/logger"; +import { maybeAskAgent, maybeAskUrl } from "./promptUtils"; import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util"; import { AgentTreeItem, @@ -58,129 +59,6 @@ export class Commands { this.contextManager = serviceContainer.getContextManager(); } - /** - * Find the requested agent if specified, otherwise return the agent if there - * is only one or ask the user to pick if there are multiple. Return - * undefined if the user cancels. - */ - public async maybeAskAgent( - agents: WorkspaceAgent[], - filter?: string, - ): Promise { - const filteredAgents = filter - ? agents.filter((agent) => agent.name === filter) - : agents; - if (filteredAgents.length === 0) { - throw new Error("Workspace has no matching agents"); - } else if (filteredAgents.length === 1) { - return filteredAgents[0]; - } else { - const quickPick = vscode.window.createQuickPick(); - quickPick.title = "Select an agent"; - quickPick.busy = true; - const agentItems: vscode.QuickPickItem[] = filteredAgents.map((agent) => { - let icon = "$(debug-start)"; - if (agent.status !== "connected") { - icon = "$(debug-stop)"; - } - return { - alwaysShow: true, - label: `${icon} ${agent.name}`, - detail: `${agent.name} • Status: ${agent.status}`, - }; - }); - quickPick.items = agentItems; - quickPick.busy = false; - quickPick.show(); - - const selected = await new Promise( - (resolve) => { - quickPick.onDidHide(() => resolve(undefined)); - quickPick.onDidChangeSelection((selected) => { - if (selected.length < 1) { - return resolve(undefined); - } - const agent = filteredAgents[quickPick.items.indexOf(selected[0])]; - resolve(agent); - }); - }, - ); - quickPick.dispose(); - return selected; - } - } - - /** - * Ask the user for the URL, letting them choose from a list of recent URLs or - * CODER_URL or enter a new one. Undefined means the user aborted. - */ - private async askURL(selection?: string): Promise { - const defaultURL = vscode.workspace - .getConfiguration() - .get("coder.defaultUrl") - ?.trim(); - const quickPick = vscode.window.createQuickPick(); - quickPick.value = - selection || defaultURL || process.env.CODER_URL?.trim() || ""; - quickPick.placeholder = "https://example.coder.com"; - quickPick.title = "Enter the URL of your Coder deployment."; - - // Initial items. - quickPick.items = this.mementoManager - .withUrlHistory(defaultURL, process.env.CODER_URL) - .map((url) => ({ - alwaysShow: true, - label: url, - })); - - // Quick picks do not allow arbitrary values, so we add the value itself as - // an option in case the user wants to connect to something that is not in - // the list. - quickPick.onDidChangeValue((value) => { - quickPick.items = this.mementoManager - .withUrlHistory(defaultURL, process.env.CODER_URL, value) - .map((url) => ({ - alwaysShow: true, - label: url, - })); - }); - - quickPick.show(); - - const selected = await new Promise((resolve) => { - quickPick.onDidHide(() => resolve(undefined)); - quickPick.onDidChangeSelection((selected) => resolve(selected[0]?.label)); - }); - quickPick.dispose(); - return selected; - } - - /** - * Ask the user for the URL if it was not provided, letting them choose from a - * list of recent URLs or the default URL or CODER_URL or enter a new one, and - * normalizes the returned URL. Undefined means the user aborted. - */ - public async maybeAskUrl( - providedUrl: string | undefined | null, - lastUsedUrl?: string, - ): Promise { - let url = providedUrl || (await this.askURL(lastUsedUrl)); - if (!url) { - // User aborted. - return undefined; - } - - // Normalize URL. - if (!url.startsWith("http://") && !url.startsWith("https://")) { - // Default to HTTPS if not provided so URLs can be typed more easily. - url = "https://" + url; - } - while (url.endsWith("/")) { - url = url.substring(0, url.length - 1); - } - return url; - } - /** * Log into the provided deployment. If the deployment URL is not specified, * ask for it first with a menu showing recent URLs along with the default URL @@ -197,7 +75,7 @@ export class Commands { } this.logger.info("Logging in"); - const url = await this.maybeAskUrl(args?.url); + const url = await maybeAskUrl(this.mementoManager, args?.url); if (!url) { return; // The user aborted. } @@ -488,7 +366,7 @@ export class Commands { ); } else if (item instanceof WorkspaceTreeItem) { const agents = await this.extractAgentsWithFallback(item.workspace); - const agent = await this.maybeAskAgent(agents); + const agent = await maybeAskAgent(agents); if (!agent) { // User declined to pick an agent. return; @@ -611,7 +489,7 @@ export class Commands { } const agents = await this.extractAgentsWithFallback(workspace); - const agent = await this.maybeAskAgent(agents, agentName); + const agent = await maybeAskAgent(agents, agentName); if (!agent) { // User declined to pick an agent. return; diff --git a/src/extension.ts b/src/extension.ts index aba94cfe..cbb9e62e 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -12,6 +12,7 @@ import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; import { AuthAction } from "./core/secretsManager"; import { CertificateError, getErrorDetail } from "./error"; +import { maybeAskUrl } from "./promptUtils"; import { Remote } from "./remote/remote"; import { toSafeHost } from "./util"; import { @@ -147,7 +148,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // queries will default to localhost) so ask for it if missing. // Pre-populate in case we do have the right URL so the user can just // hit enter and move on. - const url = await commands.maybeAskUrl( + const url = await maybeAskUrl( + mementoManager, params.get("url"), mementoManager.getUrl(), ); @@ -230,7 +232,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // queries will default to localhost) so ask for it if missing. // Pre-populate in case we do have the right URL so the user can just // hit enter and move on. - const url = await commands.maybeAskUrl( + const url = await maybeAskUrl( + mementoManager, params.get("url"), mementoManager.getUrl(), ); diff --git a/src/promptUtils.ts b/src/promptUtils.ts new file mode 100644 index 00000000..4d058f12 --- /dev/null +++ b/src/promptUtils.ts @@ -0,0 +1,131 @@ +import { type WorkspaceAgent } from "coder/site/src/api/typesGenerated"; +import * as vscode from "vscode"; + +import { type MementoManager } from "./core/mementoManager"; + +/** + * Find the requested agent if specified, otherwise return the agent if there + * is only one or ask the user to pick if there are multiple. Return + * undefined if the user cancels. + */ +export async function maybeAskAgent( + agents: WorkspaceAgent[], + filter?: string, +): Promise { + const filteredAgents = filter + ? agents.filter((agent) => agent.name === filter) + : agents; + if (filteredAgents.length === 0) { + throw new Error("Workspace has no matching agents"); + } else if (filteredAgents.length === 1) { + return filteredAgents[0]; + } else { + const quickPick = vscode.window.createQuickPick(); + quickPick.title = "Select an agent"; + quickPick.busy = true; + const agentItems: vscode.QuickPickItem[] = filteredAgents.map((agent) => { + let icon = "$(debug-start)"; + if (agent.status !== "connected") { + icon = "$(debug-stop)"; + } + return { + alwaysShow: true, + label: `${icon} ${agent.name}`, + detail: `${agent.name} • Status: ${agent.status}`, + }; + }); + quickPick.items = agentItems; + quickPick.busy = false; + quickPick.show(); + + const selected = await new Promise( + (resolve) => { + quickPick.onDidHide(() => resolve(undefined)); + quickPick.onDidChangeSelection((selected) => { + if (selected.length < 1) { + return resolve(undefined); + } + const agent = filteredAgents[quickPick.items.indexOf(selected[0])]; + resolve(agent); + }); + }, + ); + quickPick.dispose(); + return selected; + } +} + +/** + * Ask the user for the URL, letting them choose from a list of recent URLs or + * CODER_URL or enter a new one. Undefined means the user aborted. + */ +async function askURL( + mementoManager: MementoManager, + selection?: string, +): Promise { + const defaultURL = vscode.workspace + .getConfiguration() + .get("coder.defaultUrl") + ?.trim(); + const quickPick = vscode.window.createQuickPick(); + quickPick.value = + selection || defaultURL || process.env.CODER_URL?.trim() || ""; + quickPick.placeholder = "https://example.coder.com"; + quickPick.title = "Enter the URL of your Coder deployment."; + + // Initial items. + quickPick.items = mementoManager + .withUrlHistory(defaultURL, process.env.CODER_URL) + .map((url) => ({ + alwaysShow: true, + label: url, + })); + + // Quick picks do not allow arbitrary values, so we add the value itself as + // an option in case the user wants to connect to something that is not in + // the list. + quickPick.onDidChangeValue((value) => { + quickPick.items = mementoManager + .withUrlHistory(defaultURL, process.env.CODER_URL, value) + .map((url) => ({ + alwaysShow: true, + label: url, + })); + }); + + quickPick.show(); + + const selected = await new Promise((resolve) => { + quickPick.onDidHide(() => resolve(undefined)); + quickPick.onDidChangeSelection((selected) => resolve(selected[0]?.label)); + }); + quickPick.dispose(); + return selected; +} + +/** + * Ask the user for the URL if it was not provided, letting them choose from a + * list of recent URLs or the default URL or CODER_URL or enter a new one, and + * normalizes the returned URL. Undefined means the user aborted. + */ +export async function maybeAskUrl( + mementoManager: MementoManager, + providedUrl: string | undefined | null, + lastUsedUrl?: string, +): Promise { + let url = providedUrl || (await askURL(mementoManager, lastUsedUrl)); + if (!url) { + // User aborted. + return undefined; + } + + // Normalize URL. + if (!url.startsWith("http://") && !url.startsWith("https://")) { + // Default to HTTPS if not provided so URLs can be typed more easily. + url = "https://" + url; + } + while (url.endsWith("/")) { + url = url.substring(0, url.length - 1); + } + return url; +} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 97cb858e..79b9274c 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -5,10 +5,10 @@ import { type WorkspaceAgent, } from "coder/site/src/api/typesGenerated"; import find from "find-process"; -import * as fs from "fs/promises"; import * as jsonc from "jsonc-parser"; -import * as os from "os"; -import * as path from "path"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; import prettyBytes from "pretty-bytes"; import * as semver from "semver"; import * as vscode from "vscode"; @@ -19,13 +19,9 @@ import { formatEventLabel, formatMetadataError, } from "../api/agentMetadataHelper"; -import { createWorkspaceIdentifier, extractAgents } from "../api/api-helper"; +import { extractAgents } from "../api/api-helper"; import { CoderApi } from "../api/coderApi"; import { needToken } from "../api/utils"; -import { - startWorkspaceIfStoppedOrFailed, - waitForBuild, -} from "../api/workspace"; import { type Commands } from "../commands"; import { type CliManager } from "../core/cliManager"; import * as cliUtils from "../core/cliUtils"; @@ -47,6 +43,7 @@ import { WorkspaceMonitor } from "../workspace/workspaceMonitor"; import { SSHConfig, type SSHValues, mergeSSHConfigValues } from "./sshConfig"; import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"; +import { WorkspaceStateMachine } from "./workspaceStateMachine"; export interface RemoteDetails extends vscode.Disposable { url: string; @@ -104,147 +101,6 @@ export class Remote { } } - private async confirmStart(workspaceName: string): Promise { - const action = await this.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"; - } - - /** - * Try to get the workspace running. Return undefined if the user canceled. - */ - private async maybeWaitForRunning( - client: CoderApi, - workspace: Workspace, - label: string, - binPath: string, - featureSet: FeatureSet, - firstConnect: boolean, - ): Promise { - const workspaceName = createWorkspaceIdentifier(workspace); - - // A terminal will be used to stream the build, if one is necessary. - let writeEmitter: undefined | vscode.EventEmitter; - let terminal: undefined | vscode.Terminal; - let attempts = 0; - - function initWriteEmitterAndTerminal(): vscode.EventEmitter { - writeEmitter ??= new vscode.EventEmitter(); - if (!terminal) { - terminal = vscode.window.createTerminal({ - name: "Build Log", - location: vscode.TerminalLocation.Panel, - // Spin makes this gear icon spin! - iconPath: new vscode.ThemeIcon("gear~spin"), - pty: { - onDidWrite: writeEmitter.event, - close: () => undefined, - open: () => undefined, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as Partial as any, - }); - terminal.show(true); - } - return writeEmitter; - } - - try { - // Show a notification while we wait. - return await this.vscodeProposed.window.withProgress( - { - location: vscode.ProgressLocation.Notification, - cancellable: false, - title: "Waiting for workspace build...", - }, - async () => { - const globalConfigDir = this.pathResolver.getGlobalConfigDir(label); - while (workspace.latest_build.status !== "running") { - ++attempts; - switch (workspace.latest_build.status) { - case "pending": - case "starting": - case "stopping": - writeEmitter = initWriteEmitterAndTerminal(); - this.logger.info(`Waiting for ${workspaceName}...`); - workspace = await waitForBuild(client, writeEmitter, workspace); - break; - case "stopped": - if ( - !firstConnect && - !(await this.confirmStart(workspaceName)) - ) { - return undefined; - } - writeEmitter = initWriteEmitterAndTerminal(); - this.logger.info(`Starting ${workspaceName}...`); - workspace = await startWorkspaceIfStoppedOrFailed( - client, - globalConfigDir, - binPath, - workspace, - writeEmitter, - featureSet, - ); - break; - case "failed": - // On a first attempt, we will try starting a failed workspace - // (for example canceling a start seems to cause this state). - if (attempts === 1) { - if ( - !firstConnect && - !(await this.confirmStart(workspaceName)) - ) { - return undefined; - } - writeEmitter = initWriteEmitterAndTerminal(); - this.logger.info(`Starting ${workspaceName}...`); - workspace = await startWorkspaceIfStoppedOrFailed( - client, - globalConfigDir, - binPath, - workspace, - writeEmitter, - featureSet, - ); - break; - } - // Otherwise fall through and error. - case "canceled": - case "canceling": - case "deleted": - case "deleting": - default: { - const is = - workspace.latest_build.status === "failed" ? "has" : "is"; - throw new Error( - `${workspaceName} ${is} ${workspace.latest_build.status}`, - ); - } - } - this.logger.info( - `${workspaceName} status is now`, - workspace.latest_build.status, - ); - } - return workspace; - }, - ); - } finally { - if (writeEmitter) { - writeEmitter.dispose(); - } - if (terminal) { - terminal.dispose(); - } - } - } - /** * Ensure the workspace specified by the remote authority is ready to receive * SSH connections. Return undefined if the authority is not for a Coder @@ -427,36 +283,104 @@ export class Remote { dispose: () => labelFormatterDisposable.dispose(), }); - // If the workspace is not in a running state, try to get it running. - if (workspace.latest_build.status !== "running") { - const updatedWorkspace = await this.maybeWaitForRunning( - workspaceClient, - workspace, - parts.label, - binaryPath, - featureSet, - firstConnect, + // Watch the workspace for changes. + const monitor = await WorkspaceMonitor.create( + workspace, + workspaceClient, + this.logger, + this.vscodeProposed, + this.contextManager, + ); + disposables.push( + monitor, + monitor.onChange.event((w) => (this.commands.workspace = w)), + ); + + // Wait for workspace to be running and agent to be ready + const stateMachine = new WorkspaceStateMachine( + parts, + workspaceClient, + firstConnect, + binaryPath, + featureSet, + this.logger, + this.pathResolver, + this.vscodeProposed, + ); + disposables.push(stateMachine); + + try { + workspace = await this.vscodeProposed.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + cancellable: false, + title: "Preparing workspace", + }, + async (progress) => { + let inProgress = false; + let pendingWorkspace: Workspace | null = null; + + return new Promise((resolve, reject) => { + const processWorkspace = async (w: Workspace) => { + if (inProgress) { + // Process one workspace at a time, keeping only the last + pendingWorkspace = w; + return; + } + + inProgress = true; + try { + pendingWorkspace = null; + + const isReady = await stateMachine.processWorkspace( + w, + progress, + ); + if (isReady) { + subscription.dispose(); + resolve(w); + return; + } + } catch (error) { + subscription.dispose(); + reject(error); + } finally { + inProgress = false; + } + + if (pendingWorkspace) { + processWorkspace(pendingWorkspace); + } + }; + + processWorkspace(workspace); + const subscription = monitor.onChange.event(async (w) => + processWorkspace(w), + ); + }); + }, ); - if (!updatedWorkspace) { - // User declined to start the workspace. - await this.closeRemote(); - return; - } - workspace = updatedWorkspace; + } finally { + stateMachine.dispose(); } - this.commands.workspace = workspace; - // Pick an agent. - this.logger.info(`Finding agent for ${workspaceName}...`); + // Mark initial setup as complete so the monitor can start notifying about state changes + monitor.markInitialSetupComplete(); + const agents = extractAgents(workspace.latest_build.resources); - const gotAgent = await this.commands.maybeAskAgent(agents, parts.agent); - if (!gotAgent) { - // User declined to pick an agent. - await this.closeRemote(); - return; + const agent = agents.find( + (agent) => agent.id === stateMachine.getAgentId(), + ); + + if (!agent) { + throw new Error("Failed to get workspace or agent from state machine"); } - let agent = gotAgent; // Reassign so it cannot be undefined in callbacks. - this.logger.info(`Found agent ${agent.name} with status`, agent.status); + + this.commands.workspace = workspace; + + // Watch coder inbox for messages + const inbox = await Inbox.create(workspace, workspaceClient, this.logger); + disposables.push(inbox); // Do some janky setting manipulation. this.logger.info("Modifying settings..."); @@ -538,76 +462,6 @@ export class Remote { } } - // Watch the workspace for changes. - const monitor = await WorkspaceMonitor.create( - workspace, - workspaceClient, - this.logger, - this.vscodeProposed, - this.contextManager, - ); - disposables.push(monitor); - disposables.push( - monitor.onChange.event((w) => (this.commands.workspace = w)), - ); - - // Watch coder inbox for messages - const inbox = await Inbox.create(workspace, workspaceClient, this.logger); - disposables.push(inbox); - - // Wait for the agent to connect. - if (agent.status === "connecting") { - this.logger.info(`Waiting for ${workspaceName}/${agent.name}...`); - await vscode.window.withProgress( - { - title: "Waiting for the agent to connect...", - location: vscode.ProgressLocation.Notification, - }, - async () => { - await new Promise((resolve) => { - const updateEvent = monitor.onChange.event((workspace) => { - if (!agent) { - return; - } - const agents = extractAgents(workspace.latest_build.resources); - const found = agents.find((newAgent) => { - return newAgent.id === agent.id; - }); - if (!found) { - return; - } - agent = found; - if (agent.status === "connecting") { - return; - } - updateEvent.dispose(); - resolve(); - }); - }); - }, - ); - this.logger.info(`Agent ${agent.name} status is now`, agent.status); - } - - // Make sure the agent is connected. - // TODO: Should account for the lifecycle state as well? - if (agent.status !== "connected") { - const result = await this.vscodeProposed.window.showErrorMessage( - `${workspaceName}/${agent.name} ${agent.status}`, - { - useCustom: true, - modal: true, - detail: `The ${agent.name} agent failed to connect. Try restarting your workspace.`, - }, - ); - if (!result) { - await this.closeRemote(); - return; - } - await this.reloadWindow(); - return; - } - const logDir = this.getLogDir(featureSet); // This ensures the Remote SSH extension resolves the host to execute the diff --git a/src/remote/terminalSession.ts b/src/remote/terminalSession.ts new file mode 100644 index 00000000..358134a1 --- /dev/null +++ b/src/remote/terminalSession.ts @@ -0,0 +1,39 @@ +import * as vscode from "vscode"; + +/** + * Manages a terminal and its associated write emitter as a single unit. + * Ensures both are created together and disposed together properly. + */ +export class TerminalSession implements vscode.Disposable { + public readonly writeEmitter: vscode.EventEmitter; + public readonly terminal: vscode.Terminal; + + constructor(name: string) { + this.writeEmitter = new vscode.EventEmitter(); + this.terminal = vscode.window.createTerminal({ + name, + location: vscode.TerminalLocation.Panel, + // Spin makes this gear icon spin! + iconPath: new vscode.ThemeIcon("gear~spin"), + pty: { + onDidWrite: this.writeEmitter.event, + close: () => undefined, + open: () => undefined, + }, + }); + this.terminal.show(true); + } + + dispose(): void { + try { + this.writeEmitter.dispose(); + } catch { + // Ignore disposal errors + } + try { + this.terminal.dispose(); + } catch { + // Ignore disposal errors + } + } +} diff --git a/src/remote/workspaceStateMachine.ts b/src/remote/workspaceStateMachine.ts new file mode 100644 index 00000000..2251cd19 --- /dev/null +++ b/src/remote/workspaceStateMachine.ts @@ -0,0 +1,274 @@ +import { type AuthorityParts } from "src/util"; + +import { createWorkspaceIdentifier, extractAgents } from "../api/api-helper"; +import { + startWorkspaceIfStoppedOrFailed, + streamAgentLogs, + streamBuildLogs, +} from "../api/workspace"; +import { maybeAskAgent } from "../promptUtils"; + +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 { PathResolver } from "../core/pathResolver"; +import type { FeatureSet } from "../featureSet"; +import type { Logger } from "../logging/logger"; +import type { OneWayWebSocket } from "../websocket/oneWayWebSocket"; + +/** + * 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 agentId: string | undefined; + + private buildLogSocket: { + socket: OneWayWebSocket | null; + buildId: string | null; + } = { socket: null, buildId: null }; + + private agentLogSocket: OneWayWebSocket | null = null; + + 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 pathResolver: PathResolver, + private readonly vscodeProposed: typeof vscode, + ) { + this.terminal = new TerminalSession("Agent Log"); + } + + /** + * 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 { + const workspaceName = createWorkspaceIdentifier(workspace); + + switch (workspace.latest_build.status) { + case "running": + this.closeBuildLogSocket(); + break; + + case "stopped": + case "failed": { + this.closeBuildLogSocket(); + + if (!this.firstConnect && !(await this.confirmStart(workspaceName))) { + throw new Error(`User declined to start ${workspaceName}`); + } + + progress?.report({ message: `Starting ${workspaceName}...` }); + this.logger.info(`Starting ${workspaceName}...`); + const globalConfigDir = this.pathResolver.getGlobalConfigDir( + this.parts.label, + ); + await startWorkspaceIfStoppedOrFailed( + this.workspaceClient, + globalConfigDir, + 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 ID since it could change after a restart + this.agentId = undefined; + progress?.report({ message: "Waiting for workspace build..." }); + this.logger.info(`Waiting for ${workspaceName}...`); + + if (!this.buildLogSocket.socket) { + const socket = await streamBuildLogs( + this.workspaceClient, + this.terminal.writeEmitter, + workspace, + ); + this.buildLogSocket = { + socket, + buildId: workspace.latest_build.id, + }; + } + return false; + + case "deleted": + case "deleting": + case "canceled": + case "canceling": + this.closeBuildLogSocket(); + throw new Error(`${workspaceName} is ${workspace.latest_build.status}`); + + default: + this.closeBuildLogSocket(); + throw new Error( + `${workspaceName} unknown status: ${workspace.latest_build.status}`, + ); + } + + const agents = extractAgents(workspace.latest_build.resources); + if (this.agentId === 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("User declined to pick an agent"); + } + this.agentId = gotAgent.id; + this.logger.info( + `Found agent ${gotAgent.name} with status`, + gotAgent.status, + ); + } + const agent = agents.find((a) => a.id === this.agentId); + if (!agent) { + throw new Error(`Agent not found in ${workspaceName} resources`); + } + + switch (agent.status) { + case "connecting": + progress?.report({ + message: `Waiting for agent ${agent.name} to connect...`, + }); + this.logger.debug(`Waiting for agent ${agent.name}...`); + return false; + + case "disconnected": + throw new Error(`${workspaceName}/${agent.name} disconnected`); + + case "timeout": + progress?.report({ + message: `Agent ${agent.name} timed out, continuing to wait...`, + }); + this.logger.debug( + `Agent ${agent.name} timed out, continuing to wait...`, + ); + return false; + + case "connected": + break; + + default: + throw new Error( + `${workspaceName}/${agent.name} unknown status: ${agent.status}`, + ); + } + + switch (agent.lifecycle_state) { + case "ready": + this.closeAgentLogSocket(); + return true; + + case "starting": { + const isBlocking = agent.scripts.some( + (script) => script.start_blocks_login, + ); + if (!isBlocking) { + return true; + } + + progress?.report({ + message: `Waiting for agent ${agent.name} startup scripts...`, + }); + this.logger.debug(`Waiting for agent ${agent.name} startup scripts...`); + + this.agentLogSocket ??= await streamAgentLogs( + this.workspaceClient, + this.terminal.writeEmitter, + agent, + ); + return false; + } + + case "created": + progress?.report({ + message: `Waiting for agent ${agent.name} to start...`, + }); + this.logger.debug( + `Waiting for ${workspaceName}/${agent.name} to start...`, + ); + return false; + + case "start_error": + this.closeAgentLogSocket(); + this.logger.info( + `Agent ${agent.name} startup script failed, but continuing...`, + ); + return true; + + case "start_timeout": + this.closeAgentLogSocket(); + this.logger.info( + `Agent ${agent.name} startup script timed out, but continuing...`, + ); + return true; + + case "off": + this.closeAgentLogSocket(); + throw new Error(`${workspaceName}/${agent.name} is off`); + + default: + this.closeAgentLogSocket(); + throw new Error( + `${workspaceName}/${agent.name} unknown lifecycle state: ${agent.lifecycle_state}`, + ); + } + } + + private closeBuildLogSocket(): void { + if (this.buildLogSocket.socket) { + this.buildLogSocket.socket.close(); + this.buildLogSocket = { socket: null, buildId: null }; + } + } + + private closeAgentLogSocket(): void { + if (this.agentLogSocket) { + this.agentLogSocket.close(); + this.agentLogSocket = null; + } + } + + private async confirmStart(workspaceName: string): Promise { + const action = await this.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.agentId; + } + + dispose(): void { + this.closeBuildLogSocket(); + this.closeAgentLogSocket(); + this.terminal.dispose(); + } +} diff --git a/src/workspace/workspaceMonitor.ts b/src/workspace/workspaceMonitor.ts index ceea8a91..d1bdefaa 100644 --- a/src/workspace/workspaceMonitor.ts +++ b/src/workspace/workspaceMonitor.ts @@ -29,6 +29,7 @@ export class WorkspaceMonitor implements vscode.Disposable { private notifiedDeletion = false; private notifiedOutdated = false; private notifiedNotRunning = false; + private completedInitialSetup = false; readonly onChange = new vscode.EventEmitter(); private readonly statusBarItem: vscode.StatusBarItem; @@ -110,6 +111,10 @@ export class WorkspaceMonitor implements vscode.Disposable { return monitor; } + public markInitialSetupComplete(): void { + this.completedInitialSetup = true; + } + /** * Permanently close the websocket. */ @@ -131,7 +136,10 @@ export class WorkspaceMonitor implements vscode.Disposable { this.maybeNotifyOutdated(workspace); this.maybeNotifyAutostop(workspace); this.maybeNotifyDeletion(workspace); - this.maybeNotifyNotRunning(workspace); + if (this.completedInitialSetup) { + // This instance might be created before the workspace is running + this.maybeNotifyNotRunning(workspace); + } } private maybeNotifyAutostop(workspace: Workspace) {