Add profile code lens

lib/ruby_lsp/server.rb

@@ -1520,6 +1520,9 @@ def code_lens_resolve(message)
           Interface::Command.new(title: "▶ Run in terminal", command: "rubyLsp.runTestInTerminal", arguments: args)
       when "debug_test"
         code_lens[:command] = Interface::Command.new(title: "⚙ Debug", command: "rubyLsp.debugTest", arguments: args)
+      when "profile_test"
+        code_lens[:command] =
+          Interface::Command.new(title: "⏱ Profile", command: "rubyLsp.profileTest", arguments: args)
       end
 
       send_message(Result.new(

vscode/package.json

@@ -191,6 +191,11 @@
         "command": "rubyLsp.showOutput",
         "title": "Show output channel",
         "category": "Ruby LSP"
+      },
+      {
+        "command": "rubyLsp.profileTest",
+        "title": "Profile current test file",
+        "category": "Ruby LSP"
       }
     ],
     "configuration": {

vscode/src/common.ts

@@ -19,6 +19,7 @@ export enum Command {
   RunTest = "rubyLsp.runTest",
   RunTestInTerminal = "rubyLsp.runTestInTerminal",
   DebugTest = "rubyLsp.debugTest",
+  ProfileTest = "rubyLsp.profileTest",
   ShowSyntaxTree = "rubyLsp.showSyntaxTree",
   DiagnoseState = "rubyLsp.diagnoseState",
   DisplayAddons = "rubyLsp.displayAddons",

vscode/src/rubyLsp.ts

@@ -477,6 +477,12 @@
             : this.testController.debugTest(path, name, command);
         },
       ),
