Adjust LSP Document Lifetimes

Author: thecoolwinter

CodeEdit/Features/Editor/Views/CodeFileView.swift

@@ -21,8 +21,7 @@ struct CodeFileView: View {
 
     /// Any coordinators passed to the view.
     private var textViewCoordinators: [TextViewCoordinator]
-
-    @State private var highlightProviders: [any HighlightProviding] = []
+    private var highlightProviders: [any HighlightProviding] = []
 
     @AppSettings(\.textEditing.defaultTabWidth)
     var defaultTabWidth
@@ -86,15 +85,15 @@ struct CodeFileView: View {
         self.textViewCoordinators = textViewCoordinators
             + [editorInstance.rangeTranslator]
             + [codeFile.contentCoordinator]
-            + [codeFile.languageServerObjects.textCoordinator].compactMap({ $0 })
+            + [codeFile.languageServerObjects.textCoordinator]
         self.isEditable = isEditable
 
         if let openOptions = codeFile.openOptions {
             codeFile.openOptions = nil
             editorInstance.cursorPositions = openOptions.cursorPositions
         }
 
-        updateHighlightProviders()
+        highlightProviders = [codeFile.languageServerObjects.highlightProvider] + [treeSitterClient]
 
         codeFile
             .contentCoordinator
@@ -186,10 +185,6 @@ struct CodeFileView: View {
         .onChange(of: settingsFont) { newFontSetting in
             font = newFontSetting.current
         }
-        .onReceive(codeFile.$languageServerObjects) { languageServerObjects in
-            // This will not be called in single-file views (for now) but is safe to listen to either way
-            updateHighlightProviders(lspHighlightProvider: languageServerObjects.highlightProvider)
-        }
     }
 
     /// Determines the style of bracket emphasis based on the `bracketEmphasis` setting and the current theme.
@@ -212,12 +207,6 @@ struct CodeFileView: View {
             return .underline(color: color)
         }
     }
-
-    /// Updates the highlight providers array.
-    /// - Parameter lspHighlightProvider: The language server provider, if available.
-    private func updateHighlightProviders(lspHighlightProvider: HighlightProviding? = nil) {
-        highlightProviders = [lspHighlightProvider].compactMap({ $0 }) + [treeSitterClient]
-    }
 }
 
 // This extension is kept here because it should not be used elsewhere in the app and may cause confusion

CodeEdit/Features/LSP/Features/DocumentSync/LSPContentCoordinator.swift

@@ -32,16 +32,22 @@ class LSPContentCoordinator<DocumentType: LanguageServerDocument>: TextViewCoord
     private var task: Task<Void, Never>?
 
     weak var languageServer: LanguageServer<DocumentType>?
-    var documentURI: String
+    var documentURI: String?
 
     /// Initializes a content coordinator, and begins an async stream of updates
