Hide pro chat tabs from opensource code

intitni · intitni · commit 192669e5d4ee · 2024-09-26T22:12:21.000+08:00
diff --git a/Core/Sources/Service/GUI/ChatTabFactory.swift b/Core/Sources/Service/GUI/ChatTabFactory.swift
@@ -8,26 +8,20 @@ import SuggestionBasic
 import SuggestionWidget
 import XcodeInspector
 
-#if canImport(ProChatTabs)
-import ProChatTabs
-#endif
-
 enum ChatTabFactory {
     static func chatTabBuilderCollection() -> [ChatTabBuilderCollection] {
-        #if canImport(ProChatTabs)
-        _ = lazyLoadDependency
-        let collection = [
-            folderIfNeeded(ChatGPTChatTab.chatBuilders(), title: ChatGPTChatTab.name),
-            folderIfNeeded(BrowserChatTab.chatBuilders(), title: BrowserChatTab.name),
-            folderIfNeeded(TerminalChatTab.chatBuilders(), title: TerminalChatTab.name),
-        ]
-        #else
-        let collection = [
-            folderIfNeeded(ChatGPTChatTab.chatBuilders(), title: ChatGPTChatTab.name),
-        ]
-        #endif
-
-        return collection.compactMap { $0 } + chatTabsFromExtensions()
+        let chatGPTChatTab = folderIfNeeded(
+            ChatGPTChatTab.chatBuilders(),
+            title: ChatGPTChatTab.name
+        )
+        
+        let (defaultChatTab, othersChatTabs) = chatTabsFromExtensions()
+        
+        if let defaultChatTab {
+            return [defaultChatTab] + othersChatTabs + [chatGPTChatTab].compactMap(\.self)
+        } else {
+            return [chatGPTChatTab].compactMap(\.self) + othersChatTabs
+        }
     }
 
     private static func folderIfNeeded(
@@ -41,70 +35,85 @@ enum ChatTabFactory {
         return nil
     }
 
-    static func chatTabsFromExtensions() -> [ChatTabBuilderCollection] {
+    static func chatTabsFromExtensions()
+        -> (default: ChatTabBuilderCollection?, others: [ChatTabBuilderCollection])
+    {
         let extensions = BuiltinExtensionManager.shared.extensions
         let chatTabTypes = extensions.flatMap(\.chatTabTypes)
-        return chatTabTypes.compactMap { folderIfNeeded($0.chatBuilders(), title: $0.name) }
+        var defaultChatTab: ChatTabBuilderCollection?
+        var otherChatTabs = [ChatTabBuilderCollection]()
+        for chatTabType in chatTabTypes {
+            if chatTabType.isDefaultChatTabReplacement {
+                defaultChatTab = folderIfNeeded(chatTabType.chatBuilders(), title: chatTabType.name)
+            } else if let tab = folderIfNeeded(
+                chatTabType.chatBuilders(),
+                title: chatTabType.name
+            ) {
+                otherChatTabs.append(tab)
+            }
+        }
+        return (defaultChatTab, otherChatTabs)
     }
 }
 
-#if canImport(ProChatTabs)
-let lazyLoadDependency: () = {
-    BrowserChatTab.externalDependency = .init(
-        handleCustomCommand: { command, prompt in
-           switch command.feature {
-           case let .chatWithSelection(extraSystemPrompt, _, useExtraSystemPrompt):
-               let service = ChatService()
-               return try await service.processMessage(
-                   systemPrompt: nil,
-                   extraSystemPrompt: (useExtraSystemPrompt ?? false) ? extraSystemPrompt :
-                       nil,
-                   prompt: prompt
-               )
-           case let .customChat(systemPrompt, _):
-               let service = ChatService()
-               return try await service.processMessage(
-                   systemPrompt: systemPrompt,
-                   extraSystemPrompt: nil,
-                   prompt: prompt
-               )
-           case let .singleRoundDialog(
-               systemPrompt,
-               overwriteSystemPrompt,
-               _,
-               _
-           ):
-               let service = ChatService()
-               return try await service.handleSingleRoundDialogCommand(
-                   systemPrompt: systemPrompt,
-                   overwriteSystemPrompt: overwriteSystemPrompt ?? false,
-                   prompt: prompt
-               )
-           case let .promptToCode(extraSystemPrompt, instruction, _, _):
-               let service = OpenAIPromptToCodeService()
+// #if canImport(ProChatTabs)
+// let lazyLoadDependency: () = {
+//    BrowserChatTab.externalDependency = .init(
+//        handleCustomCommand: { command, prompt in
+//           switch command.feature {
+//           case let .chatWithSelection(extraSystemPrompt, _, useExtraSystemPrompt):
+//               let service = ChatService()
+//               return try await service.processMessage(
+//                   systemPrompt: nil,
+//                   extraSystemPrompt: (useExtraSystemPrompt ?? false) ? extraSystemPrompt :
+//                       nil,
+//                   prompt: prompt
+//               )
+//           case let .customChat(systemPrompt, _):
+//               let service = ChatService()
+//               return try await service.processMessage(
+//                   systemPrompt: systemPrompt,
+//                   extraSystemPrompt: nil,
+//                   prompt: prompt
+//               )
+//           case let .singleRoundDialog(
+//               systemPrompt,
+//               overwriteSystemPrompt,
+//               _,
+//               _
+//           ):
+//               let service = ChatService()
+//               return try await service.handleSingleRoundDialogCommand(
+//                   systemPrompt: systemPrompt,
+//                   overwriteSystemPrompt: overwriteSystemPrompt ?? false,
+//                   prompt: prompt
+//               )
+//           case let .promptToCode(extraSystemPrompt, instruction, _, _):
+//               let service = OpenAIPromptToCodeService()
+//
+//               let result = try await service.modifyCode(
+//                   code: prompt,
+//                   requirement: instruction ?? "Modify content.",
+//                   source: .init(
+//                       language: .plaintext,
+//                       documentURL: .init(fileURLWithPath: "/"),
+//                       projectRootURL: .init(fileURLWithPath: "/"),
+//                       content: prompt,
+//                       lines: prompt.breakLines(),
+//                       range: .outOfScope
+//                   ),
+//                   isDetached: true,
+//                   extraSystemPrompt: extraSystemPrompt,
+//                   generateDescriptionRequirement: false
+//               )
+//               var code = ""
+//               for try await (newCode, _) in result {
+//                   code = newCode
+//               }
+//               return code
+//           }
+//       }
+//    )
+// }()
+// #endif
 
-               let result = try await service.modifyCode(
-                   code: prompt,
-                   requirement: instruction ?? "Modify content.",
-                   source: .init(
-                       language: .plaintext,
-                       documentURL: .init(fileURLWithPath: "/"),
-                       projectRootURL: .init(fileURLWithPath: "/"),
-                       content: prompt,
-                       lines: prompt.breakLines(),
-                       range: .outOfScope
-                   ),
-                   isDetached: true,
-                   extraSystemPrompt: extraSystemPrompt,
-                   generateDescriptionRequirement: false
-               )
-               var code = ""
-               for try await (newCode, _) in result {
-                   code = newCode
-               }
-               return code
-           }
-       }
-    )
-}()
-#endif
diff --git a/Core/Sources/Service/GUI/GraphicalUserInterfaceController.swift b/Core/Sources/Service/GUI/GraphicalUserInterfaceController.swift
@@ -10,10 +10,6 @@ import Preferences
 import SuggestionBasic
 import SuggestionWidget
 
-#if canImport(ProChatTabs)
-import ProChatTabs
-#endif
-
 #if canImport(ChatTabPersistent)
 import ChatTabPersistent
 #endif
@@ -188,7 +184,7 @@ struct GUI {
                             )
                         }
                     }
