Add ServerNotificationHandler to handle notifications

Author: intitni

Tool/Sources/GitHubCopilotService/LanguageServer/CopilotLocalProcessServer.swift

@@ -18,35 +18,43 @@ class CopilotLocalProcessServer {
     public convenience init(
         path: String,
         arguments: [String],
-        environment: [String: String]? = nil
+        environment: [String: String]? = nil,
+        serverNotificationHandler: ServerNotificationHandler
     ) {
         let params = Process.ExecutionParameters(
             path: path,
             arguments: arguments,
             environment: environment
         )
 
-        self.init(executionParameters: params)
+        self.init(executionParameters: params, serverNotificationHandler: serverNotificationHandler)
     }
 
-    init(executionParameters parameters: Process.ExecutionParameters) {
+    init(
+        executionParameters parameters: Process.ExecutionParameters,
+        serverNotificationHandler: ServerNotificationHandler
+    ) {
         transport = StdioDataTransport()
         let framing = SeperatedHTTPHeaderMessageFraming()
         let messageTransport = MessageTransport(
             dataTransport: transport,
             messageProtocol: framing
         )
         customTransport = CustomDataTransport(nextTransport: messageTransport)
-        wrappedServer = CustomJSONRPCLanguageServer(dataTransport: customTransport)
+        wrappedServer = CustomJSONRPCLanguageServer(
+            dataTransport: customTransport,
+            serverNotificationHandler: serverNotificationHandler
+        )
 
         process = Process()
 
         // Because the implementation of LanguageClient is so closed,
         // we need to get the request IDs from a custom transport before the data
         // is written to the language server.
         customTransport.onWriteRequest = { [weak self] request in
-            if request.method == "getCompletionsCycling" 
-                || request.method ==  "textDocument/inlineCompletion" {
+            if request.method == "getCompletionsCycling"
+                || request.method == "textDocument/inlineCompletion"
+            {
                 Task { @MainActor [weak self] in
                     self?.ongoingCompletionRequestIDs.append(request.id)
                 }
@@ -85,7 +93,7 @@ class CopilotLocalProcessServer {
         get { return wrappedServer?.logMessages ?? false }
         set { wrappedServer?.logMessages = newValue }
     }
-    
+
     func terminate() {
         process.terminate()
     }
@@ -155,36 +163,47 @@ final class CustomJSONRPCLanguageServer: Server {
 
     private let protocolTransport: ProtocolTransport
 
-    public var requestHandler: RequestHandler?
-    public var notificationHandler: NotificationHandler?
+    var requestHandler: RequestHandler?
+    var serverNotificationHandler: ServerNotificationHandler
+
+    @available(*, deprecated, message: "Use `serverNotificationHandler` instead.")
+    var notificationHandler: NotificationHandler? {
+        get { nil }
+        set {}
+    }
 
     private var outOfBandError: Error?
 
-    init(protocolTransport: ProtocolTransport) {
+    init(
+        protocolTransport: ProtocolTransport,
+        serverNotificationHandler: ServerNotificationHandler
+    ) {
+        self.serverNotificationHandler = serverNotificationHandler
         self.protocolTransport = protocolTransport
         internalServer = JSONRPCLanguageServer(protocolTransport: protocolTransport)
 
-        let previouseRequestHandler = protocolTransport.requestHandler
-        let previouseNotificationHandler = protocolTransport.notificationHandler
+        let previousRequestHandler = protocolTransport.requestHandler
 
-        protocolTransport
-            .requestHandler = { [weak self] in
-                guard let self else { return }
-                if !self.handleRequest($0, data: $1, callback: $2) {
-                    previouseRequestHandler?($0, $1, $2)
-                }
-            }
-        protocolTransport
-            .notificationHandler = { [weak self] in
-                guard let self else { return }
-                if !self.handleNotification($0, data: $1, block: $2) {
-                    previouseNotificationHandler?($0, $1, $2)
-                }
+        protocolTransport.requestHandler = { [weak self] in
+            guard let self else { return }
+            if !self.handleRequest($0, data: $1, callback: $2) {
+                previousRequestHandler?($0, $1, $2)
             }
+        }
+        protocolTransport.notificationHandler = { [weak self] in
+            guard let self else { return }
+            self.handleNotification($0, data: $1, block: $2)
+        }
     }
 
-    convenience init(dataTransport: DataTransport) {
-        self.init(protocolTransport: ProtocolTransport(dataTransport: dataTransport))
+    convenience init(
+        dataTransport: DataTransport,
+        serverNotificationHandler: ServerNotificationHandler
+    ) {
+        self.init(
+            protocolTransport: ProtocolTransport(dataTransport: dataTransport),
+            serverNotificationHandler: serverNotificationHandler
+        )
     }
 
     deinit {
@@ -203,7 +222,72 @@ extension CustomJSONRPCLanguageServer {
         _ anyNotification: AnyJSONRPCNotification,
         data: Data,
         block: @escaping (Error?) -> Void
+    ) {
+        Task {
+            do {
+                try await serverNotificationHandler.handleNotification(
+                    anyNotification,
+                    data: data
+                )
+                block(nil)
+            } catch {
+                block(error)
+            }
+        }
+    }
+
+    func sendNotification(
+        _ notif: ClientNotification,
+        completionHandler: @escaping (ServerError?) -> Void
+    ) {
+        internalServer.sendNotification(notif, completionHandler: completionHandler)
+    }
+}
+
+extension CustomJSONRPCLanguageServer {
+    private func handleRequest(
+        _ request: AnyJSONRPCRequest,
+        data: Data,
+        callback: @escaping (AnyJSONRPCResponse) -> Void
     ) -> Bool {
+        return false
+    }
+}
+
+extension CustomJSONRPCLanguageServer {
+    public func sendRequest<Response: Codable>(
+        _ request: ClientRequest,
+        completionHandler: @escaping (ServerResult<Response>) -> Void
+    ) {
+        internalServer.sendRequest(request, completionHandler: completionHandler)
+    }
+}
+
+@GitHubCopilotSuggestionActor
+final class ServerNotificationHandler {
+    typealias Handler = (
+        _ anyNotification: AnyJSONRPCNotification,
+        _ data: Data
+    ) async throws -> Bool
+
+    var handlers = [AnyHashable: Handler]()
+    nonisolated init() {}
+
+    func handleNotification(
+        _ anyNotification: AnyJSONRPCNotification,
+        data: Data
+    ) async throws {
+        for handler in handlers.values {
+            do {
+                let handled = try await handler(anyNotification, data)
+                if handled {
+                    return
+                }
+            } catch {
+                throw ServerError.notificationDispatchFailed(error)
+            }
+        }
+
         let methodName = anyNotification.method
         let debugDescription = {
             if let params = anyNotification.params {
@@ -223,65 +307,29 @@ extension CustomJSONRPCLanguageServer {
                 Logger.gitHubCopilot
                     .info("\(anyNotification.method): \(debugDescription)")
             }
-            block(nil)
-            return true
         case "LogMessage":
             if UserDefaults.shared.value(for: \.gitHubCopilotVerboseLog) {
                 Logger.gitHubCopilot
                     .info("\(anyNotification.method): \(debugDescription)")
             }
-            block(nil)// 
-            return true
         case "statusNotification":
             if UserDefaults.shared.value(for: \.gitHubCopilotVerboseLog) {
                 Logger.gitHubCopilot
                     .info("\(anyNotification.method): \(debugDescription)")
             }
-            block(nil)
-            return true
         case "featureFlagsNotification":
             if UserDefaults.shared.value(for: \.gitHubCopilotVerboseLog) {
                 Logger.gitHubCopilot
                     .info("\(anyNotification.method): \(debugDescription)")
             }
-            block(nil)
-            return true
         case "conversation/preconditionsNotification":
             if UserDefaults.shared.value(for: \.gitHubCopilotVerboseLog) {
                 Logger.gitHubCopilot
                     .info("\(anyNotification.method): \(debugDescription)")
             }
-            block(nil)
-            return true
         default:
-            return false
+            throw ServerError.handlerUnavailable(methodName)
         }
     }
-
-    public func sendNotification(
-        _ notif: ClientNotification,
-        completionHandler: @escaping (ServerError?) -> Void
-    ) {
-        internalServer.sendNotification(notif, completionHandler: completionHandler)
-    }
-}
-
-extension CustomJSONRPCLanguageServer {
-    private func handleRequest(
-        _ request: AnyJSONRPCRequest,
-        data: Data,
-        callback: @escaping (AnyJSONRPCResponse) -> Void
-    ) -> Bool {
-        return false
-    }
-}
-
-extension CustomJSONRPCLanguageServer {
-    public func sendRequest<Response: Codable>(
-        _ request: ClientRequest,
-        completionHandler: @escaping (ServerResult<Response>) -> Void
-    ) {
-        internalServer.sendRequest(request, completionHandler: completionHandler)
-    }
 }