Ensure LSPs Exit In Reasonable Time On Quit

Author: thecoolwinter

CodeEdit/AppDelegate.swift

@@ -121,6 +121,8 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
         }
     }
 
+    // MARK: - Should Terminate
+
     /// Defers the application terminate message until we've finished cleanup.
     ///
     /// All paths _must_ call `NSApplication.shared.reply(toApplicationShouldTerminate: true)` as soon as possible.
@@ -255,20 +257,56 @@ final class AppDelegate: NSObject, NSApplicationDelegate, ObservableObject {
 
     /// Terminates running language servers. Used during app termination to ensure resources are freed.
     private func terminateLanguageServers() {
-        Task {
-            await lspService.stopAllServers()
-            await MainActor.run {
-                NSApplication.shared.reply(toApplicationShouldTerminate: true)
+        Task { @MainActor in
+            let task = TaskNotificationModel(
+                id: "appdelegate.terminate_language_servers",
+                title: "Stopping Language Servers",
+                message: "Stopping running language server processes...",
+                isLoading: true
+            )
+
+            if !lspService.languageClients.isEmpty {
+                TaskNotificationHandler.postTask(action: .create, model: task)
             }
+
+            try await withTimeout(
+                duration: .seconds(5.0),
+                onTimeout: {
+                    // Stop-gap measure to ensure we don't hang on CMD-Q
+                    await self.lspService.killAllServers()
+                },
+                operation: {
+                    await self.lspService.stopAllServers()
+                }
+            )
+
+            TaskNotificationHandler.postTask(action: .delete, model: task)
+            NSApplication.shared.reply(toApplicationShouldTerminate: true)
         }
     }
 
     /// Terminates all running tasks. Used during app termination to ensure resources are freed.
     private func terminateTasks() {
-        let documents = CodeEditDocumentController.shared.documents.compactMap({ $0 as? WorkspaceDocument })
-        documents.forEach { workspace in
-            workspace.taskManager?.stopAllTasks()
+        let task = TaskNotificationModel(
+            id: "appdelegate.terminate_tasks",
+            title: "Terminating Tasks",
+            message: "Interrupting all running tasks before quitting...",
+            isLoading: true
+        )
+
+        let taskManagers = CodeEditDocumentController.shared.documents
+            .compactMap({ $0 as? WorkspaceDocument })
+            .compactMap({ $0.taskManager })
+
+        if taskManagers.reduce(0, { $0 + $1.activeTasks.count }) > 0 {
+            TaskNotificationHandler.postTask(action: .create, model: task)
         }
+
+        taskManagers.forEach { manager in
+            manager.stopAllTasks()
+        }
+
+        TaskNotificationHandler.postTask(action: .delete, model: task)
     }
 }
 

CodeEdit/Features/ActivityViewer/Notifications/TaskNotificationHandler.swift

@@ -25,28 +25,35 @@ import Combine
 /// Remember to manage your task notifications appropriately. You should either delete task 
 /// notifications manually or schedule their deletion in advance using the `deleteWithDelay` method.
 ///
+/// Some tasks should be restricted to a specific workspace. To do this, specify the `workspace` attribute in the
+/// notification's `userInfo` dictionary as a `URL`, or use the `toWorkspace` parameter on
+/// ``TaskNotificationHandler/postTask(toWorkspace:action:model:)``.
+///
 /// ## Available Methods
 /// - `create`:
 ///     Creates a new Task Notification. 
 ///     Required fields: `id` (String), `action` (String), `title` (String). 
-///     Optional fields: `message` (String), `percentage` (Double), `isLoading` (Bool).
+///     Optional fields: `message` (String), `percentage` (Double), `isLoading` (Bool), `workspace` (URL).
 /// - `createWithPriority`:
 ///     Creates a new Task Notification and inserts it at the start of the array.
 ///     This ensures it appears in the activity viewer even if there are other task notifications before it.
 ///     **Note:** This should only be used for important notifications!
 ///     Required fields: `id` (String), `action` (String), `title` (String). 
-///     Optional fields: `message` (String), `percentage` (Double), `isLoading` (Bool).
+///     Optional fields: `message` (String), `percentage` (Double), `isLoading` (Bool), `workspace` (URL).
 /// - `update`:
 ///     Updates an existing task notification. It's important to pass the same `id` to update the correct task.
 ///     Required fields: `id` (String), `action` (String). 
-///     Optional fields: `title` (String), `message` (String), `percentage` (Double), `isLoading` (Bool).
+///     Optional fields: `title` (String), `message` (String), `percentage` (Double), `isLoading` (Bool),
+///     `workspace` (URL).
 /// - `delete`:
 ///     Deletes an existing task notification.
 ///     Required fields: `id` (String), `action` (String).
+///     Optional field: `workspace` (URL).
 /// - `deleteWithDelay`:
 ///     Deletes an existing task notification after a certain `TimeInterval`.
 ///     Required fields: `id` (String), `action` (String), `delay` (Double).
-///     **Important:** When specifying the delay, ensure it's a double. 
+///     Optional field: `workspace` (URL).
+///     **Important:** When specifying the delay, ensure it's a double.
 ///     For example, '2' would be invalid because it would count as an integer, use '2.0' instead.
 ///
 /// ## Example Usage:
@@ -101,13 +108,46 @@ import Combine
 /// }
 /// ```
 ///
+/// You can also use the static helper method instead of creating dictionaries manually:
+/// ```swift
+/// TaskNotificationHandler.postTask(action: .create, model: .init(id: "task_id", "title": "New Task"))
+/// ```
+///
 /// - Important: Please refer to ``CodeEdit/TaskNotificationModel`` and ensure you pass the correct values.
 final class TaskNotificationHandler: ObservableObject {
     @Published private(set) var notifications: [TaskNotificationModel] = []
+    var workspaceURL: URL?
     var cancellables: Set<AnyCancellable> = []
 
+    enum Action: String {
+        case create
+        case createWithPriority
+        case update
+        case delete
+        case deleteWithDelay
+    }
+
+    /// Post a new task.
+    /// - Parameters:
+    ///   - toWorkspace: The workspace to restrict the task to. Defaults to `nil`, which is received by all workspaces.
+    ///   - action: The action being taken on the task.
+    ///   - model: The task contents.
+    static func postTask(toWorkspace: URL? = nil, action: Action, model: TaskNotificationModel) {
+        NotificationCenter.default.post(name: .taskNotification, object: nil, userInfo: [
+            "id": model.id,
+            "title": model.title,
+            "message": model.message as Any,
+            "percentage": model.percentage as Any,
+            "isLoading": model.isLoading,
+            "action": action.rawValue,
+            "workspace": toWorkspace as Any
+        ])
+    }
+
     /// Initialises a new `TaskNotificationHandler` and starts observing for task notifications.
-    init() {
+    init(workspaceURL: URL? = nil) {
+        self.workspaceURL = workspaceURL
+
         NotificationCenter.default
             .publisher(for: .taskNotification)
             .receive(on: DispatchQueue.main)
@@ -127,21 +167,25 @@ final class TaskNotificationHandler: ObservableObject {
     private func handleNotification(_ notification: Notification) {
         guard let userInfo = notification.userInfo,
               let taskID = userInfo["id"] as? String,
-              let action = userInfo["action"] as? String else { return }
+              let actionRaw = userInfo["action"] as? String,
+              let action = Action(rawValue: actionRaw) else { return }
+
+        /// If a workspace is specified, don't do anything with this task.
+        if let workspaceURL = userInfo["workspace"] as? URL, workspaceURL != self.workspaceURL {
+            return
+        }
 
         switch action {
-        case "create", "createWithPriority":
+        case .create, .createWithPriority:
             createTask(task: userInfo)
-        case "update":
+        case .update:
             updateTask(task: userInfo)
-        case "delete":
+        case .delete:
             deleteTask(taskID: taskID)
-        case "deleteWithDelay":
+        case .deleteWithDelay:
             if let delay = userInfo["delay"] as? Double {
                 deleteTaskAfterDelay(taskID: taskID, delay: delay)
             }
-        default:
-            break
         }
     }
 

CodeEdit/Features/Documents/WorkspaceDocument/WorkspaceDocument.swift

@@ -161,6 +161,7 @@ final class WorkspaceDocument: NSDocument, ObservableObject, NSToolbarDelegate {
                 workspaceURL: url
             )
         }
+        self.taskNotificationHandler.workspaceURL = url
 
         editorManager?.restoreFromState(self)
         utilityAreaModel?.restoreFromState(self)

CodeEdit/Features/LSP/LanguageServer/DataChannel+localProcess.swift

@@ -0,0 +1,69 @@
+//
+//  DataChannel+localProcess.swift
+//  CodeEdit
+//
+//  Created by Khan Winter on 7/8/25.
+//
+
+import Foundation
+import LanguageServerProtocol
+import JSONRPC
+import ProcessEnv
+
+// This is almost exactly the same as the extension provided by `LanguageServerProtocol`, except it returns the
+// process ID as well as the channel, which we need.
+
+extension DataChannel {
+	@available(macOS 12.0, *)
+	public static func localProcessChannel(
+		parameters: Process.ExecutionParameters,
+		terminationHandler: @escaping @Sendable () -> Void
+    ) throws -> (pid: pid_t, channel: DataChannel) {
+		let process = Process()
+
+		let stdinPipe = Pipe()
+		let stdoutPipe = Pipe()
+		let stderrPipe = Pipe()
+
+		process.standardInput = stdinPipe
+		process.standardOutput = stdoutPipe
+		process.standardError = stderrPipe
+
+		process.parameters = parameters
+
+		let (stream, continuation) = DataSequence.makeStream()
+
+		process.terminationHandler = { _ in
+			continuation.finish()
+			terminationHandler()
+		}
+
+		Task {
+			let dataStream = stdoutPipe.fileHandleForReading.dataStream
+
+			for try await data in dataStream {
+				continuation.yield(data)
+			}
+
+			continuation.finish()
+		}
+
+		Task {
+			for try await line in stderrPipe.fileHandleForReading.bytes.lines {
+				print("stderr: ", line)
+			}
+		}
+
+		try process.run()
+
+		let handler: DataChannel.WriteHandler = {
+			// this is wacky, but we need the channel to hold a strong reference to the process
+			// to prevent it from being deallocated
+			_ = process
+
+			try stdinPipe.fileHandleForWriting.write(contentsOf: $0)
+		}
+
+        return (process.processIdentifier, DataChannel(writeHandler: handler, dataSequence: stream))
+	}
+}

CodeEdit/Features/LSP/LanguageServer/LanguageServer.swift

@@ -40,17 +40,21 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
     private(set) var lspInstance: InitializingServer
     /// The path to the root of the project
     private(set) var rootPath: URL
+    /// The PID of the running language server process.
+    private(set) var pid: pid_t
 
     init(
         languageId: LanguageIdentifier,
         binary: LanguageServerBinary,
         lspInstance: InitializingServer,
+        lspPid: pid_t,
         serverCapabilities: ServerCapabilities,
         rootPath: URL
     ) {
         self.languageId = languageId
         self.binary = binary
         self.lspInstance = lspInstance
+        self.pid = lspPid
         self.serverCapabilities = serverCapabilities
         self.rootPath = rootPath
         self.openFiles = LanguageServerFileMap()
@@ -82,15 +86,17 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
             environment: binary.env
         )
 
+        let (pid, connection) = try makeLocalServerConnection(languageId: languageId, executionParams: executionParams)
         let server = InitializingServer(
-            server: try makeLocalServerConnection(languageId: languageId, executionParams: executionParams),
+            server: connection,
             initializeParamsProvider: getInitParams(workspacePath: workspacePath)
         )
         let capabilities = try await server.initializeIfNeeded()
         return LanguageServer(
             languageId: languageId,
             binary: binary,
             lspInstance: server,
+            lspPid: pid,
             serverCapabilities: capabilities,
             rootPath: URL(filePath: workspacePath)
         )
@@ -106,15 +112,15 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
     static func makeLocalServerConnection(
         languageId: LanguageIdentifier,
         executionParams: Process.ExecutionParameters
-    ) throws -> JSONRPCServerConnection {
+    ) throws -> (pid: pid_t, connection: JSONRPCServerConnection) {
         do {
-            let channel = try DataChannel.localProcessChannel(
+            let (pid, channel) = try DataChannel.localProcessChannel(
                 parameters: executionParams,
                 terminationHandler: {
                     logger.debug("Terminated data channel for \(languageId.rawValue)")
                 }
             )
-            return JSONRPCServerConnection(dataChannel: channel)
+            return (pid, JSONRPCServerConnection(dataChannel: channel))
         } catch {
             logger.warning("Failed to initialize data channel for \(languageId.rawValue)")
             throw error
@@ -232,10 +238,14 @@ class LanguageServer<DocumentType: LanguageServerDocument> {
         // swiftlint:enable function_body_length
     }
 
+    // MARK: - Shutdown
+
     /// Shuts down the language server and exits it.
     public func shutdown() async throws {
         self.logger.info("Shutting down language server")
-        try await lspInstance.shutdownAndExit()
+        try await withTimeout(duration: .seconds(1.0)) {
+            try await self.lspInstance.shutdownAndExit()
+        }
     }
 }
 

CodeEdit/Features/LSP/Service/LSPService.swift

@@ -300,16 +300,13 @@ final class LSPService: ObservableObject {
 
     /// Goes through all active language servers and attempts to shut them down.
     func stopAllServers() async {
-        await withThrowingTaskGroup(of: Void.self) { group in
-            for (key, server) in languageClients {
-                group.addTask {
-                    do {
-                        try await server.shutdown()
-                    } catch {
-                        self.logger.error("Shutting down \(key.languageId.rawValue): Error \(error)")
-                        throw error
-                    }
-                }
+        // Note: This is no longer a task group for a *REASON*
+        //       The task group for some reason would never return from the `await` suspension point.
+        for (key, server) in languageClients {
+            do {
+                try await server.shutdown()
+            } catch {
+                self.logger.warning("Shutting down \(key.languageId.rawValue): Error \(error)")
             }
         }
         languageClients.removeAll()
@@ -318,4 +315,11 @@ final class LSPService: ObservableObject {
         }
         eventListeningTasks.removeAll()
     }
+
+    /// Call this when a server is refusing to terminate itself. Sends the `SIGKILL` signal to all lsp processes.
+    func killAllServers() {
+        for (_, server) in languageClients {
+            kill(server.pid, SIGKILL)
+        }
+    }
 }