Adjust service

Author: intitni

Core/Sources/Service/RealtimeSuggestionController.swift

@@ -163,10 +163,10 @@ public actor RealtimeSuggestionController {
 
     func cancelInFlightTasks(excluding: Task<Void, Never>? = nil) async {
         inflightPrefetchTask?.cancel()
-
+        let workspaces = await Service.shared.workspacePool.workspaces
         // cancel in-flight tasks
         await withTaskGroup(of: Void.self) { group in
-            for (_, workspace) in Service.shared.workspacePool.workspaces {
+            for (_, workspace) in workspaces {
                 group.addTask {
                     await workspace.cancelInFlightRealtimeSuggestionRequests()
                 }

Core/Sources/Service/Service.swift

@@ -25,12 +25,14 @@ import ProService
 
 /// The running extension service.
 public final class Service {
+    @MainActor
     public static let shared = Service()
 
-    let workspacePool = WorkspacePool()
+    @Dependency(\.workspacePool) var workspacePool
     @MainActor
-    public let guiController = GraphicalUserInterfaceController()
-    public let realtimeSuggestionController = RealtimeSuggestionController()
+    public let guiController: GraphicalUserInterfaceController
+    public let commandHandler: CommandHandler
+    public let realtimeSuggestionController: RealtimeSuggestionController
     public let scheduledCleaner: ScheduledCleaner
     let globalShortcutManager: GlobalShortcutManager
     let keyBindingManager: KeyBindingManager
@@ -43,15 +45,47 @@ public final class Service {
     @Dependency(\.toast) var toast
     var cancellable = Set<AnyCancellable>()
 
+    @MainActor
     private init() {
-        WorkspacePoolDependencyKey.liveValue = workspacePool
-        CommandHandlerDependencyKey.liveValue = PseudoCommandHandler()
+        @Dependency(\.workspacePool) var workspacePool
+        let commandHandler = PseudoCommandHandler()
+        self.commandHandler = commandHandler
+
+        func setup<R>(_ operation: () -> R) -> R {
+            withDependencies { values in
+                values.commandHandler = commandHandler
+            } operation: {
+                operation() //
+            }
+        }
+
+        realtimeSuggestionController = setup { .init() }
+        scheduledCleaner = setup { .init() }
+        let guiController = setup { GraphicalUserInterfaceController() }
+        self.guiController = guiController
+        globalShortcutManager = setup { .init(guiController: guiController) }
+
+        keyBindingManager = setup {
+            .init(
+                workspacePool: workspacePool,
+                acceptSuggestion: {
+                    Task { await PseudoCommandHandler().acceptSuggestion() }
+                },
+                dismissSuggestion: {
+                    Task { await PseudoCommandHandler().dismissSuggestion() }
+                }
+            )
+        }
+
+        #if canImport(ProService)
+        proService = setup { ProService() }
+        #endif
 
         BuiltinExtensionManager.shared.setupExtensions([
             GitHubCopilotExtension(workspacePool: workspacePool),
             CodeiumExtension(workspacePool: workspacePool),
         ])
-        scheduledCleaner = .init()
+
         workspacePool.registerPlugin {
             SuggestionServiceWorkspacePlugin(workspace: $0) { SuggestionService.service() }
         }
@@ -64,21 +98,6 @@ public final class Service {
         workspacePool.registerPlugin {
             BuiltinExtensionWorkspacePlugin(workspace: $0)
         }
-        
-        globalShortcutManager = .init(guiController: guiController)
-        keyBindingManager = .init(
-            workspacePool: workspacePool,
-            acceptSuggestion: {
-                Task { await PseudoCommandHandler().acceptSuggestion() }
-            },
-            dismissSuggestion: {
-                Task { await PseudoCommandHandler().dismissSuggestion() }
-            }
-        )
-
-        #if canImport(ProService)
-        proService = ProService()
-        #endif
 
         scheduledCleaner.service = self
     }
@@ -101,9 +120,10 @@ public final class Service {
                 .removeDuplicates()
                 .filter { $0 != .init(fileURLWithPath: "/") }
                 .compactMap { $0 }
-                .sink { [weak self] fileURL in
+                .sink { fileURL in
                     Task {
-                        try await self?.workspacePool
+                        @Dependency(\.workspacePool) var workspacePool
+                        return try await workspacePool
                             .fetchOrCreateWorkspaceAndFilespace(fileURL: fileURL)
                     }
                 }.store(in: &cancellable)

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -326,18 +326,17 @@ struct PseudoCommandHandler: CommandHandler {
 
     func dismissSuggestion() async {
         guard let documentURL = await XcodeInspector.shared.safe.activeDocumentURL else { return }
+        PresentInWindowSuggestionPresenter().discardSuggestion(fileURL: documentURL)
         guard let (_, filespace) = try? await Service.shared.workspacePool
             .fetchOrCreateWorkspaceAndFilespace(fileURL: documentURL) else { return }
-
         await filespace.reset()
-        PresentInWindowSuggestionPresenter().discardSuggestion(fileURL: documentURL)
     }
 
     func openChat(forceDetach: Bool) {
         switch UserDefaults.shared.value(for: \.openChatMode) {
         case .chatPanel:
-            let store = Service.shared.guiController.store
             Task { @MainActor in
+                let store = Service.shared.guiController.store
                 await store.send(.createAndSwitchToChatGPTChatTabIfNeeded).finish()
                 store.send(.openChatPanel(forceDetach: forceDetach))
             }
@@ -360,8 +359,8 @@ struct PseudoCommandHandler: CommandHandler {
 
             if openInApp {
                 #if canImport(BrowserChatTab)
-                let store = Service.shared.guiController.store
                 Task { @MainActor in
+                    let store = Service.shared.guiController.store
                     await store.send(.createAndSwitchToChatTabIfNeededMatching(
                         check: {
                             func match(_ tabURL: URL?) -> Bool {
@@ -386,8 +385,8 @@ struct PseudoCommandHandler: CommandHandler {
                 }
             }
         case .codeiumChat:
-            let store = Service.shared.guiController.store
             Task { @MainActor in
+                let store = Service.shared.guiController.store
                 await store.send(
                     .createAndSwitchToChatTabIfNeededMatching(
                         check: { $0 is CodeiumChatTab },
@@ -400,7 +399,7 @@ struct PseudoCommandHandler: CommandHandler {
     }
 
     func sendChatMessage(_ message: String) async {
-        let store = Service.shared.guiController.store
+        let store = await Service.shared.guiController.store
         await store.send(.sendCustomCommandToActiveChat(CustomCommand(
             commandId: "",
             name: "",

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -178,7 +178,7 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         var cursorPosition = editor.cursorPosition
         var extraInfo = SuggestionInjector.ExtraInfo()
 
-        let store = Service.shared.guiController.store
+        let store = await Service.shared.guiController.store
 
         if let promptToCode = store.state.promptToCodeGroup.activePromptToCode {
             if promptToCode.isAttachedToSelectionRange, promptToCode.documentURL != fileURL {
@@ -388,7 +388,7 @@ extension WindowBaseCommandHandler {
             )
         }() as (String, CursorRange)
 
-        let store = Service.shared.guiController.store
+        let store = await Service.shared.guiController.store
 
         let customCommandTemplateProcessor = CustomCommandTemplateProcessor()