+      vscode.commands.registerCommand(
+        Command.ProfileTest,
+        (path, name, _command) => {
+          this.testController.runViaCommand(path, name, Mode.Profile)
+        },
+      ),
       vscode.commands.registerCommand(
         Command.RunTask,
         async (command: string) => {

vscode/src/streamingRunner.ts

@@ -29,6 +29,7 @@ export enum Mode {
   Run = "run",
   RunInTerminal = "runInTerminal",
   Debug = "debug",
+  Profile = "profile",
 }
 
 // The StreamingRunner class is responsible for executing the test process or launching the debugger while handling the

vscode/src/testController.ts

@@ -1,5 +1,6 @@
 import { exec } from "child_process";
 import { promisify } from "util";
+import * as os from "os";
 import path from "path";
 
 import * as vscode from "vscode";
@@ -116,12 +117,12 @@
       this.fullDiscovery
         ? this.runTest.bind(this)
         : async () => {
-            await vscode.window.showInformationMessage(
-              `Running tests with coverage requires the new explorer implementation,
+          await vscode.window.showInformationMessage(
+            `Running tests with coverage requires the new explorer implementation,
                which is currently under development.
                If you wish to enable it, set the "fullTestDiscovery" feature flag to "true"`,
-            );
-          },
+          );
+        },
       false,
     );
 
@@ -381,8 +382,8 @@
   }
 
   // Method to run tests in any profile through code lens buttons
-  async runViaCommand(path: string, name: string, mode: Mode) {
-    const uri = vscode.Uri.file(path);
+  async runViaCommand(filePath: string, name: string, mode: Mode) {
+    const uri = vscode.Uri.file(filePath);
     const testItem = await this.findTestItem(name, uri);
     if (!testItem) return;
 
@@ -401,6 +402,93 @@
 
     let profile;
 
+    if (mode === Mode.Profile) {
+      if (this.terminal === undefined) {
+        this.terminal = this.getTerminal();
+      }
+
+      let commandToExecute: string | undefined;
+      let workspace: Workspace | undefined;
+
+      if (this.fullDiscovery) {
+        const workspaceFolder = vscode.workspace.getWorkspaceFolder(
+          testItem.uri!,
+        );
+        if (!workspaceFolder) {
+          vscode.window.showErrorMessage(
+            "Could not find workspace for the test item.",
+          );
+          return;
+        }
+        workspace = await this.getOrActivateWorkspace(workspaceFolder);
+
+        if (
+          workspace.lspClient &&
+          workspace.lspClient.initializeResult?.capabilities.experimental
+            ?.full_test_discovery
+        ) {
+          const requestTestItems = [this.testItemToServerItem(testItem)];
+          try {
+            const response =
+              await workspace.lspClient.resolveTestCommands(requestTestItems);
+
+            if (response && response.commands && response.commands.length > 0) {
+              commandToExecute = response.commands[0];
+            } else {
+              vscode.window.showErrorMessage(
+                "LSP server did not return a command for profiling.",
+              );
+              return;
+            }
+          } catch (error: any) {
+            vscode.window.showErrorMessage(
+              `Error resolving test command for profiling: ${error.message}`,
+            );
+            this.currentWorkspace()?.outputChannel.error(
+              `Error resolving test command for profiling: ${error.message}`,
+            );
+            return;
+          }
+        } else {
+          vscode.window.showErrorMessage(
+            "Cannot profile test: Ruby LSP server does not support necessary test discovery " +
+            "features or is not ready. Please update the Ruby LSP gem or check server status.",
+          );
+          return;
+        }
+      } else {
+        // Old discovery path
+        commandToExecute = this.testCommands.get(testItem);
+      }
+
+      if (!commandToExecute) {
+        vscode.window.showErrorMessage("No test command found to profile.");
+        return;
+      }
+
+      await vscode.window.withProgress(
+        {
+          location: vscode.ProgressLocation.Notification,
+          title: "Profiling in progress...",
+          cancellable: false,
+        },
+        async () => {
+          const profileUri = vscode.Uri.file(
+            path.join(os.tmpdir(), `profile-${Date.now()}.cpuprofile`),
+          );
+
+          await workspace!.execute(
+            `vernier run --output ${profileUri.fsPath} --format cpuprofile -- ${commandToExecute}`,
+          );
+
+          await vscode.commands.executeCommand("vscode.open", profileUri, {
+            viewColumn: vscode.ViewColumn.Beside,
+          });
+        },
+      );
+      return;
+    }
+
     switch (mode) {
       case Mode.Debug:
         profile = this.testDebugProfile;
@@ -636,6 +724,14 @@
 
     run.end();
     linkedCancellationSource.dispose();
+
+    this.telemetry.logUsage("ruby_lsp.code_lens", {
+      type: "counter",
+      attributes: {
+        label: "test",
+        vscodemachineid: vscode.env.machineId,
+      },
+    });
   }
 
   // When trying to a test file or directory, we need to know which framework is used by tests inside of it to resolve
@@ -760,6 +856,17 @@
         linkedCancellationSource,
       );
     }
+
+    run.end();
+    linkedCancellationSource.dispose();
+
+    this.telemetry.logUsage("ruby_lsp.code_lens", {
+      type: "counter",
+      attributes: {
+        label: "debug",
+        vscodemachineid: vscode.env.machineId,
+      },
+    });
   }
 
   private findTestInGroup(
@@ -836,8 +943,8 @@
     return previousTerminal
       ? previousTerminal
       : vscode.window.createTerminal({
-          name,
-        });
+        name,
+      });
   }
 
   private async debugHandler(
@@ -854,7 +961,10 @@
 
     this.telemetry.logUsage("ruby_lsp.code_lens", {
       type: "counter",
-      attributes: { label: "debug", vscodemachineid: vscode.env.machineId },
+      attributes: {
+        label: "debug",
+        vscodemachineid: vscode.env.machineId,
+      },
     });
   }
 
@@ -952,7 +1062,10 @@
 
     this.telemetry.logUsage("ruby_lsp.code_lens", {
       type: "counter",
-      attributes: { label: "test", vscodemachineid: vscode.env.machineId },
+      attributes: {
+        label: "test",
+        vscodemachineid: vscode.env.machineId,
+      },
     });
   }