758 lines · 27.1 KB · Swift
RawDownload
#if os(macOS)import AppKitimport Foundationimport WebKit struct OpenAIDashboardWebViewLease {    let webView: WKWebView    let log: (String) -> Void    let setPreserveLoadedPageOnRelease: (Bool) -> Void    let release: () -> Void} @MainActorfinal class OpenAIDashboardWebViewCache {    static let shared = OpenAIDashboardWebViewCache()    fileprivate static let log = CodexBarLog.logger(LogCategories.openAIWebview)     private final class ReleaseState {        var preserveLoadedPageOnRelease: Bool         init(preserveLoadedPageOnRelease: Bool) {            self.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease        }    }     private struct AcquireOptions {        let allowTimeoutRetry: Bool        let preserveLoadedPageOnRelease: Bool    }     @MainActor    private final class NavigationCancellationState {        private weak var webView: WKWebView?        private var delegate: NavigationDelegate?        private var isCancelled = false         func install(webView: WKWebView, delegate: NavigationDelegate) {            self.webView = webView            self.delegate = delegate            if self.isCancelled {                self.cancel()            }        }         func cancel() {            self.isCancelled = true            guard let webView, let delegate else { return }            delegate.cancel()            if webView.codexNavigationDelegate === delegate {                webView.stopLoading()                webView.navigationDelegate = nil                webView.codexNavigationDelegate = nil            }            self.delegate = nil            self.webView = nil        }    }     private final class Entry {        let webView: WKWebView        let host: OffscreenWebViewHost        var lastUsedAt: Date        var isBusy: Bool        var preservedPageExpiresAt: Date?        var preservedPageExpiryWorkItem: DispatchWorkItem?         init(            webView: WKWebView,            host: OffscreenWebViewHost,            lastUsedAt: Date,            isBusy: Bool,            preservedPageExpiresAt: Date? = nil)        {            self.webView = webView            self.host = host            self.lastUsedAt = lastUsedAt            self.isBusy = isBusy            self.preservedPageExpiresAt = preservedPageExpiresAt        }         func armPreservedPage(until expiry: Date) {            self.preservedPageExpiresAt = expiry        }         func setPreservedPageExpiryWorkItem(_ workItem: DispatchWorkItem?) {            self.preservedPageExpiryWorkItem?.cancel()            self.preservedPageExpiryWorkItem = workItem        }         func clearPreservedPage() {            self.preservedPageExpiresAt = nil            self.preservedPageExpiryWorkItem?.cancel()            self.preservedPageExpiryWorkItem = nil        }         func consumePreservedPageReuseIfAvailable(now: Date) -> Bool {            guard let preservedPageExpiresAt else { return false }            self.preservedPageExpiresAt = nil            self.preservedPageExpiryWorkItem?.cancel()            self.preservedPageExpiryWorkItem = nil            return preservedPageExpiresAt > now        }         func hasExpiredPreservedPage(now: Date) -> Bool {            guard let preservedPageExpiresAt else { return false }            return preservedPageExpiresAt <= now        }    }     private var entries: [ObjectIdentifier: Entry] = [:]    /// Keep the WebView alive only long enough for immediate retries/menu reopens.    /// Long-lived hidden ChatGPT tabs still consume noticeable energy on some setups.    private let idleTimeout: TimeInterval    private var idlePruneWorkItem: DispatchWorkItem?    private var idlePruneGeneration = 0    #if DEBUG    private(set) var idlePruneDeadlineForTesting: Date?    #endif    /// Reuse the validated analytics page only for the immediate next handoff.    private let preservedPageHandoffTimeout: TimeInterval = 5    private let blankURL = URL(string: "about:blank")!    private let idlePageClearScript = """    (() => {      try {        document.documentElement.innerHTML = '';        return true;      } catch {        return false;      }    })();    """    private let reusablePageResetScript = """    (() => {      try {        delete window.__codexbarDidScrollToCredits;        delete window.__codexbarUsageBreakdownJSON;        delete window.__codexbarUsageBreakdownDebug;        return true;      } catch {        return false;      }    })();    """    private let preferredLanguageScript = """    (() => {      const define = (target, name, value) => {        try {          Object.defineProperty(target, name, {            get: () => value,            configurable: true          });        } catch {}      };      define(Navigator.prototype, 'language', 'en-US');      define(Navigator.prototype, 'languages', ['en-US', 'en']);      define(navigator, 'language', 'en-US');      define(navigator, 'languages', ['en-US', 'en']);    })();    """     init(idleTimeout: TimeInterval = 60) {        self.idleTimeout = idleTimeout    }     private func releaseCachedEntry(_ entry: Entry, preserveLoadedPage: Bool) {        entry.isBusy = false        let now = Date()        entry.lastUsedAt = now        self.updatePreservedPageState(for: entry, preserveLoadedPage: preserveLoadedPage)        self.prepareCachedWebViewForIdle(            entry.webView,            host: entry.host,            preserveLoadedPage: preserveLoadedPage)        self.prune(now: now)        self.scheduleNextIdlePrune(now: now)    }     private func releaseNewEntry(_ entry: Entry, webView: WKWebView, preserveLoadedPage: Bool) {        entry.isBusy = false        let now = Date()        entry.lastUsedAt = now        self.updatePreservedPageState(for: entry, preserveLoadedPage: preserveLoadedPage)        self.prepareCachedWebViewForIdle(            webView,            host: entry.host,            preserveLoadedPage: preserveLoadedPage)        self.prune(now: now)        self.scheduleNextIdlePrune(now: now)    }     // MARK: - Testing support     #if DEBUG    /// Number of cached WebView entries (for testing).    var entryCount: Int {        self.entries.count    }     /// Check if a WebView is cached for the given data store (for testing).    func hasCachedEntry(for websiteDataStore: WKWebsiteDataStore) -> Bool {        let key = ObjectIdentifier(websiteDataStore)        return self.entries[key] != nil    }     /// Force prune with a custom "now" timestamp (for testing idle timeout).    func pruneForTesting(now: Date) {        self.prune(now: now)    }     var idleTimeoutForTesting: TimeInterval {        self.idleTimeout    }     var preservedPageHandoffTimeoutForTesting: TimeInterval {        self.preservedPageHandoffTimeout    }     func hasPreservedPageForTesting(for websiteDataStore: WKWebsiteDataStore) -> Bool {        let key = ObjectIdentifier(websiteDataStore)        return self.entries[key]?.preservedPageExpiresAt != nil    }     func markPreservedPageForTesting(        websiteDataStore: WKWebsiteDataStore,        expiresAt: Date = .init().addingTimeInterval(5))    {        let key = ObjectIdentifier(websiteDataStore)        guard let entry = self.entries[key] else { return }        entry.armPreservedPage(until: expiresAt)        self.schedulePreservedPageExpiry(for: key, entry: entry, expiresAt: expiresAt)    }     func consumePreservedPageForTesting(websiteDataStore: WKWebsiteDataStore, now: Date = Date()) -> Bool {        let key = ObjectIdentifier(websiteDataStore)        guard let entry = self.entries[key] else { return false }        return entry.consumePreservedPageReuseIfAvailable(now: now)    }     /// Seed a cached entry without navigating a real page (for test stability).    @discardableResult    func cacheEntryForTesting(        websiteDataStore: WKWebsiteDataStore,        lastUsedAt: Date = Date(),        isBusy: Bool = false) -> WKWebView    {        let key = ObjectIdentifier(websiteDataStore)        if let existing = self.entries.removeValue(forKey: key) {            existing.host.close()        }         let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore)        let entry = Entry(webView: webView, host: host, lastUsedAt: lastUsedAt, isBusy: isBusy)        self.entries[key] = entry        if isBusy {            host.show()        } else {            host.hide()        }        return webView    }     /// Clear all cached entries (for test isolation).    func clearAllForTesting() {        self.cancelIdlePrune()        for (_, entry) in self.entries {            entry.clearPreservedPage()            entry.host.close()        }        self.entries.removeAll()    }     func resetReusablePageStateForTesting(_ webView: WKWebView) async -> Bool {        await self.resetReusablePageState(webView)    }    #endif     func acquire(        websiteDataStore: WKWebsiteDataStore,        usageURL: URL,        logger: ((String) -> Void)?,        navigationTimeout: TimeInterval = 15,        allowTimeoutRetry: Bool = true,        preserveLoadedPageOnRelease: Bool = false) async throws -> OpenAIDashboardWebViewLease    {        let deadline = Date().addingTimeInterval(max(navigationTimeout, 1))        return try await self.acquire(            websiteDataStore: websiteDataStore,            usageURL: usageURL,            logger: logger,            deadline: deadline,            options: .init(                allowTimeoutRetry: allowTimeoutRetry,                preserveLoadedPageOnRelease: preserveLoadedPageOnRelease))    }     private func acquire(        websiteDataStore: WKWebsiteDataStore,        usageURL: URL,        logger: ((String) -> Void)?,        deadline: Date,        options: AcquireOptions) async throws -> OpenAIDashboardWebViewLease    {        let now = Date()        self.prune(now: now)         let log: (String) -> Void = { message in            logger?("[webview] \(message)")        }        let key = ObjectIdentifier(websiteDataStore)        let remainingTimeout = max(0.5, deadline.timeIntervalSince(now))         if let entry = self.entries[key] {            if entry.isBusy {                log("Cached WebView busy; using a temporary WebView.")                let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore)                host.show()                do {                    try await self.prepareWebView(                        webView,                        usageURL: usageURL,                        timeout: remainingTimeout,                        canReuseLoadedPage: false)                } catch {                    if options.allowTimeoutRetry, Self.isPrepareTimeout(error) {                        host.close()                        log("Temporary OpenAI WebView timed out; retrying with a fresh WebView.")                        return try await self.acquireTemporaryWebView(                            websiteDataStore: websiteDataStore,                            usageURL: usageURL,                            log: log,                            deadline: deadline)                    }                    host.close()                    throw error                }                return OpenAIDashboardWebViewLease(                    webView: webView,                    log: log,                    setPreserveLoadedPageOnRelease: { _ in },                    release: { host.close() })            }             entry.isBusy = true            entry.lastUsedAt = now            let canReuseLoadedPage = entry.consumePreservedPageReuseIfAvailable(now: now)            let releaseState = ReleaseState(preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease)            entry.host.show()            do {                try await self.prepareWebView(                    entry.webView,                    usageURL: usageURL,                    timeout: remainingTimeout,                    canReuseLoadedPage: canReuseLoadedPage)            } catch {                if options.allowTimeoutRetry, Self.isPrepareTimeout(error) {                    entry.isBusy = false                    entry.lastUsedAt = Date()                    entry.clearPreservedPage()                    entry.host.close()                    self.entries.removeValue(forKey: key)                    log("Cached OpenAI WebView timed out; recreating it.")                    return try await self.acquire(                        websiteDataStore: websiteDataStore,                        usageURL: usageURL,                        logger: logger,                        deadline: deadline,                        options: .init(                            allowTimeoutRetry: false,                            preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease))                }                entry.isBusy = false                entry.lastUsedAt = Date()                entry.clearPreservedPage()                entry.host.close()                self.entries.removeValue(forKey: key)                Self.log.warning("OpenAI webview prepare failed")                throw error            }             return OpenAIDashboardWebViewLease(                webView: entry.webView,                log: log,                setPreserveLoadedPageOnRelease: { preserveLoadedPageOnRelease in                    releaseState.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease                },                release: { [weak self, weak entry] in                    guard let self, let entry else { return }                    self.releaseCachedEntry(                        entry,                        preserveLoadedPage: releaseState.preserveLoadedPageOnRelease)                })        }         let (webView, host) = self.makeWebView(websiteDataStore: websiteDataStore)        let entry = Entry(webView: webView, host: host, lastUsedAt: now, isBusy: true)        self.entries[key] = entry        host.show()        let releaseState = ReleaseState(preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease)         do {            try await self.prepareWebView(                webView,                usageURL: usageURL,                timeout: remainingTimeout,                canReuseLoadedPage: false)        } catch {            if options.allowTimeoutRetry, Self.isPrepareTimeout(error) {                self.entries.removeValue(forKey: key)                host.close()                log("OpenAI WebView timed out during prepare; retrying once.")                return try await self.acquire(                    websiteDataStore: websiteDataStore,                    usageURL: usageURL,                    logger: logger,                    deadline: deadline,                    options: .init(                        allowTimeoutRetry: false,                        preserveLoadedPageOnRelease: options.preserveLoadedPageOnRelease))            }            self.entries.removeValue(forKey: key)            host.close()            Self.log.warning("OpenAI webview prepare failed")            throw error        }         return OpenAIDashboardWebViewLease(            webView: webView,            log: log,            setPreserveLoadedPageOnRelease: { preserveLoadedPageOnRelease in                releaseState.preserveLoadedPageOnRelease = preserveLoadedPageOnRelease            },            release: { [weak self, weak entry] in                guard let self, let entry else { return }                self.releaseNewEntry(                    entry,                    webView: webView,                    preserveLoadedPage: releaseState.preserveLoadedPageOnRelease)            })    }     func evict(websiteDataStore: WKWebsiteDataStore) {        let key = ObjectIdentifier(websiteDataStore)        guard let entry = self.entries.removeValue(forKey: key) else { return }        entry.clearPreservedPage()        Self.log.debug("OpenAI webview evicted")        entry.host.close()        self.scheduleNextIdlePrune()    }     func evictAll() {        self.cancelIdlePrune()        let existing = self.entries        self.entries.removeAll()        for (_, entry) in existing {            entry.clearPreservedPage()            entry.host.close()        }        if !existing.isEmpty {            Self.log.debug("OpenAI webview evicted all")        }    }     private func prepareCachedWebViewForIdle(        _ webView: WKWebView,        host: OffscreenWebViewHost,        preserveLoadedPage: Bool)    {        webView.navigationDelegate = nil        webView.codexNavigationDelegate = nil        if preserveLoadedPage {            host.hide()            return        }         // Detach the heavyweight ChatGPT SPA as soon as a scrape completes. Keeping the WebView object around        // still helps with immediate reuse, but letting chatgpt.com remain the active document is too expensive.        webView.stopLoading()        webView.evaluateJavaScript(self.idlePageClearScript, completionHandler: nil)        _ = webView.load(URLRequest(url: self.blankURL))        host.hide()    }     /// Schedule against the oldest idle entry so later releases cannot postpone its eviction.    private func scheduleNextIdlePrune(now: Date = Date()) {        self.cancelIdlePrune()         guard let nextExpiry = self.entries.values            .filter({ !$0.isBusy })            .map({ $0.lastUsedAt.addingTimeInterval(self.idleTimeout) })            .min()        else { return }         let generation = self.idlePruneGeneration        let workItem = DispatchWorkItem { [weak self] in            MainActor.assumeIsolated {                guard let self, self.idlePruneGeneration == generation else { return }                self.idlePruneWorkItem = nil                #if DEBUG                self.idlePruneDeadlineForTesting = nil                #endif

Lines 1–500 of 758.