From ad71da06985e65ab2664594602f09e2dcaa94a1d Mon Sep 17 00:00:00 2001 From: pickaxe <54486432+ProspectOre@users.noreply.github.com> Date: Wed, 10 Jun 2026 12:30:49 -0700 Subject: [PATCH 1/2] Keep codex account file reads off the menu-build path populateMenu -> codexAccountMenuDisplay loaded the codex account reconciliation snapshot synchronously whenever the 2s freshness cache had lapsed, paying auth.json reads, JWT parsing, and SHA256 fingerprinting on the main thread inside menu open and tracking. Menu display now tolerates a stale cached snapshot and revalidates the cache off the menu-build path; account changes land on the next rebuild. --- CHANGELOG.md | 1 + .../Providers/Codex/CodexSettingsStore.swift | 39 ++++++ Sources/CodexBar/SettingsStore.swift | 1 + ...tusItemController+AccountMenuDisplay.swift | 2 +- ...CodexAccountMenuDisplaySnapshotTests.swift | 114 ++++++++++++++++++ 5 files changed, 156 insertions(+), 1 deletion(-) create mode 100644 Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ac0eae4cc..a578f8992a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ - Menu bar: defer merged status-icon redraws until the tracked menu closes while preserving animation lifecycle and quota-warning timing, reducing WindowServer churn during long menu sessions (#1409, fixes #1399). Thanks @kiranmagic7! - Provider status: decode status feeds on the concurrent executor and reuse ISO8601 formatters, removing a measured main-thread stall during refreshes (#1406). Thanks @ProspectOre! - Menu bar: keep one stable width across merged provider tabs and resize every hosted card row to AppKit's final menu width so provider switching no longer leaves a widened menu with inset submenu arrows (#1410). +- Menu bar: keep Codex `auth.json` reads, JWT parsing, and fingerprint hashing off the menu-build path by rendering a cached account snapshot and revalidating it asynchronously (#1401). Thanks @ProspectOre! - Menu bar: defer Overview-row provider transitions out of AppKit's click callback so opening provider detail no longer performs a full synchronous menu rebuild (#1325). - Menu bar: open cached menus immediately after data-only invalidations, then refresh missing or stale provider data asynchronously without queuing redundant work on close (#1398). Thanks @joshuavial! - Menu bar: recycle SwiftUI card hosting views across data refreshes and provider switches, and reconcile matching menu rows in place instead of removing and reinserting every row, cutting open-click, switch, and idle rebuild cost (#1394). Thanks @bcssewl! diff --git a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift index 1f7987b160..b7d2650eaa 100644 --- a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift +++ b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift @@ -240,10 +240,49 @@ extension SettingsStore { return snapshot } + /// Menu rendering must not block on live `auth.json` reads, JWT parsing, and fingerprint + /// hashing. This returns the cached reconciliation snapshot even when it is older than the + /// freshness interval and refreshes the cache off the menu-build path instead, so only the + /// first menu build after launch (no cache yet) or an active-source change pays the + /// synchronous load. Account changes land on the next menu rebuild. + var codexAccountReconciliationSnapshotForMenuDisplay: CodexAccountReconciliationSnapshot { + let activeSource = self.codexPersistedActiveSource + guard Self.codexAccountReconciliationSnapshotCacheInterval > 0, + let cached = self.cachedCodexAccountReconciliationSnapshot, + cached.activeSource == activeSource + else { + return self.codexAccountReconciliationSnapshot(activeSourceOverride: nil) + } + if Date().timeIntervalSince(cached.loadedAt) >= Self.codexAccountReconciliationSnapshotCacheInterval { + self.scheduleCodexAccountReconciliationSnapshotRevalidation() + } + return cached.snapshot + } + + private func scheduleCodexAccountReconciliationSnapshotRevalidation() { + guard self.codexAccountSnapshotRevalidationTask == nil else { return } + self.codexAccountSnapshotRevalidationTask = Task { @MainActor [weak self] in + // The main dispatch queue does not drain while AppKit runs the menu-tracking run loop + // mode, so this hop keeps the reload from landing inside an open tracking session. + await withCheckedContinuation { continuation in + DispatchQueue.main.async { continuation.resume() } + } + guard let self else { return } + defer { self.codexAccountSnapshotRevalidationTask = nil } + guard !Task.isCancelled else { return } + self.invalidateCodexAccountReconciliationSnapshotCache() + _ = self.codexAccountReconciliationSnapshot(activeSourceOverride: nil) + } + } + var codexVisibleAccountProjection: CodexVisibleAccountProjection { CodexVisibleAccountProjection.make(from: self.codexAccountReconciliationSnapshot) } + var codexVisibleAccountProjectionForMenuDisplay: CodexVisibleAccountProjection { + CodexVisibleAccountProjection.make(from: self.codexAccountReconciliationSnapshotForMenuDisplay) + } + var codexVisibleAccounts: [CodexVisibleAccount] { self.codexVisibleAccountProjection.visibleAccounts } diff --git a/Sources/CodexBar/SettingsStore.swift b/Sources/CodexBar/SettingsStore.swift index c6135391ce..6bf989c264 100644 --- a/Sources/CodexBar/SettingsStore.swift +++ b/Sources/CodexBar/SettingsStore.swift @@ -140,6 +140,7 @@ final class SettingsStore { @ObservationIgnored var tokenAccountsLoaded = false @ObservationIgnored var cachedCodexAccountReconciliationSnapshot: CachedCodexAccountReconciliationSnapshot? + @ObservationIgnored var codexAccountSnapshotRevalidationTask: Task? @ObservationIgnored var mergedMenuLastSelectedWasOverviewStorage = false @ObservationIgnored var selectedMenuProviderRawStorage: String? var defaultsState: SettingsDefaultsState diff --git a/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift b/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift index ea0bb9d2df..824d5e0cf4 100644 --- a/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift +++ b/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift @@ -35,7 +35,7 @@ extension StatusItemController { func codexAccountMenuDisplay(for provider: UsageProvider) -> CodexAccountMenuDisplay? { guard provider == .codex else { return nil } - let projection = self.settings.codexVisibleAccountProjection + let projection = self.settings.codexVisibleAccountProjectionForMenuDisplay guard projection.visibleAccounts.count > 1 else { return nil } let showAll = self.settings.multiAccountMenuLayout == .stacked let accounts = showAll diff --git a/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift b/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift new file mode 100644 index 0000000000..891b0d4c13 --- /dev/null +++ b/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift @@ -0,0 +1,114 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@Suite(.serialized) +struct CodexAccountMenuDisplaySnapshotTests { + @MainActor + private static func makeSettings(suite: String) throws -> SettingsStore { + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + defaults.set(true, forKey: "providerDetectionCompleted") + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings.providerDetectionCompleted = true + return settings + } + + private static func writeCodexAuthFile(homeURL: URL, email: String, plan: String) throws { + try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true) + let auth = ["tokens": [ + "accessToken": "access-token", + "refreshToken": "refresh-token", + "idToken": Self.fakeJWT(email: email, plan: plan), + ]] + let data = try JSONSerialization.data(withJSONObject: auth) + try data.write(to: homeURL.appendingPathComponent("auth.json")) + } + + private static func fakeJWT(email: String, plan: String) -> String { + let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data() + let payload = (try? JSONSerialization.data(withJSONObject: [ + "email": email, + "chatgpt_plan_type": plan, + ])) ?? Data() + + func base64URL(_ data: Data) -> String { + data.base64EncodedString() + .replacingOccurrences(of: "=", with: "") + .replacingOccurrences(of: "+", with: "-") + .replacingOccurrences(of: "/", with: "_") + } + + return "\(base64URL(header)).\(base64URL(payload))." + } + + @Test + @MainActor + func `menu display snapshot tolerates stale cache and revalidates off the menu path`() async throws { + let suite = "CodexAccountMenuDisplaySnapshotTests-stale-cache" + let settings = try Self.makeSettings(suite: suite) + let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + try Self.writeCodexAuthFile(homeURL: ambientHome, email: "before@example.com", plan: "pro") + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": ambientHome.path] + SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60 + defer { + SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: ambientHome) + } + + let primed = settings.codexAccountReconciliationSnapshot + #expect(primed.liveSystemAccount?.email == "before@example.com") + + // Simulate a cache that has outlived the freshness interval while the auth file changed. + try Self.writeCodexAuthFile(homeURL: ambientHome, email: "after@example.com", plan: "pro") + let cached = try #require(settings.cachedCodexAccountReconciliationSnapshot) + settings.cachedCodexAccountReconciliationSnapshot = CachedCodexAccountReconciliationSnapshot( + activeSource: cached.activeSource, + loadedAt: Date(timeIntervalSinceNow: -3600), + snapshot: cached.snapshot) + + // The menu path returns the stale cache without a synchronous reload. + let menuSnapshot = settings.codexAccountReconciliationSnapshotForMenuDisplay + #expect(menuSnapshot.liveSystemAccount?.email == "before@example.com") + + // The scheduled revalidation refreshes the cache off the menu path. + let revalidation = try #require(settings.codexAccountSnapshotRevalidationTask) + await revalidation.value + #expect(settings.codexAccountSnapshotRevalidationTask == nil) + #expect( + settings.codexAccountReconciliationSnapshotForMenuDisplay.liveSystemAccount?.email == + "after@example.com") + } + + @Test + @MainActor + func `menu display snapshot loads synchronously without a cache`() throws { + let suite = "CodexAccountMenuDisplaySnapshotTests-cold-cache" + let settings = try Self.makeSettings(suite: suite) + let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + try Self.writeCodexAuthFile(homeURL: ambientHome, email: "cold@example.com", plan: "pro") + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": ambientHome.path] + SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60 + defer { + SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: ambientHome) + } + + #expect(settings.cachedCodexAccountReconciliationSnapshot == nil) + let menuSnapshot = settings.codexAccountReconciliationSnapshotForMenuDisplay + #expect(menuSnapshot.liveSystemAccount?.email == "cold@example.com") + #expect(settings.codexAccountSnapshotRevalidationTask == nil) + #expect(settings.cachedCodexAccountReconciliationSnapshot != nil) + } +} From f1c1b6cdf3eadaa8cfd0ba15aa99c78bfe3c08bc Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 11 Jun 2026 04:08:32 +0100 Subject: [PATCH 2/2] perf: keep Codex account reconciliation off menu tracking --- .../Providers/Codex/CodexSettingsStore.swift | 130 ++++-- Sources/CodexBar/SettingsStore.swift | 20 +- ...tusItemController+AccountMenuDisplay.swift | 51 ++- .../CodexBar/StatusItemController+Menu.swift | 2 + .../StatusItemController+MenuTracking.swift | 6 +- .../StatusItemController+Shutdown.swift | 2 + Sources/CodexBar/StatusItemController.swift | 3 + .../Codex/CodexAccountReconciliation.swift | 2 +- ...CodexAccountMenuDisplaySnapshotTests.swift | 376 ++++++++++++++---- .../StatusMenuCodexSwitcherTests.swift | 3 + 10 files changed, 476 insertions(+), 119 deletions(-) diff --git a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift index b7d2650eaa..a0be2ddb38 100644 --- a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift +++ b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift @@ -141,6 +141,7 @@ extension SettingsStore { } set { self.invalidateCodexAccountReconciliationSnapshotCache() + self.cachedCodexAccountMenuProjection = nil self.updateProviderConfig(provider: .codex) { entry in entry.codexActiveSource = newValue } @@ -166,6 +167,7 @@ extension SettingsStore { @discardableResult func refreshCodexAccountReconciliationAfterManagedAccountsDidChange() -> Bool { self.invalidateCodexAccountReconciliationSnapshotCache() + self.cachedCodexAccountMenuProjection = nil return self.persistResolvedCodexActiveSourceCorrectionIfNeeded() } @@ -210,6 +212,7 @@ extension SettingsStore { func invalidateCodexAccountReconciliationSnapshotCache() { self.cachedCodexAccountReconciliationSnapshot = nil + self.codexAccountReconciliationGeneration &+= 1 } var codexAccountReconciliationSnapshot: CodexAccountReconciliationSnapshot { @@ -230,57 +233,86 @@ extension SettingsStore { return cached.snapshot } - let snapshot = self.codexAccountReconciler(activeSource: activeSource).loadSnapshot() + let snapshot = self.codexAccountSnapshotLoader(activeSource: activeSource)() + let loadedAt = Date() if cacheInterval > 0 { self.cachedCodexAccountReconciliationSnapshot = CachedCodexAccountReconciliationSnapshot( activeSource: activeSource, - loadedAt: now, + loadedAt: loadedAt, snapshot: snapshot) } + if activeSource == self.codexPersistedActiveSource { + self.cachedCodexAccountMenuProjection = CachedCodexAccountMenuProjection( + activeSource: activeSource, + loadedAt: loadedAt, + projection: CodexVisibleAccountProjection.make(from: snapshot)) + } return snapshot } - /// Menu rendering must not block on live `auth.json` reads, JWT parsing, and fingerprint - /// hashing. This returns the cached reconciliation snapshot even when it is older than the - /// freshness interval and refreshes the cache off the menu-build path instead, so only the - /// first menu build after launch (no cache yet) or an active-source change pays the - /// synchronous load. Account changes land on the next menu rebuild. - var codexAccountReconciliationSnapshotForMenuDisplay: CodexAccountReconciliationSnapshot { + /// Menu rendering must stay side-effect free: no `auth.json` reads, JWT parsing, or fingerprint hashing. + var codexVisibleAccountProjectionForMenuDisplay: CodexVisibleAccountProjection? { let activeSource = self.codexPersistedActiveSource - guard Self.codexAccountReconciliationSnapshotCacheInterval > 0, - let cached = self.cachedCodexAccountReconciliationSnapshot, + guard let cached = self.cachedCodexAccountMenuProjection, cached.activeSource == activeSource else { - return self.codexAccountReconciliationSnapshot(activeSourceOverride: nil) + return nil } - if Date().timeIntervalSince(cached.loadedAt) >= Self.codexAccountReconciliationSnapshotCacheInterval { - self.scheduleCodexAccountReconciliationSnapshotRevalidation() + return cached.projection + } + + var codexAccountMenuProjectionNeedsRevalidation: Bool { + let activeSource = self.codexPersistedActiveSource + guard let cached = self.cachedCodexAccountMenuProjection, + cached.activeSource == activeSource + else { + return true } - return cached.snapshot + return Date().timeIntervalSince(cached.loadedAt) >= Self.codexAccountReconciliationSnapshotCacheInterval } - private func scheduleCodexAccountReconciliationSnapshotRevalidation() { - guard self.codexAccountSnapshotRevalidationTask == nil else { return } - self.codexAccountSnapshotRevalidationTask = Task { @MainActor [weak self] in - // The main dispatch queue does not drain while AppKit runs the menu-tracking run loop - // mode, so this hop keeps the reload from landing inside an open tracking session. - await withCheckedContinuation { continuation in - DispatchQueue.main.async { continuation.resume() } - } - guard let self else { return } - defer { self.codexAccountSnapshotRevalidationTask = nil } - guard !Task.isCancelled else { return } - self.invalidateCodexAccountReconciliationSnapshotCache() - _ = self.codexAccountReconciliationSnapshot(activeSourceOverride: nil) + func revalidateCodexAccountMenuProjection() async -> CodexAccountMenuProjectionRevalidationResult { + guard self.codexAccountMenuProjectionNeedsRevalidation else { return .skipped } + + let activeSource = self.codexPersistedActiveSource + let generation = self.codexAccountReconciliationGeneration + let loader = self.codexAccountSnapshotLoader(activeSource: activeSource) + let snapshot = await Self.loadCodexAccountSnapshot(loader) + + guard generation == self.codexAccountReconciliationGeneration, + activeSource == self.codexPersistedActiveSource + else { + return .discarded + } + + let now = Date() + let projection = CodexVisibleAccountProjection.make(from: snapshot) + let previousProjection = self.cachedCodexAccountMenuProjection.flatMap { cached in + cached.activeSource == activeSource ? cached.projection : nil + } + self.cachedCodexAccountMenuProjection = CachedCodexAccountMenuProjection( + activeSource: activeSource, + loadedAt: now, + projection: projection) + if Self.codexAccountReconciliationSnapshotCacheInterval > 0 { + self.cachedCodexAccountReconciliationSnapshot = CachedCodexAccountReconciliationSnapshot( + activeSource: activeSource, + loadedAt: now, + snapshot: snapshot) } + return previousProjection == projection ? .unchanged : .updated } - var codexVisibleAccountProjection: CodexVisibleAccountProjection { - CodexVisibleAccountProjection.make(from: self.codexAccountReconciliationSnapshot) + @concurrent + private nonisolated static func loadCodexAccountSnapshot( + _ loader: @escaping @Sendable () -> CodexAccountReconciliationSnapshot) + async -> CodexAccountReconciliationSnapshot + { + loader() } - var codexVisibleAccountProjectionForMenuDisplay: CodexVisibleAccountProjection { - CodexVisibleAccountProjection.make(from: self.codexAccountReconciliationSnapshotForMenuDisplay) + var codexVisibleAccountProjection: CodexVisibleAccountProjection { + CodexVisibleAccountProjection.make(from: self.codexAccountReconciliationSnapshot) } var codexVisibleAccounts: [CodexVisibleAccount] { @@ -296,11 +328,8 @@ extension SettingsStore { } func selectDisplayedCodexVisibleAccount(_ account: CodexVisibleAccount) { - if self.selectCodexVisibleAccount(id: account.id) { - return - } - // An open menu can preserve a previously rendered account row while the live projection is briefly incomplete. - self.invalidateCodexAccountReconciliationSnapshotCache() + // The row already carries the exact source it represented. Re-resolving its ID would synchronously + // reload auth state from the menu click callback and can also fail after a stale snapshot is rendered. self.codexActiveSource = account.selectionSource } @@ -322,6 +351,18 @@ extension SettingsStore { self.codexVisibleAccountProjection.source(forVisibleAccountID: id) } + private func codexAccountSnapshotLoader( + activeSource: CodexActiveSource) -> @Sendable () -> CodexAccountReconciliationSnapshot + { + #if DEBUG + if let loader = self._test_codexAccountSnapshotLoader { + return { loader(activeSource) } + } + #endif + let reconciler = self.codexAccountReconciler(activeSource: activeSource) + return { reconciler.loadSnapshot() } + } + private func codexAccountReconciler(activeSource: CodexActiveSource) -> DefaultCodexAccountReconciler { let baseEnvironment = self.codexReconciliationEnvironment() #if DEBUG @@ -557,10 +598,15 @@ private struct CodexManagedRemoteHomeTestingSystemObserver: CodexSystemAccountOb } extension SettingsStore { + private func invalidateCodexAccountReconciliationCachesForTesting() { + self.invalidateCodexAccountReconciliationSnapshotCache() + self.cachedCodexAccountMenuProjection = nil + } + var _test_activeManagedCodexRemoteHomePath: String? { get { CodexManagedRemoteHomeTestingOverride.homePath(for: self) } set { - self.invalidateCodexAccountReconciliationSnapshotCache() + self.invalidateCodexAccountReconciliationCachesForTesting() CodexManagedRemoteHomeTestingOverride.setHomePath(newValue, for: self) } } @@ -568,7 +614,7 @@ extension SettingsStore { var _test_activeManagedCodexAccount: ManagedCodexAccount? { get { CodexManagedRemoteHomeTestingOverride.account(for: self) } set { - self.invalidateCodexAccountReconciliationSnapshotCache() + self.invalidateCodexAccountReconciliationCachesForTesting() CodexManagedRemoteHomeTestingOverride.setAccount(newValue, for: self) } } @@ -576,7 +622,7 @@ extension SettingsStore { var _test_unreadableManagedCodexAccountStore: Bool { get { CodexManagedRemoteHomeTestingOverride.isUnreadable(for: self) } set { - self.invalidateCodexAccountReconciliationSnapshotCache() + self.invalidateCodexAccountReconciliationCachesForTesting() CodexManagedRemoteHomeTestingOverride.setUnreadable(newValue, for: self) } } @@ -584,7 +630,7 @@ extension SettingsStore { var _test_managedCodexAccountStoreURL: URL? { get { CodexManagedRemoteHomeTestingOverride.managedStoreURL(for: self) } set { - self.invalidateCodexAccountReconciliationSnapshotCache() + self.invalidateCodexAccountReconciliationCachesForTesting() CodexManagedRemoteHomeTestingOverride.setManagedStoreURL(newValue, for: self) } } @@ -592,7 +638,7 @@ extension SettingsStore { var _test_liveSystemCodexAccount: ObservedSystemCodexAccount? { get { CodexManagedRemoteHomeTestingOverride.liveSystemAccount(for: self) } set { - self.invalidateCodexAccountReconciliationSnapshotCache() + self.invalidateCodexAccountReconciliationCachesForTesting() CodexManagedRemoteHomeTestingOverride.setLiveSystemAccount(newValue, for: self) } } @@ -600,7 +646,7 @@ extension SettingsStore { var _test_codexReconciliationEnvironment: [String: String]? { get { CodexManagedRemoteHomeTestingOverride.reconciliationEnvironment(for: self) } set { - self.invalidateCodexAccountReconciliationSnapshotCache() + self.invalidateCodexAccountReconciliationCachesForTesting() CodexManagedRemoteHomeTestingOverride.setReconciliationEnvironment(newValue, for: self) } } diff --git a/Sources/CodexBar/SettingsStore.swift b/Sources/CodexBar/SettingsStore.swift index 6bf989c264..7a05a305bc 100644 --- a/Sources/CodexBar/SettingsStore.swift +++ b/Sources/CodexBar/SettingsStore.swift @@ -114,6 +114,19 @@ struct CachedCodexAccountReconciliationSnapshot { let snapshot: CodexAccountReconciliationSnapshot } +struct CachedCodexAccountMenuProjection: Equatable { + let activeSource: CodexActiveSource + let loadedAt: Date + let projection: CodexVisibleAccountProjection +} + +enum CodexAccountMenuProjectionRevalidationResult: Equatable { + case skipped + case discarded + case unchanged + case updated +} + @MainActor @Observable final class SettingsStore { @@ -140,7 +153,12 @@ final class SettingsStore { @ObservationIgnored var tokenAccountsLoaded = false @ObservationIgnored var cachedCodexAccountReconciliationSnapshot: CachedCodexAccountReconciliationSnapshot? - @ObservationIgnored var codexAccountSnapshotRevalidationTask: Task? + @ObservationIgnored var cachedCodexAccountMenuProjection: CachedCodexAccountMenuProjection? + @ObservationIgnored var codexAccountReconciliationGeneration: UInt = 0 + #if DEBUG + @ObservationIgnored var _test_codexAccountSnapshotLoader: + (@Sendable (CodexActiveSource) -> CodexAccountReconciliationSnapshot)? + #endif @ObservationIgnored var mergedMenuLastSelectedWasOverviewStorage = false @ObservationIgnored var selectedMenuProviderRawStorage: String? var defaultsState: SettingsDefaultsState diff --git a/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift b/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift index 824d5e0cf4..2c268511f4 100644 --- a/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift +++ b/Sources/CodexBar/StatusItemController+AccountMenuDisplay.swift @@ -2,6 +2,30 @@ import AppKit import CodexBarCore extension StatusItemController { + private static let defaultCodexAccountMenuProjectionRevalidationEnabled = !SettingsStore.isRunningTests + + #if DEBUG + private static var codexAccountMenuProjectionRevalidationEnabledForTesting = + defaultCodexAccountMenuProjectionRevalidationEnabled + + static func setCodexAccountMenuProjectionRevalidationEnabledForTesting(_ enabled: Bool) { + self.codexAccountMenuProjectionRevalidationEnabledForTesting = enabled + } + + static func resetCodexAccountMenuProjectionRevalidationEnabledForTesting() { + self.codexAccountMenuProjectionRevalidationEnabledForTesting = + self.defaultCodexAccountMenuProjectionRevalidationEnabled + } + #endif + + private static var codexAccountMenuProjectionRevalidationEnabled: Bool { + #if DEBUG + self.codexAccountMenuProjectionRevalidationEnabledForTesting + #else + self.defaultCodexAccountMenuProjectionRevalidationEnabled + #endif + } + func tokenAccountMenuDisplay(for provider: UsageProvider) -> TokenAccountMenuDisplay? { guard TokenAccountSupportCatalog.support(for: provider) != nil else { return nil } let accounts = self.settings.tokenAccounts(for: provider) @@ -35,7 +59,7 @@ extension StatusItemController { func codexAccountMenuDisplay(for provider: UsageProvider) -> CodexAccountMenuDisplay? { guard provider == .codex else { return nil } - let projection = self.settings.codexVisibleAccountProjectionForMenuDisplay + guard let projection = self.settings.codexVisibleAccountProjectionForMenuDisplay else { return nil } guard projection.visibleAccounts.count > 1 else { return nil } let showAll = self.settings.multiAccountMenuLayout == .stacked let accounts = showAll @@ -52,6 +76,31 @@ extension StatusItemController { layout: showAll ? .stacked : .segmented) } + func scheduleCodexAccountMenuProjectionRevalidationIfNeeded(for providers: [UsageProvider]) { + guard Self.codexAccountMenuProjectionRevalidationEnabled else { return } + guard providers.contains(.codex) else { return } + guard self.settings.codexAccountMenuProjectionNeedsRevalidation else { return } + guard self.codexAccountMenuProjectionRevalidationTask == nil else { return } + + self.codexAccountMenuProjectionRevalidationTask = Task { @MainActor [weak self] in + guard let settings = self?.settings else { return } + let result = await settings.revalidateCodexAccountMenuProjection() + guard let self else { return } + guard !Task.isCancelled else { + self.codexAccountMenuProjectionRevalidationTask = nil + return + } + self.codexAccountMenuProjectionRevalidationTask = nil + + switch result { + case .updated: + self.invalidateMenus(refreshOpenMenus: false) + case .discarded, .skipped, .unchanged: + break + } + } + } + private func codexAccountSnapshots(matching accounts: [CodexVisibleAccount]) -> [CodexAccountUsageSnapshot] { var snapshotsByID: [String: CodexAccountUsageSnapshot] = [:] for snapshot in self.store.codexAccountSnapshots { diff --git a/Sources/CodexBar/StatusItemController+Menu.swift b/Sources/CodexBar/StatusItemController+Menu.swift index 2e4f9c520c..d1a491b794 100644 --- a/Sources/CodexBar/StatusItemController+Menu.swift +++ b/Sources/CodexBar/StatusItemController+Menu.swift @@ -121,6 +121,8 @@ extension StatusItemController { let menuWasFreshBeforeOpen = !self.menuNeedsRefresh(menu) self.refreshMenuForOpenIfNeeded(menu, provider: provider) + self.scheduleCodexAccountMenuProjectionRevalidationIfNeeded( + for: self.renderedProviders(for: menu)) if self.isMenuRefreshEnabled { // Intentionally skip open-menu tracking when refresh is disabled (tests). // If refresh is re-enabled while this menu stays open, it will not be backfilled until next open. diff --git a/Sources/CodexBar/StatusItemController+MenuTracking.swift b/Sources/CodexBar/StatusItemController+MenuTracking.swift index f7acd4a4cd..7601eeb031 100644 --- a/Sources/CodexBar/StatusItemController+MenuTracking.swift +++ b/Sources/CodexBar/StatusItemController+MenuTracking.swift @@ -303,7 +303,7 @@ extension StatusItemController { parts.append(target.rawValue) parts.append(self.providerIdentitySignature(self.store.snapshot(for: target)?.identity(for: target))) - if self.store.metadata(for: target).usesAccountFallback { + if target != .codex, self.store.metadata(for: target).usesAccountFallback { let account = self.store.accountInfo(for: target) parts.append(Self.menuIdentityField(account.email)) parts.append(Self.menuIdentityField(account.plan)) @@ -316,7 +316,9 @@ extension StatusItemController { } if target == .codex { - for account in self.settings.codexVisibleAccountProjection.visibleAccounts { + parts.append(Self.menuIdentityField(self.account.email)) + parts.append(Self.menuIdentityField(self.account.plan)) + for account in self.settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts ?? [] { parts.append(Self.menuIdentityField(account.id)) parts.append(Self.menuIdentityField(account.email)) parts.append(Self.menuIdentityField(account.workspaceLabel)) diff --git a/Sources/CodexBar/StatusItemController+Shutdown.swift b/Sources/CodexBar/StatusItemController+Shutdown.swift index 5ca3f397bf..6494086c96 100644 --- a/Sources/CodexBar/StatusItemController+Shutdown.swift +++ b/Sources/CodexBar/StatusItemController+Shutdown.swift @@ -52,6 +52,8 @@ extension StatusItemController { } self.openMenuInvalidationRetryTask?.cancel() self.openMenuInvalidationRetryTask = nil + self.codexAccountMenuProjectionRevalidationTask?.cancel() + self.codexAccountMenuProjectionRevalidationTask = nil self.providerSelectionUIRefreshTask?.cancel() self.providerSelectionUIRefreshTask = nil self.deferredMergedIconRenderAfterTracking = false diff --git a/Sources/CodexBar/StatusItemController.swift b/Sources/CodexBar/StatusItemController.swift index ca719199ad..60ad9d553a 100644 --- a/Sources/CodexBar/StatusItemController.swift +++ b/Sources/CodexBar/StatusItemController.swift @@ -139,6 +139,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin var openMenuRebuildTokens: [ObjectIdentifier: Int] = [:] var openMenuRebuildTokenCounter = 0 var menuIdentitySignatures: [ObjectIdentifier: String] = [:] + var codexAccountMenuProjectionRevalidationTask: Task? var openMenuRebuildsClosingHostedSubviewMenus: Set = [] var parentMenuRebuildsDeferredDuringTracking: Set = [] var deferredMenuInteractionRefreshProviders: Set = [] @@ -415,6 +416,8 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin self.wireBindings() self.updateVisibility() self.updateIcons() + self.scheduleCodexAccountMenuProjectionRevalidationIfNeeded( + for: self.store.enabledProvidersForDisplay()) self.scheduleStartupStatusItemVisibilityCheck() NotificationCenter.default.addObserver( self, diff --git a/Sources/CodexBarCore/Providers/Codex/CodexAccountReconciliation.swift b/Sources/CodexBarCore/Providers/Codex/CodexAccountReconciliation.swift index 25dcae7ac5..3b2f8a0758 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexAccountReconciliation.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexAccountReconciliation.swift @@ -119,7 +119,7 @@ public struct CodexAccountReconciliationSnapshot: Equatable, Sendable { } } -public struct DefaultCodexAccountReconciler { +public struct DefaultCodexAccountReconciler: Sendable { public let storeLoader: @Sendable () throws -> ManagedCodexAccountSet public let systemObserver: any CodexSystemAccountObserving public let activeSource: CodexActiveSource diff --git a/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift b/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift index 891b0d4c13..5941c8b000 100644 --- a/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift +++ b/Tests/CodexBarTests/CodexAccountMenuDisplaySnapshotTests.swift @@ -1,13 +1,15 @@ +import AppKit import CodexBarCore import Foundation import Testing @testable import CodexBar @Suite(.serialized) +@MainActor struct CodexAccountMenuDisplaySnapshotTests { - @MainActor - private static func makeSettings(suite: String) throws -> SettingsStore { - let defaults = try #require(UserDefaults(suiteName: suite)) + private func makeSettings() -> SettingsStore { + let suite = "CodexAccountMenuDisplaySnapshotTests-\(UUID().uuidString)" + let defaults = UserDefaults(suiteName: suite)! defaults.removePersistentDomain(forName: suite) defaults.set(true, forKey: "providerDetectionCompleted") let settings = SettingsStore( @@ -19,96 +21,326 @@ struct CodexAccountMenuDisplaySnapshotTests { return settings } - private static func writeCodexAuthFile(homeURL: URL, email: String, plan: String) throws { - try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true) - let auth = ["tokens": [ - "accessToken": "access-token", - "refreshToken": "refresh-token", - "idToken": Self.fakeJWT(email: email, plan: plan), - ]] - let data = try JSONSerialization.data(withJSONObject: auth) - try data.write(to: homeURL.appendingPathComponent("auth.json")) - } - - private static func fakeJWT(email: String, plan: String) -> String { - let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data() - let payload = (try? JSONSerialization.data(withJSONObject: [ - "email": email, - "chatgpt_plan_type": plan, - ])) ?? Data() - - func base64URL(_ data: Data) -> String { - data.base64EncodedString() - .replacingOccurrences(of: "=", with: "") - .replacingOccurrences(of: "+", with: "-") - .replacingOccurrences(of: "/", with: "_") + private func enableOnlyCodex(_ settings: SettingsStore) { + for provider in UsageProvider.allCases { + guard let metadata = ProviderRegistry.shared.metadata[provider] else { continue } + settings.setProviderEnabled(provider: provider, metadata: metadata, enabled: provider == .codex) } + } + + private func liveSnapshot(email: String) -> CodexAccountReconciliationSnapshot { + CodexAccountReconciliationSnapshot( + storedAccounts: [], + activeStoredAccount: nil, + liveSystemAccount: ObservedSystemCodexAccount( + email: email, + codexHomePath: "/tmp/\(email)", + observedAt: Date()), + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .liveSystem, + hasUnreadableAddedAccountStore: false) + } + + private func cachedProjection( + snapshot: CodexAccountReconciliationSnapshot, + loadedAt: Date = Date(timeIntervalSinceNow: -3600)) -> CachedCodexAccountMenuProjection + { + CachedCodexAccountMenuProjection( + activeSource: snapshot.activeSource, + loadedAt: loadedAt, + projection: CodexVisibleAccountProjection.make(from: snapshot)) + } + + @Test + func `cold menu projection read never loads auth state`() async { + let settings = self.makeSettings() + let probe = CodexAccountSnapshotLoaderProbe(snapshot: self.liveSnapshot(email: "loaded@example.com")) + settings._test_codexAccountSnapshotLoader = { _ in probe.load() } + defer { settings._test_codexAccountSnapshotLoader = nil } + + #expect(settings.codexVisibleAccountProjectionForMenuDisplay == nil) + #expect(probe.callCount == 0) + + let result = await settings.revalidateCodexAccountMenuProjection() - return "\(base64URL(header)).\(base64URL(payload))." + #expect(result == .updated) + #expect(probe.callCount == 1) + #expect(probe.loadedOffMainThread) + #expect( + settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email == + "loaded@example.com") } @Test - @MainActor - func `menu display snapshot tolerates stale cache and revalidates off the menu path`() async throws { - let suite = "CodexAccountMenuDisplaySnapshotTests-stale-cache" - let settings = try Self.makeSettings(suite: suite) - let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( - UUID().uuidString, - isDirectory: true) - try Self.writeCodexAuthFile(homeURL: ambientHome, email: "before@example.com", plan: "pro") - settings._test_codexReconciliationEnvironment = ["CODEX_HOME": ambientHome.path] + func `override snapshot load preserves persisted account menu projection`() { + let settings = self.makeSettings() + let activeSnapshot = self.liveSnapshot(email: "active@example.com") + settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: activeSnapshot) + + let otherID = UUID() + let otherAccount = ManagedCodexAccount( + id: otherID, + email: "other@example.com", + managedHomePath: "/tmp/other", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let overrideSnapshot = CodexAccountReconciliationSnapshot( + storedAccounts: [otherAccount], + activeStoredAccount: otherAccount, + liveSystemAccount: nil, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .managedAccount(id: otherID), + hasUnreadableAddedAccountStore: false) + settings._test_codexAccountSnapshotLoader = { _ in overrideSnapshot } + defer { settings._test_codexAccountSnapshotLoader = nil } + + _ = settings.codexAccountReconciliationSnapshot(activeSourceOverride: .managedAccount(id: otherID)) + + #expect( + settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email == + "active@example.com") + } + + @Test + func `managed account change refreshes account menu projection`() { + let settings = self.makeSettings() + let activeSnapshot = self.liveSnapshot(email: "active@example.com") + settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: activeSnapshot, loadedAt: Date()) + + let addedAccount = ManagedCodexAccount( + id: UUID(), + email: "added@example.com", + managedHomePath: "/tmp/added", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let refreshedSnapshot = CodexAccountReconciliationSnapshot( + storedAccounts: [addedAccount], + activeStoredAccount: nil, + liveSystemAccount: activeSnapshot.liveSystemAccount, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .liveSystem, + hasUnreadableAddedAccountStore: false) + settings._test_codexAccountSnapshotLoader = { _ in refreshedSnapshot } + defer { settings._test_codexAccountSnapshotLoader = nil } + + settings.refreshCodexAccountReconciliationAfterManagedAccountsDidChange() + + #expect( + settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.contains { + $0.email == "added@example.com" + } == true) + } + + @Test + func `stale menu projection returns immediately then refreshes concurrently`() async { + let settings = self.makeSettings() + let staleSnapshot = self.liveSnapshot(email: "before@example.com") + let probe = CodexAccountSnapshotLoaderProbe(snapshot: self.liveSnapshot(email: "after@example.com")) + settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: staleSnapshot) + settings._test_codexAccountSnapshotLoader = { _ in probe.load() } SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60 defer { SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil - settings._test_codexReconciliationEnvironment = nil - try? FileManager.default.removeItem(at: ambientHome) + settings._test_codexAccountSnapshotLoader = nil } - let primed = settings.codexAccountReconciliationSnapshot - #expect(primed.liveSystemAccount?.email == "before@example.com") - - // Simulate a cache that has outlived the freshness interval while the auth file changed. - try Self.writeCodexAuthFile(homeURL: ambientHome, email: "after@example.com", plan: "pro") - let cached = try #require(settings.cachedCodexAccountReconciliationSnapshot) - settings.cachedCodexAccountReconciliationSnapshot = CachedCodexAccountReconciliationSnapshot( - activeSource: cached.activeSource, - loadedAt: Date(timeIntervalSinceNow: -3600), - snapshot: cached.snapshot) + #expect( + settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email == + "before@example.com") + #expect(probe.callCount == 0) + #expect(settings.codexAccountMenuProjectionNeedsRevalidation) - // The menu path returns the stale cache without a synchronous reload. - let menuSnapshot = settings.codexAccountReconciliationSnapshotForMenuDisplay - #expect(menuSnapshot.liveSystemAccount?.email == "before@example.com") + let result = await settings.revalidateCodexAccountMenuProjection() - // The scheduled revalidation refreshes the cache off the menu path. - let revalidation = try #require(settings.codexAccountSnapshotRevalidationTask) - await revalidation.value - #expect(settings.codexAccountSnapshotRevalidationTask == nil) + #expect(result == .updated) + #expect(probe.callCount == 1) + #expect(probe.loadedOffMainThread) #expect( - settings.codexAccountReconciliationSnapshotForMenuDisplay.liveSystemAccount?.email == + settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email == "after@example.com") } @Test - @MainActor - func `menu display snapshot loads synchronously without a cache`() throws { - let suite = "CodexAccountMenuDisplaySnapshotTests-cold-cache" - let settings = try Self.makeSettings(suite: suite) - let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( - UUID().uuidString, - isDirectory: true) - try Self.writeCodexAuthFile(homeURL: ambientHome, email: "cold@example.com", plan: "pro") - settings._test_codexReconciliationEnvironment = ["CODEX_HOME": ambientHome.path] + func `revalidation discards result after reconciliation generation changes`() async { + let settings = self.makeSettings() + let staleSnapshot = self.liveSnapshot(email: "before@example.com") + let probe = CodexAccountSnapshotLoaderProbe( + snapshot: self.liveSnapshot(email: "discarded@example.com"), + blocks: true) + settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: staleSnapshot) + settings._test_codexAccountSnapshotLoader = { _ in probe.load() } SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60 defer { + probe.release() SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil - settings._test_codexReconciliationEnvironment = nil - try? FileManager.default.removeItem(at: ambientHome) + settings._test_codexAccountSnapshotLoader = nil + } + + let task = Task { await settings.revalidateCodexAccountMenuProjection() } + await probe.waitUntilCalled() + settings.invalidateCodexAccountReconciliationSnapshotCache() + probe.release() + + #expect(await task.value == .discarded) + #expect( + settings.codexVisibleAccountProjectionForMenuDisplay?.visibleAccounts.first?.email == + "before@example.com") + } + + @Test + func `fresh menu open coalesces account projection revalidation and identity stays read only`() async throws { + StatusItemController.menuCardRenderingEnabled = false + StatusItemController.setMenuRefreshEnabledForTesting(false) + StatusItemController.setCodexAccountMenuProjectionRevalidationEnabledForTesting(true) + defer { + StatusItemController.resetCodexAccountMenuProjectionRevalidationEnabledForTesting() + StatusItemController.resetMenuRefreshEnabledForTesting() } - #expect(settings.cachedCodexAccountReconciliationSnapshot == nil) - let menuSnapshot = settings.codexAccountReconciliationSnapshotForMenuDisplay - #expect(menuSnapshot.liveSystemAccount?.email == "cold@example.com") - #expect(settings.codexAccountSnapshotRevalidationTask == nil) - #expect(settings.cachedCodexAccountReconciliationSnapshot != nil) + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodex(settings) + let store = UsageStore( + fetcher: UsageFetcher(), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: AccountInfo(email: nil, plan: nil), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: .system) + defer { controller.releaseStatusItemsForTesting() } + + let staleSnapshot = self.liveSnapshot(email: "before@example.com") + let probe = CodexAccountSnapshotLoaderProbe( + snapshot: self.liveSnapshot(email: "after@example.com"), + blocks: true) + settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: staleSnapshot) + settings._test_codexAccountSnapshotLoader = { _ in probe.load() } + SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = 60 + defer { + probe.release() + SettingsStore.codexAccountReconciliationSnapshotCacheIntervalOverrideForTesting = nil + settings._test_codexAccountSnapshotLoader = nil + } + + let menu = NSMenu() + controller.menuProviders[ObjectIdentifier(menu)] = .codex + controller.markMenuFresh(menu) + #expect(controller.codexAccountMenuDisplay(for: .codex) == nil) + #expect(probe.callCount == 0) + + let versionBeforeOpen = controller.menuContentVersion + controller.menuWillOpen(menu) + let revalidation = try #require(controller.codexAccountMenuProjectionRevalidationTask) + controller.menuWillOpen(menu) + await probe.waitUntilCalled() + + #expect(probe.callCount == 1) + #expect(probe.loadedOffMainThread) + probe.release() + await revalidation.value + + #expect(controller.codexAccountMenuProjectionRevalidationTask == nil) + #expect(controller.menuContentVersion == versionBeforeOpen + 1) + } + + @Test + func `selecting displayed account uses captured source without reconciliation`() throws { + let settings = self.makeSettings() + let firstID = UUID() + let secondID = UUID() + let first = ManagedCodexAccount( + id: firstID, + email: "first@example.com", + managedHomePath: "/tmp/first", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let second = ManagedCodexAccount( + id: secondID, + email: "second@example.com", + managedHomePath: "/tmp/second", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let snapshot = CodexAccountReconciliationSnapshot( + storedAccounts: [first, second], + activeStoredAccount: first, + liveSystemAccount: nil, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .managedAccount(id: firstID), + hasUnreadableAddedAccountStore: false) + let projection = CodexVisibleAccountProjection.make(from: snapshot) + let displayedAccount = try #require(projection.visibleAccounts.first { + $0.selectionSource == .managedAccount(id: secondID) + }) + let probe = CodexAccountSnapshotLoaderProbe(snapshot: snapshot) + settings.cachedCodexAccountMenuProjection = self.cachedProjection(snapshot: snapshot) + settings._test_codexAccountSnapshotLoader = { _ in probe.load() } + defer { settings._test_codexAccountSnapshotLoader = nil } + + settings.selectDisplayedCodexVisibleAccount(displayedAccount) + + #expect(probe.callCount == 0) + #expect(settings.codexActiveSource == .managedAccount(id: secondID)) + #expect(settings.cachedCodexAccountMenuProjection == nil) + } +} + +private final class CodexAccountSnapshotLoaderProbe: @unchecked Sendable { + private let lock = NSLock() + private let snapshot: CodexAccountReconciliationSnapshot + private let blocks: Bool + private let releaseSemaphore = DispatchSemaphore(value: 0) + private var _callCount = 0 + private var _loadedOffMainThread = false + private var released = false + + init(snapshot: CodexAccountReconciliationSnapshot, blocks: Bool = false) { + self.snapshot = snapshot + self.blocks = blocks + } + + var callCount: Int { + self.lock.withLock { self._callCount } + } + + var loadedOffMainThread: Bool { + self.lock.withLock { self._loadedOffMainThread } + } + + func load() -> CodexAccountReconciliationSnapshot { + self.lock.withLock { + self._callCount += 1 + self._loadedOffMainThread = self._loadedOffMainThread || !Thread.isMainThread + } + if self.blocks { + self.releaseSemaphore.wait() + } + return self.snapshot + } + + func waitUntilCalled() async { + while self.callCount == 0 { + await Task.yield() + } + } + + func release() { + let shouldSignal = self.lock.withLock { + guard !self.released else { return false } + self.released = true + return true + } + if shouldSignal { + self.releaseSemaphore.signal() + } } } diff --git a/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift index 4ad6ce8033..c7feea4789 100644 --- a/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift +++ b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift @@ -223,6 +223,7 @@ struct StatusMenuCodexSwitcherTests { codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem + _ = settings.codexVisibleAccountProjection let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) @@ -285,6 +286,7 @@ struct StatusMenuCodexSwitcherTests { codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem + _ = settings.codexVisibleAccountProjection let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) @@ -449,6 +451,7 @@ struct StatusMenuCodexSwitcherTests { codexHomePath: "/Users/test/.codex", observedAt: Date()) settings.codexActiveSource = .liveSystem + _ = settings.codexVisibleAccountProjection let fetcher = UsageFetcher() let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings)