Move chat services to WidgetDataSource

Author: intitni

Core/Sources/Service/GUI/WidgetDataSource.swift

@@ -1,14 +1,31 @@
 import ChatService
 import Foundation
+import OpenAIService
 import SuggestionWidget
 
-final class WidgetDataSource: SuggestionWidgetDataSource {
+final class WidgetDataSource {
     static let shared = WidgetDataSource()
-    
-    var globalChat: ChatService? = nil
+
+    var globalChat: ChatService?
+    var chats = [URL: ChatService]()
 
     private init() {}
 
+    func createChatIfNeeded(for url: URL) -> ChatService {
+        let useGlobalChat = UserDefaults.standard.value(for: \.useGlobalChat)
+        let chat: ChatService
+        if useGlobalChat {
+            chat = globalChat ?? ChatService(chatGPTService: ChatGPTService())
+            globalChat = chat
+        } else {
+            chat = chats[url] ?? ChatService(chatGPTService: ChatGPTService())
+            chats[url] = chat
+        }
+        return chat
+    }
+}
+
+extension WidgetDataSource: SuggestionWidgetDataSource {
     func suggestionForFile(at url: URL) async -> SuggestionProvider? {
         for workspace in await workspaces.values {
             if let filespace = await workspace.filespaces[url],
@@ -63,30 +80,17 @@ final class WidgetDataSource: SuggestionWidgetDataSource {
                 }
             )
         }
-        
-        for workspace in await workspaces.values {
-            if let filespace = await workspace.filespaces[url],
-               let service = await filespace.chatService
-            {
-                return .init(
-                    service: service,
-                    fileURL: url,
-                    onCloseChat: { [weak filespace] in
-                        Task { @ServiceActor [weak filespace] in
-                            filespace?.chatService = nil
-                        }
-                    }
-                )
-            }
-        }
-        return nil
-    }
 
-    func chatServiceForFile(at url: URL) async -> ChatService? {
-        for workspace in await workspaces.values {
-            if let filespace = await workspace.filespaces[url] {
-                return await filespace.chatService
-            }
+        if let service = chats[url] {
+            return .init(
+                service: service,
+                fileURL: url,
+                onCloseChat: { [weak self] in
+                    Task { @ServiceActor [weak self] in
+                        self?.chats[url] = nil
+                    }
+                }
+            )
         }
         return nil
     }

Core/Sources/Service/ScheduledCleaner.swift

@@ -2,15 +2,23 @@ import Foundation
 
 public final class ScheduledCleaner {
     public init() {
-        // Occasionally cleanup workspaces.
+        // occasionally cleanup workspaces.
         Task { @ServiceActor in
             while !Task.isCancelled {
                 try await Task.sleep(nanoseconds: 8 * 60 * 60 * 1_000_000_000)
                 for (url, workspace) in workspaces {
                     if workspace.isExpired {
                         workspaces[url] = nil
                     } else {
-                        workspaces[url]?.cleanUp()
+                        // cleanup chats for unused files
+                        let filespaces = workspace.filespaces
+                        for (url, filespace) in filespaces {
+                            if filespace.isExpired {
+                                WidgetDataSource.shared.chats[url] = nil
+                            }
+                        }
+                        // cleanup workspace
+                        workspace.cleanUp()
                     }
                 }
             }

Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift

@@ -188,13 +188,11 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         defer { presenter.markAsProcessing(false) }
 
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let (_, filespace) = try await Workspace
-            .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
         let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
         let codeLanguage = languageIdentifierFromFileURL(fileURL)
         guard let selection = editor.selections.last else { return }
 
-        let chat = createChatServiceIfNeeded(filespace: filespace)
+        let chat = WidgetDataSource.shared.createChatIfNeeded(for: fileURL)
 
         await chat.mutateSystemPrompt(
             """
@@ -234,8 +232,6 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         defer { presenter.markAsProcessing(false) }
 
         let fileURL = try await Environment.fetchCurrentFileURL()
-        let (_, filespace) = try await Workspace
-            .fetchOrCreateWorkspaceIfNeeded(fileURL: fileURL)
         let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
         let codeLanguage = languageIdentifierFromFileURL(fileURL)
 
@@ -261,7 +257,7 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
             """
         }()
 
-        let chat = createChatServiceIfNeeded(filespace: filespace)
+        let chat = WidgetDataSource.shared.createChatIfNeeded(for: fileURL)
 
         await chat.mutateSystemPrompt(prompt)
 
@@ -280,19 +276,3 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
         presenter.presentChatGPTConversation(fileURL: fileURL)
     }
 }
-
-extension SuggestionCommandHandler {
-    @ServiceActor
-    func createChatServiceIfNeeded(filespace: Filespace) -> ChatService {
-        let chat: ChatService
-        if UserDefaults.shared.value(for: \.useGlobalChat) {
-            chat = WidgetDataSource.shared.globalChat
-                ?? ChatService(chatGPTService: ChatGPTService())
-            WidgetDataSource.shared.globalChat = chat
-        } else {
-            chat = filespace.chatService ?? ChatService(chatGPTService: ChatGPTService())
-            filespace.chatService = chat
-        }
-        return chat
-    }
-}

Core/Sources/Service/Workspace.swift

@@ -39,8 +39,6 @@ final class Filespace {
         Environment.now().timeIntervalSince(lastSuggestionUpdateTime) > 60 * 60 * 8
     }
 
-    var chatService: ChatService? = nil
-
     init(fileURL: URL) {
         self.fileURL = fileURL
     }