-    init(documentURI: String, languageServer: LanguageServer<DocumentType>) {
+    init(documentURI: String? = nil, languageServer: LanguageServer<DocumentType>? = nil) {
         self.documentURI = documentURI
         self.languageServer = languageServer
 
         setUpUpdatesTask()
     }
 
+    func setUp(server: LanguageServer<DocumentType>, document: DocumentType) {
+        languageServer = server
+        documentURI = document.languageServerURI
+    }
+
+
     func setUpUpdatesTask() {
         task?.cancel()
         // Create this stream here so it's always set up when the text view is set up, rather than only once on init.
@@ -76,7 +82,7 @@ class LSPContentCoordinator<DocumentType: LanguageServerDocument>: TextViewCoord
     }
 
     func textView(_ textView: TextView, didReplaceContentsIn range: NSRange, with string: String) {
-        guard let lspRange = editedRange else {
+        guard let lspRange = editedRange, let documentURI else {
             return
         }
         self.editedRange = nil

CodeEdit/Features/LSP/Features/SemanticTokens/SemanticTokenHighlightProvider.swift

@@ -32,8 +32,8 @@ final class SemanticTokenHighlightProvider<
     typealias EditCallback = @MainActor (Result<IndexSet, any Error>) -> Void
     typealias HighlightCallback = @MainActor (Result<[HighlightRange], any Error>) -> Void
 
-    private let tokenMap: SemanticTokenMap
-    private let documentURI: String
+    private var tokenMap: SemanticTokenMap?
+    private var documentURI: String?
     private weak var languageServer: LanguageServer<DocumentType>?
     private weak var textView: TextView?
 
@@ -45,13 +45,23 @@ final class SemanticTokenHighlightProvider<
         textView?.documentRange ?? .zero
     }
 
-    init(tokenMap: SemanticTokenMap, languageServer: LanguageServer<DocumentType>, documentURI: String) {
+    init(
+        tokenMap: SemanticTokenMap? = nil,
+        languageServer: LanguageServer<DocumentType>? = nil,
+        documentURI: String? = nil
+    ) {
         self.tokenMap = tokenMap
         self.languageServer = languageServer
         self.documentURI = documentURI
         self.storage = Storage()
     }
 
+    func setUp(server: LanguageServer<DocumentType>, document: DocumentType) {
+        languageServer = server
+        documentURI = document.languageServerURI
+        tokenMap = server.highlightMap
+    }
+
     // MARK: - Language Server Content Lifecycle
 
     /// Called when the language server finishes sending a document update.
@@ -95,7 +105,8 @@ final class SemanticTokenHighlightProvider<
         textView: TextView,
         lastResultId: String
     ) async throws {
-        guard let response = try await languageServer.requestSemanticTokens(
+        guard let documentURI,
+              let response = try await languageServer.requestSemanticTokens(
             for: documentURI,
             previousResultId: lastResultId
         ) else {
@@ -112,7 +123,7 @@ final class SemanticTokenHighlightProvider<
     /// Requests and applies tokens for an entire document. This does not require a previous response id, and should be
     /// used in place of `requestDeltaTokens` when that's the case.
     private func requestTokens(languageServer: LanguageServer<DocumentType>, textView: TextView) async throws {
-        guard let response = try await languageServer.requestSemanticTokens(for: documentURI) else {
+        guard let documentURI, let response = try await languageServer.requestSemanticTokens(for: documentURI) else {
             return
         }
         await applyEntireResponse(response, callback: lastEditCallback)
@@ -159,7 +170,7 @@ final class SemanticTokenHighlightProvider<
             return
         }
 
-        guard let lspRange = textView.lspRangeFrom(nsRange: range) else {
+        guard let lspRange = textView.lspRangeFrom(nsRange: range), let tokenMap else {
             completion(.failure(HighlightError.lspRangeFailure))
             return
         }

CodeEdit/Features/LSP/LanguageServer/Capabilities/LanguageServer+DocumentSync.swift

@@ -86,6 +86,7 @@ extension LanguageServer {
             switch resolveDocumentSyncKind() {
             case .full:
                 guard let content = await getIsolatedDocumentContent(document) else {
+                    logger.error("Failed to get isolated document content")
                     return
                 }
                 let changeEvent = TextDocumentContentChangeEvent(range: nil, rangeLength: nil, text: content.string)
@@ -107,7 +108,7 @@ extension LanguageServer {
 
             // Let the semantic token provider know about the update.
             // Note for future: If a related LSP object need notifying about document changes, do it here.
-            try await document.languageServerObjects.highlightProvider?.documentDidChange()
+            try await document.languageServerObjects.highlightProvider.documentDidChange()
         } catch {
             logger.warning("closeDocument: Error \(error)")
             throw error
@@ -128,10 +129,7 @@ extension LanguageServer {
 
     @MainActor
     private func updateIsolatedDocument(_ document: DocumentType) {
-        document.languageServerObjects = LanguageServerDocumentObjects(
-            textCoordinator: openFiles.contentCoordinator(for: document),
-            highlightProvider: openFiles.semanticHighlighter(for: document)
-        )
+        document.languageServerObjects.setUp(server: self, document: document)
     }
 
     @MainActor

CodeEdit/Features/LSP/LanguageServer/LanguageServer.swift

@@ -51,7 +51,8 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
         lspInstance: InitializingServer,
         lspPid: pid_t,
         serverCapabilities: ServerCapabilities,
-        rootPath: URL
+        rootPath: URL,
+        logContainer: LanguageServerLogContainer
     ) {
         self.languageId = languageId
         self.binary = binary
@@ -60,7 +61,7 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
         self.serverCapabilities = serverCapabilities
         self.rootPath = rootPath
         self.openFiles = LanguageServerFileMap()
-        self.logContainer = LanguageServerLogContainer(language: languageId)
+        self.logContainer = logContainer
         self.logger = Logger(
             subsystem: Bundle.main.bundleIdentifier ?? "",
             category: "LanguageServer.\(languageId.rawValue)"
@@ -89,9 +90,11 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
             environment: binary.env
         )
 
+        let logContainer = LanguageServerLogContainer(language: languageId)
         let (connection, process) = try makeLocalServerConnection(
             languageId: languageId,
-            executionParams: executionParams
+            executionParams: executionParams,
+            logContainer: logContainer
         )
         let server = InitializingServer(
             server: connection,
@@ -105,7 +108,8 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
             lspInstance: server,
             lspPid: process.processIdentifier,
             serverCapabilities: initializationResponse.capabilities,
-            rootPath: URL(filePath: workspacePath)
+            rootPath: URL(filePath: workspacePath),
+            logContainer: logContainer
         )
     }
 
@@ -118,13 +122,17 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
     /// - Returns: A new connection to the language server.
     static func makeLocalServerConnection(
         languageId: LanguageIdentifier,
-        executionParams: Process.ExecutionParameters
+        executionParams: Process.ExecutionParameters,
+        logContainer: LanguageServerLogContainer
     ) throws -> (connection: JSONRPCServerConnection, process: Process) {
         do {
             let (channel, process) = try DataChannel.localProcessChannel(
                 parameters: executionParams,
-                terminationHandler: {
+                terminationHandler: { [weak logContainer] in
                     logger.debug("Terminated data channel for \(languageId.rawValue)")
+                    logContainer?.appendLog(
+                        LogMessageParams(type: .error, message: "Data Channel Terminated Unexpectedly")
+                    )
                 }
             )
             return (JSONRPCServerConnection(dataChannel: channel), process)

CodeEdit/Features/LSP/LanguageServer/LanguageServerFileMap.swift

@@ -16,8 +16,6 @@ class LanguageServerFileMap<DocumentType: LanguageServerDocument> {
     private struct DocumentObject {
         let uri: String
         var documentVersion: Int
-        var contentCoordinator: LSPContentCoordinator<DocumentType>
-        var semanticHighlighter: HighlightProviderType?
     }
 
     private var trackedDocuments: NSMapTable<NSString, DocumentType>
@@ -32,24 +30,7 @@ class LanguageServerFileMap<DocumentType: LanguageServerDocument> {
     func addDocument(_ document: DocumentType, for server: LanguageServer<DocumentType>) {
         guard let uri = document.languageServerURI else { return }
         trackedDocuments.setObject(document, forKey: uri as NSString)
-        var docData = DocumentObject(
-            uri: uri,
-            documentVersion: 0,
-            contentCoordinator: LSPContentCoordinator(
-                documentURI: uri,
-                languageServer: server
-            ),
-            semanticHighlighter: nil
-        )
-
-        if let tokenMap = server.highlightMap {
-            docData.semanticHighlighter = HighlightProviderType(
-                tokenMap: tokenMap,
-                languageServer: server,
-                documentURI: uri
-            )
-        }
-
+        let docData = DocumentObject(uri: uri, documentVersion: 0)
         trackedDocumentData[uri] = docData
     }
 
@@ -87,22 +68,4 @@ class LanguageServerFileMap<DocumentType: LanguageServerDocument> {
     func documentVersion(for uri: DocumentUri) -> Int? {
         return trackedDocumentData[uri]?.documentVersion
     }
-
-    // MARK: - Content Coordinator
-
-    func contentCoordinator(for document: DocumentType) -> LSPContentCoordinator<DocumentType>? {
-        guard let uri = document.languageServerURI else { return nil }
-        return contentCoordinator(for: uri)
-    }
-
-    func contentCoordinator(for uri: DocumentUri) -> LSPContentCoordinator<DocumentType>? {
-        trackedDocumentData[uri]?.contentCoordinator
-    }
-
-    // MARK: - Semantic Highlighter
-
-    func semanticHighlighter(for document: DocumentType) -> HighlightProviderType? {
-        guard let uri = document.languageServerURI else { return nil }
-        return trackedDocumentData[uri]?.semanticHighlighter
-    }
 }

CodeEdit/Features/LSP/LanguageServerDocument.swift

@@ -10,8 +10,15 @@ import CodeEditLanguages
 
 /// A set of properties a language server sets when a document is registered.
 struct LanguageServerDocumentObjects<DocumentType: LanguageServerDocument> {
-    var textCoordinator: LSPContentCoordinator<DocumentType>?
-    var highlightProvider: SemanticTokenHighlightProvider<SemanticTokenStorage, DocumentType>?
+    var textCoordinator: LSPContentCoordinator<DocumentType> = LSPContentCoordinator()
+    // swiftlint:disable:next line_length
+    var highlightProvider: SemanticTokenHighlightProvider<SemanticTokenStorage, DocumentType> = SemanticTokenHighlightProvider()
+
+    @MainActor
+    func setUp(server: LanguageServer<DocumentType>, document: DocumentType) {
+        textCoordinator.setUp(server: server, document: document)
+        highlightProvider.setUp(server: server, document: document)
+    }
 }
 
 /// A protocol that allows a language server to register objects on a text document.