Merge branch 'feature/open-file-aware-suggestion' into develop

intitni · intitni · commit 06811a979e12 · 2023-04-29T23:21:26.000+08:00
diff --git a/Core/Sources/CopilotService/CopilotService.swift b/Core/Sources/CopilotService/CopilotService.swift
@@ -2,6 +2,7 @@ import CopilotModel
 import Foundation
 import LanguageClient
 import LanguageServerProtocol
+import Logger
 import Preferences
 import XPCShared
 
@@ -25,10 +26,15 @@ public protocol CopilotSuggestionServiceType {
     ) async throws -> [CopilotCompletion]
     func notifyAccepted(_ completion: CopilotCompletion) async
     func notifyRejected(_ completions: [CopilotCompletion]) 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
 }
 
 protocol CopilotLSP {
     func sendRequest<E: CopilotRequestType>(_ endpoint: E) async throws -> E.Response
+    func sendNotification(_ notif: ClientNotification) async throws
 }
 
 public class CopilotBaseService {
@@ -57,7 +63,7 @@ public class CopilotBaseService {
             }
             let executionParams: Process.ExecutionParameters
             let runner = UserDefaults.shared.value(for: \.runNodeWith)
-            
+
             switch runner {
             case .bash:
                 let nodePath = UserDefaults.shared.value(for: \.nodePath)
@@ -248,6 +254,57 @@ public final class CopilotSuggestionService: CopilotBaseService, CopilotSuggesti
             CopilotRequest.NotifyRejected(completionUUIDs: completions.map(\.uuid))
         )
     }
