From 5e0f0cef3067dd2c9cc6e0ebde32c240e1f7c053 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 17 Oct 2025 16:11:01 +0300 Subject: [PATCH 1/9] Show Remote SSH Output panel on workspace start --- src/extension.ts | 20 ++++++++++++++++++++ src/remote/remote.ts | 4 ++++ 2 files changed, 24 insertions(+) diff --git a/src/extension.ts b/src/extension.ts index aba94cfe..3b52b21c 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -353,6 +353,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { }), ); + let shouldShowSshOutput = false; // Since the "onResolveRemoteAuthority:ssh-remote" activation event exists // in package.json we're able to perform actions before the authority is // resolved by the remote SSH extension. @@ -370,6 +371,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); if (details) { ctx.subscriptions.push(details); + shouldShowSshOutput = details.startedWorkspace; // Authenticate the plugin client which is used in the sidebar to display // workspaces belonging to this deployment. client.setHost(details.url); @@ -460,9 +462,27 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { } } } + + if (shouldShowSshOutput) { + showSshOutput(); + } } async function showTreeViewSearch(id: string): Promise { await vscode.commands.executeCommand(`${id}.focus`); await vscode.commands.executeCommand("list.find"); } + +function showSshOutput(): void { + for (const command of [ + "opensshremotes.showLog", + "windsurf-remote-openssh.showLog", + ]) { + /** + * We must not await this command because + * 1) it may not exist + * 2) it might cause the Remote SSH extension to be loaded synchronously + */ + void vscode.commands.executeCommand(command); + } +} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 97cb858e..3ff7ef9c 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -51,6 +51,7 @@ import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"; export interface RemoteDetails extends vscode.Disposable { url: string; token: string; + startedWorkspace: boolean; } export class Remote { @@ -415,6 +416,7 @@ export class Remote { } } + let startedWorkspace = false; const disposables: vscode.Disposable[] = []; try { // Register before connection so the label still displays! @@ -442,6 +444,7 @@ export class Remote { await this.closeRemote(); return; } + startedWorkspace = true; workspace = updatedWorkspace; } this.commands.workspace = workspace; @@ -681,6 +684,7 @@ export class Remote { return { url: baseUrlRaw, token, + startedWorkspace, dispose: () => { disposables.forEach((d) => d.dispose()); }, From 6e86cfee6b72a92dfd8bceacf5098260db8feda2 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 23 Oct 2025 16:49:29 +0300 Subject: [PATCH 2/9] Wait for startup script when launching the workspace --- src/api/coderApi.ts | 19 +++++++++ src/api/workspace.ts | 43 +++++++++++++++++++- src/extension.ts | 20 ---------- src/remote/remote.ts | 95 +++++++++++++++++++++++++++++++------------- 4 files changed, 129 insertions(+), 48 deletions(-) diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index da624bad..697bf9d2 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"; @@ -122,6 +123,24 @@ export class CoderApi extends Api { }); }; + watchWorkspaceAgentLogs = async ( + agentId: string, + logs: WorkspaceAgentLog[], + 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/workspaceagents/${agentId}/logs`, + searchParams, + options, + }); + }; + private async createWebSocket( configs: Omit, ) { diff --git a/src/api/workspace.ts b/src/api/workspace.ts index cb03d9fc..958a55a0 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -1,11 +1,15 @@ 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 Workspace, +} from "coder/site/src/api/typesGenerated"; 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"; @@ -81,6 +85,43 @@ 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. + */ +export async function writeAgentLogs( + client: CoderApi, + writeEmitter: vscode.EventEmitter, + agentId: string, +): Promise> { + // This fetches the initial bunch of logs. + const logs = await client.getWorkspaceAgentLogs(agentId); + logs.forEach((log) => writeEmitter.fire(log.output + "\r\n")); + + const socket = await client.watchWorkspaceAgentLogs(agentId, logs); + + socket.addEventListener("message", (data) => { + if (data.parseError) { + writeEmitter.fire( + errToStr(data.parseError, "Failed to parse message") + "\r\n", + ); + } else { + data.parsedMessage.forEach((message) => + writeEmitter.fire(message.output + "\r\n"), + ); + } + }); + + socket.addEventListener("error", (error) => { + const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; + throw new Error( + `Failed to watch workspace build on ${baseUrlRaw}: ${errToStr(error, "no further details")}`, + ); + }); + return socket; +} + /** * Wait for the latest build to finish while streaming logs to the emitter. * diff --git a/src/extension.ts b/src/extension.ts index 3b52b21c..aba94cfe 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -353,7 +353,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { }), ); - let shouldShowSshOutput = false; // Since the "onResolveRemoteAuthority:ssh-remote" activation event exists // in package.json we're able to perform actions before the authority is // resolved by the remote SSH extension. @@ -371,7 +370,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); if (details) { ctx.subscriptions.push(details); - shouldShowSshOutput = details.startedWorkspace; // Authenticate the plugin client which is used in the sidebar to display // workspaces belonging to this deployment. client.setHost(details.url); @@ -462,27 +460,9 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { } } } - - if (shouldShowSshOutput) { - showSshOutput(); - } } async function showTreeViewSearch(id: string): Promise { await vscode.commands.executeCommand(`${id}.focus`); await vscode.commands.executeCommand("list.find"); } - -function showSshOutput(): void { - for (const command of [ - "opensshremotes.showLog", - "windsurf-remote-openssh.showLog", - ]) { - /** - * We must not await this command because - * 1) it may not exist - * 2) it might cause the Remote SSH extension to be loaded synchronously - */ - void vscode.commands.executeCommand(command); - } -} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 3ff7ef9c..16d6cf00 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -25,6 +25,7 @@ import { needToken } from "../api/utils"; import { startWorkspaceIfStoppedOrFailed, waitForBuild, + writeAgentLogs, } from "../api/workspace"; import { type Commands } from "../commands"; import { type CliManager } from "../core/cliManager"; @@ -51,7 +52,6 @@ import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"; export interface RemoteDetails extends vscode.Disposable { url: string; token: string; - startedWorkspace: boolean; } export class Remote { @@ -131,29 +131,18 @@ export class Remote { 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 writeEmitter: vscode.EventEmitter | undefined; + let terminal: vscode.Terminal | undefined; 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); + const initBuildLogTerminal = () => { + if (!writeEmitter) { + const init = this.initWriteEmitterAndTerminal("Build Log"); + writeEmitter = init.writeEmitter; + terminal = init.terminal; } return writeEmitter; - } + }; try { // Show a notification while we wait. @@ -171,7 +160,7 @@ export class Remote { case "pending": case "starting": case "stopping": - writeEmitter = initWriteEmitterAndTerminal(); + writeEmitter = initBuildLogTerminal(); this.logger.info(`Waiting for ${workspaceName}...`); workspace = await waitForBuild(client, writeEmitter, workspace); break; @@ -182,7 +171,7 @@ export class Remote { ) { return undefined; } - writeEmitter = initWriteEmitterAndTerminal(); + writeEmitter = initBuildLogTerminal(); this.logger.info(`Starting ${workspaceName}...`); workspace = await startWorkspaceIfStoppedOrFailed( client, @@ -203,7 +192,7 @@ export class Remote { ) { return undefined; } - writeEmitter = initWriteEmitterAndTerminal(); + writeEmitter = initBuildLogTerminal(); this.logger.info(`Starting ${workspaceName}...`); workspace = await startWorkspaceIfStoppedOrFailed( client, @@ -246,6 +235,27 @@ export class Remote { } } + private initWriteEmitterAndTerminal(name: string): { + writeEmitter: vscode.EventEmitter; + terminal: vscode.Terminal; + } { + const writeEmitter = new vscode.EventEmitter(); + const terminal = vscode.window.createTerminal({ + name, + location: vscode.TerminalLocation.Panel, + // Spin makes this gear icon spin! + iconPath: new vscode.ThemeIcon("gear~spin"), + pty: { + onDidWrite: writeEmitter.event, + close: () => undefined, + open: () => undefined, + }, + }); + terminal.show(true); + + return { writeEmitter, terminal }; + } + /** * Ensure the workspace specified by the remote authority is ready to receive * SSH connections. Return undefined if the authority is not for a Coder @@ -416,7 +426,6 @@ export class Remote { } } - let startedWorkspace = false; const disposables: vscode.Disposable[] = []; try { // Register before connection so the label still displays! @@ -444,7 +453,6 @@ export class Remote { await this.closeRemote(); return; } - startedWorkspace = true; workspace = updatedWorkspace; } this.commands.workspace = workspace; @@ -593,7 +601,6 @@ export class Remote { } // 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}`, @@ -611,6 +618,41 @@ export class Remote { return; } + if (agent.lifecycle_state === "starting") { + const isBlocking = agent.scripts.some( + (script) => script.start_blocks_login, + ); + if (isBlocking) { + const { writeEmitter, terminal } = + this.initWriteEmitterAndTerminal("Agent Log"); + const socket = await writeAgentLogs( + workspaceClient, + writeEmitter, + agent.id, + ); + await new Promise((resolve) => { + const updateEvent = monitor.onChange.event((workspace) => { + const agents = extractAgents(workspace.latest_build.resources); + const found = agents.find((newAgent) => { + return newAgent.id === agent.id; + }); + if (!found) { + return; + } + agent = found; + if (agent.lifecycle_state === "starting") { + return; + } + updateEvent.dispose(); + resolve(); + }); + }); + writeEmitter.dispose(); + terminal.dispose(); + socket.close(); + } + } + const logDir = this.getLogDir(featureSet); // This ensures the Remote SSH extension resolves the host to execute the @@ -684,7 +726,6 @@ export class Remote { return { url: baseUrlRaw, token, - startedWorkspace, dispose: () => { disposables.forEach((d) => d.dispose()); }, From d545ba5fd50f2e0beb4fcb97768bd796a6bc7f77 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 24 Oct 2025 12:00:04 +0300 Subject: [PATCH 3/9] Refactoring and cleanup of agent startup script waiting --- src/api/workspace.ts | 122 +++++++++------- src/remote/remote.ts | 268 ++++++++++++++++++++++------------ src/remote/terminalSession.ts | 39 +++++ 3 files changed, 280 insertions(+), 149 deletions(-) create mode 100644 src/remote/terminalSession.ts diff --git a/src/api/workspace.ts b/src/api/workspace.ts index 958a55a0..c662e2af 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -1,9 +1,10 @@ -import { spawn } from "child_process"; import { type Api } from "coder/site/src/api/api"; import { type WorkspaceAgentLog, 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"; @@ -40,7 +41,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 @@ -48,27 +49,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) => { @@ -85,43 +84,6 @@ 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. - */ -export async function writeAgentLogs( - client: CoderApi, - writeEmitter: vscode.EventEmitter, - agentId: string, -): Promise> { - // This fetches the initial bunch of logs. - const logs = await client.getWorkspaceAgentLogs(agentId); - logs.forEach((log) => writeEmitter.fire(log.output + "\r\n")); - - const socket = await client.watchWorkspaceAgentLogs(agentId, logs); - - socket.addEventListener("message", (data) => { - if (data.parseError) { - writeEmitter.fire( - errToStr(data.parseError, "Failed to parse message") + "\r\n", - ); - } else { - data.parsedMessage.forEach((message) => - writeEmitter.fire(message.output + "\r\n"), - ); - } - }); - - socket.addEventListener("error", (error) => { - const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; - throw new Error( - `Failed to watch workspace build on ${baseUrlRaw}: ${errToStr(error, "no further details")}`, - ); - }); - return socket; -} - /** * Wait for the latest build to finish while streaming logs to the emitter. * @@ -134,7 +96,9 @@ export async function waitForBuild( ): 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")); + for (const log of logs) { + writeEmitter.fire(log.output + "\r\n"); + } const socket = await client.watchBuildLogsByBuildId( workspace.latest_build.id, @@ -171,3 +135,55 @@ export async function waitForBuild( ); return updatedWorkspace; } + +/** + * Streams agent logs to the emitter in real-time. + * Fetches existing logs and subscribes to new logs via websocket. + * Returns the websocket and a completion promise that rejects on error. + */ +export async function streamAgentLogs( + client: CoderApi, + writeEmitter: vscode.EventEmitter, + agent: WorkspaceAgent, +): Promise<{ + socket: OneWayWebSocket; + completion: Promise; +}> { + // This fetches the initial bunch of logs. + const logs = await client.getWorkspaceAgentLogs(agent.id); + for (const log of logs) { + writeEmitter.fire(log.output + "\r\n"); + } + + const socket = await client.watchWorkspaceAgentLogs(agent.id, logs); + + const completion = new Promise((resolve, reject) => { + 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("error", (error) => { + const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; + writeEmitter.fire( + `Error watching agent logs on ${baseUrlRaw}: ${errToStr(error, "no further details")}\r\n`, + ); + return reject( + new Error( + `Failed to watch agent logs on ${baseUrlRaw}: ${errToStr(error, "no further details")}`, + ), + ); + }); + + socket.addEventListener("close", () => resolve()); + }); + + return { socket, completion }; +} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 16d6cf00..833ed8e8 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -1,14 +1,15 @@ import { isAxiosError } from "axios"; import { type Api } from "coder/site/src/api/api"; import { + type WorkspaceAgentLog, type Workspace, 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"; @@ -24,8 +25,8 @@ import { CoderApi } from "../api/coderApi"; import { needToken } from "../api/utils"; import { startWorkspaceIfStoppedOrFailed, + streamAgentLogs, waitForBuild, - writeAgentLogs, } from "../api/workspace"; import { type Commands } from "../commands"; import { type CliManager } from "../core/cliManager"; @@ -44,10 +45,12 @@ import { findPort, parseRemoteAuthority, } from "../util"; +import { type OneWayWebSocket } from "../websocket/oneWayWebSocket"; import { WorkspaceMonitor } from "../workspace/workspaceMonitor"; import { SSHConfig, type SSHValues, mergeSSHConfigValues } from "./sshConfig"; import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"; +import { TerminalSession } from "./terminalSession"; export interface RemoteDetails extends vscode.Disposable { url: string; @@ -131,17 +134,12 @@ export class Remote { const workspaceName = createWorkspaceIdentifier(workspace); // A terminal will be used to stream the build, if one is necessary. - let writeEmitter: vscode.EventEmitter | undefined; - let terminal: vscode.Terminal | undefined; + let terminalSession: TerminalSession | undefined; let attempts = 0; - const initBuildLogTerminal = () => { - if (!writeEmitter) { - const init = this.initWriteEmitterAndTerminal("Build Log"); - writeEmitter = init.writeEmitter; - terminal = init.terminal; - } - return writeEmitter; + const getOrCreateTerminal = () => { + terminalSession ??= new TerminalSession("Build Log"); + return terminalSession.writeEmitter; }; try { @@ -160,9 +158,12 @@ export class Remote { case "pending": case "starting": case "stopping": - writeEmitter = initBuildLogTerminal(); this.logger.info(`Waiting for ${workspaceName}...`); - workspace = await waitForBuild(client, writeEmitter, workspace); + workspace = await waitForBuild( + client, + getOrCreateTerminal(), + workspace, + ); break; case "stopped": if ( @@ -171,14 +172,13 @@ export class Remote { ) { return undefined; } - writeEmitter = initBuildLogTerminal(); this.logger.info(`Starting ${workspaceName}...`); workspace = await startWorkspaceIfStoppedOrFailed( client, globalConfigDir, binPath, workspace, - writeEmitter, + getOrCreateTerminal(), featureSet, ); break; @@ -192,14 +192,13 @@ export class Remote { ) { return undefined; } - writeEmitter = initBuildLogTerminal(); this.logger.info(`Starting ${workspaceName}...`); workspace = await startWorkspaceIfStoppedOrFailed( client, globalConfigDir, binPath, workspace, - writeEmitter, + getOrCreateTerminal(), featureSet, ); break; @@ -226,36 +225,10 @@ export class Remote { }, ); } finally { - if (writeEmitter) { - writeEmitter.dispose(); - } - if (terminal) { - terminal.dispose(); - } + terminalSession?.dispose(); } } - private initWriteEmitterAndTerminal(name: string): { - writeEmitter: vscode.EventEmitter; - terminal: vscode.Terminal; - } { - const writeEmitter = new vscode.EventEmitter(); - const terminal = vscode.window.createTerminal({ - name, - location: vscode.TerminalLocation.Panel, - // Spin makes this gear icon spin! - iconPath: new vscode.ThemeIcon("gear~spin"), - pty: { - onDidWrite: writeEmitter.event, - close: () => undefined, - open: () => undefined, - }, - }); - terminal.show(true); - - return { writeEmitter, terminal }; - } - /** * Ensure the workspace specified by the remote authority is ready to receive * SSH connections. Return undefined if the authority is not for a Coder @@ -569,34 +542,25 @@ export class Remote { // Wait for the agent to connect. if (agent.status === "connecting") { this.logger.info(`Waiting for ${workspaceName}/${agent.name}...`); - await vscode.window.withProgress( + const updatedAgent = await this.waitForAgentWithProgress( + monitor, + agent, { - title: "Waiting for the agent to connect...", - location: vscode.ProgressLocation.Notification, + progressTitle: "Waiting for the agent to connect...", + timeoutMs: 3 * 60 * 1000, + errorDialogTitle: `Failed to wait for ${agent.name} connection`, }, - 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(); - }); - }); + (foundAgent) => { + if (foundAgent.status !== "connecting") { + return foundAgent; + } + return undefined; }, ); + if (!updatedAgent) { + return; + } + agent = updatedAgent; this.logger.info(`Agent ${agent.name} status is now`, agent.status); } @@ -623,33 +587,56 @@ export class Remote { (script) => script.start_blocks_login, ); if (isBlocking) { - const { writeEmitter, terminal } = - this.initWriteEmitterAndTerminal("Agent Log"); - const socket = await writeAgentLogs( - workspaceClient, - writeEmitter, - agent.id, + this.logger.info( + `Waiting for ${workspaceName}/${agent.name} startup...`, ); - await new Promise((resolve) => { - const updateEvent = monitor.onChange.event((workspace) => { - const agents = extractAgents(workspace.latest_build.resources); - const found = agents.find((newAgent) => { - return newAgent.id === agent.id; - }); - if (!found) { - return; - } - agent = found; - if (agent.lifecycle_state === "starting") { - return; - } - updateEvent.dispose(); - resolve(); - }); - }); - writeEmitter.dispose(); - terminal.dispose(); - socket.close(); + + let terminalSession: TerminalSession | undefined; + let socket: OneWayWebSocket | undefined; + try { + terminalSession = new TerminalSession("Agent Log"); + const { socket: agentSocket, completion: logsCompletion } = + await streamAgentLogs( + workspaceClient, + terminalSession.writeEmitter, + agent, + ); + socket = agentSocket; + + const agentStatePromise = this.waitForAgentWithProgress( + monitor, + agent, + { + progressTitle: "Waiting for agent startup scripts...", + timeoutMs: 5 * 60 * 1000, + errorDialogTitle: `Failed to wait for ${agent.name} startup`, + }, + (foundAgent) => { + if (foundAgent.lifecycle_state !== "starting") { + return foundAgent; + } + return undefined; + }, + ); + + // Race between logs completion (which fails on socket error) and agent state change + const updatedAgent = await Promise.race([ + agentStatePromise, + logsCompletion.then(() => agentStatePromise), + ]); + + if (!updatedAgent) { + return; + } + agent = updatedAgent; + this.logger.info( + `Agent ${agent.name} lifecycle state is now`, + agent.lifecycle_state, + ); + } finally { + terminalSession?.dispose(); + socket?.close(); + } } } @@ -778,6 +765,95 @@ export class Remote { return ` --log-dir ${escapeCommandArg(logDir)} -v`; } + /** + * Waits for an agent condition with progress notification and error handling. + * Returns the updated agent or undefined if the user chooses to close remote. + */ + private async waitForAgentWithProgress( + monitor: WorkspaceMonitor, + agent: WorkspaceAgent, + options: { + progressTitle: string; + timeoutMs: number; + errorDialogTitle: string; + }, + checker: (agent: WorkspaceAgent) => WorkspaceAgent | undefined, + ): Promise { + try { + const foundAgent = await vscode.window.withProgress( + { + title: options.progressTitle, + location: vscode.ProgressLocation.Notification, + }, + async () => + this.waitForAgentCondition( + monitor, + agent, + checker, + options.timeoutMs, + `Timeout: ${options.errorDialogTitle}`, + ), + ); + return foundAgent; + } catch (error) { + this.logger.error(options.errorDialogTitle, error); + const result = await this.vscodeProposed.window.showErrorMessage( + options.errorDialogTitle, + { + useCustom: true, + modal: true, + detail: error instanceof Error ? error.message : String(error), + }, + "Close Remote", + ); + if (result === "Close Remote") { + await this.closeRemote(); + } + return undefined; + } + } + + /** + * Waits for an agent condition to be met by monitoring workspace changes. + * Returns the result when the condition is met or throws an error on timeout. + */ + private async waitForAgentCondition( + monitor: WorkspaceMonitor, + agent: WorkspaceAgent, + checker: (agent: WorkspaceAgent) => T | undefined, + timeoutMs: number, + timeoutMessage: string, + ): Promise { + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + updateEvent.dispose(); + reject(new Error(timeoutMessage)); + }, timeoutMs); + + const updateEvent = monitor.onChange.event((workspace) => { + try { + const agents = extractAgents(workspace.latest_build.resources); + const foundAgent = agents.find((a) => a.id === agent.id); + if (!foundAgent) { + throw new Error( + `Agent ${agent.name} not found in workspace resources`, + ); + } + const result = checker(foundAgent); + if (result !== undefined) { + clearTimeout(timeout); + updateEvent.dispose(); + resolve(result); + } + } catch (error) { + clearTimeout(timeout); + updateEvent.dispose(); + reject(error); + } + }); + }); + } + // updateSSHConfig updates the SSH configuration with a wildcard that handles // all Coder entries. private async updateSSHConfig( 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 + } + } +} From 9bf0cd7f9712dda671467cf578644704db696002 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 24 Oct 2025 12:22:21 +0300 Subject: [PATCH 4/9] More refactoring --- src/api/coderApi.ts | 32 +++++++++++++++++++------------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index 697bf9d2..ef120ce4 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -110,17 +110,11 @@ export class CoderApi extends Api { logs: ProvisionerJobLog[], 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`, - searchParams, + return this.watchLogs( + `/api/v2/workspacebuilds/${buildId}/logs`, + logs, options, - }); + ); }; watchWorkspaceAgentLogs = async ( @@ -128,18 +122,30 @@ export class CoderApi extends Api { 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/workspaceagents/${agentId}/logs`, + return this.createWebSocket({ + apiRoute, searchParams, options, }); - }; + } private async createWebSocket( configs: Omit, From e4bd57f8e71f3593af3a601ca3f5a6d57a1ebe52 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 28 Oct 2025 11:33:36 +0300 Subject: [PATCH 5/9] Review comments --- src/remote/remote.ts | 107 +++++++++---------------------------------- 1 file changed, 22 insertions(+), 85 deletions(-) diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 833ed8e8..e1f4495b 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -545,21 +545,9 @@ export class Remote { const updatedAgent = await this.waitForAgentWithProgress( monitor, agent, - { - progressTitle: "Waiting for the agent to connect...", - timeoutMs: 3 * 60 * 1000, - errorDialogTitle: `Failed to wait for ${agent.name} connection`, - }, - (foundAgent) => { - if (foundAgent.status !== "connecting") { - return foundAgent; - } - return undefined; - }, + `Waiting for agent ${agent.name} to connect...`, + (foundAgent) => foundAgent.status !== "connecting", ); - if (!updatedAgent) { - return; - } agent = updatedAgent; this.logger.info(`Agent ${agent.name} status is now`, agent.status); } @@ -606,28 +594,15 @@ export class Remote { const agentStatePromise = this.waitForAgentWithProgress( monitor, agent, - { - progressTitle: "Waiting for agent startup scripts...", - timeoutMs: 5 * 60 * 1000, - errorDialogTitle: `Failed to wait for ${agent.name} startup`, - }, - (foundAgent) => { - if (foundAgent.lifecycle_state !== "starting") { - return foundAgent; - } - return undefined; - }, + `Waiting for agent ${agent.name} startup scripts...`, + (foundAgent) => foundAgent.lifecycle_state !== "starting", ); - // Race between logs completion (which fails on socket error) and agent state change + // Race between logs completion and agent state change const updatedAgent = await Promise.race([ agentStatePromise, logsCompletion.then(() => agentStatePromise), ]); - - if (!updatedAgent) { - return; - } agent = updatedAgent; this.logger.info( `Agent ${agent.name} lifecycle state is now`, @@ -766,70 +741,34 @@ export class Remote { } /** - * Waits for an agent condition with progress notification and error handling. - * Returns the updated agent or undefined if the user chooses to close remote. + * Waits for an agent condition with progress notification. */ private async waitForAgentWithProgress( monitor: WorkspaceMonitor, agent: WorkspaceAgent, - options: { - progressTitle: string; - timeoutMs: number; - errorDialogTitle: string; - }, - checker: (agent: WorkspaceAgent) => WorkspaceAgent | undefined, - ): Promise { - try { - const foundAgent = await vscode.window.withProgress( - { - title: options.progressTitle, - location: vscode.ProgressLocation.Notification, - }, - async () => - this.waitForAgentCondition( - monitor, - agent, - checker, - options.timeoutMs, - `Timeout: ${options.errorDialogTitle}`, - ), - ); - return foundAgent; - } catch (error) { - this.logger.error(options.errorDialogTitle, error); - const result = await this.vscodeProposed.window.showErrorMessage( - options.errorDialogTitle, - { - useCustom: true, - modal: true, - detail: error instanceof Error ? error.message : String(error), - }, - "Close Remote", - ); - if (result === "Close Remote") { - await this.closeRemote(); - } - return undefined; - } + progressTitle: string, + checker: (agent: WorkspaceAgent) => boolean, + ): Promise { + const foundAgent = await vscode.window.withProgress( + { + title: progressTitle, + location: vscode.ProgressLocation.Notification, + }, + async () => this.waitForAgentCondition(monitor, agent, checker), + ); + return foundAgent; } /** * Waits for an agent condition to be met by monitoring workspace changes. * Returns the result when the condition is met or throws an error on timeout. */ - private async waitForAgentCondition( + private async waitForAgentCondition( monitor: WorkspaceMonitor, agent: WorkspaceAgent, - checker: (agent: WorkspaceAgent) => T | undefined, - timeoutMs: number, - timeoutMessage: string, - ): Promise { - return new Promise((resolve, reject) => { - const timeout = setTimeout(() => { - updateEvent.dispose(); - reject(new Error(timeoutMessage)); - }, timeoutMs); - + checker: (agent: WorkspaceAgent) => boolean, + ): Promise { + return new Promise((resolve, reject) => { const updateEvent = monitor.onChange.event((workspace) => { try { const agents = extractAgents(workspace.latest_build.resources); @@ -841,12 +780,10 @@ export class Remote { } const result = checker(foundAgent); if (result !== undefined) { - clearTimeout(timeout); updateEvent.dispose(); - resolve(result); + resolve(foundAgent); } } catch (error) { - clearTimeout(timeout); updateEvent.dispose(); reject(error); } From 2bc26f9dc75545afa61fbc519d4ffef8198a31d1 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 28 Oct 2025 12:58:53 +0300 Subject: [PATCH 6/9] Add agent state machine --- src/api/workspace.ts | 1 - src/remote/remote.ts | 224 ++++++++++++++++++++++++++++--------------- 2 files changed, 146 insertions(+), 79 deletions(-) diff --git a/src/api/workspace.ts b/src/api/workspace.ts index c662e2af..80c67099 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -138,7 +138,6 @@ export async function waitForBuild( /** * Streams agent logs to the emitter in real-time. - * Fetches existing logs and subscribes to new logs via websocket. * Returns the websocket and a completion promise that rejects on error. */ export async function streamAgentLogs( diff --git a/src/remote/remote.ts b/src/remote/remote.ts index e1f4495b..386ebbe7 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -539,81 +539,13 @@ export class Remote { 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}...`); - const updatedAgent = await this.waitForAgentWithProgress( - monitor, - agent, - `Waiting for agent ${agent.name} to connect...`, - (foundAgent) => foundAgent.status !== "connecting", - ); - agent = updatedAgent; - this.logger.info(`Agent ${agent.name} status is now`, agent.status); - } - - // Make sure the agent is connected. - 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; - } - - if (agent.lifecycle_state === "starting") { - const isBlocking = agent.scripts.some( - (script) => script.start_blocks_login, - ); - if (isBlocking) { - this.logger.info( - `Waiting for ${workspaceName}/${agent.name} startup...`, - ); - - let terminalSession: TerminalSession | undefined; - let socket: OneWayWebSocket | undefined; - try { - terminalSession = new TerminalSession("Agent Log"); - const { socket: agentSocket, completion: logsCompletion } = - await streamAgentLogs( - workspaceClient, - terminalSession.writeEmitter, - agent, - ); - socket = agentSocket; - - const agentStatePromise = this.waitForAgentWithProgress( - monitor, - agent, - `Waiting for agent ${agent.name} startup scripts...`, - (foundAgent) => foundAgent.lifecycle_state !== "starting", - ); - - // Race between logs completion and agent state change - const updatedAgent = await Promise.race([ - agentStatePromise, - logsCompletion.then(() => agentStatePromise), - ]); - agent = updatedAgent; - this.logger.info( - `Agent ${agent.name} lifecycle state is now`, - agent.lifecycle_state, - ); - } finally { - terminalSession?.dispose(); - socket?.close(); - } - } - } + // Ensure agent is ready by handling all status and lifecycle states + agent = await this.ensureAgentReady( + monitor, + agent, + workspaceName, + workspaceClient, + ); const logDir = this.getLogDir(featureSet); @@ -740,6 +672,139 @@ export class Remote { return ` --log-dir ${escapeCommandArg(logDir)} -v`; } + /** + * Ensures agent is ready to connect by handling all status and lifecycle states. + * Throws an error if the agent cannot be made ready. + */ + private async ensureAgentReady( + monitor: WorkspaceMonitor, + agent: WorkspaceAgent, + workspaceName: string, + workspaceClient: CoderApi, + ): Promise { + let currentAgent = agent; + + while ( + currentAgent.status !== "connected" || + currentAgent.lifecycle_state !== "ready" + ) { + switch (currentAgent.status) { + case "connecting": + this.logger.info(`Waiting for agent ${currentAgent.name}...`); + currentAgent = await this.waitForAgentWithProgress( + monitor, + currentAgent, + `Waiting for agent ${currentAgent.name} to connect...`, + (foundAgent) => foundAgent.status !== "connecting", + ); + this.logger.info( + `Agent ${currentAgent.name} status is now`, + currentAgent.status, + ); + continue; + + case "connected": + // Agent connected, now handle lifecycle state + break; + + case "disconnected": + case "timeout": + throw new Error( + `${workspaceName}/${currentAgent.name} ${currentAgent.status}`, + ); + + default: + throw new Error( + `${workspaceName}/${currentAgent.name} unknown status: ${currentAgent.status}`, + ); + } + + // Handle agent lifecycle state (only when status is "connected") + switch (currentAgent.lifecycle_state) { + case "ready": + return currentAgent; + + case "starting": { + const isBlocking = currentAgent.scripts.some( + (script) => script.start_blocks_login, + ); + if (!isBlocking) { + return currentAgent; + } + + const logMsg = `Waiting for agent ${currentAgent.name} startup scripts...`; + this.logger.info(logMsg); + + let terminalSession: TerminalSession | undefined; + let socket: OneWayWebSocket | undefined; + try { + terminalSession = new TerminalSession("Agent Log"); + const { socket: agentSocket, completion: logsCompletion } = + await streamAgentLogs( + workspaceClient, + terminalSession.writeEmitter, + currentAgent, + ); + socket = agentSocket; + + const agentStatePromise = this.waitForAgentWithProgress( + monitor, + currentAgent, + logMsg, + (foundAgent) => foundAgent.lifecycle_state !== "starting", + ); + + // Race between logs completion and agent state change + currentAgent = await Promise.race([ + agentStatePromise, + logsCompletion.then(() => agentStatePromise), + ]); + this.logger.info( + `Agent ${currentAgent.name} lifecycle state is now`, + currentAgent.lifecycle_state, + ); + } finally { + terminalSession?.dispose(); + socket?.close(); + } + continue; + } + + case "created": + this.logger.info( + `Waiting for ${workspaceName}/${currentAgent.name} to start...`, + ); + currentAgent = await this.waitForAgentWithProgress( + monitor, + currentAgent, + `Waiting for agent ${currentAgent.name} to start...`, + (foundAgent) => foundAgent.lifecycle_state !== "created", + ); + this.logger.info( + `Agent ${currentAgent.name} lifecycle state is now`, + currentAgent.lifecycle_state, + ); + continue; + + case "off": + case "start_error": + case "start_timeout": + case "shutdown_error": + case "shutdown_timeout": + case "shutting_down": + throw new Error( + `${workspaceName}/${currentAgent.name} lifecycle state: ${currentAgent.lifecycle_state}`, + ); + + default: + throw new Error( + `${workspaceName}/${currentAgent.name} unknown lifecycle state: ${currentAgent.lifecycle_state}`, + ); + } + } + return currentAgent; + } + /** * Waits for an agent condition with progress notification. */ @@ -749,7 +814,7 @@ export class Remote { progressTitle: string, checker: (agent: WorkspaceAgent) => boolean, ): Promise { - const foundAgent = await vscode.window.withProgress( + const foundAgent = await this.vscodeProposed.window.withProgress( { title: progressTitle, location: vscode.ProgressLocation.Notification, @@ -770,6 +835,10 @@ export class Remote { ): Promise { return new Promise((resolve, reject) => { const updateEvent = monitor.onChange.event((workspace) => { + if (workspace.latest_build.status !== "running") { + const workspaceName = createWorkspaceIdentifier(workspace); + reject(new Error(`Workspace ${workspaceName} is not running.`)); + } try { const agents = extractAgents(workspace.latest_build.resources); const foundAgent = agents.find((a) => a.id === agent.id); @@ -778,8 +847,7 @@ export class Remote { `Agent ${agent.name} not found in workspace resources`, ); } - const result = checker(foundAgent); - if (result !== undefined) { + if (checker(foundAgent)) { updateEvent.dispose(); resolve(foundAgent); } From b18dd3357a715a38392ee992a2133862a648135b Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 28 Oct 2025 13:02:39 +0300 Subject: [PATCH 7/9] Rely on the websocket to get all the logs instead of fetching the initial ones beforehand --- src/api/workspace.ts | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/src/api/workspace.ts b/src/api/workspace.ts index 80c67099..d9ac928d 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -94,15 +94,9 @@ export async function waitForBuild( writeEmitter: vscode.EventEmitter, workspace: Workspace, ): Promise { - // This fetches the initial bunch of logs. - const logs = await client.getWorkspaceBuildLogs(workspace.latest_build.id); - for (const log of logs) { - writeEmitter.fire(log.output + "\r\n"); - } - const socket = await client.watchBuildLogsByBuildId( workspace.latest_build.id, - logs, + [], ); await new Promise((resolve, reject) => { @@ -148,13 +142,7 @@ export async function streamAgentLogs( socket: OneWayWebSocket; completion: Promise; }> { - // This fetches the initial bunch of logs. - const logs = await client.getWorkspaceAgentLogs(agent.id); - for (const log of logs) { - writeEmitter.fire(log.output + "\r\n"); - } - - const socket = await client.watchWorkspaceAgentLogs(agent.id, logs); + const socket = await client.watchWorkspaceAgentLogs(agent.id, []); const completion = new Promise((resolve, reject) => { socket.addEventListener("message", (data) => { From bc4ba72e58d2f9b0a7eeabbb811c4fabce20e785 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 30 Oct 2025 11:57:47 +0300 Subject: [PATCH 8/9] Add workspace state and agent state machines --- src/api/workspace.ts | 99 +++--- src/commands.ts | 130 +------- src/extension.ts | 7 +- src/promptUtils.ts | 131 ++++++++ src/remote/remote.ts | 468 ++++++---------------------- src/remote/workspaceStateMachine.ts | 272 ++++++++++++++++ src/workspace/workspaceMonitor.ts | 13 +- 7 files changed, 565 insertions(+), 555 deletions(-) create mode 100644 src/promptUtils.ts create mode 100644 src/remote/workspaceStateMachine.ts diff --git a/src/api/workspace.ts b/src/api/workspace.ts index d9ac928d..a24d3a64 100644 --- a/src/api/workspace.ts +++ b/src/api/workspace.ts @@ -1,6 +1,7 @@ import { type Api } from "coder/site/src/api/api"; import { type WorkspaceAgentLog, + type ProvisionerJobLog, type Workspace, type WorkspaceAgent, } from "coder/site/src/api/typesGenerated"; @@ -85,92 +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 { +): Promise> { const socket = await client.watchBuildLogsByBuildId( workspace.latest_build.id, [], ); - 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("error", (error) => { - const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; - return reject( - new Error( - `Failed to watch workspace build on ${baseUrlRaw}: ${errToStr(error, "no further details")}`, - ), + 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("close", () => resolve()); + 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`, + ); }); - 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; + socket.addEventListener("close", () => { + writeEmitter.fire("Build complete\r\n"); + }); + + return socket; } /** * Streams agent logs to the emitter in real-time. - * Returns the websocket and a completion promise that rejects on error. + * Returns the websocket for lifecycle management. */ export async function streamAgentLogs( client: CoderApi, writeEmitter: vscode.EventEmitter, agent: WorkspaceAgent, -): Promise<{ - socket: OneWayWebSocket; - completion: Promise; -}> { +): Promise> { const socket = await client.watchWorkspaceAgentLogs(agent.id, []); - const completion = new Promise((resolve, reject) => { - 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("error", (error) => { - const baseUrlRaw = client.getAxiosInstance().defaults.baseURL; + socket.addEventListener("message", (data) => { + if (data.parseError) { writeEmitter.fire( - `Error watching agent logs on ${baseUrlRaw}: ${errToStr(error, "no further details")}\r\n`, + errToStr(data.parseError, "Failed to parse message") + "\r\n", ); - return reject( - new Error( - `Failed to watch agent logs on ${baseUrlRaw}: ${errToStr(error, "no further details")}`, - ), - ); - }); + } 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`, + ); }); - return { socket, completion }; + 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 386ebbe7..3782a9ad 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -1,7 +1,6 @@ import { isAxiosError } from "axios"; import { type Api } from "coder/site/src/api/api"; import { - type WorkspaceAgentLog, type Workspace, type WorkspaceAgent, } from "coder/site/src/api/typesGenerated"; @@ -20,14 +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, - streamAgentLogs, - waitForBuild, -} from "../api/workspace"; import { type Commands } from "../commands"; import { type CliManager } from "../core/cliManager"; import * as cliUtils from "../core/cliUtils"; @@ -45,12 +39,11 @@ import { findPort, parseRemoteAuthority, } from "../util"; -import { type OneWayWebSocket } from "../websocket/oneWayWebSocket"; import { WorkspaceMonitor } from "../workspace/workspaceMonitor"; import { SSHConfig, type SSHValues, mergeSSHConfigValues } from "./sshConfig"; import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"; -import { TerminalSession } from "./terminalSession"; +import { WorkspaceStateMachine } from "./workspaceStateMachine"; export interface RemoteDetails extends vscode.Disposable { url: string; @@ -108,127 +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 terminalSession: TerminalSession | undefined; - let attempts = 0; - - const getOrCreateTerminal = () => { - terminalSession ??= new TerminalSession("Build Log"); - return terminalSession.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": - this.logger.info(`Waiting for ${workspaceName}...`); - workspace = await waitForBuild( - client, - getOrCreateTerminal(), - workspace, - ); - break; - case "stopped": - if ( - !firstConnect && - !(await this.confirmStart(workspaceName)) - ) { - return undefined; - } - this.logger.info(`Starting ${workspaceName}...`); - workspace = await startWorkspaceIfStoppedOrFailed( - client, - globalConfigDir, - binPath, - workspace, - getOrCreateTerminal(), - 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; - } - this.logger.info(`Starting ${workspaceName}...`); - workspace = await startWorkspaceIfStoppedOrFailed( - client, - globalConfigDir, - binPath, - workspace, - getOrCreateTerminal(), - 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 { - terminalSession?.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 @@ -411,36 +283,110 @@ 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); + disposables.push( + 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 { + await this.vscodeProposed.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + cancellable: false, + title: "Preparing workspace", + }, + async (progress) => { + let inProgress = false; + let pendingWorkspace: Workspace | null = null; + + await new Promise((resolve, reject) => { + const processWorkspace = async (w: Workspace) => { + workspace = w; + + 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(); + return; + } + + if (pendingWorkspace) { + const isReadyAfter = await stateMachine.processWorkspace( + pendingWorkspace, + progress, + ); + if (isReadyAfter) { + subscription.dispose(); + resolve(); + } + } + } catch (error) { + subscription.dispose(); + reject(error); + } finally { + inProgress = false; + } + }; + + 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}...`); 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..."); @@ -522,31 +468,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); - - // Ensure agent is ready by handling all status and lifecycle states - agent = await this.ensureAgentReady( - monitor, - agent, - workspaceName, - workspaceClient, - ); - const logDir = this.getLogDir(featureSet); // This ensures the Remote SSH extension resolves the host to execute the @@ -672,193 +593,6 @@ export class Remote { return ` --log-dir ${escapeCommandArg(logDir)} -v`; } - /** - * Ensures agent is ready to connect by handling all status and lifecycle states. - * Throws an error if the agent cannot be made ready. - */ - private async ensureAgentReady( - monitor: WorkspaceMonitor, - agent: WorkspaceAgent, - workspaceName: string, - workspaceClient: CoderApi, - ): Promise { - let currentAgent = agent; - - while ( - currentAgent.status !== "connected" || - currentAgent.lifecycle_state !== "ready" - ) { - switch (currentAgent.status) { - case "connecting": - this.logger.info(`Waiting for agent ${currentAgent.name}...`); - currentAgent = await this.waitForAgentWithProgress( - monitor, - currentAgent, - `Waiting for agent ${currentAgent.name} to connect...`, - (foundAgent) => foundAgent.status !== "connecting", - ); - this.logger.info( - `Agent ${currentAgent.name} status is now`, - currentAgent.status, - ); - continue; - - case "connected": - // Agent connected, now handle lifecycle state - break; - - case "disconnected": - case "timeout": - throw new Error( - `${workspaceName}/${currentAgent.name} ${currentAgent.status}`, - ); - - default: - throw new Error( - `${workspaceName}/${currentAgent.name} unknown status: ${currentAgent.status}`, - ); - } - - // Handle agent lifecycle state (only when status is "connected") - switch (currentAgent.lifecycle_state) { - case "ready": - return currentAgent; - - case "starting": { - const isBlocking = currentAgent.scripts.some( - (script) => script.start_blocks_login, - ); - if (!isBlocking) { - return currentAgent; - } - - const logMsg = `Waiting for agent ${currentAgent.name} startup scripts...`; - this.logger.info(logMsg); - - let terminalSession: TerminalSession | undefined; - let socket: OneWayWebSocket | undefined; - try { - terminalSession = new TerminalSession("Agent Log"); - const { socket: agentSocket, completion: logsCompletion } = - await streamAgentLogs( - workspaceClient, - terminalSession.writeEmitter, - currentAgent, - ); - socket = agentSocket; - - const agentStatePromise = this.waitForAgentWithProgress( - monitor, - currentAgent, - logMsg, - (foundAgent) => foundAgent.lifecycle_state !== "starting", - ); - - // Race between logs completion and agent state change - currentAgent = await Promise.race([ - agentStatePromise, - logsCompletion.then(() => agentStatePromise), - ]); - this.logger.info( - `Agent ${currentAgent.name} lifecycle state is now`, - currentAgent.lifecycle_state, - ); - } finally { - terminalSession?.dispose(); - socket?.close(); - } - continue; - } - - case "created": - this.logger.info( - `Waiting for ${workspaceName}/${currentAgent.name} to start...`, - ); - currentAgent = await this.waitForAgentWithProgress( - monitor, - currentAgent, - `Waiting for agent ${currentAgent.name} to start...`, - (foundAgent) => foundAgent.lifecycle_state !== "created", - ); - this.logger.info( - `Agent ${currentAgent.name} lifecycle state is now`, - currentAgent.lifecycle_state, - ); - continue; - - case "off": - case "start_error": - case "start_timeout": - case "shutdown_error": - case "shutdown_timeout": - case "shutting_down": - throw new Error( - `${workspaceName}/${currentAgent.name} lifecycle state: ${currentAgent.lifecycle_state}`, - ); - - default: - throw new Error( - `${workspaceName}/${currentAgent.name} unknown lifecycle state: ${currentAgent.lifecycle_state}`, - ); - } - } - return currentAgent; - } - - /** - * Waits for an agent condition with progress notification. - */ - private async waitForAgentWithProgress( - monitor: WorkspaceMonitor, - agent: WorkspaceAgent, - progressTitle: string, - checker: (agent: WorkspaceAgent) => boolean, - ): Promise { - const foundAgent = await this.vscodeProposed.window.withProgress( - { - title: progressTitle, - location: vscode.ProgressLocation.Notification, - }, - async () => this.waitForAgentCondition(monitor, agent, checker), - ); - return foundAgent; - } - - /** - * Waits for an agent condition to be met by monitoring workspace changes. - * Returns the result when the condition is met or throws an error on timeout. - */ - private async waitForAgentCondition( - monitor: WorkspaceMonitor, - agent: WorkspaceAgent, - checker: (agent: WorkspaceAgent) => boolean, - ): Promise { - return new Promise((resolve, reject) => { - const updateEvent = monitor.onChange.event((workspace) => { - if (workspace.latest_build.status !== "running") { - const workspaceName = createWorkspaceIdentifier(workspace); - reject(new Error(`Workspace ${workspaceName} is not running.`)); - } - try { - const agents = extractAgents(workspace.latest_build.resources); - const foundAgent = agents.find((a) => a.id === agent.id); - if (!foundAgent) { - throw new Error( - `Agent ${agent.name} not found in workspace resources`, - ); - } - if (checker(foundAgent)) { - updateEvent.dispose(); - resolve(foundAgent); - } - } catch (error) { - updateEvent.dispose(); - reject(error); - } - }); - }); - } - // updateSSHConfig updates the SSH configuration with a wildcard that handles // all Coder entries. private async updateSSHConfig( diff --git a/src/remote/workspaceStateMachine.ts b/src/remote/workspaceStateMachine.ts new file mode 100644 index 00000000..0b9e8626 --- /dev/null +++ b/src/remote/workspaceStateMachine.ts @@ -0,0 +1,272 @@ +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": + 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..f21e1c40 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 isReady = false; readonly onChange = new vscode.EventEmitter(); private readonly statusBarItem: vscode.StatusBarItem; @@ -123,6 +124,7 @@ export class WorkspaceMonitor implements vscode.Disposable { } private update(workspace: Workspace) { + this.updateReadyState(workspace); this.updateContext(workspace); this.updateStatusBar(workspace); } @@ -131,7 +133,10 @@ export class WorkspaceMonitor implements vscode.Disposable { this.maybeNotifyOutdated(workspace); this.maybeNotifyAutostop(workspace); this.maybeNotifyDeletion(workspace); - this.maybeNotifyNotRunning(workspace); + if (this.isReady) { + // This instance might be created before the workspace is running + this.maybeNotifyNotRunning(workspace); + } } private maybeNotifyAutostop(workspace: Workspace) { @@ -244,6 +249,12 @@ export class WorkspaceMonitor implements vscode.Disposable { this.logger.error(message); } + private updateReadyState(workspace: Workspace): void { + if (workspace.latest_build.status === "running") { + this.isReady = true; + } + } + private updateContext(workspace: Workspace) { this.contextManager.set("coder.workspace.updatable", workspace.outdated); } From b342a11c6b347405e76cbe3494c6747026c264f2 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 31 Oct 2025 17:43:12 +0300 Subject: [PATCH 9/9] Review comments --- src/remote/remote.ts | 28 +++++++++++----------------- src/remote/workspaceStateMachine.ts | 2 ++ src/workspace/workspaceMonitor.ts | 15 ++++++--------- 3 files changed, 19 insertions(+), 26 deletions(-) diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 3782a9ad..79b9274c 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -291,8 +291,8 @@ export class Remote { this.vscodeProposed, this.contextManager, ); - disposables.push(monitor); disposables.push( + monitor, monitor.onChange.event((w) => (this.commands.workspace = w)), ); @@ -310,7 +310,7 @@ export class Remote { disposables.push(stateMachine); try { - await this.vscodeProposed.window.withProgress( + workspace = await this.vscodeProposed.window.withProgress( { location: vscode.ProgressLocation.Notification, cancellable: false, @@ -320,10 +320,8 @@ export class Remote { let inProgress = false; let pendingWorkspace: Workspace | null = null; - await new Promise((resolve, reject) => { + return new Promise((resolve, reject) => { const processWorkspace = async (w: Workspace) => { - workspace = w; - if (inProgress) { // Process one workspace at a time, keeping only the last pendingWorkspace = w; @@ -340,26 +338,19 @@ export class Remote { ); if (isReady) { subscription.dispose(); - resolve(); + resolve(w); return; } - - if (pendingWorkspace) { - const isReadyAfter = await stateMachine.processWorkspace( - pendingWorkspace, - progress, - ); - if (isReadyAfter) { - subscription.dispose(); - resolve(); - } - } } catch (error) { subscription.dispose(); reject(error); } finally { inProgress = false; } + + if (pendingWorkspace) { + processWorkspace(pendingWorkspace); + } }; processWorkspace(workspace); @@ -373,6 +364,9 @@ export class Remote { stateMachine.dispose(); } + // Mark initial setup as complete so the monitor can start notifying about state changes + monitor.markInitialSetupComplete(); + const agents = extractAgents(workspace.latest_build.resources); const agent = agents.find( (agent) => agent.id === stateMachine.getAgentId(), diff --git a/src/remote/workspaceStateMachine.ts b/src/remote/workspaceStateMachine.ts index 0b9e8626..2251cd19 100644 --- a/src/remote/workspaceStateMachine.ts +++ b/src/remote/workspaceStateMachine.ts @@ -95,6 +95,8 @@ export class WorkspaceStateMachine implements vscode.Disposable { 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}...`); diff --git a/src/workspace/workspaceMonitor.ts b/src/workspace/workspaceMonitor.ts index f21e1c40..d1bdefaa 100644 --- a/src/workspace/workspaceMonitor.ts +++ b/src/workspace/workspaceMonitor.ts @@ -29,7 +29,7 @@ export class WorkspaceMonitor implements vscode.Disposable { private notifiedDeletion = false; private notifiedOutdated = false; private notifiedNotRunning = false; - private isReady = false; + private completedInitialSetup = false; readonly onChange = new vscode.EventEmitter(); private readonly statusBarItem: vscode.StatusBarItem; @@ -111,6 +111,10 @@ export class WorkspaceMonitor implements vscode.Disposable { return monitor; } + public markInitialSetupComplete(): void { + this.completedInitialSetup = true; + } + /** * Permanently close the websocket. */ @@ -124,7 +128,6 @@ export class WorkspaceMonitor implements vscode.Disposable { } private update(workspace: Workspace) { - this.updateReadyState(workspace); this.updateContext(workspace); this.updateStatusBar(workspace); } @@ -133,7 +136,7 @@ export class WorkspaceMonitor implements vscode.Disposable { this.maybeNotifyOutdated(workspace); this.maybeNotifyAutostop(workspace); this.maybeNotifyDeletion(workspace); - if (this.isReady) { + if (this.completedInitialSetup) { // This instance might be created before the workspace is running this.maybeNotifyNotRunning(workspace); } @@ -249,12 +252,6 @@ export class WorkspaceMonitor implements vscode.Disposable { this.logger.error(message); } - private updateReadyState(workspace: Workspace): void { - if (workspace.latest_build.status === "running") { - this.isReady = true; - } - } - private updateContext(workspace: Workspace) { this.contextManager.set("coder.workspace.updatable", workspace.outdated); }