Write TCP server port to a tempfile instead of using env var

lib/ruby_lsp/test_reporters/lsp_reporter.rb

@@ -4,6 +4,7 @@
 require "json"
 require "socket"
 require "singleton"
+require "tmpdir"
 
 module RubyLsp
   class LspReporter
@@ -14,11 +15,24 @@ class LspReporter
 
     #: -> void
     def initialize
+      dir_path = File.join(Dir.tmpdir, "ruby-lsp")
+      FileUtils.mkdir_p(dir_path)
+
+      port_path = File.join(dir_path, "test_reporter_port")
       port = ENV["RUBY_LSP_REPORTER_PORT"]
-      @io = if port
-        TCPSocket.new("localhost", port)
-      else
-        # For tests that don't spawn the TCP server
+
+      @io = begin
+        # The environment variable is only used for tests. The extension always writes to the temporary file
+        if port
+          TCPSocket.new("localhost", port)
+        elsif File.exist?(port_path)
+          TCPSocket.new("localhost", File.read(port_path))
+        else
+          # For tests that don't spawn the TCP server
+          require "stringio"
+          StringIO.new
+        end
+      rescue
         require "stringio"
         StringIO.new
       end #: IO | StringIO
@@ -168,7 +182,7 @@ def send_message(method_name, **params)
     File.write(File.join(".ruby-lsp", "coverage_result.json"), coverage_results.to_json)
     RubyLsp::LspReporter.instance.internal_shutdown
   end
-elsif ENV["RUBY_LSP_TEST_RUNNER"] && !ENV["RUBY_LSP_ENV"] == "test"
+elsif ENV["RUBY_LSP_TEST_RUNNER"] && ENV["RUBY_LSP_ENV"] != "test"
   at_exit do
     # If the test process crashed immediately without finishing the tests, we still need to tell the extension that the
     # execution ended so that it can clean up

vscode/src/rubyLsp.ts

@@ -161,6 +161,7 @@ export class RubyLsp {
     }
 
     STATUS_EMITTER.fire(this.currentActiveWorkspace());
+    await this.testController.activate();
   }
 
   // Deactivate the extension, which should stop all language servers. Notice that this just stops anything that is

vscode/src/streamingRunner.ts

@@ -1,5 +1,7 @@
 import { spawn } from "child_process";
 import net from "net";
+import os from "os";
+import path from "path";
 
 import * as rpc from "vscode-jsonrpc/node";
 import * as vscode from "vscode";
