Merge branch 'feature/global-chat' into develop

Author: intitni

Core/Package.swift

@@ -140,6 +140,7 @@ let package = Package(
             name: "OpenAIService",
             dependencies: [
                 "Logger",
+                "Preferences",
                 .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
             ]
         ),

Core/Sources/Preferences/Keys.swift

@@ -109,6 +109,13 @@ public struct UserDefaultPreferenceKeys {
         .init()
     }
 
+    public struct UseGlobalChat: UserDefaultPreferenceKey {
+        public let defaultValue = false
+        public let key = "UseGlobalChat"
+    }
+
+    public var useGlobalChat: UseGlobalChat { .init() }
+
     public var disableLazyVStack: FeatureFlags.DisableLazyVStack { .init() }
 }
 

Core/Sources/Service/GUI/ChatProvider+Service.swift

@@ -4,7 +4,12 @@ import OpenAIService
 import SuggestionWidget
 
 extension ChatProvider {
-    convenience init(service: ChatService, fileURL: URL, onCloseChat: @escaping () -> Void) {
+    convenience init(
+        service: ChatService,
+        fileURL: URL,
+        onCloseChat: @escaping () -> Void,
+        onSwitchContext: @escaping () -> Void
+    ) {
         self.init()
         let cancellable = service.objectWillChange.sink { [weak self] in
             guard let self else { return }
@@ -19,7 +24,7 @@ extension ChatProvider {
                 self.isReceivingMessage = await service.chatGPTService.isReceivingMessage
             }
         }
-        
+
         service.objectWillChange.send()
 
         onMessageSend = { [cancellable] message in
@@ -47,9 +52,12 @@ extension ChatProvider {
         onClose = {
             Task {
                 await service.stopReceivingMessage()
-                PresentInWindowSuggestionPresenter().closeChatRoom(fileURL: fileURL)
                 onCloseChat()
             }
         }
+        
+        self.onSwitchContext = {
+            onSwitchContext()
+        }
     }
 }

Core/Sources/Service/GUI/WidgetDataSource.swift

@@ -1,12 +1,32 @@
 import ChatService
 import Foundation
+import OpenAIService
 import SuggestionWidget
 
-final class WidgetDataSource: SuggestionWidgetDataSource {
+final class WidgetDataSource {
     static let shared = WidgetDataSource()
 
+    var globalChat: ChatService?
+    var chats = [URL: ChatService]()
+
     private init() {}
 
+    @discardableResult
+    func createChatIfNeeded(for url: URL) -> ChatService {
+        let useGlobalChat = UserDefaults.shared.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],
@@ -49,30 +69,40 @@ final class WidgetDataSource: SuggestionWidgetDataSource {
     }
 
     func chatForFile(at url: URL) async -> ChatProvider? {
-        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
-                        }
+        let useGlobalChat = UserDefaults.shared.value(for: \.useGlobalChat)
+        let buildChatProvider = { (service: ChatService) in
+            ChatProvider(
+                service: service,
+                fileURL: url,
+                onCloseChat: { [weak self] in
+                    if UserDefaults.shared.value(for: \.useGlobalChat) {
+                        self?.globalChat = nil
+                    } else {
+                        self?.chats[url] = nil
                     }
-                )
-            }
+                    let presenter = PresentInWindowSuggestionPresenter()
+                    presenter.closeChatRoom(fileURL: url)
+                },
+                onSwitchContext: { [weak self] in
+                    let useGlobalChat = UserDefaults.shared.value(for: \.useGlobalChat)
+                    UserDefaults.shared.set(!useGlobalChat, for: \.useGlobalChat)
+                    self?.createChatIfNeeded(for: url)
+                    let presenter = PresentInWindowSuggestionPresenter()
+                    presenter.presentChatRoom(fileURL: url)
+                }
+            )
         }
-        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 useGlobalChat {
+            if let globalChat {
+                return buildChatProvider(globalChat)
+            }
+        } else {
+            if let service = chats[url] {
+                return buildChatProvider(service)
             }
         }
+
         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,14 +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 = filespace.chatService ?? ChatService(chatGPTService: ChatGPTService())
-        filespace.chatService = chat
+        let chat = WidgetDataSource.shared.createChatIfNeeded(for: fileURL)
 
         await chat.mutateSystemPrompt(
             """
@@ -212,11 +209,11 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
                 \(code)
                 ```
                 """,
-                summary: "Explain selected code from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character + 1)`."
+                summary: "Explain selected code in `\(fileURL.lastPathComponent)` from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character + 1)`."
             )
         }
 
-        presenter.presentChatGPTConversation(fileURL: fileURL)
+        presenter.presentChatRoom(fileURL: fileURL)
     }
 
     func chatWithSelection(editor: EditorContent) async throws -> UpdatedContent? {
@@ -235,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)
 
@@ -262,8 +257,7 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
             """
         }()
 
-        let chat = filespace.chatService ?? ChatService(chatGPTService: ChatGPTService())
-        filespace.chatService = chat
+        let chat = WidgetDataSource.shared.createChatIfNeeded(for: fileURL)
 
         await chat.mutateSystemPrompt(prompt)
 
@@ -273,12 +267,12 @@ struct WindowBaseCommandHandler: SuggestionCommandHandler {
                     history.append(.init(
                         role: .user,
                         content: "",
-                        summary: "Chat about selected code from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character)`.\nThe code will persist in the conversation."
+                        summary: "Chat about selected code in `\(fileURL.lastPathComponent)` from `\(selection.start.line + 1):\(selection.start.character + 1)` to `\(selection.end.line + 1):\(selection.end.character)`.\nThe code will persist in the conversation."
                     ))
                 }
             }
         }
 
-        presenter.presentChatGPTConversation(fileURL: fileURL)
+        presenter.presentChatRoom(fileURL: fileURL)
     }
 }

Core/Sources/Service/SuggestionPresenter/PresentInWindowSuggestionPresenter.swift

@@ -48,7 +48,7 @@ struct PresentInWindowSuggestionPresenter {
         }
     }
 
-    func presentChatGPTConversation(fileURL: URL) {
+    func presentChatRoom(fileURL: URL) {
         Task { @MainActor in
             let controller = GraphicalUserInterfaceController.shared.suggestionWidget
             controller.presentChatRoom(fileURL: fileURL)

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
     }

Core/Sources/SuggestionWidget/ChatProvider.swift

@@ -1,37 +1,38 @@
 import Foundation
+import SwiftUI
 
-public final class ChatProvider: ObservableObject, Equatable {
+public final class ChatProvider: ObservableObject {
     @Published public var history: [ChatMessage] = []
     @Published public var isReceivingMessage = false
     public var onMessageSend: (String) -> Void
     public var onStop: () -> Void
     public var onClear: () -> Void
     public var onClose: () -> Void
+    public var onSwitchContext: () -> Void
 
     public init(
         history: [ChatMessage] = [],
         isReceivingMessage: Bool = false,
         onMessageSend: @escaping (String) -> Void = { _ in },
         onStop: @escaping () -> Void = {},
         onClear: @escaping () -> Void = {},
-        onClose: @escaping () -> Void = {}
+        onClose: @escaping () -> Void = {},
+        onSwitchContext: @escaping () -> Void = {}
     ) {
         self.history = history
         self.isReceivingMessage = isReceivingMessage
         self.onMessageSend = onMessageSend
         self.onStop = onStop
         self.onClear = onClear
         self.onClose = onClose
+        self.onSwitchContext = onSwitchContext
     }
 
-    public static func == (lhs: ChatProvider, rhs: ChatProvider) -> Bool {
-        lhs.history == rhs.history && lhs.isReceivingMessage == rhs.isReceivingMessage
-    }
-    
     public func send(_ message: String) { onMessageSend(message) }
     public func stop() { onStop() }
     public func clear() { onClear() }
     public func close() { onClose() }
+    public func switchContext() { onSwitchContext() }
 }
 
 public struct ChatMessage: Equatable {