Adjust definition of SuggestionServiceProvider

Author: intitni

Tool/Sources/SuggestionProvider/SuggestionProvider.swift

@@ -1,5 +1,6 @@
 import AppKit
 import struct CopilotForXcodeKit.SuggestionServiceConfiguration
+import struct CopilotForXcodeKit.WorkspaceInfo
 import Foundation
 import Preferences
 import SuggestionModel
@@ -55,15 +56,19 @@ public struct RelevantCodeSnippet: Codable {
 }
 
 public protocol SuggestionServiceProvider {
-    func getSuggestions(_ request: SuggestionRequest) async throws -> [CodeSuggestion]
-    func notifyAccepted(_ suggestion: CodeSuggestion) async
-    func notifyRejected(_ suggestions: [CodeSuggestion]) async
-    func notifyOpenTextDocument(fileURL: URL, content: String) async throws
-    func notifyChangeTextDocument(fileURL: URL, content: String) async throws
-    func notifyCloseTextDocument(fileURL: URL) async throws
-    func notifySaveTextDocument(fileURL: URL) async throws
-    func cancelRequest() async
-    func terminate() async
+    func getSuggestions(
+        _ request: SuggestionRequest,
+        workspaceInfo: CopilotForXcodeKit.WorkspaceInfo
+    ) async throws -> [CodeSuggestion]
+    func notifyAccepted(
+        _ suggestion: CodeSuggestion,
+        workspaceInfo: CopilotForXcodeKit.WorkspaceInfo
+    ) async
+    func notifyRejected(
+        _ suggestions: [CodeSuggestion],
+        workspaceInfo: CopilotForXcodeKit.WorkspaceInfo
+    ) async
+    func cancelRequest(workspaceInfo: CopilotForXcodeKit.WorkspaceInfo) async
 
     var configuration: SuggestionServiceConfiguration { get async }
 }

Tool/Sources/WorkspaceSuggestionService/SuggestionWorkspacePlugin.swift

@@ -1,29 +1,36 @@
+import BuiltinExtension
 import Foundation
 import Preferences
 import SuggestionModel
 import SuggestionProvider
 import UserDefaultsObserver
 import Workspace
 
