Skip to content

Commit 38047b7

Browse files
intitniintitni
committed
Support heartbeat
1 parent bf43bb9 commit 38047b7

File tree

4 files changed

+99
-28
lines changed

4 files changed

+99
-28
lines changed

‎Core/Sources/CodeiumService/CodeiumLanguageServer.swift‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ protocol CodeiumLSP {
88
func sendRequest<E: CodeiumRequestType>(_ endpoint: E) async throws -> E.Response
99
}
1010

11-
class CodeiumLanguageServer: CodeiumLSP {
11+
final class CodeiumLanguageServer {
1212
let languageServerExecutableURL: URL
1313
let managerDirectoryURL: URL
1414
let supportURL: URL
@@ -17,6 +17,7 @@ class CodeiumLanguageServer: CodeiumLSP {
1717
var terminationHandler: (() -> Void)?
1818
var launchHandler: (() -> Void)?
1919
var port: String?
20+
var heartbeatTask: Task<Void, Error>?
2021

2122
init(
2223
languageServerExecutableURL: URL,
@@ -88,10 +89,8 @@ class CodeiumLanguageServer: CodeiumLSP {
8889
while true {
8990
try await Task.sleep(nanoseconds: 1_000_000_000)
9091
waited += 1
91-
port = findPort()
92-
if port != nil {
93-
Logger.codeium.info("Language server started.")
94-
launchHandler?()
92+
if let port = findPort() {
93+
finishStarting(port: port)
9594
return
9695
}
9796
if waited >= 60 {
@@ -116,6 +115,14 @@ class CodeiumLanguageServer: CodeiumLSP {
116115
terminationHandler?()
117116
}
118117

118+
private func finishStarting(port: String) {
119+
Logger.codeium.info("Language server started.")
120+
self.port = port
121+
launchHandler?()
122+
}
123+
}
124+
125+
extension CodeiumLanguageServer: CodeiumLSP {
119126
func sendRequest<E>(_ request: E) async throws -> E.Response where E: CodeiumRequestType {
120127
guard let port else { throw CancellationError() }
121128

@@ -126,6 +133,7 @@ class CodeiumLanguageServer: CodeiumLSP {
126133
let response = try JSONDecoder().decode(E.Response.self, from: data)
127134
return response
128135
} catch {
136+
dump(error)
129137
Logger.codeium.error(error.localizedDescription)
130138
throw error
131139
}

‎Core/Sources/CodeiumService/CodeiumModels.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,15 @@ struct CodeiumCompletionItem: Codable {
2121
var suffix: Suffix?
2222
var range: Range
2323
var source: CompletionSource
24-
var completionParts: [CompletionPart]
24+
var completionParts: [CompletionPart]?
2525
}
2626

2727
struct Suffix: Codable {
2828
/// Text to insert after the cursor when accepting the completion.
2929
var text: String
3030
/// Cursor position delta (as signed offset) from the end of the inserted
3131
/// completion (including the suffix).
32-
var deltaCursorOffset: Int
32+
var deltaCursorOffset: String
3333
}
3434

3535
struct Range: Codable {

‎Core/Sources/CodeiumService/CodeiumRequest.swift‎

Lines changed: 49 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,18 @@ protocol CodeiumRequestType {
88
func makeURLRequest(server: String) -> URLRequest
99
}
1010

11+
extension CodeiumRequestType {
12+
func assembleURLRequest(server: String, method: String, body: Data?) -> URLRequest {
13+
var request = URLRequest(url: .init(
14+
string: "\(server)/exa.language_server_pb.LanguageServerService/\(method)"
15+
)!)
16+
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
17+
request.httpMethod = "POST"
18+
request.httpBody = body
19+
return request
20+
}
21+
}
22+
1123
struct CodeiumResponseError: Codable, Error, LocalizedError {
1224
var code: String
1325
var message: String
@@ -21,26 +33,53 @@ enum CodeiumRequest {
2133
var completionItems: [CodeiumCompletionItem]?
2234
}
2335

24-
struct Request: Codable {
36+
struct RequestBody: Codable {
2537
var metadata: Metadata
2638
var document: CodeiumDocument
2739
var editor_options: CodeiumEditorOptions
2840
var other_documents: [CodeiumDocument]
2941
}
3042

31-
var requestBody: Request
43+
var requestBody: RequestBody
44+
45+
func makeURLRequest(server: String) -> URLRequest {
46+
let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
47+
return assembleURLRequest(server: server, method: "GetCompletions", body: data)
48+
}
49+
}
50+
51+
struct AcceptCompletion: CodeiumRequestType {
52+
struct Response: Codable {
53+
var state: State
54+
var completionItems: [CodeiumCompletionItem]?
55+
}
56+
57+
struct RequestBody: Codable {
58+
var metadata: Metadata
59+
var completion_id: String
60+
}
61+
62+
var requestBody: RequestBody
3263

3364
func makeURLRequest(server: String) -> URLRequest {
34-
var request = URLRequest(url: .init(string: "\(server)/exa.language_server_pb.LanguageServerService/GetCompletions")!)
35-
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
36-
let data = (try? JSONEncoder().encode(requestBody)) ?? Data() //
37-
38-
request.httpMethod = "POST"
39-
request.httpBody = data
40-
41-
return request
65+
let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
66+
return assembleURLRequest(server: server, method: "AcceptCompletion", body: data)
4267
}
68+
}
69+
70+
struct Heartbeat: CodeiumRequestType {
71+
struct Response: Codable {}
4372

73+
struct RequestBody: Codable {
74+
var metadata: Metadata
75+
}
76+
77+
var requestBody: RequestBody
78+
79+
func makeURLRequest(server: String) -> URLRequest {
80+
let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
81+
return assembleURLRequest(server: server, method: "Heartbeat", body: data)
82+
}
4483
}
4584
}
4685

‎Core/Sources/CodeiumService/CodeiumService.swift‎

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import Foundation
22
import LanguageClient
33
import LanguageServerProtocol
4+
import Logger
45
import SuggestionModel
56

67
public protocol CodeiumSuggestionServiceType {
@@ -26,11 +27,13 @@ enum CodeiumError: Error, LocalizedError {
2627
}
2728
}
2829

29-
let token = ""
3030

3131
public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
32+
static let sessionId = UUID().uuidString
3233
let projectRootURL: URL
3334
var server: CodeiumLSP
35+
var heartbeatTask: Task<Void, Error>?
36+
var requestCounter: UInt64 = 0
3437

3538
init(designatedServer: CodeiumLSP) {
3639
projectRootURL = URL(fileURLWithPath: "/")
@@ -62,11 +65,22 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
6265
)
6366

6467
self.server = server
65-
server.terminationHandler = {
66-
print("terminated")
68+
server.terminationHandler = { [weak self] in
69+
Logger.codeium.info("Language server is terminated")
70+
guard let self else { return }
6771
}
68-
server.launchHandler = {
69-
print("launched")
72+
server.launchHandler = { [weak self] in
73+
guard let self else { return }
74+
let metadata = self.getMetadata()
75+
self.heartbeatTask = Task { [weak self] in
76+
while true {
77+
try Task.checkCancellation()
78+
_ = try? await self?.server.sendRequest(
79+
CodeiumRequest.Heartbeat(requestBody: .init(metadata: metadata))
80+
)
81+
try await Task.sleep(nanoseconds: 5_000_000_000)
82+
}
83+
}
7084
}
7185
server.start()
7286
}
@@ -80,6 +94,7 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
8094
usesTabsForIndentation: Bool,
8195
ignoreSpaceOnlySuggestions: Bool
8296
) async throws -> [CodeSuggestion] {
97+
requestCounter += 1
8398
let languageId = languageIdentifierFromFileURL(fileURL)
8499

85100
let relativePath = {
@@ -98,12 +113,7 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
98113
}()
99114

100115
let request = CodeiumRequest.GetCompletion(requestBody: .init(
101-
metadata: .init(
102-
ide_name: "jetbrains", ide_version: "14.3", extension_name: "Copilot for Xcode",
103-
extension_version: "14.0.0",
104-
api_key: token,
105-
session_id: UUID().uuidString, request_id: 100
106-
),
116+
metadata: getMetadata(),
107117
document: .init(
108118
absolute_path: fileURL.path,
109119
relative_path: relativePath,
@@ -184,3 +194,17 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
184194
}
185195
}
186196

197+
extension CodeiumSuggestionService {
198+
func getMetadata() -> Metadata {
199+
Metadata(
200+
ide_name: "jetbrains",
201+
ide_version: "14.3",
202+
extension_name: "Copilot for Xcode",
203+
extension_version: "14.0.0",
204+
api_key: token,
205+
session_id: CodeiumSuggestionService.sessionId,
206+
request_id: requestCounter
207+
)
208+
}
209+
}
210+

0 commit comments

Comments
 (0)