Skip to content

Commit c3bae3c

Browse files
intitniintitni
committed
Add CodeiumExtension
1 parent d62284f commit c3bae3c

File tree

10 files changed

+479
-192
lines changed

10 files changed

+479
-192
lines changed

‎Core/Sources/Service/Service.swift‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import BuiltinExtension
2+
import CodeiumService
23
import Combine
34
import Dependencies
45
import Foundation
@@ -43,6 +44,7 @@ public final class Service {
4344

4445
BuiltinExtensionManager.shared.setupExtensions([
4546
GitHubCopilotExtension(workspacePool: workspacePool),
47+
CodeiumExtension(workspacePool: workspacePool),
4648
])
4749
scheduledCleaner = .init()
4850
workspacePool.registerPlugin {

‎Core/Sources/SuggestionService/SuggestionService.swift‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import BuiltinExtension
2+
import CodeiumService
23
import struct CopilotForXcodeKit.WorkspaceInfo
34
import Foundation
45
import GitHubCopilotService
@@ -44,12 +45,10 @@ public actor SuggestionService: SuggestionServiceType {
4445

4546
switch serviceType {
4647
case .builtIn(.codeium):
47-
fatalError()
48-
// let provider = CodeiumSuggestionProvider(
49-
// projectRootURL: projectRootURL,
50-
// onServiceLaunched: onServiceLaunched
51-
// )
52-
// return SuggestionService(provider: provider)
48+
let provider = BuiltinExtensionSuggestionServiceProvider(
49+
extension: CodeiumExtension.self
50+
)
51+
return SuggestionService(provider: provider)
5352
case .builtIn(.gitHubCopilot), .extension:
5453
let provider = BuiltinExtensionSuggestionServiceProvider(
5554
extension: GitHubCopilotExtension.self

‎Tool/Sources/CodeiumService/CodeiumExtension.swift‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
import CopilotForXcodeKit
2+
import Foundation
3+
import Workspace
4+
import Logger
5+
import BuiltinExtension
6+
7+
public final class CodeiumExtension: BuiltinExtension {
8+
public var suggestionService: SuggestionServiceType? { _suggestionService }
9+
public var chatService: ChatServiceType? { nil }
10+
public var promptToCodeService: PromptToCodeServiceType? { nil }
11+
let workspacePool: WorkspacePool
12+
13+
let serviceLocator: ServiceLocator
14+
let _suggestionService: CodeiumSuggestionService
15+
16+
public init(workspacePool: WorkspacePool) {
17+
self.workspacePool = workspacePool
18+
serviceLocator = .init(workspacePool: workspacePool)
19+
_suggestionService = .init(serviceLocator: serviceLocator)
20+
}
21+
22+
public func workspaceDidOpen(_: WorkspaceInfo) {}
23+
24+
public func workspaceDidClose(_: WorkspaceInfo) {}
25+
26+
public func workspace(_ workspace: WorkspaceInfo, didOpenDocumentAt documentURL: URL) {
27+
// check if file size is larger than 15MB, if so, return immediately
28+
if let attrs = try? FileManager.default
29+
.attributesOfItem(atPath: documentURL.path),
30+
let fileSize = attrs[FileAttributeKey.size] as? UInt64,
31+
fileSize > 15 * 1024 * 1024
32+
{ return }
33+
34+
Task {
35+
do {
36+
let content = try String(contentsOf: documentURL, encoding: .utf8)
37+
guard let service = await serviceLocator.getService(from: workspace) else { return }
38+
try await service.notifyOpenTextDocument(fileURL: documentURL, content: content)
39+
} catch {
40+
Logger.gitHubCopilot.error(error.localizedDescription)
41+
}
42+
}
43+
}
44+
45+
public func workspace(_ workspace: WorkspaceInfo, didSaveDocumentAt documentURL: URL) {
46+
// unimplemented
47+
}
48+
49+
public func workspace(_ workspace: WorkspaceInfo, didCloseDocumentAt documentURL: URL) {
50+
Task {
51+
do {
52+
guard let service = await serviceLocator.getService(from: workspace) else { return }
53+
try await service.notifyCloseTextDocument(fileURL: documentURL)
54+
} catch {
55+
Logger.gitHubCopilot.error(error.localizedDescription)
56+
}
57+
}
58+
}
59+
60+
public func workspace(
61+
_ workspace: WorkspaceInfo,
62+
didUpdateDocumentAt documentURL: URL,
63+
content: String
64+
) {
65+
// check if file size is larger than 15MB, if so, return immediately
66+
if let attrs = try? FileManager.default
67+
.attributesOfItem(atPath: documentURL.path),
68+
let fileSize = attrs[FileAttributeKey.size] as? UInt64,
69+
fileSize > 15 * 1024 * 1024
70+
{ return }
71+
72+
Task {
73+
do {
74+
let content = try String(contentsOf: documentURL, encoding: .utf8)
75+
guard let service = await serviceLocator.getService(from: workspace) else { return }
76+
try await service.notifyOpenTextDocument(fileURL: documentURL, content: content)
77+
} catch {
78+
Logger.gitHubCopilot.error(error.localizedDescription)
79+
}
80+
}
81+
}
82+
83+
public func appConfigurationDidChange(_ configuration: AppConfiguration) {
84+
if !configuration.chatServiceInUse && !configuration.suggestionServiceInUse {
85+
for workspace in workspacePool.workspaces.values {
86+
guard let plugin = workspace.plugin(for: CodeiumWorkspacePlugin.self)
87+
else { continue }
88+
plugin.terminate()
89+
}
90+
}
91+
}
92+
93+
public func terminate() {
94+
for workspace in workspacePool.workspaces.values {
95+
guard let plugin = workspace.plugin(for: CodeiumWorkspacePlugin.self)
96+
else { continue }
97+
plugin.terminate()
98+
}
99+
}
100+
}
101+
102+
final class ServiceLocator {
103+
let workspacePool: WorkspacePool
104+
105+
init(workspacePool: WorkspacePool) {
106+
self.workspacePool = workspacePool
107+
}
108+
109+
func getService(from workspace: WorkspaceInfo) async -> CodeiumService? {
110+
guard let workspace = workspacePool.workspaces[workspace.workspaceURL],
111+
let plugin = workspace.plugin(for: CodeiumWorkspacePlugin.self)
112+
else { return nil }
113+
return plugin.codeiumService
114+
}
115+
}
116+

‎Tool/Sources/CodeiumService/CodeiumInstallationManager.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ public struct CodeiumInstallationManager {
1515
}
1616

1717
public func checkInstallation() -> InstallationStatus {
18-
guard let urls = try? CodeiumSuggestionService.createFoldersIfNeeded()
18+
guard let urls = try? CodeiumService.createFoldersIfNeeded()
1919
else { return .notInstalled }
2020
let executableFolderURL = urls.executableURL
2121
let binaryURL = executableFolderURL.appendingPathComponent("language_server")
@@ -60,7 +60,7 @@ public struct CodeiumInstallationManager {
6060
defer { CodeiumInstallationManager.isInstalling = false }
6161
do {
6262
continuation.yield(.downloading)
63-
let urls = try CodeiumSuggestionService.createFoldersIfNeeded()
63+
let urls = try CodeiumService.createFoldersIfNeeded()
6464
let urlString =
6565
"https://github.com/Exafunction/codeium/releases/download/language-server-v\(Self.latestSupportedVersion)/language_server_macos_\(isAppleSilicon() ? "arm" : "x64").gz"
6666
guard let url = URL(string: urlString) else { return }
@@ -108,7 +108,7 @@ public struct CodeiumInstallationManager {
108108
}
109109

110110
public func uninstall() async throws {
111-
guard let urls = try? CodeiumSuggestionService.createFoldersIfNeeded()
111+
guard let urls = try? CodeiumService.createFoldersIfNeeded()
112112
else { return }
113113
let executableFolderURL = urls.executableURL
114114
let binaryURL = executableFolderURL.appendingPathComponent("language_server")

‎Tool/Sources/CodeiumService/CodeiumService.swift‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ enum CodeiumError: Error, LocalizedError {
3939
}
4040
}
4141

42-
public class CodeiumSuggestionService {
42+
public class CodeiumService {
4343
static let sessionId = UUID().uuidString
4444
let projectRootURL: URL
4545
var server: CodeiumLSP?
@@ -70,7 +70,7 @@ public class CodeiumSuggestionService {
7070
public init(projectRootURL: URL, onServiceLaunched: @escaping () -> Void) throws {
7171
self.projectRootURL = projectRootURL
7272
self.onServiceLaunched = onServiceLaunched
73-
let urls = try CodeiumSuggestionService.createFoldersIfNeeded()
73+
let urls = try CodeiumService.createFoldersIfNeeded()
7474
languageServerURL = urls.executableURL.appendingPathComponent("language_server")
7575
supportURL = urls.supportURL
7676
Task {
@@ -177,7 +177,7 @@ public class CodeiumSuggestionService {
177177
}
178178
}
179179

180-
extension CodeiumSuggestionService {
180+
extension CodeiumService {
181181
func getMetadata() async throws -> Metadata {
182182
guard let key = authService.key else {
183183
struct E: Error, LocalizedError {
@@ -198,7 +198,7 @@ extension CodeiumSuggestionService {
198198
ide_version: ideVersion,
199199
extension_version: languageServerVersion,
200200
api_key: key,
201-
session_id: CodeiumSuggestionService.sessionId,
201+
session_id: CodeiumService.sessionId,
202202
request_id: requestCounter
203203
)
204204
}
@@ -219,7 +219,7 @@ extension CodeiumSuggestionService {
219219
}
220220
}
221221

222-
extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
222+
extension CodeiumService: CodeiumSuggestionServiceType {
223223
public func getCompletions(
224224
fileURL: URL,
225225
content: String,
@@ -298,7 +298,7 @@ extension CodeiumSuggestionService: CodeiumSuggestionServiceType {
298298
_ = try? await server?.sendRequest(
299299
CodeiumRequest.CancelRequest(requestBody: .init(
300300
request_id: requestCounter,
301-
session_id: CodeiumSuggestionService.sessionId
301+
session_id: CodeiumService.sessionId
302302
))
303303
)
304304
}

‎Tool/Sources/CodeiumService/CodeiumSuggestionService.swift‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
import CopilotForXcodeKit
2+
import Foundation
3+
import SuggestionModel
4+
import Workspace
5+
6+
class CodeiumSuggestionService: SuggestionServiceType {
7+
var configuration: SuggestionServiceConfiguration {
8+
.init(
9+
acceptsRelevantCodeSnippets: true,
10+
mixRelevantCodeSnippetsInSource: true,
11+
acceptsRelevantSnippetsFromOpenedFiles: false
12+
)
13+
}
14+
15+
let serviceLocator: ServiceLocator
16+
17+
init(serviceLocator: ServiceLocator) {
18+
self.serviceLocator = serviceLocator
19+
}
20+
21+
func getSuggestions(
22+
_ request: SuggestionRequest,
23+
workspace: WorkspaceInfo
24+
) async throws -> [CopilotForXcodeKit.CodeSuggestion] {
25+
guard let service = await serviceLocator.getService(from: workspace) else { return [] }
26+
return try await service.getCompletions(
27+
fileURL: request.fileURL,
28+
content: request.content,
29+
cursorPosition: .init(
30+
line: request.cursorPosition.line,
31+
character: request.cursorPosition.character
32+
),
33+
tabSize: request.tabSize,
34+
indentSize: request.indentSize,
35+
usesTabsForIndentation: request.usesTabsForIndentation
36+
).map(Self.convert)
37+
}
38+
39+
func notifyAccepted(
40+
_ suggestion: CopilotForXcodeKit.CodeSuggestion,
41+
workspace: WorkspaceInfo
42+
) async {
43+
guard let service = await serviceLocator.getService(from: workspace) else { return }
44+
await service.notifyAccepted(Self.convert(suggestion))
45+
}
46+
47+
func notifyRejected(
48+
_ suggestions: [CopilotForXcodeKit.CodeSuggestion],
49+
workspace: WorkspaceInfo
50+
) async {
51+
// unimplemented
52+
}
53+
54+
func cancelRequest(workspace: WorkspaceInfo) async {
55+
guard let service = await serviceLocator.getService(from: workspace) else { return }
56+
await service.cancelRequest()
57+
}
58+
59+
static func convert(
60+
_ suggestion: SuggestionModel.CodeSuggestion
61+
) -> CopilotForXcodeKit.CodeSuggestion {
62+
.init(
63+
id: suggestion.id,
64+
text: suggestion.text,
65+
position: .init(
66+
line: suggestion.position.line,
67+
character: suggestion.position.character
68+
),
69+
range: .init(
70+
start: .init(
71+
line: suggestion.range.start.line,
72+
character: suggestion.range.start.character
73+
),
74+
end: .init(
75+
line: suggestion.range.end.line,
76+
character: suggestion.range.end.character
77+
)
78+
)
79+
)
80+
}
81+
82+
static func convert(
83+
_ suggestion: CopilotForXcodeKit.CodeSuggestion
84+
) -> SuggestionModel.CodeSuggestion {
85+
.init(
86+
id: suggestion.id,
87+
text: suggestion.text,
88+
position: .init(
89+
line: suggestion.position.line,
90+
character: suggestion.position.character
91+
),
92+
range: .init(
93+
start: .init(
94+
line: suggestion.range.start.line,
95+
character: suggestion.range.start.character
96+
),
97+
end: .init(
98+
line: suggestion.range.end.line,
99+
character: suggestion.range.end.character
100+
)
101+
)
102+
)
103+
}
104+
}
105+

‎Tool/Sources/CodeiumService/CodeiumWorkspacePlugin.swift‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import Foundation
2+
import Logger
3+
import Workspace
4+
5+
public final class CodeiumWorkspacePlugin: WorkspacePlugin {
6+
var _codeiumService: CodeiumService?
7+
var codeiumService: CodeiumService? {
8+
if let service = _codeiumService { return service }
9+
do {
10+
return try createCodeiumService()
11+
} catch {
12+
Logger.codeium.error("Failed to create Codeium service: \(error)")
13+
return nil
14+
}
15+
}
16+
17+
deinit {
18+
if let codeiumService {
19+
codeiumService.terminate()
20+
}
21+
}
22+
23+
func createCodeiumService() throws -> CodeiumService {
24+
let newService = try CodeiumService(
25+
projectRootURL: projectRootURL,
26+
onServiceLaunched: {
27+
[weak self] in
28+
self?.finishLaunchingService()
29+
}
30+
)
31+
_codeiumService = newService
32+
return newService
33+
}
34+
35+
func finishLaunchingService() {
36+
guard let workspace, let _codeiumService else { return }
37+
Task {
38+
for (_, filespace) in workspace.filespaces {
39+
let documentURL = filespace.fileURL
40+
guard let content = try? String(contentsOf: documentURL) else { continue }
41+
try? await _codeiumService.notifyOpenTextDocument(
42+
fileURL: documentURL,
43+
content: content
44+
)
45+
}
46+
}
47+
}
48+
49+
func terminate() {
50+
_codeiumService = nil
51+
}
52+
}
53+

0 commit comments

Comments
 (0)