Added Indexing Changes for Codeium

Author: RoshanNagaram-eng

Core/Sources/HostApp/AccountSettings/CodeiumView.swift

@@ -14,6 +14,7 @@ struct CodeiumView: View {
         @AppStorage(\.codeiumEnterpriseMode) var codeiumEnterpriseMode
         @AppStorage(\.codeiumPortalUrl) var codeiumPortalUrl
         @AppStorage(\.codeiumApiUrl) var codeiumApiUrl
+        @AppStorage(\.indexEnabled) var indexEnabled
 
         init() {
             isSignedIn = codeiumAuthService.isSignedIn
@@ -205,6 +206,12 @@ struct CodeiumView: View {
                     }
                 }
             }
+            
+            SubSection(title: Text("Indexing")) {
+                Form {
+                    Toggle("Enable Indexing", isOn: $viewModel.indexEnabled)
+                }
+            }
 
             SubSection(title: Text("Enterprise")) {
                 Form {

Tool/Sources/CodeiumService/CodeiumWorkspacePlugin.swift

@@ -26,8 +26,7 @@ public final class CodeiumWorkspacePlugin: WorkspacePlugin {
         let newService = try CodeiumService(
             projectRootURL: projectRootURL,
             onServiceLaunched: {
-                [weak self] in
-                self?.finishLaunchingService()
+
             },
             onServiceTerminated: {
                 // start handled in the service.

Tool/Sources/CodeiumService/LanguageServer/CodeiumInstallationManager.swift

@@ -20,6 +20,15 @@ public struct CodeiumInstallationManager {
             }
         }
     }
+    
+    public func getLatestSupportedVersion() -> String {
+        if isEnterprise{
+            return UserDefaults.shared.value(for: \.codeiumEnterpriseVersion);
+        }
+        
+        return Self.latestSupportedVersion;
+        
+    }
 
     func getEnterprisePortalVersion() async throws -> String {
         let enterprisePortalUrl = UserDefaults.shared.value(for: \.codeiumPortalUrl)

Tool/Sources/CodeiumService/LanguageServer/CodeiumLanguageServer.swift

@@ -4,9 +4,11 @@ import LanguageClient
 import LanguageServerProtocol
 import Logger
 import Preferences
+import XcodeInspector
 
 protocol CodeiumLSP {
     func sendRequest<E: CodeiumRequestType>(_ endpoint: E) async throws -> E.Response
+    func updateIndexing() async
     func terminate()
 }
 
@@ -20,6 +22,7 @@ final class CodeiumLanguageServer {
     var launchHandler: (() -> Void)?
     var port: String?
     var heartbeatTask: Task<Void, Error>?
+    var project_paths: [String]
 
     init(
         languageServerExecutableURL: URL,
@@ -33,6 +36,7 @@ final class CodeiumLanguageServer {
         self.supportURL = supportURL
         self.terminationHandler = terminationHandler
         self.launchHandler = launchHandler
+        self.project_paths = []
         process = Process()
         transport = IOTransport()
 
@@ -60,6 +64,18 @@ final class CodeiumLanguageServer {
         if isEnterpriseMode {
             process.arguments?.append("--enterprise_mode")
         }
+        
+        let indexEnabled = UserDefaults.shared.value(for: \.indexEnabled)
+        if indexEnabled {
+            let indexingMaxFileSize = UserDefaults.shared.value(for: \.indexingMaxFileSize)
+            if (indexEnabled) {
+              process.arguments?.append("--enable_local_search")
+              process.arguments?.append("--enable_index_service")
+              process.arguments?.append("--search_max_workspace_file_count")
+              process.arguments?.append("\(indexingMaxFileSize)")
+              Logger.codeium.info("Indexing Enabled")
+            }
+        }
 
         process.currentDirectoryURL = supportURL
 
@@ -177,6 +193,38 @@ extension CodeiumLanguageServer: CodeiumLSP {
             }
         }
     }
+    
+    func updateIndexing() async {
+        let indexEnabled = UserDefaults.shared.value(for: \.indexEnabled)
+        if !indexEnabled {
+            return
+        }
+        
+
+        let curr_proj_paths = await getProjectPaths()
+        
+        // Add all workspaces that are in the curr_proj_paths but not in the previous project paths
+        for curr_proj_path in curr_proj_paths {
+            if !self.project_paths.contains(curr_proj_path) && FileManager.default.fileExists(atPath: curr_proj_path) {
+                _ = try? await self.sendRequest(CodeiumRequest.AddTrackedWorkspace(requestBody: .init(
+                    workspace: curr_proj_path
+                )))
+            }
+        }
+        
+        // Remove all workspaces that are in previous project paths but not in the curr_proj_paths
+        for proj_path in self.project_paths {
+            if !curr_proj_paths.contains(proj_path) && FileManager.default.fileExists(atPath: proj_path) {
+                _ = try? await self.sendRequest(CodeiumRequest.RemoveTrackedWorkspace(requestBody: .init(
+                    workspace: proj_path
+                )))
+            }
+        }
+        // These should be identical now
+        self.project_paths = curr_proj_paths
+        
+    }
+
 }
 
 final class IOTransport {
@@ -272,3 +320,70 @@ final class IOTransport {
     }
 }
 
+class WorkspaceParser: NSObject, XMLParserDelegate {
+    var projectPaths: [String] = []
+    var workspaceFileURL: URL
+    var workspaceBaseURL: URL
+
+    init(workspaceFileURL: URL, workspaceBaseURL: URL) {
+        self.workspaceFileURL = workspaceFileURL
+        self.workspaceBaseURL = workspaceBaseURL
+    }
+    
+    func parse() -> [String] {
+        guard let parser = XMLParser(contentsOf: workspaceFileURL) else {
+            print("Failed to create XML parser for file: \(workspaceFileURL.path)")
+            return []
+        }
+        parser.delegate = self
+        parser.parse()
+        return projectPaths
+    }
+    
+    // XMLParserDelegate methods
+    func parser(_ parser: XMLParser, didStartElement elementName: String, namespaceURI: String?, qualifiedName qName: String?, attributes attributeDict: [String : String]) {
+        if elementName == "FileRef", let location = attributeDict["location"] {
+            var project_path: String
+            if location.starts(with: "group:") && pathEndsWithXcodeproj(location) {
+                let curr_path = String(location.dropFirst("group:".count))
+                guard let relative_project_url = URL(string: curr_path) else {
+                    return
+                }
+                let relative_base_path = relative_project_url.deletingLastPathComponent()
+                project_path = (self.workspaceBaseURL.appendingPathComponent(relative_base_path.relativePath)).standardized.path
+            } else if location.starts(with: "absolute:") && pathEndsWithXcodeproj(location){
+                let abs_url = URL(fileURLWithPath: String(location.dropFirst("absolute:".count)))
+                project_path = abs_url.deletingLastPathComponent().standardized.path
+            } else {
+                return
+            }
+            if FileManager.default.fileExists(atPath: project_path) {
+                
+                projectPaths.append(project_path)
+            }
+        }
+    }
+        
+    func parser(_ parser: XMLParser, parseErrorOccurred parseError: Error) {
+        print("Failed to parse XML: \(parseError.localizedDescription)")
+    }
+    
+    func pathEndsWithXcodeproj(_ path: String) -> Bool {
+        return path.hasSuffix(".xcodeproj")
+    }
+    
+}
+
+public func getProjectPaths() async -> [String] {
+    guard let workspaceURL = await XcodeInspector.shared.safe.realtimeActiveWorkspaceURL else {
+        return []
+    }
+
+    let workspacebaseURL = workspaceURL.deletingLastPathComponent()
+    
+    let workspaceContentsURL = workspaceURL.appendingPathComponent("contents.xcworkspacedata")
+    
+    let parser = WorkspaceParser(workspaceFileURL: workspaceContentsURL, workspaceBaseURL: workspacebaseURL)
+    let absolutePaths = parser.parse()
+    return absolutePaths
+}

Tool/Sources/CodeiumService/Services/CodeiumService.swift

@@ -49,6 +49,7 @@ public class CodeiumService {
     let projectRootURL: URL
     var server: CodeiumLSP?
     var heartbeatTask: Task<Void, Error>?
+    var workspaceTask: Task<Void, Error>?
     var requestCounter: UInt64 = 0
     var cancellationCounter: UInt64 = 0
     let openedDocumentPool = OpenedDocumentPool()
@@ -127,6 +128,7 @@ public class CodeiumService {
         server.terminationHandler = { [weak self] in
             self?.server = nil
             self?.heartbeatTask?.cancel()
+            self?.workspaceTask?.cancel()
             self?.requestCounter = 0
             self?.cancellationCounter = 0
             self?.onServiceTerminated()
@@ -145,6 +147,14 @@ public class CodeiumService {
                     try await Task.sleep(nanoseconds: 5_000_000_000)
                 }
             }
+            
+            self.workspaceTask = Task { [weak self] in
+                while true {
+                    try Task.checkCancellation()
+                    _ = await self?.server?.updateIndexing()
+                    try await Task.sleep(nanoseconds: 5_000_000_000)
+                }
+            }
         }
 
         self.server = server

Tool/Sources/Preferences/Keys.swift

@@ -219,6 +219,14 @@ public extension UserDefaultPreferenceKeys {
     var codeiumEnterpriseVersion: PreferenceKey<String> {
         .init(defaultValue: "", key: "CodeiumEnterpriseVersion")
     }
+
+    var indexEnabled: PreferenceKey<Bool> {
+        .init(defaultValue: false, key: "indexEnabled")
+    }
+    
+    var indexingMaxFileSize: PreferenceKey<Int> {
+        .init(defaultValue: 5000, key: "indexingMaxFileSize")
+    }
 
 }