-                    
+
                 case let .sendCustomCommandToActiveChat(command):
                     @Sendable func stopAndHandleCommand(_ tab: ChatGPTChatTab) async {
                         if tab.service.isReceivingMessage {
@@ -298,7 +294,7 @@ public final class GraphicalUserInterfaceController {
             dependencies.suggestionWidgetUserDefaultsObservers = .init()
             dependencies.chatTabPool = chatTabPool
             dependencies.chatTabBuilderCollection = ChatTabFactory.chatTabBuilderCollection
-        
+
             #if canImport(ChatTabPersistent) && canImport(ProChatTabs)
             dependencies.restoreChatTabInPool = {
                 await chatTabPool.restore($0)
@@ -383,20 +379,15 @@ extension ChatTabPool {
         return (chatTap, info)
     }
 
-    #if canImport(ChatTabPersistent) && canImport(ProChatTabs)
+    #if canImport(ChatTabPersistent)
     @MainActor
     func restore(
         _ data: ChatTabPersistent.RestorableTabData
     ) async -> (any ChatTab, ChatTabInfo)? {
         switch data.name {
         case ChatGPTChatTab.name:
-            guard let builder = try? await ChatGPTChatTab.restore(from: data.data) else { break }
-            return await createTab(id: data.id, from: builder)
-        case BrowserChatTab.name:
-            guard let builder = try? BrowserChatTab.restore(from: data.data) else { break }
-            return await createTab(id: data.id, from: builder)
-        case TerminalChatTab.name:
-            guard let builder = try? await TerminalChatTab.restore(from: data.data) else { break }
+            guard let builder = try? await ChatGPTChatTab.restore(from: data.data)
+            else { fallthrough }
             return await createTab(id: data.id, from: builder)
         default:
             let chatTabTypes = BuiltinExtensionManager.shared.extensions.flatMap(\.chatTabTypes)
diff --git a/Core/Sources/Service/Service.swift b/Core/Sources/Service/Service.swift
@@ -54,20 +54,21 @@ public final class Service {
 
         realtimeSuggestionController = .init()
         scheduledCleaner = .init()
-        let guiController = GraphicalUserInterfaceController()
-        self.guiController = guiController
-        globalShortcutManager = .init(guiController: guiController)
-        keyBindingManager = .init()
 
         #if canImport(ProService)
         proService = ProService()
         #endif
 
-        BuiltinExtensionManager.shared.setupExtensions([
+        BuiltinExtensionManager.shared.addExtensions([
             GitHubCopilotExtension(workspacePool: workspacePool),
             CodeiumExtension(workspacePool: workspacePool),
         ])
 
+        let guiController = GraphicalUserInterfaceController()
+        self.guiController = guiController
+        globalShortcutManager = .init(guiController: guiController)
+        keyBindingManager = .init()
+
         workspacePool.registerPlugin {
             SuggestionServiceWorkspacePlugin(workspace: $0) { SuggestionService.service() }
         }
diff --git a/Tool/Package.swift b/Tool/Package.swift
@@ -51,6 +51,7 @@ let package = Package(
         .library(name: "CustomAsyncAlgorithms", targets: ["CustomAsyncAlgorithms"]),
         .library(name: "CommandHandler", targets: ["CommandHandler"]),
         .library(name: "CodeDiff", targets: ["CodeDiff"]),
+        .library(name: "BuiltinExtension", targets: ["BuiltinExtension"]),
     ],
     dependencies: [
         // A fork of https://github.com/aespinilla/Tiktoken to allow loading from local files.
diff --git a/Tool/Sources/BuiltinExtension/BuiltinExtension.swift b/Tool/Sources/BuiltinExtension/BuiltinExtension.swift
@@ -6,7 +6,7 @@ import Preferences
 
 public protocol BuiltinExtension: CopilotForXcodeExtensionCapability {
     /// An id that let the extension manager determine whether the extension is in use.
-    var suggestionServiceId: BuiltInSuggestionFeatureProvider { get }
+    var suggestionServiceId: BuiltInSuggestionFeatureProvider? { get }
     /// An identifier for the extension.
     var extensionIdentifier: String { get }
 
@@ -21,6 +21,7 @@ public protocol BuiltinExtension: CopilotForXcodeExtensionCapability {
 // MARK: - Default Implementation
 
 public extension BuiltinExtension {
+    var suggestionServiceId: BuiltInSuggestionFeatureProvider? { nil }
     var chatTabTypes: [any ChatTab.Type] { [] }
 }
 
diff --git a/Tool/Sources/BuiltinExtension/BuiltinExtensionManager.swift b/Tool/Sources/BuiltinExtension/BuiltinExtensionManager.swift
@@ -22,6 +22,11 @@ public final class BuiltinExtensionManager {
         checkAppConfiguration()
     }
 
+    public func addExtensions(_ extensions: [any BuiltinExtension]) {
+        self.extensions.append(contentsOf: extensions)
+        checkAppConfiguration()
+    }
+
     public func terminate() {
         for ext in extensions {
             ext.terminate()
@@ -33,8 +38,11 @@ extension BuiltinExtensionManager {
     func checkAppConfiguration() {
         let suggestionFeatureProvider = UserDefaults.shared.value(for: \.suggestionFeatureProvider)
         for ext in extensions {
-            let isSuggestionFeatureInUse = suggestionFeatureProvider ==
-                .builtIn(ext.suggestionServiceId)
+            let isSuggestionFeatureInUse = if let id = ext.suggestionServiceId {
+                suggestionFeatureProvider == .builtIn(id)
+            } else {
+                false
+            }
             let isChatFeatureInUse = false
             ext.extensionUsageDidChange(.init(
                 isSuggestionServiceInUse: isSuggestionFeatureInUse,
diff --git a/Tool/Sources/ChatTab/ChatTab.swift b/Tool/Sources/ChatTab/ChatTab.swift
@@ -45,6 +45,13 @@ public protocol ChatTabType {
     /// It will be called only once so long as you don't call it yourself.
     /// It will be called from MainActor.
     func start()
+    /// Whenever the user close the tab, this method will be called.
+    func close()
+    
+    /// Whether this chat tab should be the default chat tab replacement.
+    static var isDefaultChatTabReplacement: Bool { get }
+    /// Whether this chat tab can handle open chat command.
+    static var canHandleOpenChatCommand: Bool { get }
 }
 
 /// The base class for all chat tabs.
@@ -67,10 +74,10 @@ open class BaseChatTab {
 
     public init(store: StoreOf<ChatTabItem>) {
         chatTabStore = store
-        self.id = store.id
-        self.title = store.title
-        
+
         Task { @MainActor in
+            self.title = store.title
+            self.id = store.id
             storeObserver.observe { [weak self] in
                 guard let self else { return }
                 self.title = store.title
@@ -166,6 +173,12 @@ public struct DisabledChatTabBuilder: ChatTabBuilder {
 public extension ChatTabType {
     /// The name of this chat tab type.
     var name: String { Self.name }
+
+    /// Default implementation that does nothing.
+    func close() {}
+    
+    static var canHandleOpenChatCommand: Bool { false }
+    static var isDefaultChatTabReplacement: Bool { false }
 }
 
 /// A chat tab that does nothing.
