Support heartbeat

Author: intitni

Core/Sources/CodeiumService/CodeiumLanguageServer.swift

@@ -8,7 +8,7 @@ protocol CodeiumLSP {
     func sendRequest<E: CodeiumRequestType>(_ endpoint: E) async throws -> E.Response
 }
 
-class CodeiumLanguageServer: CodeiumLSP {
+final class CodeiumLanguageServer {
     let languageServerExecutableURL: URL
     let managerDirectoryURL: URL
     let supportURL: URL
@@ -17,6 +17,7 @@ class CodeiumLanguageServer: CodeiumLSP {
     var terminationHandler: (() -> Void)?
     var launchHandler: (() -> Void)?
     var port: String?
+    var heartbeatTask: Task<Void, Error>?
 
     init(
         languageServerExecutableURL: URL,
@@ -88,10 +89,8 @@ class CodeiumLanguageServer: CodeiumLSP {
                 while true {
                     try await Task.sleep(nanoseconds: 1_000_000_000)
                     waited += 1
-                    port = findPort()
-                    if port != nil {
-                        Logger.codeium.info("Language server started.")
-                        launchHandler?()
+                    if let port = findPort() {
+                        finishStarting(port: port)
                         return
                     }
                     if waited >= 60 {
@@ -116,6 +115,14 @@ class CodeiumLanguageServer: CodeiumLSP {
         terminationHandler?()
     }
 
+    private func finishStarting(port: String) {
+        Logger.codeium.info("Language server started.")
+        self.port = port
+        launchHandler?()
+    }
+}
+
+extension CodeiumLanguageServer: CodeiumLSP {
     func sendRequest<E>(_ request: E) async throws -> E.Response where E: CodeiumRequestType {
         guard let port else { throw CancellationError() }
 
@@ -126,6 +133,7 @@ class CodeiumLanguageServer: CodeiumLSP {
                 let response = try JSONDecoder().decode(E.Response.self, from: data)
                 return response
             } catch {
+                dump(error)
                 Logger.codeium.error(error.localizedDescription)
                 throw error
             }

Core/Sources/CodeiumService/CodeiumModels.swift

@@ -21,15 +21,15 @@ struct CodeiumCompletionItem: Codable {
     var suffix: Suffix?
     var range: Range
     var source: CompletionSource
-    var completionParts: [CompletionPart]
+    var completionParts: [CompletionPart]?
 }
 
 struct Suffix: Codable {
     /// Text to insert after the cursor when accepting the completion.
     var text: String
     /// Cursor position delta (as signed offset) from the end of the inserted
     /// completion (including the suffix).
-    var deltaCursorOffset: Int
+    var deltaCursorOffset: String
 }
 
 struct Range: Codable {

Core/Sources/CodeiumService/CodeiumRequest.swift

@@ -8,6 +8,18 @@ protocol CodeiumRequestType {
     func makeURLRequest(server: String) -> URLRequest
 }
 
+extension CodeiumRequestType {
+    func assembleURLRequest(server: String, method: String, body: Data?) -> URLRequest {
+        var request = URLRequest(url: .init(
+            string: "\(server)/exa.language_server_pb.LanguageServerService/\(method)"
+        )!)
+        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
+        request.httpMethod = "POST"
+        request.httpBody = body
+        return request
+    }
+}
+
 struct CodeiumResponseError: Codable, Error, LocalizedError {
     var code: String
     var message: String
@@ -21,26 +33,53 @@ enum CodeiumRequest {
             var completionItems: [CodeiumCompletionItem]?
         }
 
-        struct Request: Codable {
+        struct RequestBody: Codable {
             var metadata: Metadata
             var document: CodeiumDocument
             var editor_options: CodeiumEditorOptions
             var other_documents: [CodeiumDocument]
         }
 
-        var requestBody: Request
+        var requestBody: RequestBody
+
+        func makeURLRequest(server: String) -> URLRequest {
+            let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
+            return assembleURLRequest(server: server, method: "GetCompletions", body: data)
+        }
+    }
+    
+    struct AcceptCompletion: CodeiumRequestType {
+        struct Response: Codable {
+            var state: State
+            var completionItems: [CodeiumCompletionItem]?
+        }
+
+        struct RequestBody: Codable {
+            var metadata: Metadata
+            var completion_id: String
+        }
+
+        var requestBody: RequestBody
 
         func makeURLRequest(server: String) -> URLRequest {
-            var request = URLRequest(url: .init(string: "\(server)/exa.language_server_pb.LanguageServerService/GetCompletions")!)
-            request.setValue("application/json", forHTTPHeaderField: "Content-Type")
-            let data = (try? JSONEncoder().encode(requestBody)) ?? Data() //
-            
-            request.httpMethod = "POST"
-            request.httpBody = data
-            
-            return request
+            let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
+            return assembleURLRequest(server: server, method: "AcceptCompletion", body: data)
         }
+    }
+    
+    struct Heartbeat: CodeiumRequestType {
+        struct Response: Codable {}
 
+        struct RequestBody: Codable {
+            var metadata: Metadata
+        }
+
+        var requestBody: RequestBody
+
+        func makeURLRequest(server: String) -> URLRequest {
+            let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
+            return assembleURLRequest(server: server, method: "Heartbeat", body: data)
+        }
     }
 }
 

Core/Sources/CodeiumService/CodeiumService.swift

@@ -1,6 +1,7 @@
 import Foundation
 import LanguageClient
 import LanguageServerProtocol
+import Logger
 import SuggestionModel
 
 public protocol CodeiumSuggestionServiceType {
@@ -26,11 +27,13 @@ enum CodeiumError: Error, LocalizedError {
     }
 }
 
-let token = ""
 
 public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
+    static let sessionId = UUID().uuidString
     let projectRootURL: URL
     var server: CodeiumLSP
+    var heartbeatTask: Task<Void, Error>?
+    var requestCounter: UInt64 = 0
 
     init(designatedServer: CodeiumLSP) {
         projectRootURL = URL(fileURLWithPath: "/")
@@ -62,11 +65,22 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
         )
 
         self.server = server
-        server.terminationHandler = {
-            print("terminated")
+        server.terminationHandler = { [weak self] in
+            Logger.codeium.info("Language server is terminated")
+            guard let self else { return }
         }
-        server.launchHandler = {
-            print("launched")
+        server.launchHandler = { [weak self] in
+            guard let self else { return }
+            let metadata = self.getMetadata()
+            self.heartbeatTask = Task { [weak self] in
+                while true {
+                    try Task.checkCancellation()
+                    _ = try? await self?.server.sendRequest(
+                        CodeiumRequest.Heartbeat(requestBody: .init(metadata: metadata))
+                    )
+                    try await Task.sleep(nanoseconds: 5_000_000_000)
+                }
+            }
         }
         server.start()
     }
@@ -80,6 +94,7 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
         usesTabsForIndentation: Bool,
         ignoreSpaceOnlySuggestions: Bool
     ) async throws -> [CodeSuggestion] {
+        requestCounter += 1
         let languageId = languageIdentifierFromFileURL(fileURL)
 
         let relativePath = {
@@ -98,12 +113,7 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
         }()
 
         let request = CodeiumRequest.GetCompletion(requestBody: .init(
-            metadata: .init(
-                ide_name: "jetbrains", ide_version: "14.3", extension_name: "Copilot for Xcode",
-                extension_version: "14.0.0",
-                api_key: token,
-                session_id: UUID().uuidString, request_id: 100
-            ),
+            metadata: getMetadata(),
             document: .init(
                 absolute_path: fileURL.path,
                 relative_path: relativePath,
@@ -184,3 +194,17 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
     }
 }
 
+extension CodeiumSuggestionService {
+    func getMetadata() -> Metadata {
+        Metadata(
+            ide_name: "jetbrains",
+            ide_version: "14.3",
+            extension_name: "Copilot for Xcode",
+            extension_version: "14.0.0",
+            api_key: token,
+            session_id: CodeiumSuggestionService.sessionId,
+            request_id: requestCounter
+        )
+    }
+}
+