+
+    public func notifyOpenTextDocument(
+        fileURL: URL,
+        content: String
+    ) async throws {
+        let languageId = languageIdentifierFromFileURL(fileURL)
+        let uri = "file://\(fileURL.path)"
+//        Logger.service.debug("Open \(uri)")
+        try await server.sendNotification(
+            .didOpenTextDocument(
+                DidOpenTextDocumentParams(
+                    textDocument: .init(
+                        uri: uri,
+                        languageId: languageId.rawValue,
+                        version: 0,
+                        text: content
+                    )
+                )
+            )
+        )
+    }
+
+    public func notifyChangeTextDocument(fileURL: URL, content: String) async throws {
+        let uri = "file://\(fileURL.path)"
+//        Logger.service.debug("Change \(uri)")
+        try await server.sendNotification(
+            .didChangeTextDocument(
+                DidChangeTextDocumentParams(
+                    uri: uri,
+                    version: 0,
+                    contentChange: .init(
+                        range: nil,
+                        rangeLength: nil,
+                        text: content
+                    )
+                )
+            )
+        )
+    }
+
+    public func notifySaveTextDocument(fileURL: URL) async throws {
+        let uri = "file://\(fileURL.path)"
+//        Logger.service.debug("Save \(uri)")
+        try await server.sendNotification(.didSaveTextDocument(.init(uri: uri)))
+    }
+
+    public func notifyCloseTextDocument(fileURL: URL) async throws {
+        let uri = "file://\(fileURL.path)"
+//        Logger.service.debug("Close \(uri)")
+        try await server.sendNotification(.didCloseTextDocument(.init(uri: uri)))
+    }
 }
 
 extension InitializingServer: CopilotLSP {
diff --git a/Core/Sources/Service/RealtimeSuggestionController.swift b/Core/Sources/Service/RealtimeSuggestionController.swift
@@ -131,6 +131,7 @@ public class RealtimeSuggestionController {
                 switch notification.name {
                 case kAXValueChangedNotification:
                     self.triggerPrefetchDebounced()
+                    await self.notifyEditingFileChange(editor: focusElement)
                 default:
                     continue
                 }
@@ -207,7 +208,7 @@ public class RealtimeSuggestionController {
 
             guard UserDefaults.shared.value(for: \.realtimeSuggestionToggle)
             else { return }
-            
+
             if UserDefaults.shared.value(for: \.disableSuggestionFeatureGlobally),
                let fileURL = try? await Environment.fetchCurrentFileURL(),
                let (workspace, _) = try? await Workspace
@@ -216,7 +217,7 @@ public class RealtimeSuggestionController {
                 let isEnabled = workspace.isSuggestionFeatureEnabled
                 if !isEnabled { return }
             }
-            
+
             if Task.isCancelled { return }
 
             Logger.service.info("Prefetch suggestions.")
@@ -264,4 +265,12 @@ public class RealtimeSuggestionController {
         let application = AXUIElementCreateApplication(activeXcode.processIdentifier)
         return application.focusedWindow?.child(identifier: "_XC_COMPLETION_TABLE_") != nil
     }
+
+    func notifyEditingFileChange(editor: AXUIElement) async {
+        guard let fileURL = try? await Environment.fetchCurrentFileURL(),
+            let (workspace, filespace) = try? await Workspace
+            .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+        else { return }
+        workspace.notifyUpdateFile(filespace: filespace, content: editor.value)
+    }
 }
diff --git a/Core/Sources/Service/ScheduledCleaner.swift b/Core/Sources/Service/ScheduledCleaner.swift
@@ -5,7 +5,7 @@ public final class ScheduledCleaner {
         // occasionally cleanup workspaces.
         Task { @ServiceActor in
             while !Task.isCancelled {
-                try await Task.sleep(nanoseconds: 8 * 60 * 60 * 1_000_000_000)
+                try await Task.sleep(nanoseconds: 2 * 60 * 60 * 1_000_000_000)
                 for (url, workspace) in workspaces {
                     if workspace.isExpired {
                         workspaces[url] = nil
diff --git a/Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift b/Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift
@@ -147,7 +147,7 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         defer { presenter.markAsProcessing(false) }
 
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let (workspace, _) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
+        let (workspace, filespace) = try await Workspace.fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
 
         let injector = SuggestionInjector()
         var lines = editor.lines
diff --git a/Core/Sources/Service/Workspace.swift b/Core/Sources/Service/Workspace.swift
@@ -3,6 +3,7 @@ import CopilotModel
 import CopilotService
 import Environment
 import Foundation
+import Logger
 import Preferences
 import SuggestionInjector
 import XPCShared
@@ -17,7 +18,7 @@ final class Filespace {
     let fileURL: URL
     private(set) lazy var language: String = languageIdentifierFromFileURL(fileURL).rawValue
     var suggestions: [CopilotCompletion] = [] {
-        didSet { lastSuggestionUpdateTime = Environment.now() }
+        didSet { refreshUpdateTime() }
     }
 
     // stored for pseudo command handler
@@ -39,8 +40,15 @@ final class Filespace {
         Environment.now().timeIntervalSince(lastSuggestionUpdateTime) > 60 * 60 * 8
     }
 
-    init(fileURL: URL) {
+    let fileSaveWatcher: FileSaveWatcher
+
+    fileprivate init(fileURL: URL, onSave: @escaping (Filespace) -> Void) {
         self.fileURL = fileURL
+        fileSaveWatcher = .init(fileURL: fileURL)
+        fileSaveWatcher.changeHandler = { [weak self] in
+            guard let self else { return }
+            onSave(self)
+        }
     }
 
     func reset(resetSnapshot: Bool = true) {
@@ -50,6 +58,10 @@ final class Filespace {
             suggestionSourceSnapshot = .init(linesHash: -1, cursorPosition: .outOfScope)
         }
     }
+
+    func refreshUpdateTime() {
+        lastSuggestionUpdateTime = Environment.now()
+    }
 }
 
 @ServiceActor
@@ -66,7 +78,7 @@ final class Workspace {
             onChange?()
         }
     }
-    
+
     struct SuggestionFeatureDisabledError: Error, LocalizedError {
         var errorDescription: String? {
             "Suggestion feature is disabled for this project."
@@ -107,7 +119,7 @@ final class Workspace {
         }
         return _copilotSuggestionService
     }
-    
+
     var isSuggestionFeatureEnabled: Bool {
         let isSuggestionDisabledGlobally = UserDefaults.shared
             .value(for: \.disableSuggestionFeatureGlobally)
@@ -122,7 +134,7 @@ final class Workspace {
 
     private init(projectRootURL: URL) {
         self.projectRootURL = projectRootURL
-        
+
         Task {
             userDefaultsObserver.onChange = { [weak self] in
                 guard let self else { return }
@@ -135,7 +147,7 @@ final class Workspace {
                 options: .new,
                 context: nil
             )
-            
+
             UserDefaults.shared.addObserver(
                 userDefaultsObserver,
                 forKeyPath: UserDefaultPreferenceKeys().disableSuggestionFeatureGlobally.key,
@@ -166,17 +178,32 @@ final class Workspace {
                 return (workspace, filespace)
             }
         }
-        
+
         let projectURL = try await Environment.fetchCurrentProjectRootURL(fileURL)
         let workspaceURL = projectURL ?? fileURL
         let workspace = workspaces[workspaceURL] ?? Workspace(projectRootURL: workspaceURL)
-        let filespace = workspace.filespaces[fileURL] ?? .init(fileURL: fileURL)
-        if workspace.filespaces[fileURL] == nil {
-            workspace.filespaces[fileURL] = filespace
-        }
+        let filespace = workspace.createFilespaceIfNeeded(fileURL: fileURL)
         workspaces[workspaceURL] = workspace
         return (workspace, filespace)
     }
+
+    func createFilespaceIfNeeded(fileURL: URL) -> Filespace {
+        let existedFilespace = filespaces[fileURL]
+        let filespace = existedFilespace ?? .init(fileURL: fileURL, onSave: { [weak self]
+            filespace in
+                guard let self else { return }
+                notifySaveFile(filespace: filespace)
+        })
+        if filespaces[fileURL] == nil {
+            filespaces[fileURL] = filespace
+        }
+        if existedFilespace == nil {
+            notifyOpenFile(filespace: filespace)
+        } else {
+            filespace.refreshUpdateTime()
+        }
+        return filespace
+    }
 }
 
 extension Workspace {
@@ -191,7 +218,8 @@ extension Workspace {
         }
         lastTriggerDate = Environment.now()
 
-        let filespace = filespaces[fileURL] ?? .init(fileURL: fileURL)
+        let filespace = createFilespaceIfNeeded(fileURL: fileURL)
+
         if filespaces[fileURL] == nil {
             filespaces[fileURL] = filespace
         }
@@ -295,12 +323,39 @@ extension Workspace {
 
         return suggestion
     }
+
+    func notifyOpenFile(filespace: Filespace) {
+        Task {
+            try await copilotSuggestionService?.notifyOpenTextDocument(
+                fileURL: filespace.fileURL,
+                content: try String(contentsOf: filespace.fileURL, encoding: .utf8)
+            )
+        }
+    }
+
+    func notifyUpdateFile(filespace: Filespace, content: String) {
+        Task {
+            try await copilotSuggestionService?.notifyChangeTextDocument(
+                fileURL: filespace.fileURL,
+                content: content
+            )
+        }
+    }
+
+    func notifySaveFile(filespace: Filespace) {
+        Task {
+            try await copilotSuggestionService?.notifySaveTextDocument(fileURL: filespace.fileURL)
+        }
+    }
 }
 
 extension Workspace {
     func cleanUp() {
         for (fileURL, filespace) in filespaces {
             if filespace.isExpired {
+                Task {
+                    try await copilotSuggestionService?.notifyCloseTextDocument(fileURL: fileURL)
+                }
                 filespaces[fileURL] = nil
             }
         }
@@ -313,3 +368,41 @@ extension Workspace {
         realtimeSuggestionRequests = []
     }
 }
+
+final class FileSaveWatcher {
+    let url: URL
+    var fileHandle: FileHandle?
+    var source: DispatchSourceFileSystemObject?
+    var changeHandler: () -> Void = {}
+    
+    init(fileURL: URL) {
+        url = fileURL
+        startup()
+    }
+    
+    deinit {
+        source?.cancel()
+    }
+    
+    func startup() {
+        if let source = source {
+            source.cancel()
+        }
+        
+        fileHandle = try? FileHandle(forReadingFrom: url)
+        if let fileHandle {
+            source = DispatchSource.makeFileSystemObjectSource(
+                fileDescriptor: fileHandle.fileDescriptor,
+                eventMask: .link,
+                queue: .main
+            )
+
+            source?.setEventHandler { [weak self] in
+                self?.changeHandler()
+                self?.startup()
+            }
+            
+            source?.resume()
+        }
+    }
+}
