From e1fb04377d3792c587f5f1ecc92349b0cc229759 Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Mon, 3 Jun 2024 14:19:57 -0400
Subject: [PATCH] Handle unsaved files in main language server client

---
 vscode/src/client.ts    | 15 ++++++++++++++-
 vscode/src/rubyLsp.ts   |  1 +
 vscode/src/workspace.ts |  4 ++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/vscode/src/client.ts b/vscode/src/client.ts
index 4cd8a73f41..e37e1a068b 100644
--- a/vscode/src/client.ts
+++ b/vscode/src/client.ts
@@ -12,6 +12,7 @@ import {
   ExecutableOptions,
   ServerOptions,
   MessageSignature,
+  DocumentSelector,
 } from "vscode-languageclient/node";
 
 import { LSP_NAME, ClientInterface } from "./common";
@@ -92,6 +93,7 @@ function collectClientOptions(
   workspaceFolder: vscode.WorkspaceFolder,
   outputChannel: WorkspaceChannel,
   ruby: Ruby,
+  mainWorkspace: boolean,
 ): LanguageClientOptions {
   const pullOn: "change" | "save" | "both" =
     configuration.get("pullDiagnosticsOn")!;
@@ -105,13 +107,22 @@ function collectClientOptions(
   const enabledFeatures = Object.keys(features).filter((key) => features[key]);
 
   const fsPath = workspaceFolder.uri.fsPath.replace(/\/$/, "");
-  const documentSelector = [
+  const documentSelector: DocumentSelector = [
     {
       language: "ruby",
       pattern: `${fsPath}/**/*`,
     },
   ];
 
+  // Only the first language server we spawn should handle unsaved files, otherwise requests will be duplicated across
+  // all workspaces
+  if (mainWorkspace) {
+    documentSelector.push({
+      language: "ruby",
+      scheme: "untitled",
+    });
+  }
+
   if (ruby.env.GEM_PATH) {
     const parts = ruby.env.GEM_PATH.split(path.delimiter);
 
@@ -161,6 +172,7 @@ export default class Client extends LanguageClient implements ClientInterface {
     createTestItems: (response: CodeLens[]) => void,
     workspaceFolder: vscode.WorkspaceFolder,
     outputChannel: WorkspaceChannel,
+    mainWorkspace = false,
   ) {
     super(
       LSP_NAME,
@@ -170,6 +182,7 @@ export default class Client extends LanguageClient implements ClientInterface {
         workspaceFolder,
         outputChannel,
         ruby,
+        mainWorkspace,
       ),
     );
 
diff --git a/vscode/src/rubyLsp.ts b/vscode/src/rubyLsp.ts
index d64fa5f5e8..38e3a9cc2d 100644
--- a/vscode/src/rubyLsp.ts
+++ b/vscode/src/rubyLsp.ts
@@ -161,6 +161,7 @@ export class RubyLsp {
       workspaceFolder,
       this.telemetry,
       this.testController.createTestItems.bind(this.testController),
+      eager,
     );
     this.workspaces.set(workspaceFolder.uri.toString(), workspace);
 
diff --git a/vscode/src/workspace.ts b/vscode/src/workspace.ts
index 458d351a09..1adb7928f5 100644
--- a/vscode/src/workspace.ts
+++ b/vscode/src/workspace.ts
@@ -21,6 +21,7 @@ export class Workspace implements WorkspaceInterface {
   private readonly context: vscode.ExtensionContext;
   private readonly telemetry: Telemetry;
   private readonly outputChannel: WorkspaceChannel;
+  private readonly mainWorkspace: boolean;
   private needsRestart = false;
   #rebaseInProgress = false;
   #error = false;
@@ -30,6 +31,7 @@ export class Workspace implements WorkspaceInterface {
     workspaceFolder: vscode.WorkspaceFolder,
     telemetry: Telemetry,
     createTestItems: (response: CodeLens[]) => void,
+    mainWorkspace = false,
   ) {
     this.context = context;
     this.workspaceFolder = workspaceFolder;
@@ -40,6 +42,7 @@ export class Workspace implements WorkspaceInterface {
     this.telemetry = telemetry;
     this.ruby = new Ruby(context, workspaceFolder, this.outputChannel);
     this.createTestItems = createTestItems;
+    this.mainWorkspace = mainWorkspace;
 
     this.registerRestarts(context);
     this.registerRebaseWatcher(context);
@@ -102,6 +105,7 @@ export class Workspace implements WorkspaceInterface {
       this.createTestItems,
       this.workspaceFolder,
       this.outputChannel,
+      this.mainWorkspace,
     );
 
     try {