Skip to content

Commit c16e7ef

Browse files
intitniintitni
committed
Support other_documents
1 parent 38047b7 commit c16e7ef

File tree

2 files changed

+181
-62
lines changed

2 files changed

+181
-62
lines changed

‎Core/Sources/CodeiumService/CodeiumService.swift‎

Lines changed: 111 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ public protocol CodeiumSuggestionServiceType {
1414
usesTabsForIndentation: Bool,
1515
ignoreSpaceOnlySuggestions: Bool
1616
) async throws -> [CodeSuggestion]
17+
func notifyAccepted(_ suggestion: CodeSuggestion) async
18+
func notifyOpenTextDocument(fileURL: URL, content: String) async throws
19+
func notifyChangeTextDocument(fileURL: URL, content: String) async throws
20+
func notifyCloseTextDocument(fileURL: URL) async throws
1721
}
1822

1923
enum CodeiumError: Error, LocalizedError {
@@ -28,19 +32,20 @@ enum CodeiumError: Error, LocalizedError {
2832
}
2933

3034

31-
public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
35+
public class CodeiumSuggestionService {
3236
static let sessionId = UUID().uuidString
3337
let projectRootURL: URL
3438
var server: CodeiumLSP
3539
var heartbeatTask: Task<Void, Error>?
3640
var requestCounter: UInt64 = 0
41+
let openedDocumentPool = OpenedDocumentPool()
3742

3843
init(designatedServer: CodeiumLSP) {
3944
projectRootURL = URL(fileURLWithPath: "/")
4045
server = designatedServer
4146
}
4247

43-
public init(projectRootURL: URL) throws {
48+
public init(projectRootURL: URL, onServiceLaunched: @escaping () -> Void) throws {
4449
self.projectRootURL = projectRootURL
4550

4651
let urls = try CodeiumSuggestionService.createFoldersIfNeeded()
@@ -72,6 +77,7 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
7277
server.launchHandler = { [weak self] in
7378
guard let self else { return }
7479
let metadata = self.getMetadata()
80+
onServiceLaunched()
7581
self.heartbeatTask = Task { [weak self] in
7682
while true {
7783
try Task.checkCancellation()
@@ -85,6 +91,74 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
8591
server.start()
8692
}
8793

94+
public static func createFoldersIfNeeded() throws -> (
95+
applicationSupportURL: URL,
96+
gitHubCopilotURL: URL,
97+
executableURL: URL,
98+
supportURL: URL
99+
) {
100+
let supportURL = FileManager.default.urls(
101+
for: .applicationSupportDirectory,
102+
in: .userDomainMask
103+
).first!.appendingPathComponent(
104+
Bundle.main
105+
.object(forInfoDictionaryKey: "APPLICATION_SUPPORT_FOLDER") as! String
106+
)
107+
108+
if !FileManager.default.fileExists(atPath: supportURL.path) {
109+
try? FileManager.default
110+
.createDirectory(at: supportURL, withIntermediateDirectories: false)
111+
}
112+
let gitHubCopilotFolderURL = supportURL.appendingPathComponent("Codeium")
113+
if !FileManager.default.fileExists(atPath: gitHubCopilotFolderURL.path) {
114+
try? FileManager.default
115+
.createDirectory(at: gitHubCopilotFolderURL, withIntermediateDirectories: false)
116+
}
117+
let supportFolderURL = gitHubCopilotFolderURL.appendingPathComponent("support")
118+
if !FileManager.default.fileExists(atPath: supportFolderURL.path) {
119+
try? FileManager.default
120+
.createDirectory(at: supportFolderURL, withIntermediateDirectories: false)
121+
}
122+
let executableFolderURL = gitHubCopilotFolderURL.appendingPathComponent("executable")
123+
if !FileManager.default.fileExists(atPath: executableFolderURL.path) {
124+
try? FileManager.default
125+
.createDirectory(at: executableFolderURL, withIntermediateDirectories: false)
126+
}
127+
128+
return (supportURL, gitHubCopilotFolderURL, executableFolderURL, supportFolderURL)
129+
}
130+
}
131+
132+
extension CodeiumSuggestionService {
133+
func getMetadata() -> Metadata {
134+
Metadata(
135+
ide_name: "jetbrains",
136+
ide_version: "14.3",
137+
extension_name: "Copilot for Xcode",
138+
extension_version: "14.0.0",
139+
api_key: token,
140+
session_id: CodeiumSuggestionService.sessionId,
141+
request_id: requestCounter
142+
)
143+
}
144+
145+
func getRelativePath(of fileURL: URL) -> String {
146+
let filePath = fileURL.path
147+
let rootPath = projectRootURL.path
148+
if let range = filePath.range(of: rootPath),
149+
range.lowerBound == filePath.startIndex
150+
{
151+
let relativePath = filePath.replacingCharacters(
152+
in: filePath.startIndex..<range.upperBound,
153+
with: ""
154+
)
155+
return relativePath
156+
}
157+
return filePath
158+
}
159+
}
160+
161+
extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
88162
public func getCompletions(
89163
fileURL: URL,
90164
content: String,
@@ -97,20 +171,7 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
97171
requestCounter += 1
98172
let languageId = languageIdentifierFromFileURL(fileURL)
99173

100-
let relativePath = {
101-
let filePath = fileURL.path
102-
let rootPath = projectRootURL.path
103-
if let range = filePath.range(of: rootPath),
104-
range.lowerBound == filePath.startIndex
105-
{
106-
let relativePath = filePath.replacingCharacters(
107-
in: filePath.startIndex..<range.upperBound,
108-
with: ""
109-
)
110-
return relativePath
111-
}
112-
return filePath
113-
}()
174+
let relativePath = getRelativePath(of: fileURL)
114175

115176
let request = CodeiumRequest.GetCompletion(requestBody: .init(
116177
metadata: getMetadata(),
@@ -126,7 +187,17 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
126187
)
127188
),
128189
editor_options: .init(tab_size: indentSize, insert_spaces: !usesTabsForIndentation),
129-
other_documents: []
190+
other_documents: openedDocumentPool.getOtherDocuments(exceptURL: fileURL)
191+
.map { openedDocument in
192+
let languageId = languageIdentifierFromFileURL(openedDocument.url)
193+
return .init(
194+
absolute_path: openedDocument.url.path,
195+
relative_path: openedDocument.relativePath,
196+
text: openedDocument.content,
197+
editor_language: languageId.rawValue,
198+
language: .init(codeLanguage: languageId)
199+
)
200+
}
130201
))
131202

132203
let result = try await server.sendRequest(request)
@@ -156,55 +227,33 @@ public class CodeiumSuggestionService: CodeiumSuggestionServiceType {
156227
} ?? []
157228
}
158229

159-
public static func createFoldersIfNeeded() throws -> (
160-
applicationSupportURL: URL,
161-
gitHubCopilotURL: URL,
162-
executableURL: URL,
163-
supportURL: URL
164-
) {
165-
let supportURL = FileManager.default.urls(
166-
for: .applicationSupportDirectory,
167-
in: .userDomainMask
168-
).first!.appendingPathComponent(
169-
Bundle.main
170-
.object(forInfoDictionaryKey: "APPLICATION_SUPPORT_FOLDER") as! String
171-
)
172-
173-
if !FileManager.default.fileExists(atPath: supportURL.path) {
174-
try? FileManager.default
175-
.createDirectory(at: supportURL, withIntermediateDirectories: false)
176-
}
177-
let gitHubCopilotFolderURL = supportURL.appendingPathComponent("Codeium")
178-
if !FileManager.default.fileExists(atPath: gitHubCopilotFolderURL.path) {
179-
try? FileManager.default
180-
.createDirectory(at: gitHubCopilotFolderURL, withIntermediateDirectories: false)
181-
}
182-
let supportFolderURL = gitHubCopilotFolderURL.appendingPathComponent("support")
183-
if !FileManager.default.fileExists(atPath: supportFolderURL.path) {
184-
try? FileManager.default
185-
.createDirectory(at: supportFolderURL, withIntermediateDirectories: false)
186-
}
187-
let executableFolderURL = gitHubCopilotFolderURL.appendingPathComponent("executable")
188-
if !FileManager.default.fileExists(atPath: executableFolderURL.path) {
189-
try? FileManager.default
190-
.createDirectory(at: executableFolderURL, withIntermediateDirectories: false)
191-
}
230+
public func notifyAccepted(_ suggestion: CodeSuggestion) async {
231+
_ = try? await server.sendRequest(CodeiumRequest.AcceptCompletion(requestBody: .init(
232+
metadata: getMetadata(),
233+
completion_id: suggestion.uuid
234+
)))
235+
}
192236

193-
return (supportURL, gitHubCopilotFolderURL, executableFolderURL, supportFolderURL)
237+
public func notifyOpenTextDocument(fileURL: URL, content: String) async throws {
238+
let relativePath = getRelativePath(of: fileURL)
239+
openedDocumentPool.openDocument(
240+
url: fileURL,
241+
relativePath: relativePath,
242+
content: content
243+
)
194244
}
195-
}
196245

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
246+
public func notifyChangeTextDocument(fileURL: URL, content: String) async throws {
247+
let relativePath = getRelativePath(of: fileURL)
248+
openedDocumentPool.updateDocument(
249+
url: fileURL,
250+
relativePath: relativePath,
251+
content: content
207252
)
208253
}
254+
255+
public func notifyCloseTextDocument(fileURL: URL) async throws {
256+
openedDocumentPool.closeDocument(url: fileURL)
257+
}
209258
}
210259

‎Core/Sources/CodeiumService/OpendDocumentPool.swift‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
import Foundation
2+
3+
private let maxSize: Int = 1_000_000 // Byte
4+
5+
final class OpenedDocumentPool {
6+
var openedDocuments = [URL: OpenedDocument]()
7+
8+
func getOtherDocuments(exceptURL: URL) -> [OpenedDocument] {
9+
let ordered = openedDocuments.values.sorted { $0.updateTime > $1.updateTime }
10+
var documents = [OpenedDocument]()
11+
var size = 0
12+
for document in ordered where document.url != exceptURL {
13+
size += document.size
14+
if size > maxSize {
15+
break
16+
}
17+
documents.append(document)
18+
}
19+
20+
return documents
21+
}
22+
23+
func openDocument(url: URL, relativePath: String, content: String) {
24+
let document = OpenedDocument(url: url, relativePath: relativePath, content: content)
25+
openedDocuments[url] = document
26+
}
27+
28+
func updateDocument(url: URL, relativePath: String, content: String) {
29+
if let document = openedDocuments[url] {
30+
document.update(content: content)
31+
} else {
32+
openDocument(url: url, relativePath: relativePath, content: content)
33+
}
34+
}
35+
36+
func closeDocument(url: URL) {
37+
openedDocuments[url] = nil
38+
}
39+
}
40+
41+
final class OpenedDocument {
42+
var url: URL
43+
var relativePath: String
44+
var updateTime: Date
45+
var content: String
46+
var size: Int
47+
48+
public init(url: URL, relativePath: String, content: String) {
49+
self.url = url
50+
self.relativePath = relativePath
51+
updateTime = Date()
52+
size = content.utf8.count
53+
if size > maxSize {
54+
self.content = ""
55+
} else {
56+
self.content = content
57+
}
58+
}
59+
60+
func update(content: String) {
61+
updateTime = Date()
62+
size = content.utf8.count
63+
if size > maxSize {
64+
self.content = ""
65+
} else {
66+
self.content = content
67+
}
68+
}
69+
}
70+

0 commit comments

Comments
 (0)