Skip to content

Commit 317b09b

Browse files
intitniintitni
committed
Add new Codeium requests
1 parent d556b91 commit 317b09b

File tree

4 files changed

+148
-14
lines changed

4 files changed

+148
-14
lines changed

‎Tool/Sources/CodeiumService/CodeiumExtension.swift‎

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,29 @@ public final class CodeiumExtension: BuiltinExtension {
4343
suggestionService = .init(serviceLocator: serviceLocator)
4444
}
4545

46-
public func workspaceDidOpen(_: WorkspaceInfo) {}
46+
public func workspaceDidOpen(_ workspace: WorkspaceInfo) {
47+
Task {
48+
do {
49+
guard await isLanguageServerInUse else { return }
50+
guard let service = await serviceLocator.getService(from: workspace) else { return }
51+
try await service.notifyOpenWorkspace(workspaceURL: workspace.workspaceURL)
52+
} catch {
53+
Logger.codeium.error(error.localizedDescription)
54+
}
55+
}
56+
}
4757

48-
public func workspaceDidClose(_: WorkspaceInfo) {}
58+
public func workspaceDidClose(_ workspace: WorkspaceInfo) {
59+
Task {
60+
do {
61+
guard await isLanguageServerInUse else { return }
62+
guard let service = await serviceLocator.getService(from: workspace) else { return }
63+
try await service.notifyCloseWorkspace(workspaceURL: workspace.workspaceURL)
64+
} catch {
65+
Logger.codeium.error(error.localizedDescription)
66+
}
67+
}
68+
}
4969

5070
public func workspace(_ workspace: WorkspaceInfo, didOpenDocumentAt documentURL: URL) {
5171
Task {
@@ -62,7 +82,7 @@ public final class CodeiumExtension: BuiltinExtension {
6282
guard let service = await serviceLocator.getService(from: workspace) else { return }
6383
try await service.notifyOpenTextDocument(fileURL: documentURL, content: content)
6484
} catch {
65-
Logger.gitHubCopilot.error(error.localizedDescription)
85+
Logger.codeium.error(error.localizedDescription)
6686
}
6787
}
6888
}
@@ -78,7 +98,7 @@ public final class CodeiumExtension: BuiltinExtension {
7898
guard let service = await serviceLocator.getService(from: workspace) else { return }
7999
try await service.notifyCloseTextDocument(fileURL: documentURL)
80100
} catch {
81-
Logger.gitHubCopilot.error(error.localizedDescription)
101+
Logger.codeium.error(error.localizedDescription)
82102
}
83103
}
84104
}
@@ -100,8 +120,15 @@ public final class CodeiumExtension: BuiltinExtension {
100120
guard let content else { return }
101121
guard let service = await serviceLocator.getService(from: workspace) else { return }
102122
try await service.notifyChangeTextDocument(fileURL: documentURL, content: content)
123+
try await service.refreshIDEContext(
124+
fileURL: documentURL,
125+
content: content,
126+
cursorPosition: .zero,
127+
tabSize: 4, indentSize: 4, usesTabsForIndentation: false,
128+
workspaceURL: workspace.workspaceURL
129+
)
103130
} catch {
104-
Logger.gitHubCopilot.error(error.localizedDescription)
131+
Logger.codeium.error(error.localizedDescription)
105132
}
106133
}
107134
}

