Larry Hao(郝卓远)fix: scope provider refresh state (#1911)e3c2ade4 months agoHistory
791 lines · 30.4 KB · Swift
RawDownload
import AppKitimport CodexBarCore extension StatusItemController {    /// Identifies which manual refresh a task belongs to, so per-provider refreshes stay independent    /// of each other and of the all-providers refresh.    enum ManualRefreshScope: Hashable {        case global        case provider(UsageProvider)    }} enum LoginNotificationLogic {    static func notificationCopy(providerName: String) -> (title: String, body: String) {        (            L("login_success_notification_title", providerName),            L("login_success_notification_body"))    }} extension StatusItemController: StatusItemMenuPersistentActionDelegate {    // MARK: - Actions reachable from menus     func refreshStore(        forceTokenUsage: Bool,        refreshOpenMenusWhenComplete: Bool = true,        interaction: ProviderInteraction = .userInitiated)    {        Task {            await self.performStoreRefresh(                forceTokenUsage: forceTokenUsage,                refreshOpenMenusWhenComplete: refreshOpenMenusWhenComplete,                interaction: interaction)        }    }     func performStoreRefresh(        forceTokenUsage: Bool,        refreshOpenMenusWhenComplete: Bool,        interaction: ProviderInteraction) async    {        await ProviderInteractionContext.$current.withValue(interaction) {            await self.store.refresh(forceTokenUsage: forceTokenUsage)            guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }            self.store.scheduleStorageFootprintRefreshForOverview(force: true)            if refreshOpenMenusWhenComplete {                self.refreshOpenMenusAfterExplicitStoreAction()            } else {                self.invalidateMenus()            }        }    }     func performStoreRefresh(        for provider: UsageProvider,        refreshOpenMenusWhenComplete: Bool,        interaction: ProviderInteraction) async    {        await ProviderInteractionContext.$current.withValue(interaction) {            let refreshStartedAt = Date()            await self.store.refreshProvider(provider)            guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }            await self.store.refreshProviderStatus(provider)            guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }            await self.store.refreshTokenUsageNow(for: provider, force: true)            guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }            if provider == .codex {                await self.store.refreshCreditsNow(minimumSnapshotUpdatedAt: refreshStartedAt)                guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }                await self.store.refreshOpenAIDashboardIfNeeded(                    force: true,                    expectedGuard: self.store.freshCodexOpenAIWebRefreshGuard())                guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }                if self.store.openAIDashboardRequiresLogin {                    await self.store.refreshProvider(.codex)                    guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }                    await self.store.refreshCreditsNow(minimumSnapshotUpdatedAt: refreshStartedAt)                    guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }                }            }            self.store.scheduleStorageFootprintRefresh(for: [provider], force: true)            self.store.persistWidgetSnapshot(reason: "provider-refresh")            if refreshOpenMenusWhenComplete {                self.refreshOpenMenusAfterExplicitStoreAction()            } else {                self.invalidateMenus()            }        }    }     func refreshOpenMenusAfterExplicitStoreAction() {        self.invalidateMenus(            refreshOpenMenus: true,            deferOpenParentMenuRebuild: true)    }     @objc func refreshNow() {        self.startManualRefresh(for: nil)    }     @objc func refreshMenuItem(_ sender: NSMenuItem) {        self.refreshMenuProviderNow(in: sender.menu)    }     func refreshMenuProviderNow(in menu: NSMenu?) {        guard let provider = self.manualRefreshProvider(for: menu) else {            self.startManualRefresh(for: nil)            return        }        self.startManualRefresh(for: provider)    }     private func refreshMenuProviderNow(menuID: ObjectIdentifier) {        if let menu = self.openMenus[menuID] {            self.refreshMenuProviderNow(in: menu)        } else if let mergedMenu = self.mergedMenu, ObjectIdentifier(mergedMenu) == menuID {            self.refreshMenuProviderNow(in: mergedMenu)        } else if let provider = self.menuProviders[menuID] {            self.startManualRefresh(for: provider)        } else {            self.startManualRefresh(for: nil)        }    }     private func startManualRefresh(for provider: UsageProvider?) {        let scope: ManualRefreshScope = provider.map(ManualRefreshScope.provider) ?? .global        let scopedRefreshInFlight = provider.map { self.store.refreshingProviders.contains($0) }            ?? !self.store.refreshingProviders.isEmpty        // Two different providers may refresh concurrently, but an all-providers (.global) refresh must        // not overlap a per-provider one (or vice versa) — that would duplicate the shared fetch work.        let conflictsWithOtherScope = scope == .global            ? self.manualRefreshTasks.contains { $0.key != .global }            : self.manualRefreshTasks[.global] != nil        guard !self.hasPreparedForAppShutdown,              self.manualRefreshTasks[scope] == nil,              !conflictsWithOtherScope,              !self.store.isRefreshing,              !scopedRefreshInFlight        else { return }         let frozenModels = self.frozenManualRefreshMenuCardModels()        let task = Task { @MainActor [weak self] in            guard let self else { return }            defer {                self.manualRefreshTasks[scope] = nil                self.menuCardRefreshMonitor.endManualRefresh(for: provider)                self.updatePersistentRefreshItemsEnabled()                self.prepareAttachedClosedMenusIfNeeded()            }            guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }            #if DEBUG            if let operation = self._test_manualRefreshOperation {                await operation()                guard !Task.isCancelled, !self.hasPreparedForAppShutdown else { return }                return            }            #endif            if let provider {                await self.performStoreRefresh(                    for: provider,                    refreshOpenMenusWhenComplete: true,                    interaction: .userInitiated)            } else {                await self.performStoreRefresh(                    forceTokenUsage: true,                    refreshOpenMenusWhenComplete: true,                    interaction: .userInitiated)            }        }        self.manualRefreshTasks[scope] = task        self.menuCardRefreshMonitor.beginManualRefresh(frozenModels: frozenModels, provider: provider)        self.updatePersistentRefreshItemsEnabled()    }     private func manualRefreshProvider(for menu: NSMenu?) -> UsageProvider? {        guard let menu else { return nil }        if self.shouldMergeIcons {            guard self.mergedMenu == nil || menu === self.mergedMenu else { return nil }            guard !self.isMergedOverviewSelected(in: menu) else { return nil }            return self.resolvedMenuProvider()        }        return self.menuProviders[ObjectIdentifier(menu)]    }     private func frozenManualRefreshMenuCardModels() -> [UsageProvider: UsageMenuCardView.Model] {        var providers = self.store.enabledProvidersForDisplay()        if let lastMenuProvider,           !providers.contains(lastMenuProvider)        {            providers.append(lastMenuProvider)        }        if providers.isEmpty,           let defaultProvider = self.settings.orderedProviders().first ?? UsageProvider.allCases.first        {            providers.append(defaultProvider)        }         var models: [UsageProvider: UsageMenuCardView.Model] = [:]        for provider in providers {            models[provider] = self.menuCardModel(for: provider)        }        return models    }     nonisolated func performPersistentRefreshAction(in menuID: ObjectIdentifier) {        Task { @MainActor [weak self] in            guard let self else { return }            self.refreshMenuProviderNow(menuID: menuID)        }    }     nonisolated func performPersistentSettingsAction() {        Task { @MainActor [weak self] in            guard let self else { return }            self.closeOpenMenusFromShortcutIfNeeded()            self.showSettingsGeneral()        }    }     nonisolated func performPersistentQuitAction() {        Task { @MainActor [weak self] in            guard let self else { return }            self.closeOpenMenusFromShortcutIfNeeded()            self.quit()        }    }     nonisolated func performProviderNavigation(_ direction: StatusItemMenuProviderNavigationDirection) {        Task { @MainActor [weak self] in            self?.navigateProviderSwitcher(direction)        }    }     @objc func refreshAugmentSession() {        Task {            await self.store.forceRefreshAugmentSession()            // Also trigger a full refresh to update the menu and clear any stale errors            await ProviderInteractionContext.$current.withValue(.userInitiated) {                await self.store.refresh(forceTokenUsage: false)            }            self.refreshOpenMenusAfterExplicitStoreAction()        }    }     @objc func installUpdate() {        self.updater.installUpdate()    }     @objc func openDashboard() {        let preferred = self.lastMenuProvider            ?? (self.store.isEnabled(.codex) ? .codex : self.store.enabledProviders().first)         let provider = preferred ?? .codex        guard let url = self.dashboardURL(for: provider) else { return }        NSWorkspace.shared.open(url)    }     func dashboardURL(        for provider: UsageProvider,        environment: [String: String] = ProcessInfo.processInfo.environment) -> URL?    {        if provider == .alibaba {            return self.settings.alibabaCodingPlanAPIRegion.dashboardURL        }        if provider == .minimax {            return self.settings.minimaxAPIRegion.dashboardURL        }         if provider == .opencodego {            return self.settings.opencodegoDashboardURL        }         if provider == .zai {            return ZaiUsageFetcher.resolveDashboardURL(                region: self.settings.zaiAPIRegion,                environment: environment,                usageScope: self.settings.zaiEffectiveUsageScope())        }         if provider == .qoder {            return QoderProviderDescriptor.dashboardURL(                settings: self.settings.qoderSettingsSnapshot(tokenOverride: nil),                sourceLabel: self.store.sourceLabel(for: .qoder))        }         let meta = self.store.metadata(for: provider)        let urlString: String? = if provider == .claude, self.store.isClaudeSubscription() {            meta.subscriptionDashboardURL ?? meta.dashboardURL        } else {            meta.dashboardURL        }         guard let urlString else { return nil }        return URL(string: urlString)    }     @objc func openCreditsPurchase() {        let preferred = self.lastMenuProvider            ?? (self.store.isEnabled(.codex) ? .codex : self.store.enabledProviders().first)        let provider = preferred ?? .codex        guard provider == .codex else { return }         let dashboardURL = self.store.metadata(for: .codex).dashboardURL        let purchaseURL = Self.sanitizedCreditsPurchaseURL(self.store.openAIDashboard?.creditsPurchaseURL)        let urlString = purchaseURL ?? dashboardURL        guard let urlString,              let url = URL(string: urlString) else { return }         let autoStart = true        let accountEmail = self.store.codexAccountEmailForOpenAIDashboard()        let cacheScope = self.store.codexCookieCacheScopeForOpenAIWeb()        guard OpenAICreditsPurchaseWindowController.canOpenPurchaseWindow(            accountEmail: accountEmail,            cacheScope: cacheScope)        else {            self.creditsPurchaseWindow?.close()            self.creditsPurchaseWindow = nil            return        }        let controller = self.creditsPurchaseWindow ?? OpenAICreditsPurchaseWindowController()        controller.show(            purchaseURL: url,            accountEmail: accountEmail,            cacheScope: cacheScope,            autoStartPurchase: autoStart)        self.creditsPurchaseWindow = controller    }     static func sanitizedCreditsPurchaseURL(_ raw: String?) -> String? {        guard let raw, let url = URL(string: raw) else { return nil }        guard Self.isAllowedChatGPTPurchaseHost(url) else { return nil }        let pathComponents = url.pathComponents.map { $0.lowercased() }        let allowed = ["settings", "usage", "billing", "credits"]        guard pathComponents.contains(where: { allowed.contains($0) }) else { return nil }        var components = URLComponents(url: url, resolvingAgainstBaseURL: false)        components?.query = nil        components?.fragment = nil        return components?.url?.absoluteString ?? url.absoluteString    }     private static func isAllowedChatGPTPurchaseHost(_ url: URL) -> Bool {        guard url.scheme?.lowercased() == "https" else { return false }        guard let host = url.host?.lowercased() else { return false }        return host == "chatgpt.com" || host.hasSuffix(".chatgpt.com")    }     @objc func openStatusPage() {        let preferred = self.lastMenuProvider            ?? (self.store.isEnabled(.codex) ? .codex : self.store.enabledProviders().first)         self.openStatusPage(for: preferred ?? .codex)    }     @objc func openStatusPageFromMenuItem(_ sender: NSMenuItem) {        let provider = (sender.identifier?.rawValue).flatMap(UsageProvider.init(rawValue:))            ?? self.lastMenuProvider            ?? .codex        self.openStatusPage(for: provider)    }     private func openStatusPage(for provider: UsageProvider) {        let meta = self.store.metadata(for: provider)        let urlString = meta.statusPageURL ?? meta.statusLinkURL        guard let urlString, let url = URL(string: urlString) else { return }        NSWorkspace.shared.open(url)    }     @objc func openChangelog() {        let preferred = self.lastMenuProvider            ?? (self.store.isEnabled(.codex) ? .codex : self.store.enabledProviders().first)         let provider = preferred ?? .codex        let meta = self.store.metadata(for: provider)        guard let urlString = meta.changelogURL, let url = URL(string: urlString) else { return }        NSWorkspace.shared.open(url)    }     @objc func openTerminalCommand(_ sender: NSMenuItem) {        let command = sender.representedObject as? String ?? "claude"        self.openTerminal(command: command)    }     @objc func openLoginToProvider(_ sender: NSMenuItem) {        guard let urlString = sender.representedObject as? String,              let url = URL(string: urlString) else { return }        NSWorkspace.shared.open(url)    }     @objc func addManagedCodexAccountFromMenu(_: NSMenuItem) {        guard self.codexAccountPromotionCoordinator.isInteractionBlocked() == false else {            self.loginLogger.info("Add Account tap ignored: Codex account change already in-flight")            return        }        guard self.settings.hasUnreadableManagedCodexAccountStore == false else {            self.presentLoginAlert(                title: L("Managed Codex accounts unavailable"),                message: L(                    "CodexBar could not read managed account storage. " +                        "Recover the store before adding another account."))            return        }         Task { @MainActor [weak self] in            guard let self else { return }            do {                let account = try await self.managedCodexAccountCoordinator.authenticateManagedAccount()                self.settings.selectAuthenticatedManagedCodexAccount(account)                await ProviderInteractionContext.$current.withValue(.userInitiated) {                    await self.store.refreshCodexAccountScopedState(allowDisabled: true)                }            } catch {                self.presentManagedCodexAccountError(error)            }        }    }     @objc func requestCodexSystemPromotionFromMenu(_ sender: NSMenuItem) {        guard let rawManagedAccountID = sender.representedObject as? String,              let managedAccountID = UUID(uuidString: rawManagedAccountID)        else {            return        }         Task { @MainActor [weak self] in            guard let self else { return }            let result = await self.codexAccountPromotionCoordinator.promote(managedAccountID: managedAccountID)            if case let .failure(error) = result {                self.presentLoginAlert(title: error.title, message: error.message)            }        }    }     @objc func runSwitchAccount(_ sender: NSMenuItem) {        if self.loginTask != nil {            self.loginLogger.info("Switch Account tap ignored: login already in-flight")            return        }         let rawProvider = sender.representedObject as? String        let provider = rawProvider.flatMap(UsageProvider.init(rawValue:)) ?? self.lastMenuProvider ?? .codex        self.loginLogger.info("Switch Account tapped", metadata: ["provider": provider.rawValue])        self.startLoginFlow(provider: provider)    }     func runLoginFlowFromSettings(provider: UsageProvider) async {        guard self.loginTask == nil else {            self.loginLogger.info(                "Settings login tap ignored: login already in-flight",                metadata: ["provider": provider.rawValue])            return        }        self.startLoginFlow(provider: provider)        await self.loginTask?.value    }     @objc func showSettingsGeneral() {        // Restore the last selected pane; only About navigates explicitly.        self.openSettings(pane: nil)    }     @objc func showSettingsAbout() {        self.openSettings(pane: .about)    }     func openMenuFromShortcut() {        if self.closeOpenMenusFromShortcutIfNeeded() {            return        }         if self.shouldMergeIcons {            self.statusItem.button?.performClick(nil)            return        }         let provider = self.resolvedShortcutProvider()        // Use the lazy accessor to ensure the item exists        let item = self.lazyStatusItem(for: provider)        item.button?.performClick(nil)    }     @discardableResult    func closeOpenMenusFromShortcutIfNeeded() -> Bool {        guard !self.openMenus.isEmpty else { return false }         let menus = Array(self.openMenus.values)        for menu in menus {            menu.cancelTrackingWithoutAnimation()            self.forgetClosedMenu(menu)        }        return true    }     func celebrationOriginPoint(for provider: UsageProvider?) -> CGPoint? {        let item: NSStatusItem = if self.shouldMergeIcons {            self.statusItem        } else if let provider, let existing = self.statusItems[provider], existing.isVisible {            existing        } else {            self.lazyStatusItem(for: provider ?? .codex)        }

Lines 1–500 of 791.