Refactoring and cleanup of agent startup script waiting

Author: EhabY

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,35 +41,33 @@ 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
 		const cmd = `${escapeCommandArg(binPath)} ${startArgs.join(" ")}`;
 		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<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.
  *
@@ -134,7 +96,9 @@ export async function waitForBuild(
 ): Promise<Workspace> {
 	// 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<string>,
+	agent: WorkspaceAgent,
+): Promise<{
+	socket: OneWayWebSocket<WorkspaceAgentLog[]>;
+	completion: Promise<void>;
+}> {
+	// 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<void>((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 };
+}