@@ -32,15 +34,21 @@ export enum Mode {
 // The StreamingRunner class is responsible for executing the test process or launching the debugger while handling the
 // streaming events to update the test explorer status
 export class StreamingRunner implements vscode.Disposable {
+  tcpPort: string | undefined;
   private promises: Promise<void>[] = [];
   private disposables: vscode.Disposable[] = [];
   private readonly findTestItem: (
     id: string,
     uri: vscode.Uri,
   ) => Promise<vscode.TestItem | undefined>;
 
-  private readonly tcpServer: net.Server;
-  private tcpPort: string | undefined;
+  private readonly createTestRun: (
+    request: vscode.TestRunRequest,
+    name?: string,
+    persist?: boolean,
+  ) => vscode.TestRun;
+
+  private tcpServer: net.Server | undefined;
   private connection: rpc.MessageConnection | undefined;
   private executionPromise:
     | { resolve: () => void; reject: (error: Error) => void }
@@ -56,9 +64,14 @@ export class StreamingRunner implements vscode.Disposable {
       id: string,
       uri: vscode.Uri,
     ) => Promise<vscode.TestItem | undefined>,
+    createTestRun: (
+      request: vscode.TestRunRequest,
+      name?: string,
+      persist?: boolean,
+    ) => vscode.TestRun,
   ) {
     this.findTestItem = findTestItem;
-    this.tcpServer = this.startServer();
+    this.createTestRun = createTestRun;
 
     context.subscriptions.push(
       vscode.window.onDidCloseTerminal((terminal) => {
@@ -67,6 +80,10 @@ export class StreamingRunner implements vscode.Disposable {
     );
   }
 
+  async activate() {
+    this.tcpServer = await this.startServer();
+  }
+
   async execute(
     currentRun: vscode.TestRun,
     command: string,
@@ -105,7 +122,7 @@ export class StreamingRunner implements vscode.Disposable {
   }
 
   dispose() {
-    this.tcpServer.close();
+    this.tcpServer?.close();
     this.connection?.dispose();
   }
 
@@ -122,11 +139,7 @@ export class StreamingRunner implements vscode.Disposable {
         name: "Debug",
         request: "launch",
         program: command,
-        env: {
-          ...env,
-          DISABLE_SPRING: "1",
-          RUBY_LSP_REPORTER_PORT: this.tcpPort,
-        },
+        env: { ...env, DISABLE_SPRING: "1" },
       },
       { testRun: this.run },
     );
@@ -164,18 +177,11 @@ export class StreamingRunner implements vscode.Disposable {
       });
     }
 
-    // Set the TCP port information every time even if there's an existing terminal. The user can close the editor
-    // window or reload extensions, which will assign a new port but maintain the same terminal.
-    //
-    // We also send RUBYOPT since that hooks up the custom LSP test reporters and the user's shell may override it
+    // We need to send RUBYOPT since that hooks up the custom LSP test reporters and the user's shell may override it
     if (process.platform === "win32") {
-      terminal.sendText(
-        `$env:RUBY_LSP_REPORTER_PORT="${this.tcpPort}"; $env:RUBYOPT="${env.RUBYOPT}"; Clear-Host`,
-      );
+      terminal.sendText(`$env:RUBYOPT="${env.RUBYOPT}"; Clear-Host`);
     } else {
-      terminal.sendText(
-        `export RUBY_LSP_REPORTER_PORT="${this.tcpPort}"; export RUBYOPT="${env.RUBYOPT}"; clear`,
-      );
+      terminal.sendText(`export RUBYOPT="${env.RUBYOPT}"; clear`);
     }
 
     this.terminals.set(name, terminal);
@@ -193,7 +199,7 @@ export class StreamingRunner implements vscode.Disposable {
   ) {
     const promise = new Promise<void>((resolve, _reject) => {
       const testProcess = spawn(command, {
-        env: { ...env, RUBY_LSP_REPORTER_PORT: this.tcpPort },
+        env,
         stdio: ["pipe", "pipe", "pipe"],
         shell: true,
         signal: abortController.signal,
@@ -216,36 +222,54 @@ export class StreamingRunner implements vscode.Disposable {
     this.promises.push(promise);
   }
 
-  private startServer() {
-    const server = net.createServer();
-    server.on("error", (error) => {
-      throw error;
-    });
-    server.unref();
+  private async startServer(): Promise<net.Server> {
+    // Listening on the TCP connection is asynchronous. We can only resolve the promise once we know what port has been
+    // assigned, otherwise we risk trying to start tests without a port
+    return new Promise((resolve, reject) => {
+      const server = net.createServer();
+      server.on("error", reject);
+      server.unref();
+
+      server.listen(0, "localhost", async () => {
+        const address = server.address();
+
+        if (!address) {
+          throw new Error("Failed setup TCP server for streaming updates");
+        }
+        this.tcpPort =
+          typeof address === "string" ? address : address.port.toString();
+
+        const tempDirUri = vscode.Uri.file(path.join(os.tmpdir(), "ruby-lsp"));
+
+        await vscode.workspace.fs.createDirectory(tempDirUri);
+        await vscode.workspace.fs.writeFile(
+          vscode.Uri.joinPath(tempDirUri, "test_reporter_port"),
+          Buffer.from(this.tcpPort!.toString()),
+        );
 
-    server.listen(0, "localhost", () => {
-      const address = server.address();
+        // On any new connection to the TCP server, attach the JSON RPC reader and the events we defined
+        server.on("connection", (socket) => {
+          this.connection = rpc.createMessageConnection(
+            new rpc.StreamMessageReader(socket),
+            new rpc.StreamMessageWriter(socket),
+          );
 
-      if (!address) {
-        throw new Error("Failed setup TCP server for streaming updates");
-      }
-      this.tcpPort =
-        typeof address === "string" ? address : address.port.toString();
-
-      // On any new connection to the TCP server, attach the JSON RPC reader and the events we defined
-      server.on("connection", (socket) => {
-        this.connection = rpc.createMessageConnection(
-          new rpc.StreamMessageReader(socket),
-          new rpc.StreamMessageWriter(socket),
-        );
+          // Register and start listening for events
+          this.registerStreamingEvents();
+
+          if (!this.run) {
+            this.run = this.createTestRun(
+              new vscode.TestRunRequest(),
+              "on_demand_run_in_terminal",
+            );
+          }
+
+          this.connection.listen();
+        });
 
-        // Register and start listening for events
-        this.registerStreamingEvents();
-        this.connection.listen();
+        resolve(server);
       });
     });
-
-    return server;
   }
 
   private async finalize(cancellation: boolean) {
@@ -272,6 +296,11 @@ export class StreamingRunner implements vscode.Disposable {
       this.connection.dispose();
     }
 
+    if (this.run!.name === "on_demand_run_in_terminal") {
+      this.run!.end();
+    }
+    this.run = undefined;
+
     this.executionPromise!.resolve();
   }
 

vscode/src/test/suite/testController.test.ts

@@ -32,7 +32,7 @@ suite("TestController", () => {
   const serverTestUri = vscode.Uri.joinPath(testDirUri, "server_test.rb");
   const storeTestUri = vscode.Uri.joinPath(testDirUri, "store_test.rb");
 
-  beforeEach(() => {
+  beforeEach(async () => {
     sandbox = sinon.createSandbox();
     workspaceStubs = [];
 
@@ -56,6 +56,8 @@ suite("TestController", () => {
 
     setupLspClientStub(workspace);
     stubWorkspaceOperations(LSP_WORKSPACE_FOLDER);
+
+    await controller.activate();
   });
 
   afterEach(() => {
@@ -594,8 +596,11 @@ suite("TestController", () => {
       "fakeTestServer.js",
     );
 
-    // eslint-disable-next-line no-process-env
-    workspace.ruby.mergeComposedEnvironment(process.env as any);
+    workspace.ruby.mergeComposedEnvironment({
+      // eslint-disable-next-line no-process-env
+      ...process.env,
+      RUBY_LSP_REPORTER_PORT: controller.streamingPort!,
+    });
 
     sandbox.stub(workspace, "lspClient").value({
       resolveTestCommands: sinon.stub().resolves({
@@ -646,6 +651,9 @@ suite("TestController", () => {
     }
 
     await workspace.ruby.activateRuby(manager);
+    workspace.ruby.mergeComposedEnvironment({
+      RUBY_LSP_REPORTER_PORT: controller.streamingPort!,
+    });
 
     const testItem = (await controller.findTestItem(
       "ServerTest::NestedTest#test_something",
@@ -739,8 +747,11 @@ suite("TestController", () => {
       "fakeTestServer.js",
     );
 
-    // eslint-disable-next-line no-process-env
-    workspace.ruby.mergeComposedEnvironment(process.env as any);
+    workspace.ruby.mergeComposedEnvironment({
+      // eslint-disable-next-line no-process-env
+      ...process.env,
+      RUBY_LSP_REPORTER_PORT: controller.streamingPort!,
+    });
     sandbox.stub(workspace, "lspClient").value({
       resolveTestCommands: sinon.stub().resolves({
         commands: [`node ${fakeServerPath}`],
@@ -825,8 +836,11 @@ suite("TestController", () => {
       "fakeTestServer.js",
     );
 
-    // eslint-disable-next-line no-process-env
-    workspace.ruby.mergeComposedEnvironment(process.env as any);
+    workspace.ruby.mergeComposedEnvironment({
+      // eslint-disable-next-line no-process-env
+      ...process.env,
+      RUBY_LSP_REPORTER_PORT: controller.streamingPort!,
+    });
     sandbox.stub(workspace, "lspClient").value({
       resolveTestCommands: sinon.stub().resolves({
         commands: [`node ${fakeServerPath}`],

vscode/src/testController.ts

@@ -72,13 +72,17 @@ export class TestController {
     ) => Promise<Workspace>,
   ) {
     this.telemetry = telemetry;
-    this.runner = new StreamingRunner(context, this.findTestItem.bind(this));
     this.currentWorkspace = currentWorkspace;
     this.getOrActivateWorkspace = getOrActivateWorkspace;
     this.testController = vscode.tests.createTestController(
       "rubyTests",
       "Ruby Tests",
     );
+    this.runner = new StreamingRunner(
+      context,
+      this.findTestItem.bind(this),
+      this.testController.createTestRun.bind(this.testController),
+    );
 
     if (this.fullDiscovery) {
       this.testController.resolveHandler = this.resolveHandler.bind(this);
@@ -498,6 +502,14 @@ export class TestController {
     return this.findTestInGroup(id, testFileItem);
   }
 
+  async activate() {
+    await this.runner.activate();
+  }
+
+  get streamingPort() {
+    return this.runner.tcpPort;
+  }
+
   private async handleTests(
     request: vscode.TestRunRequest,
     token: vscode.CancellationToken,