+#if canImport(ProExtension)
+import ProExtension
+#endif
+
 public final class SuggestionServiceWorkspacePlugin: WorkspacePlugin {
-    public typealias SuggestionServiceFactory = (
-        _ projectRootURL: URL,
-        _ onServiceLaunched: @escaping (any SuggestionServiceProvider) -> Void
-    ) -> any SuggestionServiceProvider
-    
-    let userDefaultsObserver = UserDefaultsObserver(
+    public typealias SuggestionServiceFactory = () -> any SuggestionServiceProvider
+    let suggestionServiceFactory: SuggestionServiceFactory
+
+    let suggestionFeatureUsabilityObserver = UserDefaultsObserver(
         object: UserDefaults.shared, forKeyPaths: [
             UserDefaultPreferenceKeys().suggestionFeatureEnabledProjectList.key,
             UserDefaultPreferenceKeys().disableSuggestionFeatureGlobally.key,
         ], context: nil
     )
 
+    let providerChangeObserver = UserDefaultsObserver(
+        object: UserDefaults.shared,
+        forKeyPaths: [UserDefaultPreferenceKeys().suggestionFeatureProvider.key],
+        context: nil
+    )
+
     public var isRealtimeSuggestionEnabled: Bool {
         UserDefaults.shared.value(for: \.realtimeSuggestionToggle)
     }
 
-    let suggestionServiceFactory: SuggestionServiceFactory
-
     private var _suggestionService: SuggestionServiceProvider?
 
     public var suggestionService: SuggestionServiceProvider? {
@@ -40,13 +47,7 @@ public final class SuggestionServiceWorkspacePlugin: WorkspacePlugin {
         }
 
         if _suggestionService == nil {
-            _suggestionService = suggestionServiceFactory(projectRootURL) {
-                [weak self] _ in
-                guard let self else { return }
-                for (_, filespace) in filespaces {
-                    notifyOpenFile(filespace: filespace)
-                }
-            }
+            _suggestionService = suggestionServiceFactory()
         }
         return _suggestionService
     }
@@ -62,77 +63,37 @@ public final class SuggestionServiceWorkspacePlugin: WorkspacePlugin {
         }
         return true
     }
-    
+
     public init(
         workspace: Workspace,
         suggestionProviderFactory: @escaping SuggestionServiceFactory
     ) {
-        self.suggestionServiceFactory = suggestionProviderFactory
+        suggestionServiceFactory = suggestionProviderFactory
         super.init(workspace: workspace)
 
-        userDefaultsObserver.onChange = { [weak self] in
+        suggestionFeatureUsabilityObserver.onChange = { [weak self] in
             guard let self else { return }
             _ = self.suggestionService
         }
-    }
 
-    override public func didOpenFilespace(_ filespace: Filespace) {
-        notifyOpenFile(filespace: filespace)
-    }
-
-    override public func didSaveFilespace(_ filespace: Filespace) {
-        notifySaveFile(filespace: filespace)
-    }
-
-    override public func didUpdateFilespace(_ filespace: Filespace, content: String) {
-        notifyUpdateFile(filespace: filespace, content: content)
-    }
-
-    override public func didCloseFilespace(_ fileURL: URL) {
-        Task {
-            try await suggestionService?.notifyCloseTextDocument(fileURL: fileURL)
-        }
-    }
-
-    public func notifyOpenFile(filespace: Filespace) {
-        Task {
-            guard filespace.isTextReadable else { return }
-            guard !(await filespace.isGitIgnored) else { return }
-            // check if file size is larger than 15MB, if so, return immediately
-            if let attrs = try? FileManager.default
-                .attributesOfItem(atPath: filespace.fileURL.path),
-                let fileSize = attrs[FileAttributeKey.size] as? UInt64,
-                fileSize > 15 * 1024 * 1024
-            { return }
-
-            try await suggestionService?.notifyOpenTextDocument(
-                fileURL: filespace.fileURL,
-                content: String(contentsOf: filespace.fileURL, encoding: .utf8)
-            )
-        }
-    }
-
-    public func notifyUpdateFile(filespace: Filespace, content: String) {
-        Task {
-            guard filespace.isTextReadable else { return }
-            guard !(await filespace.isGitIgnored) else { return }
-            try await suggestionService?.notifyChangeTextDocument(
-                fileURL: filespace.fileURL,
-                content: content
-            )
+        providerChangeObserver.onChange = { [weak self] in
+            guard let self else { return }
+            self._suggestionService = nil
         }
     }
 
-    public func notifySaveFile(filespace: Filespace) {
-        Task {
-            guard filespace.isTextReadable else { return }
-            guard !(await filespace.isGitIgnored) else { return }
-            try await suggestionService?.notifySaveTextDocument(fileURL: filespace.fileURL)
-        }
+    func notifyAccepted(_ suggestion: CodeSuggestion) async {
+        await suggestionService?.notifyAccepted(
+            suggestion,
+            workspaceInfo: .init(workspaceURL: workspaceURL, projectURL: projectRootURL)
+        )
     }
 
-    public func terminateSuggestionService() async {
-        await _suggestionService?.terminate()
+    func notifyRejected(_ suggestions: [CodeSuggestion]) async {
+        await suggestionService?.notifyRejected(
+            suggestions,
+            workspaceInfo: .init(workspaceURL: workspaceURL, projectURL: projectRootURL)
+        )
     }
 }
 

Tool/Sources/WorkspaceSuggestionService/Workspace+SuggestionService.swift

@@ -60,13 +60,14 @@ public extension Workspace {
                 relativePath: fileURL.path.replacingOccurrences(of: projectRootURL.path, with: ""),
                 content: editor.lines.joined(separator: ""),
                 lines: editor.lines,
-                cursorPosition: editor.cursorPosition, 
+                cursorPosition: editor.cursorPosition,
                 cursorOffset: editor.cursorOffset,
                 tabSize: editor.tabSize,
                 indentSize: editor.indentSize,
                 usesTabsForIndentation: editor.usesTabsForIndentation,
                 relevantCodeSnippets: []
-            )
+            ),
+            workspaceInfo: .init(workspaceURL: projectRootURL, projectURL: projectRootURL)
         )
 
         filespace.setSuggestions(completions)
@@ -102,9 +103,15 @@ public extension Workspace {
             filespaces[fileURL]?.codeMetadata.indentSize = editor.indentSize
             filespaces[fileURL]?.codeMetadata.usesTabsForIndentation = editor.usesTabsForIndentation
         }
-        
+
         Task {
-            await suggestionService?.notifyRejected(filespaces[fileURL]?.suggestions ?? [])
+            await suggestionService?.notifyRejected(
+                filespaces[fileURL]?.suggestions ?? [],
+                workspaceInfo: .init(
+                    workspaceURL: projectRootURL,
+                    projectURL: projectRootURL
+                )
+            )
         }
         filespaces[fileURL]?.reset()
     }
@@ -128,9 +135,14 @@ public extension Workspace {
         var allSuggestions = filespace.suggestions
         let suggestion = allSuggestions.remove(at: filespace.suggestionIndex)
 
-        Task { [allSuggestions] in
-            await suggestionService?.notifyAccepted(suggestion)
-            await suggestionService?.notifyRejected(allSuggestions)
+        Task {
+            await suggestionService?.notifyAccepted(
+                suggestion,
+                workspaceInfo: .init(
+                    workspaceURL: projectRootURL,
+                    projectURL: projectRootURL
+                )
+            )
         }
 
         filespaces[fileURL]?.reset()