feat(vscode): communicate with worker using ipc instead of websocket

.vscode/launch.json

@@ -15,7 +15,8 @@
         "${workspaceFolder}/packages/vscode/sample"
       ],
       "outFiles": ["${workspaceFolder}/packages/vscode/dist/**/*.js"],
-      "preLaunchTask": "npm: build:local"
+      "preLaunchTask": "npm: build:local",
+      "autoAttachChildProcesses": true
     }
   ]
 }

packages/vscode/package.json

@@ -70,15 +70,13 @@
     "@types/mocha": "^10.0.10",
     "@types/node": "^22.16.5",
     "@types/vscode": "1.97.0",
-    "@types/ws": "^8.18.1",
     "@vscode/test-cli": "^0.0.12",
     "@vscode/test-electron": "^2.5.2",
     "@vscode/vsce": "3.6.2",
-    "get-port": "^7.1.0",
+    "birpc": "2.6.1",
     "glob": "^7.2.3",
     "mocha": "^11.7.4",
     "ovsx": "^0.10.6",
-    "typescript": "^5.9.3",
-    "ws": "^8.18.3"
+    "typescript": "^5.9.3"
   }
 }

packages/vscode/src/logger.ts

@@ -1,40 +1,21 @@
-import { formatWithOptions } from 'node:util';
 import vscode from 'vscode';
+import { BaseLogger, type LogLevel } from './shared/logger';
 