‎Tool/Sources/CodeiumService/CodeiumWorkspacePlugin.swift‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ public final class CodeiumWorkspacePlugin: WorkspacePlugin {
4141
func finishLaunchingService() {
4242
guard let workspace, let _codeiumService else { return }
4343
Task {
44+
try? await _codeiumService.notifyOpenWorkspace(workspaceURL: workspaceURL)
45+
4446
for (_, filespace) in workspace.filespaces {
4547
let documentURL = filespace.fileURL
4648
guard let content = try? String(contentsOf: documentURL) else { continue }

‎Tool/Sources/CodeiumService/LanguageServer/CodeiumRequest.swift‎

Lines changed: 67 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -33,17 +33,17 @@ enum CodeiumRequest {
3333
var chatWebServerPort: UInt32
3434
var chatClientPort: UInt32
3535
}
36-
36+
3737
struct RequestBody: Codable {}
38-
38+
3939
var requestBody: RequestBody
40-
40+
4141
func makeURLRequest(server: String) -> URLRequest {
4242
let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
4343
return assembleURLRequest(server: server, method: "GetProcesses", body: data)
4444
}
4545
}
46-
46+
4747
struct GetCompletion: CodeiumRequestType {
4848
struct Response: Codable {
4949
var state: State
@@ -64,7 +64,7 @@ enum CodeiumRequest {
6464
return assembleURLRequest(server: server, method: "GetCompletions", body: data)
6565
}
6666
}
67-
67+
6868
struct CancelRequest: CodeiumRequestType {
6969
struct Response: Codable {}
7070

@@ -80,7 +80,7 @@ enum CodeiumRequest {
8080
return assembleURLRequest(server: server, method: "CancelRequest", body: data)
8181
}
8282
}
83-
83+
8484
struct AcceptCompletion: CodeiumRequestType {
8585
struct Response: Codable {}
8686

@@ -96,7 +96,67 @@ enum CodeiumRequest {
9696
return assembleURLRequest(server: server, method: "AcceptCompletion", body: data)
9797
}
9898
}
99-
99+
100+
struct RemoveTrackedWorkspace: CodeiumRequestType {
101+
struct Response: Codable {}
102+
103+
struct RequestBody: Codable {
104+
var workspace: String
105+
}
106+
107+
var requestBody: RequestBody
108+
109+
func makeURLRequest(server: String) -> URLRequest {
110+
let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
111+
return assembleURLRequest(
112+
server: server,
113+
method: "RemoveTrackedWorkspace",
114+
body: data
115+
)
116+
}
117+
}
118+
119+
struct AddTrackedWorkspace: CodeiumRequestType {
120+
struct Response: Codable {}
121+
122+
struct RequestBody: Codable {
123+
var workspace: String
124+
}
125+
126+
var requestBody: RequestBody
127+
128+
func makeURLRequest(server: String) -> URLRequest {
129+
let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
130+
return assembleURLRequest(
131+
server: server,
132+
method: "AddTrackedWorkspace",
133+
body: data
134+
)
135+
}
136+
}
137+
138+
struct RefreshContextForIdeAction: CodeiumRequestType {
139+
struct Response: Codable {}
140+
141+
struct RequestBody: Codable {
142+
var active_document: CodeiumDocument
143+
var open_document_filepaths: [String]
144+
var workspace_paths: [String]
145+
var blocking: Bool = false
146+
}
147+
148+
var requestBody: RequestBody
149+
150+
func makeURLRequest(server: String) -> URLRequest {
151+
let data = (try? JSONEncoder().encode(requestBody)) ?? Data()
152+
return assembleURLRequest(
153+
server: server,
154+
method: "RefreshContextForIdeAction",
155+
body: data
156+
)
157+
}
158+
}
159+
100160
struct Heartbeat: CodeiumRequestType {
101161
struct Response: Codable {}
102162

‎Tool/Sources/CodeiumService/Services/CodeiumService.swift‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -247,7 +247,7 @@ extension CodeiumService: CodeiumSuggestionServiceType {
247247

248248
requestCounter += 1
249249
let languageId = languageIdentifierFromFileURL(fileURL)
250-
let relativePath = getRelativePath(of: fileURL)
250+
let relativePath = getRelativePath(of: fileURL)
251251

252252
let task = Task {
253253
let request = try await CodeiumRequest.GetCompletion(requestBody: .init(
@@ -340,6 +340,7 @@ extension CodeiumService: CodeiumSuggestionServiceType {
340340
URLQueryItem(name: "ide_name", value: metadata.ide_name),
341341
URLQueryItem(name: "ide_version", value: metadata.ide_version),
342342
URLQueryItem(name: "web_server_url", value: webServerUrl),
343+
URLQueryItem(name: "ide_telemetry_enabled", value: "true")
343344
]
344345

345346
if let url = components.url {
@@ -356,7 +357,7 @@ extension CodeiumService: CodeiumSuggestionServiceType {
356357
metadata: getMetadata(),
357358
completion_id: suggestion.id
358359
)))
359-
}
360+
}
360361

361362
public func notifyOpenTextDocument(fileURL: URL, content: String) async throws {
362363
let relativePath = getRelativePath(of: fileURL)
@@ -380,6 +381,50 @@ extension CodeiumService: CodeiumSuggestionServiceType {
380381
await openedDocumentPool.closeDocument(url: fileURL)
381382
}
382383

384+
public func notifyOpenWorkspace(workspaceURL: URL) async throws {
385+
_ = try await (try setupServerIfNeeded()).sendRequest(
386+
CodeiumRequest
387+
.AddTrackedWorkspace(requestBody: .init(workspace: workspaceURL.path))
388+
)
389+
}
390+
391+
public func notifyCloseWorkspace(workspaceURL: URL) async throws {
392+
_ = try await (try setupServerIfNeeded()).sendRequest(
393+
CodeiumRequest
394+
.RemoveTrackedWorkspace(requestBody: .init(workspace: workspaceURL.path))
395+
)
396+
}
397+
398+
public func refreshIDEContext(
399+
fileURL: URL,
400+
content: String,
401+
cursorPosition: CursorPosition,
402+
tabSize: Int,
403+
indentSize: Int,
404+
usesTabsForIndentation: Bool,
405+
workspaceURL: URL
406+
) async throws {
407+
let languageId = languageIdentifierFromFileURL(fileURL)
408+
let relativePath = getRelativePath(of: fileURL)
409+
let request = await CodeiumRequest.RefreshContextForIdeAction(requestBody: .init(
410+
active_document: .init(
411+
absolute_path: fileURL.path,
412+
relative_path: relativePath,
413+
text: content,
414+
editor_language: languageId.rawValue,
415+
language: .init(codeLanguage: languageId),
416+
cursor_position: .init(
417+
row: cursorPosition.line,
418+
col: cursorPosition.character
419+
)
420+
),
421+
open_document_filepaths: openedDocumentPool.getOtherDocuments(exceptURL: fileURL)
422+
.map(\.url.path),
423+
workspace_paths: [workspaceURL.path]
424+
))
425+
_ = try await (try setupServerIfNeeded()).sendRequest(request)
426+
}
427+
383428
public func terminate() {
384429
server?.terminate()
385430
server = nil

0 commit comments

Comments
 (0)