From 7437d924538e229dab98e3b6d87ae51db467d196 Mon Sep 17 00:00:00 2001
From: loveucifer <134506987+loveucifer@users.noreply.github.com>
Date: Sat, 3 Jan 2026 17:15:13 +0530
Subject: [PATCH] Fix language services lifecycle management in Workspace

Closes #2209

- Only call `setLanguageServices` from `openDocument` to avoid race conditions
- Remove language services when documents are closed via `removeLanguageServices`
- `SourceKitLSPServer.languageServices` now just returns services without storing them
- Fixed a small typo (serveer -> server) while I was in there

This way the languageServices dictionary only tracks documents that are actually open,
and we avoid race conditions since openDocument is a blocking request for that document.
---
 Sources/SourceKitLSP/SourceKitLSPServer.swift |  5 +++-
 Sources/SourceKitLSP/Workspace.swift          | 26 +++++++++----------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/Sources/SourceKitLSP/SourceKitLSPServer.swift b/Sources/SourceKitLSP/SourceKitLSPServer.swift
index 98aed803..7cda59c5 100644
--- a/Sources/SourceKitLSP/SourceKitLSPServer.swift
+++ b/Sources/SourceKitLSP/SourceKitLSPServer.swift
@@ -592,7 +592,7 @@ package actor SourceKitLSPServer {
       """
     )
 
-    return workspace.setLanguageServices(for: uri, languageServices)
+    return languageServices
   }
 
   /// The language service with the highest precedence that can handle the given document.
@@ -1345,6 +1345,7 @@ extension SourceKitLSPServer {
     let language = textDocument.language
 
     let languageServices = await languageServices(for: uri, language, in: workspace)
+    workspace.setLanguageServices(for: uri, languageServices)
 
     if languageServices.isEmpty {
       // If we can't create a service, this document is unsupported and we can bail here.
@@ -1399,6 +1400,8 @@ extension SourceKitLSPServer {
       await languageService.closeDocument(notification)
     }
 
+    workspace.removeLanguageServices(for: uri)
+
     workspaceQueue.async {
       self.workspaceForUri[notification.textDocument.uri] = nil
     }
diff --git a/Sources/SourceKitLSP/Workspace.swift b/Sources/SourceKitLSP/Workspace.swift
index f3badb11..89641e7b 100644
--- a/Sources/SourceKitLSP/Workspace.swift
+++ b/Sources/SourceKitLSP/Workspace.swift
@@ -424,23 +424,23 @@ package final class Workspace: Sendable, BuildServerManagerDelegate {
     return languageServices(for: uri).first
   }
 
-  /// Set a language service for a document uri and returns if none exists already.
+  /// Set the language services for a document URI.
   ///
-  /// If language services already exist for this document, eg. because two requests start creating a language
-  /// service for a document and race, `newLanguageServices` is dropped and the existing language services for the
-  /// document are returned.
-  func setLanguageServices(for uri: DocumentURI, _ newLanguageService: [any LanguageService]) -> [any LanguageService] {
-    return languageServices.withLock { languageServices in
-      if let languageService = languageServices[uri] {
-        return languageService
-      }
-
-      languageServices[uri] = newLanguageService
-      return newLanguageService
+  /// This should only be called from `openDocument` to ensure there are no race conditions.
+  func setLanguageServices(for uri: DocumentURI, _ newLanguageService: [any LanguageService]) {
+    languageServices.withLock { languageServices in
+      languageServices[uri.buildSettingsFile] = newLanguageService
     }
   }
 
-  /// Handle a build settings change notification from the build serveer.
+  /// Remove the language services association for a document when it is closed.
+  func removeLanguageServices(for uri: DocumentURI) {
+    languageServices.withLock { languageServices in
+      languageServices[uri.buildSettingsFile] = nil
+    }
+  }
+
+  /// Handle a build settings change notification from the build server.
   /// This has two primary cases:
   /// - Initial settings reported for a given file, now we can fully open it
   /// - Changed settings for an already open file