Wait for startup script when launching the workspace

Author: EhabY

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<WorkspaceAgentLog[]>({
+			apiRoute: `/api/v2/workspaceagents/${agentId}/logs`,
+			searchParams,
+			options,
+		});
+	};
+
 	private async createWebSocket<TData = unknown>(
 		configs: Omit<OneWayWebSocketInit, "location">,
 	) {

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<string>,
+	agentId: string,
+): Promise<OneWayWebSocket<WorkspaceAgentLog[]>> {
+	// 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.
  *

src/extension.ts

@@ -353,7 +353,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		}),
 	);
 
-	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<void> {
 			);
 			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<void> {
 			}
 		}
 	}
-
-	if (shouldShowSshOutput) {
-		showSshOutput();
-	}
 }
 
 async function showTreeViewSearch(id: string): Promise<void> {
 	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);
-	}
-}

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<string>;
-		let terminal: undefined | vscode.Terminal;
+		let writeEmitter: vscode.EventEmitter<string> | undefined;
+		let terminal: vscode.Terminal | undefined;
 		let attempts = 0;
 
-		function initWriteEmitterAndTerminal(): vscode.EventEmitter<string> {
-			writeEmitter ??= new vscode.EventEmitter<string>();
-			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<vscode.Pseudoterminal> 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<string>;
+		terminal: vscode.Terminal;
+	} {
+		const writeEmitter = new vscode.EventEmitter<string>();
+		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<void>((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());
 			},