Add SuggestionServiceMiddleware

Tool/Sources/SuggestionService/CodeiumSuggestionProvider.swift

@@ -28,23 +28,17 @@ actor CodeiumSuggestionProvider: SuggestionServiceProvider {
 }
 
 extension CodeiumSuggestionProvider {
-    func getSuggestions(
-        fileURL: URL,
-        content: String,
-        cursorPosition: SuggestionModel.CursorPosition,
-        tabSize: Int,
-        indentSize: Int,
-        usesTabsForIndentation: Bool,
-        ignoreSpaceOnlySuggestions: Bool
-    ) async throws -> [SuggestionModel.CodeSuggestion] {
-        try await (try createCodeiumServiceIfNeeded()).getCompletions(
-            fileURL: fileURL,
-            content: content,
-            cursorPosition: cursorPosition,
-            tabSize: tabSize,
-            indentSize: indentSize,
-            usesTabsForIndentation: usesTabsForIndentation,
-            ignoreSpaceOnlySuggestions: ignoreSpaceOnlySuggestions
+    func getSuggestions(_ request: SuggestionRequest) async throws
+        -> [SuggestionModel.CodeSuggestion]
+    {
+        try await (createCodeiumServiceIfNeeded()).getCompletions(
+            fileURL: request.fileURL,
+            content: request.content,
+            cursorPosition: request.cursorPosition,
+            tabSize: request.tabSize,
+            indentSize: request.indentSize,
+            usesTabsForIndentation: request.usesTabsForIndentation,
+            ignoreSpaceOnlySuggestions: request.ignoreSpaceOnlySuggestions
         )
     }
 
@@ -70,12 +64,12 @@ extension CodeiumSuggestionProvider {
     }
 
     func notifySaveTextDocument(fileURL: URL) async throws {}
-    
+
     func cancelRequest() async {
         await (try? createCodeiumServiceIfNeeded())?
             .cancelRequest()
     }
-    
+
     func terminate() async {
         (try? createCodeiumServiceIfNeeded())?.terminate()
     }

Tool/Sources/SuggestionService/GitHubCopilotSuggestionProvider.swift

@@ -26,23 +26,17 @@ actor GitHubCopilotSuggestionProvider: SuggestionServiceProvider {
 }
 
 extension GitHubCopilotSuggestionProvider {
-    func getSuggestions(
-        fileURL: URL,
-        content: String,
-        cursorPosition: SuggestionModel.CursorPosition,
-        tabSize: Int,
-        indentSize: Int,
-        usesTabsForIndentation: Bool,
-        ignoreSpaceOnlySuggestions: Bool
-    ) async throws -> [SuggestionModel.CodeSuggestion] {
-        try await (try createGitHubCopilotServiceIfNeeded()).getCompletions(
-            fileURL: fileURL,
-            content: content,
-            cursorPosition: cursorPosition,
-            tabSize: tabSize,
-            indentSize: indentSize,
-            usesTabsForIndentation: usesTabsForIndentation,
-            ignoreSpaceOnlySuggestions: ignoreSpaceOnlySuggestions,
+    func getSuggestions(_ request: SuggestionRequest) async throws
+        -> [SuggestionModel.CodeSuggestion]
+    {
+        try await (createGitHubCopilotServiceIfNeeded()).getCompletions(
+            fileURL: request.fileURL,
+            content: request.content,
+            cursorPosition: request.cursorPosition,
+            tabSize: request.tabSize,
+            indentSize: request.indentSize,
+            usesTabsForIndentation: request.usesTabsForIndentation,
+            ignoreSpaceOnlySuggestions: request.ignoreSpaceOnlySuggestions,
             ignoreTrailingNewLinesAndSpaces: UserDefaults.shared
                 .value(for: \.gitHubCopilotIgnoreTrailingNewLines)
         )

Tool/Sources/SuggestionService/SuggestionService.swift

@@ -4,16 +4,36 @@ import Preferences
 import SuggestionModel
 import UserDefaultsObserver
 
-public protocol SuggestionServiceType {
-    func getSuggestions(
+public struct SuggestionRequest {
+    public var fileURL: URL
+    public var content: String
+    public var cursorPosition: CursorPosition
+    public var tabSize: Int
+    public var indentSize: Int
+    public var usesTabsForIndentation: Bool
+    public var ignoreSpaceOnlySuggestions: Bool
+
+    public init(
         fileURL: URL,
         content: String,
         cursorPosition: CursorPosition,
         tabSize: Int,
         indentSize: Int,
         usesTabsForIndentation: Bool,
         ignoreSpaceOnlySuggestions: Bool
-    ) async throws -> [CodeSuggestion]
+    ) {
+        self.fileURL = fileURL
+        self.content = content
+        self.cursorPosition = cursorPosition
+        self.tabSize = tabSize
+        self.indentSize = indentSize
+        self.usesTabsForIndentation = usesTabsForIndentation
+        self.ignoreSpaceOnlySuggestions = ignoreSpaceOnlySuggestions
+    }
+}
+
+public protocol SuggestionServiceType {
+    func getSuggestions(_ request: SuggestionRequest) async throws -> [CodeSuggestion]
 
     func notifyAccepted(_ suggestion: CodeSuggestion) async
     func notifyRejected(_ suggestions: [CodeSuggestion]) async
@@ -25,9 +45,45 @@ public protocol SuggestionServiceType {
     func terminate() async
 }
 
+public extension SuggestionServiceType {
+    func getSuggestions(
+        fileURL: URL,
+        content: String,
+        cursorPosition: CursorPosition,
+        tabSize: Int,
+        indentSize: Int,
+        usesTabsForIndentation: Bool,
+        ignoreSpaceOnlySuggestions: Bool
+    ) async throws -> [CodeSuggestion] {
+        return try await getSuggestions(.init(
+            fileURL: fileURL,
+            content: content,
+            cursorPosition: cursorPosition,
+            tabSize: tabSize,
+            indentSize: indentSize,
+            usesTabsForIndentation: usesTabsForIndentation,
+            ignoreSpaceOnlySuggestions: ignoreSpaceOnlySuggestions
+        ))
+    }
+}
+
 protocol SuggestionServiceProvider: SuggestionServiceType {}
 
 public actor SuggestionService: SuggestionServiceType {
+    static var builtInMiddlewares: [SuggestionServiceMiddleware] = [
+        DisabledLanguageSuggestionServiceMiddleware()
+    ]
+
+    static var customMiddlewares: [SuggestionServiceMiddleware] = []
+
+    static var middlewares: [SuggestionServiceMiddleware] {
+        builtInMiddlewares + customMiddlewares
+    }
+
+    public static func addMiddleware(_ middleware: SuggestionServiceMiddleware) {
+        customMiddlewares.append(middleware)
+    }
+
     let projectRootURL: URL
     let onServiceLaunched: (SuggestionServiceType) -> Void
     let providerChangeObserver = UserDefaultsObserver(
@@ -75,31 +131,16 @@ public actor SuggestionService: SuggestionServiceType {
 }
 
 public extension SuggestionService {
-    func getSuggestions(
-        fileURL: URL,
-        content: String,
-        cursorPosition: SuggestionModel.CursorPosition,
-        tabSize: Int,
-        indentSize: Int,
-        usesTabsForIndentation: Bool,
-        ignoreSpaceOnlySuggestions: Bool
-    ) async throws -> [SuggestionModel.CodeSuggestion] {
-        let language = languageIdentifierFromFileURL(fileURL)
-        if UserDefaults.shared.value(for: \.suggestionFeatureDisabledLanguageList)
-            .contains(where: { $0 == language.rawValue })
-        {
-            return []
+    func getSuggestions(_ request: SuggestionRequest) async throws -> [SuggestionModel.CodeSuggestion] {
+        var getSuggestion = suggestionProvider.getSuggestions
+
+        for middleware in Self.middlewares.reversed() {
+            getSuggestion = { request in
+                try await middleware.getSuggestion(request, next: getSuggestion)
+            }
         }
 
-        return try await suggestionProvider.getSuggestions(
-            fileURL: fileURL,
-            content: content,
-            cursorPosition: cursorPosition,
-            tabSize: tabSize,
-            indentSize: indentSize,
-            usesTabsForIndentation: usesTabsForIndentation,
-            ignoreSpaceOnlySuggestions: ignoreSpaceOnlySuggestions
-        )
+        return try await getSuggestion(request)
     }
 
     func notifyAccepted(_ suggestion: SuggestionModel.CodeSuggestion) async {

Tool/Sources/SuggestionService/SuggestionServiceMiddleware.swift

@@ -0,0 +1,41 @@
+import Foundation
+import SuggestionModel
+import Logger
+
+public protocol SuggestionServiceMiddleware {
+    typealias Next = (SuggestionRequest) async throws -> [CodeSuggestion]
+
+    func getSuggestion(_ request: SuggestionRequest, next: Next) async throws -> [CodeSuggestion]
+}
+
+struct DisabledLanguageSuggestionServiceMiddleware: SuggestionServiceMiddleware {
+    func getSuggestion(_ request: SuggestionRequest, next: Next) async throws -> [CodeSuggestion] {
+        let language = languageIdentifierFromFileURL(request.fileURL)
+        if UserDefaults.shared.value(for: \.suggestionFeatureDisabledLanguageList)
+            .contains(where: { $0 == language.rawValue })
+        {
+            #if DEBUG
+            Logger.service.info("Suggestion service is disabled for \(language).")
+            #endif
+            return []
+        }
+
+        return try await next(request)
+    }
+}
+
+public struct DebugSuggestionServiceMiddleware: SuggestionServiceMiddleware {
+    public init() {}
+
+    public func getSuggestion(_ request: SuggestionRequest, next: Next) async throws -> [CodeSuggestion] {
+        Logger.service.debug("""
+        Get suggestion for \(request.fileURL) at \(request.cursorPosition)
+        """)
+        let suggestions = try await next(request)
+        Logger.service.debug("""
+        Receive \(suggestions.count) suggestions for \(request.fileURL) at \(request.cursorPosition)
+        """)
+
+        return suggestions
+    }
+}