-function formatValues(values: unknown[]): string {
-  return formatWithOptions({ depth: 4 }, ...values);
-}
-
-export class Logger implements vscode.Disposable {
+export class MasterLogger extends BaseLogger implements vscode.Disposable {
   readonly #channel: vscode.LogOutputChannel;
 
   constructor(private readonly name = 'Rstest') {
+    super();
     this.#channel = vscode.window.createOutputChannel(this.name, { log: true });
   }
 
-  public trace(...values: unknown[]) {
-    this.#channel.trace(formatValues(values));
-  }
-
-  public debug(...values: unknown[]) {
-    this.#channel.debug(formatValues(values));
-  }
-
-  public info(...values: unknown[]) {
-    this.#channel.info(formatValues(values));
-  }
-
-  public warn(...values: unknown[]) {
-    this.#channel.warn(formatValues(values));
-  }
-
-  public error(...values: unknown[]) {
-    this.#channel.error(formatValues(values));
+  override log(level: LogLevel, message: string) {
+    this.#channel[level](message);
   }
 
   public dispose() {
     this.#channel.dispose();
   }
 }
 
-export const logger = new Logger();
+export const logger = new MasterLogger();

packages/vscode/src/master.ts

@@ -1,24 +1,16 @@
 import { spawn } from 'node:child_process';
-import { createServer } from 'node:http';
 import path, { dirname } from 'node:path';
-import getPort from 'get-port';
+import { createBirpc } from 'birpc';
 import vscode from 'vscode';
-import type { WebSocket } from 'ws';
-import { WebSocketServer } from 'ws';
 import { getConfigValue } from './config';
 import { logger } from './logger';
-import type {
-  WorkerEvent,
-  WorkerEventFinish,
-  WorkerRunTestData,
-} from './types';
+import type { LogLevel } from './shared/logger';
+import type { WorkerRunTestData } from './types';
+import { promiseWithTimeout } from './utils';
+import type { Worker } from './worker';
 
 export class RstestApi {
-  public ws: WebSocket | null = null;
-  private testPromises: Map<
-    string,
-    { resolve: (value: any) => void; reject: (reason?: any) => void }
-  > = new Map();
+  public worker: Pick<Worker, 'initRstest' | 'runTest'> | null = null;
   private versionMismatchWarned = false;
 
   public resolveRstestPath(): { cwd: string; rstestPath: string }[] {
@@ -116,140 +108,87 @@ export class RstestApi {
   }
 
   public async runTest(item: vscode.TestItem) {
-    if (this.ws) {
+    if (this.worker) {
       const data: WorkerRunTestData = {
-        type: 'runTest',
         id: item.id,
         fileFilters: [item.uri!.fsPath],
         testNamePattern: item.label,
       };
 
-      // Create a promise that will be resolved when we get a response with the matching ID
-      const promise = new Promise<any>((resolve, reject) => {
-        this.testPromises.set(item.id, { resolve, reject });
-
-        // Set a timeout to prevent hanging indefinitely
-        setTimeout(() => {
-          const promiseObj = this.testPromises.get(item.id);
-          if (promiseObj) {
-            this.testPromises.delete(item.id);
-            reject(new Error(`Test execution timed out for ${item.label}`));
-          }
-        }, 10000); // 10 seconds timeout
-      });
-
-      this.ws.send(JSON.stringify(data));
-      return promise;
+      return promiseWithTimeout(
+        this.worker.runTest(data),
+        10_000,
+        new Error(`Test execution timed out for ${item.label}`),
+      ); // 10 seconds timeout
     }
   }
 
   public async runFileTests(fileItem: vscode.TestItem) {
-    if (this.ws) {
+    if (this.worker) {
       const fileId = `file_${fileItem.id}`;
       const data: WorkerRunTestData = {
-        type: 'runTest',
         id: fileId,
         fileFilters: [fileItem.uri!.fsPath],
         testNamePattern: '', // Empty pattern to run all tests in the file
       };
 
-      // Create a promise that will be resolved when we get a response with the matching ID
-      const promise = new Promise<WorkerEventFinish>((resolve, reject) => {
-        this.testPromises.set(fileId, { resolve, reject });
-
-        // Set a timeout to prevent hanging indefinitely
-        setTimeout(() => {
-          const promiseObj = this.testPromises.get(fileId);
-          if (promiseObj) {
-            this.testPromises.delete(fileId);
-            reject(
-              new Error(
-                `File test execution timed out for ${fileItem.uri!.fsPath}`,
-              ),
-            );
-          }
-        }, 30000); // 30 seconds timeout for file-level tests
-      });
-
-      this.ws.send(JSON.stringify(data));
-      return promise;
+      return promiseWithTimeout(
+        this.worker.runTest(data),
+        30_000,
+        new Error(`File test execution timed out for ${fileItem.uri!.fsPath}`),
+      ); // 30 seconds timeout for file-level tests
     }
   }
 
   public async createChildProcess() {
+    const { cwd, rstestPath } = this.resolveRstestPath()[0];
+    if (!cwd || !rstestPath) {
+      logger.error('Failed to resolve rstest path or cwd');
+      return;
+    }
+
     const execArgv: string[] = [];
     const workerPath = path.resolve(__dirname, 'worker.js');
-    const port = await getPort();
-    const wsAddress = `ws://localhost:${port}`;
     logger.debug('Spawning worker process', {
       workerPath,
-      wsAddress,
     });
     const rstestProcess = spawn('node', [...execArgv, workerPath], {
-      stdio: 'pipe',
+      cwd,
+      stdio: ['pipe', 'pipe', 'pipe', 'ipc'],
+      serialization: 'advanced',
       env: {
         ...process.env,
         TEST: 'true',
-        RSTEST_WS_ADDRESS: wsAddress,
       },
     });
 
     rstestProcess.stdout?.on('data', (d) => {
       const content = d.toString();
-      logger.debug('worker stdout', content.trimEnd());
+      logger.debug('[worker stdout]', content.trimEnd());
     });
 
     rstestProcess.stderr?.on('data', (d) => {
       const content = d.toString();
-      logger.error('worker stderr', content.trimEnd());
+      logger.error('[worker stderr]', content.trimEnd());
     });
 
-    const server = createServer().listen(port).unref();
-    const wss = new WebSocketServer({ server });
-
-    wss.once('connection', (ws) => {
-      this.ws = ws;
-      logger.debug('Worker connected', { wsAddress });
-      const { cwd, rstestPath } = this.resolveRstestPath()[0];
-      if (!cwd || !rstestPath) {
-        logger.error('Failed to resolve rstest path or cwd');
-        return;
-      }
-
-      ws.send(
-        JSON.stringify({
-          type: 'init',
-          rstestPath,
-          cwd,
-        }),
-      );
-      logger.debug('Sent init payload to worker', { cwd, rstestPath });
-
-      ws.on('message', (_data) => {
-        const _message = JSON.parse(_data.toString()) as WorkerEvent;
-        if (_message.type === 'finish') {
-          const message: WorkerEventFinish = _message;
-          logger.debug('Received worker completion event', {
-            id: message.id,
-            testResult: message.testResults,
-            testFileResult: message.testFileResults,
-          });
-          // Check if we have a pending promise for this test ID
-          const promiseObj = this.testPromises.get(message.id);
-          if (promiseObj) {
-            // Resolve the promise with the message data
-            promiseObj.resolve(message);
-            // Remove the promise from the map
-            this.testPromises.delete(message.id);
-          }
-        }
-      });
+    this.worker = createBirpc<Worker, RstestApi>(this, {
+      post: (data) => rstestProcess.send(data),
+      on: (fn) => rstestProcess.on('message', fn),
+      bind: 'functions',
     });
 
+    await this.worker.initRstest({ cwd, rstestPath });
+    logger.debug('Sent init payload to worker', { cwd, rstestPath });
+
     rstestProcess.on('exit', (code, signal) => {
       logger.debug('Worker process exited', { code, signal });
     });
   }
 
   public async createRstestWorker() {}
+
+  async log(level: LogLevel, message: string) {
+    logger[level](message);
+  }
 }

packages/vscode/src/shared/logger.ts

@@ -0,0 +1,39 @@
+import { formatWithOptions } from 'node:util';
+
+export type LogLevel = 'trace' | 'debug' | 'info' | 'warn' | 'error';
+
+export type Logger = {
+  [K in LogLevel]: (...params: unknown[]) => void;
+};
+
+export abstract class BaseLogger implements Logger {
+  constructor(private prefix?: string) {}
+  protected log(level: LogLevel, message: string): void {
+    console[level](message);
+  }
+  private logWithFormat(level: LogLevel, params: unknown[]) {
+    this.log(
+      level,
+      formatWithOptions(
+        { depth: 4 },
+        ...(this.prefix ? [`[${this.prefix}]`] : []),
+        ...params,
+      ),
+    );
+  }
+  trace(...params: unknown[]) {
+    this.logWithFormat('trace', params);
+  }
+  debug(...params: unknown[]) {
+    this.logWithFormat('debug', params);
+  }
+  info(...params: unknown[]) {
+    this.logWithFormat('info', params);
+  }
+  warn(...params: unknown[]) {
+    this.logWithFormat('warn', params);
+  }
+  error(...params: unknown[]) {
+    this.logWithFormat('error', params);
+  }
+}

packages/vscode/src/types.ts

@@ -2,25 +2,19 @@ import type { TestFileResult, TestResult } from '@rstest/core';
 
 //#region master -> worker
 export type WorkerInitData = {
-  type: 'init';
   rstestPath: string;
   cwd: string;
 };
 
 export type WorkerRunTestData = {
   id: string;
-  type: 'runTest';
   fileFilters: string[];
   testNamePattern: string;
 };
 // #endregion
 
 //#region worker -> master
-export type WorkerEvent = WorkerEventFinish;
-
 export type WorkerEventFinish = {
-  type: 'finish';
-  id: string;
   testResults: TestResult[];
   testFileResults?: TestFileResult[];
 };

packages/vscode/src/utils.ts

@@ -33,3 +33,14 @@ export function getWorkspaceTestPatterns(): WorkspaceTestPattern[] {
     }));
   });
 }
+
+export function promiseWithTimeout<T>(
+  promise: Promise<T>,
+  timeout: number,
+  error: Error,
+) {
+  return Promise.race<T>([
+    promise,
+    new Promise((_, reject) => setTimeout(reject, timeout, error)),
+  ]);
+}