Support recovering Codeium chat tab

Author: intitni

Core/Sources/ChatGPTChatTab/ChatGPTChatTab.swift

@@ -76,10 +76,7 @@ public class ChatGPTChatTab: ChatTab {
         return (try? JSONEncoder().encode(state)) ?? Data()
     }
 
-    public static func restore(
-        from data: Data,
-        externalDependency: Void
-    ) async throws -> any ChatTabBuilder {
+    public static func restore(from data: Data) async throws -> any ChatTabBuilder {
         let state = try JSONDecoder().decode(RestorableState.self, from: data)
         let builder = Builder(title: "Chat") { @MainActor tab in
             tab.service.configuration.overriding = state.configuration
@@ -96,7 +93,7 @@ public class ChatGPTChatTab: ChatTab {
         return builder
     }
 
-    public static func chatBuilders(externalDependency: Void) -> [ChatTabBuilder] {
+    public static func chatBuilders() -> [ChatTabBuilder] {
         let customCommands = UserDefaults.shared.value(for: \.customCommands).compactMap {
             command in
             if case .customChat = command.feature {

Core/Sources/Service/GUI/ChatTabFactory.swift

@@ -15,14 +15,10 @@ import ProChatTabs
 enum ChatTabFactory {
     static func chatTabBuilderCollection() -> [ChatTabBuilderCollection] {
         #if canImport(ProChatTabs)
+        _ = lazyLoadDependency
         let collection = [
             folderIfNeeded(ChatGPTChatTab.chatBuilders(), title: ChatGPTChatTab.name),
-            folderIfNeeded(
-                BrowserChatTab.chatBuilders(
-                    externalDependency: externalDependenciesForBrowserChatTab()
-                ),
-                title: BrowserChatTab.name
-            ),
+            folderIfNeeded(BrowserChatTab.chatBuilders(), title: BrowserChatTab.name),
             folderIfNeeded(TerminalChatTab.chatBuilders(), title: TerminalChatTab.name),
         ]
         #else
@@ -47,83 +43,68 @@ enum ChatTabFactory {
 
     static func chatTabsFromExtensions() -> [ChatTabBuilderCollection] {
         let extensions = BuiltinExtensionManager.shared.extensions
-        let chatBuilders = extensions.flatMap { $0.chatBuilders }
-        return chatBuilders.compactMap { folderIfNeeded($0.value, title: $0.key) }
+        let chatTabTypes = extensions.flatMap(\.chatTabTypes)
+        return chatTabTypes.compactMap { folderIfNeeded($0.chatBuilders(), title: $0.name) }
     }
-
-    #if canImport(ProChatTabs)
-    static func externalDependenciesForBrowserChatTab() -> BrowserChatTab.ExternalDependency {
-        .init(
-            getEditorContent: {
-                guard let editor = XcodeInspector.shared.focusedEditor else {
-                    return .init(selectedText: "", language: "", fileContent: "")
-                }
-                let content = editor.getContent()
-                return .init(
-                    selectedText: content.selectedContent,
-                    language: (
-                        XcodeInspector.shared.activeDocumentURL
-                            .map(languageIdentifierFromFileURL) ?? .plaintext
-                    ).rawValue,
-                    fileContent: content.content
-                )
-            },
-            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
 }
 
+#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

Core/Sources/Service/GUI/GraphicalUserInterfaceController.swift

@@ -1,6 +1,7 @@
 import ActiveApplicationMonitor
 import AppActivator
 import AppKit
+import BuiltinExtension
 import ChatGPTChatTab
 import ChatTab
 import ComposableArchitecture
@@ -405,38 +406,25 @@ extension ChatTabPool {
     ) async -> (any ChatTab, ChatTabInfo)? {
         switch data.name {
         case ChatGPTChatTab.name:
-            guard let builder = try? await ChatGPTChatTab.restore(
-                from: data.data,
-                externalDependency: ()
-            ) else { break }
-            return await createTab(id: data.id, from: builder)
-        case EmptyChatTab.name:
-            guard let builder = try? await EmptyChatTab.restore(
-                from: data.data,
-                externalDependency: ()
-            ) else { break }
+            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,
-                externalDependency: ChatTabFactory.externalDependenciesForBrowserChatTab()
-            ) else { break }
+            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,
-                externalDependency: ()
-            ) else { break }
+            guard let builder = try? await TerminalChatTab.restore(from: data.data) else { break }
             return await createTab(id: data.id, from: builder)
         default:
-            break
+            let chatTabTypes = BuiltinExtensionManager.shared.extensions.flatMap(\.chatTabTypes)
+            for type in chatTabTypes {
+                if type.name == data.name {
+                    guard let builder = try? await type.restore(from: data.data) else { break }
+                    return await createTab(id: data.id, from: builder)
+                }
+            }
         }
 
-        guard let builder = try? await EmptyChatTab.restore(
-            from: data.data, externalDependency: ()
-        ) else {
-            return nil
-        }
+        guard let builder = try? await EmptyChatTab.restore(from: data.data) else { return nil }
         return await createTab(id: data.id, from: builder)
     }
     #endif

Core/Sources/Service/SuggestionCommandHandler/PseudoCommandHandler.swift

@@ -362,11 +362,7 @@ struct PseudoCommandHandler {
                                   match(tab.url) else { return false }
                             return true
                         },
-                        kind: .init(BrowserChatTab.urlChatBuilder(
-                            url: url,
-                            externalDependency: ChatTabFactory
-                                .externalDependenciesForBrowserChatTab()
-                        ))
+                        kind: .init(BrowserChatTab.urlChatBuilder(url: url))
                     )).finish()
                     store.send(.openChatPanel(forceDetach: forceDetach))
                 }

Pro

@@ -8,7 +8,7 @@ public protocol BuiltinExtension: CopilotForXcodeExtensionCapability {
     var suggestionServiceId: BuiltInSuggestionFeatureProvider { get }
 
     /// All chat builders provided by this extension.
-    var chatBuilders: [String: [ChatTabBuilder]] { get }
+    var chatTabTypes: [any ChatTab.Type] { get }
 
     /// It's usually called when the app is about to quit,
     /// you should clean up all the resources here.
@@ -18,6 +18,6 @@ public protocol BuiltinExtension: CopilotForXcodeExtensionCapability {
 // MARK: - Default Implementation
 
 public extension BuiltinExtension {
-    var chatBuilders: [String: [ChatTabBuilder]] { [:] }
+    var chatTabTypes: [any ChatTab.Type] { [] }
 }
 

Tool/Sources/BuiltinExtension/BuiltinExtension.swift

@@ -20,8 +20,6 @@ public typealias ChatTab = BaseChatTab & ChatTabType
 
 /// Defines a bunch of things a chat tab should implement.
 public protocol ChatTabType {
-    /// The type of the external dependency required by this chat tab.
-    associatedtype ExternalDependency
     /// Build the view for this chat tab.
     @ViewBuilder
     func buildView() -> any View
@@ -38,14 +36,11 @@ public protocol ChatTabType {
     static var name: String { get }
     /// Available builders for this chat tab.
     /// It's used to generate a list of tab types for user to create.
-    static func chatBuilders(externalDependency: ExternalDependency) -> [ChatTabBuilder]
+    static func chatBuilders() -> [ChatTabBuilder]
     /// Restorable state
     func restorableState() async -> Data
     /// Restore state
-    static func restore(
-        from data: Data,
-        externalDependency: ExternalDependency
-    ) async throws -> any ChatTabBuilder
+    static func restore(from data: Data) async throws -> any ChatTabBuilder
     /// Whenever the body or menu is accessed, this method will be called.
     /// It will be called only once so long as you don't call it yourself.
     /// It will be called from MainActor.
@@ -66,7 +61,7 @@ open class BaseChatTab {
     public var title: String = ""
     /// The store for chat tab info. You should only access it after `start` is called.
     public let chatTabStore: StoreOf<ChatTabItem>
-    
+
     private var didStart = false
     private let storeObserver = NSObject()
 
@@ -109,7 +104,7 @@ open class BaseChatTab {
             EmptyView().id(id)
         }
     }
-    
+
     /// The icon for this chat tab.
     @ViewBuilder
     public var icon: some View {
@@ -120,7 +115,7 @@ open class BaseChatTab {
             EmptyView().id(id)
         }
     }
-    
+
     /// The tab item for this chat tab.
     @ViewBuilder
     public var menu: some View {
@@ -171,14 +166,6 @@ public extension ChatTabType {
     var name: String { Self.name }
 }
 
-public extension ChatTabType where ExternalDependency == Void {
-    /// Available builders for this chat tab.
-    /// It's used to generate a list of tab types for user to create.
-    static func chatBuilders() -> [ChatTabBuilder] {
-        chatBuilders(externalDependency: ())
-    }
-}
-
 /// A chat tab that does nothing.
 public class EmptyChatTab: ChatTab {
     public static var name: String { "Empty" }
@@ -190,7 +177,7 @@ public class EmptyChatTab: ChatTab {
         }
     }
 
-    public static func chatBuilders(externalDependency: Void) -> [ChatTabBuilder] {
+    public static func chatBuilders() -> [ChatTabBuilder] {
         [Builder(title: "Empty")]
     }
 
@@ -204,11 +191,11 @@ public class EmptyChatTab: ChatTab {
     public func buildTabItem() -> any View {
         Text("Empty-\(id)")
     }
-    
+
     public func buildIcon() -> any View {
         Image(systemName: "square")
     }
-    
+
     public func buildMenu() -> any View {
         Text("Empty-\(id)")
     }
@@ -217,10 +204,7 @@ public class EmptyChatTab: ChatTab {
         return Data()
     }
 
-    public static func restore(
-        from data: Data,
-        externalDependency: Void
-    ) async throws -> any ChatTabBuilder {
+    public static func restore(from data: Data) async throws -> any ChatTabBuilder {
         return Builder(title: "Empty")
     }