diff --git a/Sources/CodexBar/CodexAccountReconciliation.swift b/Sources/CodexBar/CodexAccountReconciliation.swift new file mode 100644 index 000000000..3b0683b97 --- /dev/null +++ b/Sources/CodexBar/CodexAccountReconciliation.swift @@ -0,0 +1,259 @@ +import CodexBarCore +import Foundation + +struct CodexVisibleAccount: Equatable, Sendable, Identifiable { + let id: String + let email: String + let storedAccountID: UUID? + let selectionSource: CodexActiveSource + let isActive: Bool + let isLive: Bool + let canReauthenticate: Bool + let canRemove: Bool +} + +struct CodexVisibleAccountProjection: Equatable, Sendable { + let visibleAccounts: [CodexVisibleAccount] + let activeVisibleAccountID: String? + let liveVisibleAccountID: String? + let hasUnreadableAddedAccountStore: Bool + + func source(forVisibleAccountID id: String) -> CodexActiveSource? { + self.visibleAccounts.first { $0.id == id }?.selectionSource + } +} + +struct CodexResolvedActiveSource: Equatable, Sendable { + let persistedSource: CodexActiveSource + let resolvedSource: CodexActiveSource + + var requiresPersistenceCorrection: Bool { + self.persistedSource != self.resolvedSource + } +} + +enum CodexActiveSourceResolver { + static func resolve(from snapshot: CodexAccountReconciliationSnapshot) -> CodexResolvedActiveSource { + let persistedSource = snapshot.activeSource + let resolvedSource: CodexActiveSource = switch persistedSource { + case .liveSystem: + .liveSystem + case let .managedAccount(id): + if let activeStoredAccount = snapshot.activeStoredAccount { + self.matchesLiveSystemAccountEmail( + storedAccount: activeStoredAccount, + liveSystemAccount: snapshot.liveSystemAccount) ? .liveSystem : .managedAccount(id: id) + } else { + snapshot.liveSystemAccount != nil ? .liveSystem : .managedAccount(id: id) + } + } + + return CodexResolvedActiveSource( + persistedSource: persistedSource, + resolvedSource: resolvedSource) + } + + private static func matchesLiveSystemAccountEmail( + storedAccount: ManagedCodexAccount, + liveSystemAccount: ObservedSystemCodexAccount?) -> Bool + { + guard let liveSystemAccount else { return false } + return Self.normalizeEmail(storedAccount.email) == Self.normalizeEmail(liveSystemAccount.email) + } + + private static func normalizeEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } +} + +struct CodexAccountReconciliationSnapshot: Equatable, Sendable { + let storedAccounts: [ManagedCodexAccount] + let activeStoredAccount: ManagedCodexAccount? + let liveSystemAccount: ObservedSystemCodexAccount? + let matchingStoredAccountForLiveSystemAccount: ManagedCodexAccount? + let activeSource: CodexActiveSource + let hasUnreadableAddedAccountStore: Bool + + static func == (lhs: CodexAccountReconciliationSnapshot, rhs: CodexAccountReconciliationSnapshot) -> Bool { + lhs.storedAccounts.map(AccountIdentity.init) == rhs.storedAccounts.map(AccountIdentity.init) + && lhs.activeStoredAccount.map(AccountIdentity.init) == rhs.activeStoredAccount.map(AccountIdentity.init) + && lhs.liveSystemAccount == rhs.liveSystemAccount + && lhs.matchingStoredAccountForLiveSystemAccount.map(AccountIdentity.init) + == rhs.matchingStoredAccountForLiveSystemAccount.map(AccountIdentity.init) + && lhs.activeSource == rhs.activeSource + && lhs.hasUnreadableAddedAccountStore == rhs.hasUnreadableAddedAccountStore + } +} + +struct DefaultCodexAccountReconciler { + let storeLoader: @Sendable () throws -> ManagedCodexAccountSet + let systemObserver: any CodexSystemAccountObserving + let activeSource: CodexActiveSource + + init( + storeLoader: @escaping @Sendable () throws -> ManagedCodexAccountSet = { + try FileManagedCodexAccountStore().loadAccounts() + }, + systemObserver: any CodexSystemAccountObserving = DefaultCodexSystemAccountObserver(), + activeSource: CodexActiveSource = .liveSystem) + { + self.storeLoader = storeLoader + self.systemObserver = systemObserver + self.activeSource = activeSource + } + + func loadSnapshot(environment: [String: String]) -> CodexAccountReconciliationSnapshot { + let liveSystemAccount = self.loadLiveSystemAccount(environment: environment) + + do { + let accounts = try self.storeLoader() + let activeStoredAccount: ManagedCodexAccount? = switch self.activeSource { + case let .managedAccount(id): + accounts.account(id: id) + case .liveSystem: + nil + } + let matchingStoredAccountForLiveSystemAccount = liveSystemAccount.flatMap { + accounts.account(email: $0.email) + } + + return CodexAccountReconciliationSnapshot( + storedAccounts: accounts.accounts, + activeStoredAccount: activeStoredAccount, + liveSystemAccount: liveSystemAccount, + matchingStoredAccountForLiveSystemAccount: matchingStoredAccountForLiveSystemAccount, + activeSource: self.activeSource, + hasUnreadableAddedAccountStore: false) + } catch { + return CodexAccountReconciliationSnapshot( + storedAccounts: [], + activeStoredAccount: nil, + liveSystemAccount: liveSystemAccount, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: self.activeSource, + hasUnreadableAddedAccountStore: true) + } + } + + func loadVisibleAccounts(environment: [String: String]) -> CodexVisibleAccountProjection { + CodexVisibleAccountProjection.make(from: self.loadSnapshot(environment: environment)) + } + + private func loadLiveSystemAccount(environment: [String: String]) -> ObservedSystemCodexAccount? { + do { + guard let account = try self.systemObserver.loadSystemAccount(environment: environment) else { + return nil + } + let normalizedEmail = Self.normalizeEmail(account.email) + guard !normalizedEmail.isEmpty else { + return nil + } + return ObservedSystemCodexAccount( + email: normalizedEmail, + codexHomePath: account.codexHomePath, + observedAt: account.observedAt) + } catch { + return nil + } + } + + private static func normalizeEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } +} + +extension CodexVisibleAccountProjection { + static func make(from snapshot: CodexAccountReconciliationSnapshot) -> CodexVisibleAccountProjection { + let resolvedActiveSource = CodexActiveSourceResolver.resolve(from: snapshot).resolvedSource + var visibleByEmail: [String: CodexVisibleAccount] = [:] + + for storedAccount in snapshot.storedAccounts { + let normalizedEmail = Self.normalizeVisibleEmail(storedAccount.email) + visibleByEmail[normalizedEmail] = CodexVisibleAccount( + id: normalizedEmail, + email: normalizedEmail, + storedAccountID: storedAccount.id, + selectionSource: .managedAccount(id: storedAccount.id), + isActive: false, + isLive: false, + canReauthenticate: true, + canRemove: true) + } + + if let liveSystemAccount = snapshot.liveSystemAccount { + let normalizedEmail = Self.normalizeVisibleEmail(liveSystemAccount.email) + if let existing = visibleByEmail[normalizedEmail] { + visibleByEmail[normalizedEmail] = CodexVisibleAccount( + id: existing.id, + email: existing.email, + storedAccountID: existing.storedAccountID, + selectionSource: .liveSystem, + isActive: existing.isActive, + isLive: true, + canReauthenticate: existing.canReauthenticate, + canRemove: existing.canRemove) + } else { + visibleByEmail[normalizedEmail] = CodexVisibleAccount( + id: normalizedEmail, + email: normalizedEmail, + storedAccountID: nil, + selectionSource: .liveSystem, + isActive: false, + isLive: true, + canReauthenticate: true, + canRemove: false) + } + } + + let activeEmail: String? = switch resolvedActiveSource { + case let .managedAccount(id): + snapshot.storedAccounts.first { $0.id == id }.map { Self.normalizeVisibleEmail($0.email) } + case .liveSystem: + snapshot.liveSystemAccount.map { Self.normalizeVisibleEmail($0.email) } + } + + if let activeEmail, let current = visibleByEmail[activeEmail] { + visibleByEmail[activeEmail] = CodexVisibleAccount( + id: current.id, + email: current.email, + storedAccountID: current.storedAccountID, + selectionSource: current.selectionSource, + isActive: true, + isLive: current.isLive, + canReauthenticate: current.canReauthenticate, + canRemove: current.canRemove) + } + + let visibleAccounts = visibleByEmail.values.sorted { lhs, rhs in + lhs.email < rhs.email + } + + return CodexVisibleAccountProjection( + visibleAccounts: visibleAccounts, + activeVisibleAccountID: visibleAccounts.first { $0.isActive }?.id, + liveVisibleAccountID: visibleAccounts.first { $0.isLive }?.id, + hasUnreadableAddedAccountStore: snapshot.hasUnreadableAddedAccountStore) + } + + private static func normalizeVisibleEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } +} + +private struct AccountIdentity: Equatable { + let id: UUID + let email: String + let managedHomePath: String + let createdAt: TimeInterval + let updatedAt: TimeInterval + let lastAuthenticatedAt: TimeInterval? + + init(_ account: ManagedCodexAccount) { + self.id = account.id + self.email = account.email + self.managedHomePath = account.managedHomePath + self.createdAt = account.createdAt + self.updatedAt = account.updatedAt + self.lastAuthenticatedAt = account.lastAuthenticatedAt + } +} diff --git a/Sources/CodexBar/CodexLoginAlertPresentation.swift b/Sources/CodexBar/CodexLoginAlertPresentation.swift new file mode 100644 index 000000000..9f1e79650 --- /dev/null +++ b/Sources/CodexBar/CodexLoginAlertPresentation.swift @@ -0,0 +1,38 @@ +import Foundation + +struct CodexLoginAlertInfo: Equatable { + let title: String + let message: String +} + +enum CodexLoginAlertPresentation { + static func alertInfo(for result: CodexLoginRunner.Result) -> CodexLoginAlertInfo? { + switch result.outcome { + case .success: + return nil + case .missingBinary: + return CodexLoginAlertInfo( + title: "Codex CLI not found", + message: "Install the Codex CLI (npm i -g @openai/codex) and try again.") + case let .launchFailed(message): + return CodexLoginAlertInfo(title: "Could not start codex login", message: message) + case .timedOut: + return CodexLoginAlertInfo( + title: "Codex login timed out", + message: self.trimmedOutput(result.output)) + case let .failed(status): + let statusLine = "codex login exited with status \(status)." + let message = self.trimmedOutput(result.output.isEmpty ? statusLine : result.output) + return CodexLoginAlertInfo(title: "Codex login failed", message: message) + } + } + + private static func trimmedOutput(_ text: String) -> String { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + let limit = 600 + if trimmed.isEmpty { return "No output captured." } + if trimmed.count <= limit { return trimmed } + let idx = trimmed.index(trimmed.startIndex, offsetBy: limit) + return "\(trimmed[.. Result { + static func run(homePath: String? = nil, timeout: TimeInterval = 120) async -> Result { await Task(priority: .userInitiated) { var env = ProcessInfo.processInfo.environment env["PATH"] = PathBuilder.effectivePATH( purposes: [.rpc, .tty, .nodeTooling], env: env, loginPATH: LoginShellPathCache.shared.current) + env = CodexHomeScope.scopedEnvironment(base: env, codexHome: homePath) guard let executable = BinaryLocator.resolveCodexBinary( env: env, diff --git a/Sources/CodexBar/CodexSystemAccountObserver.swift b/Sources/CodexBar/CodexSystemAccountObserver.swift new file mode 100644 index 000000000..584672c69 --- /dev/null +++ b/Sources/CodexBar/CodexSystemAccountObserver.swift @@ -0,0 +1,31 @@ +import CodexBarCore +import Foundation + +struct ObservedSystemCodexAccount: Equatable, Sendable { + let email: String + let codexHomePath: String + let observedAt: Date +} + +protocol CodexSystemAccountObserving: Sendable { + func loadSystemAccount(environment: [String: String]) throws -> ObservedSystemCodexAccount? +} + +struct DefaultCodexSystemAccountObserver: CodexSystemAccountObserving { + func loadSystemAccount(environment: [String: String]) throws -> ObservedSystemCodexAccount? { + let homeURL = CodexHomeScope.ambientHomeURL(env: environment) + let fetcher = UsageFetcher(environment: environment) + let info = fetcher.loadAccountInfo() + + guard let rawEmail = info.email?.trimmingCharacters(in: .whitespacesAndNewlines), + !rawEmail.isEmpty + else { + return nil + } + + return ObservedSystemCodexAccount( + email: rawEmail.lowercased(), + codexHomePath: homeURL.path, + observedAt: Date()) + } +} diff --git a/Sources/CodexBar/CodexbarApp.swift b/Sources/CodexBar/CodexbarApp.swift index fca4af153..f5097b0e4 100644 --- a/Sources/CodexBar/CodexbarApp.swift +++ b/Sources/CodexBar/CodexbarApp.swift @@ -11,6 +11,7 @@ struct CodexBarApp: App { @NSApplicationDelegateAdaptor(AppDelegate.self) private var appDelegate @State private var settings: SettingsStore @State private var store: UsageStore + @State private var managedCodexAccountCoordinator: ManagedCodexAccountCoordinator private let preferencesSelection: PreferencesSelection private let account: AccountInfo @@ -41,6 +42,11 @@ struct CodexBarApp: App { let preferencesSelection = PreferencesSelection() let settings = SettingsStore() + let managedCodexAccountCoordinator = ManagedCodexAccountCoordinator() + managedCodexAccountCoordinator.onManagedAccountsDidChange = { + _ = settings.persistResolvedCodexActiveSourceCorrectionIfNeeded() + } + _ = settings.persistResolvedCodexActiveSourceCorrectionIfNeeded() let fetcher = UsageFetcher() let browserDetection = BrowserDetection(cacheTTL: BrowserDetection.defaultCacheTTL) let account = fetcher.loadAccountInfo() @@ -48,13 +54,15 @@ struct CodexBarApp: App { self.preferencesSelection = preferencesSelection _settings = State(wrappedValue: settings) _store = State(wrappedValue: store) + _managedCodexAccountCoordinator = State(wrappedValue: managedCodexAccountCoordinator) self.account = account CodexBarLog.setLogLevel(settings.debugLogLevel) self.appDelegate.configure( store: store, settings: settings, account: account, - selection: preferencesSelection) + selection: preferencesSelection, + managedCodexAccountCoordinator: managedCodexAccountCoordinator) } @SceneBuilder @@ -72,7 +80,8 @@ struct CodexBarApp: App { settings: self.settings, store: self.store, updater: self.appDelegate.updaterController, - selection: self.preferencesSelection) + selection: self.preferencesSelection, + managedCodexAccountCoordinator: self.managedCodexAccountCoordinator) } .defaultSize(width: PreferencesTab.general.preferredWidth, height: PreferencesTab.general.preferredHeight) .windowResizability(.contentSize) @@ -257,12 +266,20 @@ final class AppDelegate: NSObject, NSApplicationDelegate { private var settings: SettingsStore? private var account: AccountInfo? private var preferencesSelection: PreferencesSelection? - - func configure(store: UsageStore, settings: SettingsStore, account: AccountInfo, selection: PreferencesSelection) { + private var managedCodexAccountCoordinator: ManagedCodexAccountCoordinator? + + func configure( + store: UsageStore, + settings: SettingsStore, + account: AccountInfo, + selection: PreferencesSelection, + managedCodexAccountCoordinator: ManagedCodexAccountCoordinator) + { self.store = store self.settings = settings self.account = account self.preferencesSelection = selection + self.managedCodexAccountCoordinator = managedCodexAccountCoordinator } func applicationWillFinishLaunching(_ notification: Notification) { @@ -311,13 +328,19 @@ final class AppDelegate: NSObject, NSApplicationDelegate { private func ensureStatusController() { if self.statusController != nil { return } - if let store, let settings, let account, let selection = self.preferencesSelection { + if let store, + let settings, + let account, + let selection = self.preferencesSelection, + let managedCodexAccountCoordinator + { self.statusController = StatusItemController.factory( store, settings, account, self.updaterController, - selection) + selection, + managedCodexAccountCoordinator) return } @@ -330,11 +353,13 @@ final class AppDelegate: NSObject, NSApplicationDelegate { let browserDetection = BrowserDetection(cacheTTL: BrowserDetection.defaultCacheTTL) let fallbackAccount = fetcher.loadAccountInfo() let fallbackStore = UsageStore(fetcher: fetcher, browserDetection: browserDetection, settings: fallbackSettings) + let fallbackManagedCodexAccountCoordinator = ManagedCodexAccountCoordinator() self.statusController = StatusItemController.factory( fallbackStore, fallbackSettings, fallbackAccount, self.updaterController, - PreferencesSelection()) + PreferencesSelection(), + fallbackManagedCodexAccountCoordinator) } } diff --git a/Sources/CodexBar/ManagedCodexAccountCoordinator.swift b/Sources/CodexBar/ManagedCodexAccountCoordinator.swift new file mode 100644 index 000000000..eb00b8fc7 --- /dev/null +++ b/Sources/CodexBar/ManagedCodexAccountCoordinator.swift @@ -0,0 +1,48 @@ +import CodexBarCore +import Foundation +import Observation + +enum ManagedCodexAccountCoordinatorError: Error, Equatable, Sendable { + case authenticationInProgress +} + +@MainActor +@Observable +final class ManagedCodexAccountCoordinator { + let service: ManagedCodexAccountService + private(set) var isAuthenticatingManagedAccount: Bool = false + private(set) var authenticatingManagedAccountID: UUID? + var onManagedAccountsDidChange: (@MainActor () -> Void)? + + init(service: ManagedCodexAccountService = ManagedCodexAccountService()) { + self.service = service + } + + func authenticateManagedAccount( + existingAccountID: UUID? = nil, + timeout: TimeInterval = 120) + async throws -> ManagedCodexAccount + { + guard self.isAuthenticatingManagedAccount == false else { + throw ManagedCodexAccountCoordinatorError.authenticationInProgress + } + + self.isAuthenticatingManagedAccount = true + self.authenticatingManagedAccountID = existingAccountID + defer { + self.isAuthenticatingManagedAccount = false + self.authenticatingManagedAccountID = nil + } + + let account = try await self.service.authenticateManagedAccount( + existingAccountID: existingAccountID, + timeout: timeout) + self.onManagedAccountsDidChange?() + return account + } + + func removeManagedAccount(id: UUID) async throws { + try await self.service.removeManagedAccount(id: id) + self.onManagedAccountsDidChange?() + } +} diff --git a/Sources/CodexBar/ManagedCodexAccountService.swift b/Sources/CodexBar/ManagedCodexAccountService.swift new file mode 100644 index 000000000..91f9d32a7 --- /dev/null +++ b/Sources/CodexBar/ManagedCodexAccountService.swift @@ -0,0 +1,196 @@ +import CodexBarCore +import Foundation + +protocol ManagedCodexHomeProducing: Sendable { + func makeHomeURL() -> URL + func validateManagedHomeForDeletion(_ url: URL) throws +} + +protocol ManagedCodexLoginRunning: Sendable { + func run(homePath: String, timeout: TimeInterval) async -> CodexLoginRunner.Result +} + +protocol ManagedCodexIdentityReading: Sendable { + func loadAccountInfo(homePath: String) throws -> AccountInfo +} + +enum ManagedCodexAccountServiceError: Error, Equatable, Sendable { + case loginFailed + case missingEmail + case unsafeManagedHome(String) +} + +struct ManagedCodexHomeFactory: ManagedCodexHomeProducing, Sendable { + let root: URL + + init(root: URL = Self.defaultRootURL(), fileManager: FileManager = .default) { + let standardizedRoot = root.standardizedFileURL + if standardizedRoot.path != root.path { + self.root = standardizedRoot + } else { + self.root = root + } + _ = fileManager + } + + func makeHomeURL() -> URL { + self.root.appendingPathComponent(UUID().uuidString, isDirectory: true) + } + + func validateManagedHomeForDeletion(_ url: URL) throws { + let rootPath = self.root.standardizedFileURL.path + let targetPath = url.standardizedFileURL.path + let rootPrefix = rootPath.hasSuffix("/") ? rootPath : rootPath + "/" + guard targetPath.hasPrefix(rootPrefix), targetPath != rootPath else { + throw ManagedCodexAccountServiceError.unsafeManagedHome(url.path) + } + } + + static func defaultRootURL(fileManager: FileManager = .default) -> URL { + let base = fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? fileManager.homeDirectoryForCurrentUser + return base + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("managed-codex-homes", isDirectory: true) + } +} + +struct DefaultManagedCodexLoginRunner: ManagedCodexLoginRunning { + func run(homePath: String, timeout: TimeInterval) async -> CodexLoginRunner.Result { + await CodexLoginRunner.run(homePath: homePath, timeout: timeout) + } +} + +struct DefaultManagedCodexIdentityReader: ManagedCodexIdentityReading { + func loadAccountInfo(homePath: String) throws -> AccountInfo { + let env = CodexHomeScope.scopedEnvironment( + base: ProcessInfo.processInfo.environment, + codexHome: homePath) + return UsageFetcher(environment: env).loadAccountInfo() + } +} + +@MainActor +final class ManagedCodexAccountService { + private let store: any ManagedCodexAccountStoring + private let homeFactory: any ManagedCodexHomeProducing + private let loginRunner: any ManagedCodexLoginRunning + private let identityReader: any ManagedCodexIdentityReading + private let fileManager: FileManager + + init( + store: any ManagedCodexAccountStoring, + homeFactory: any ManagedCodexHomeProducing, + loginRunner: any ManagedCodexLoginRunning, + identityReader: any ManagedCodexIdentityReading, + fileManager: FileManager = .default) + { + self.store = store + self.homeFactory = homeFactory + self.loginRunner = loginRunner + self.identityReader = identityReader + self.fileManager = fileManager + } + + convenience init(fileManager: FileManager = .default) { + self.init( + store: FileManagedCodexAccountStore(fileManager: fileManager), + homeFactory: ManagedCodexHomeFactory(fileManager: fileManager), + loginRunner: DefaultManagedCodexLoginRunner(), + identityReader: DefaultManagedCodexIdentityReader(), + fileManager: fileManager) + } + + func authenticateManagedAccount( + existingAccountID: UUID? = nil, + timeout: TimeInterval = 120) + async throws -> ManagedCodexAccount + { + let snapshot = try self.store.loadAccounts() + let homeURL = self.homeFactory.makeHomeURL() + try self.fileManager.createDirectory(at: homeURL, withIntermediateDirectories: true) + let account: ManagedCodexAccount + let existingHomePathToDelete: String? + + do { + let result = await self.loginRunner.run(homePath: homeURL.path, timeout: timeout) + guard case .success = result.outcome else { throw ManagedCodexAccountServiceError.loginFailed } + + let info = try self.identityReader.loadAccountInfo(homePath: homeURL.path) + guard let rawEmail = info.email?.trimmingCharacters(in: .whitespacesAndNewlines), !rawEmail.isEmpty else { + throw ManagedCodexAccountServiceError.missingEmail + } + + let now = Date().timeIntervalSince1970 + let existing = self.reconciledExistingAccount( + authenticatedEmail: rawEmail, + existingAccountID: existingAccountID, + snapshot: snapshot) + + account = ManagedCodexAccount( + id: existing?.id ?? UUID(), + email: rawEmail, + managedHomePath: homeURL.path, + createdAt: existing?.createdAt ?? now, + updatedAt: now, + lastAuthenticatedAt: now) + existingHomePathToDelete = existing?.managedHomePath + + let updatedSnapshot = ManagedCodexAccountSet( + version: snapshot.version, + accounts: snapshot.accounts.filter { $0.id != account.id && $0.email != account.email } + [account]) + try self.store.storeAccounts(updatedSnapshot) + } catch { + try? self.removeManagedHomeIfSafe(atPath: homeURL.path) + throw error + } + + if let existingHomePathToDelete, existingHomePathToDelete != homeURL.path { + try? self.removeManagedHomeIfSafe(atPath: existingHomePathToDelete) + } + return account + } + + func removeManagedAccount(id: UUID) async throws { + let snapshot = try self.store.loadAccounts() + guard let account = snapshot.account(id: id) else { return } + + let homeURL = URL(fileURLWithPath: account.managedHomePath, isDirectory: true) + try self.homeFactory.validateManagedHomeForDeletion(homeURL) + + let remaining = snapshot.accounts.filter { $0.id != id } + try self.store.storeAccounts(ManagedCodexAccountSet( + version: snapshot.version, + accounts: remaining)) + + if self.fileManager.fileExists(atPath: homeURL.path) { + try? self.fileManager.removeItem(at: homeURL) + } + } + + private func removeManagedHomeIfSafe(atPath path: String) throws { + let homeURL = URL(fileURLWithPath: path, isDirectory: true) + try self.homeFactory.validateManagedHomeForDeletion(homeURL) + if self.fileManager.fileExists(atPath: homeURL.path) { + try self.fileManager.removeItem(at: homeURL) + } + } + + private func reconciledExistingAccount( + authenticatedEmail: String, + existingAccountID: UUID?, + snapshot: ManagedCodexAccountSet) + -> ManagedCodexAccount? + { + if let existingByEmail = snapshot.account(email: authenticatedEmail) { + return existingByEmail + } + guard let existingAccountID else { return nil } + guard let existingByID = snapshot.account(id: existingAccountID) else { return nil } + return existingByID.email == Self.normalizeEmail(authenticatedEmail) ? existingByID : nil + } + + private static func normalizeEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } +} diff --git a/Sources/CodexBar/MenuContent.swift b/Sources/CodexBar/MenuContent.swift index fa41695f5..c8e113267 100644 --- a/Sources/CodexBar/MenuContent.swift +++ b/Sources/CodexBar/MenuContent.swift @@ -86,6 +86,8 @@ struct MenuContent: View { self.actions.openDashboard() case .statusPage: self.actions.openStatusPage() + case .addCodexAccount: + self.actions.addCodexAccount() case let .switchAccount(provider): self.actions.switchAccount(provider) case let .openTerminal(command): @@ -112,6 +114,7 @@ struct MenuActions { let refreshAugmentSession: () -> Void let openDashboard: () -> Void let openStatusPage: () -> Void + let addCodexAccount: () -> Void let switchAccount: (UsageProvider) -> Void let openTerminal: (String) -> Void let openSettings: () -> Void diff --git a/Sources/CodexBar/MenuDescriptor.swift b/Sources/CodexBar/MenuDescriptor.swift index 05aa55fff..01675205f 100644 --- a/Sources/CodexBar/MenuDescriptor.swift +++ b/Sources/CodexBar/MenuDescriptor.swift @@ -17,6 +17,7 @@ struct MenuDescriptor { case refresh = "arrow.clockwise" case dashboard = "chart.bar" case statusPage = "waveform.path.ecg" + case addAccount = "plus" case switchAccount = "key" case openTerminal = "terminal" case loginToProvider = "arrow.right.square" @@ -38,6 +39,7 @@ struct MenuDescriptor { case refreshAugmentSession case dashboard case statusPage + case addCodexAccount case switchAccount(UsageProvider) case openTerminal(command: String) case loginToProvider(url: String) @@ -60,12 +62,13 @@ struct MenuDescriptor { var sections: [Section] = [] if let provider { + let fallbackAccount = store.accountInfo(for: provider) sections.append(Self.usageSection(for: provider, store: store, settings: settings)) if let accountSection = Self.accountSection( for: provider, store: store, settings: settings, - account: account) + account: fallbackAccount) { sections.append(accountSection) } @@ -78,11 +81,12 @@ struct MenuDescriptor { } if addedUsage { if let accountProvider = Self.accountProviderForCombined(store: store), + let fallbackAccount = Optional(store.accountInfo(for: accountProvider)), let accountSection = Self.accountSection( for: accountProvider, store: store, settings: settings, - account: account) + account: fallbackAccount) { sections.append(accountSection) } @@ -307,12 +311,13 @@ struct MenuDescriptor { var entries: [Entry] = [] let targetProvider = provider ?? store.enabledProviders().first let metadata = targetProvider.map { store.metadata(for: $0) } + let fallbackAccount = targetProvider.map { store.accountInfo(for: $0) } ?? account let loginContext = targetProvider.map { ProviderMenuLoginContext( provider: $0, store: store, settings: store.settings, - account: account) + account: fallbackAccount) } // Show "Add Account" if no account, "Switch Account" if logged in @@ -326,7 +331,7 @@ struct MenuDescriptor { entries.append(.action(override.label, override.action)) } else { let loginAction = self.switchAccountTarget(for: provider, store: store) - let hasAccount = self.hasAccount(for: provider, store: store, account: account) + let hasAccount = self.hasAccount(for: provider, store: store, account: fallbackAccount) let accountLabel = hasAccount ? "Switch Account..." : "Add Account..." entries.append(.action(accountLabel, loginAction)) } @@ -337,7 +342,7 @@ struct MenuDescriptor { provider: targetProvider, store: store, settings: store.settings, - account: account) + account: fallbackAccount) ProviderCatalog.implementation(for: targetProvider)? .appendActionMenuEntries(context: actionContext, entries: &entries) } @@ -456,6 +461,7 @@ extension MenuDescriptor.MenuAction { case .refreshAugmentSession: MenuDescriptor.MenuActionSystemImage.refresh.rawValue case .dashboard: MenuDescriptor.MenuActionSystemImage.dashboard.rawValue case .statusPage: MenuDescriptor.MenuActionSystemImage.statusPage.rawValue + case .addCodexAccount: MenuDescriptor.MenuActionSystemImage.addAccount.rawValue case .switchAccount: MenuDescriptor.MenuActionSystemImage.switchAccount.rawValue case .openTerminal: MenuDescriptor.MenuActionSystemImage.openTerminal.rawValue case .loginToProvider: MenuDescriptor.MenuActionSystemImage.loginToProvider.rawValue diff --git a/Sources/CodexBar/PreferencesCodexAccountsSection.swift b/Sources/CodexBar/PreferencesCodexAccountsSection.swift new file mode 100644 index 000000000..c1b620e07 --- /dev/null +++ b/Sources/CodexBar/PreferencesCodexAccountsSection.swift @@ -0,0 +1,225 @@ +import Foundation +import SwiftUI + +protocol CodexAmbientLoginRunning: Sendable { + func run(timeout: TimeInterval) async -> CodexLoginRunner.Result +} + +struct DefaultCodexAmbientLoginRunner: CodexAmbientLoginRunning { + func run(timeout: TimeInterval) async -> CodexLoginRunner.Result { + await CodexLoginRunner.run(timeout: timeout) + } +} + +struct CodexAccountsSectionNotice: Equatable { + enum Tone: Equatable { + case secondary + case warning + } + + let text: String + let tone: Tone +} + +struct CodexAccountsSectionState: Equatable { + let visibleAccounts: [CodexVisibleAccount] + let activeVisibleAccountID: String? + let hasUnreadableManagedAccountStore: Bool + let isAuthenticatingManagedAccount: Bool + let authenticatingManagedAccountID: UUID? + let isAuthenticatingLiveAccount: Bool + let notice: CodexAccountsSectionNotice? + + var showsActivePicker: Bool { + self.visibleAccounts.count > 1 + } + + var singleVisibleAccount: CodexVisibleAccount? { + self.visibleAccounts.count == 1 ? self.visibleAccounts.first : nil + } + + var canAddAccount: Bool { + !self.hasUnreadableManagedAccountStore && + !self.isAuthenticatingManagedAccount && + !self.isAuthenticatingLiveAccount + } + + var addAccountTitle: String { + if self.isAuthenticatingManagedAccount, self.authenticatingManagedAccountID == nil { + return "Adding Account…" + } + return "Add Account" + } + + func showsLiveBadge(for account: CodexVisibleAccount) -> Bool { + self.visibleAccounts.count > 1 && account.isLive && account.storedAccountID == nil + } + + func canReauthenticate(_ account: CodexVisibleAccount) -> Bool { + guard account.canReauthenticate else { return false } + guard self.isAuthenticatingManagedAccount == false else { return false } + guard self.isAuthenticatingLiveAccount == false else { return false } + if account.storedAccountID != nil { + return self.hasUnreadableManagedAccountStore == false + } + return true + } + + func canRemove(_ account: CodexVisibleAccount) -> Bool { + guard account.canRemove else { return false } + guard self.isAuthenticatingManagedAccount == false else { return false } + guard self.isAuthenticatingLiveAccount == false else { return false } + return self.hasUnreadableManagedAccountStore == false + } + + func reauthenticateTitle(for account: CodexVisibleAccount) -> String { + if let accountID = account.storedAccountID, + self.isAuthenticatingManagedAccount, + self.authenticatingManagedAccountID == accountID + { + return "Re-authenticating…" + } + if account.storedAccountID == nil, self.isAuthenticatingLiveAccount { + return "Re-authenticating…" + } + return "Re-auth" + } +} + +@MainActor +struct CodexAccountsSectionView: View { + let state: CodexAccountsSectionState + let setActiveVisibleAccount: (String) -> Void + let reauthenticateAccount: (CodexVisibleAccount) -> Void + let removeAccount: (CodexVisibleAccount) -> Void + let addAccount: () -> Void + + var body: some View { + ProviderSettingsSection(title: "Accounts") { + if let selection = self.activeSelectionBinding { + VStack(alignment: .leading, spacing: 6) { + HStack(alignment: .firstTextBaseline, spacing: 10) { + Text("Active") + .font(.subheadline.weight(.semibold)) + .frame(width: ProviderSettingsMetrics.pickerLabelWidth, alignment: .leading) + + Picker("", selection: selection) { + ForEach(self.state.visibleAccounts) { account in + Text(account.email).tag(account.id) + } + } + .labelsHidden() + .pickerStyle(.menu) + .controlSize(.small) + + Spacer(minLength: 0) + } + + Text("Choose which Codex account CodexBar should follow.") + .font(.footnote) + .foregroundStyle(.secondary) + } + .disabled(self.state.isAuthenticatingManagedAccount || self.state.isAuthenticatingLiveAccount) + } else if let account = self.state.singleVisibleAccount { + VStack(alignment: .leading, spacing: 6) { + HStack(alignment: .firstTextBaseline, spacing: 10) { + Text("Account") + .font(.subheadline.weight(.semibold)) + .frame(width: ProviderSettingsMetrics.pickerLabelWidth, alignment: .leading) + + Text(account.email) + .font(.subheadline) + + Spacer(minLength: 0) + } + } + } + + if self.state.visibleAccounts.isEmpty { + Text("No Codex accounts detected yet.") + .font(.footnote) + .foregroundStyle(.secondary) + } else { + VStack(alignment: .leading, spacing: 10) { + ForEach(self.state.visibleAccounts) { account in + CodexAccountsSectionRowView( + account: account, + showsLiveBadge: self.state.showsLiveBadge(for: account), + reauthenticateTitle: self.state.reauthenticateTitle(for: account), + canReauthenticate: self.state.canReauthenticate(account), + canRemove: self.state.canRemove(account), + onReauthenticate: { self.reauthenticateAccount(account) }, + onRemove: { self.removeAccount(account) }) + } + } + } + + if let notice = self.state.notice { + Text(notice.text) + .font(.footnote) + .foregroundStyle(notice.tone == .warning ? .red : .secondary) + .fixedSize(horizontal: false, vertical: true) + } + + Button(self.state.addAccountTitle) { + self.addAccount() + } + .buttonStyle(.bordered) + .controlSize(.small) + .disabled(self.state.canAddAccount == false) + } + } + + private var activeSelectionBinding: Binding? { + guard self.state.showsActivePicker else { return nil } + let fallbackID = self.state.activeVisibleAccountID ?? self.state.visibleAccounts.first?.id + guard let fallbackID else { return nil } + return Binding( + get: { self.state.activeVisibleAccountID ?? fallbackID }, + set: { self.setActiveVisibleAccount($0) }) + } +} + +private struct CodexAccountsSectionRowView: View { + let account: CodexVisibleAccount + let showsLiveBadge: Bool + let reauthenticateTitle: String + let canReauthenticate: Bool + let canRemove: Bool + let onReauthenticate: () -> Void + let onRemove: () -> Void + + var body: some View { + HStack(alignment: .center, spacing: 12) { + HStack(alignment: .firstTextBaseline, spacing: 6) { + Text(self.account.email) + .font(.subheadline.weight(.semibold)) + if self.showsLiveBadge { + Text("(Live)") + .font(.caption.weight(.semibold)) + .foregroundStyle(.secondary) + } + } + + Spacer(minLength: 8) + + if self.account.canReauthenticate { + Button(self.reauthenticateTitle) { + self.onReauthenticate() + } + .buttonStyle(.bordered) + .controlSize(.small) + .disabled(self.canReauthenticate == false) + } + + if self.account.canRemove { + Button("Remove") { + self.onRemove() + } + .buttonStyle(.bordered) + .controlSize(.small) + .disabled(self.canRemove == false) + } + } + } +} diff --git a/Sources/CodexBar/PreferencesProviderDetailView.swift b/Sources/CodexBar/PreferencesProviderDetailView.swift index 58a55deb5..edcc1d0dc 100644 --- a/Sources/CodexBar/PreferencesProviderDetailView.swift +++ b/Sources/CodexBar/PreferencesProviderDetailView.swift @@ -2,7 +2,7 @@ import CodexBarCore import SwiftUI @MainActor -struct ProviderDetailView: View { +struct ProviderDetailView: View { let provider: UsageProvider @Bindable var store: UsageStore @Binding var isEnabled: Bool @@ -16,6 +16,42 @@ struct ProviderDetailView: View { @Binding var isErrorExpanded: Bool let onCopyError: (String) -> Void let onRefresh: () -> Void + let supplementarySettingsContent: SupplementaryContent + let showsSupplementarySettingsContent: Bool + + init( + provider: UsageProvider, + store: UsageStore, + isEnabled: Binding, + subtitle: String, + model: UsageMenuCardView.Model, + settingsPickers: [ProviderSettingsPickerDescriptor], + settingsToggles: [ProviderSettingsToggleDescriptor], + settingsFields: [ProviderSettingsFieldDescriptor], + settingsTokenAccounts: ProviderSettingsTokenAccountsDescriptor?, + errorDisplay: ProviderErrorDisplay?, + isErrorExpanded: Binding, + onCopyError: @escaping (String) -> Void, + onRefresh: @escaping () -> Void, + showsSupplementarySettingsContent: Bool = false, + @ViewBuilder supplementarySettingsContent: () -> SupplementaryContent) + { + self.provider = provider + self.store = store + self._isEnabled = isEnabled + self.subtitle = subtitle + self.model = model + self.settingsPickers = settingsPickers + self.settingsToggles = settingsToggles + self.settingsFields = settingsFields + self.settingsTokenAccounts = settingsTokenAccounts + self.errorDisplay = errorDisplay + self._isErrorExpanded = isErrorExpanded + self.onCopyError = onCopyError + self.onRefresh = onRefresh + self.showsSupplementarySettingsContent = showsSupplementarySettingsContent + self.supplementarySettingsContent = supplementarySettingsContent() + } static func metricTitle(provider: UsageProvider, metric: UsageMenuCardView.Model.Metric) -> String { UsageMenuCardView.popupMetricTitle(provider: provider, metric: metric) @@ -85,6 +121,10 @@ struct ProviderDetailView: View { } } + if self.showsSupplementarySettingsContent { + self.supplementarySettingsContent + } + if !self.settingsToggles.isEmpty { ProviderSettingsSection(title: "Options") { ForEach(self.settingsToggles) { toggle in @@ -256,7 +296,10 @@ private struct ProviderDetailInfoGrid: View { ProviderDetailInfoRow(label: "Account", value: email, labelWidth: self.labelWidth) } - if let planRow = ProviderDetailView.planRow(provider: self.provider, planText: self.model.planText) { + if let planRow = ProviderDetailView.planRow( + provider: self.provider, + planText: self.model.planText) + { ProviderDetailInfoRow(label: planRow.label, value: planRow.value, labelWidth: self.labelWidth) } } @@ -317,7 +360,7 @@ struct ProviderMetricsInlineView: View { ForEach(self.model.metrics, id: \.id) { metric in ProviderMetricInlineRow( metric: metric, - title: ProviderDetailView.metricTitle(provider: self.provider, metric: metric), + title: ProviderDetailView.metricTitle(provider: self.provider, metric: metric), progressColor: self.model.progressColor, labelWidth: self.labelWidth) } diff --git a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift index e2dce0a7f..437cb6a93 100644 --- a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift +++ b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift @@ -22,6 +22,26 @@ extension ProvidersPane { func _test_menuCardModel(for provider: UsageProvider) -> UsageMenuCardView.Model { self.menuCardModel(for: provider) } + + func _test_providerErrorDisplay(for provider: UsageProvider) -> ProviderErrorDisplay? { + self.providerErrorDisplay(provider) + } + + func _test_codexAccountsSectionState() -> CodexAccountsSectionState? { + self.codexAccountsSectionState(for: .codex) + } + + func _test_selectCodexVisibleAccount(id: String) async { + await self.selectCodexVisibleAccount(id: id) + } + + func _test_addManagedCodexAccount() async { + await self.addManagedCodexAccount() + } + + func _test_reauthenticateCodexAccount(_ account: CodexVisibleAccount) async { + await self.reauthenticateCodexAccount(account) + } } @MainActor @@ -100,7 +120,13 @@ enum ProvidersPaneTestHarness { errorDisplay: ProviderErrorDisplay(preview: "Preview", full: "Full"), isErrorExpanded: expandedBinding, onCopyError: { _ in }, - onRefresh: {}).body + onRefresh: {}, + showsSupplementarySettingsContent: true, + supplementarySettingsContent: { + ProviderSettingsSection(title: "Accounts") { + Text("Supplementary") + } + }).body } private static func makeDescriptors() -> ProviderListTestDescriptors { diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index f3d5bc112..53943cf53 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -6,16 +6,32 @@ import SwiftUI struct ProvidersPane: View { @Bindable var settings: SettingsStore @Bindable var store: UsageStore + let managedCodexAccountCoordinator: ManagedCodexAccountCoordinator + let codexAmbientLoginRunner: any CodexAmbientLoginRunning @State private var expandedErrors: Set = [] @State private var settingsStatusTextByID: [String: String] = [:] @State private var settingsLastAppActiveRunAtByID: [String: Date] = [:] @State private var activeConfirmation: ProviderSettingsConfirmationState? + @State private var codexAccountsNotice: CodexAccountsSectionNotice? + @State private var isAuthenticatingLiveCodexAccount = false @State private var selectedProvider: UsageProvider? private var providers: [UsageProvider] { self.settings.orderedProviders() } + init( + settings: SettingsStore, + store: UsageStore, + managedCodexAccountCoordinator: ManagedCodexAccountCoordinator = ManagedCodexAccountCoordinator(), + codexAmbientLoginRunner: any CodexAmbientLoginRunning = DefaultCodexAmbientLoginRunner()) + { + self.settings = settings + self.store = store + self.managedCodexAccountCoordinator = managedCodexAccountCoordinator + self.codexAmbientLoginRunner = codexAmbientLoginRunner + } + var body: some View { HStack(alignment: .top, spacing: 16) { ProviderSidebarListView( @@ -45,9 +61,38 @@ struct ProvidersPane: View { onRefresh: { Task { @MainActor in await ProviderInteractionContext.$current.withValue(.userInitiated) { - await self.store.refreshProvider(provider, allowDisabled: true) + if provider == .codex { + await self.store.refreshCodexAccountScopedState(allowDisabled: true) + } else { + await self.store.refreshProvider(provider, allowDisabled: true) + } } } + }, + showsSupplementarySettingsContent: self.codexAccountsSectionState(for: provider) != nil, + supplementarySettingsContent: { + if let state = self.codexAccountsSectionState(for: provider) { + CodexAccountsSectionView( + state: state, + setActiveVisibleAccount: { visibleAccountID in + Task { @MainActor in + await self.selectCodexVisibleAccount(id: visibleAccountID) + } + }, + reauthenticateAccount: { account in + Task { @MainActor in + await self.reauthenticateCodexAccount(account) + } + }, + removeAccount: { account in + self.requestManagedCodexAccountRemoval(account) + }, + addAccount: { + Task { @MainActor in + await self.addManagedCodexAccount() + } + }) + } }) } else { Text("Select a provider") @@ -133,11 +178,110 @@ struct ProvidersPane: View { return "\(detailLine)\n\(usageText)" } - private func providerErrorDisplay(_ provider: UsageProvider) -> ProviderErrorDisplay? { - guard let raw = self.store.error(for: provider), !raw.isEmpty else { return nil } + func codexAccountsSectionState(for provider: UsageProvider) -> CodexAccountsSectionState? { + guard provider == .codex else { return nil } + let projection = self.settings.codexVisibleAccountProjection + let degradedNotice: CodexAccountsSectionNotice? = if projection.hasUnreadableAddedAccountStore { + CodexAccountsSectionNotice( + text: "Managed account storage is unreadable. Live account access is still available, " + + "but managed add, re-auth, and remove actions are disabled until the store is recoverable.", + tone: .warning) + } else { + nil + } + + return CodexAccountsSectionState( + visibleAccounts: projection.visibleAccounts, + activeVisibleAccountID: projection.activeVisibleAccountID, + hasUnreadableManagedAccountStore: projection.hasUnreadableAddedAccountStore, + isAuthenticatingManagedAccount: self.managedCodexAccountCoordinator.isAuthenticatingManagedAccount, + authenticatingManagedAccountID: self.managedCodexAccountCoordinator.authenticatingManagedAccountID, + isAuthenticatingLiveAccount: self.isAuthenticatingLiveCodexAccount, + notice: self.codexAccountsNotice ?? degradedNotice) + } + + func selectCodexVisibleAccount(id: String) async { + self.codexAccountsNotice = nil + guard self.settings.selectCodexVisibleAccount(id: id) else { return } + await self.refreshCodexProvider() + } + + func addManagedCodexAccount() async { + self.codexAccountsNotice = nil + guard let state = self.codexAccountsSectionState(for: .codex), state.canAddAccount else { + return + } + + do { + let account = try await self.managedCodexAccountCoordinator.authenticateManagedAccount() + self.selectCodexVisibleAccountForAuthenticatedManagedAccount(account) + await self.refreshCodexProvider() + } catch { + self.codexAccountsNotice = self.codexAccountsNotice(for: error) + } + } + + func reauthenticateCodexAccount(_ account: CodexVisibleAccount) async { + self.codexAccountsNotice = nil + if let accountID = account.storedAccountID { + guard let state = self.codexAccountsSectionState(for: .codex), state.canReauthenticate(account) else { + return + } + do { + _ = try await self.managedCodexAccountCoordinator + .authenticateManagedAccount(existingAccountID: accountID) + await self.refreshCodexProvider() + } catch { + self.codexAccountsNotice = self.codexAccountsNotice(for: error) + } + return + } + + guard let state = self.codexAccountsSectionState(for: .codex), state.canReauthenticate(account) else { + return + } + + self.isAuthenticatingLiveCodexAccount = true + defer { self.isAuthenticatingLiveCodexAccount = false } + + let result = await self.codexAmbientLoginRunner.run(timeout: 120) + if let info = CodexLoginAlertPresentation.alertInfo(for: result) { + self.presentLoginAlert(title: info.title, message: info.message) + return + } + + await self.refreshCodexProvider() + } + + func removeManagedCodexAccount(id: UUID) async { + self.codexAccountsNotice = nil + do { + try await self.managedCodexAccountCoordinator.removeManagedAccount(id: id) + await self.refreshCodexProvider() + } catch { + self.codexAccountsNotice = self.codexAccountsNotice(for: error) + } + } + + func requestManagedCodexAccountRemoval(_ account: CodexVisibleAccount) { + guard let accountID = account.storedAccountID else { return } + self.activeConfirmation = ProviderSettingsConfirmationState( + title: "Remove Codex account?", + message: "Remove \(account.email) from CodexBar? Its managed Codex home will be deleted.", + confirmTitle: "Remove", + onConfirm: { + Task { @MainActor in + await self.removeManagedCodexAccount(id: accountID) + } + }) + } + + func providerErrorDisplay(_ provider: UsageProvider) -> ProviderErrorDisplay? { + guard let full = self.store.error(for: provider), !full.isEmpty else { return nil } + let preview = self.store.userFacingError(for: provider) ?? full return ProviderErrorDisplay( - preview: self.truncated(raw, prefix: ""), - full: raw) + preview: self.truncated(preview, prefix: ""), + full: full) } private func extraSettingsToggles(for provider: UsageProvider) -> [ProviderSettingsToggleDescriptor] { @@ -329,9 +473,9 @@ struct ProvidersPane: View { let tokenError: String? if provider == .codex { credits = self.store.credits - creditsError = self.store.lastCreditsError + creditsError = self.store.userFacingLastCreditsError dashboard = self.store.openAIDashboardRequiresLogin ? nil : self.store.openAIDashboard - dashboardError = self.store.lastOpenAIDashboardError + dashboardError = self.store.userFacingLastOpenAIDashboardError tokenSnapshot = self.store.tokenSnapshot(for: provider) tokenError = self.store.tokenError(for: provider) } else if provider == .claude || provider == .vertexai { @@ -364,9 +508,9 @@ struct ProvidersPane: View { dashboardError: dashboardError, tokenSnapshot: tokenSnapshot, tokenError: tokenError, - account: self.store.accountInfo(), + account: self.store.accountInfo(for: provider), isRefreshing: self.store.refreshingProviders.contains(provider), - lastError: self.store.error(for: provider), + lastError: self.store.userFacingError(for: provider), usageBarsShowUsed: self.settings.usageBarsShowUsed, resetTimeDisplayStyle: self.settings.resetTimeDisplayStyle, tokenCostUsageEnabled: self.settings.isCostUsageEffectivelyEnabled(for: provider), @@ -377,6 +521,52 @@ struct ProvidersPane: View { return UsageMenuCardView.Model.make(input) } + private func refreshCodexProvider() async { + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await self.store.refreshCodexAccountScopedState(allowDisabled: true) + } + } + + private func selectCodexVisibleAccountForAuthenticatedManagedAccount(_ account: ManagedCodexAccount) { + self.settings.selectAuthenticatedManagedCodexAccount(account) + } + + private func codexAccountsNotice(for error: Error) -> CodexAccountsSectionNotice { + if let error = error as? ManagedCodexAccountCoordinatorError, + error == .authenticationInProgress + { + return CodexAccountsSectionNotice( + text: "A managed Codex login is already running. Wait for it to finish before adding " + + "or re-authenticating another account.", + tone: .warning) + } + + if let error = error as? ManagedCodexAccountServiceError { + let message = switch error { + case .loginFailed: + "Managed Codex login did not complete. Try again after finishing the browser login flow." + case .missingEmail: + "Codex login completed, but no account email was available. Try again after confirming " + + "the account is fully signed in." + case let .unsafeManagedHome(path): + "CodexBar refused to modify an unexpected managed home path: \(path)" + } + return CodexAccountsSectionNotice(text: message, tone: .warning) + } + + return CodexAccountsSectionNotice( + text: error.localizedDescription, + tone: .warning) + } + + private func presentLoginAlert(title: String, message: String) { + let alert = NSAlert() + alert.messageText = title + alert.informativeText = message + alert.alertStyle = .warning + alert.runModal() + } + private func runSettingsDidBecomeActiveHooks() { for provider in UsageProvider.allCases { for toggle in self.extraSettingsToggles(for: provider) { @@ -424,6 +614,18 @@ struct ProviderSettingsConfirmationState: Identifiable { let confirmTitle: String let onConfirm: () -> Void + init( + title: String, + message: String, + confirmTitle: String, + onConfirm: @escaping () -> Void) + { + self.title = title + self.message = message + self.confirmTitle = confirmTitle + self.onConfirm = onConfirm + } + init(confirmation: ProviderSettingsConfirmation) { self.title = confirmation.title self.message = confirmation.message diff --git a/Sources/CodexBar/PreferencesView.swift b/Sources/CodexBar/PreferencesView.swift index a6f893950..3f2e58ec9 100644 --- a/Sources/CodexBar/PreferencesView.swift +++ b/Sources/CodexBar/PreferencesView.swift @@ -28,16 +28,34 @@ struct PreferencesView: View { @Bindable var store: UsageStore let updater: UpdaterProviding @Bindable var selection: PreferencesSelection + let managedCodexAccountCoordinator: ManagedCodexAccountCoordinator @State private var contentWidth: CGFloat = PreferencesTab.general.preferredWidth @State private var contentHeight: CGFloat = PreferencesTab.general.preferredHeight + init( + settings: SettingsStore, + store: UsageStore, + updater: UpdaterProviding, + selection: PreferencesSelection, + managedCodexAccountCoordinator: ManagedCodexAccountCoordinator = ManagedCodexAccountCoordinator()) + { + self.settings = settings + self.store = store + self.updater = updater + self.selection = selection + self.managedCodexAccountCoordinator = managedCodexAccountCoordinator + } + var body: some View { TabView(selection: self.$selection.tab) { GeneralPane(settings: self.settings, store: self.store) .tabItem { Label("General", systemImage: "gearshape") } .tag(PreferencesTab.general) - ProvidersPane(settings: self.settings, store: self.store) + ProvidersPane( + settings: self.settings, + store: self.store, + managedCodexAccountCoordinator: self.managedCodexAccountCoordinator) .tabItem { Label("Providers", systemImage: "square.grid.2x2") } .tag(PreferencesTab.providers) diff --git a/Sources/CodexBar/ProviderRegistry.swift b/Sources/CodexBar/ProviderRegistry.swift index 1e26c4c86..e3934d288 100644 --- a/Sources/CodexBar/ProviderRegistry.swift +++ b/Sources/CodexBar/ProviderRegistry.swift @@ -45,6 +45,7 @@ struct ProviderRegistry { provider: provider, settings: settings, tokenOverride: nil) + let fetcher = Self.makeFetcher(base: codexFetcher, provider: provider, env: env) let verbose = settings.isVerboseLoggingEnabled return ProviderFetchContext( runtime: .app, @@ -55,7 +56,7 @@ struct ProviderRegistry { verbose: verbose, env: env, settings: snapshot, - fetcher: codexFetcher, + fetcher: fetcher, claudeFetcher: claudeFetcher, browserDetection: browserDetection) }) @@ -107,6 +108,21 @@ struct ProviderRegistry { env[key] = value } } + // Managed Codex routing only scopes remote account fetches such as identity, plan, + // quotas, and dashboard data, and only when the active source is a managed account. + // Token-cost/session history is intentionally not routed through the managed home + // because that data is currently treated as provider-level local telemetry from this + // Mac's Codex sessions, not as account-owned remote state. If we later want + // account-scoped token history in the UI, that needs an explicit product decision and + // presentation change so the two concepts are not conflated. + if provider == .codex, let managedHomePath = settings.activeManagedCodexRemoteHomePath { + env = CodexHomeScope.scopedEnvironment(base: env, codexHome: managedHomePath) + } return env } + + static func makeFetcher(base: UsageFetcher, provider: UsageProvider, env: [String: String]) -> UsageFetcher { + guard provider == .codex else { return base } + return UsageFetcher(environment: env) + } } diff --git a/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift b/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift index c8847374c..3491df990 100644 --- a/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift +++ b/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift @@ -3,6 +3,9 @@ import CodexBarCore @MainActor extension StatusItemController { func runCodexLoginFlow() async { + // This menu action still follows the ambient Codex login behavior. Managed-account authentication is + // implemented separately, but wiring add/switch/re-auth UI through that service needs its own account-aware + // flow so this entry point does not silently change what "Switch Account" means for existing users. let result = await CodexLoginRunner.run(timeout: 120) guard !Task.isCancelled else { return } self.loginPhase = .idle diff --git a/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift b/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift index 61aa3a501..b710c1077 100644 --- a/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift @@ -186,11 +186,18 @@ struct CodexProviderImplementation: ProviderImplementation { entries.append(.text("Last spend: \(UsageFormatter.creditEventSummary(latest))", .secondary)) } } else { - let hint = context.store.lastCreditsError ?? context.metadata.creditsHint + let hint = context.store.userFacingLastCreditsError ?? context.metadata.creditsHint entries.append(.text(hint, .secondary)) } } + @MainActor + func loginMenuAction(context _: ProviderMenuLoginContext) + -> (label: String, action: MenuDescriptor.MenuAction)? + { + ("Add Account...", .addCodexAccount) + } + @MainActor func runLoginFlow(context: ProviderLoginContext) async -> Bool { await context.controller.runCodexLoginFlow() diff --git a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift index 335dbf411..ae6ec617a 100644 --- a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift +++ b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift @@ -2,6 +2,138 @@ import CodexBarCore import Foundation extension SettingsStore { + private var codexPersistedActiveSource: CodexActiveSource { + if let persistedSource = self.providerConfig(for: .codex)?.codexActiveSource { + return persistedSource + } + let source = CodexActiveSource.liveSystem + self.updateProviderConfig(provider: .codex) { entry in + entry.codexActiveSource = source + } + return source + } + + private enum ManagedCodexAccountStoreState { + case none + case selected(ManagedCodexAccount) + case unreadable + } + + private static func failClosedManagedCodexHomePath(fileManager: FileManager = .default) -> String { + ManagedCodexHomeFactory.defaultRootURL(fileManager: fileManager) + .appendingPathComponent("managed-store-unreadable", isDirectory: true) + .path + } + + private func loadManagedCodexAccounts() throws -> ManagedCodexAccountSet { + #if DEBUG + if CodexManagedRemoteHomeTestingOverride.isUnreadable(for: self) { + throw CodexManagedRemoteHomeTestingOverrideError.unreadableManagedStore + } + if let override = CodexManagedRemoteHomeTestingOverride.account(for: self) { + return ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [override]) + } + let store = if let storeURL = CodexManagedRemoteHomeTestingOverride.managedStoreURL(for: self) { + FileManagedCodexAccountStore(fileURL: storeURL) + } else { + FileManagedCodexAccountStore() + } + #else + let store = FileManagedCodexAccountStore() + #endif + + return try store.loadAccounts() + } + + private func managedCodexAccountStoreState() -> ManagedCodexAccountStoreState { + guard case let .managedAccount(id) = self.codexResolvedActiveSource else { + return .none + } + do { + let accounts = try self.loadManagedCodexAccounts() + guard let account = accounts.account(id: id) + else { + return .none + } + return .selected(account) + } catch { + return .unreadable + } + } + + var activeManagedCodexAccount: ManagedCodexAccount? { + guard case let .selected(account) = self.managedCodexAccountStoreState() else { + return nil + } + return account + } + + var activeManagedCodexRemoteHomePath: String? { + guard case .managedAccount = self.codexResolvedActiveSource else { + return nil + } + + #if DEBUG + if let override = CodexManagedRemoteHomeTestingOverride.homePath(for: self) { + return override + } + #endif + + guard case let .managedAccount(id) = self.codexResolvedActiveSource else { + return nil + } + + do { + let accounts = try self.loadManagedCodexAccounts() + // A selected managed source must never fall back to ambient ~/.codex. + return accounts.account(id: id)?.managedHomePath ?? Self.failClosedManagedCodexHomePath() + } catch { + return Self.failClosedManagedCodexHomePath() + } + } + + var activeManagedCodexCookieCacheScope: CookieHeaderCache.Scope? { + switch self.managedCodexAccountStoreState() { + case let .selected(account): + .managedAccount(account.id) + case .unreadable: + .managedStoreUnreadable + case .none: + nil + } + } + + var hasUnreadableManagedCodexAccountStore: Bool { + self.codexAccountReconciliationSnapshot.hasUnreadableAddedAccountStore + } + + private var hasUnreadableSelectedManagedCodexAccountStore: Bool { + guard case .managedAccount = self.codexResolvedActiveSource else { + return false + } + if case .unreadable = self.managedCodexAccountStoreState() { + return true + } + return false + } + + private var hasUnavailableSelectedManagedCodexAccount: Bool { + guard case let .managedAccount(id) = self.codexResolvedActiveSource else { + return false + } + guard self.hasUnreadableManagedCodexAccountStore == false else { + return false + } + do { + let accounts = try self.loadManagedCodexAccounts() + return accounts.account(id: id) == nil + } catch { + return false + } + } + var codexUsageDataSource: CodexUsageDataSource { get { let source = self.configSnapshot.providerConfig(for: .codex)?.source @@ -20,9 +152,38 @@ extension SettingsStore { } } + var codexActiveSource: CodexActiveSource { + get { + self.codexPersistedActiveSource + } + set { + self.updateProviderConfig(provider: .codex) { entry in + entry.codexActiveSource = newValue + } + } + } + + var codexResolvedActiveSource: CodexActiveSource { + self.codexResolvedActiveSourceState.resolvedSource + } + + var codexResolvedActiveSourceState: CodexResolvedActiveSource { + CodexActiveSourceResolver.resolve(from: self.codexAccountReconciliationSnapshot) + } + + @discardableResult + func persistResolvedCodexActiveSourceCorrectionIfNeeded() -> Bool { + let resolution = self.codexResolvedActiveSourceState + guard resolution.requiresPersistenceCorrection else { return false } + self.codexActiveSource = resolution.resolvedSource + return true + } + var codexCookieHeader: String { get { self.configSnapshot.providerConfig(for: .codex)?.sanitizedCookieHeader ?? "" } set { + // This is intentionally provider-scoped today. A per-managed-account manual cookie override would need + // its own storage and UI semantics so editing one account's header does not silently rewrite another's. self.updateProviderConfig(provider: .codex) { entry in entry.cookieHeader = self.normalizedConfigValue(newValue) } @@ -47,12 +208,268 @@ extension SettingsStore { func ensureCodexCookieLoaded() {} } +extension SettingsStore { + var codexAccountReconciliationSnapshot: CodexAccountReconciliationSnapshot { + self.codexAccountReconciler().loadSnapshot(environment: self.codexReconciliationEnvironment()) + } + + var codexVisibleAccountProjection: CodexVisibleAccountProjection { + CodexVisibleAccountProjection.make(from: self.codexAccountReconciliationSnapshot) + } + + var codexVisibleAccounts: [CodexVisibleAccount] { + self.codexVisibleAccountProjection.visibleAccounts + } + + @discardableResult + func selectCodexVisibleAccount(id: String) -> Bool { + guard let source = self.codexSource(forVisibleAccountID: id) else { return false } + self.codexActiveSource = source + return true + } + + func selectAuthenticatedManagedCodexAccount(_ account: ManagedCodexAccount) { + let visibleAccountID = Self.codexVisibleAccountID(for: account.email) + if self.selectCodexVisibleAccount(id: visibleAccountID) { + return + } + + self.codexActiveSource = .managedAccount(id: account.id) + _ = self.persistResolvedCodexActiveSourceCorrectionIfNeeded() + } + + func codexSource(forVisibleAccountID id: String) -> CodexActiveSource? { + self.codexVisibleAccountProjection.source(forVisibleAccountID: id) + } + + private func codexAccountReconciler() -> DefaultCodexAccountReconciler { + #if DEBUG + let liveSystemAccountOverride = CodexManagedRemoteHomeTestingOverride.liveSystemAccount(for: self) + let reconciliationEnvironmentOverride = CodexManagedRemoteHomeTestingOverride + .reconciliationEnvironment(for: self) + let managedAccountOverride = CodexManagedRemoteHomeTestingOverride.account(for: self) + let managedStoreURLOverride = CodexManagedRemoteHomeTestingOverride.managedStoreURL(for: self) + let unreadableStoreOverride = CodexManagedRemoteHomeTestingOverride.isUnreadable(for: self) + guard CodexManagedRemoteHomeTestingOverride.hasAnyOverride(for: self) else { + return DefaultCodexAccountReconciler(activeSource: self.codexPersistedActiveSource) + } + + let storeLoader: @Sendable () throws -> ManagedCodexAccountSet + if unreadableStoreOverride { + storeLoader = { throw CodexManagedRemoteHomeTestingOverrideError.unreadableManagedStore } + } else if let managedAccountOverride { + let accounts = ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [managedAccountOverride]) + storeLoader = { accounts } + } else if let managedStoreURLOverride { + let store = FileManagedCodexAccountStore(fileURL: managedStoreURLOverride) + storeLoader = { try store.loadAccounts() } + } else { + let accounts = ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: []) + storeLoader = { accounts } + } + + return DefaultCodexAccountReconciler( + storeLoader: storeLoader, + systemObserver: CodexManagedRemoteHomeTestingSystemObserver( + overrideAccount: liveSystemAccountOverride, + usesInjectedEnvironment: reconciliationEnvironmentOverride != nil), + activeSource: self.codexPersistedActiveSource) + #else + return DefaultCodexAccountReconciler(activeSource: self.codexPersistedActiveSource) + #endif + } + + private func codexReconciliationEnvironment() -> [String: String] { + #if DEBUG + if let override = CodexManagedRemoteHomeTestingOverride.reconciliationEnvironment(for: self) { + return override + } + #endif + return ProcessInfo.processInfo.environment + } + + private static func codexVisibleAccountID(for email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } +} + +#if DEBUG +private enum CodexManagedRemoteHomeTestingOverride { + private struct Override { + var account: ManagedCodexAccount? + var homePath: String? + var unreadableStore: Bool = false + var managedStoreURL: URL? + var liveSystemAccount: ObservedSystemCodexAccount? + var reconciliationEnvironment: [String: String]? + + var isEmpty: Bool { + self.account == nil && self.homePath == nil && self.unreadableStore == false && self + .managedStoreURL == nil && self.liveSystemAccount == nil && self + .reconciliationEnvironment == nil + } + } + + @MainActor + private static var values: [ObjectIdentifier: Override] = [:] + + @MainActor + private static func store(_ override: Override, for key: ObjectIdentifier) { + if override.isEmpty { + self.values.removeValue(forKey: key) + } else { + self.values[key] = override + } + } + + @MainActor + static func account(for settings: SettingsStore) -> ManagedCodexAccount? { + self.values[ObjectIdentifier(settings)]?.account + } + + @MainActor + static func setAccount(_ account: ManagedCodexAccount?, for settings: SettingsStore) { + let key = ObjectIdentifier(settings) + var override = self.values[key] ?? Override() + override.account = account + self.store(override, for: key) + } + + @MainActor + static func homePath(for settings: SettingsStore) -> String? { + self.values[ObjectIdentifier(settings)]?.homePath + } + + @MainActor + static func setHomePath(_ value: String?, for settings: SettingsStore) { + let key = ObjectIdentifier(settings) + var override = self.values[key] ?? Override() + override.homePath = value + self.store(override, for: key) + } + + @MainActor + static func isUnreadable(for settings: SettingsStore) -> Bool { + self.values[ObjectIdentifier(settings)]?.unreadableStore == true + } + + @MainActor + static func setUnreadable(_ value: Bool, for settings: SettingsStore) { + let key = ObjectIdentifier(settings) + var override = self.values[key] ?? Override() + override.unreadableStore = value + self.store(override, for: key) + } + + @MainActor + static func liveSystemAccount(for settings: SettingsStore) -> ObservedSystemCodexAccount? { + self.values[ObjectIdentifier(settings)]?.liveSystemAccount + } + + @MainActor + static func managedStoreURL(for settings: SettingsStore) -> URL? { + self.values[ObjectIdentifier(settings)]?.managedStoreURL + } + + @MainActor + static func setManagedStoreURL(_ value: URL?, for settings: SettingsStore) { + let key = ObjectIdentifier(settings) + var override = self.values[key] ?? Override() + override.managedStoreURL = value + self.store(override, for: key) + } + + @MainActor + static func setLiveSystemAccount(_ account: ObservedSystemCodexAccount?, for settings: SettingsStore) { + let key = ObjectIdentifier(settings) + var override = self.values[key] ?? Override() + override.liveSystemAccount = account + self.store(override, for: key) + } + + @MainActor + static func reconciliationEnvironment(for settings: SettingsStore) -> [String: String]? { + self.values[ObjectIdentifier(settings)]?.reconciliationEnvironment + } + + @MainActor + static func setReconciliationEnvironment(_ environment: [String: String]?, for settings: SettingsStore) { + let key = ObjectIdentifier(settings) + var override = self.values[key] ?? Override() + override.reconciliationEnvironment = environment + self.store(override, for: key) + } + + @MainActor + static func hasAnyOverride(for settings: SettingsStore) -> Bool { + self.values[ObjectIdentifier(settings)]?.isEmpty == false + } +} + +private enum CodexManagedRemoteHomeTestingOverrideError: Error { + case unreadableManagedStore +} + +private struct CodexManagedRemoteHomeTestingSystemObserver: CodexSystemAccountObserving { + let overrideAccount: ObservedSystemCodexAccount? + let usesInjectedEnvironment: Bool + + func loadSystemAccount(environment: [String: String]) throws -> ObservedSystemCodexAccount? { + if let overrideAccount { + return overrideAccount + } + guard self.usesInjectedEnvironment else { + return nil + } + return try DefaultCodexSystemAccountObserver().loadSystemAccount(environment: environment) + } +} + +extension SettingsStore { + var _test_activeManagedCodexRemoteHomePath: String? { + get { CodexManagedRemoteHomeTestingOverride.homePath(for: self) } + set { CodexManagedRemoteHomeTestingOverride.setHomePath(newValue, for: self) } + } + + var _test_activeManagedCodexAccount: ManagedCodexAccount? { + get { CodexManagedRemoteHomeTestingOverride.account(for: self) } + set { CodexManagedRemoteHomeTestingOverride.setAccount(newValue, for: self) } + } + + var _test_unreadableManagedCodexAccountStore: Bool { + get { CodexManagedRemoteHomeTestingOverride.isUnreadable(for: self) } + set { CodexManagedRemoteHomeTestingOverride.setUnreadable(newValue, for: self) } + } + + var _test_managedCodexAccountStoreURL: URL? { + get { CodexManagedRemoteHomeTestingOverride.managedStoreURL(for: self) } + set { CodexManagedRemoteHomeTestingOverride.setManagedStoreURL(newValue, for: self) } + } + + var _test_liveSystemCodexAccount: ObservedSystemCodexAccount? { + get { CodexManagedRemoteHomeTestingOverride.liveSystemAccount(for: self) } + set { CodexManagedRemoteHomeTestingOverride.setLiveSystemAccount(newValue, for: self) } + } + + var _test_codexReconciliationEnvironment: [String: String]? { + get { CodexManagedRemoteHomeTestingOverride.reconciliationEnvironment(for: self) } + set { CodexManagedRemoteHomeTestingOverride.setReconciliationEnvironment(newValue, for: self) } + } +} +#endif + extension SettingsStore { func codexSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.CodexProviderSettings { ProviderSettingsSnapshot.CodexProviderSettings( usageDataSource: self.codexUsageDataSource, cookieSource: self.codexSnapshotCookieSource(tokenOverride: tokenOverride), - manualCookieHeader: self.codexSnapshotCookieHeader(tokenOverride: tokenOverride)) + manualCookieHeader: self.codexSnapshotCookieHeader(tokenOverride: tokenOverride), + managedAccountStoreUnreadable: self.hasUnreadableSelectedManagedCodexAccountStore, + managedAccountTargetUnavailable: self.hasUnavailableSelectedManagedCodexAccount) } private static func codexUsageDataSource(from source: ProviderSourceMode?) -> CodexUsageDataSource { diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift new file mode 100644 index 000000000..c6f476717 --- /dev/null +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift @@ -0,0 +1,278 @@ +import CodexBarCore +import Foundation + +enum CodexAccountScopedRefreshPhase: Sendable { + case invalidated + case usage + case credits + case dashboard + case completed +} + +struct CodexAccountScopedRefreshGuard: Equatable, Sendable { + let source: CodexActiveSource + let accountKey: String? +} + +@MainActor +extension UsageStore { + func refreshCodexAccountScopedState( + allowDisabled: Bool = false, + phaseDidChange: (@MainActor (CodexAccountScopedRefreshPhase) -> Void)? = nil) + async + { + let refreshStartedAt = Date() + self.prepareRefreshState(for: .codex) + if self.prepareCodexAccountScopedRefreshIfNeeded() { + phaseDidChange?(.invalidated) + } + + await self.refreshProvider(.codex, allowDisabled: allowDisabled) + phaseDidChange?(.usage) + await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) + phaseDidChange?(.credits) + + if self.settings.codexCookieSource.isEnabled { + let expectedGuard = self.currentCodexAccountScopedRefreshGuard() + await self.refreshOpenAIDashboardIfNeeded( + force: true, + expectedGuard: expectedGuard, + allowCodexUsageBackfill: true) + phaseDidChange?(.dashboard) + } + + if self.openAIDashboardRequiresLogin { + await self.refreshProvider(.codex, allowDisabled: allowDisabled) + phaseDidChange?(.usage) + await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) + phaseDidChange?(.credits) + } + + self.persistWidgetSnapshot(reason: "codex-account-refresh") + phaseDidChange?(.completed) + } + + @discardableResult + func prepareCodexAccountScopedRefreshIfNeeded() -> Bool { + let currentGuard = self.currentCodexAccountScopedRefreshGuard( + preferCurrentSnapshot: false, + allowLastKnownLiveFallback: false) + let previousGuard = self.lastCodexAccountScopedRefreshGuard + self.lastCodexAccountScopedRefreshGuard = currentGuard + + guard previousGuard != nil, previousGuard != currentGuard else { return false } + + self.snapshots.removeValue(forKey: .codex) + self.errors[.codex] = nil + self.lastSourceLabels.removeValue(forKey: .codex) + self.lastFetchAttempts.removeValue(forKey: .codex) + self.accountSnapshots.removeValue(forKey: .codex) + self.failureGates[.codex]?.reset() + self.lastKnownSessionRemaining.removeValue(forKey: .codex) + self.lastKnownSessionWindowSource.removeValue(forKey: .codex) + + self.credits = nil + self.lastCreditsError = nil + self.lastCreditsSnapshot = nil + self.lastCreditsSnapshotAccountKey = nil + self.creditsFailureStreak = 0 + + self.clearCodexOpenAIWebStateForAccountTransition(targetEmail: self.codexAccountEmailForOpenAIDashboard()) + + self.persistWidgetSnapshot(reason: "codex-account-invalidate") + return true + } + + func seedCodexAccountScopedRefreshGuard( + source: CodexActiveSource? = nil, + accountEmail: String?) + { + guard let accountKey = Self.normalizeCodexAccountScopedKey(accountEmail) else { return } + self.lastCodexAccountScopedRefreshGuard = CodexAccountScopedRefreshGuard( + source: source ?? self.settings.codexResolvedActiveSource, + accountKey: accountKey) + } + + func currentCodexAccountScopedRefreshGuard( + preferCurrentSnapshot: Bool = true, + allowLastKnownLiveFallback: Bool = true) -> CodexAccountScopedRefreshGuard + { + CodexAccountScopedRefreshGuard( + source: self.settings.codexResolvedActiveSource, + accountKey: self.codexAccountScopedRefreshKey( + preferCurrentSnapshot: preferCurrentSnapshot, + allowLastKnownLiveFallback: allowLastKnownLiveFallback)) + } + + func currentCodexOpenAIWebRefreshGuard() -> CodexAccountScopedRefreshGuard { + let accountKey: String? = switch self.settings.codexResolvedActiveSource { + case .liveSystem: + Self + .normalizeCodexAccountScopedKey(self.settings.codexAccountReconciliationSnapshot.liveSystemAccount? + .email) + case .managedAccount: + Self.normalizeCodexAccountScopedKey(self.settings.activeManagedCodexAccount?.email) + } + return CodexAccountScopedRefreshGuard( + source: self.settings.codexResolvedActiveSource, + accountKey: accountKey) + } + + func shouldApplyCodexUsageResult( + expectedGuard: CodexAccountScopedRefreshGuard, + usage: UsageSnapshot) -> Bool + { + let currentGuard = self.currentCodexAccountScopedRefreshGuard() + guard currentGuard.source == expectedGuard.source else { return false } + + if let expectedKey = expectedGuard.accountKey { + return currentGuard.accountKey == expectedKey + } + + let resultKey = Self.normalizeCodexAccountScopedKey(usage.accountEmail(for: .codex)) + if let currentKey = currentGuard.accountKey { + return resultKey == currentKey + } + + switch currentGuard.source { + case .liveSystem: + return resultKey != nil + case .managedAccount: + return false + } + } + + func shouldApplyCodexScopedFailure(expectedGuard: CodexAccountScopedRefreshGuard) -> Bool { + let currentGuard = self.currentCodexAccountScopedRefreshGuard() + guard currentGuard.source == expectedGuard.source else { return false } + + if let expectedKey = expectedGuard.accountKey { + return currentGuard.accountKey == expectedKey + } + + return currentGuard.accountKey == nil + } + + func shouldApplyCodexScopedNonUsageResult(expectedGuard: CodexAccountScopedRefreshGuard) -> Bool { + let currentGuard = self.currentCodexAccountScopedRefreshGuard() + guard currentGuard.source == expectedGuard.source else { return false } + guard let expectedKey = expectedGuard.accountKey else { return false } + return currentGuard.accountKey == expectedKey + } + + func shouldApplyOpenAIDashboardResult( + expectedGuard: CodexAccountScopedRefreshGuard, + dashboardAccountEmail: String?) -> Bool + { + if let expectedKey = expectedGuard.accountKey { + let currentGuard = self.currentCodexAccountScopedRefreshGuard() + guard currentGuard.source == expectedGuard.source else { return false } + return currentGuard.accountKey == expectedKey + } + + let currentGuard = self.currentCodexOpenAIWebRefreshGuard() + guard currentGuard.source == expectedGuard.source else { return false } + guard case .liveSystem = expectedGuard.source else { return false } + guard currentGuard.accountKey == nil else { return false } + guard let dashboardKey = Self.normalizeCodexAccountScopedKey(dashboardAccountEmail) else { return false } + let currentTargetKey = Self.normalizeCodexAccountScopedKey(self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: true, + allowLastKnownLiveFallback: false)) + if let currentTargetKey { + return dashboardKey == currentTargetKey + } + return true + } + + func rememberLiveSystemCodexEmailIfNeeded(_ email: String?) { + guard case .liveSystem = self.settings.codexResolvedActiveSource else { return } + guard let normalized = Self.normalizeCodexAccountScopedEmail(email) else { return } + self.lastKnownLiveSystemCodexEmail = normalized + } + + func codexAccountScopedRefreshKey( + preferCurrentSnapshot: Bool = true, + allowLastKnownLiveFallback: Bool = true) -> String? + { + Self.normalizeCodexAccountScopedKey( + self.codexAccountScopedRefreshEmail( + preferCurrentSnapshot: preferCurrentSnapshot, + allowLastKnownLiveFallback: allowLastKnownLiveFallback)) + } + + func codexAccountScopedRefreshEmail( + preferCurrentSnapshot: Bool = true, + allowLastKnownLiveFallback: Bool = true) -> String? + { + switch self.settings.codexResolvedActiveSource { + case .liveSystem: + let liveSystem = Self.normalizeCodexAccountScopedEmail( + self.settings.codexAccountReconciliationSnapshot.liveSystemAccount?.email) + if let liveSystem { + self.lastKnownLiveSystemCodexEmail = liveSystem + return liveSystem + } + + if preferCurrentSnapshot, + let snapshotEmail = Self + .normalizeCodexAccountScopedEmail(self.snapshots[.codex]?.accountEmail(for: .codex)) + { + self.lastKnownLiveSystemCodexEmail = snapshotEmail + return snapshotEmail + } + + if allowLastKnownLiveFallback, + let lastKnown = Self.normalizeCodexAccountScopedEmail(self.lastKnownLiveSystemCodexEmail) + { + return lastKnown + } + + return nil + case .managedAccount: + if self.settings.codexSettingsSnapshot(tokenOverride: nil).managedAccountStoreUnreadable { + return nil + } + return Self.normalizeCodexAccountScopedEmail(self.settings.activeManagedCodexAccount?.email) + } + } + + private func clearCodexOpenAIWebStateForAccountTransition(targetEmail: String?) { + self.invalidateOpenAIDashboardRefreshTask() + if self.settings.codexCookieSource.isEnabled, + let normalizedTarget = Self.normalizeCodexAccountScopedEmail(targetEmail) + { + let previous = self.lastOpenAIDashboardTargetEmail + self.lastOpenAIDashboardTargetEmail = normalizedTarget + if let previous, !previous.isEmpty, previous != normalizedTarget { + self.openAIWebAccountDidChange = true + self.openAIDashboardCookieImportStatus = "Codex account changed; importing browser cookies…" + } else { + self.openAIDashboardCookieImportStatus = nil + } + self.openAIDashboardRequiresLogin = true + } else { + self.lastOpenAIDashboardTargetEmail = Self.normalizeCodexAccountScopedEmail(targetEmail) + self.openAIWebAccountDidChange = false + self.openAIDashboardRequiresLogin = false + self.openAIDashboardCookieImportStatus = nil + } + + self.openAIDashboard = nil + self.lastOpenAIDashboardSnapshot = nil + self.lastOpenAIDashboardError = nil + self.openAIDashboardCookieImportDebugLog = nil + self.lastOpenAIDashboardCookieImportAttemptAt = nil + self.lastOpenAIDashboardCookieImportEmail = nil + } + + static func normalizeCodexAccountScopedEmail(_ email: String?) -> String? { + guard let trimmed = email?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } + + static func normalizeCodexAccountScopedKey(_ email: String?) -> String? { + self.normalizeCodexAccountScopedEmail(email)?.lowercased() + } +} diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift index 35b184897..7b65f943d 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift @@ -4,18 +4,36 @@ import Foundation @MainActor extension UsageStore { nonisolated static let codexSnapshotWaitTimeoutSeconds: TimeInterval = 6 + nonisolated static let codexRefreshStartGraceSeconds: TimeInterval = 0.25 nonisolated static let codexSnapshotPollIntervalNanoseconds: UInt64 = 100_000_000 + func codexCreditsFetcher() -> UsageFetcher { + // Credits are remote Codex account state, so they need the same managed-home routing as the + // primary Codex usage fetch. Local token-cost scanning intentionally stays ambient-system scoped. + self.makeFetchContext(provider: .codex, override: nil).fetcher + } + func refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: Date? = nil) async { guard self.isEnabled(.codex) else { return } + var expectedGuard = self.currentCodexAccountScopedRefreshGuard() + if expectedGuard.accountKey == nil, + let minimumSnapshotUpdatedAt, + case .liveSystem = expectedGuard.source + { + _ = await self.waitForCodexSnapshotOrRefreshCompletion(minimumUpdatedAt: minimumSnapshotUpdatedAt) + expectedGuard = self.currentCodexAccountScopedRefreshGuard() + } + guard expectedGuard.accountKey != nil else { return } do { - let credits = try await self.codexFetcher.loadLatestCredits( - keepCLISessionsAlive: self.settings.debugKeepCLISessionsAlive) + let credits = try await self.loadLatestCodexCredits() + guard self.shouldApplyCodexScopedNonUsageResult(expectedGuard: expectedGuard) else { return } await MainActor.run { self.credits = credits self.lastCreditsError = nil self.lastCreditsSnapshot = credits + self.lastCreditsSnapshotAccountKey = expectedGuard.accountKey self.creditsFailureStreak = 0 + self.lastCodexAccountScopedRefreshGuard = expectedGuard } let codexSnapshot = await MainActor.run { self.snapshots[.codex] @@ -37,10 +55,14 @@ extension UsageStore { } catch { let message = error.localizedDescription if message.localizedCaseInsensitiveContains("data not available yet") { + guard self.shouldApplyCodexScopedNonUsageResult(expectedGuard: expectedGuard) else { return } await MainActor.run { - if let cached = self.lastCreditsSnapshot { + if let cached = self.lastCreditsSnapshot, + self.lastCreditsSnapshotAccountKey == expectedGuard.accountKey + { self.credits = cached self.lastCreditsError = nil + self.lastCodexAccountScopedRefreshGuard = expectedGuard } else { self.credits = nil self.lastCreditsError = "Codex credits are still loading; will retry shortly." @@ -49,13 +71,17 @@ extension UsageStore { return } + guard self.shouldApplyCodexScopedNonUsageResult(expectedGuard: expectedGuard) else { return } await MainActor.run { self.creditsFailureStreak += 1 - if let cached = self.lastCreditsSnapshot { + if let cached = self.lastCreditsSnapshot, + self.lastCreditsSnapshotAccountKey == expectedGuard.accountKey + { self.credits = cached let stamp = cached.updatedAt.formatted(date: .abbreviated, time: .shortened) self.lastCreditsError = "Last Codex credits refresh failed: \(message). Cached values from \(stamp)." + self.lastCodexAccountScopedRefreshGuard = expectedGuard } else { self.lastCreditsError = message self.credits = nil @@ -64,6 +90,14 @@ extension UsageStore { } } + private func loadLatestCodexCredits() async throws -> CreditsSnapshot { + if let override = self._test_codexCreditsLoaderOverride { + return try await override() + } + return try await self.codexCreditsFetcher().loadLatestCredits( + keepCLISessionsAlive: self.settings.debugKeepCLISessionsAlive) + } + func waitForCodexSnapshot(minimumUpdatedAt: Date) async -> UsageSnapshot? { let deadline = Date().addingTimeInterval(Self.codexSnapshotWaitTimeoutSeconds) @@ -80,6 +114,38 @@ extension UsageStore { return nil } + func waitForCodexSnapshotOrRefreshCompletion(minimumUpdatedAt: Date) async -> UsageSnapshot? { + let deadline = Date().addingTimeInterval(Self.codexSnapshotWaitTimeoutSeconds) + let refreshStartDeadline = Date().addingTimeInterval(Self.codexRefreshStartGraceSeconds) + + while Date() < deadline { + if Task.isCancelled { return nil } + let state = await MainActor.run { + ( + snapshot: self.snapshots[.codex], + isRefreshing: self.refreshingProviders.contains(.codex), + hasAttempts: !(self.lastFetchAttempts[.codex] ?? []).isEmpty, + hasError: self.errors[.codex] != nil) + } + if let snapshot = state.snapshot, snapshot.updatedAt >= minimumUpdatedAt { + return snapshot + } + if !state.isRefreshing, state.hasAttempts || state.hasError { + return nil + } + if !state.isRefreshing, + !state.hasAttempts, + !state.hasError, + Date() >= refreshStartDeadline + { + return nil + } + try? await Task.sleep(nanoseconds: Self.codexSnapshotPollIntervalNanoseconds) + } + + return nil + } + func scheduleCodexPlanHistoryBackfill( minimumSnapshotUpdatedAt: Date) { diff --git a/Sources/CodexBar/SettingsStore+MenuObservation.swift b/Sources/CodexBar/SettingsStore+MenuObservation.swift index f01cc49fa..f799822ef 100644 --- a/Sources/CodexBar/SettingsStore+MenuObservation.swift +++ b/Sources/CodexBar/SettingsStore+MenuObservation.swift @@ -28,6 +28,7 @@ extension SettingsStore { _ = self.showOptionalCreditsAndExtraUsage _ = self.openAIWebAccessEnabled _ = self.codexUsageDataSource + _ = self.codexActiveSource _ = self.claudeUsageDataSource _ = self.kiloUsageDataSource _ = self.kiloExtrasEnabled diff --git a/Sources/CodexBar/StatusItemController+Actions.swift b/Sources/CodexBar/StatusItemController+Actions.swift index ad65b0879..930e42c00 100644 --- a/Sources/CodexBar/StatusItemController+Actions.swift +++ b/Sources/CodexBar/StatusItemController+Actions.swift @@ -105,6 +105,33 @@ extension StatusItemController { NSWorkspace.shared.open(url) } + @objc func addManagedCodexAccountFromMenu(_: NSMenuItem) { + guard self.managedCodexAccountCoordinator.isAuthenticatingManagedAccount == false else { + self.loginLogger.info("Add Account tap ignored: managed Codex login already in-flight") + return + } + guard self.settings.hasUnreadableManagedCodexAccountStore == false else { + self.presentLoginAlert( + title: "Managed Codex accounts unavailable", + message: "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 runSwitchAccount(_ sender: NSMenuItem) { if self.loginTask != nil { self.loginLogger.info("Switch Account tap ignored: login already in-flight") @@ -210,24 +237,34 @@ extension StatusItemController { } func presentCodexLoginResult(_ result: CodexLoginRunner.Result) { - switch result.outcome { - case .success: - return - case .missingBinary: - self.presentLoginAlert( - title: "Codex CLI not found", - message: "Install the Codex CLI (npm i -g @openai/codex) and try again.") - case let .launchFailed(message): - self.presentLoginAlert(title: "Could not start codex login", message: message) - case .timedOut: - self.presentLoginAlert( - title: "Codex login timed out", - message: self.trimmedLoginOutput(result.output)) - case let .failed(status): - let statusLine = "codex login exited with status \(status)." - let message = self.trimmedLoginOutput(result.output.isEmpty ? statusLine : result.output) - self.presentLoginAlert(title: "Codex login failed", message: message) + guard let info = CodexLoginAlertPresentation.alertInfo(for: result) else { return } + self.presentLoginAlert(title: info.title, message: info.message) + } + + private func presentManagedCodexAccountError(_ error: Error) { + let info: LoginAlertInfo + if let error = error as? ManagedCodexAccountCoordinatorError, + error == .authenticationInProgress + { + info = LoginAlertInfo( + title: "Codex account login already running", + message: "Wait for the current managed Codex login to finish before adding another account.") + } else if let error = error as? ManagedCodexAccountServiceError { + let message = switch error { + case .loginFailed: + "Managed Codex login did not complete. Try again after finishing the browser login flow." + case .missingEmail: + "Codex login completed, but no account email was available. " + + "Try again after confirming the account is fully signed in." + case let .unsafeManagedHome(path): + "CodexBar refused to modify an unexpected managed home path: \(path)" + } + info = LoginAlertInfo(title: "Could not add Codex account", message: message) + } else { + info = LoginAlertInfo(title: "Could not add Codex account", message: error.localizedDescription) } + + self.presentLoginAlert(title: info.title, message: info.message) } func presentClaudeLoginResult(_ result: ClaudeLoginRunner.Result) { diff --git a/Sources/CodexBar/StatusItemController+Menu.swift b/Sources/CodexBar/StatusItemController+Menu.swift index 484be310a..7436e381f 100644 --- a/Sources/CodexBar/StatusItemController+Menu.swift +++ b/Sources/CodexBar/StatusItemController+Menu.swift @@ -63,6 +63,11 @@ extension StatusItemController { let showSwitcher: Bool } + private struct CodexAccountMenuDisplay { + let accounts: [CodexVisibleAccount] + let activeVisibleAccountID: String? + } + private func menuCardWidth(for providers: [UsageProvider], menu: NSMenu? = nil) -> CGFloat { _ = menu return Self.menuCardBaseWidth @@ -162,13 +167,16 @@ extension StatusItemController { } let menuWidth = self.menuCardWidth(for: enabledProviders, menu: menu) let currentProvider = selectedProvider ?? enabledProviders.first ?? .codex + let codexAccountDisplay = isOverviewSelected ? nil : self.codexAccountMenuDisplay(for: currentProvider) let tokenAccountDisplay = isOverviewSelected ? nil : self.tokenAccountMenuDisplay(for: currentProvider) let showAllTokenAccounts = tokenAccountDisplay?.showAll ?? false let openAIContext = self.openAIWebContext( currentProvider: currentProvider, showAllTokenAccounts: showAllTokenAccounts) - let hasTokenAccountSwitcher = menu.items.contains { $0.view is TokenAccountSwitcherView } + let hasAuxiliarySwitcher = menu.items.contains { + $0.view is TokenAccountSwitcherView || $0.view is CodexAccountSwitcherView + } let switcherProvidersMatch = enabledProviders == self.lastSwitcherProviders let switcherUsageBarsShowUsedMatch = self.settings.usageBarsShowUsed == self.lastSwitcherUsageBarsShowUsed let switcherSelectionMatches = switcherSelection == self.lastMergedSwitcherSelection @@ -180,8 +188,9 @@ extension StatusItemController { switcherUsageBarsShowUsedMatch && switcherSelectionMatches && switcherOverviewAvailabilityMatches && + codexAccountDisplay == nil && tokenAccountDisplay == nil && - !hasTokenAccountSwitcher && + !hasAuxiliarySwitcher && !menu.items.isEmpty && menu.items.first?.view is ProviderSwitcherView @@ -217,6 +226,7 @@ extension StatusItemController { self.lastMergedSwitcherSelection = switcherSelection self.lastSwitcherIncludesOverview = includesOverview } + self.addCodexAccountSwitcherIfNeeded(to: menu, display: codexAccountDisplay) self.addTokenAccountSwitcherIfNeeded(to: menu, display: tokenAccountDisplay) let menuContext = MenuCardContext( currentProvider: currentProvider, @@ -359,6 +369,13 @@ extension StatusItemController { menu.addItem(.separator()) } + private func addCodexAccountSwitcherIfNeeded(to menu: NSMenu, display: CodexAccountMenuDisplay?) { + guard let display else { return } + let switcherItem = self.makeCodexAccountSwitcherItem(display: display, menu: menu) + menu.addItem(switcherItem) + menu.addItem(.separator()) + } + @discardableResult private func addOverviewRows( to menu: NSMenu, @@ -536,6 +553,11 @@ extension StatusItemController { { item.isEnabled = false self.applySubtitle(subtitle, to: item, title: title) + } else if case .addCodexAccount = action, + let subtitle = self.codexAddAccountSubtitle() + { + item.isEnabled = false + self.applySubtitle(subtitle, to: item, title: title) } menu.addItem(item) case .divider: @@ -667,6 +689,42 @@ extension StatusItemController { return item } + private func makeCodexAccountSwitcherItem( + display: CodexAccountMenuDisplay, + menu: NSMenu) -> NSMenuItem + { + let view = CodexAccountSwitcherView( + accounts: display.accounts, + selectedAccountID: display.activeVisibleAccountID, + width: self.menuCardWidth(for: self.store.enabledProvidersForDisplay(), menu: menu), + onSelect: { [weak self, weak menu] visibleAccountID in + guard let self, let menu else { return } + guard self.settings.selectCodexVisibleAccount(id: visibleAccountID) else { return } + if self.store.prepareCodexAccountScopedRefreshIfNeeded() { + self.refreshOpenMenuIfStillVisible(menu, provider: .codex) + } + Task { @MainActor in + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await self.store.refreshCodexAccountScopedState( + allowDisabled: true, + phaseDidChange: { [weak self, weak menu] _ in + guard let self, let menu else { return } + guard self.settings.codexVisibleAccountProjection.activeVisibleAccountID == + visibleAccountID + else { + return + } + self.refreshOpenMenuIfStillVisible(menu, provider: .codex) + }) + } + } + }) + let item = NSMenuItem() + item.view = view + item.isEnabled = false + return item + } + private func resolvedMenuProvider(enabledProviders: [UsageProvider]? = nil) -> UsageProvider? { let enabled = enabledProviders ?? self.store.enabledProvidersForDisplay() if enabled.isEmpty { return .codex } @@ -710,6 +768,15 @@ extension StatusItemController { showSwitcher: !showAll) } + private func codexAccountMenuDisplay(for provider: UsageProvider) -> CodexAccountMenuDisplay? { + guard provider == .codex else { return nil } + let projection = self.settings.codexVisibleAccountProjection + guard projection.visibleAccounts.count > 1 else { return nil } + return CodexAccountMenuDisplay( + accounts: projection.visibleAccounts, + activeVisibleAccountID: projection.activeVisibleAccountID) + } + private func menuNeedsRefresh(_ menu: NSMenu) -> Bool { let key = ObjectIdentifier(menu) return self.menuVersions[key] != self.menuContentVersion @@ -759,6 +826,33 @@ extension StatusItemController { return self.store.enabledProvidersForDisplay().first ?? .codex } + func refreshOpenMenuIfStillVisible(_ menu: NSMenu, provider: UsageProvider?) { + self.rebuildOpenMenuIfStillVisible(menu, provider: provider) + Task { @MainActor [weak self, weak menu] in + guard let self, let menu else { return } + #if DEBUG + if let override = self._test_openMenuRefreshYieldOverride { + await override() + } else { + await Task.yield() + } + #else + await Task.yield() + #endif + self.rebuildOpenMenuIfStillVisible(menu, provider: provider) + } + } + + private func rebuildOpenMenuIfStillVisible(_ menu: NSMenu, provider: UsageProvider?) { + guard self.openMenus[ObjectIdentifier(menu)] != nil else { return } + self.populateMenu(menu, provider: provider) + self.markMenuFresh(menu) + self.applyIcon(phase: nil) + #if DEBUG + self._test_openMenuRebuildObserver?(menu) + #endif + } + private func scheduleOpenMenuRefresh(for menu: NSMenu) { // Kick off a user-initiated refresh on open (non-forced) and re-check after a delay. // NEVER block menu opening with network requests. @@ -1054,6 +1148,7 @@ extension StatusItemController { case .refreshAugmentSession: (#selector(self.refreshAugmentSession), nil) case .dashboard: (#selector(self.openDashboard), nil) case .statusPage: (#selector(self.openStatusPage), nil) + case .addCodexAccount: (#selector(self.addManagedCodexAccountFromMenu(_:)), nil) case let .switchAccount(provider): (#selector(self.runSwitchAccount(_:)), provider.rawValue) case let .openTerminal(command): (#selector(self.openTerminalCommand(_:)), command) case let .loginToProvider(url): (#selector(self.openLoginToProvider(_:)), url) @@ -1064,6 +1159,14 @@ extension StatusItemController { } } + private func codexAddAccountSubtitle() -> String? { + if self.settings.hasUnreadableManagedCodexAccountStore { + return "Managed account storage unavailable" + } + guard self.managedCodexAccountCoordinator.isAuthenticatingManagedAccount else { return nil } + return "Managed Codex login in progress…" + } + @MainActor private protocol MenuCardHighlighting: AnyObject { func setHighlighted(_ highlighted: Bool) @@ -1435,9 +1538,9 @@ extension StatusItemController { let tokenError: String? if target == .codex, snapshotOverride == nil { credits = self.store.credits - creditsError = self.store.lastCreditsError + creditsError = self.store.userFacingLastCreditsError dashboard = self.store.openAIDashboardRequiresLogin ? nil : self.store.openAIDashboard - dashboardError = self.store.lastOpenAIDashboardError + dashboardError = self.store.userFacingLastOpenAIDashboardError tokenSnapshot = self.store.tokenSnapshot(for: target) tokenError = self.store.tokenError(for: target) } else if target == .claude || target == .vertexai, snapshotOverride == nil { @@ -1472,9 +1575,9 @@ extension StatusItemController { dashboardError: dashboardError, tokenSnapshot: tokenSnapshot, tokenError: tokenError, - account: self.account, + account: self.store.accountInfo(for: target), isRefreshing: self.store.shouldShowRefreshingMenuCard(for: target), - lastError: errorOverride ?? self.store.error(for: target), + lastError: errorOverride ?? self.store.userFacingError(for: target), usageBarsShowUsed: self.settings.usageBarsShowUsed, resetTimeDisplayStyle: self.settings.resetTimeDisplayStyle, tokenCostUsageEnabled: self.settings.isCostUsageEffectivelyEnabled(for: target), diff --git a/Sources/CodexBar/StatusItemController+SwitcherViews.swift b/Sources/CodexBar/StatusItemController+SwitcherViews.swift index 92c334231..4c56de1ab 100644 --- a/Sources/CodexBar/StatusItemController+SwitcherViews.swift +++ b/Sources/CodexBar/StatusItemController+SwitcherViews.swift @@ -895,3 +895,110 @@ final class TokenAccountSwitcherView: NSView { self.onSelect(index) } } + +final class CodexAccountSwitcherView: NSView { + private let accounts: [CodexVisibleAccount] + private let onSelect: (String) -> Void + private var selectedAccountID: String + private var buttons: [NSButton] = [] + private let rowSpacing: CGFloat = 4 + private let rowHeight: CGFloat = 26 + private let selectedBackground = NSColor.controlAccentColor.cgColor + private let unselectedBackground = NSColor.clear.cgColor + private let selectedTextColor = NSColor.white + private let unselectedTextColor = NSColor.secondaryLabelColor + + init( + accounts: [CodexVisibleAccount], + selectedAccountID: String?, + width: CGFloat, + onSelect: @escaping (String) -> Void) + { + self.accounts = accounts + self.onSelect = onSelect + self.selectedAccountID = selectedAccountID ?? accounts.first?.id ?? "" + let useTwoRows = accounts.count > 3 + let rows = useTwoRows ? 2 : 1 + let height = self.rowHeight * CGFloat(rows) + (useTwoRows ? self.rowSpacing : 0) + super.init(frame: NSRect(x: 0, y: 0, width: width, height: height)) + self.wantsLayer = true + self.buildButtons(useTwoRows: useTwoRows) + self.updateButtonStyles() + } + + @available(*, unavailable) + required init?(coder: NSCoder) { + nil + } + + private func buildButtons(useTwoRows: Bool) { + let perRow = useTwoRows ? Int(ceil(Double(self.accounts.count) / 2.0)) : self.accounts.count + let rows: [[CodexVisibleAccount]] = { + if !useTwoRows { return [self.accounts] } + let first = Array(self.accounts.prefix(perRow)) + let second = Array(self.accounts.dropFirst(perRow)) + return [first, second] + }() + + let stack = NSStackView() + stack.orientation = .vertical + stack.alignment = .centerX + stack.spacing = self.rowSpacing + stack.translatesAutoresizingMaskIntoConstraints = false + + for rowAccounts in rows { + let row = NSStackView() + row.orientation = .horizontal + row.alignment = .centerY + row.distribution = .fillEqually + row.spacing = self.rowSpacing + row.translatesAutoresizingMaskIntoConstraints = false + + for account in rowAccounts { + let button = PaddedToggleButton( + title: account.email, + target: self, + action: #selector(self.handleSelect)) + button.identifier = NSUserInterfaceItemIdentifier(account.id) + button.toolTip = account.email + button.isBordered = false + button.setButtonType(.toggle) + button.controlSize = .small + button.font = NSFont.systemFont(ofSize: NSFont.smallSystemFontSize) + button.wantsLayer = true + button.layer?.cornerRadius = 6 + row.addArrangedSubview(button) + self.buttons.append(button) + } + + stack.addArrangedSubview(row) + } + + self.addSubview(stack) + NSLayoutConstraint.activate([ + stack.leadingAnchor.constraint(equalTo: self.leadingAnchor, constant: 6), + stack.trailingAnchor.constraint(equalTo: self.trailingAnchor, constant: -6), + stack.topAnchor.constraint(equalTo: self.topAnchor), + stack.bottomAnchor.constraint(equalTo: self.bottomAnchor), + stack.heightAnchor.constraint(equalToConstant: self.rowHeight * CGFloat(rows.count) + + (useTwoRows ? self.rowSpacing : 0)), + ]) + } + + private func updateButtonStyles() { + for button in self.buttons { + let selected = button.identifier?.rawValue == self.selectedAccountID + button.state = selected ? .on : .off + button.layer?.backgroundColor = selected ? self.selectedBackground : self.unselectedBackground + button.contentTintColor = selected ? self.selectedTextColor : self.unselectedTextColor + } + } + + @objc private func handleSelect(_ sender: NSButton) { + guard let accountID = sender.identifier?.rawValue else { return } + guard self.accounts.contains(where: { $0.id == accountID }) else { return } + self.selectedAccountID = accountID + self.updateButtonStyles() + self.onSelect(accountID) + } +} diff --git a/Sources/CodexBar/StatusItemController.swift b/Sources/CodexBar/StatusItemController.swift index 1d95ccd91..26330607f 100644 --- a/Sources/CodexBar/StatusItemController.swift +++ b/Sources/CodexBar/StatusItemController.swift @@ -16,16 +16,24 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin // Disable SwiftUI menu cards + menu refresh work in tests to avoid swiftpm-testing-helper crashes. static var menuCardRenderingEnabled = !SettingsStore.isRunningTests static var menuRefreshEnabled = !SettingsStore.isRunningTests - typealias Factory = (UsageStore, SettingsStore, AccountInfo, UpdaterProviding, PreferencesSelection) + typealias Factory = ( + UsageStore, + SettingsStore, + AccountInfo, + UpdaterProviding, + PreferencesSelection, + ManagedCodexAccountCoordinator) -> StatusItemControlling - static let defaultFactory: Factory = { store, settings, account, updater, selection in - StatusItemController( - store: store, - settings: settings, - account: account, - updater: updater, - preferencesSelection: selection) - } + static let defaultFactory: Factory = + { store, settings, account, updater, selection, managedCodexAccountCoordinator in + StatusItemController( + store: store, + settings: settings, + account: account, + updater: updater, + preferencesSelection: selection, + managedCodexAccountCoordinator: managedCodexAccountCoordinator) + } static var factory: Factory = StatusItemController.defaultFactory @@ -33,6 +41,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin let settings: SettingsStore let account: AccountInfo let updater: UpdaterProviding + let managedCodexAccountCoordinator: ManagedCodexAccountCoordinator private let statusBar: NSStatusBar var statusItem: NSStatusItem var statusItems: [UsageProvider: NSStatusItem] = [:] @@ -45,6 +54,10 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin var fallbackMenu: NSMenu? var openMenus: [ObjectIdentifier: NSMenu] = [:] var menuRefreshTasks: [ObjectIdentifier: Task] = [:] + #if DEBUG + var _test_openMenuRefreshYieldOverride: (@MainActor () async -> Void)? + var _test_openMenuRebuildObserver: (@MainActor (NSMenu) -> Void)? + #endif var blinkTask: Task? var loginTask: Task? { didSet { self.refreshMenusForLoginStateChange() } @@ -135,6 +148,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin account: AccountInfo, updater: UpdaterProviding, preferencesSelection: PreferencesSelection, + managedCodexAccountCoordinator: ManagedCodexAccountCoordinator = ManagedCodexAccountCoordinator(), statusBar: NSStatusBar = .system) { if SettingsStore.isRunningTests { @@ -145,6 +159,7 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin self.account = account self.updater = updater self.preferencesSelection = preferencesSelection + self.managedCodexAccountCoordinator = managedCodexAccountCoordinator self.lastConfigRevision = settings.configRevision self.lastProviderOrder = settings.providerOrder self.lastMergeIcons = settings.mergeIcons @@ -178,11 +193,30 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin object: nil) } + convenience init( + store: UsageStore, + settings: SettingsStore, + account: AccountInfo, + updater: UpdaterProviding, + preferencesSelection: PreferencesSelection, + statusBar: NSStatusBar = .system) + { + self.init( + store: store, + settings: settings, + account: account, + updater: updater, + preferencesSelection: preferencesSelection, + managedCodexAccountCoordinator: ManagedCodexAccountCoordinator(), + statusBar: statusBar) + } + private func wireBindings() { self.observeStoreChanges() self.observeDebugForceAnimation() self.observeSettingsChanges() self.observeUpdaterChanges() + self.observeManagedCodexCoordinatorChanges() } private func observeStoreChanges() { @@ -254,6 +288,19 @@ final class StatusItemController: NSObject, NSMenuDelegate, StatusItemControllin } } + private func observeManagedCodexCoordinatorChanges() { + withObservationTracking { + _ = self.managedCodexAccountCoordinator.isAuthenticatingManagedAccount + _ = self.managedCodexAccountCoordinator.authenticatingManagedAccountID + } onChange: { [weak self] in + Task { @MainActor [weak self] in + guard let self else { return } + self.observeManagedCodexCoordinatorChanges() + self.refreshMenusForLoginStateChange() + } + } + } + private func invalidateMenus() { self.menuContentVersion &+= 1 // Don't refresh menus while they're open - wait until they close and reopen diff --git a/Sources/CodexBar/UsageStore+Accessors.swift b/Sources/CodexBar/UsageStore+Accessors.swift index 98d77c6d8..1a1f0650d 100644 --- a/Sources/CodexBar/UsageStore+Accessors.swift +++ b/Sources/CodexBar/UsageStore+Accessors.swift @@ -14,6 +14,18 @@ extension UsageStore { self.errors[.codex] } + var userFacingLastCodexError: String? { + self.userFacingError(for: .codex) + } + + var userFacingLastCreditsError: String? { + self.userFacingCodexUIError(self.lastCreditsError) + } + + var userFacingLastOpenAIDashboardError: String? { + self.userFacingCodexUIError(self.lastOpenAIDashboardError) + } + var lastClaudeError: String? { self.errors[.claude] } @@ -22,6 +34,12 @@ extension UsageStore { self.errors[provider] } + func userFacingError(for provider: UsageProvider) -> String? { + let raw = self.errors[provider] + guard provider == .codex else { return raw } + return self.userFacingCodexUIError(raw) + } + func status(for provider: UsageProvider) -> ProviderStatus? { guard self.statusChecksEnabled else { return nil } return self.statuses[provider] @@ -31,7 +49,99 @@ extension UsageStore { self.status(for: provider)?.indicator ?? .none } - func accountInfo() -> AccountInfo { - self.codexFetcher.loadAccountInfo() + func accountInfo(for provider: UsageProvider) -> AccountInfo { + guard provider == .codex else { + return self.codexFetcher.loadAccountInfo() + } + let env = ProviderRegistry.makeEnvironment( + base: ProcessInfo.processInfo.environment, + provider: .codex, + settings: self.settings, + tokenOverride: nil) + let fetcher = ProviderRegistry.makeFetcher(base: self.codexFetcher, provider: .codex, env: env) + return fetcher.loadAccountInfo() + } + + private func userFacingCodexUIError(_ raw: String?) -> String? { + guard let raw, !raw.isEmpty else { return nil } + let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return nil } + + let lower = trimmed.lowercased() + if self.codexErrorIsAlreadyUserFacing(lower: lower) { + return trimmed + } + + if let cachedMessage = self.userFacingCachedCodexError(trimmed, lower: lower) { + return cachedMessage + } + + if self.codexErrorLooksExpired(lower: lower) { + return "Codex session expired. Sign in again." + } + + if lower.contains("frame load interrupted") { + return "OpenAI web refresh was interrupted. Refresh OpenAI cookies and try again." + } + + if self.codexErrorLooksInternalTransport(lower: lower) { + return "Codex usage is temporarily unavailable. Try refreshing." + } + + return trimmed + } + + private func userFacingCachedCodexError(_ raw: String, lower: String) -> String? { + let cachedMarker = " Cached values from " + guard let suffixRange = raw.range(of: cachedMarker) else { return nil } + + let suffix = String(raw[suffixRange.lowerBound...]).trimmingCharacters(in: .whitespacesAndNewlines) + if lower.hasPrefix("last codex credits refresh failed:"), + let base = self.userFacingCodexUIError(String(raw[.. Bool { + lower.contains("openai cookies are for") + || lower.contains("sign in to chatgpt.com") + || lower.contains("requires a signed-in chatgpt.com session") + || lower.contains("managed codex account data is unavailable") + || lower.contains("selected managed codex account is unavailable") + || lower.contains("codex credits are still loading") + || lower.contains("codex account changed; importing browser cookies") + || lower.contains("codex session expired. sign in again.") + || lower.contains("codex usage is temporarily unavailable. try refreshing.") + } + + private func codexErrorLooksExpired(lower: String) -> Bool { + lower.contains("token_expired") + || lower.contains("authentication token is expired") + || lower.contains("oauth token has expired") + || lower.contains("provided authentication token is expired") + || lower.contains("please try signing in again") + || lower.contains("please sign in again") + || (lower.contains("401") && lower.contains("unauthorized")) + } + + private func codexErrorLooksInternalTransport(lower: String) -> Bool { + lower.contains("codex connection failed") + || lower.contains("failed to fetch codex rate limits") + || lower.contains("/backend-api/") + || lower.contains("content-type=") + || lower.contains("body={") + || lower.contains("body=") + || lower.contains("get https://") + || lower.contains("get http://") + || lower.contains("returned invalid data") } } diff --git a/Sources/CodexBar/UsageStore+OpenAIWeb.swift b/Sources/CodexBar/UsageStore+OpenAIWeb.swift index af164cc62..238bcdfa9 100644 --- a/Sources/CodexBar/UsageStore+OpenAIWeb.swift +++ b/Sources/CodexBar/UsageStore+OpenAIWeb.swift @@ -1,5 +1,960 @@ +import CodexBarCore import Foundation +// MARK: - OpenAI web lifecycle + +extension UsageStore { + private struct OpenAIDashboardRefreshContext { + let targetEmail: String? + let allowCurrentSnapshotFallback: Bool + let expectedGuard: CodexAccountScopedRefreshGuard? + let refreshTaskToken: UUID + let allowCodexUsageBackfill: Bool + } + + private static let openAIWebRefreshMultiplier: TimeInterval = 5 + private static let openAIWebPrimaryFetchTimeout: TimeInterval = 15 + private static let openAIWebRetryFetchTimeout: TimeInterval = 8 + + private func openAIWebRefreshIntervalSeconds() -> TimeInterval { + let base = max(self.settings.refreshFrequency.seconds ?? 0, 120) + return base * Self.openAIWebRefreshMultiplier + } + + func requestOpenAIDashboardRefreshIfStale(reason: String) { + guard self.isEnabled(.codex), self.settings.codexCookieSource.isEnabled else { return } + let now = Date() + let refreshInterval = self.openAIWebRefreshIntervalSeconds() + let lastUpdatedAt = self.openAIDashboard?.updatedAt ?? self.lastOpenAIDashboardSnapshot?.updatedAt + if let lastUpdatedAt, now.timeIntervalSince(lastUpdatedAt) < refreshInterval { return } + let stamp = now.formatted(date: .abbreviated, time: .shortened) + self.logOpenAIWeb("[\(stamp)] OpenAI web refresh request: \(reason)") + let expectedGuard = self.currentCodexOpenAIWebRefreshGuard() + Task { await self.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) } + } + + func applyOpenAIDashboard( + _ dash: OpenAIDashboardSnapshot, + targetEmail: String?, + expectedGuard: CodexAccountScopedRefreshGuard? = nil, + refreshTaskToken: UUID? = nil, + allowCodexUsageBackfill: Bool = true) async + { + guard self.shouldApplyOpenAIDashboardRefreshTask(token: refreshTaskToken) else { return } + let resolvedAccountEmail = targetEmail ?? dash.signedInEmail + let resolvedAccountKey = Self.normalizeCodexAccountScopedKey(resolvedAccountEmail) + if let expectedGuard, + !self.shouldApplyOpenAIDashboardResult( + expectedGuard: expectedGuard, + dashboardAccountEmail: resolvedAccountEmail) + { + return + } + + await MainActor.run { + self.openAIDashboard = dash + self.lastOpenAIDashboardError = nil + self.lastOpenAIDashboardSnapshot = dash + self.openAIDashboardRequiresLogin = false + // Only fill gaps; OAuth/CLI remain the primary sources for usage + credits. + if allowCodexUsageBackfill, + self.snapshots[.codex] == nil, + let usage = dash.toUsageSnapshot(provider: .codex, accountEmail: targetEmail) + { + self.snapshots[.codex] = usage + self.errors[.codex] = nil + self.failureGates[.codex]?.recordSuccess() + self.lastSourceLabels[.codex] = "openai-web" + self.rememberLiveSystemCodexEmailIfNeeded(usage.accountEmail(for: .codex)) + } + if self.credits == nil, let credits = dash.toCreditsSnapshot() { + self.credits = credits + self.lastCreditsSnapshot = credits + self.lastCreditsSnapshotAccountKey = resolvedAccountKey + self.lastCreditsError = nil + self.creditsFailureStreak = 0 + } + self.seedCodexAccountScopedRefreshGuard(accountEmail: resolvedAccountEmail) + } + + if let email = targetEmail, !email.isEmpty { + OpenAIDashboardCacheStore.save(OpenAIDashboardCache(accountEmail: email, snapshot: dash)) + } + self.backfillCodexHistoricalFromDashboardIfNeeded(dash) + } + + func applyOpenAIDashboardFailure( + message: String, + expectedGuard: CodexAccountScopedRefreshGuard? = nil, + refreshTaskToken: UUID? = nil) async + { + guard self.shouldApplyOpenAIDashboardRefreshTask(token: refreshTaskToken) else { return } + if let expectedGuard, + !self.shouldApplyOpenAIWebNonSuccessResult(expectedGuard: expectedGuard) + { + return + } + if self.openAIWebManagedTargetStoreIsUnreadable() { + await self.failClosedRefreshForUnreadableManagedCodexStore() + return + } + if self.openAIWebManagedTargetIsMissing() { + await self.failClosedRefreshForMissingManagedCodexTarget() + return + } + + await MainActor.run { + if let cached = self.lastOpenAIDashboardSnapshot { + self.openAIDashboard = cached + let stamp = cached.updatedAt.formatted(date: .abbreviated, time: .shortened) + self.lastOpenAIDashboardError = + "Last OpenAI dashboard refresh failed: \(message). Cached values from \(stamp)." + } else { + self.lastOpenAIDashboardError = message + self.openAIDashboard = nil + } + } + } + + func applyOpenAIDashboardMismatchFailure( + signedInEmail: String, + expectedEmail: String?, + expectedGuard: CodexAccountScopedRefreshGuard? = nil, + refreshTaskToken: UUID? = nil) async + { + guard self.shouldApplyOpenAIDashboardRefreshTask(token: refreshTaskToken) else { return } + if let expectedGuard, + !self.shouldApplyOpenAIWebNonSuccessResult(expectedGuard: expectedGuard) + { + return + } + await MainActor.run { + self.failClosedOpenAIDashboardSnapshot() + self.lastOpenAIDashboardError = [ + "OpenAI dashboard signed in as \(signedInEmail), but Codex uses \(expectedEmail ?? "unknown").", + "Switch accounts in your browser and update OpenAI cookies in Providers → Codex.", + ].joined(separator: " ") + } + } + + func applyOpenAIDashboardLoginRequiredFailure( + expectedGuard: CodexAccountScopedRefreshGuard? = nil, + refreshTaskToken: UUID? = nil) async + { + guard self.shouldApplyOpenAIDashboardRefreshTask(token: refreshTaskToken) else { return } + if let expectedGuard, + !self.shouldApplyOpenAIWebNonSuccessResult(expectedGuard: expectedGuard) + { + return + } + if self.openAIWebManagedTargetStoreIsUnreadable() { + await self.failClosedRefreshForUnreadableManagedCodexStore() + return + } + if self.openAIWebManagedTargetIsMissing() { + await self.failClosedRefreshForMissingManagedCodexTarget() + return + } + + await MainActor.run { + self.lastOpenAIDashboardError = [ + "OpenAI web access requires a signed-in chatgpt.com session.", + "Sign in using \(self.codexBrowserCookieOrder.loginHint), " + + "then update OpenAI cookies in Providers → Codex.", + ].joined(separator: " ") + self.openAIDashboard = self.lastOpenAIDashboardSnapshot + self.openAIDashboardRequiresLogin = true + } + } + + private func failClosedOpenAIDashboardSnapshot() { + self.openAIDashboard = nil + self.lastOpenAIDashboardSnapshot = nil + self.openAIDashboardRequiresLogin = true + } + + func refreshOpenAIDashboardIfNeeded( + force: Bool = false, + expectedGuard: CodexAccountScopedRefreshGuard? = nil, + bypassCoalescing: Bool = false, + allowCodexUsageBackfill: Bool = true) async + { + guard self.isEnabled(.codex), self.settings.codexCookieSource.isEnabled else { + self.resetOpenAIWebState() + return + } + if self.openAIWebManagedTargetStoreIsUnreadable() { + await self.failClosedRefreshForUnreadableManagedCodexStore() + return + } + if self.openAIWebManagedTargetIsMissing() { + await self.failClosedRefreshForMissingManagedCodexTarget() + return + } + + let allowCurrentSnapshotFallback = expectedGuard?.source == .liveSystem && expectedGuard?.accountKey == nil + let targetEmail = self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: allowCurrentSnapshotFallback, + allowLastKnownLiveFallback: expectedGuard?.accountKey != nil) + let refreshKey = self.openAIDashboardRefreshKey(targetEmail: targetEmail, expectedGuard: expectedGuard) + if !bypassCoalescing, + let task = self.openAIDashboardRefreshTask, + self.openAIDashboardRefreshTaskKey == refreshKey + { + await task.value + return + } + self.handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: targetEmail) + + let now = Date() + let minInterval = self.openAIWebRefreshIntervalSeconds() + if !force, + !self.openAIWebAccountDidChange, + self.lastOpenAIDashboardError == nil, + let snapshot = self.lastOpenAIDashboardSnapshot, + now.timeIntervalSince(snapshot.updatedAt) < minInterval + { + return + } + + let taskToken = UUID() + let context = OpenAIDashboardRefreshContext( + targetEmail: targetEmail, + allowCurrentSnapshotFallback: allowCurrentSnapshotFallback, + expectedGuard: expectedGuard, + refreshTaskToken: taskToken, + allowCodexUsageBackfill: allowCodexUsageBackfill) + let task = Task { [weak self] in + guard let self else { return } + await self.performOpenAIDashboardRefreshIfNeeded(context) + } + self.openAIDashboardRefreshTask = task + self.openAIDashboardRefreshTaskKey = refreshKey + self.openAIDashboardRefreshTaskToken = taskToken + await task.value + if self.openAIDashboardRefreshTaskToken == taskToken { + self.openAIDashboardRefreshTask = nil + self.openAIDashboardRefreshTaskKey = nil + self.openAIDashboardRefreshTaskToken = nil + } + } + + private func performOpenAIDashboardRefreshIfNeeded(_ context: OpenAIDashboardRefreshContext) async { + self.openAIDashboardCookieImportStatus = nil + var latestCookieImportStatus: String? + if self.openAIWebDebugLines.isEmpty { + self.resetOpenAIWebDebugLog(context: "refresh") + } else { + let stamp = Date().formatted(date: .abbreviated, time: .shortened) + self.logOpenAIWeb("[\(stamp)] OpenAI web refresh start") + } + let log: (String) -> Void = { [weak self] line in + guard let self else { return } + self.logOpenAIWeb(line) + } + + do { + let normalized = context.targetEmail? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + var effectiveEmail = context.targetEmail + + // Use a per-email persistent `WKWebsiteDataStore` so multiple dashboard sessions can coexist. + // Strategy: + // - Try the existing per-email WebKit cookie store first (fast; avoids Keychain prompts). + // - On login-required or account mismatch, import cookies from the configured browser order and retry once. + if self.openAIWebAccountDidChange, let targetEmail = context.targetEmail, !targetEmail.isEmpty { + // On account switches, proactively re-import cookies so we don't show stale data from the previous + // user. + let imported = await self.importOpenAIDashboardCookiesIfNeeded( + targetEmail: targetEmail, + force: true) + latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() + if await self.abortOpenAIDashboardRetryAfterImportFailure( + importedEmail: imported, + targetEmail: targetEmail, + expectedGuard: context.expectedGuard, + cookieImportStatus: latestCookieImportStatus, + refreshTaskToken: context.refreshTaskToken) + { + self.openAIWebAccountDidChange = false + return + } + if let imported { + effectiveEmail = imported + } + self.openAIWebAccountDidChange = false + } + + var dash = try await self.loadLatestOpenAIDashboard( + accountEmail: effectiveEmail, + logger: log, + timeout: Self.openAIWebPrimaryFetchTimeout) + + if self.dashboardEmailMismatch(expected: normalized, actual: dash.signedInEmail) { + if let imported = await self.importOpenAIDashboardCookiesIfNeeded( + targetEmail: context.targetEmail, + force: true) + { + effectiveEmail = imported + } + latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() + dash = try await self.loadLatestOpenAIDashboard( + accountEmail: effectiveEmail, + logger: log, + timeout: Self.openAIWebRetryFetchTimeout) + } + + if self.dashboardEmailMismatch(expected: normalized, actual: dash.signedInEmail) { + let signedIn = dash.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "unknown" + await self.applyOpenAIDashboardMismatchFailure( + signedInEmail: signedIn, + expectedEmail: normalized, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken) + return + } + + await self.applyOpenAIDashboard( + dash, + targetEmail: effectiveEmail, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken, + allowCodexUsageBackfill: context.allowCodexUsageBackfill) + } catch let OpenAIDashboardFetcher.FetchError.noDashboardData(body) { + await self.retryOpenAIDashboardAfterNoData( + body: body, + context: context, + latestCookieImportStatus: &latestCookieImportStatus, + logger: log) + } catch OpenAIDashboardFetcher.FetchError.loginRequired { + await self.retryOpenAIDashboardAfterLoginRequired( + context: context, + latestCookieImportStatus: &latestCookieImportStatus, + logger: log) + } catch { + let message = self.preferredOpenAIDashboardFailureMessage( + error: error, + targetEmail: context.targetEmail, + cookieImportStatus: latestCookieImportStatus) + await self.applyOpenAIDashboardFailure( + message: message, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken) + } + } + + private func retryOpenAIDashboardAfterNoData( + body: String, + context: OpenAIDashboardRefreshContext, + latestCookieImportStatus: inout String?, + logger: @escaping (String) -> Void) async + { + let targetEmail = self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: context.allowCurrentSnapshotFallback, + allowLastKnownLiveFallback: context.expectedGuard?.accountKey != nil) + var effectiveEmail = targetEmail + let imported = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) + latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() + if await self.abortOpenAIDashboardRetryAfterImportFailure( + importedEmail: imported, + targetEmail: targetEmail, + expectedGuard: context.expectedGuard, + cookieImportStatus: latestCookieImportStatus, + refreshTaskToken: context.refreshTaskToken) + { + return + } + if let imported { + effectiveEmail = imported + } + do { + let dash = try await self.loadLatestOpenAIDashboard( + accountEmail: effectiveEmail, + logger: logger, + timeout: Self.openAIWebRetryFetchTimeout) + await self.applyOpenAIDashboard( + dash, + targetEmail: effectiveEmail, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken, + allowCodexUsageBackfill: context.allowCodexUsageBackfill) + } catch let OpenAIDashboardFetcher.FetchError.noDashboardData(retryBody) { + let finalBody = retryBody.isEmpty ? body : retryBody + let message = self.openAIDashboardFriendlyError( + body: finalBody, + targetEmail: targetEmail, + cookieImportStatus: latestCookieImportStatus) + ?? OpenAIDashboardFetcher.FetchError.noDashboardData(body: finalBody).localizedDescription + await self.applyOpenAIDashboardFailure( + message: message, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken) + } catch { + let message = self.preferredOpenAIDashboardFailureMessage( + error: error, + targetEmail: targetEmail, + cookieImportStatus: latestCookieImportStatus) + await self.applyOpenAIDashboardFailure( + message: message, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken) + } + } + + private func retryOpenAIDashboardAfterLoginRequired( + context: OpenAIDashboardRefreshContext, + latestCookieImportStatus: inout String?, + logger: @escaping (String) -> Void) async + { + let targetEmail = self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: context.allowCurrentSnapshotFallback, + allowLastKnownLiveFallback: context.expectedGuard?.accountKey != nil) + var effectiveEmail = targetEmail + let imported = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) + latestCookieImportStatus = self.currentOpenAIDashboardCookieImportStatus() + if await self.abortOpenAIDashboardRetryAfterImportFailure( + importedEmail: imported, + targetEmail: targetEmail, + expectedGuard: context.expectedGuard, + cookieImportStatus: latestCookieImportStatus, + refreshTaskToken: context.refreshTaskToken) + { + return + } + if let imported { + effectiveEmail = imported + } + do { + let dash = try await self.loadLatestOpenAIDashboard( + accountEmail: effectiveEmail, + logger: logger, + timeout: Self.openAIWebRetryFetchTimeout) + await self.applyOpenAIDashboard( + dash, + targetEmail: effectiveEmail, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken, + allowCodexUsageBackfill: context.allowCodexUsageBackfill) + } catch OpenAIDashboardFetcher.FetchError.loginRequired { + await self.applyOpenAIDashboardLoginRequiredFailure( + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken) + } catch { + let message = self.preferredOpenAIDashboardFailureMessage( + error: error, + targetEmail: targetEmail, + cookieImportStatus: latestCookieImportStatus) + await self.applyOpenAIDashboardFailure( + message: message, + expectedGuard: context.expectedGuard, + refreshTaskToken: context.refreshTaskToken) + } + } + + // MARK: - OpenAI web account switching + + /// Detect Codex account email changes and clear stale OpenAI web state so the UI can't show the wrong user. + /// This does not delete other per-email WebKit cookie stores (we keep multiple accounts around). + func handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: String?) { + let normalized = targetEmail? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + + guard let normalized, !normalized.isEmpty else { return } + + let previous = self.lastOpenAIDashboardTargetEmail + self.lastOpenAIDashboardTargetEmail = normalized + + if let previous, + !previous.isEmpty, + previous != normalized + { + let stamp = Date().formatted(date: .abbreviated, time: .shortened) + self.logOpenAIWeb( + "[\(stamp)] Codex account changed: \(previous) → \(normalized); " + + "clearing OpenAI web snapshot") + self.openAIWebAccountDidChange = true + self.openAIDashboard = nil + self.lastOpenAIDashboardSnapshot = nil + self.lastOpenAIDashboardError = nil + self.openAIDashboardRequiresLogin = true + self.openAIDashboardCookieImportStatus = "Codex account changed; importing browser cookies…" + self.lastOpenAIDashboardCookieImportAttemptAt = nil + self.lastOpenAIDashboardCookieImportEmail = nil + } + } + + func importOpenAIDashboardBrowserCookiesNow() async { + self.resetOpenAIWebDebugLog(context: "manual import") + let targetEmail = self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: true, + allowLastKnownLiveFallback: false) + _ = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) + let expectedGuard = self.currentCodexOpenAIWebRefreshGuard() + await self.refreshOpenAIDashboardIfNeeded( + force: true, + expectedGuard: expectedGuard, + bypassCoalescing: true) + } + + func currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: Bool, + allowLastKnownLiveFallback: Bool) -> String? + { + switch self.settings.codexResolvedActiveSource { + case .liveSystem: + let liveSystem = self.settings.codexAccountReconciliationSnapshot.liveSystemAccount?.email + .trimmingCharacters(in: .whitespacesAndNewlines) + if let liveSystem, !liveSystem.isEmpty { + self.lastKnownLiveSystemCodexEmail = liveSystem + return liveSystem + } + + if allowCurrentSnapshotFallback, + let snapshotEmail = self.snapshots[.codex]?.accountEmail(for: .codex)? + .trimmingCharacters(in: .whitespacesAndNewlines), + !snapshotEmail.isEmpty + { + self.lastKnownLiveSystemCodexEmail = snapshotEmail + return snapshotEmail + } + + if allowLastKnownLiveFallback { + let lastKnown = self.lastKnownLiveSystemCodexEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + if let lastKnown, !lastKnown.isEmpty { return lastKnown } + } + return nil + case .managedAccount: + return self.codexAccountEmailForOpenAIDashboard() + } + } + + private func shouldApplyOpenAIWebNonSuccessResult(expectedGuard: CodexAccountScopedRefreshGuard) -> Bool { + if expectedGuard.accountKey != nil { + return self.shouldApplyCodexScopedNonUsageResult(expectedGuard: expectedGuard) + } + + guard case .liveSystem = expectedGuard.source else { return false } + let currentGuard = self.currentCodexOpenAIWebRefreshGuard() + guard currentGuard.source == expectedGuard.source else { return false } + guard currentGuard.accountKey == nil else { return false } + return self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: true, + allowLastKnownLiveFallback: false) != nil + } + + private func openAIDashboardRefreshKey( + targetEmail: String?, + expectedGuard: CodexAccountScopedRefreshGuard?) -> String + { + let source = String(describing: expectedGuard?.source ?? self.settings.codexResolvedActiveSource) + let accountKey = Self.normalizeCodexAccountScopedKey(targetEmail ?? expectedGuard?.accountKey) ?? "unknown" + return "\(source)|\(accountKey)" + } + + private func actionableOpenAIDashboardImportFailure(targetEmail: String?) -> String? { + self.actionableOpenAIDashboardImportFailure( + targetEmail: targetEmail, + cookieImportStatus: self.openAIDashboardCookieImportStatus) + } + + private func actionableOpenAIDashboardImportFailure( + targetEmail: String?, + cookieImportStatus: String?) -> String? + { + let status = cookieImportStatus?.trimmingCharacters(in: .whitespacesAndNewlines) + guard let status, !status.isEmpty else { return nil } + + if status.localizedCaseInsensitiveContains("openai cookies are for") { + return "\(status) Switch chatgpt.com account, then refresh OpenAI cookies." + } + if status.localizedCaseInsensitiveContains("no signed-in openai web session found") { + let targetLabel = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + let accountLabel = (targetLabel?.isEmpty == false) ? targetLabel! : "your OpenAI account" + return "\(status) Sign in to chatgpt.com as \(accountLabel), then refresh OpenAI cookies." + } + if status.localizedCaseInsensitiveContains("openai cookie import failed") + || status.localizedCaseInsensitiveContains("browser cookie import failed") + { + return status + } + return nil + } + + private func preferredOpenAIDashboardFailureMessage( + error: Error, + targetEmail: String?, + cookieImportStatus: String?) -> String + { + if let actionable = self.actionableOpenAIDashboardImportFailure( + targetEmail: targetEmail, + cookieImportStatus: cookieImportStatus) + { + return actionable + } + return error.localizedDescription + } + + private func abortOpenAIDashboardRetryAfterImportFailure( + importedEmail: String?, + targetEmail: String?, + expectedGuard: CodexAccountScopedRefreshGuard?, + cookieImportStatus: String?, + refreshTaskToken: UUID) async -> Bool + { + guard importedEmail == nil, + let message = self.actionableOpenAIDashboardImportFailure( + targetEmail: targetEmail, + cookieImportStatus: cookieImportStatus) + else { + return false + } + await self.applyOpenAIDashboardFailure( + message: message, + expectedGuard: expectedGuard, + refreshTaskToken: refreshTaskToken) + return true + } + + private func shouldApplyOpenAIDashboardRefreshTask(token: UUID?) -> Bool { + guard let token else { return true } + return self.openAIDashboardRefreshTaskToken == token + } + + func invalidateOpenAIDashboardRefreshTask() { + self.openAIDashboardRefreshTask?.cancel() + self.openAIDashboardRefreshTask = nil + self.openAIDashboardRefreshTaskKey = nil + self.openAIDashboardRefreshTaskToken = nil + } + + private func currentOpenAIDashboardCookieImportStatus() -> String? { + self.openAIDashboardCookieImportStatus + } + + private func loadLatestOpenAIDashboard( + accountEmail: String?, + logger: @escaping (String) -> Void, + timeout: TimeInterval) async throws -> OpenAIDashboardSnapshot + { + if let override = self._test_openAIDashboardLoaderOverride { + return try await override(accountEmail, logger, timeout) + } + return try await OpenAIDashboardFetcher().loadLatestDashboard( + accountEmail: accountEmail, + logger: logger, + debugDumpHTML: timeout != Self.openAIWebPrimaryFetchTimeout, + timeout: timeout) + } + + private func failClosedForUnreadableManagedCodexStore() async -> String? { + await MainActor.run { + self.failClosedOpenAIDashboardSnapshot() + self.openAIDashboardCookieImportStatus = [ + "Managed Codex account data is unavailable.", + "Fix the managed account store before importing OpenAI cookies.", + ].joined(separator: " ") + } + return nil + } + + private func failClosedRefreshForUnreadableManagedCodexStore() async { + await MainActor.run { + self.failClosedOpenAIDashboardSnapshot() + self.lastOpenAIDashboardError = [ + "Managed Codex account data is unavailable.", + "Fix the managed account store before refreshing OpenAI web data.", + ].joined(separator: " ") + } + } + + private func failClosedForMissingManagedCodexTarget() async -> String? { + await MainActor.run { + self.failClosedOpenAIDashboardSnapshot() + self.openAIDashboardCookieImportStatus = [ + "The selected managed Codex account is unavailable.", + "Pick another Codex account before importing OpenAI cookies.", + ].joined(separator: " ") + } + return nil + } + + private func failClosedRefreshForMissingManagedCodexTarget() async { + await MainActor.run { + self.failClosedOpenAIDashboardSnapshot() + self.lastOpenAIDashboardError = [ + "The selected managed Codex account is unavailable.", + "Pick another Codex account before refreshing OpenAI web data.", + ].joined(separator: " ") + } + } + + private func openAIWebCookieImportShouldFailClosed() async -> Bool { + if self.openAIWebManagedTargetStoreIsUnreadable() { + _ = await self.failClosedForUnreadableManagedCodexStore() + return true + } + if self.openAIWebManagedTargetIsMissing() { + _ = await self.failClosedForMissingManagedCodexTarget() + return true + } + return false + } + + func importOpenAIDashboardCookiesIfNeeded(targetEmail: String?, force: Bool) async -> String? { + if await self.openAIWebCookieImportShouldFailClosed() { + return nil + } + + let normalizedTarget = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + let allowAnyAccount = normalizedTarget == nil || normalizedTarget?.isEmpty == true + let cookieSource = self.settings.codexCookieSource + let cacheScope = self.codexCookieCacheScopeForOpenAIWeb() + + let now = Date() + let lastEmail = self.lastOpenAIDashboardCookieImportEmail + let lastAttempt = self.lastOpenAIDashboardCookieImportAttemptAt ?? .distantPast + + let shouldAttempt: Bool = if force { + true + } else { + if allowAnyAccount { + now.timeIntervalSince(lastAttempt) > 300 + } else { + self.openAIDashboardRequiresLogin && + ( + lastEmail?.lowercased() != normalizedTarget?.lowercased() || now + .timeIntervalSince(lastAttempt) > 300) + } + } + + guard shouldAttempt else { return normalizedTarget } + self.lastOpenAIDashboardCookieImportEmail = normalizedTarget + self.lastOpenAIDashboardCookieImportAttemptAt = now + + let stamp = now.formatted(date: .abbreviated, time: .shortened) + let targetLabel = normalizedTarget ?? "unknown" + self.logOpenAIWeb("[\(stamp)] import start (target=\(targetLabel))") + + do { + let log: (String) -> Void = { [weak self] message in + guard let self else { return } + self.logOpenAIWeb(message) + } + + let result: OpenAIDashboardBrowserCookieImporter.ImportResult + if let override = self._test_openAIDashboardCookieImportOverride { + result = try await override(normalizedTarget, allowAnyAccount, cookieSource, cacheScope, log) + } else { + let importer = OpenAIDashboardBrowserCookieImporter(browserDetection: self.browserDetection) + switch cookieSource { + case .manual: + self.settings.ensureCodexCookieLoaded() + // Manual OpenAI cookies still come from one provider-level setting. Auto-imported cookies are + // isolated per managed account, but a manual header is an explicit override owned by settings, + // so switching managed accounts does not currently swap it underneath the user. + let manualHeader = self.settings.codexCookieHeader + guard CookieHeaderNormalizer.normalize(manualHeader) != nil else { + throw OpenAIDashboardBrowserCookieImporter.ImportError.manualCookieHeaderInvalid + } + result = try await importer.importManualCookies( + cookieHeader: manualHeader, + intoAccountEmail: normalizedTarget, + allowAnyAccount: allowAnyAccount, + cacheScope: cacheScope, + logger: log) + case .auto: + result = try await importer.importBestCookies( + intoAccountEmail: normalizedTarget, + allowAnyAccount: allowAnyAccount, + cacheScope: cacheScope, + logger: log) + case .off: + result = OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Off", + cookieCount: 0, + signedInEmail: normalizedTarget, + matchesCodexEmail: true) + } + } + let effectiveEmail = result.signedInEmail? + .trimmingCharacters(in: .whitespacesAndNewlines) + .isEmpty == false + ? result.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + : normalizedTarget + self.lastOpenAIDashboardCookieImportEmail = effectiveEmail ?? normalizedTarget + await MainActor.run { + let signed = result.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + let matchText = result.matchesCodexEmail ? "matches Codex" : "does not match Codex" + let sourceLabel = switch cookieSource { + case .manual: + "Manual cookie header" + case .auto: + "\(result.sourceLabel) cookies" + case .off: + "OpenAI cookies disabled" + } + if let signed, !signed.isEmpty { + self.openAIDashboardCookieImportStatus = + allowAnyAccount + ? [ + "Using \(sourceLabel) (\(result.cookieCount)).", + "Signed in as \(signed).", + ].joined(separator: " ") + : [ + "Using \(sourceLabel) (\(result.cookieCount)).", + "Signed in as \(signed) (\(matchText)).", + ].joined(separator: " ") + } else { + self.openAIDashboardCookieImportStatus = + "Using \(sourceLabel) (\(result.cookieCount))." + } + } + return effectiveEmail + } catch let err as OpenAIDashboardBrowserCookieImporter.ImportError { + switch err { + case let .noMatchingAccount(found): + let foundText: String = if found.isEmpty { + "no signed-in session detected in \(self.codexBrowserCookieOrder.loginHint)" + } else { + found + .sorted { lhs, rhs in + if lhs.sourceLabel == rhs.sourceLabel { return lhs.email < rhs.email } + return lhs.sourceLabel < rhs.sourceLabel + } + .map { "\($0.sourceLabel): \($0.email)" } + .joined(separator: " • ") + } + self.logOpenAIWeb("[\(stamp)] import mismatch: \(foundText)") + await MainActor.run { + self.openAIDashboardCookieImportStatus = allowAnyAccount + ? [ + "No signed-in OpenAI web session found.", + "Found \(foundText).", + ].joined(separator: " ") + : Self.conciseOpenAICookieMismatchStatus( + found: found.map(\.email), + targetEmail: normalizedTarget) + self.failClosedOpenAIDashboardSnapshot() + } + case .noCookiesFound, + .browserAccessDenied, + .dashboardStillRequiresLogin, + .manualCookieHeaderInvalid: + self.logOpenAIWeb("[\(stamp)] import failed: \(err.localizedDescription)") + await MainActor.run { + self.openAIDashboardCookieImportStatus = + "OpenAI cookie import failed: \(err.localizedDescription)" + self.openAIDashboardRequiresLogin = true + } + } + } catch { + self.logOpenAIWeb("[\(stamp)] import failed: \(error.localizedDescription)") + await MainActor.run { + self.openAIDashboardCookieImportStatus = + "Browser cookie import failed: \(error.localizedDescription)" + } + } + return nil + } + + private func resetOpenAIWebDebugLog(context: String) { + let stamp = Date().formatted(date: .abbreviated, time: .shortened) + self.openAIWebDebugLines.removeAll(keepingCapacity: true) + self.openAIDashboardCookieImportDebugLog = nil + self.logOpenAIWeb("[\(stamp)] OpenAI web \(context) start") + } + + private func logOpenAIWeb(_ message: String) { + let safeMessage = LogRedactor.redact(message) + self.openAIWebLogger.debug(safeMessage) + self.openAIWebDebugLines.append(safeMessage) + if self.openAIWebDebugLines.count > 240 { + self.openAIWebDebugLines.removeFirst(self.openAIWebDebugLines.count - 240) + } + self.openAIDashboardCookieImportDebugLog = self.openAIWebDebugLines.joined(separator: "\n") + } + + func resetOpenAIWebState() { + self.invalidateOpenAIDashboardRefreshTask() + self.openAIDashboard = nil + self.lastOpenAIDashboardError = nil + self.lastOpenAIDashboardSnapshot = nil + self.lastOpenAIDashboardTargetEmail = nil + self.openAIDashboardRequiresLogin = false + self.openAIDashboardCookieImportStatus = nil + self.openAIDashboardCookieImportDebugLog = nil + self.lastOpenAIDashboardCookieImportAttemptAt = nil + self.lastOpenAIDashboardCookieImportEmail = nil + self.lastKnownLiveSystemCodexEmail = nil + } + + private func dashboardEmailMismatch(expected: String?, actual: String?) -> Bool { + guard let expected, !expected.isEmpty else { return false } + guard let raw = actual?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return false } + return raw.lowercased() != expected.lowercased() + } + + private func openAIWebManagedTargetStoreIsUnreadable() -> Bool { + guard case .managedAccount = self.settings.codexResolvedActiveSource else { + return false + } + return self.settings.codexSettingsSnapshot(tokenOverride: nil).managedAccountStoreUnreadable + } + + private func openAIWebManagedTargetIsMissing() -> Bool { + guard case .managedAccount = self.settings.codexResolvedActiveSource else { + return false + } + return self.selectedManagedCodexAccountForOpenAIWeb() == nil + } + + private func selectedManagedCodexAccountForOpenAIWeb() -> ManagedCodexAccount? { + guard case let .managedAccount(id) = self.settings.codexResolvedActiveSource else { + return nil + } + + let snapshot = self.settings.codexAccountReconciliationSnapshot + return snapshot.storedAccounts.first { $0.id == id } + } + + func codexAccountEmailForOpenAIDashboard(allowLastKnownLiveFallback: Bool = true) -> String? { + switch self.settings.codexResolvedActiveSource { + case .liveSystem: + let liveSystem = self.settings.codexAccountReconciliationSnapshot.liveSystemAccount?.email + .trimmingCharacters(in: .whitespacesAndNewlines) + if let liveSystem, !liveSystem.isEmpty { + self.lastKnownLiveSystemCodexEmail = liveSystem + return liveSystem + } + + guard allowLastKnownLiveFallback else { return nil } + let lastKnown = self.lastKnownLiveSystemCodexEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + if let lastKnown, !lastKnown.isEmpty { return lastKnown } + return nil + case .managedAccount: + if self.openAIWebManagedTargetStoreIsUnreadable() { + return nil + } + + let managed = self.selectedManagedCodexAccountForOpenAIWeb()?.email + .trimmingCharacters(in: .whitespacesAndNewlines) + if let managed, !managed.isEmpty { return managed } + return nil + } + } + + func codexCookieCacheScopeForOpenAIWeb() -> CookieHeaderCache.Scope? { + switch self.settings.codexResolvedActiveSource { + case .liveSystem: + nil + case let .managedAccount(id): + self.openAIWebManagedTargetStoreIsUnreadable() ? .managedStoreUnreadable : .managedAccount(id) + } + } +} + // MARK: - OpenAI web error messaging extension UsageStore { @@ -32,12 +987,10 @@ extension UsageStore { let targetLabel = (emailLabel?.isEmpty == false) ? emailLabel! : "your OpenAI account" if let status, !status.isEmpty { if status.contains("cookies do not match Codex account") + || status.localizedCaseInsensitiveContains("openai cookies are for") || status.localizedCaseInsensitiveContains("cookie import failed") { - return [ - status, - "Sign in to chatgpt.com as \(targetLabel), then update OpenAI cookies in Providers → Codex.", - ].joined(separator: " ") + return "\(status) Switch chatgpt.com account, then refresh OpenAI cookies." } } return [ @@ -45,4 +998,33 @@ extension UsageStore { "Sign in to chatgpt.com as \(targetLabel), then update OpenAI cookies in Providers → Codex.", ].joined(separator: " ") } + + private static func conciseOpenAICookieMismatchStatus( + found: [String], + targetEmail: String?) + -> String + { + let normalizedFound = Array(Set( + found + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() } + .filter { !$0.isEmpty })) + .sorted() + + let foundLabel: String = switch normalizedFound.count { + case 0: + "another account" + case 1: + normalizedFound[0] + case 2: + "\(normalizedFound[0]) or \(normalizedFound[1])" + default: + "\(normalizedFound[0]) or \(normalizedFound.count - 1) other accounts" + } + + let targetLabel = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + guard let targetLabel, !targetLabel.isEmpty else { + return "OpenAI cookies are for \(foundLabel)." + } + return "OpenAI cookies are for \(foundLabel), not \(targetLabel)." + } } diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index 06fccc774..c292f210a 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -2,13 +2,20 @@ import CodexBarCore import Foundation extension UsageStore { + func prepareRefreshState(for provider: UsageProvider? = nil) { + guard provider == nil || provider == .codex else { return } + _ = self.settings.persistResolvedCodexActiveSourceCorrectionIfNeeded() + } + /// Force refresh Augment session (called from UI button) func forceRefreshAugmentSession() async { await self.performRuntimeAction(.forceSessionRefresh, for: .augment) } func refreshProvider(_ provider: UsageProvider, allowDisabled: Bool = false) async { + self.prepareRefreshState(for: provider) guard let spec = self.providerSpecs[provider] else { return } + let codexExpectedGuard = provider == .codex ? self.currentCodexAccountScopedRefreshGuard() : nil if !spec.isEnabled(), !allowDisabled { self.refreshingProviders.remove(provider) @@ -78,12 +85,22 @@ extension UsageStore { switch outcome.result { case let .success(result): let scoped = result.usage.scoped(to: provider) + if provider == .codex, + let codexExpectedGuard, + !self.shouldApplyCodexUsageResult(expectedGuard: codexExpectedGuard, usage: scoped) + { + return + } await MainActor.run { self.handleSessionQuotaTransition(provider: provider, snapshot: scoped) self.snapshots[provider] = scoped self.lastSourceLabels[provider] = result.sourceLabel self.errors[provider] = nil self.failureGates[provider]?.recordSuccess() + if provider == .codex { + self.rememberLiveSystemCodexEmailIfNeeded(scoped.accountEmail(for: .codex)) + self.seedCodexAccountScopedRefreshGuard(accountEmail: scoped.accountEmail(for: .codex)) + } } await self.recordPlanUtilizationHistorySample( provider: provider, @@ -97,6 +114,12 @@ extension UsageStore { self.recordCodexHistoricalSampleIfNeeded(snapshot: scoped) } case let .failure(error): + if provider == .codex, + let codexExpectedGuard, + !self.shouldApplyCodexScopedFailure(expectedGuard: codexExpectedGuard) + { + return + } await MainActor.run { let hadPriorData = self.snapshots[provider] != nil let shouldSurface = diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 80a41b3d9..4940eeec5 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -88,6 +88,14 @@ extension UsageStore { override: TokenAccountOverride?) async -> ProviderFetchOutcome { let descriptor = ProviderDescriptorRegistry.descriptor(for: provider) + let context = self.makeFetchContext(provider: provider, override: override) + return await descriptor.fetchOutcome(context: context) + } + + func makeFetchContext( + provider: UsageProvider, + override: TokenAccountOverride?) -> ProviderFetchContext + { let sourceMode = self.sourceMode(for: provider) let snapshot = ProviderRegistry.makeSettingsSnapshot(settings: self.settings, tokenOverride: override) let env = ProviderRegistry.makeEnvironment( @@ -95,8 +103,9 @@ extension UsageStore { provider: provider, settings: self.settings, tokenOverride: override) + let fetcher = ProviderRegistry.makeFetcher(base: self.codexFetcher, provider: provider, env: env) let verbose = self.settings.isVerboseLoggingEnabled - let context = ProviderFetchContext( + return ProviderFetchContext( runtime: .app, sourceMode: sourceMode, includeCredits: false, @@ -105,10 +114,9 @@ extension UsageStore { verbose: verbose, env: env, settings: snapshot, - fetcher: self.codexFetcher, + fetcher: fetcher, claudeFetcher: self.claudeFetcher, browserDetection: self.browserDetection) - return await descriptor.fetchOutcome(context: context) } func sourceMode(for provider: UsageProvider) -> ProviderSourceMode { diff --git a/Sources/CodexBar/UsageStore+WidgetSnapshot.swift b/Sources/CodexBar/UsageStore+WidgetSnapshot.swift index f8699128d..44258056e 100644 --- a/Sources/CodexBar/UsageStore+WidgetSnapshot.swift +++ b/Sources/CodexBar/UsageStore+WidgetSnapshot.swift @@ -7,12 +7,20 @@ import WidgetKit extension UsageStore { func persistWidgetSnapshot(reason: String) { let snapshot = self.makeWidgetSnapshot() - Task.detached(priority: .utility) { - WidgetSnapshotStore.save(snapshot) - #if canImport(WidgetKit) - await MainActor.run { - WidgetCenter.shared.reloadAllTimelines() + let previousTask = self.widgetSnapshotPersistTask + self.widgetSnapshotPersistTask = Task { @MainActor in + _ = await previousTask?.result + + if let override = self._test_widgetSnapshotSaveOverride { + await override(snapshot) + return } + + await Task.detached(priority: .utility) { + WidgetSnapshotStore.save(snapshot) + }.value + #if canImport(WidgetKit) + WidgetCenter.shared.reloadAllTimelines() #endif } } diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 47efc5b63..128891b51 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -119,12 +119,31 @@ final class UsageStore { var probeLogs: [UsageProvider: String] = [:] var historicalPaceRevision: Int = 0 @ObservationIgnored var lastCreditsSnapshot: CreditsSnapshot? + @ObservationIgnored var lastCreditsSnapshotAccountKey: String? @ObservationIgnored var creditsFailureStreak: Int = 0 - @ObservationIgnored private var lastOpenAIDashboardSnapshot: OpenAIDashboardSnapshot? - @ObservationIgnored private var lastOpenAIDashboardTargetEmail: String? - @ObservationIgnored private var lastOpenAIDashboardCookieImportAttemptAt: Date? - @ObservationIgnored private var lastOpenAIDashboardCookieImportEmail: String? - @ObservationIgnored private var openAIWebAccountDidChange: Bool = false + @ObservationIgnored var lastOpenAIDashboardSnapshot: OpenAIDashboardSnapshot? + @ObservationIgnored var lastOpenAIDashboardTargetEmail: String? + @ObservationIgnored var lastOpenAIDashboardCookieImportAttemptAt: Date? + @ObservationIgnored var lastOpenAIDashboardCookieImportEmail: String? + @ObservationIgnored var lastCodexAccountScopedRefreshGuard: CodexAccountScopedRefreshGuard? + @ObservationIgnored var lastKnownLiveSystemCodexEmail: String? + @ObservationIgnored var openAIWebAccountDidChange: Bool = false + @ObservationIgnored var openAIDashboardRefreshTask: Task? + @ObservationIgnored var openAIDashboardRefreshTaskKey: String? + @ObservationIgnored var openAIDashboardRefreshTaskToken: UUID? + @ObservationIgnored var _test_openAIDashboardCookieImportOverride: (@MainActor ( + String?, + Bool, + ProviderCookieSource, + CookieHeaderCache.Scope?, + @escaping (String) -> Void) async throws -> OpenAIDashboardBrowserCookieImporter.ImportResult)? + @ObservationIgnored var _test_openAIDashboardLoaderOverride: (@MainActor ( + String?, + @escaping (String) -> Void, + TimeInterval) async throws -> OpenAIDashboardSnapshot)? + @ObservationIgnored var _test_codexCreditsLoaderOverride: (@MainActor () async throws -> CreditsSnapshot)? + @ObservationIgnored var _test_widgetSnapshotSaveOverride: (@MainActor (WidgetSnapshot) async -> Void)? + @ObservationIgnored var widgetSnapshotPersistTask: Task? @ObservationIgnored let codexFetcher: UsageFetcher @ObservationIgnored let claudeFetcher: any ClaudeUsageFetching @@ -134,11 +153,11 @@ final class UsageStore { @ObservationIgnored let settings: SettingsStore @ObservationIgnored private let sessionQuotaNotifier: any SessionQuotaNotifying @ObservationIgnored private let sessionQuotaLogger = CodexBarLog.logger(LogCategories.sessionQuota) - @ObservationIgnored private let openAIWebLogger = CodexBarLog.logger(LogCategories.openAIWeb) + @ObservationIgnored let openAIWebLogger = CodexBarLog.logger(LogCategories.openAIWeb) @ObservationIgnored private let tokenCostLogger = CodexBarLog.logger(LogCategories.tokenCost) @ObservationIgnored let augmentLogger = CodexBarLog.logger(LogCategories.augment) @ObservationIgnored let providerLogger = CodexBarLog.logger(LogCategories.providers) - @ObservationIgnored private var openAIWebDebugLines: [String] = [] + @ObservationIgnored var openAIWebDebugLines: [String] = [] @ObservationIgnored var failureGates: [UsageProvider: ConsecutiveFailureGate] = [:] @ObservationIgnored var tokenFailureGates: [UsageProvider: ConsecutiveFailureGate] = [:] @ObservationIgnored var providerSpecs: [UsageProvider: ProviderSpec] = [:] @@ -305,7 +324,7 @@ final class UsageStore { self.providerMetadata[provider]! } - private var codexBrowserCookieOrder: BrowserCookieImportOrder { + var codexBrowserCookieOrder: BrowserCookieImportOrder { self.metadata(for: .codex).browserCookieOrder ?? Browser.defaultImportOrder } @@ -400,6 +419,7 @@ final class UsageStore { func refresh(forceTokenUsage: Bool = false) async { guard !self.isRefreshing else { return } + self.prepareRefreshState() let refreshPhase: ProviderRefreshPhase = self.hasCompletedInitialRefresh ? .regular : .startup let refreshStartedAt = Date() @@ -423,7 +443,10 @@ final class UsageStore { // OpenAI web scrape depends on the current Codex account email (which can change after login/account // switch). Run this after Codex usage refresh so we don't accidentally scrape with stale credentials. - await self.refreshOpenAIDashboardIfNeeded(force: forceTokenUsage) + let codexDashboardGuard = self.currentCodexOpenAIWebRefreshGuard() + await self.refreshOpenAIDashboardIfNeeded( + force: forceTokenUsage, + expectedGuard: codexDashboardGuard) if self.openAIDashboardRequiresLogin { await self.refreshProvider(.codex) @@ -630,449 +653,6 @@ final class UsageStore { } } -extension UsageStore { - private static let openAIWebRefreshMultiplier: TimeInterval = 5 - private static let openAIWebPrimaryFetchTimeout: TimeInterval = 15 - private static let openAIWebRetryFetchTimeout: TimeInterval = 8 - - private func openAIWebRefreshIntervalSeconds() -> TimeInterval { - let base = max(self.settings.refreshFrequency.seconds ?? 0, 120) - return base * Self.openAIWebRefreshMultiplier - } - - func requestOpenAIDashboardRefreshIfStale(reason: String) { - guard self.isEnabled(.codex), self.settings.codexCookieSource.isEnabled else { return } - let now = Date() - let refreshInterval = self.openAIWebRefreshIntervalSeconds() - let lastUpdatedAt = self.openAIDashboard?.updatedAt ?? self.lastOpenAIDashboardSnapshot?.updatedAt - if let lastUpdatedAt, now.timeIntervalSince(lastUpdatedAt) < refreshInterval { return } - let stamp = now.formatted(date: .abbreviated, time: .shortened) - self.logOpenAIWeb("[\(stamp)] OpenAI web refresh request: \(reason)") - Task { await self.refreshOpenAIDashboardIfNeeded(force: true) } - } - - private func applyOpenAIDashboard(_ dash: OpenAIDashboardSnapshot, targetEmail: String?) async { - await MainActor.run { - self.openAIDashboard = dash - self.lastOpenAIDashboardError = nil - self.lastOpenAIDashboardSnapshot = dash - self.openAIDashboardRequiresLogin = false - // Only fill gaps; OAuth/CLI remain the primary sources for usage + credits. - if self.snapshots[.codex] == nil, - let usage = dash.toUsageSnapshot(provider: .codex, accountEmail: targetEmail) - { - self.snapshots[.codex] = usage - self.errors[.codex] = nil - self.failureGates[.codex]?.recordSuccess() - self.lastSourceLabels[.codex] = "openai-web" - } - if self.credits == nil, let credits = dash.toCreditsSnapshot() { - self.credits = credits - self.lastCreditsSnapshot = credits - self.lastCreditsError = nil - self.creditsFailureStreak = 0 - } - } - - if let email = targetEmail, !email.isEmpty { - OpenAIDashboardCacheStore.save(OpenAIDashboardCache(accountEmail: email, snapshot: dash)) - } - self.backfillCodexHistoricalFromDashboardIfNeeded(dash) - } - - private func applyOpenAIDashboardFailure(message: String) async { - await MainActor.run { - if let cached = self.lastOpenAIDashboardSnapshot { - self.openAIDashboard = cached - let stamp = cached.updatedAt.formatted(date: .abbreviated, time: .shortened) - self.lastOpenAIDashboardError = - "Last OpenAI dashboard refresh failed: \(message). Cached values from \(stamp)." - } else { - self.lastOpenAIDashboardError = message - self.openAIDashboard = nil - } - } - } - - private func refreshOpenAIDashboardIfNeeded(force: Bool = false) async { - guard self.isEnabled(.codex), self.settings.codexCookieSource.isEnabled else { - self.resetOpenAIWebState() - return - } - - let targetEmail = self.codexAccountEmailForOpenAIDashboard() - self.handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: targetEmail) - - let now = Date() - let minInterval = self.openAIWebRefreshIntervalSeconds() - if !force, - !self.openAIWebAccountDidChange, - self.lastOpenAIDashboardError == nil, - let snapshot = self.lastOpenAIDashboardSnapshot, - now.timeIntervalSince(snapshot.updatedAt) < minInterval - { - return - } - - if self.openAIWebDebugLines.isEmpty { - self.resetOpenAIWebDebugLog(context: "refresh") - } else { - let stamp = Date().formatted(date: .abbreviated, time: .shortened) - self.logOpenAIWeb("[\(stamp)] OpenAI web refresh start") - } - let log: (String) -> Void = { [weak self] line in - guard let self else { return } - self.logOpenAIWeb(line) - } - - do { - let normalized = targetEmail? - .trimmingCharacters(in: .whitespacesAndNewlines) - .lowercased() - var effectiveEmail = targetEmail - - // Use a per-email persistent `WKWebsiteDataStore` so multiple dashboard sessions can coexist. - // Strategy: - // - Try the existing per-email WebKit cookie store first (fast; avoids Keychain prompts). - // - On login-required or account mismatch, import cookies from the configured browser order and retry once. - if self.openAIWebAccountDidChange, let targetEmail, !targetEmail.isEmpty { - // On account switches, proactively re-import cookies so we don't show stale data from the previous - // user. - if let imported = await self.importOpenAIDashboardCookiesIfNeeded( - targetEmail: targetEmail, - force: true) - { - effectiveEmail = imported - } - self.openAIWebAccountDidChange = false - } - - var dash = try await OpenAIDashboardFetcher().loadLatestDashboard( - accountEmail: effectiveEmail, - logger: log, - debugDumpHTML: false, - timeout: Self.openAIWebPrimaryFetchTimeout) - - if self.dashboardEmailMismatch(expected: normalized, actual: dash.signedInEmail) { - if let imported = await self.importOpenAIDashboardCookiesIfNeeded( - targetEmail: targetEmail, - force: true) - { - effectiveEmail = imported - } - dash = try await OpenAIDashboardFetcher().loadLatestDashboard( - accountEmail: effectiveEmail, - logger: log, - debugDumpHTML: false, - timeout: Self.openAIWebRetryFetchTimeout) - } - - if self.dashboardEmailMismatch(expected: normalized, actual: dash.signedInEmail) { - let signedIn = dash.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "unknown" - await MainActor.run { - self.openAIDashboard = nil - self.lastOpenAIDashboardError = [ - "OpenAI dashboard signed in as \(signedIn), but Codex uses \(normalized ?? "unknown").", - "Switch accounts in your browser and update OpenAI cookies in Providers → Codex.", - ].joined(separator: " ") - self.openAIDashboardRequiresLogin = true - } - return - } - - await self.applyOpenAIDashboard(dash, targetEmail: effectiveEmail) - } catch let OpenAIDashboardFetcher.FetchError.noDashboardData(body) { - // Often indicates a missing/stale session without an obvious login prompt. Retry once after - // importing cookies from the user's browser. - let targetEmail = self.codexAccountEmailForOpenAIDashboard() - var effectiveEmail = targetEmail - if let imported = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) { - effectiveEmail = imported - } - do { - let dash = try await OpenAIDashboardFetcher().loadLatestDashboard( - accountEmail: effectiveEmail, - logger: log, - debugDumpHTML: true, - timeout: Self.openAIWebRetryFetchTimeout) - await self.applyOpenAIDashboard(dash, targetEmail: effectiveEmail) - } catch let OpenAIDashboardFetcher.FetchError.noDashboardData(retryBody) { - let finalBody = retryBody.isEmpty ? body : retryBody - let message = self.openAIDashboardFriendlyError( - body: finalBody, - targetEmail: targetEmail, - cookieImportStatus: self.openAIDashboardCookieImportStatus) - ?? OpenAIDashboardFetcher.FetchError.noDashboardData(body: finalBody).localizedDescription - await self.applyOpenAIDashboardFailure(message: message) - } catch { - await self.applyOpenAIDashboardFailure(message: error.localizedDescription) - } - } catch OpenAIDashboardFetcher.FetchError.loginRequired { - let targetEmail = self.codexAccountEmailForOpenAIDashboard() - var effectiveEmail = targetEmail - if let imported = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) { - effectiveEmail = imported - } - do { - let dash = try await OpenAIDashboardFetcher().loadLatestDashboard( - accountEmail: effectiveEmail, - logger: log, - debugDumpHTML: true, - timeout: Self.openAIWebRetryFetchTimeout) - await self.applyOpenAIDashboard(dash, targetEmail: effectiveEmail) - } catch OpenAIDashboardFetcher.FetchError.loginRequired { - await MainActor.run { - self.lastOpenAIDashboardError = [ - "OpenAI web access requires a signed-in chatgpt.com session.", - "Sign in using \(self.codexBrowserCookieOrder.loginHint), " + - "then update OpenAI cookies in Providers → Codex.", - ].joined(separator: " ") - self.openAIDashboard = self.lastOpenAIDashboardSnapshot - self.openAIDashboardRequiresLogin = true - } - } catch { - await self.applyOpenAIDashboardFailure(message: error.localizedDescription) - } - } catch { - await self.applyOpenAIDashboardFailure(message: error.localizedDescription) - } - } - - // MARK: - OpenAI web account switching - - /// Detect Codex account email changes and clear stale OpenAI web state so the UI can't show the wrong user. - /// This does not delete other per-email WebKit cookie stores (we keep multiple accounts around). - func handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: String?) { - let normalized = targetEmail? - .trimmingCharacters(in: .whitespacesAndNewlines) - .lowercased() - - guard let normalized, !normalized.isEmpty else { return } - - let previous = self.lastOpenAIDashboardTargetEmail - self.lastOpenAIDashboardTargetEmail = normalized - - if let previous, - !previous.isEmpty, - previous != normalized - { - let stamp = Date().formatted(date: .abbreviated, time: .shortened) - self.logOpenAIWeb( - "[\(stamp)] Codex account changed: \(previous) → \(normalized); " + - "clearing OpenAI web snapshot") - self.openAIWebAccountDidChange = true - self.openAIDashboard = nil - self.lastOpenAIDashboardSnapshot = nil - self.lastOpenAIDashboardError = nil - self.openAIDashboardRequiresLogin = true - self.openAIDashboardCookieImportStatus = "Codex account changed; importing browser cookies…" - self.lastOpenAIDashboardCookieImportAttemptAt = nil - self.lastOpenAIDashboardCookieImportEmail = nil - } - } - - func importOpenAIDashboardBrowserCookiesNow() async { - self.resetOpenAIWebDebugLog(context: "manual import") - let targetEmail = self.codexAccountEmailForOpenAIDashboard() - _ = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) - await self.refreshOpenAIDashboardIfNeeded(force: true) - } - - private func importOpenAIDashboardCookiesIfNeeded(targetEmail: String?, force: Bool) async -> String? { - let normalizedTarget = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) - let allowAnyAccount = normalizedTarget == nil || normalizedTarget?.isEmpty == true - let cookieSource = self.settings.codexCookieSource - - let now = Date() - let lastEmail = self.lastOpenAIDashboardCookieImportEmail - let lastAttempt = self.lastOpenAIDashboardCookieImportAttemptAt ?? .distantPast - - let shouldAttempt: Bool = if force { - true - } else { - if allowAnyAccount { - now.timeIntervalSince(lastAttempt) > 300 - } else { - self.openAIDashboardRequiresLogin && - ( - lastEmail?.lowercased() != normalizedTarget?.lowercased() || now - .timeIntervalSince(lastAttempt) > 300) - } - } - - guard shouldAttempt else { return normalizedTarget } - self.lastOpenAIDashboardCookieImportEmail = normalizedTarget - self.lastOpenAIDashboardCookieImportAttemptAt = now - - let stamp = now.formatted(date: .abbreviated, time: .shortened) - let targetLabel = normalizedTarget ?? "unknown" - self.logOpenAIWeb("[\(stamp)] import start (target=\(targetLabel))") - - do { - let log: (String) -> Void = { [weak self] message in - guard let self else { return } - self.logOpenAIWeb(message) - } - - let importer = OpenAIDashboardBrowserCookieImporter(browserDetection: self.browserDetection) - let result: OpenAIDashboardBrowserCookieImporter.ImportResult - switch cookieSource { - case .manual: - self.settings.ensureCodexCookieLoaded() - let manualHeader = self.settings.codexCookieHeader - guard CookieHeaderNormalizer.normalize(manualHeader) != nil else { - throw OpenAIDashboardBrowserCookieImporter.ImportError.manualCookieHeaderInvalid - } - result = try await importer.importManualCookies( - cookieHeader: manualHeader, - intoAccountEmail: normalizedTarget, - allowAnyAccount: allowAnyAccount, - logger: log) - case .auto: - result = try await importer.importBestCookies( - intoAccountEmail: normalizedTarget, - allowAnyAccount: allowAnyAccount, - logger: log) - case .off: - result = OpenAIDashboardBrowserCookieImporter.ImportResult( - sourceLabel: "Off", - cookieCount: 0, - signedInEmail: normalizedTarget, - matchesCodexEmail: true) - } - let effectiveEmail = result.signedInEmail? - .trimmingCharacters(in: .whitespacesAndNewlines) - .isEmpty == false - ? result.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) - : normalizedTarget - self.lastOpenAIDashboardCookieImportEmail = effectiveEmail ?? normalizedTarget - await MainActor.run { - let signed = result.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) - let matchText = result.matchesCodexEmail ? "matches Codex" : "does not match Codex" - let sourceLabel = switch cookieSource { - case .manual: - "Manual cookie header" - case .auto: - "\(result.sourceLabel) cookies" - case .off: - "OpenAI cookies disabled" - } - if let signed, !signed.isEmpty { - self.openAIDashboardCookieImportStatus = - allowAnyAccount - ? [ - "Using \(sourceLabel) (\(result.cookieCount)).", - "Signed in as \(signed).", - ].joined(separator: " ") - : [ - "Using \(sourceLabel) (\(result.cookieCount)).", - "Signed in as \(signed) (\(matchText)).", - ].joined(separator: " ") - } else { - self.openAIDashboardCookieImportStatus = - "Using \(sourceLabel) (\(result.cookieCount))." - } - } - return effectiveEmail - } catch let err as OpenAIDashboardBrowserCookieImporter.ImportError { - switch err { - case let .noMatchingAccount(found): - let foundText: String = if found.isEmpty { - "no signed-in session detected in \(self.codexBrowserCookieOrder.loginHint)" - } else { - found - .sorted { lhs, rhs in - if lhs.sourceLabel == rhs.sourceLabel { return lhs.email < rhs.email } - return lhs.sourceLabel < rhs.sourceLabel - } - .map { "\($0.sourceLabel): \($0.email)" } - .joined(separator: " • ") - } - self.logOpenAIWeb("[\(stamp)] import mismatch: \(foundText)") - await MainActor.run { - self.openAIDashboardCookieImportStatus = allowAnyAccount - ? [ - "No signed-in OpenAI web session found.", - "Found \(foundText).", - ].joined(separator: " ") - : [ - "Browser cookies do not match Codex account (\(normalizedTarget ?? "unknown")).", - "Found \(foundText).", - ].joined(separator: " ") - // Treat mismatch like "not logged in" for the current Codex account. - self.openAIDashboardRequiresLogin = true - self.openAIDashboard = nil - } - case .noCookiesFound, - .browserAccessDenied, - .dashboardStillRequiresLogin, - .manualCookieHeaderInvalid: - self.logOpenAIWeb("[\(stamp)] import failed: \(err.localizedDescription)") - await MainActor.run { - self.openAIDashboardCookieImportStatus = - "OpenAI cookie import failed: \(err.localizedDescription)" - self.openAIDashboardRequiresLogin = true - } - } - } catch { - self.logOpenAIWeb("[\(stamp)] import failed: \(error.localizedDescription)") - await MainActor.run { - self.openAIDashboardCookieImportStatus = - "Browser cookie import failed: \(error.localizedDescription)" - } - } - return nil - } - - private func resetOpenAIWebDebugLog(context: String) { - let stamp = Date().formatted(date: .abbreviated, time: .shortened) - self.openAIWebDebugLines.removeAll(keepingCapacity: true) - self.openAIDashboardCookieImportDebugLog = nil - self.logOpenAIWeb("[\(stamp)] OpenAI web \(context) start") - } - - private func logOpenAIWeb(_ message: String) { - let safeMessage = LogRedactor.redact(message) - self.openAIWebLogger.debug(safeMessage) - self.openAIWebDebugLines.append(safeMessage) - if self.openAIWebDebugLines.count > 240 { - self.openAIWebDebugLines.removeFirst(self.openAIWebDebugLines.count - 240) - } - self.openAIDashboardCookieImportDebugLog = self.openAIWebDebugLines.joined(separator: "\n") - } - - func resetOpenAIWebState() { - self.openAIDashboard = nil - self.lastOpenAIDashboardError = nil - self.lastOpenAIDashboardSnapshot = nil - self.lastOpenAIDashboardTargetEmail = nil - self.openAIDashboardRequiresLogin = false - self.openAIDashboardCookieImportStatus = nil - self.openAIDashboardCookieImportDebugLog = nil - self.lastOpenAIDashboardCookieImportAttemptAt = nil - self.lastOpenAIDashboardCookieImportEmail = nil - } - - private func dashboardEmailMismatch(expected: String?, actual: String?) -> Bool { - guard let expected, !expected.isEmpty else { return false } - guard let raw = actual?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return false } - return raw.lowercased() != expected.lowercased() - } - - func codexAccountEmailForOpenAIDashboard() -> String? { - let direct = self.snapshots[.codex]?.accountEmail(for: .codex)? - .trimmingCharacters(in: .whitespacesAndNewlines) - if let direct, !direct.isEmpty { return direct } - let fallback = self.codexFetcher.loadAccountInfo().email?.trimmingCharacters(in: .whitespacesAndNewlines) - if let fallback, !fallback.isEmpty { return fallback } - let cached = self.openAIDashboard?.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) - if let cached, !cached.isEmpty { return cached } - let imported = self.lastOpenAIDashboardCookieImportEmail?.trimmingCharacters(in: .whitespacesAndNewlines) - if let imported, !imported.isEmpty { return imported } - return nil - } -} - extension UsageStore { func debugDumpClaude() async { let fetcher = ClaudeUsageFetcher( @@ -1568,6 +1148,11 @@ extension UsageStore { do { let fetcher = self.costUsageFetcher let timeoutSeconds = self.tokenFetchTimeout + // CostUsageFetcher scans local Codex session logs from this machine. That data is + // intentionally presented as provider-level local telemetry rather than managed-account + // remote state, so managed Codex account selection does not retarget this fetch. + // If the UI later needs account-scoped token history, it should label and source that + // separately instead of silently changing the meaning of this section. let snapshot = try await withThrowingTaskGroup(of: CostUsageTokenSnapshot.self) { group in group.addTask(priority: .utility) { try await fetcher.loadTokenSnapshot( diff --git a/Sources/CodexBarCore/CodexHomeScope.swift b/Sources/CodexBarCore/CodexHomeScope.swift new file mode 100644 index 000000000..5da3d65eb --- /dev/null +++ b/Sources/CodexBarCore/CodexHomeScope.swift @@ -0,0 +1,21 @@ +import Foundation + +public enum CodexHomeScope { + public static func ambientHomeURL( + env: [String: String], + fileManager: FileManager = .default) + -> URL + { + if let raw = env["CODEX_HOME"]?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty { + return URL(fileURLWithPath: raw, isDirectory: true) + } + return fileManager.homeDirectoryForCurrentUser.appendingPathComponent(".codex", isDirectory: true) + } + + public static func scopedEnvironment(base: [String: String], codexHome: String?) -> [String: String] { + guard let codexHome, !codexHome.isEmpty else { return base } + var env = base + env["CODEX_HOME"] = codexHome + return env + } +} diff --git a/Sources/CodexBarCore/CodexManagedAccounts.swift b/Sources/CodexBarCore/CodexManagedAccounts.swift new file mode 100644 index 000000000..bfae7ce50 --- /dev/null +++ b/Sources/CodexBarCore/CodexManagedAccounts.swift @@ -0,0 +1,82 @@ +import Foundation + +public struct ManagedCodexAccount: Codable, Identifiable, Sendable { + public let id: UUID + public let email: String + public let managedHomePath: String + public let createdAt: TimeInterval + public let updatedAt: TimeInterval + public let lastAuthenticatedAt: TimeInterval? + + public init( + id: UUID, + email: String, + managedHomePath: String, + createdAt: TimeInterval, + updatedAt: TimeInterval, + lastAuthenticatedAt: TimeInterval?) + { + self.id = id + self.email = Self.normalizeEmail(email) + self.managedHomePath = managedHomePath + self.createdAt = createdAt + self.updatedAt = updatedAt + self.lastAuthenticatedAt = lastAuthenticatedAt + } + + static func normalizeEmail(_ email: String) -> String { + email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + try self.init( + id: container.decode(UUID.self, forKey: .id), + email: container.decode(String.self, forKey: .email), + managedHomePath: container.decode(String.self, forKey: .managedHomePath), + createdAt: container.decode(TimeInterval.self, forKey: .createdAt), + updatedAt: container.decode(TimeInterval.self, forKey: .updatedAt), + lastAuthenticatedAt: container.decodeIfPresent(TimeInterval.self, forKey: .lastAuthenticatedAt)) + } +} + +public struct ManagedCodexAccountSet: Codable, Sendable { + public let version: Int + public let accounts: [ManagedCodexAccount] + + public init(version: Int, accounts: [ManagedCodexAccount]) { + self.version = version + self.accounts = Self.sanitizedAccounts(accounts) + } + + public func account(id: UUID) -> ManagedCodexAccount? { + self.accounts.first { $0.id == id } + } + + public func account(email: String) -> ManagedCodexAccount? { + let normalizedEmail = ManagedCodexAccount.normalizeEmail(email) + return self.accounts.first { $0.email == normalizedEmail } + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + try self.init( + version: container.decode(Int.self, forKey: .version), + accounts: container.decode([ManagedCodexAccount].self, forKey: .accounts)) + } + + private static func sanitizedAccounts(_ accounts: [ManagedCodexAccount]) -> [ManagedCodexAccount] { + var seenIDs: Set = [] + var seenEmails: Set = [] + var sanitized: [ManagedCodexAccount] = [] + sanitized.reserveCapacity(accounts.count) + + for account in accounts { + guard seenIDs.insert(account.id).inserted else { continue } + guard seenEmails.insert(account.email).inserted else { continue } + sanitized.append(account) + } + + return sanitized + } +} diff --git a/Sources/CodexBarCore/Config/CodexActiveSource.swift b/Sources/CodexBarCore/Config/CodexActiveSource.swift new file mode 100644 index 000000000..e009e0d84 --- /dev/null +++ b/Sources/CodexBarCore/Config/CodexActiveSource.swift @@ -0,0 +1,38 @@ +import Foundation + +public enum CodexActiveSource: Codable, Equatable, Sendable { + case liveSystem + case managedAccount(id: UUID) + + private enum CodingKeys: String, CodingKey { + case kind + case accountID + } + + private enum Kind: String, Codable { + case liveSystem + case managedAccount + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + switch try container.decode(Kind.self, forKey: .kind) { + case .liveSystem: + self = .liveSystem + case .managedAccount: + let id = try container.decode(UUID.self, forKey: .accountID) + self = .managedAccount(id: id) + } + } + + public func encode(to encoder: any Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + switch self { + case .liveSystem: + try container.encode(Kind.liveSystem, forKey: .kind) + case let .managedAccount(id): + try container.encode(Kind.managedAccount, forKey: .kind) + try container.encode(id, forKey: .accountID) + } + } +} diff --git a/Sources/CodexBarCore/Config/CodexBarConfig.swift b/Sources/CodexBarCore/Config/CodexBarConfig.swift index 04c85409d..c20759676 100644 --- a/Sources/CodexBarCore/Config/CodexBarConfig.swift +++ b/Sources/CodexBarCore/Config/CodexBarConfig.swift @@ -83,6 +83,7 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { public var region: String? public var workspaceID: String? public var tokenAccounts: ProviderTokenAccountData? + public var codexActiveSource: CodexActiveSource? public init( id: UsageProvider, @@ -94,7 +95,8 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { cookieSource: ProviderCookieSource? = nil, region: String? = nil, workspaceID: String? = nil, - tokenAccounts: ProviderTokenAccountData? = nil) + tokenAccounts: ProviderTokenAccountData? = nil, + codexActiveSource: CodexActiveSource? = nil) { self.id = id self.enabled = enabled @@ -106,6 +108,7 @@ public struct ProviderConfig: Codable, Sendable, Identifiable { self.region = region self.workspaceID = workspaceID self.tokenAccounts = tokenAccounts + self.codexActiveSource = codexActiveSource } public var sanitizedAPIKey: String? { diff --git a/Sources/CodexBarCore/CookieHeaderCache.swift b/Sources/CodexBarCore/CookieHeaderCache.swift index c00610070..0d24cde93 100644 --- a/Sources/CodexBarCore/CookieHeaderCache.swift +++ b/Sources/CodexBarCore/CookieHeaderCache.swift @@ -1,6 +1,20 @@ import Foundation public enum CookieHeaderCache { + public enum Scope: Sendable, Equatable { + case managedAccount(UUID) + case managedStoreUnreadable + + fileprivate var keychainIdentifier: String { + switch self { + case let .managedAccount(accountID): + "managed.\(accountID.uuidString.lowercased())" + case .managedStoreUnreadable: + "managed-store-unreadable" + } + } + } + public struct Entry: Codable, Sendable { public let cookieHeader: String public let storedAt: Date @@ -16,8 +30,8 @@ public enum CookieHeaderCache { private static let log = CodexBarLog.logger(LogCategories.cookieCache) private nonisolated(unsafe) static var legacyBaseURLOverride: URL? - public static func load(provider: UsageProvider) -> Entry? { - let key = KeychainCacheStore.Key.cookie(provider: provider) + public static func load(provider: UsageProvider, scope: Scope? = nil) -> Entry? { + let key = self.key(for: provider, scope: scope) switch KeychainCacheStore.load(key: key, as: Entry.self) { case let .found(entry): self.log.debug("Cookie cache hit", metadata: ["provider": provider.rawValue]) @@ -29,6 +43,7 @@ public enum CookieHeaderCache { self.log.debug("Cookie cache miss", metadata: ["provider": provider.rawValue]) } + guard scope == nil else { return nil } guard let legacy = self.loadLegacyEntry(for: provider) else { return nil } KeychainCacheStore.store(key: key, entry: legacy) self.removeLegacyEntry(for: provider) @@ -38,26 +53,31 @@ public enum CookieHeaderCache { public static func store( provider: UsageProvider, + scope: Scope? = nil, cookieHeader: String, sourceLabel: String, now: Date = Date()) { let trimmed = cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines) guard let normalized = CookieHeaderNormalizer.normalize(trimmed), !normalized.isEmpty else { - self.clear(provider: provider) + self.clear(provider: provider, scope: scope) return } let entry = Entry(cookieHeader: normalized, storedAt: now, sourceLabel: sourceLabel) - let key = KeychainCacheStore.Key.cookie(provider: provider) + let key = self.key(for: provider, scope: scope) KeychainCacheStore.store(key: key, entry: entry) - self.removeLegacyEntry(for: provider) + if scope == nil { + self.removeLegacyEntry(for: provider) + } self.log.debug("Cookie cache stored", metadata: ["provider": provider.rawValue, "source": sourceLabel]) } - public static func clear(provider: UsageProvider) { - let key = KeychainCacheStore.Key.cookie(provider: provider) + public static func clear(provider: UsageProvider, scope: Scope? = nil) { + let key = self.key(for: provider, scope: scope) KeychainCacheStore.clear(key: key) - self.removeLegacyEntry(for: provider) + if scope == nil { + self.removeLegacyEntry(for: provider) + } self.log.debug("Cookie cache cleared", metadata: ["provider": provider.rawValue]) } @@ -110,4 +130,8 @@ public enum CookieHeaderCache { return base.appendingPathComponent("CodexBar", isDirectory: true) .appendingPathComponent("\(provider.rawValue)-cookie.json") } + + private static func key(for provider: UsageProvider, scope: Scope?) -> KeychainCacheStore.Key { + KeychainCacheStore.Key.cookie(provider: provider, scopeIdentifier: scope?.keychainIdentifier) + } } diff --git a/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift b/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift index cdab14aba..8870d48db 100644 --- a/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift +++ b/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift @@ -95,6 +95,7 @@ public struct TTYCommandRunner { public var idleTimeout: TimeInterval? public var workingDirectory: URL? public var extraArgs: [String] = [] + public var baseEnvironment: [String: String]? public var initialDelay: TimeInterval = 0.4 public var sendEnterEvery: TimeInterval? public var sendOnSubstrings: [String: String] @@ -109,6 +110,7 @@ public struct TTYCommandRunner { idleTimeout: TimeInterval? = nil, workingDirectory: URL? = nil, extraArgs: [String] = [], + baseEnvironment: [String: String]? = nil, initialDelay: TimeInterval = 0.4, sendEnterEvery: TimeInterval? = nil, sendOnSubstrings: [String: String] = [:], @@ -122,6 +124,7 @@ public struct TTYCommandRunner { self.idleTimeout = idleTimeout self.workingDirectory = workingDirectory self.extraArgs = extraArgs + self.baseEnvironment = baseEnvironment self.initialDelay = initialDelay self.sendEnterEvery = sendEnterEvery self.sendOnSubstrings = sendOnSubstrings @@ -375,7 +378,8 @@ public struct TTYCommandRunner { proc.standardError = secondaryHandle // Use login-shell PATH when available, but keep the caller’s environment (HOME, LANG, etc.) so // the CLIs can find their auth/config files. - var env = Self.enrichedEnvironment() + let baseEnv = options.baseEnvironment ?? ProcessInfo.processInfo.environment + var env = Self.enrichedEnvironment(baseEnv: baseEnv, home: baseEnv["HOME"] ?? NSHomeDirectory()) if let workingDirectory = options.workingDirectory { proc.currentDirectoryURL = workingDirectory env["PWD"] = workingDirectory.path diff --git a/Sources/CodexBarCore/KeychainCacheStore.swift b/Sources/CodexBarCore/KeychainCacheStore.swift index e77ebbd77..ebe5c45ae 100644 --- a/Sources/CodexBarCore/KeychainCacheStore.swift +++ b/Sources/CodexBarCore/KeychainCacheStore.swift @@ -244,8 +244,13 @@ public enum KeychainCacheStore { } extension KeychainCacheStore.Key { - public static func cookie(provider: UsageProvider) -> Self { - Self(category: "cookie", identifier: provider.rawValue) + public static func cookie(provider: UsageProvider, scopeIdentifier: String? = nil) -> Self { + let identifier: String = if let scopeIdentifier, !scopeIdentifier.isEmpty { + "\(provider.rawValue).\(scopeIdentifier)" + } else { + provider.rawValue + } + return Self(category: "cookie", identifier: identifier) } public static func oauth(provider: UsageProvider) -> Self { diff --git a/Sources/CodexBarCore/ManagedCodexAccountStore.swift b/Sources/CodexBarCore/ManagedCodexAccountStore.swift new file mode 100644 index 000000000..9c0baad7b --- /dev/null +++ b/Sources/CodexBarCore/ManagedCodexAccountStore.swift @@ -0,0 +1,78 @@ +import Foundation + +public enum FileManagedCodexAccountStoreError: Error, Equatable, Sendable { + case unsupportedVersion(Int) +} + +public protocol ManagedCodexAccountStoring: Sendable { + func loadAccounts() throws -> ManagedCodexAccountSet + func storeAccounts(_ accounts: ManagedCodexAccountSet) throws + func ensureFileExists() throws -> URL +} + +public struct FileManagedCodexAccountStore: ManagedCodexAccountStoring, @unchecked Sendable { + public static let currentVersion = 1 + + private let fileURL: URL + private let fileManager: FileManager + + public init(fileURL: URL = Self.defaultURL(), fileManager: FileManager = .default) { + self.fileURL = fileURL + self.fileManager = fileManager + } + + public func loadAccounts() throws -> ManagedCodexAccountSet { + guard self.fileManager.fileExists(atPath: self.fileURL.path) else { + return Self.emptyAccountSet() + } + + let data = try Data(contentsOf: self.fileURL) + let decoder = JSONDecoder() + let accounts = try decoder.decode(ManagedCodexAccountSet.self, from: data) + guard accounts.version == Self.currentVersion else { + throw FileManagedCodexAccountStoreError.unsupportedVersion(accounts.version) + } + return accounts + } + + public func storeAccounts(_ accounts: ManagedCodexAccountSet) throws { + let normalizedAccounts = ManagedCodexAccountSet( + version: Self.currentVersion, + accounts: accounts.accounts) + let encoder = JSONEncoder() + encoder.outputFormatting = [.prettyPrinted, .sortedKeys] + let data = try encoder.encode(normalizedAccounts) + let directory = self.fileURL.deletingLastPathComponent() + if !self.fileManager.fileExists(atPath: directory.path) { + try self.fileManager.createDirectory(at: directory, withIntermediateDirectories: true) + } + try data.write(to: self.fileURL, options: [.atomic]) + try self.applySecurePermissionsIfNeeded() + } + + public func ensureFileExists() throws -> URL { + if self.fileManager.fileExists(atPath: self.fileURL.path) { return self.fileURL } + try self.storeAccounts(Self.emptyAccountSet()) + return self.fileURL + } + + private func applySecurePermissionsIfNeeded() throws { + #if os(macOS) + try self.fileManager.setAttributes([ + .posixPermissions: NSNumber(value: Int16(0o600)), + ], ofItemAtPath: self.fileURL.path) + #endif + } + + private static func emptyAccountSet() -> ManagedCodexAccountSet { + ManagedCodexAccountSet(version: self.currentVersion, accounts: []) + } + + public static func defaultURL() -> URL { + let base = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? FileManager.default.homeDirectoryForCurrentUser + return base + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("managed-codex-accounts.json") + } +} diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift index cf3deda4a..514c7f5f6 100644 --- a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift @@ -78,6 +78,12 @@ public struct OpenAIDashboardBrowserCookieImporter { var accessDeniedHints: [String] = [] } + private struct ImportContext { + let targetEmail: String? + let allowAnyAccount: Bool + let cacheScope: CookieHeaderCache.Scope? + } + private static let cookieDomains = ["chatgpt.com", "openai.com"] private static let cookieClient = BrowserCookieClient() private static let cookieImportOrder: BrowserCookieImportOrder = @@ -94,6 +100,7 @@ public struct OpenAIDashboardBrowserCookieImporter { public func importBestCookies( intoAccountEmail targetEmail: String?, allowAnyAccount: Bool = false, + cacheScope: CookieHeaderCache.Scope? = nil, logger: ((String) -> Void)? = nil) async throws -> ImportResult { let log: (String) -> Void = { message in @@ -102,11 +109,15 @@ public struct OpenAIDashboardBrowserCookieImporter { let targetEmail = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) let normalizedTarget = targetEmail?.isEmpty == false ? targetEmail : nil + let context = ImportContext( + targetEmail: normalizedTarget, + allowAnyAccount: allowAnyAccount, + cacheScope: cacheScope) if normalizedTarget != nil { log("Codex email known; matching required.") } else { - guard allowAnyAccount else { + guard context.allowAnyAccount else { throw ImportError.noCookiesFound } log("Codex email unknown; importing any signed-in session.") @@ -114,20 +125,21 @@ public struct OpenAIDashboardBrowserCookieImporter { var diagnostics = ImportDiagnostics() - if let cached = CookieHeaderCache.load(provider: .codex), + if let cached = CookieHeaderCache.load(provider: .codex, scope: cacheScope), !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { log("Using cached cookie header from \(cached.sourceLabel)") do { return try await self.importManualCookies( cookieHeader: cached.cookieHeader, - intoAccountEmail: normalizedTarget, - allowAnyAccount: allowAnyAccount, + intoAccountEmail: context.targetEmail, + allowAnyAccount: context.allowAnyAccount, + cacheScope: cacheScope, logger: log) } catch let error as ImportError { switch error { case .manualCookieHeaderInvalid, .noMatchingAccount, .dashboardStillRequiresLogin: - CookieHeaderCache.clear(provider: .codex) + CookieHeaderCache.clear(provider: .codex, scope: cacheScope) default: throw error } @@ -141,8 +153,7 @@ public struct OpenAIDashboardBrowserCookieImporter { for browserSource in installedBrowsers { if let match = await self.trySource( browserSource, - targetEmail: normalizedTarget, - allowAnyAccount: allowAnyAccount, + context: context, log: log, diagnostics: &diagnostics) { @@ -177,6 +188,7 @@ public struct OpenAIDashboardBrowserCookieImporter { cookieHeader: String, intoAccountEmail targetEmail: String?, allowAnyAccount: Bool = false, + cacheScope _: CookieHeaderCache.Scope? = nil, logger: ((String) -> Void)? = nil) async throws -> ImportResult { let log: (String) -> Void = { message in @@ -217,8 +229,7 @@ public struct OpenAIDashboardBrowserCookieImporter { } private func trySafari( - targetEmail: String?, - allowAnyAccount: Bool, + context: ImportContext, log: @escaping (String) -> Void, diagnostics: inout ImportDiagnostics) async -> ImportResult? { @@ -245,8 +256,7 @@ public struct OpenAIDashboardBrowserCookieImporter { let candidate = Candidate(label: source.label, cookies: cookies) if let match = await self.applyCandidate( candidate, - targetEmail: targetEmail, - allowAnyAccount: allowAnyAccount, + context: context, log: log, diagnostics: &diagnostics) { @@ -268,8 +278,7 @@ public struct OpenAIDashboardBrowserCookieImporter { } private func tryChrome( - targetEmail: String?, - allowAnyAccount: Bool, + context: ImportContext, log: @escaping (String) -> Void, diagnostics: inout ImportDiagnostics) async -> ImportResult? { @@ -290,8 +299,7 @@ public struct OpenAIDashboardBrowserCookieImporter { let candidate = Candidate(label: source.label, cookies: cookies) if let match = await self.applyCandidate( candidate, - targetEmail: targetEmail, - allowAnyAccount: allowAnyAccount, + context: context, log: log, diagnostics: &diagnostics) { @@ -313,8 +321,7 @@ public struct OpenAIDashboardBrowserCookieImporter { } private func tryFirefox( - targetEmail: String?, - allowAnyAccount: Bool, + context: ImportContext, log: @escaping (String) -> Void, diagnostics: inout ImportDiagnostics) async -> ImportResult? { @@ -335,8 +342,7 @@ public struct OpenAIDashboardBrowserCookieImporter { let candidate = Candidate(label: source.label, cookies: cookies) if let match = await self.applyCandidate( candidate, - targetEmail: targetEmail, - allowAnyAccount: allowAnyAccount, + context: context, log: log, diagnostics: &diagnostics) { @@ -359,28 +365,24 @@ public struct OpenAIDashboardBrowserCookieImporter { private func trySource( _ source: Browser, - targetEmail: String?, - allowAnyAccount: Bool, + context: ImportContext, log: @escaping (String) -> Void, diagnostics: inout ImportDiagnostics) async -> ImportResult? { switch source { case .safari: await self.trySafari( - targetEmail: targetEmail, - allowAnyAccount: allowAnyAccount, + context: context, log: log, diagnostics: &diagnostics) case .chrome: await self.tryChrome( - targetEmail: targetEmail, - allowAnyAccount: allowAnyAccount, + context: context, log: log, diagnostics: &diagnostics) case .firefox: await self.tryFirefox( - targetEmail: targetEmail, - allowAnyAccount: allowAnyAccount, + context: context, log: log, diagnostics: &diagnostics) default: @@ -390,22 +392,21 @@ public struct OpenAIDashboardBrowserCookieImporter { private func applyCandidate( _ candidate: Candidate, - targetEmail: String?, - allowAnyAccount: Bool, + context: ImportContext, log: @escaping (String) -> Void, diagnostics: inout ImportDiagnostics) async -> ImportResult? { switch await self.evaluateCandidate( candidate, - targetEmail: targetEmail, - allowAnyAccount: allowAnyAccount, + targetEmail: context.targetEmail, + allowAnyAccount: context.allowAnyAccount, log: log) { case let .match(candidate, signedInEmail): log("Selected \(candidate.label) (matches Codex: \(signedInEmail))") - guard let targetEmail else { return nil } + guard let targetEmail = context.targetEmail else { return nil } if let result = try? await self.persist(candidate: candidate, targetEmail: targetEmail, logger: log) { - self.cacheCookies(candidate: candidate) + self.cacheCookies(candidate: candidate, scope: context.cacheScope) return result } return nil @@ -419,15 +420,15 @@ public struct OpenAIDashboardBrowserCookieImporter { case let .loggedIn(candidate, signedInEmail): log("Selected \(candidate.label) (signed in: \(signedInEmail))") if let result = try? await self.persist(candidate: candidate, targetEmail: signedInEmail, logger: log) { - self.cacheCookies(candidate: candidate) + self.cacheCookies(candidate: candidate, scope: context.cacheScope) return result } return nil case .unknown: - if allowAnyAccount { + if context.allowAnyAccount { log("Selected \(candidate.label) (signed in: unknown)") if let result = try? await self.persistToDefaultStore(candidate: candidate, logger: log) { - self.cacheCookies(candidate: candidate) + self.cacheCookies(candidate: candidate, scope: context.cacheScope) return result } return nil @@ -598,6 +599,11 @@ public struct OpenAIDashboardBrowserCookieImporter { // Validate against the persistent store (login + email sync). do { + defer { + // The probe is only a validation step. Start the real dashboard scrape with a + // fresh WKWebView instead of reusing the probe instance. + OpenAIDashboardWebViewCache.shared.evict(websiteDataStore: persistent) + } let probe = try await OpenAIDashboardFetcher().probeUsagePage( websiteDataStore: persistent, logger: logger, @@ -631,6 +637,11 @@ public struct OpenAIDashboardBrowserCookieImporter { await self.setCookies(candidate.cookies, into: persistent) do { + defer { + // The probe is only a validation step. Start the real dashboard scrape with a + // fresh WKWebView instead of reusing the probe instance. + OpenAIDashboardWebViewCache.shared.evict(websiteDataStore: persistent) + } let probe = try await OpenAIDashboardFetcher().probeUsagePage( websiteDataStore: persistent, logger: logger, @@ -669,10 +680,10 @@ public struct OpenAIDashboardBrowserCookieImporter { return cookies } - private func cacheCookies(candidate: Candidate) { + private func cacheCookies(candidate: Candidate, scope: CookieHeaderCache.Scope?) { let header = self.cookieHeader(from: candidate.cookies) guard !header.isEmpty else { return } - CookieHeaderCache.store(provider: .codex, cookieHeader: header, sourceLabel: candidate.label) + CookieHeaderCache.store(provider: .codex, scope: scope, cookieHeader: header, sourceLabel: candidate.label) } private func cookieHeader(from cookies: [HTTPCookie]) -> String { @@ -798,6 +809,7 @@ public struct OpenAIDashboardBrowserCookieImporter { public func importBestCookies( intoAccountEmail _: String?, allowAnyAccount _: Bool = false, + cacheScope _: CookieHeaderCache.Scope? = nil, logger _: ((String) -> Void)? = nil) async throws -> ImportResult { throw ImportError.browserAccessDenied(details: "OpenAI web cookie import is only supported on macOS.") @@ -807,6 +819,7 @@ public struct OpenAIDashboardBrowserCookieImporter { cookieHeader _: String, intoAccountEmail _: String?, allowAnyAccount _: Bool = false, + cacheScope _: CookieHeaderCache.Scope? = nil, logger _: ((String) -> Void)? = nil) async throws -> ImportResult { throw ImportError.browserAccessDenied(details: "OpenAI web cookie import is only supported on macOS.") diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardNavigationDelegate.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardNavigationDelegate.swift index cc7d93709..50581caa6 100644 --- a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardNavigationDelegate.swift +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardNavigationDelegate.swift @@ -41,7 +41,15 @@ final class NavigationDelegate: NSObject, WKNavigationDelegate { nonisolated static func shouldIgnoreNavigationError(_ error: Error) -> Bool { let nsError = error as NSError - return nsError.domain == NSURLErrorDomain && nsError.code == NSURLErrorCancelled + if nsError.domain == NSURLErrorDomain, nsError.code == NSURLErrorCancelled { + return true + } + + if nsError.domain == "WebKitErrorDomain", nsError.code == 102 { + return true + } + + return false } private func completeOnce(_ result: Result) { diff --git a/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift b/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift index 8c7f3b245..8a6c811e8 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift @@ -31,6 +31,7 @@ actor CodexCLISession { private var startedAt: Date? private var ptyRows: UInt16 = 0 private var ptyCols: UInt16 = 0 + private var sessionEnvironment: [String: String]? private struct RollingBuffer { private let maxNeedle: Int @@ -81,8 +82,14 @@ actor CodexCLISession { } // swiftlint:disable cyclomatic_complexity - func captureStatus(binary: String, timeout: TimeInterval, rows: UInt16, cols: UInt16) async throws -> String { - try self.ensureStarted(binary: binary, rows: rows, cols: cols) + func captureStatus( + binary: String, + timeout: TimeInterval, + rows: UInt16, + cols: UInt16, + environment: [String: String]) async throws -> String + { + try self.ensureStarted(binary: binary, rows: rows, cols: cols, environment: environment) if let startedAt { let sinceStart = Date().timeIntervalSince(startedAt) if sinceStart < 0.4 { @@ -243,12 +250,18 @@ actor CodexCLISession { self.cleanup() } - private func ensureStarted(binary: String, rows: UInt16, cols: UInt16) throws { + private func ensureStarted( + binary: String, + rows: UInt16, + cols: UInt16, + environment: [String: String]) throws + { if let proc = self.process, proc.isRunning, self.binaryPath == binary, self.ptyRows == rows, - self.ptyCols == cols + self.ptyCols == cols, + self.sessionEnvironment == environment { return } @@ -273,7 +286,9 @@ actor CodexCLISession { proc.standardOutput = secondaryHandle proc.standardError = secondaryHandle - let env = TTYCommandRunner.enrichedEnvironment() + let env = TTYCommandRunner.enrichedEnvironment( + baseEnv: environment, + home: environment["HOME"] ?? NSHomeDirectory()) proc.environment = env do { @@ -299,6 +314,7 @@ actor CodexCLISession { self.startedAt = Date() self.ptyRows = rows self.ptyCols = cols + self.sessionEnvironment = environment } private func cleanup() { @@ -336,6 +352,7 @@ actor CodexCLISession { self.startedAt = nil self.ptyRows = 0 self.ptyCols = 0 + self.sessionEnvironment = nil } private func readChunk() -> Data { diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthCredentials.swift b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthCredentials.swift index f7f5c3191..c6c5698b2 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthCredentials.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthCredentials.swift @@ -46,19 +46,19 @@ public enum CodexOAuthCredentialsError: LocalizedError, Sendable { } public enum CodexOAuthCredentialsStore { - private static var authFilePath: URL { - let home = FileManager.default.homeDirectoryForCurrentUser - if let codexHome = ProcessInfo.processInfo.environment["CODEX_HOME"]?.trimmingCharacters( - in: .whitespacesAndNewlines), - !codexHome.isEmpty - { - return URL(fileURLWithPath: codexHome).appendingPathComponent("auth.json") - } - return home.appendingPathComponent(".codex").appendingPathComponent("auth.json") + private static func authFilePath( + env: [String: String] = ProcessInfo.processInfo.environment, + fileManager: FileManager = .default) -> URL + { + CodexHomeScope + .ambientHomeURL(env: env, fileManager: fileManager) + .appendingPathComponent("auth.json") } - public static func load() throws -> CodexOAuthCredentials { - let url = self.authFilePath + public static func load(env: [String: String] = ProcessInfo.processInfo + .environment) throws -> CodexOAuthCredentials + { + let url = self.authFilePath(env: env) guard FileManager.default.fileExists(atPath: url.path) else { throw CodexOAuthCredentialsError.notFound } @@ -86,15 +86,18 @@ public enum CodexOAuthCredentialsStore { guard let tokens = json["tokens"] as? [String: Any] else { throw CodexOAuthCredentialsError.missingTokens } - guard let accessToken = tokens["access_token"] as? String, - let refreshToken = tokens["refresh_token"] as? String, + guard let accessToken = Self.stringValue(in: tokens, snakeCaseKey: "access_token", camelCaseKey: "accessToken"), + let refreshToken = Self.stringValue( + in: tokens, + snakeCaseKey: "refresh_token", + camelCaseKey: "refreshToken"), !accessToken.isEmpty else { throw CodexOAuthCredentialsError.missingTokens } - let idToken = tokens["id_token"] as? String - let accountId = tokens["account_id"] as? String + let idToken = Self.stringValue(in: tokens, snakeCaseKey: "id_token", camelCaseKey: "idToken") + let accountId = Self.stringValue(in: tokens, snakeCaseKey: "account_id", camelCaseKey: "accountId") let lastRefresh = Self.parseLastRefresh(from: json["last_refresh"]) return CodexOAuthCredentials( @@ -105,8 +108,11 @@ public enum CodexOAuthCredentialsStore { lastRefresh: lastRefresh) } - public static func save(_ credentials: CodexOAuthCredentials) throws { - let url = self.authFilePath + public static func save( + _ credentials: CodexOAuthCredentials, + env: [String: String] = ProcessInfo.processInfo.environment) throws + { + let url = self.authFilePath(env: env) var json: [String: Any] = [:] if let data = try? Data(contentsOf: url), @@ -143,4 +149,27 @@ public enum CodexOAuthCredentialsStore { formatter.formatOptions = [.withInternetDateTime] return formatter.date(from: value) } + + private static func stringValue( + in dictionary: [String: Any], + snakeCaseKey: String, + camelCaseKey: String) + -> String? + { + if let value = dictionary[snakeCaseKey] as? String, !value.isEmpty { + return value + } + if let value = dictionary[camelCaseKey] as? String, !value.isEmpty { + return value + } + return nil + } +} + +#if DEBUG +extension CodexOAuthCredentialsStore { + static func _authFileURLForTesting(env: [String: String]) -> URL { + self.authFilePath(env: env) + } } +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthUsageFetcher.swift b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthUsageFetcher.swift index 174c08c54..7bcfe079e 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthUsageFetcher.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexOAuth/CodexOAuthUsageFetcher.swift @@ -150,8 +150,12 @@ public enum CodexOAuthUsageFetcher { private static let chatGPTUsagePath = "/wham/usage" private static let codexUsagePath = "/api/codex/usage" - public static func fetchUsage(accessToken: String, accountId: String?) async throws -> CodexUsageResponse { - var request = URLRequest(url: Self.resolveUsageURL()) + public static func fetchUsage( + accessToken: String, + accountId: String?, + env: [String: String] = ProcessInfo.processInfo.environment) async throws -> CodexUsageResponse + { + var request = URLRequest(url: Self.resolveUsageURL(env: env)) request.httpMethod = "GET" request.timeoutInterval = 30 request.setValue("Bearer \(accessToken)", forHTTPHeaderField: "Authorization") @@ -188,8 +192,8 @@ public enum CodexOAuthUsageFetcher { } } - private static func resolveUsageURL() -> URL { - self.resolveUsageURL(env: ProcessInfo.processInfo.environment, configContents: nil) + private static func resolveUsageURL(env: [String: String]) -> URL { + self.resolveUsageURL(env: env, configContents: nil) } private static func resolveUsageURL(env: [String: String], configContents: String?) -> URL { diff --git a/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift index 12b68cd22..6e86f800a 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift @@ -75,13 +75,11 @@ public enum CodexProviderDescriptor { } private static func noDataMessage() -> String { - let fm = FileManager.default - let home = fm.homeDirectoryForCurrentUser.path - let base = ProcessInfo.processInfo.environment["CODEX_HOME"].flatMap { raw -> String? in - let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines) - guard !trimmed.isEmpty else { return nil } - return trimmed - } ?? "\(home)/.codex" + self.noDataMessage(env: ProcessInfo.processInfo.environment) + } + + private static func noDataMessage(env: [String: String], fileManager: FileManager = .default) -> String { + let base = CodexHomeScope.ambientHomeURL(env: env, fileManager: fileManager).path let sessions = "\(base)/sessions" let archived = "\(base)/archived_sessions" return "No Codex sessions found in \(sessions) or \(archived)." @@ -134,21 +132,22 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy { let id: String = "codex.oauth" let kind: ProviderFetchKind = .oauth - func isAvailable(_: ProviderFetchContext) async -> Bool { - (try? CodexOAuthCredentialsStore.load()) != nil + func isAvailable(_ context: ProviderFetchContext) async -> Bool { + (try? CodexOAuthCredentialsStore.load(env: context.env)) != nil } - func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { - var credentials = try CodexOAuthCredentialsStore.load() + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + var credentials = try CodexOAuthCredentialsStore.load(env: context.env) if credentials.needsRefresh, !credentials.refreshToken.isEmpty { credentials = try await CodexTokenRefresher.refresh(credentials) - try CodexOAuthCredentialsStore.save(credentials) + try CodexOAuthCredentialsStore.save(credentials, env: context.env) } let usage = try await CodexOAuthUsageFetcher.fetchUsage( accessToken: credentials.accessToken, - accountId: credentials.accountId) + accountId: credentials.accountId, + env: context.env) return self.makeResult( usage: Self.mapUsage(usage, credentials: credentials), @@ -227,4 +226,10 @@ extension CodexOAuthFetchStrategy { return Self.mapUsage(usage, credentials: credentials) } } + +extension CodexProviderDescriptor { + static func _noDataMessageForTesting(env: [String: String]) -> String { + self.noDataMessage(env: env) + } +} #endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexStatusProbe.swift b/Sources/CodexBarCore/Providers/Codex/CodexStatusProbe.swift index 6e8667036..3226674b3 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexStatusProbe.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexStatusProbe.swift @@ -59,21 +59,24 @@ public struct CodexStatusProbe { public var codexBinary: String = "codex" public var timeout: TimeInterval = Self.defaultTimeoutSeconds public var keepCLISessionsAlive: Bool = false + public var environment: [String: String] = ProcessInfo.processInfo.environment public init() {} public init( codexBinary: String = "codex", timeout: TimeInterval = 8.0, - keepCLISessionsAlive: Bool = false) + keepCLISessionsAlive: Bool = false, + environment: [String: String] = ProcessInfo.processInfo.environment) { self.codexBinary = codexBinary self.timeout = timeout self.keepCLISessionsAlive = keepCLISessionsAlive + self.environment = environment } public func fetch() async throws -> CodexStatusSnapshot { - let env = ProcessInfo.processInfo.environment + let env = self.environment let resolved = BinaryLocator.resolveCodexBinary(env: env, loginPATH: LoginShellPathCache.shared.current) ?? self.codexBinary guard FileManager.default.isExecutableFile(atPath: resolved) || TTYCommandRunner.which(resolved) != nil else { @@ -200,7 +203,8 @@ public struct CodexStatusProbe { binary: binary, timeout: timeout, rows: rows, - cols: cols) + cols: cols, + environment: self.environment) } catch CodexCLISession.SessionError.processExited { throw CodexStatusProbeError.timedOut } catch CodexCLISession.SessionError.timedOut { @@ -218,7 +222,8 @@ public struct CodexStatusProbe { rows: rows, cols: cols, timeout: timeout, - extraArgs: ["-s", "read-only", "-a", "untrusted"])) + extraArgs: ["-s", "read-only", "-a", "untrusted"], + baseEnvironment: self.environment)) text = result.text } return try Self.parse(text: text) diff --git a/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift index 2e03d51a3..5a402ea5c 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift @@ -9,10 +9,23 @@ public struct CodexWebDashboardStrategy: ProviderFetchStrategy { public init() {} public func isAvailable(_ context: ProviderFetchContext) async -> Bool { - context.sourceMode.usesWeb + context.sourceMode.usesWeb && + !Self.managedAccountStoreIsUnreadable(context) && + !Self.managedAccountTargetIsUnavailable(context) } public func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + guard !Self.managedAccountStoreIsUnreadable(context) else { + // A fail-closed placeholder CODEX_HOME does not identify a target account. If the managed store + // itself is unreadable, web import must not fall back to "any signed-in browser account". + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + guard !Self.managedAccountTargetIsUnavailable(context) else { + // If the selected managed account no longer exists in a readable store, web import must not + // fall back to "any signed-in browser account" for that stale selection. + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + // Ensure AppKit is initialized before using WebKit in a CLI. await MainActor.run { _ = NSApplication.shared @@ -41,6 +54,14 @@ public struct CodexWebDashboardStrategy: ProviderFetchStrategy { _ = error return true } + + private static func managedAccountStoreIsUnreadable(_ context: ProviderFetchContext) -> Bool { + context.settings?.codex?.managedAccountStoreUnreadable == true + } + + private static func managedAccountTargetIsUnavailable(_ context: ProviderFetchContext) -> Bool { + context.settings?.codex?.managedAccountTargetUnavailable == true + } } private struct OpenAIWebCodexResult { diff --git a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift index c5d9af4f7..ac667b803 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -44,15 +44,21 @@ public struct ProviderSettingsSnapshot: Sendable { public let usageDataSource: CodexUsageDataSource public let cookieSource: ProviderCookieSource public let manualCookieHeader: String? + public let managedAccountStoreUnreadable: Bool + public let managedAccountTargetUnavailable: Bool public init( usageDataSource: CodexUsageDataSource, cookieSource: ProviderCookieSource, - manualCookieHeader: String?) + manualCookieHeader: String?, + managedAccountStoreUnreadable: Bool = false, + managedAccountTargetUnavailable: Bool = false) { self.usageDataSource = usageDataSource self.cookieSource = cookieSource self.manualCookieHeader = manualCookieHeader + self.managedAccountStoreUnreadable = managedAccountStoreUnreadable + self.managedAccountTargetUnavailable = managedAccountTargetUnavailable } } diff --git a/Sources/CodexBarCore/UsageFetcher.swift b/Sources/CodexBarCore/UsageFetcher.swift index 98859b18e..077331168 100644 --- a/Sources/CodexBarCore/UsageFetcher.swift +++ b/Sources/CodexBarCore/UsageFetcher.swift @@ -350,7 +350,8 @@ private final class CodexRPCClient: @unchecked Sendable { init( executable: String = "codex", - arguments: [String] = ["-s", "read-only", "-a", "untrusted", "app-server"]) throws + arguments: [String] = ["-s", "read-only", "-a", "untrusted", "app-server"], + environment: [String: String] = ProcessInfo.processInfo.environment) throws { var stdoutContinuation: AsyncStream.Continuation! self.stdoutLineStream = AsyncStream { continuation in @@ -358,7 +359,7 @@ private final class CodexRPCClient: @unchecked Sendable { } self.stdoutLineContinuation = stdoutContinuation - let resolvedExec = BinaryLocator.resolveCodexBinary() + let resolvedExec = BinaryLocator.resolveCodexBinary(env: environment) ?? TTYCommandRunner.which(executable) guard let resolvedExec else { @@ -366,7 +367,7 @@ private final class CodexRPCClient: @unchecked Sendable { throw RPCWireError.startFailed( "Codex CLI not found. Install with `npm i -g @openai/codex` (or bun) then relaunch CodexBar.") } - var env = ProcessInfo.processInfo.environment + var env = environment env["PATH"] = PathBuilder.effectivePATH( purposes: [.rpc, .nodeTooling], env: env) @@ -534,7 +535,7 @@ public struct UsageFetcher: Sendable { } private func loadRPCUsage() async throws -> UsageSnapshot { - let rpc = try CodexRPCClient() + let rpc = try CodexRPCClient(environment: self.environment) defer { rpc.shutdown() } try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4") @@ -568,7 +569,10 @@ public struct UsageFetcher: Sendable { } private func loadTTYUsage(keepCLISessionsAlive: Bool) async throws -> UsageSnapshot { - let status = try await CodexStatusProbe(keepCLISessionsAlive: keepCLISessionsAlive).fetch() + let status = try await CodexStatusProbe( + keepCLISessionsAlive: keepCLISessionsAlive, + environment: self.environment) + .fetch() guard let fiveLeft = status.fiveHourPercentLeft, let weekLeft = status.weeklyPercentLeft else { throw UsageError.noRateLimitsFound } @@ -599,7 +603,7 @@ public struct UsageFetcher: Sendable { } private func loadRPCCredits() async throws -> CreditsSnapshot { - let rpc = try CodexRPCClient() + let rpc = try CodexRPCClient(environment: self.environment) defer { rpc.shutdown() } try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4") let limits = try await rpc.fetchRateLimits().rateLimits @@ -609,7 +613,10 @@ public struct UsageFetcher: Sendable { } private func loadTTYCredits(keepCLISessionsAlive: Bool) async throws -> CreditsSnapshot { - let status = try await CodexStatusProbe(keepCLISessionsAlive: keepCLISessionsAlive).fetch() + let status = try await CodexStatusProbe( + keepCLISessionsAlive: keepCLISessionsAlive, + environment: self.environment) + .fetch() guard let credits = status.credits else { throw UsageError.noRateLimitsFound } return CreditsSnapshot(remaining: credits, events: [], updatedAt: Date()) } @@ -632,7 +639,7 @@ public struct UsageFetcher: Sendable { public func debugRawRateLimits() async -> String { do { - let rpc = try CodexRPCClient() + let rpc = try CodexRPCClient(environment: self.environment) defer { rpc.shutdown() } try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4") let limits = try await rpc.fetchRateLimits() @@ -645,11 +652,9 @@ public struct UsageFetcher: Sendable { public func loadAccountInfo() -> AccountInfo { // Keep using auth.json for quick startup (non-blocking, no RPC spin-up required). - let authURL = URL(fileURLWithPath: self.environment["CODEX_HOME"] ?? "\(NSHomeDirectory())/.codex") - .appendingPathComponent("auth.json") - guard let data = try? Data(contentsOf: authURL), - let auth = try? JSONDecoder().decode(AuthFile.self, from: data), - let idToken = auth.tokens?.idToken + guard let credentials = try? CodexOAuthCredentialsStore.load(env: self.environment), + let idToken = credentials.idToken, + !idToken.isEmpty else { return AccountInfo(email: nil, plan: nil) } @@ -704,9 +709,3 @@ public struct UsageFetcher: Sendable { return json } } - -/// Minimal auth.json struct preserved from previous implementation -private struct AuthFile: Decodable { - struct Tokens: Decodable { let idToken: String? } - let tokens: Tokens? -} diff --git a/Tests/CodexBarTests/AppDelegateTests.swift b/Tests/CodexBarTests/AppDelegateTests.swift index 4efe641b8..c8b784c9b 100644 --- a/Tests/CodexBarTests/AppDelegateTests.swift +++ b/Tests/CodexBarTests/AppDelegateTests.swift @@ -9,10 +9,12 @@ struct AppDelegateTests { func `builds status controller after launch`() { let appDelegate = AppDelegate() var factoryCalls = 0 + let managedCodexAccountCoordinator = ManagedCodexAccountCoordinator() // Install a test factory that records invocations without touching NSStatusBar. - StatusItemController.factory = { _, _, _, _, _ in + StatusItemController.factory = { _, _, _, _, _, receivedCoordinator in factoryCalls += 1 + #expect(receivedCoordinator === managedCodexAccountCoordinator) return DummyStatusController() } defer { StatusItemController.factory = StatusItemController.defaultFactory } @@ -26,7 +28,12 @@ struct AppDelegateTests { let account = fetcher.loadAccountInfo() // configure should not eagerly construct the status controller - appDelegate.configure(store: store, settings: settings, account: account, selection: PreferencesSelection()) + appDelegate.configure( + store: store, + settings: settings, + account: account, + selection: PreferencesSelection(), + managedCodexAccountCoordinator: managedCodexAccountCoordinator) #expect(factoryCalls == 0) // construction happens once after launch diff --git a/Tests/CodexBarTests/CLIWebFallbackTests.swift b/Tests/CodexBarTests/CLIWebFallbackTests.swift index 10fcf5396..d590103a2 100644 --- a/Tests/CodexBarTests/CLIWebFallbackTests.swift +++ b/Tests/CodexBarTests/CLIWebFallbackTests.swift @@ -62,6 +62,34 @@ struct CLIWebFallbackTests { context: context)) } + @Test + func `codex web strategy is unavailable when managed account store is unreadable`() async { + let context = self.makeContext(settings: ProviderSettingsSnapshot.make( + codex: .init( + usageDataSource: .auto, + cookieSource: .auto, + manualCookieHeader: nil, + managedAccountStoreUnreadable: true))) + let strategy = CodexWebDashboardStrategy() + let available = await strategy.isAvailable(context) + + #expect(!available) + } + + @Test + func `codex web strategy is unavailable when selected managed target is unavailable`() async { + let context = self.makeContext(settings: ProviderSettingsSnapshot.make( + codex: .init( + usageDataSource: .auto, + cookieSource: .auto, + manualCookieHeader: nil, + managedAccountTargetUnavailable: true))) + let strategy = CodexWebDashboardStrategy() + let available = await strategy.isAvailable(context) + + #expect(!available) + } + @Test func `claude falls back when no session key`() { let context = self.makeContext() diff --git a/Tests/CodexBarTests/CodexAccountReconciliationTests.swift b/Tests/CodexBarTests/CodexAccountReconciliationTests.swift new file mode 100644 index 000000000..904dd72ae --- /dev/null +++ b/Tests/CodexBarTests/CodexAccountReconciliationTests.swift @@ -0,0 +1,669 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@Suite(.serialized) +struct CodexAccountReconciliationTests { + @Test + @MainActor + func `settings store exposes codex reconciliation accessors using managed and live overrides`() throws { + let suite = "CodexAccountReconciliationTests-settings-store" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let managed = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let live = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings._test_activeManagedCodexAccount = managed + settings._test_liveSystemCodexAccount = live + settings.codexActiveSource = .managedAccount(id: managed.id) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let snapshot = settings.codexAccountReconciliationSnapshot + let projection = settings.codexVisibleAccountProjection + + #expect(settings.codexActiveSource == .managedAccount(id: managed.id)) + #expect(snapshot.storedAccounts.map(\.id) == [managed.id]) + #expect(snapshot.storedAccounts.map(\.email) == [managed.email]) + #expect(snapshot.activeStoredAccount?.id == managed.id) + #expect(snapshot.activeStoredAccount?.email == managed.email) + #expect(snapshot.liveSystemAccount == live) + #expect(snapshot.matchingStoredAccountForLiveSystemAccount == nil) + #expect(snapshot.activeSource == .managedAccount(id: managed.id)) + #expect(snapshot.hasUnreadableAddedAccountStore == false) + #expect(Set(projection.visibleAccounts.map(\.email)) == ["managed@example.com", "system@example.com"]) + #expect(settings.codexVisibleAccounts == projection.visibleAccounts) + #expect(projection.activeVisibleAccountID == "managed@example.com") + #expect(projection.liveVisibleAccountID == "system@example.com") + } + + @Test + @MainActor + func `settings store managed override does not leak ambient live system account`() throws { + let suite = "CodexAccountReconciliationTests-managed-only" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let managed = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + settings._test_activeManagedCodexAccount = managed + settings.codexActiveSource = .managedAccount(id: managed.id) + defer { + settings._test_activeManagedCodexAccount = nil + } + + let snapshot = settings.codexAccountReconciliationSnapshot + let projection = settings.codexVisibleAccountProjection + + #expect(settings.codexActiveSource == .managedAccount(id: managed.id)) + #expect(snapshot.liveSystemAccount == nil) + #expect(snapshot.matchingStoredAccountForLiveSystemAccount == nil) + #expect(snapshot.activeSource == .managedAccount(id: managed.id)) + #expect(projection.visibleAccounts.map(\.email) == ["managed@example.com"]) + #expect(projection.activeVisibleAccountID == "managed@example.com") + #expect(projection.liveVisibleAccountID == nil) + } + + @Test + @MainActor + func `settings store reconciliation environment override drives live observation with synthetic store`() throws { + let suite = "CodexAccountReconciliationTests-environment-only" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + try Self.writeCodexAuthFile(homeURL: ambientHome, email: "ambient@example.com", plan: "pro") + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": ambientHome.path] + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: ambientHome) + } + + let snapshot = settings.codexAccountReconciliationSnapshot + let projection = settings.codexVisibleAccountProjection + + #expect(settings.codexActiveSource == .liveSystem) + #expect(snapshot.storedAccounts.isEmpty) + #expect(snapshot.activeStoredAccount == nil) + #expect(snapshot.liveSystemAccount?.email == "ambient@example.com") + #expect(snapshot.liveSystemAccount?.codexHomePath == ambientHome.path) + #expect(snapshot.matchingStoredAccountForLiveSystemAccount == nil) + #expect(snapshot.activeSource == .liveSystem) + #expect(projection.visibleAccounts.map(\.email) == ["ambient@example.com"]) + #expect(projection.activeVisibleAccountID == "ambient@example.com") + #expect(projection.liveVisibleAccountID == "ambient@example.com") + } + + @Test + @MainActor + func `settings store home path override also keeps reconciliation hermetic`() throws { + let suite = "CodexAccountReconciliationTests-home-path-only" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings._test_activeManagedCodexRemoteHomePath = "/tmp/managed-route-home" + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = nil + defer { + settings._test_activeManagedCodexRemoteHomePath = nil + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = nil + } + + let snapshot = settings.codexAccountReconciliationSnapshot + let projection = settings.codexVisibleAccountProjection + + #expect(snapshot.storedAccounts.isEmpty) + #expect(snapshot.activeStoredAccount == nil) + #expect(snapshot.liveSystemAccount == nil) + #expect(snapshot.matchingStoredAccountForLiveSystemAccount == nil) + #expect(projection.visibleAccounts.isEmpty) + #expect(projection.activeVisibleAccountID == nil) + #expect(projection.liveVisibleAccountID == nil) + } + + @Test + @MainActor + func `settings store home path override keeps active source hermetic without persisted source`() throws { + let suite = "CodexAccountReconciliationTests-home-path-hermetic-source" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let ambient = ManagedCodexAccount( + id: UUID(), + email: "ambient-managed@example.com", + managedHomePath: "/tmp/ambient-managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let accounts = ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [ambient]) + let storeURL = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-managed-store-\(UUID().uuidString).json") + try Self.writeManagedCodexStore(accounts, to: storeURL) + + settings._test_managedCodexAccountStoreURL = storeURL + settings._test_activeManagedCodexRemoteHomePath = "/tmp/managed-route-home" + defer { + settings._test_managedCodexAccountStoreURL = nil + settings._test_activeManagedCodexRemoteHomePath = nil + try? FileManager.default.removeItem(at: storeURL) + } + + let snapshot = settings.codexAccountReconciliationSnapshot + + #expect(settings.codexActiveSource == .liveSystem) + #expect(settings.providerConfig(for: .codex)?.codexActiveSource == .liveSystem) + #expect(snapshot.storedAccounts.map(\.id) == [ambient.id]) + #expect(snapshot.storedAccounts.map(\.email) == [ambient.email]) + #expect(snapshot.activeStoredAccount == nil) + #expect(snapshot.activeSource == .liveSystem) + } + + @Test + @MainActor + func `settings store normal reconciliation path honors persisted active source`() throws { + let suite = "CodexAccountReconciliationTests-normal-path-active-source" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let persistedSource = CodexActiveSource.managedAccount(id: UUID()) + settings.codexActiveSource = persistedSource + + let snapshot = settings.codexAccountReconciliationSnapshot + + #expect(snapshot.activeSource == persistedSource) + } + + @Test + @MainActor + func `settings store debug managed store U R L override loads on disk accounts`() throws { + let suite = "CodexAccountReconciliationTests-debug-store-url" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let stored = ManagedCodexAccount( + id: UUID(), + email: "stored@example.com", + managedHomePath: "/tmp/stored-managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let accounts = ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [stored]) + let storeURL = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-managed-store-\(UUID().uuidString).json") + try Self.writeManagedCodexStore(accounts, to: storeURL) + + settings._test_managedCodexAccountStoreURL = storeURL + settings.codexActiveSource = .managedAccount(id: stored.id) + defer { + settings._test_managedCodexAccountStoreURL = nil + try? FileManager.default.removeItem(at: storeURL) + } + + let snapshot = settings.codexAccountReconciliationSnapshot + + #expect(snapshot.storedAccounts.map(\.id) == [stored.id]) + #expect(snapshot.storedAccounts.map(\.email) == [stored.email]) + #expect(snapshot.activeStoredAccount?.id == stored.id) + #expect(snapshot.activeStoredAccount?.email == stored.email) + #expect(snapshot.activeSource == .managedAccount(id: stored.id)) + } + + @Test + func `live only visible account is active when active source is live system`() { + let live = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let projection = CodexVisibleAccountProjection.make(from: CodexAccountReconciliationSnapshot( + storedAccounts: [], + activeStoredAccount: nil, + liveSystemAccount: live, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .liveSystem, + hasUnreadableAddedAccountStore: false)) + + #expect(projection.visibleAccounts.map(\.email) == ["live@example.com"]) + #expect(projection.activeVisibleAccountID == "live@example.com") + #expect(projection.liveVisibleAccountID == "live@example.com") + } + + @Test + func `matching live system account does not duplicate stored identity`() { + let stored = ManagedCodexAccount( + id: UUID(), + email: "user@example.com", + managedHomePath: "/tmp/managed-a", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let accounts = ManagedCodexAccountSet(version: 1, accounts: [stored]) + let live = ObservedSystemCodexAccount( + email: "USER@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let reconciler = DefaultCodexAccountReconciler( + storeLoader: { accounts }, + systemObserver: StubSystemObserver(account: live), + activeSource: .managedAccount(id: stored.id)) + + let projection = reconciler.loadVisibleAccounts(environment: [:]) + + #expect(projection.visibleAccounts.count == 1) + #expect(projection.activeVisibleAccountID == "user@example.com") + #expect(projection.liveVisibleAccountID == "user@example.com") + } + + @Test + func `matching live system account resolves merged row selection to live system`() { + let stored = ManagedCodexAccount( + id: UUID(), + email: "user@example.com", + managedHomePath: "/tmp/managed-a", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let live = ObservedSystemCodexAccount( + email: "USER@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let snapshot = CodexAccountReconciliationSnapshot( + storedAccounts: [stored], + activeStoredAccount: stored, + liveSystemAccount: live, + matchingStoredAccountForLiveSystemAccount: stored, + activeSource: .managedAccount(id: stored.id), + hasUnreadableAddedAccountStore: false) + + let resolution = CodexActiveSourceResolver.resolve(from: snapshot) + let projection = CodexVisibleAccountProjection.make(from: snapshot) + + #expect(resolution.persistedSource == .managedAccount(id: stored.id)) + #expect(resolution.resolvedSource == .liveSystem) + #expect(resolution.requiresPersistenceCorrection) + #expect(projection.activeVisibleAccountID == "user@example.com") + #expect(projection.source(forVisibleAccountID: "user@example.com") == .liveSystem) + } + + @Test + func `missing managed source resolves to live system when live account exists`() { + let live = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let missingID = UUID() + let snapshot = CodexAccountReconciliationSnapshot( + storedAccounts: [], + activeStoredAccount: nil, + liveSystemAccount: live, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .managedAccount(id: missingID), + hasUnreadableAddedAccountStore: false) + + let resolution = CodexActiveSourceResolver.resolve(from: snapshot) + + #expect(resolution.persistedSource == .managedAccount(id: missingID)) + #expect(resolution.resolvedSource == .liveSystem) + #expect(resolution.requiresPersistenceCorrection) + } + + @Test + func `unreadable managed source resolves to live system when live account exists`() { + let live = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let unreadableID = UUID() + let snapshot = CodexAccountReconciliationSnapshot( + storedAccounts: [], + activeStoredAccount: nil, + liveSystemAccount: live, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .managedAccount(id: unreadableID), + hasUnreadableAddedAccountStore: true) + + let resolution = CodexActiveSourceResolver.resolve(from: snapshot) + + #expect(resolution.persistedSource == .managedAccount(id: unreadableID)) + #expect(resolution.resolvedSource == .liveSystem) + #expect(resolution.requiresPersistenceCorrection) + } + + @Test + func `managed account remains active when active source stays managed while live account changes`() { + let managed = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let live = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let projection = CodexVisibleAccountProjection.make(from: CodexAccountReconciliationSnapshot( + storedAccounts: [managed], + activeStoredAccount: managed, + liveSystemAccount: live, + matchingStoredAccountForLiveSystemAccount: nil, + activeSource: .managedAccount(id: managed.id), + hasUnreadableAddedAccountStore: false)) + + #expect(Set(projection.visibleAccounts.map(\.email)) == [ + "managed@example.com", + "system@example.com", + ]) + #expect(projection.activeVisibleAccountID == "managed@example.com") + #expect(projection.liveVisibleAccountID == "system@example.com") + } + + @Test + func `live system account that differs from active stored account remains visible`() { + let active = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-a", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let accounts = ManagedCodexAccountSet(version: 1, accounts: [active]) + let live = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let reconciler = DefaultCodexAccountReconciler( + storeLoader: { accounts }, + systemObserver: StubSystemObserver(account: live), + activeSource: .managedAccount(id: active.id)) + + let projection = reconciler.loadVisibleAccounts(environment: [:]) + + #expect(Set(projection.visibleAccounts.map(\.email)) == ["managed@example.com", "system@example.com"]) + #expect(projection.activeVisibleAccountID == "managed@example.com") + #expect(projection.liveVisibleAccountID == "system@example.com") + } + + @Test + func `inactive stored account still appears as visible`() { + let active = ManagedCodexAccount( + id: UUID(), + email: "active@example.com", + managedHomePath: "/tmp/managed-a", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let inactive = ManagedCodexAccount( + id: UUID(), + email: "inactive@example.com", + managedHomePath: "/tmp/managed-b", + createdAt: 4, + updatedAt: 5, + lastAuthenticatedAt: 6) + let accounts = ManagedCodexAccountSet( + version: 1, + accounts: [active, inactive]) + let live = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let reconciler = DefaultCodexAccountReconciler( + storeLoader: { accounts }, + systemObserver: StubSystemObserver(account: live), + activeSource: .managedAccount(id: active.id)) + + let projection = reconciler.loadVisibleAccounts(environment: [:]) + + #expect(Set(projection.visibleAccounts.map(\.email)) == [ + "active@example.com", + "inactive@example.com", + "system@example.com", + ]) + #expect(projection.activeVisibleAccountID == "active@example.com") + #expect(projection.liveVisibleAccountID == "system@example.com") + } + + @Test + func `unreadable account store still exposes live system account and degraded flag`() { + let live = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let reconciler = DefaultCodexAccountReconciler( + storeLoader: { throw FileManagedCodexAccountStoreError.unsupportedVersion(999) }, + systemObserver: StubSystemObserver(account: live)) + + let projection = reconciler.loadVisibleAccounts(environment: [:]) + + #expect(projection.visibleAccounts.map(\.email) == ["live@example.com"]) + #expect(projection.activeVisibleAccountID == "live@example.com") + #expect(projection.liveVisibleAccountID == "live@example.com") + #expect(projection.hasUnreadableAddedAccountStore) + } + + @Test + func `whitespace only live email is ignored`() { + let accounts = ManagedCodexAccountSet(version: 1, accounts: []) + let live = ObservedSystemCodexAccount( + email: " \n\t ", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let reconciler = DefaultCodexAccountReconciler( + storeLoader: { accounts }, + systemObserver: StubSystemObserver(account: live)) + + let projection = reconciler.loadVisibleAccounts(environment: [:]) + + #expect(projection.visibleAccounts.isEmpty) + #expect(projection.activeVisibleAccountID == nil) + #expect(projection.liveVisibleAccountID == nil) + } + + @Test + @MainActor + func `settings store can override active source to live system`() throws { + let suite = "CodexAccountReconciliationTests-live-source-override" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let managed = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let live = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings._test_activeManagedCodexAccount = managed + settings._test_liveSystemCodexAccount = live + settings.codexActiveSource = .liveSystem + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let snapshot = settings.codexAccountReconciliationSnapshot + let projection = settings.codexVisibleAccountProjection + + #expect(settings.codexActiveSource == .liveSystem) + #expect(snapshot.activeSource == .liveSystem) + #expect(projection.activeVisibleAccountID == "system@example.com") + #expect(projection.liveVisibleAccountID == "system@example.com") + } + + @Test + @MainActor + func `selecting merged visible account persists live system source`() throws { + let suite = "CodexAccountReconciliationTests-select-merged-visible-account" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let managed = ManagedCodexAccount( + id: UUID(), + email: "same@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let live = ObservedSystemCodexAccount( + email: "SAME@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings._test_activeManagedCodexAccount = managed + settings._test_liveSystemCodexAccount = live + settings.codexActiveSource = .managedAccount(id: managed.id) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let didSelect = settings.selectCodexVisibleAccount(id: "same@example.com") + + #expect(didSelect) + #expect(settings.codexActiveSource == .liveSystem) + #expect(settings.codexResolvedActiveSource == .liveSystem) + } + + @Test + @MainActor + func `selecting authenticated managed account prefers live system when visible row is merged`() throws { + let suite = "CodexAccountReconciliationTests-select-authenticated-managed-merged" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + let managed = ManagedCodexAccount( + id: UUID(), + email: "same@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let live = ObservedSystemCodexAccount( + email: "SAME@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings._test_activeManagedCodexAccount = managed + settings._test_liveSystemCodexAccount = live + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + settings.selectAuthenticatedManagedCodexAccount(managed) + + #expect(settings.codexActiveSource == .liveSystem) + #expect(settings.codexResolvedActiveSource == .liveSystem) + } +} + +private struct StubSystemObserver: CodexSystemAccountObserving { + let account: ObservedSystemCodexAccount? + + func loadSystemAccount(environment _: [String: String]) throws -> ObservedSystemCodexAccount? { + self.account + } +} + +extension CodexAccountReconciliationTests { + private static func writeManagedCodexStore(_ accounts: ManagedCodexAccountSet, to storeURL: URL) throws { + let store = FileManagedCodexAccountStore(fileURL: storeURL) + try store.storeAccounts(accounts) + } + + 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))." + } +} diff --git a/Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift b/Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift new file mode 100644 index 000000000..303cac032 --- /dev/null +++ b/Tests/CodexBarTests/CodexAccountScopedRefreshTests.swift @@ -0,0 +1,748 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@Suite(.serialized) +@MainActor +struct CodexAccountScopedRefreshTests { + @Test + func `account transition invalidates codex scoped state and preserves token usage`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-invalidate") + settings.refreshFrequency = .manual + settings.codexCookieSource = .auto + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + let staleSnapshot = self.codexSnapshot(email: "alpha@example.com", usedPercent: 10) + let staleCredits = self.credits(remaining: 42) + let staleDashboard = self.dashboard(email: "alpha@example.com", creditsRemaining: 42, usedPercent: 20) + let tokenSnapshot = CostUsageTokenSnapshot( + sessionTokens: 120, + sessionCostUSD: 1.2, + last30DaysTokens: 900, + last30DaysCostUSD: 9.0, + daily: [], + updatedAt: Date()) + var widgetSnapshots: [WidgetSnapshot] = [] + + store._setSnapshotForTesting(staleSnapshot, provider: .codex) + store.credits = staleCredits + store.lastCreditsSnapshot = staleCredits + store.lastCreditsSnapshotAccountKey = "alpha@example.com" + store.openAIDashboard = staleDashboard + store.lastOpenAIDashboardSnapshot = staleDashboard + store.lastOpenAIDashboardTargetEmail = "alpha@example.com" + store._setTokenSnapshotForTesting(tokenSnapshot, provider: .codex) + store.lastCodexAccountScopedRefreshGuard = store + .currentCodexAccountScopedRefreshGuard(preferCurrentSnapshot: false) + store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) } + defer { store._test_widgetSnapshotSaveOverride = nil } + + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + + let didInvalidate = store.prepareCodexAccountScopedRefreshIfNeeded() + await store.widgetSnapshotPersistTask?.value + + #expect(didInvalidate) + #expect(store.snapshots[.codex] == nil) + #expect(store.credits == nil) + #expect(store.lastCreditsSnapshot == nil) + #expect(store.lastCreditsSnapshotAccountKey == nil) + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardSnapshot == nil) + #expect(store.tokenSnapshots[.codex] == tokenSnapshot) + #expect(widgetSnapshots.count == 1) + #expect(widgetSnapshots[0].entries.contains(where: { $0.provider == .codex }) == false) + } + + @Test + func `first switch invalidates after codex refresh seeds the previous account guard`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-first-switch") + settings.refreshFrequency = .manual + settings.codexCookieSource = .off + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + self.installImmediateCodexProvider( + on: store, + snapshot: self.codexSnapshot(email: "alpha@example.com", usedPercent: 10)) + + await store.refreshProvider(.codex, allowDisabled: true) + + #expect(store.lastCodexAccountScopedRefreshGuard?.accountKey == "alpha@example.com") + + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + + let didInvalidate = store.prepareCodexAccountScopedRefreshIfNeeded() + + #expect(didInvalidate) + #expect(store.snapshots[.codex] == nil) + } + + @Test + func `stale codex usage success is discarded after account switch`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-stale-success") + settings.refreshFrequency = .manual + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + let blocker = BlockingCodexFetchStrategy() + self.installBlockingCodexProvider(on: store, blocker: blocker) + + let refreshTask = Task { await store.refreshProvider(.codex, allowDisabled: true) } + await blocker.waitUntilStarted() + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + await blocker.resume(with: .success(self.codexSnapshot(email: "alpha@example.com", usedPercent: 25))) + await refreshTask.value + + #expect(store.snapshots[.codex] == nil) + #expect(store.errors[.codex] == nil) + } + + @Test + func `stale codex usage failure does not clear newer account snapshot`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-stale-failure") + settings.refreshFrequency = .manual + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + let blocker = BlockingCodexFetchStrategy() + self.installBlockingCodexProvider(on: store, blocker: blocker) + + let refreshTask = Task { await store.refreshProvider(.codex, allowDisabled: true) } + await blocker.waitUntilStarted() + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + let freshSnapshot = self.codexSnapshot(email: "beta@example.com", usedPercent: 5) + store._setSnapshotForTesting(freshSnapshot, provider: .codex) + await blocker.resume(with: .failure(TestRefreshError(message: "stale failure"))) + await refreshTask.value + + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "beta@example.com") + #expect(store.errors[.codex] == nil) + } + + @Test + func `credits fallback only reuses cache for the same codex account`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-credits") + settings.refreshFrequency = .manual + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + let cachedCredits = self.credits(remaining: 12) + store._setSnapshotForTesting(self.codexSnapshot(email: "alpha@example.com", usedPercent: 10), provider: .codex) + store.lastCreditsSnapshot = cachedCredits + store.lastCreditsSnapshotAccountKey = "alpha@example.com" + store._test_codexCreditsLoaderOverride = { + throw TestRefreshError(message: "Codex credits data not available yet") + } + defer { store._test_codexCreditsLoaderOverride = nil } + + await store.refreshCreditsIfNeeded() + #expect(store.credits == cachedCredits) + #expect(store.lastCreditsError == nil) + + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + store._setSnapshotForTesting(self.codexSnapshot(email: "beta@example.com", usedPercent: 10), provider: .codex) + + await store.refreshCreditsIfNeeded() + #expect(store.credits == nil) + #expect(store.lastCreditsError == "Codex credits are still loading; will retry shortly.") + } + + @Test + func `credits refresh returns quickly when no live codex account is available`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-credits-no-live-account") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-credits-no-live-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + settings.refreshFrequency = .manual + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = self.makeUsageStore(settings: settings) + var loaderCalled = false + store._test_codexCreditsLoaderOverride = { + loaderCalled = true + return self.credits(remaining: 1) + } + defer { store._test_codexCreditsLoaderOverride = nil } + + let startedAt = ContinuousClock.now + await store.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: Date()) + let elapsed = startedAt.duration(to: .now) + + #expect(loaderCalled == false) + #expect(elapsed < .seconds(2)) + } + + @Test + func `stale dashboard apply is discarded after account switch`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-dashboard") + settings.refreshFrequency = .manual + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + let expectedGuard = store.currentCodexAccountScopedRefreshGuard() + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + + await store.applyOpenAIDashboard( + self.dashboard(email: "alpha@example.com", creditsRemaining: 11, usedPercent: 35), + targetEmail: "alpha@example.com", + expectedGuard: expectedGuard, + allowCodexUsageBackfill: true) + + #expect(store.openAIDashboard == nil) + #expect(store.snapshots[.codex] == nil) + #expect(store.credits == nil) + } + + @Test + func `dashboard refresh can seed unknown live codex account`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-dashboard-seed-unknown-live") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-seed-unknown-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + settings.refreshFrequency = .manual + settings.codexCookieSource = .auto + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = self.makeUsageStore(settings: settings) + store.lastKnownLiveSystemCodexEmail = nil + store._test_openAIDashboardLoaderOverride = { _, _, _ in + self.dashboard(email: "seeded@example.com", creditsRemaining: 33, usedPercent: 12) + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let expectedGuard = store.currentCodexAccountScopedRefreshGuard() + #expect(expectedGuard.source == .liveSystem) + #expect(expectedGuard.accountKey == nil) + + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + + #expect(store.openAIDashboard?.signedInEmail == "seeded@example.com") + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "seeded@example.com") + #expect(store.credits?.remaining == 33) + #expect(store.lastCreditsSnapshotAccountKey == "seeded@example.com") + #expect(store.lastKnownLiveSystemCodexEmail == "seeded@example.com") + #expect(store.lastCodexAccountScopedRefreshGuard?.accountKey == "seeded@example.com") + } + + @Test + func `dashboard refresh rejects stale completion during live account reconciliation lag`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-dashboard-reject-stale-live-lag") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-stale-live-lag-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + settings.refreshFrequency = .manual + settings.codexCookieSource = .auto + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + settings.codexActiveSource = .liveSystem + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = self.makeUsageStore(settings: settings) + store._setSnapshotForTesting(self.codexSnapshot(email: "alpha@example.com", usedPercent: 12), provider: .codex) + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + #expect(expectedGuard.accountKey == nil) + + store._setSnapshotForTesting(self.codexSnapshot(email: "beta@example.com", usedPercent: 18), provider: .codex) + + await store.applyOpenAIDashboard( + self.dashboard(email: "alpha@example.com", creditsRemaining: 40, usedPercent: 20), + targetEmail: nil, + expectedGuard: expectedGuard, + allowCodexUsageBackfill: true) + + #expect(store.openAIDashboard == nil) + #expect(store.credits == nil) + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "beta@example.com") + } + + @Test + func `default dashboard refresh path discards stale completion after account switch`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-dashboard-guard") + settings.refreshFrequency = .manual + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + self.installImmediateCodexProvider( + on: store, + snapshot: self.codexSnapshot(email: "alpha@example.com", usedPercent: 18)) + let dashboardBlocker = BlockingOpenAIDashboardLoader() + store._test_openAIDashboardLoaderOverride = { _, _, _ in + try await dashboardBlocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let refreshTask = Task { await store.refresh() } + await dashboardBlocker.waitUntilStarted() + + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + store._setSnapshotForTesting(self.codexSnapshot(email: "beta@example.com", usedPercent: 7), provider: .codex) + store.openAIDashboard = nil + store.credits = nil + + await dashboardBlocker.resume(with: .success( + self.dashboard(email: "alpha@example.com", creditsRemaining: 44, usedPercent: 21))) + await refreshTask.value + + #expect(store.openAIDashboard == nil) + #expect(store.credits == nil) + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "beta@example.com") + } + + @Test + func `live switch invalidates stale codex state even when only last known live email remains`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-invalidate-with-stale-last-known") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-invalidate-stale-last-known-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + settings.refreshFrequency = .manual + settings.codexCookieSource = .auto + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = self.makeUsageStore(settings: settings) + store._setSnapshotForTesting(self.codexSnapshot(email: "alpha@example.com", usedPercent: 10), provider: .codex) + store.credits = self.credits(remaining: 12) + store.lastCreditsSnapshot = self.credits(remaining: 12) + store.lastCreditsSnapshotAccountKey = "alpha@example.com" + store.openAIDashboard = self.dashboard(email: "alpha@example.com", creditsRemaining: 12, usedPercent: 20) + store.lastOpenAIDashboardSnapshot = store.openAIDashboard + store.lastKnownLiveSystemCodexEmail = "alpha@example.com" + store.lastCodexAccountScopedRefreshGuard = store.currentCodexAccountScopedRefreshGuard( + preferCurrentSnapshot: false, + allowLastKnownLiveFallback: true) + + settings._test_liveSystemCodexAccount = nil + store.snapshots.removeValue(forKey: .codex) + + let didInvalidate = store.prepareCodexAccountScopedRefreshIfNeeded() + await store.widgetSnapshotPersistTask?.value + + #expect(didInvalidate) + #expect(store.snapshots[.codex] == nil) + #expect(store.credits == nil) + #expect(store.openAIDashboard == nil) + #expect(store.lastCodexAccountScopedRefreshGuard?.accountKey == nil) + } + + @Test + func `codex account refresh persists widget snapshots on invalidation and completion`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-widgets") + settings.refreshFrequency = .manual + settings.codexCookieSource = .off + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + store._setSnapshotForTesting(self.codexSnapshot(email: "alpha@example.com", usedPercent: 18), provider: .codex) + store.lastCodexAccountScopedRefreshGuard = store + .currentCodexAccountScopedRefreshGuard(preferCurrentSnapshot: false) + + let blocker = BlockingCodexFetchStrategy() + self.installBlockingCodexProvider(on: store, blocker: blocker) + store._test_codexCreditsLoaderOverride = { self.credits(remaining: 77) } + defer { store._test_codexCreditsLoaderOverride = nil } + + var widgetSnapshots: [WidgetSnapshot] = [] + store._test_widgetSnapshotSaveOverride = { widgetSnapshots.append($0) } + defer { store._test_widgetSnapshotSaveOverride = nil } + + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + let refreshTask = Task { await store.refreshCodexAccountScopedState(allowDisabled: true) } + await blocker.waitUntilStarted() + await blocker.resume(with: .success(self.codexSnapshot(email: "beta@example.com", usedPercent: 8))) + await refreshTask.value + await store.widgetSnapshotPersistTask?.value + + #expect(widgetSnapshots.count == 2) + #expect(widgetSnapshots[0].entries.contains(where: { $0.provider == .codex }) == false) + #expect(widgetSnapshots[1].entries.first { $0.provider == .codex }?.creditsRemaining == 77) + } + + @Test + func `widget snapshot saves stay ordered across codex account invalidation and completion`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-widget-order") + settings.refreshFrequency = .manual + + let store = self.makeUsageStore(settings: settings) + let saver = BlockingWidgetSnapshotSaver() + store._test_widgetSnapshotSaveOverride = { snapshot in + await saver.save(snapshot) + } + defer { store._test_widgetSnapshotSaveOverride = nil } + + store.persistWidgetSnapshot(reason: "codex-account-invalidate") + await saver.waitUntilStarted(count: 1) + #expect(await saver.startedCount() == 1) + + store._setSnapshotForTesting(self.codexSnapshot(email: "beta@example.com", usedPercent: 8), provider: .codex) + store.credits = self.credits(remaining: 77) + store.persistWidgetSnapshot(reason: "codex-account-refresh") + + try? await Task.sleep(nanoseconds: 50_000_000) + #expect(await saver.startedCount() == 1) + + await saver.resumeNext() + await saver.waitUntilStarted(count: 2) + await saver.resumeNext() + await store.widgetSnapshotPersistTask?.value + + let snapshots = await saver.savedSnapshots() + #expect(snapshots.count == 2) + #expect(snapshots[0].entries.contains(where: { $0.provider == .codex }) == false) + #expect(snapshots[1].entries.first { $0.provider == .codex }?.creditsRemaining == 77) + } + + @Test + func `codex account refresh reports usage and credits phases before completion`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-phases") + settings.refreshFrequency = .manual + settings.codexCookieSource = .off + settings._test_liveSystemCodexAccount = self.liveAccount(email: "alpha@example.com") + + let store = self.makeUsageStore(settings: settings) + store._setSnapshotForTesting(self.codexSnapshot(email: "alpha@example.com", usedPercent: 18), provider: .codex) + store.lastCodexAccountScopedRefreshGuard = store + .currentCodexAccountScopedRefreshGuard(preferCurrentSnapshot: false) + + let blocker = BlockingCodexFetchStrategy() + self.installBlockingCodexProvider(on: store, blocker: blocker) + store._test_codexCreditsLoaderOverride = { self.credits(remaining: 77) } + defer { store._test_codexCreditsLoaderOverride = nil } + + var phases: [CodexAccountScopedRefreshPhase] = [] + settings._test_liveSystemCodexAccount = self.liveAccount(email: "beta@example.com") + + let refreshTask = Task { + await store.refreshCodexAccountScopedState( + allowDisabled: true, + phaseDidChange: { phases.append($0) }) + } + + await blocker.waitUntilStarted() + #expect(phases == [.invalidated]) + + await blocker.resume(with: .success(self.codexSnapshot(email: "beta@example.com", usedPercent: 8))) + await refreshTask.value + + #expect(phases == [.invalidated, .usage, .credits, .completed]) + } + + @Test + func `refresh loads credits when codex email is discovered by usage in the same cycle`() async { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-refresh-credits") + settings.refreshFrequency = .manual + settings.codexCookieSource = .off + + let store = self.makeUsageStore(settings: settings) + let blocker = BlockingCodexFetchStrategy() + self.installBlockingCodexProvider(on: store, blocker: blocker) + store._test_codexCreditsLoaderOverride = { self.credits(remaining: 55) } + defer { store._test_codexCreditsLoaderOverride = nil } + + let refreshTask = Task { await store.refresh() } + await blocker.waitUntilStarted() + await blocker.resume(with: .success(self.codexSnapshot(email: "alpha@example.com", usedPercent: 12))) + await refreshTask.value + + #expect(store.credits?.remaining == 55) + #expect(store.lastCodexAccountScopedRefreshGuard?.accountKey == "alpha@example.com") + } + + @Test + func `settings codex account selection refreshes credits on the first switch`() async throws { + let settings = self.makeSettingsStore(suite: "CodexAccountScopedRefreshTests-settings-selection") + settings.refreshFrequency = .manual + settings.codexCookieSource = .off + let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let managedAccount = ManagedCodexAccount( + id: managedAccountID, + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) + defer { + settings._test_managedCodexAccountStoreURL = nil + settings._test_liveSystemCodexAccount = nil + try? FileManager.default.removeItem(at: storeURL) + } + + settings._test_managedCodexAccountStoreURL = storeURL + settings._test_liveSystemCodexAccount = self.liveAccount(email: "live@example.com") + + let store = self.makeUsageStore(settings: settings) + store._setSnapshotForTesting(self.codexSnapshot(email: "live@example.com", usedPercent: 30), provider: .codex) + store.lastCodexAccountScopedRefreshGuard = store + .currentCodexAccountScopedRefreshGuard(preferCurrentSnapshot: false) + self.installImmediateCodexProvider( + on: store, + snapshot: self.codexSnapshot(email: "managed@example.com", usedPercent: 9)) + store._test_codexCreditsLoaderOverride = { self.credits(remaining: 55) } + defer { store._test_codexCreditsLoaderOverride = nil } + + let pane = ProvidersPane(settings: settings, store: store) + await pane._test_selectCodexVisibleAccount(id: "managed@example.com") + + #expect(settings.codexActiveSource == .managedAccount(id: managedAccountID)) + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "managed@example.com") + #expect(store.credits?.remaining == 55) + } + + private func makeSettingsStore(suite: String) -> SettingsStore { + let defaults = UserDefaults(suiteName: suite)! + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings._test_activeManagedCodexAccount = nil + settings._test_activeManagedCodexRemoteHomePath = nil + settings._test_unreadableManagedCodexAccountStore = false + settings._test_managedCodexAccountStoreURL = nil + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = nil + return settings + } + + private func makeUsageStore(settings: SettingsStore) -> UsageStore { + UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + } + + private func liveAccount(email: String) -> ObservedSystemCodexAccount { + ObservedSystemCodexAccount( + email: email, + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + } + + private func codexSnapshot(email: String, usedPercent: Double) -> UsageSnapshot { + UsageSnapshot( + primary: RateWindow(usedPercent: usedPercent, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: email, + accountOrganization: nil, + loginMethod: "Pro")) + } + + private func credits(remaining: Double) -> CreditsSnapshot { + CreditsSnapshot(remaining: remaining, events: [], updatedAt: Date()) + } + + private func dashboard(email: String, creditsRemaining: Double, usedPercent: Double) -> OpenAIDashboardSnapshot { + OpenAIDashboardSnapshot( + signedInEmail: email, + codeReviewRemainingPercent: 88, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + primaryLimit: RateWindow( + usedPercent: usedPercent, + windowMinutes: 300, + resetsAt: nil, + resetDescription: nil), + secondaryLimit: nil, + creditsRemaining: creditsRemaining, + accountPlan: "Pro", + updatedAt: Date()) + } + + private func makeManagedAccountStoreURL(accounts: [ManagedCodexAccount]) throws -> URL { + let storeURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + let store = FileManagedCodexAccountStore(fileURL: storeURL) + try store.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: accounts)) + return storeURL + } + + private func installBlockingCodexProvider(on store: UsageStore, blocker: BlockingCodexFetchStrategy) { + let baseSpec = store.providerSpecs[.codex]! + store.providerSpecs[.codex] = Self.makeCodexProviderSpec(baseSpec: baseSpec) { + try await blocker.awaitResult() + } + } + + private func installImmediateCodexProvider(on store: UsageStore, snapshot: UsageSnapshot) { + let baseSpec = store.providerSpecs[.codex]! + store.providerSpecs[.codex] = Self.makeCodexProviderSpec(baseSpec: baseSpec) { + snapshot + } + } + + private static func makeCodexProviderSpec( + baseSpec: ProviderSpec, + loader: @escaping @Sendable () async throws -> UsageSnapshot) -> ProviderSpec + { + let baseDescriptor = baseSpec.descriptor + let strategy = TestCodexFetchStrategy(loader: loader) + let descriptor = ProviderDescriptor( + id: .codex, + metadata: baseDescriptor.metadata, + branding: baseDescriptor.branding, + tokenCost: baseDescriptor.tokenCost, + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli, .oauth], + pipeline: ProviderFetchPipeline { _ in [strategy] }), + cli: baseDescriptor.cli) + return ProviderSpec( + style: baseSpec.style, + isEnabled: baseSpec.isEnabled, + descriptor: descriptor, + makeFetchContext: baseSpec.makeFetchContext) + } +} + +private struct TestRefreshError: LocalizedError, Equatable { + let message: String + + var errorDescription: String? { + self.message + } +} + +private struct TestCodexFetchStrategy: ProviderFetchStrategy { + let loader: @Sendable () async throws -> UsageSnapshot + + var id: String { + "test-codex" + } + + var kind: ProviderFetchKind { + .cli + } + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try await self.loader() + return self.makeResult(usage: snapshot, sourceLabel: "test-codex") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} + +private actor BlockingCodexFetchStrategy { + private var waiters: [CheckedContinuation, Never>] = [] + private var startedWaiters: [CheckedContinuation] = [] + private var didStart = false + + func awaitResult() async throws -> UsageSnapshot { + self.didStart = true + self.startedWaiters.forEach { $0.resume() } + self.startedWaiters.removeAll() + let result = await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + return try result.get() + } + + func waitUntilStarted() async { + if self.didStart { return } + await withCheckedContinuation { continuation in + self.startedWaiters.append(continuation) + } + } + + func resume(with result: Result) { + self.waiters.forEach { $0.resume(returning: result) } + self.waiters.removeAll() + } +} + +private actor BlockingOpenAIDashboardLoader { + private var waiters: [CheckedContinuation, Never>] = [] + private var startedWaiters: [CheckedContinuation] = [] + private var didStart = false + + func awaitResult() async throws -> OpenAIDashboardSnapshot { + self.didStart = true + self.startedWaiters.forEach { $0.resume() } + self.startedWaiters.removeAll() + let result = await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + return try result.get() + } + + func waitUntilStarted() async { + if self.didStart { return } + await withCheckedContinuation { continuation in + self.startedWaiters.append(continuation) + } + } + + func resume(with result: Result) { + self.waiters.forEach { $0.resume(returning: result) } + self.waiters.removeAll() + } +} + +private actor BlockingWidgetSnapshotSaver { + private var snapshots: [WidgetSnapshot] = [] + private var waiters: [CheckedContinuation] = [] + private var startedWaiters: [CheckedContinuation] = [] + + func save(_ snapshot: WidgetSnapshot) async { + self.snapshots.append(snapshot) + self.startedWaiters.forEach { $0.resume() } + self.startedWaiters.removeAll() + await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + } + + func waitUntilStarted(count: Int) async { + if self.snapshots.count >= count { return } + await withCheckedContinuation { continuation in + self.startedWaiters.append(continuation) + } + } + + func startedCount() -> Int { + self.snapshots.count + } + + func resumeNext() { + guard !self.waiters.isEmpty else { return } + let waiter = self.waiters.removeFirst() + waiter.resume() + } + + func savedSnapshots() -> [WidgetSnapshot] { + self.snapshots + } +} diff --git a/Tests/CodexBarTests/CodexAccountsSettingsSectionTests.swift b/Tests/CodexBarTests/CodexAccountsSettingsSectionTests.swift new file mode 100644 index 000000000..2b78dc0ea --- /dev/null +++ b/Tests/CodexBarTests/CodexAccountsSettingsSectionTests.swift @@ -0,0 +1,324 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@MainActor +struct CodexAccountsSettingsSectionTests { + @Test + func `codex accounts section shows live badge only for live only multi account row`() throws { + let settings = Self.makeSettingsStore(suite: "CodexAccountsSettingsSectionTests-live-badge") + let store = Self.makeUsageStore(settings: settings) + let managedStoreURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: managedStoreURL) } + + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let managedStore = FileManagedCodexAccountStore(fileURL: managedStoreURL) + try managedStore.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [managedAccount])) + + settings._test_managedCodexAccountStoreURL = managedStoreURL + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + + let pane = ProvidersPane(settings: settings, store: store) + let state = try #require(pane._test_codexAccountsSectionState()) + let liveAccount = try #require(state.visibleAccounts.first { $0.email == "live@example.com" }) + let managedVisibleAccount = try #require(state.visibleAccounts.first { $0.email == "managed@example.com" }) + + #expect(state.showsLiveBadge(for: liveAccount)) + #expect(state.showsLiveBadge(for: managedVisibleAccount) == false) + } + + @Test + func `single account codex settings uses simple account view instead of picker`() throws { + let settings = Self.makeSettingsStore(suite: "CodexAccountsSettingsSectionTests-single-account") + let store = Self.makeUsageStore(settings: settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "solo@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + + let pane = ProvidersPane(settings: settings, store: store) + let state = try #require(pane._test_codexAccountsSectionState()) + + #expect(state.visibleAccounts.count == 1) + #expect(state.showsActivePicker == false) + #expect(state.singleVisibleAccount?.email == "solo@example.com") + } + + @Test + func `codex accounts section disables managed mutations when store is unreadable`() throws { + let settings = Self.makeSettingsStore(suite: "CodexAccountsSettingsSectionTests-unreadable") + let store = Self.makeUsageStore(settings: settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings._test_unreadableManagedCodexAccountStore = true + defer { settings._test_unreadableManagedCodexAccountStore = false } + + let pane = ProvidersPane(settings: settings, store: store) + let state = try #require(pane._test_codexAccountsSectionState()) + let liveAccount = try #require(state.visibleAccounts.first) + + #expect(state.hasUnreadableManagedAccountStore) + #expect(state.canAddAccount == false) + #expect(state.notice?.tone == .warning) + #expect(state.canReauthenticate(liveAccount)) + } + + @Test + func `selecting merged visible account from settings keeps live system source`() async throws { + let settings = Self.makeSettingsStore(suite: "CodexAccountsSettingsSectionTests-select-merged") + let store = Self.makeUsageStore(settings: settings) + let managedStoreURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: managedStoreURL) } + + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "same@example.com", + managedHomePath: "/tmp/managed", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let managedStore = FileManagedCodexAccountStore(fileURL: managedStoreURL) + try managedStore.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [managedAccount])) + + settings._test_managedCodexAccountStoreURL = managedStoreURL + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "SAME@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + + let pane = ProvidersPane(settings: settings, store: store) + await pane._test_selectCodexVisibleAccount(id: "same@example.com") + + #expect(settings.codexActiveSource == .liveSystem) + let state = try #require(pane._test_codexAccountsSectionState()) + #expect(state.activeVisibleAccountID == "same@example.com") + } + + @Test + func `codex accounts section disables add and reauth while managed authentication is in flight`() async throws { + let settings = Self.makeSettingsStore(suite: "CodexAccountsSettingsSectionTests-in-flight") + let store = Self.makeUsageStore(settings: settings) + let managedStoreURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: managedStoreURL) } + + let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let managedAccount = ManagedCodexAccount( + id: managedAccountID, + email: "managed@example.com", + managedHomePath: "/tmp/managed", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let managedStore = FileManagedCodexAccountStore(fileURL: managedStoreURL) + try managedStore.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [managedAccount])) + settings._test_managedCodexAccountStoreURL = managedStoreURL + + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + let runner = BlockingManagedCodexLoginRunnerForSettingsSectionTests() + let service = ManagedCodexAccountService( + store: managedStore, + homeFactory: TestManagedCodexHomeFactoryForSettingsSectionTests(root: root), + loginRunner: runner, + identityReader: StubManagedCodexIdentityReaderForSettingsSectionTests(emails: ["managed@example.com"])) + let coordinator = ManagedCodexAccountCoordinator(service: service) + let authTask = Task { try await coordinator.authenticateManagedAccount() } + await runner.waitUntilStarted() + + let pane = ProvidersPane( + settings: settings, + store: store, + managedCodexAccountCoordinator: coordinator) + let state = try #require(pane._test_codexAccountsSectionState()) + let visibleAccount = try #require(state.visibleAccounts.first { $0.email == "managed@example.com" }) + + #expect(state.canAddAccount == false) + #expect(state.addAccountTitle == "Adding Account…") + #expect(state.canReauthenticate(visibleAccount) == false) + + await runner.resume() + _ = try await authTask.value + } + + @Test + func `adding managed codex account auto selects the merged live row`() async throws { + let settings = Self.makeSettingsStore(suite: "CodexAccountsSettingsSectionTests-add-merged") + let store = Self.makeUsageStore(settings: settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "same@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + + let coordinator = Self.makeManagedCoordinator(settings: settings, email: "same@example.com") + let pane = ProvidersPane( + settings: settings, + store: store, + managedCodexAccountCoordinator: coordinator) + + await pane._test_addManagedCodexAccount() + + #expect(settings.codexActiveSource == .liveSystem) + let state = try #require(pane._test_codexAccountsSectionState()) + #expect(state.activeVisibleAccountID == "same@example.com") + } + + @Test + func `adding managed codex account selects the new managed account when email differs`() async throws { + let settings = Self.makeSettingsStore(suite: "CodexAccountsSettingsSectionTests-add-managed") + let store = Self.makeUsageStore(settings: settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + + let coordinator = Self.makeManagedCoordinator(settings: settings, email: "managed@example.com") + let pane = ProvidersPane( + settings: settings, + store: store, + managedCodexAccountCoordinator: coordinator) + + await pane._test_addManagedCodexAccount() + + guard case .managedAccount = settings.codexActiveSource else { + Issue.record("Expected the new managed account to become active") + return + } + let state = try #require(pane._test_codexAccountsSectionState()) + #expect(state.activeVisibleAccountID == "managed@example.com") + } + + private static func makeManagedCoordinator( + settings: SettingsStore, + email: String) + -> ManagedCodexAccountCoordinator + { + let storeURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + let store = FileManagedCodexAccountStore(fileURL: storeURL) + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + settings._test_managedCodexAccountStoreURL = storeURL + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactoryForSettingsSectionTests(root: root), + loginRunner: StubManagedCodexLoginRunnerForSettingsSectionTests.success, + identityReader: StubManagedCodexIdentityReaderForSettingsSectionTests(emails: [email])) + return ManagedCodexAccountCoordinator(service: service) + } + + private static func makeSettingsStore(suite: String) -> SettingsStore { + let defaults = UserDefaults(suiteName: suite)! + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + + return SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore(), + codexCookieStore: InMemoryCookieHeaderStore(), + claudeCookieStore: InMemoryCookieHeaderStore(), + cursorCookieStore: InMemoryCookieHeaderStore(), + opencodeCookieStore: InMemoryCookieHeaderStore(), + factoryCookieStore: InMemoryCookieHeaderStore(), + minimaxCookieStore: InMemoryMiniMaxCookieStore(), + minimaxAPITokenStore: InMemoryMiniMaxAPITokenStore(), + kimiTokenStore: InMemoryKimiTokenStore(), + kimiK2TokenStore: InMemoryKimiK2TokenStore(), + augmentCookieStore: InMemoryCookieHeaderStore(), + ampCookieStore: InMemoryCookieHeaderStore(), + copilotTokenStore: InMemoryCopilotTokenStore(), + tokenAccountStore: InMemoryTokenAccountStore()) + } + + private static func makeUsageStore(settings: SettingsStore) -> UsageStore { + UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + } +} + +private struct TestManagedCodexHomeFactoryForSettingsSectionTests: ManagedCodexHomeProducing, Sendable { + let root: URL + private let nextID = UUID().uuidString + + func makeHomeURL() -> URL { + self.root.appendingPathComponent(self.nextID, isDirectory: true) + } + + func validateManagedHomeForDeletion(_ url: URL) throws { + try ManagedCodexHomeFactory(root: self.root).validateManagedHomeForDeletion(url) + } +} + +private struct StubManagedCodexLoginRunnerForSettingsSectionTests: ManagedCodexLoginRunning, Sendable { + let result: CodexLoginRunner.Result + + func run(homePath _: String, timeout _: TimeInterval) async -> CodexLoginRunner.Result { + self.result + } + + static let success = StubManagedCodexLoginRunnerForSettingsSectionTests( + result: CodexLoginRunner.Result(outcome: .success, output: "ok")) +} + +private actor BlockingManagedCodexLoginRunnerForSettingsSectionTests: ManagedCodexLoginRunning { + private var waiters: [CheckedContinuation] = [] + private var startedWaiters: [CheckedContinuation] = [] + private var didStart = false + + func run(homePath _: String, timeout _: TimeInterval) async -> CodexLoginRunner.Result { + self.didStart = true + self.startedWaiters.forEach { $0.resume() } + self.startedWaiters.removeAll() + return await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + } + + func waitUntilStarted() async { + if self.didStart { return } + await withCheckedContinuation { continuation in + self.startedWaiters.append(continuation) + } + } + + func resume() { + let result = CodexLoginRunner.Result(outcome: .success, output: "ok") + self.waiters.forEach { $0.resume(returning: result) } + self.waiters.removeAll() + } +} + +private final class StubManagedCodexIdentityReaderForSettingsSectionTests: ManagedCodexIdentityReading, +@unchecked Sendable { + private var emails: [String] + + init(emails: [String]) { + self.emails = emails + } + + func loadAccountInfo(homePath _: String) throws -> AccountInfo { + let email = self.emails.isEmpty ? nil : self.emails.removeFirst() + return AccountInfo(email: email, plan: "Pro") + } +} diff --git a/Tests/CodexBarTests/CodexActiveSourceConfigTests.swift b/Tests/CodexBarTests/CodexActiveSourceConfigTests.swift new file mode 100644 index 000000000..85ef3a406 --- /dev/null +++ b/Tests/CodexBarTests/CodexActiveSourceConfigTests.swift @@ -0,0 +1,98 @@ +import CodexBarCore +import Foundation +import Testing + +@Suite(.serialized) +struct CodexActiveSourceConfigTests { + @Test + func `legacy config without codex active source decodes to nil`() throws { + let legacyJSON = """ + { + "version": 1, + "providers": [ + { + "id": "codex" + } + ] + } + """ + + let decoded = try JSONDecoder().decode( + CodexBarConfig.self, + from: Data(legacyJSON.utf8)) + + #expect(decoded.providerConfig(for: .codex)?.codexActiveSource == nil) + } + + @Test + func `provider config encodes live system active source with expected schema`() throws { + let config = CodexBarConfig( + providers: [ + ProviderConfig( + id: .codex, + codexActiveSource: .liveSystem), + ]) + + let data = try JSONEncoder().encode(config) + let object = try JSONSerialization.jsonObject(with: data) as? [String: Any] + let providers = try #require(object?["providers"] as? [[String: Any]]) + let provider = try #require(providers.first(where: { $0["id"] as? String == "codex" })) + let activeSource = try #require(provider["codexActiveSource"] as? [String: Any]) + + #expect(activeSource.count == 1) + #expect(activeSource["kind"] as? String == "liveSystem") + #expect(activeSource["accountID"] == nil) + } + + @Test + func `provider config encodes managed account active source with expected schema`() throws { + let accountID = UUID() + let config = CodexBarConfig( + providers: [ + ProviderConfig( + id: .codex, + codexActiveSource: .managedAccount(id: accountID)), + ]) + + let data = try JSONEncoder().encode(config) + let object = try JSONSerialization.jsonObject(with: data) as? [String: Any] + let providers = try #require(object?["providers"] as? [[String: Any]]) + let provider = try #require(providers.first(where: { $0["id"] as? String == "codex" })) + let activeSource = try #require(provider["codexActiveSource"] as? [String: Any]) + + #expect(activeSource.count == 2) + #expect(activeSource["kind"] as? String == "managedAccount") + #expect((activeSource["accountID"] as? String) == accountID.uuidString) + } + + @Test + func `provider config round trips live system active source`() throws { + let config = CodexBarConfig( + providers: [ + ProviderConfig( + id: .codex, + codexActiveSource: .liveSystem), + ]) + + let data = try JSONEncoder().encode(config) + let decoded = try JSONDecoder().decode(CodexBarConfig.self, from: data) + + #expect(decoded.providerConfig(for: .codex)?.codexActiveSource == .liveSystem) + } + + @Test + func `provider config round trips managed account active source`() throws { + let accountID = UUID() + let config = CodexBarConfig( + providers: [ + ProviderConfig( + id: .codex, + codexActiveSource: .managedAccount(id: accountID)), + ]) + + let data = try JSONEncoder().encode(config) + let decoded = try JSONDecoder().decode(CodexBarConfig.self, from: data) + + #expect(decoded.providerConfig(for: .codex)?.codexActiveSource == .managedAccount(id: accountID)) + } +} diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift new file mode 100644 index 000000000..0703aec3f --- /dev/null +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebRefreshTests.swift @@ -0,0 +1,320 @@ +import Foundation +import Testing +@testable import CodexBar +@testable import CodexBarCore + +@Suite(.serialized) +@MainActor +struct CodexManagedOpenAIWebRefreshTests { + @Test + func `manual cookie import bypasses same account refresh coalescing`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebRefreshTests-manual-import-bypass-coalesce") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingManagedOpenAIDashboardLoader() + store._test_openAIDashboardLoaderOverride = { _, _, _ in + try await blocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + store._test_openAIDashboardCookieImportOverride = { targetEmail, _, _, _, _ in + OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Chrome", + cookieCount: 2, + signedInEmail: targetEmail, + matchesCodexEmail: true) + } + defer { store._test_openAIDashboardCookieImportOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + let firstTask = Task { + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + } + await blocker.waitUntilStarted(count: 1) + + let manualImportTask = Task { + await store.importOpenAIDashboardBrowserCookiesNow() + } + await blocker.waitUntilStarted(count: 2) + + await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 70, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 1, + accountPlan: "Free", + updatedAt: Date()))) + await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 95, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 25, + accountPlan: "Pro", + updatedAt: Date()))) + + await firstTask.value + await manualImportTask.value + + #expect(await blocker.startedCount() == 2) + #expect(store.openAIDashboard?.creditsRemaining == 25) + #expect(store.openAIDashboard?.accountPlan == "Pro") + } + + @Test + func `stale cookie import status does not override later unrelated refresh failure`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebRefreshTests-stale-cookie-status") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store.openAIDashboardCookieImportStatus = + "OpenAI cookies are for other@example.com, not managed@example.com." + store._test_openAIDashboardLoaderOverride = { _, _, _ in + throw ManagedDashboardTestError.networkTimeout + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + + #expect(store.lastOpenAIDashboardError == ManagedDashboardTestError.networkTimeout.localizedDescription) + } + + @Test + func `reset open A I web state blocks stale in flight dashboard completion`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebRefreshTests-reset-invalidates-task") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingManagedOpenAIDashboardLoader() + store._test_openAIDashboardLoaderOverride = { _, _, _ in + try await blocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + let refreshTask = Task { + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + } + await blocker.waitUntilStarted() + + store.resetOpenAIWebState() + #expect(store.openAIDashboardRefreshTaskToken == nil) + + await blocker.resumeNext(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 85, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 12, + accountPlan: "Pro", + updatedAt: Date()))) + + await refreshTask.value + + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardError == nil) + } + + @Test + func `active refresh failure ignores stale import status from older task`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebRefreshTests-concurrent-import-status") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingManagedOpenAIDashboardLoader() + let importTracker = OpenAIDashboardImportCallTracker() + store._test_openAIDashboardLoaderOverride = { _, _, _ in + try await blocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + store._test_openAIDashboardCookieImportOverride = { _, _, _, _, _ in + let call = await importTracker.recordCall() + if call == 1 { + return OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "Chrome", + cookieCount: 2, + signedInEmail: managedAccount.email, + matchesCodexEmail: true) + } + throw OpenAIDashboardBrowserCookieImporter.ImportError.noMatchingAccount( + found: [.init(sourceLabel: "Chrome", email: "other@example.com")]) + } + defer { store._test_openAIDashboardCookieImportOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + let firstTask = Task { + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + } + await blocker.waitUntilStarted(count: 1) + + let secondTask = Task { + await store.importOpenAIDashboardBrowserCookiesNow() + } + await blocker.waitUntilStarted(count: 2) + + await blocker.resumeNext(with: .failure(OpenAIDashboardFetcher.FetchError.loginRequired)) + await importTracker.waitUntilCalls(count: 2) + await blocker.resumeNext(with: .failure(ManagedDashboardTestError.networkTimeout)) + + await firstTask.value + await secondTask.value + + #expect(store.lastOpenAIDashboardError == ManagedDashboardTestError.networkTimeout.localizedDescription) + } + + private func makeSettingsStore(suite: String) -> SettingsStore { + let defaults = UserDefaults(suiteName: suite)! + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + return SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + } +} + +private enum ManagedDashboardTestError: LocalizedError { + case networkTimeout + + var errorDescription: String? { + switch self { + case .networkTimeout: + "Network timeout" + } + } +} + +private actor BlockingManagedOpenAIDashboardLoader { + private var continuations: [CheckedContinuation, Never>] = [] + private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var started: Int = 0 + + func awaitResult() async throws -> OpenAIDashboardSnapshot { + self.started += 1 + self.resumeReadyStartWaiters() + let result = await withCheckedContinuation { continuation in + self.continuations.append(continuation) + } + return try result.get() + } + + func waitUntilStarted(count: Int = 1) async { + if self.started >= count { return } + await withCheckedContinuation { continuation in + self.startWaiters.append((count: count, continuation: continuation)) + } + } + + func startedCount() -> Int { + self.started + } + + func resumeNext(with result: Result) { + guard !self.continuations.isEmpty else { return } + let continuation = self.continuations.removeFirst() + continuation.resume(returning: result) + } + + private func resumeReadyStartWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.startWaiters { + if self.started >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.startWaiters = remaining + } +} + +private actor OpenAIDashboardImportCallTracker { + private var calls: Int = 0 + private var waiters: [(count: Int, continuation: CheckedContinuation)] = [] + + func recordCall() -> Int { + self.calls += 1 + self.resumeReadyWaiters() + return self.calls + } + + func waitUntilCalls(count: Int) async { + if self.calls >= count { return } + await withCheckedContinuation { continuation in + self.waiters.append((count: count, continuation: continuation)) + } + } + + private func resumeReadyWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.waiters { + if self.calls >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.waiters = remaining + } +} diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebTests.swift new file mode 100644 index 000000000..71d60b083 --- /dev/null +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebTests.swift @@ -0,0 +1,884 @@ +import Foundation +import Testing +@testable import CodexBar +@testable import CodexBarCore + +@Suite(.serialized) +@MainActor +struct CodexManagedOpenAIWebTests { + @Test + func `managed codex open A I web uses active managed identity and cache scope`() { + KeychainCacheStore.setTestStoreForTesting(true) + defer { KeychainCacheStore.setTestStoreForTesting(false) } + + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-managed") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let otherAccountID = UUID() + CookieHeaderCache.store( + provider: .codex, + scope: .managedAccount(otherAccountID), + cookieHeader: "auth=other-account", + sourceLabel: "Chrome") + CookieHeaderCache.store( + provider: .codex, + cookieHeader: "auth=provider-global", + sourceLabel: "Safari") + defer { + CookieHeaderCache.clear(provider: .codex, scope: .managedAccount(otherAccountID)) + CookieHeaderCache.clear(provider: .codex) + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(store.codexAccountEmailForOpenAIDashboard() == "managed@example.com") + #expect(store.codexCookieCacheScopeForOpenAIWeb() == .managedAccount(managedAccount.id)) + #expect(CookieHeaderCache.load(provider: .codex, scope: store.codexCookieCacheScopeForOpenAIWeb()) == nil) + } + + @Test + func `live system codex open A I web uses live identity and no managed cache scope`() { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-live-system") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let liveAccount = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: "/tmp/live-codex-home", + observedAt: Date()) + settings._test_activeManagedCodexAccount = managedAccount + settings._test_liveSystemCodexAccount = liveAccount + settings.codexActiveSource = .liveSystem + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: ["CODEX_HOME": liveAccount.codexHomePath]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(store.codexAccountEmailForOpenAIDashboard() == liveAccount.email) + #expect(store.codexAccountEmailForOpenAIDashboard() != managedAccount.email) + #expect(store.codexCookieCacheScopeForOpenAIWeb() == nil) + } + + @Test + func `live system codex open A I web does not reuse stale managed snapshot email after source switch`() { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-live-system-stale-managed-snapshot") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-empty-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + + settings._test_activeManagedCodexAccount = managedAccount + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + settings.codexActiveSource = .liveSystem + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._setSnapshotForTesting( + UsageSnapshot( + primary: nil, + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: managedAccount.email, + accountOrganization: nil, + loginMethod: nil)), + provider: .codex) + + #expect(store.codexAccountEmailForOpenAIDashboard() == nil) + #expect(store.codexAccountEmailForOpenAIDashboard() != managedAccount.email) + #expect(store.codexCookieCacheScopeForOpenAIWeb() == nil) + } + + @Test + func `live system codex open A I web reuses last known live email without allowing any account`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-live-system-last-known-email") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-last-known-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + let liveAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/tmp/live-codex-home", + observedAt: Date()) + settings._test_liveSystemCodexAccount = liveAccount + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + settings.codexActiveSource = .liveSystem + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(store.codexAccountEmailForOpenAIDashboard() == liveAccount.email) + + settings._test_liveSystemCodexAccount = nil + defer { + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + store.openAIDashboard = OpenAIDashboardSnapshot( + signedInEmail: "managed@example.com", + codeReviewRemainingPercent: 100, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + updatedAt: Date()) + store.lastOpenAIDashboardCookieImportEmail = "managed-import@example.com" + + var observedTargetEmail: String? + var observedAllowAnyAccount: Bool? + store._test_openAIDashboardCookieImportOverride = { targetEmail, allowAnyAccount, _, _, _ in + observedTargetEmail = targetEmail + observedAllowAnyAccount = allowAnyAccount + return OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "test", + cookieCount: 1, + signedInEmail: targetEmail, + matchesCodexEmail: true) + } + defer { store._test_openAIDashboardCookieImportOverride = nil } + + let targetEmail = store.codexAccountEmailForOpenAIDashboard() + let imported = await store.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) + + #expect(targetEmail == liveAccount.email) + #expect(targetEmail != "managed@example.com") + #expect(targetEmail != "managed-import@example.com") + #expect(imported == liveAccount.email) + #expect(observedTargetEmail == liveAccount.email) + #expect(observedAllowAnyAccount == false) + } + + @Test + func `dashboard refresh does not target stale last known live email`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-live-system-refresh-strict-target") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-refresh-strict-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + let liveAccount = ObservedSystemCodexAccount( + email: "old@example.com", + codexHomePath: "/tmp/live-codex-home", + observedAt: Date()) + settings._test_liveSystemCodexAccount = liveAccount + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + settings.codexActiveSource = .liveSystem + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(store.codexAccountEmailForOpenAIDashboard() == liveAccount.email) + + settings._test_liveSystemCodexAccount = nil + defer { + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + var observedTargetEmail: String? + store._test_openAIDashboardLoaderOverride = { accountEmail, _, _ in + observedTargetEmail = accountEmail + return OpenAIDashboardSnapshot( + signedInEmail: "new@example.com", + codeReviewRemainingPercent: 88, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + primaryLimit: RateWindow(usedPercent: 12, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondaryLimit: nil, + creditsRemaining: 22, + accountPlan: "Pro", + updatedAt: Date()) + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + #expect(expectedGuard.accountKey == nil) + + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + + #expect(observedTargetEmail == nil) + #expect(store.openAIDashboard?.signedInEmail == "new@example.com") + #expect(store.lastKnownLiveSystemCodexEmail == "new@example.com") + } + + @Test + func `dashboard refresh targets usage discovered live email before reconciliation catches up`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-live-system-usage-discovered-target") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-usage-discovered-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + settings.codexActiveSource = .liveSystem + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._setSnapshotForTesting( + UsageSnapshot( + primary: nil, + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: "usage@example.com", + accountOrganization: nil, + loginMethod: nil)), + provider: .codex) + + var observedTargetEmail: String? + store._test_openAIDashboardLoaderOverride = { accountEmail, _, _ in + observedTargetEmail = accountEmail + return OpenAIDashboardSnapshot( + signedInEmail: "usage@example.com", + codeReviewRemainingPercent: 88, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + primaryLimit: RateWindow(usedPercent: 12, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondaryLimit: nil, + creditsRemaining: 22, + accountPlan: "Pro", + updatedAt: Date()) + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + #expect(expectedGuard.accountKey == nil) + + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + + #expect(observedTargetEmail == "usage@example.com") + #expect(store.openAIDashboard?.signedInEmail == "usage@example.com") + } + + @Test + func `usage discovered live email still surfaces open A I web login guidance during reconciliation lag`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-live-system-usage-discovered-failure") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-usage-discovered-failure-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + settings.codexActiveSource = .liveSystem + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._setSnapshotForTesting( + UsageSnapshot( + primary: nil, + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: "usage@example.com", + accountOrganization: nil, + loginMethod: nil)), + provider: .codex) + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + #expect(expectedGuard.accountKey == nil) + + await store.applyOpenAIDashboardLoginRequiredFailure(expectedGuard: expectedGuard) + + #expect(store.openAIDashboardRequiresLogin == true) + #expect(store.lastOpenAIDashboardError?.contains("requires a signed-in chatgpt.com session") == true) + } + + @Test + func `open A I web import uses managed account target when live account differs`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-targeting-active-vs-live") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let liveAccount = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: "/tmp/live-codex-home", + observedAt: Date()) + let expectedScope = CookieHeaderCache.Scope.managedAccount(managedAccount.id) + let expectedEmail = managedAccount.email + var observedTargetEmail: String? + var observedScope: CookieHeaderCache.Scope? + var observedCookieSource: ProviderCookieSource? + var observedAllowAnyAccount = false + + settings._test_activeManagedCodexAccount = managedAccount + settings._test_liveSystemCodexAccount = liveAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let snapshot = UsageSnapshot( + primary: nil, + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: liveAccount.email, + accountOrganization: nil, + loginMethod: nil)) + + let store = UsageStore( + fetcher: UsageFetcher(environment: ["CODEX_HOME": liveAccount.codexHomePath]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._setSnapshotForTesting(snapshot, provider: .codex) + store._test_openAIDashboardCookieImportOverride = { targetEmail, allowAnyAccount, cookieSource, scope, _ in + observedTargetEmail = targetEmail + observedScope = scope + observedCookieSource = cookieSource + observedAllowAnyAccount = allowAnyAccount + return OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "test", + cookieCount: 1, + signedInEmail: targetEmail, + matchesCodexEmail: targetEmail == expectedEmail) + } + defer { store._test_openAIDashboardCookieImportOverride = nil } + + let importerTarget = store.codexAccountEmailForOpenAIDashboard() + let imported = await store.importOpenAIDashboardCookiesIfNeeded(targetEmail: importerTarget, force: true) + + #expect(importerTarget == expectedEmail) + #expect(importerTarget != liveAccount.email) + #expect(imported == expectedEmail) + #expect(observedTargetEmail == expectedEmail) + #expect(observedScope == expectedScope) + #expect(observedAllowAnyAccount == false) + #expect(observedCookieSource == .auto) + #expect(store.codexCookieCacheScopeForOpenAIWeb() == expectedScope) + } + + @Test + func `open A I web prefers live identity when managed and live share email`() { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-same-email-prefers-live") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "person@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let liveAccount = ObservedSystemCodexAccount( + email: "PERSON@example.com", + codexHomePath: "/tmp/live-codex-home", + observedAt: Date()) + settings._test_activeManagedCodexAccount = managedAccount + settings._test_liveSystemCodexAccount = liveAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: ["CODEX_HOME": liveAccount.codexHomePath]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(settings.codexResolvedActiveSource == .liveSystem) + #expect(store.codexAccountEmailForOpenAIDashboard() == "person@example.com") + #expect(store.codexCookieCacheScopeForOpenAIWeb() == nil) + } + + @Test + func `unmanaged codex open A I web falls back to provider global cache scope`() { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-unmanaged") + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(store.codexCookieCacheScopeForOpenAIWeb() == nil) + } + + @Test + func `unreadable managed codex store fails closed for open A I web`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-unreadable-store") + settings._test_unreadableManagedCodexAccountStore = true + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { settings._test_unreadableManagedCodexAccountStore = false } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(store.codexCookieCacheScopeForOpenAIWeb() == .managedStoreUnreadable) + #expect(store.codexAccountEmailForOpenAIDashboard() == nil) + + let imported = await store.importOpenAIDashboardCookiesIfNeeded(targetEmail: nil, force: true) + + #expect(imported == nil) + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardRequiresLogin == true) + #expect(store.openAIDashboardCookieImportStatus?.contains("Managed Codex account data is unavailable") == true) + + await store.refreshOpenAIDashboardIfNeeded(force: true) + + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardRequiresLogin == true) + #expect(store.lastOpenAIDashboardError?.contains("Managed Codex account data is unavailable") == true) + } + + @Test + func `missing managed codex open A I web target fails closed`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-missing-managed-target") + let storedAccount = ManagedCodexAccount( + id: UUID(), + email: "stored@example.com", + managedHomePath: "/tmp/stored-managed-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let storeURL = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-\(UUID().uuidString).json") + let managedStore = FileManagedCodexAccountStore(fileURL: storeURL) + try? managedStore.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [storedAccount])) + settings._test_managedCodexAccountStoreURL = storeURL + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { + settings._test_managedCodexAccountStoreURL = nil + try? FileManager.default.removeItem(at: storeURL) + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + var importWasCalled = false + store._test_openAIDashboardCookieImportOverride = { _, _, _, _, _ in + importWasCalled = true + return OpenAIDashboardBrowserCookieImporter.ImportResult( + sourceLabel: "test", + cookieCount: 1, + signedInEmail: "unexpected@example.com", + matchesCodexEmail: true) + } + defer { store._test_openAIDashboardCookieImportOverride = nil } + store.openAIDashboard = OpenAIDashboardSnapshot( + signedInEmail: "stale-dashboard@example.com", + codeReviewRemainingPercent: 100, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + updatedAt: Date()) + store.lastOpenAIDashboardCookieImportEmail = "stale-import@example.com" + + #expect(store.codexAccountEmailForOpenAIDashboard() == nil) + #expect(store.codexCookieCacheScopeForOpenAIWeb() != nil) + #expect(store.codexCookieCacheScopeForOpenAIWeb() != .managedStoreUnreadable) + + let imported = await store.importOpenAIDashboardCookiesIfNeeded(targetEmail: nil, force: true) + #expect(imported == nil) + #expect(importWasCalled == false) + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardRequiresLogin == true) + #expect(store.openAIDashboardCookieImportStatus? + .contains("selected managed Codex account is unavailable") == true) + + await store.refreshOpenAIDashboardIfNeeded(force: true) + #expect(importWasCalled == false) + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardError?.contains("selected managed Codex account is unavailable") == true) + } + + @Test + func `managed codex mismatch fail closed blocks stale dashboard restoration`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-mismatch") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let staleSnapshot = OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 100, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + updatedAt: Date()) + + await store.applyOpenAIDashboard(staleSnapshot, targetEmail: managedAccount.email) + await store.applyOpenAIDashboardMismatchFailure( + signedInEmail: "other@example.com", + expectedEmail: managedAccount.email) + + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardRequiresLogin == true) + + await store.applyOpenAIDashboardFailure(message: "No dashboard data") + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardError == "No dashboard data") + + await store.applyOpenAIDashboardLoginRequiredFailure() + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardRequiresLogin == true) + #expect(store.lastOpenAIDashboardError?.contains("requires a signed-in chatgpt.com session") == true) + } + + @Test + func `managed codex import mismatch fail closed blocks stale dashboard restoration`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-import-mismatch") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._test_openAIDashboardCookieImportOverride = { _, _, _, _, _ in + throw OpenAIDashboardBrowserCookieImporter.ImportError.noMatchingAccount( + found: [.init(sourceLabel: "Chrome", email: "other@example.com")]) + } + + let staleSnapshot = OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 100, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + updatedAt: Date()) + await store.applyOpenAIDashboard(staleSnapshot, targetEmail: managedAccount.email) + + let imported = await store.importOpenAIDashboardCookiesIfNeeded( + targetEmail: managedAccount.email, + force: true) + + #expect(imported == nil) + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardRequiresLogin == true) + #expect( + store.openAIDashboardCookieImportStatus == + "OpenAI cookies are for other@example.com, not managed@example.com.") + + await store.applyOpenAIDashboardFailure(message: "No dashboard data") + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardError == "No dashboard data") + + await store.applyOpenAIDashboardLoginRequiredFailure() + #expect(store.openAIDashboard == nil) + #expect(store.openAIDashboardRequiresLogin == true) + #expect(store.lastOpenAIDashboardError?.contains("requires a signed-in chatgpt.com session") == true) + } + + @Test + func `missing managed target failure handlers do not resurrect stale dashboard state`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-missing-target-failure-handlers") + let isolatedHome = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-openai-web-missing-target-\(UUID().uuidString)", isDirectory: true) + try? FileManager.default.createDirectory(at: isolatedHome, withIntermediateDirectories: true) + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": isolatedHome.path] + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { + settings._test_codexReconciliationEnvironment = nil + try? FileManager.default.removeItem(at: isolatedHome) + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let staleSnapshot = OpenAIDashboardSnapshot( + signedInEmail: "stale@example.com", + codeReviewRemainingPercent: 100, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + updatedAt: Date()) + + await store.applyOpenAIDashboard(staleSnapshot, targetEmail: "stale@example.com") + await store.applyOpenAIDashboardFailure(message: "No dashboard data") + + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardSnapshot == nil) + #expect(store.lastOpenAIDashboardError?.contains("selected managed Codex account is unavailable") == true) + + await store.applyOpenAIDashboard(staleSnapshot, targetEmail: "stale@example.com") + await store.applyOpenAIDashboardLoginRequiredFailure() + + #expect(store.openAIDashboard == nil) + #expect(store.lastOpenAIDashboardSnapshot == nil) + #expect(store.openAIDashboardRequiresLogin == true) + #expect(store.lastOpenAIDashboardError?.contains("selected managed Codex account is unavailable") == true) + } + + @Test + func `managed codex refresh stops after cookie mismatch instead of retrying web view`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-mismatch-aborts-retry") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "ratulsarna@gmail.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + var loaderCalls = 0 + store._test_openAIDashboardLoaderOverride = { _, _, _ in + loaderCalls += 1 + throw OpenAIDashboardFetcher.FetchError.loginRequired + } + defer { store._test_openAIDashboardLoaderOverride = nil } + store._test_openAIDashboardCookieImportOverride = { _, _, _, _, _ in + throw OpenAIDashboardBrowserCookieImporter.ImportError.noMatchingAccount( + found: [.init(sourceLabel: "Chrome", email: "rdsarna@gmail.com")]) + } + defer { store._test_openAIDashboardCookieImportOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + + #expect(loaderCalls == 1) + #expect( + store.lastOpenAIDashboardError == + "OpenAI cookies are for rdsarna@gmail.com, not ratulsarna@gmail.com. " + + "Switch chatgpt.com account, then refresh OpenAI cookies.") + #expect(store.openAIDashboard == nil) + } + + @Test + func `same account dashboard refresh requests coalesce while one is in flight`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-refresh-coalesce") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/managed-codex-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { settings._test_activeManagedCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let blocker = BlockingManagedOpenAIDashboardLoader() + store._test_openAIDashboardLoaderOverride = { _, _, _ in + try await blocker.awaitResult() + } + defer { store._test_openAIDashboardLoaderOverride = nil } + + let expectedGuard = store.currentCodexOpenAIWebRefreshGuard() + let firstTask = Task { + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + } + await blocker.waitUntilStarted() + + let secondTask = Task { + await store.refreshOpenAIDashboardIfNeeded(force: true, expectedGuard: expectedGuard) + } + + try? await Task.sleep(nanoseconds: 50_000_000) + #expect(await blocker.startedCount() == 1) + + await blocker.resume(with: .success(OpenAIDashboardSnapshot( + signedInEmail: managedAccount.email, + codeReviewRemainingPercent: 90, + creditEvents: [], + dailyBreakdown: [], + usageBreakdown: [], + creditsPurchaseURL: nil, + creditsRemaining: 10, + accountPlan: "Pro", + updatedAt: Date()))) + + await firstTask.value + await secondTask.value + + #expect(await blocker.startedCount() == 1) + #expect(store.openAIDashboard?.signedInEmail == managedAccount.email) + } + + @Test + func `friendly error shortens cookie mismatch copy`() { + let settings = self.makeSettingsStore(suite: "CodexManagedOpenAIWebTests-friendly-error-short") + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + let message = store.openAIDashboardFriendlyError( + body: "Sign in to continue", + targetEmail: "ratulsarna@gmail.com", + cookieImportStatus: "OpenAI cookies are for rdsarna@gmail.com, not ratulsarna@gmail.com.") + + #expect( + message == + "OpenAI cookies are for rdsarna@gmail.com, not ratulsarna@gmail.com. " + + "Switch chatgpt.com account, then refresh OpenAI cookies.") + } + + private func makeSettingsStore(suite: String) -> SettingsStore { + let defaults = UserDefaults(suiteName: suite)! + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + let settings = SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + settings._test_activeManagedCodexAccount = nil + settings._test_activeManagedCodexRemoteHomePath = nil + settings._test_unreadableManagedCodexAccountStore = false + settings._test_managedCodexAccountStoreURL = nil + settings._test_liveSystemCodexAccount = nil + settings._test_codexReconciliationEnvironment = nil + return settings + } +} + +private actor BlockingManagedOpenAIDashboardLoader { + private var continuations: [CheckedContinuation, Never>] = [] + private var startWaiters: [(count: Int, continuation: CheckedContinuation)] = [] + private var started: Int = 0 + + func awaitResult() async throws -> OpenAIDashboardSnapshot { + self.started += 1 + self.resumeReadyStartWaiters() + let result = await withCheckedContinuation { continuation in + self.continuations.append(continuation) + } + return try result.get() + } + + func waitUntilStarted(count: Int = 1) async { + if self.started >= count { return } + await withCheckedContinuation { continuation in + self.startWaiters.append((count: count, continuation: continuation)) + } + } + + func startedCount() -> Int { + self.started + } + + func resume(with result: Result) { + self.resumeNext(with: result) + } + + func resumeNext(with result: Result) { + guard !self.continuations.isEmpty else { return } + let continuation = self.continuations.removeFirst() + continuation.resume(returning: result) + } + + private func resumeReadyStartWaiters() { + var remaining: [(count: Int, continuation: CheckedContinuation)] = [] + for waiter in self.startWaiters { + if self.started >= waiter.count { + waiter.continuation.resume() + } else { + remaining.append(waiter) + } + } + self.startWaiters = remaining + } +} diff --git a/Tests/CodexBarTests/CodexManagedRoutingTests.swift b/Tests/CodexBarTests/CodexManagedRoutingTests.swift new file mode 100644 index 000000000..9a85e780f --- /dev/null +++ b/Tests/CodexBarTests/CodexManagedRoutingTests.swift @@ -0,0 +1,601 @@ +import Foundation +import Testing +@testable import CodexBar +@testable import CodexBarCore + +@Suite(.serialized) +@MainActor +struct CodexManagedRoutingTests { + @Test + func `provider registry injects managed home when active source is managed account`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-registry") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: "/tmp/codex-managed-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + settings._test_activeManagedCodexAccount = managedAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { + settings._test_activeManagedCodexAccount = nil + } + + let codexEnv = ProviderRegistry.makeEnvironment( + base: ["PATH": "/usr/bin"], + provider: .codex, + settings: settings, + tokenOverride: nil) + let claudeEnv = ProviderRegistry.makeEnvironment( + base: ["PATH": "/usr/bin"], + provider: .claude, + settings: settings, + tokenOverride: nil) + + #expect(codexEnv["CODEX_HOME"] == managedAccount.managedHomePath) + #expect(claudeEnv["CODEX_HOME"] == nil) + } + + @Test + func `provider registry preserves ambient live system home when active source is live system`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-live-system-routing") + let managedHomePath = "/tmp/managed-remote-home" + let liveHomePath = "/tmp/system-remote-home" + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomePath, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let liveSystemAccount = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: liveHomePath, + observedAt: Date()) + + settings._test_activeManagedCodexAccount = managedAccount + settings._test_liveSystemCodexAccount = liveSystemAccount + settings.codexActiveSource = .liveSystem + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let env = ProviderRegistry.makeEnvironment( + base: ["CODEX_HOME": liveHomePath], + provider: .codex, + settings: settings, + tokenOverride: nil) + + #expect(env["CODEX_HOME"] == liveHomePath) + #expect(env["CODEX_HOME"] != managedHomePath) + } + + @Test + func `provider registry keeps managed home when live account differs`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-active-vs-live") + let managedHomePath = "/tmp/managed-remote-home" + let liveHomePath = "/tmp/system-remote-home" + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "managed@example.com", + managedHomePath: managedHomePath, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let liveSystemAccount = ObservedSystemCodexAccount( + email: "system@example.com", + codexHomePath: liveHomePath, + observedAt: Date()) + + settings._test_activeManagedCodexAccount = managedAccount + settings._test_liveSystemCodexAccount = liveSystemAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let env = ProviderRegistry.makeEnvironment( + base: ["CODEX_HOME": liveHomePath], + provider: .codex, + settings: settings, + tokenOverride: nil) + + #expect(env["CODEX_HOME"] == managedHomePath) + #expect(env["CODEX_HOME"] != liveHomePath) + } + + @Test + func `provider registry prefers live system routing when managed and live share email`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-same-email-prefers-live") + let managedHomePath = "/tmp/managed-remote-home" + let liveHomePath = "/tmp/system-remote-home" + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "person@example.com", + managedHomePath: managedHomePath, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let liveSystemAccount = ObservedSystemCodexAccount( + email: "PERSON@example.com", + codexHomePath: liveHomePath, + observedAt: Date()) + + settings._test_activeManagedCodexAccount = managedAccount + settings._test_liveSystemCodexAccount = liveSystemAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let env = ProviderRegistry.makeEnvironment( + base: ["CODEX_HOME": liveHomePath], + provider: .codex, + settings: settings, + tokenOverride: nil) + + #expect(settings.codexResolvedActiveSource == .liveSystem) + #expect(env["CODEX_HOME"] == liveHomePath) + #expect(env["CODEX_HOME"] != managedHomePath) + } + + @Test + func `persisted managed source corrects to live system when selected row collapses with live account`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-same-email-persist-correction") + let managedAccount = ManagedCodexAccount( + id: UUID(), + email: "person@example.com", + managedHomePath: "/tmp/managed-remote-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let liveSystemAccount = ObservedSystemCodexAccount( + email: "PERSON@example.com", + codexHomePath: "/tmp/system-remote-home", + observedAt: Date()) + + settings._test_activeManagedCodexAccount = managedAccount + settings._test_liveSystemCodexAccount = liveSystemAccount + settings.codexActiveSource = .managedAccount(id: managedAccount.id) + defer { + settings._test_activeManagedCodexAccount = nil + settings._test_liveSystemCodexAccount = nil + } + + let corrected = settings.persistResolvedCodexActiveSourceCorrectionIfNeeded() + + #expect(corrected) + #expect(settings.codexActiveSource == .liveSystem) + } + + @Test + func `codex provider refresh persists live correction for stale managed source`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-provider-refresh-persists-correction") + let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: ambientHome) } + + try? self.writeCodexAuthFile(homeURL: ambientHome, email: "live@example.com", plan: "pro") + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: ambientHome.path, + observedAt: Date()) + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { settings._test_liveSystemCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: ["CODEX_HOME": ambientHome.path]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(settings.codexActiveSource != .liveSystem) + + await store.refreshProvider(.codex, allowDisabled: true) + + #expect(settings.codexActiveSource == .liveSystem) + } + + @Test + func `full refresh persists live correction for stale managed source`() async { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-full-refresh-persists-correction") + let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: ambientHome) } + + try? self.writeCodexAuthFile(homeURL: ambientHome, email: "live@example.com", plan: "pro") + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: ambientHome.path, + observedAt: Date()) + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { settings._test_liveSystemCodexAccount = nil } + + let store = UsageStore( + fetcher: UsageFetcher(environment: ["CODEX_HOME": ambientHome.path]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + + #expect(settings.codexActiveSource != .liveSystem) + + await store.refresh() + + #expect(settings.codexActiveSource == .liveSystem) + } + + @Test + func `provider registry fails closed when managed account store is unreadable`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-unreadable-store") + settings._test_unreadableManagedCodexAccountStore = true + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { settings._test_unreadableManagedCodexAccountStore = false } + + let env = ProviderRegistry.makeEnvironment( + base: ["CODEX_HOME": "/Users/example/.codex"], + provider: .codex, + settings: settings, + tokenOverride: nil) + + #expect(env["CODEX_HOME"] != nil) + #expect(env["CODEX_HOME"] != "/Users/example/.codex") + #expect(env["CODEX_HOME"]?.isEmpty == false) + } + + @Test + func `provider registry bootstraps live system source instead of inferring managed fallback`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-unreadable-legacy-source") + settings._test_unreadableManagedCodexAccountStore = true + defer { settings._test_unreadableManagedCodexAccountStore = false } + + let ambientHome = "/Users/example/.codex" + let env = ProviderRegistry.makeEnvironment( + base: ["CODEX_HOME": ambientHome], + provider: .codex, + settings: settings, + tokenOverride: nil) + let snapshot = settings.codexSettingsSnapshot(tokenOverride: nil) + + #expect(env["CODEX_HOME"] == ambientHome) + #expect(settings.providerConfig(for: .codex)?.codexActiveSource == .liveSystem) + #expect(snapshot.managedAccountStoreUnreadable == false) + #expect(snapshot.managedAccountTargetUnavailable == false) + } + + @Test + func `provider registry fails closed when selected managed source is missing from readable store`() throws { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-missing-managed-source") + let storedAccount = ManagedCodexAccount( + id: UUID(), + email: "stored@example.com", + managedHomePath: "/tmp/stored-managed-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let storeURL = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-managed-routing-\(UUID().uuidString).json") + let store = FileManagedCodexAccountStore(fileURL: storeURL) + try store.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [storedAccount])) + settings._test_managedCodexAccountStoreURL = storeURL + settings.codexActiveSource = .managedAccount(id: UUID()) + settings._test_codexReconciliationEnvironment = ["CODEX_HOME": "/Users/example/.codex"] + defer { + settings._test_codexReconciliationEnvironment = nil + settings._test_managedCodexAccountStoreURL = nil + try? FileManager.default.removeItem(at: storeURL) + } + + let ambientHome = "/Users/example/.codex" + let expectedFailClosedPath = ManagedCodexHomeFactory.defaultRootURL() + .appendingPathComponent("managed-store-unreadable", isDirectory: true) + .path + let env = ProviderRegistry.makeEnvironment( + base: ["CODEX_HOME": ambientHome], + provider: .codex, + settings: settings, + tokenOverride: nil) + + #expect(env["CODEX_HOME"] == expectedFailClosedPath) + #expect(env["CODEX_HOME"] != ambientHome) + #expect(env["CODEX_HOME"] != storedAccount.managedHomePath) + } + + @Test + func `codex settings snapshot marks missing selected managed source as unavailable`() throws { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-missing-managed-snapshot") + let storedAccount = ManagedCodexAccount( + id: UUID(), + email: "stored@example.com", + managedHomePath: "/tmp/stored-managed-home", + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let storeURL = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-managed-snapshot-\(UUID().uuidString).json") + let store = FileManagedCodexAccountStore(fileURL: storeURL) + try store.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: [storedAccount])) + settings._test_managedCodexAccountStoreURL = storeURL + settings.codexActiveSource = .managedAccount(id: UUID()) + defer { + settings._test_managedCodexAccountStoreURL = nil + try? FileManager.default.removeItem(at: storeURL) + } + + let snapshot = settings.codexSettingsSnapshot(tokenOverride: nil) + + #expect(snapshot.managedAccountStoreUnreadable == false) + #expect(snapshot.managedAccountTargetUnavailable == true) + } + + @Test + func `codex settings snapshot ignores unreadable added account store when live system is active`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-live-system-snapshot") + settings._test_unreadableManagedCodexAccountStore = true + settings.codexActiveSource = .liveSystem + defer { settings._test_unreadableManagedCodexAccountStore = false } + + let snapshot = settings.codexSettingsSnapshot(tokenOverride: nil) + + #expect(snapshot.managedAccountStoreUnreadable == false) + #expect(snapshot.managedAccountTargetUnavailable == false) + } + + @Test + func `provider registry ignores debug managed home override without explicit managed source`() throws { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-debug-home-override") + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: managedHome) } + + settings._test_activeManagedCodexRemoteHomePath = managedHome.path + defer { settings._test_activeManagedCodexRemoteHomePath = nil } + try self.writeCodexAuthFile(homeURL: managedHome, email: "managed@example.com", plan: "pro") + + let ambientHome = "/Users/example/.codex" + let env = ProviderRegistry.makeEnvironment( + base: ["CODEX_HOME": ambientHome], + provider: .codex, + settings: settings, + tokenOverride: nil) + + #expect(env["CODEX_HOME"] == ambientHome) + #expect(settings.providerConfig(for: .codex)?.codexActiveSource == .liveSystem) + } + + @Test + func `provider registry builds codex fetcher scoped to managed home`() throws { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-registry-fetcher") + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: managedHome) } + + settings._test_activeManagedCodexRemoteHomePath = managedHome.path + settings.codexActiveSource = .managedAccount(id: UUID()) + try self.writeCodexAuthFile(homeURL: managedHome, email: "managed@example.com", plan: "pro") + defer { + settings._test_activeManagedCodexRemoteHomePath = nil + } + + let browserDetection = BrowserDetection(cacheTTL: 0) + let specs = ProviderRegistry.shared.specs( + settings: settings, + metadata: ProviderDescriptorRegistry.metadata, + codexFetcher: UsageFetcher(environment: [:]), + claudeFetcher: ClaudeUsageFetcher(browserDetection: browserDetection), + browserDetection: browserDetection) + let context = try #require(specs[.codex]?.makeFetchContext()) + + let account = context.fetcher.loadAccountInfo() + #expect(account.email == "managed@example.com") + #expect(account.plan == "pro") + } + + @Test + func `usage store builds codex token account fetcher scoped to managed home`() throws { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-usage-store") + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: managedHome) } + + settings._test_activeManagedCodexRemoteHomePath = managedHome.path + settings.codexActiveSource = .managedAccount(id: UUID()) + try self.writeCodexAuthFile(homeURL: managedHome, email: "token@example.com", plan: "team") + defer { + settings._test_activeManagedCodexRemoteHomePath = nil + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let context = store.makeFetchContext(provider: .codex, override: nil) + + let account = context.fetcher.loadAccountInfo() + #expect(account.email == "token@example.com") + #expect(account.plan == "team") + } + + @Test + func `usage store builds codex credits fetcher scoped to managed home`() throws { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-credits-fetcher") + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: managedHome) } + + settings._test_activeManagedCodexRemoteHomePath = managedHome.path + settings.codexActiveSource = .managedAccount(id: UUID()) + try self.writeCodexAuthFile(homeURL: managedHome, email: "credits@example.com", plan: "enterprise") + defer { + settings._test_activeManagedCodexRemoteHomePath = nil + } + + let store = UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + let account = store.codexCreditsFetcher().loadAccountInfo() + + #expect(account.email == "credits@example.com") + #expect(account.plan == "enterprise") + } + + @Test + func `codex O auth strategy availability reads auth from context env`() async throws { + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: managedHome) } + + let credentials = CodexOAuthCredentials( + accessToken: "access-token", + refreshToken: "refresh-token", + idToken: nil, + accountId: nil, + lastRefresh: Date()) + try CodexOAuthCredentialsStore.save(credentials, env: ["CODEX_HOME": managedHome.path]) + + let strategy = CodexOAuthFetchStrategy() + let available = await strategy.isAvailable(self.makeContext(env: ["CODEX_HOME": managedHome.path])) + + #expect(available) + } + + @Test + func `codex O auth credentials store loads and saves using explicit env`() throws { + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { try? FileManager.default.removeItem(at: managedHome) } + + let credentials = CodexOAuthCredentials( + accessToken: "access-token", + refreshToken: "refresh-token", + idToken: "id-token", + accountId: "account-id", + lastRefresh: Date()) + let env = ["CODEX_HOME": managedHome.path] + + try CodexOAuthCredentialsStore.save(credentials, env: env) + + let authURL = CodexOAuthCredentialsStore._authFileURLForTesting(env: env) + #expect(authURL.path == managedHome.appendingPathComponent("auth.json").path) + + let loaded = try CodexOAuthCredentialsStore.load(env: env) + #expect(loaded.accessToken == credentials.accessToken) + #expect(loaded.refreshToken == credentials.refreshToken) + #expect(loaded.idToken == credentials.idToken) + #expect(loaded.accountId == credentials.accountId) + } + + @Test + func `codex no data message uses explicit environment home`() { + let env = ["CODEX_HOME": "/tmp/managed-codex-home"] + + let message = CodexProviderDescriptor._noDataMessageForTesting(env: env) + + #expect(message.contains("/tmp/managed-codex-home/sessions")) + #expect(message.contains("/tmp/managed-codex-home/archived_sessions")) + } + + private func makeContext(env: [String: String]) -> ProviderFetchContext { + let browserDetection = BrowserDetection(cacheTTL: 0) + return ProviderFetchContext( + runtime: .app, + sourceMode: .auto, + includeCredits: false, + webTimeout: 60, + webDebugDumpHTML: false, + verbose: false, + env: env, + settings: nil, + fetcher: UsageFetcher(), + claudeFetcher: ClaudeUsageFetcher(browserDetection: browserDetection), + browserDetection: browserDetection) + } + + private func makeSettingsStore(suite: String) -> SettingsStore { + let defaults = UserDefaults(suiteName: suite)! + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + return SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: InMemoryZaiTokenStore(), + syntheticTokenStore: InMemorySyntheticTokenStore(), + codexCookieStore: InMemoryCookieHeaderStore(), + claudeCookieStore: InMemoryCookieHeaderStore(), + cursorCookieStore: InMemoryCookieHeaderStore(), + opencodeCookieStore: InMemoryCookieHeaderStore(), + factoryCookieStore: InMemoryCookieHeaderStore(), + minimaxCookieStore: InMemoryMiniMaxCookieStore(), + minimaxAPITokenStore: InMemoryMiniMaxAPITokenStore(), + kimiTokenStore: InMemoryKimiTokenStore(), + kimiK2TokenStore: InMemoryKimiK2TokenStore(), + augmentCookieStore: InMemoryCookieHeaderStore(), + ampCookieStore: InMemoryCookieHeaderStore(), + copilotTokenStore: InMemoryCopilotTokenStore(), + tokenAccountStore: InMemoryTokenAccountStore()) + } + + private 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))." + } +} + +private final class InMemoryZaiTokenStore: ZaiTokenStoring, @unchecked Sendable { + func loadToken() throws -> String? { + nil + } + + func storeToken(_: String?) throws {} +} + +private final class InMemorySyntheticTokenStore: SyntheticTokenStoring, @unchecked Sendable { + func loadToken() throws -> String? { + nil + } + + func storeToken(_: String?) throws {} +} diff --git a/Tests/CodexBarTests/CodexOAuthTests.swift b/Tests/CodexBarTests/CodexOAuthTests.swift index fdc0a2850..feec37717 100644 --- a/Tests/CodexBarTests/CodexOAuthTests.swift +++ b/Tests/CodexBarTests/CodexOAuthTests.swift @@ -25,6 +25,28 @@ struct CodexOAuthTests { #expect(creds.lastRefresh != nil) } + @Test + func `parses legacy camel case O auth credentials`() throws { + let json = """ + { + "OPENAI_API_KEY": null, + "tokens": { + "accessToken": "access-token", + "refreshToken": "refresh-token", + "idToken": "id-token", + "accountId": "account-123" + }, + "last_refresh": "2025-12-20T12:34:56Z" + } + """ + let creds = try CodexOAuthCredentialsStore.parse(data: Data(json.utf8)) + #expect(creds.accessToken == "access-token") + #expect(creds.refreshToken == "refresh-token") + #expect(creds.idToken == "id-token") + #expect(creds.accountId == "account-123") + #expect(creds.lastRefresh != nil) + } + @Test func `parses API key credentials`() throws { let json = """ diff --git a/Tests/CodexBarTests/CodexSystemAccountObserverTests.swift b/Tests/CodexBarTests/CodexSystemAccountObserverTests.swift new file mode 100644 index 000000000..cf2928d86 --- /dev/null +++ b/Tests/CodexBarTests/CodexSystemAccountObserverTests.swift @@ -0,0 +1,76 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@Suite(.serialized) +struct CodexSystemAccountObserverTests { + @Test + func `observer reads ambient CODEX_HOME when present`() throws { + let home = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: home) } + try Self.writeCodexAuthFile(homeURL: home, email: " LIVE@Example.com ", plan: "pro") + + let observer = DefaultCodexSystemAccountObserver() + let account = try observer.loadSystemAccount(environment: ["CODEX_HOME": home.path]) + + #expect(account?.email == "live@example.com") + #expect(account?.codexHomePath == home.path) + } + + @Test + func `observer falls back to nil when ambient home has no readable email`() throws { + let home = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: home) } + try FileManager.default.createDirectory(at: home, withIntermediateDirectories: true) + + let observer = DefaultCodexSystemAccountObserver() + let account = try observer.loadSystemAccount(environment: ["CODEX_HOME": home.path]) + + #expect(account == nil) + } + + @Test + func `observer records observation timestamp`() throws { + let home = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: home) } + try Self.writeCodexAuthFile(homeURL: home, email: "user@example.com", plan: "team") + + let before = Date() + let observer = DefaultCodexSystemAccountObserver() + let account = try observer.loadSystemAccount(environment: ["CODEX_HOME": home.path]) + let observed = try #require(account) + + #expect(observed.observedAt >= before) + } + + 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))." + } +} diff --git a/Tests/CodexBarTests/CodexUserFacingErrorTests.swift b/Tests/CodexBarTests/CodexUserFacingErrorTests.swift new file mode 100644 index 000000000..9b789cde2 --- /dev/null +++ b/Tests/CodexBarTests/CodexUserFacingErrorTests.swift @@ -0,0 +1,153 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@MainActor +struct CodexUserFacingErrorTests { + @Test + func `expired codex auth is sanitized`() { + let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-expired-auth") + store.errors[.codex] = """ + Codex connection failed: failed to fetch codex rate limits: GET https://chatgpt.com/backend-api/wham/usage \ + failed: 401 Unauthorized; content-type=text/plain; body={\"error\":{\"message\":\"Provided authentication \ + token is expired. Please try signing in again.\",\"code\":\"token_expired\"}} + """ + + #expect(store.userFacingError(for: .codex) == "Codex session expired. Sign in again.") + } + + @Test + func `transport codex error is sanitized`() { + let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-transport") + store.errors[.codex] = + "Codex connection failed: failed to fetch codex rate limits: " + + "GET https://chatgpt.com/backend-api/wham/usage failed: 500" + + #expect(store.userFacingError(for: .codex) == "Codex usage is temporarily unavailable. Try refreshing.") + } + + @Test + func `cached credits failure preserves cached suffix while sanitizing body`() { + let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-cached-credits") + store.lastCreditsError = + "Last Codex credits refresh failed: Codex connection failed: failed to fetch codex rate limits: " + + "GET https://chatgpt.com/backend-api/wham/usage failed: 500; body={\"error\":{}} " + + "Cached values from 2m ago." + + #expect( + store.userFacingLastCreditsError == + "Codex usage is temporarily unavailable. Try refreshing. Cached values from 2m ago.") + } + + @Test + func `browser mismatch remains unchanged`() { + let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-browser-mismatch") + store.lastOpenAIDashboardError = + "OpenAI cookies are for ratulsarna@gmail.com, not rdsarna@gmail.com. " + + "Switch chatgpt.com account, then refresh OpenAI cookies." + + #expect( + store.userFacingLastOpenAIDashboardError == + "OpenAI cookies are for ratulsarna@gmail.com, not rdsarna@gmail.com. " + + "Switch chatgpt.com account, then refresh OpenAI cookies.") + } + + @Test + func `frame load interrupted becomes retry guidance`() { + let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-frame-load") + store.lastOpenAIDashboardError = "Frame load interrupted" + + #expect( + store.userFacingLastOpenAIDashboardError == + "OpenAI web refresh was interrupted. Refresh OpenAI cookies and try again.") + } + + @Test + func `non codex providers keep raw errors`() { + let store = self.makeUsageStore(suite: "CodexUserFacingErrorTests-non-codex") + store.errors[.claude] = "Claude probe failed with debug detail" + + #expect(store.userFacingError(for: .claude) == "Claude probe failed with debug detail") + } + + @Test + func `providers pane codex model uses sanitized values`() { + let settings = self.makeSettingsStore(suite: "CodexUserFacingErrorTests-pane-model") + let store = self.makeUsageStore(settings: settings) + store.errors[.codex] = + "Codex connection failed: failed to fetch codex rate limits: " + + "GET https://chatgpt.com/backend-api/wham/usage failed: 500" + store.lastCreditsError = + "Last Codex credits refresh failed: Codex connection failed: failed to fetch codex rate limits: " + + "GET https://chatgpt.com/backend-api/wham/usage failed: 500 " + + "Cached values from 1m ago." + store.lastOpenAIDashboardError = "Frame load interrupted" + + let pane = ProvidersPane(settings: settings, store: store) + let model = pane._test_menuCardModel(for: .codex) + + #expect(model.subtitleText == "Codex usage is temporarily unavailable. Try refreshing.") + #expect( + model.creditsHintText == + "OpenAI web refresh was interrupted. Refresh OpenAI cookies and try again.") + #expect( + model.creditsHintCopyText == + "OpenAI web refresh was interrupted. Refresh OpenAI cookies and try again.") + #expect( + model.creditsText == "Codex usage is temporarily unavailable. Try refreshing. Cached values from 1m ago.") + } + + @Test + func `providers pane codex error display keeps raw full text for copy`() { + let settings = self.makeSettingsStore(suite: "CodexUserFacingErrorTests-pane-error-display") + let store = self.makeUsageStore(settings: settings) + let raw = + "Codex connection failed: failed to fetch codex rate limits: " + + "GET https://chatgpt.com/backend-api/wham/usage failed: 500; body={\"error\":{}}" + store.errors[.codex] = raw + + let pane = ProvidersPane(settings: settings, store: store) + let display = pane._test_providerErrorDisplay(for: .codex) + + #expect(display?.preview == "Codex usage is temporarily unavailable. Try refreshing.") + #expect(display?.full == raw) + } + + private func makeUsageStore(suite: String) -> UsageStore { + let settings = self.makeSettingsStore(suite: suite) + return self.makeUsageStore(settings: settings) + } + + private func makeUsageStore(settings: SettingsStore) -> UsageStore { + UsageStore( + fetcher: UsageFetcher(environment: [:]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + } + + private func makeSettingsStore(suite: String) -> SettingsStore { + let defaults = UserDefaults(suiteName: suite)! + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + return SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore(), + codexCookieStore: InMemoryCookieHeaderStore(), + claudeCookieStore: InMemoryCookieHeaderStore(), + cursorCookieStore: InMemoryCookieHeaderStore(), + opencodeCookieStore: InMemoryCookieHeaderStore(), + factoryCookieStore: InMemoryCookieHeaderStore(), + minimaxCookieStore: InMemoryMiniMaxCookieStore(), + minimaxAPITokenStore: InMemoryMiniMaxAPITokenStore(), + kimiTokenStore: InMemoryKimiTokenStore(), + kimiK2TokenStore: InMemoryKimiK2TokenStore(), + augmentCookieStore: InMemoryCookieHeaderStore(), + ampCookieStore: InMemoryCookieHeaderStore(), + copilotTokenStore: InMemoryCopilotTokenStore(), + tokenAccountStore: InMemoryTokenAccountStore()) + } +} diff --git a/Tests/CodexBarTests/CodexbarTests.swift b/Tests/CodexBarTests/CodexbarTests.swift index 026ba43bc..b5aae529e 100644 --- a/Tests/CodexBarTests/CodexbarTests.swift +++ b/Tests/CodexBarTests/CodexbarTests.swift @@ -194,7 +194,7 @@ struct CodexBarTests { } @Test - func `account info parses auth token`() throws { + func `account info parses snake case auth token`() throws { let tmp = try FileManager.default.url( for: .itemReplacementDirectory, in: .userDomainMask, @@ -203,7 +203,28 @@ struct CodexBarTests { defer { try? FileManager.default.removeItem(at: tmp) } let token = Self.fakeJWT(email: "user@example.com", plan: "pro") - let auth = ["tokens": ["idToken": token]] + let auth = ["tokens": ["id_token": token, "access_token": "access", "refresh_token": "refresh"]] + let data = try JSONSerialization.data(withJSONObject: auth) + let authURL = tmp.appendingPathComponent("auth.json") + try data.write(to: authURL) + + let fetcher = UsageFetcher(environment: ["CODEX_HOME": tmp.path]) + let account = fetcher.loadAccountInfo() + #expect(account.email == "user@example.com") + #expect(account.plan == "pro") + } + + @Test + func `account info parses legacy camel case auth token`() throws { + let tmp = try FileManager.default.url( + for: .itemReplacementDirectory, + in: .userDomainMask, + appropriateFor: URL(fileURLWithPath: NSTemporaryDirectory()), + create: true) + defer { try? FileManager.default.removeItem(at: tmp) } + + let token = Self.fakeJWT(email: "user@example.com", plan: "pro") + let auth = ["tokens": ["idToken": token, "accessToken": "access", "refreshToken": "refresh"]] let data = try JSONSerialization.data(withJSONObject: auth) let authURL = tmp.appendingPathComponent("auth.json") try data.write(to: authURL) diff --git a/Tests/CodexBarTests/CookieHeaderCacheTests.swift b/Tests/CodexBarTests/CookieHeaderCacheTests.swift index 23dfec68c..3a8a7b50d 100644 --- a/Tests/CodexBarTests/CookieHeaderCacheTests.swift +++ b/Tests/CodexBarTests/CookieHeaderCacheTests.swift @@ -25,6 +25,54 @@ struct CookieHeaderCacheTests { #expect(loaded?.storedAt == storedAt) } + @Test + func `stores separate codex entries per managed account scope`() { + KeychainCacheStore.setTestStoreForTesting(true) + defer { KeychainCacheStore.setTestStoreForTesting(false) } + + let provider: UsageProvider = .codex + let accountA = UUID() + let accountB = UUID() + + CookieHeaderCache.store( + provider: provider, + scope: .managedAccount(accountA), + cookieHeader: "auth=account-a", + sourceLabel: "Chrome") + CookieHeaderCache.store( + provider: provider, + scope: .managedAccount(accountB), + cookieHeader: "auth=account-b", + sourceLabel: "Safari") + defer { + CookieHeaderCache.clear(provider: provider, scope: .managedAccount(accountA)) + CookieHeaderCache.clear(provider: provider, scope: .managedAccount(accountB)) + } + + #expect(CookieHeaderCache.load(provider: provider, scope: .managedAccount(accountA))? + .cookieHeader == "auth=account-a") + #expect(CookieHeaderCache.load(provider: provider, scope: .managedAccount(accountB))? + .cookieHeader == "auth=account-b") + #expect(CookieHeaderCache.load(provider: provider)?.cookieHeader == nil) + } + + @Test + func `provider global scope remains available without managed account`() { + KeychainCacheStore.setTestStoreForTesting(true) + defer { KeychainCacheStore.setTestStoreForTesting(false) } + + let provider: UsageProvider = .codex + + CookieHeaderCache.store( + provider: provider, + cookieHeader: "auth=system", + sourceLabel: "Chrome") + defer { CookieHeaderCache.clear(provider: provider) } + + #expect(CookieHeaderCache.load(provider: provider)?.cookieHeader == "auth=system") + #expect(CookieHeaderCache.load(provider: provider, scope: .managedAccount(UUID())) == nil) + } + @Test func `migrates legacy file to keychain`() { KeychainCacheStore.setTestStoreForTesting(true) diff --git a/Tests/CodexBarTests/ManagedCodexAccountCoordinatorTests.swift b/Tests/CodexBarTests/ManagedCodexAccountCoordinatorTests.swift new file mode 100644 index 000000000..4d24a69b3 --- /dev/null +++ b/Tests/CodexBarTests/ManagedCodexAccountCoordinatorTests.swift @@ -0,0 +1,118 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@Suite(.serialized) +@MainActor +struct ManagedCodexAccountCoordinatorTests { + @Test + func `coordinator exposes in flight state and rejects overlapping managed authentication`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + + let existingAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let runner = BlockingManagedCodexLoginRunner() + let service = ManagedCodexAccountService( + store: InMemoryManagedCodexAccountStoreForCoordinatorTests( + accounts: ManagedCodexAccountSet(version: 1, accounts: [])), + homeFactory: CoordinatorTestManagedCodexHomeFactory(root: root), + loginRunner: runner, + identityReader: CoordinatorStubManagedCodexIdentityReader(email: "user@example.com")) + let coordinator = ManagedCodexAccountCoordinator(service: service) + + let authTask = Task { try await coordinator.authenticateManagedAccount(existingAccountID: existingAccountID) } + await runner.waitUntilStarted() + + #expect(coordinator.isAuthenticatingManagedAccount) + #expect(coordinator.authenticatingManagedAccountID == existingAccountID) + + await #expect(throws: ManagedCodexAccountCoordinatorError.authenticationInProgress) { + try await coordinator.authenticateManagedAccount() + } + + await runner.resume() + let account = try await authTask.value + + #expect(account.email == "user@example.com") + #expect(coordinator.isAuthenticatingManagedAccount == false) + #expect(coordinator.authenticatingManagedAccountID == nil) + } +} + +private actor BlockingManagedCodexLoginRunner: ManagedCodexLoginRunning { + private var waiters: [CheckedContinuation] = [] + private var startedWaiters: [CheckedContinuation] = [] + private var didStart = false + + func run(homePath _: String, timeout _: TimeInterval) async -> CodexLoginRunner.Result { + self.didStart = true + self.startedWaiters.forEach { $0.resume() } + self.startedWaiters.removeAll() + return await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + } + + func waitUntilStarted() async { + if self.didStart { return } + await withCheckedContinuation { continuation in + self.startedWaiters.append(continuation) + } + } + + func resume() { + let result = CodexLoginRunner.Result(outcome: .success, output: "ok") + self.waiters.forEach { $0.resume(returning: result) } + self.waiters.removeAll() + } +} + +private final class InMemoryManagedCodexAccountStoreForCoordinatorTests: ManagedCodexAccountStoring, +@unchecked Sendable { + var snapshot: ManagedCodexAccountSet + + init(accounts: ManagedCodexAccountSet) { + self.snapshot = accounts + } + + func loadAccounts() throws -> ManagedCodexAccountSet { + self.snapshot + } + + func storeAccounts(_ accounts: ManagedCodexAccountSet) throws { + self.snapshot = accounts + } + + func ensureFileExists() throws -> URL { + FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + } +} + +private final class CoordinatorTestManagedCodexHomeFactory: ManagedCodexHomeProducing, @unchecked Sendable { + let root: URL + + init(root: URL) { + self.root = root + } + + func makeHomeURL() -> URL { + self.root.appendingPathComponent(UUID().uuidString, isDirectory: true) + } + + func validateManagedHomeForDeletion(_ url: URL) throws { + try ManagedCodexHomeFactory(root: self.root).validateManagedHomeForDeletion(url) + } +} + +private final class CoordinatorStubManagedCodexIdentityReader: ManagedCodexIdentityReading, @unchecked Sendable { + let email: String + + init(email: String) { + self.email = email + } + + func loadAccountInfo(homePath _: String) throws -> AccountInfo { + AccountInfo(email: self.email, plan: "Pro") + } +} diff --git a/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift b/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift new file mode 100644 index 000000000..67256f679 --- /dev/null +++ b/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift @@ -0,0 +1,428 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@Suite(.serialized) +@MainActor +struct ManagedCodexAccountServiceTests { + @Test + func `upsert preserves uuid for matching canonical email`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + + let fileURL = root.appendingPathComponent("managed.json", isDirectory: false) + let store = FileManagedCodexAccountStore(fileURL: fileURL, fileManager: .default) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails(["user@example.com", "user@example.com"])) + + let first = try await service.authenticateManagedAccount() + let second = try await service.authenticateManagedAccount() + let snapshot = try store.loadAccounts() + + #expect(first.id == second.id) + #expect(second.email == "user@example.com") + #expect(snapshot.accounts.count == 1) + #expect(second.managedHomePath.hasPrefix(root.standardizedFileURL.path + "/")) + } + + @Test + func `new authentication appends managed account without implicit selection side effect`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + + let firstID = try #require(UUID(uuidString: "11111111-1111-1111-1111-111111111111")) + let firstAccount = ManagedCodexAccount( + id: firstID, + email: "first@example.com", + managedHomePath: root.appendingPathComponent("accounts/first", isDirectory: true).path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet( + version: 1, + accounts: [firstAccount])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails(["second@example.com"])) + + let authenticated = try await service.authenticateManagedAccount() + + #expect(store.snapshot.accounts.count == 2) + #expect(authenticated.email == "second@example.com") + } + + @Test + func `reauth keeps previous home when store write fails`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let existingHome = root.appendingPathComponent("accounts/existing", isDirectory: true) + try FileManager.default.createDirectory(at: existingHome, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + let existingAccountID = try #require(UUID(uuidString: "11111111-2222-3333-4444-555555555555")) + let existingAccount = ManagedCodexAccount( + id: existingAccountID, + email: "user@example.com", + managedHomePath: existingHome.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let store = FailingManagedCodexAccountStore( + accounts: ManagedCodexAccountSet( + version: 1, + accounts: [existingAccount])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails(["user@example.com"])) + + await #expect(throws: TestManagedCodexAccountStoreError.writeFailed) { + try await service.authenticateManagedAccount() + } + + let newHome = root.appendingPathComponent("accounts/account-1", isDirectory: true) + #expect(FileManager.default.fileExists(atPath: existingHome.path)) + #expect(FileManager.default.fileExists(atPath: newHome.path) == false) + #expect(store.snapshot.accounts.count == 1) + #expect(store.snapshot.accounts.first?.managedHomePath == existingHome.path) + } + + @Test + func `reauth reconciles by canonical email before existing account id`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let alphaHome = root.appendingPathComponent("accounts/alpha", isDirectory: true) + let betaHome = root.appendingPathComponent("accounts/beta", isDirectory: true) + try FileManager.default.createDirectory(at: alphaHome, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: betaHome, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + let alphaID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let betaID = try #require(UUID(uuidString: "BBBBBBBB-CCCC-DDDD-EEEE-222222222222")) + let alphaAccount = ManagedCodexAccount( + id: alphaID, + email: "alpha@example.com", + managedHomePath: alphaHome.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let betaAccount = ManagedCodexAccount( + id: betaID, + email: "beta@example.com", + managedHomePath: betaHome.path, + createdAt: 2, + updatedAt: 2, + lastAuthenticatedAt: 2) + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet( + version: 1, + accounts: [alphaAccount, betaAccount])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails(["BETA@example.com"])) + + let account = try await service.authenticateManagedAccount(existingAccountID: alphaAccount.id) + + let storedAlpha = try #require(store.snapshot.account(id: alphaAccount.id)) + let storedBeta = try #require(store.snapshot.account(id: betaAccount.id)) + #expect(account.id == betaAccount.id) + #expect(store.snapshot.accounts.count == 2) + #expect(storedAlpha.email == "alpha@example.com") + #expect(storedAlpha.managedHomePath == alphaHome.path) + #expect(storedBeta.email == "beta@example.com") + #expect(storedBeta.managedHomePath.hasPrefix(root.standardizedFileURL.path + "/")) + #expect(storedBeta.managedHomePath != betaHome.path) + #expect(FileManager.default.fileExists(atPath: alphaHome.path)) + #expect(FileManager.default.fileExists(atPath: betaHome.path) == false) + #expect(FileManager.default.fileExists(atPath: storedBeta.managedHomePath)) + } + + @Test + func `auth failure cleanup uses managed root safety check`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let outsideHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { + try? FileManager.default.removeItem(at: root) + try? FileManager.default.removeItem(at: outsideHome) + } + + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet(version: 1, accounts: [])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: UnsafeManagedCodexHomeFactory(root: root, homeURL: outsideHome), + loginRunner: StubManagedCodexLoginRunner( + result: CodexLoginRunner.Result(outcome: .failed(status: 1), output: "nope")), + identityReader: StubManagedCodexIdentityReader.emails([])) + + await #expect(throws: ManagedCodexAccountServiceError.loginFailed) { + try await service.authenticateManagedAccount() + } + + #expect(FileManager.default.fileExists(atPath: outsideHome.path)) + #expect(store.snapshot.accounts.isEmpty) + } + + @Test + func `remove deletes managed home under managed root`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let home = root.appendingPathComponent("accounts/account-a", isDirectory: true) + try FileManager.default.createDirectory(at: home, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + let accountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-EEEEEEEEEEEE")) + let account = ManagedCodexAccount( + id: accountID, + email: "user@example.com", + managedHomePath: home.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet(version: 1, accounts: [account])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails([])) + + try await service.removeManagedAccount(id: account.id) + + #expect(store.snapshot.accounts.isEmpty) + #expect(FileManager.default.fileExists(atPath: home.path) == false) + } + + @Test + func `remove keeps remaining managed account records`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let firstHome = root.appendingPathComponent("accounts/account-a", isDirectory: true) + let secondHome = root.appendingPathComponent("accounts/account-b", isDirectory: true) + try FileManager.default.createDirectory(at: firstHome, withIntermediateDirectories: true) + try FileManager.default.createDirectory(at: secondHome, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + let firstID = try #require(UUID(uuidString: "AAAAAAAA-1111-1111-1111-111111111111")) + let secondID = try #require(UUID(uuidString: "BBBBBBBB-2222-2222-2222-222222222222")) + let first = ManagedCodexAccount( + id: firstID, + email: "first@example.com", + managedHomePath: firstHome.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let second = ManagedCodexAccount( + id: secondID, + email: "second@example.com", + managedHomePath: secondHome.path, + createdAt: 2, + updatedAt: 2, + lastAuthenticatedAt: 2) + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet( + version: 1, + accounts: [first, second])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails([])) + + try await service.removeManagedAccount(id: second.id) + + #expect(store.snapshot.accounts.count == 1) + #expect(store.snapshot.accounts.first?.id == first.id) + #expect(FileManager.default.fileExists(atPath: secondHome.path) == false) + } + + @Test + func `remove keeps persisted account when store write fails`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let home = root.appendingPathComponent("accounts/account-a", isDirectory: true) + try FileManager.default.createDirectory(at: home, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: root) } + + let accountID = try #require(UUID(uuidString: "CCCCCCCC-DDDD-EEEE-FFFF-000000000000")) + let account = ManagedCodexAccount( + id: accountID, + email: "user@example.com", + managedHomePath: home.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let store = FailingManagedCodexAccountStore( + accounts: ManagedCodexAccountSet(version: 1, accounts: [account])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails([])) + + await #expect(throws: TestManagedCodexAccountStoreError.writeFailed) { + try await service.removeManagedAccount(id: account.id) + } + + #expect(store.snapshot.accounts.count == 1) + #expect(store.snapshot.accounts.first?.managedHomePath == home.path) + #expect(FileManager.default.fileExists(atPath: home.path)) + } + + @Test + func `remove fails closed for home outside managed root`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let outsideRoot = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + try FileManager.default.createDirectory(at: outsideRoot, withIntermediateDirectories: true) + defer { + try? FileManager.default.removeItem(at: root) + try? FileManager.default.removeItem(at: outsideRoot) + } + + let accountID = try #require(UUID(uuidString: "BBBBBBBB-CCCC-DDDD-EEEE-FFFFFFFFFFFF")) + let account = ManagedCodexAccount( + id: accountID, + email: "user@example.com", + managedHomePath: outsideRoot.path, + createdAt: 1, + updatedAt: 1, + lastAuthenticatedAt: 1) + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet(version: 1, accounts: [account])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.emails([])) + + await #expect(throws: ManagedCodexAccountServiceError.unsafeManagedHome(account.managedHomePath)) { + try await service.removeManagedAccount(id: account.id) + } + + #expect(store.snapshot.accounts.count == 1) + #expect(FileManager.default.fileExists(atPath: outsideRoot.path)) + } +} + +private final class InMemoryManagedCodexAccountStore: ManagedCodexAccountStoring, @unchecked Sendable { + var snapshot: ManagedCodexAccountSet + + init(accounts: ManagedCodexAccountSet) { + self.snapshot = accounts + } + + func loadAccounts() throws -> ManagedCodexAccountSet { + self.snapshot + } + + func storeAccounts(_ accounts: ManagedCodexAccountSet) throws { + self.snapshot = accounts + } + + func ensureFileExists() throws -> URL { + FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + } +} + +private final class FailingManagedCodexAccountStore: ManagedCodexAccountStoring, @unchecked Sendable { + var snapshot: ManagedCodexAccountSet + + init(accounts: ManagedCodexAccountSet) { + self.snapshot = accounts + } + + func loadAccounts() throws -> ManagedCodexAccountSet { + self.snapshot + } + + func storeAccounts(_ accounts: ManagedCodexAccountSet) throws { + _ = accounts + throw TestManagedCodexAccountStoreError.writeFailed + } + + func ensureFileExists() throws -> URL { + FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + } +} + +private final class TestManagedCodexHomeFactory: ManagedCodexHomeProducing, @unchecked Sendable { + let root: URL + private let lock = NSLock() + private var index: Int = 0 + + init(root: URL) { + self.root = root + } + + private func nextPathComponent() -> String { + self.lock.lock() + defer { self.lock.unlock() } + self.index += 1 + return "accounts/account-\(self.index)" + } + + func makeHomeURL() -> URL { + self.root.appendingPathComponent(self.nextPathComponent(), isDirectory: true) + } + + func validateManagedHomeForDeletion(_ url: URL) throws { + try ManagedCodexHomeFactory(root: self.root).validateManagedHomeForDeletion(url) + } +} + +private struct UnsafeManagedCodexHomeFactory: ManagedCodexHomeProducing, Sendable { + let root: URL + let homeURL: URL + + func makeHomeURL() -> URL { + self.homeURL + } + + func validateManagedHomeForDeletion(_ url: URL) throws { + try ManagedCodexHomeFactory(root: self.root).validateManagedHomeForDeletion(url) + } +} + +private struct StubManagedCodexLoginRunner: ManagedCodexLoginRunning, Sendable { + let result: CodexLoginRunner.Result + + func run(homePath: String, timeout: TimeInterval) async -> CodexLoginRunner.Result { + self.result + } + + static let success = StubManagedCodexLoginRunner( + result: CodexLoginRunner.Result(outcome: .success, output: "ok")) +} + +private enum TestManagedCodexAccountStoreError: Error, Equatable { + case writeFailed +} + +private final class StubManagedCodexIdentityReader: ManagedCodexIdentityReading, @unchecked Sendable { + private let lock = NSLock() + private var emails: [String] + + init(emails: [String]) { + self.emails = emails + } + + func loadAccountInfo(homePath: String) throws -> AccountInfo { + self.lock.lock() + defer { self.lock.unlock() } + let email = self.emails.isEmpty ? nil : self.emails.removeFirst() + return AccountInfo(email: email, plan: "Pro") + } + + static func emails(_ emails: [String]) -> StubManagedCodexIdentityReader { + StubManagedCodexIdentityReader(emails: emails) + } +} diff --git a/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift b/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift new file mode 100644 index 000000000..9925ad0c6 --- /dev/null +++ b/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift @@ -0,0 +1,289 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@Test +func `FileManagedCodexAccountStore round trip`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let firstID = UUID() + let secondID = UUID() + let firstAccount = ManagedCodexAccount( + id: firstID, + email: " FIRST@Example.COM ", + managedHomePath: "/tmp/managed-home-1", + createdAt: 1000, + updatedAt: 2000, + lastAuthenticatedAt: 3000) + let secondAccount = ManagedCodexAccount( + id: secondID, + email: "second@example.com", + managedHomePath: "/tmp/managed-home-2", + createdAt: 4000, + updatedAt: 5000, + lastAuthenticatedAt: nil) + let payload = ManagedCodexAccountSet( + version: 1, + accounts: [firstAccount, secondAccount]) + let store = FileManagedCodexAccountStore(fileURL: fileURL) + + try store.storeAccounts(payload) + let contents = try String(contentsOf: fileURL, encoding: .utf8) + let loaded = try store.loadAccounts() + let accountsRange = try #require(contents.range(of: "\"accounts\"")) + let versionRange = try #require(contents.range(of: "\"version\"")) + + #expect(loaded.version == 1) + #expect(loaded.accounts.count == 2) + #expect(loaded.accounts[0].email == "first@example.com") + #expect(loaded.account(id: firstID)?.managedHomePath == "/tmp/managed-home-1") + #expect(loaded.account(email: "SECOND@example.com")?.id == secondID) + #expect(contents.contains("\n \"accounts\"")) + #expect(accountsRange.lowerBound < versionRange.lowerBound) + #expect(contents.contains("\"activeAccountID\"") == false) +} + +@Test +func `FileManagedCodexAccountStore missing file loads empty set`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-nil-active-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + try? FileManager.default.removeItem(at: fileURL) + + let store = FileManagedCodexAccountStore(fileURL: fileURL) + let initial = try store.loadAccounts() + + #expect(initial.version == 1) + #expect(initial.accounts.isEmpty) + + let account = ManagedCodexAccount( + id: UUID(), + email: "user@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 10, + updatedAt: 20, + lastAuthenticatedAt: nil) + let payload = ManagedCodexAccountSet( + version: 1, + accounts: [account]) + + try store.storeAccounts(payload) + let loaded = try store.loadAccounts() + + #expect(loaded.version == 1) + #expect(loaded.accounts.count == 1) + #expect(loaded.account(email: "USER@example.com")?.id == account.id) +} + +@Test +func `FileManagedCodexAccountStore canonicalizes decoded emails`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-decode-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let accountID = UUID() + let json = """ + { + "accounts" : [ + { + "createdAt" : 10, + "email" : " MIXED@Example.COM ", + "id" : "\(accountID.uuidString)", + "lastAuthenticatedAt" : null, + "managedHomePath" : "/tmp/managed-home", + "updatedAt" : 20 + } + ], + "version" : 1 + } + """ + + try json.write(to: fileURL, atomically: true, encoding: .utf8) + + let store = FileManagedCodexAccountStore(fileURL: fileURL) + let loaded = try store.loadAccounts() + + #expect(loaded.accounts.first?.email == "mixed@example.com") + #expect(loaded.account(email: "mixed@example.com")?.id == accountID) +} + +@Test +func `FileManagedCodexAccountStore drops duplicate canonical emails on load`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-duplicate-email-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let firstID = UUID() + let secondID = UUID() + let json = """ + { + "accounts" : [ + { + "createdAt" : 10, + "email" : " First@Example.com ", + "id" : "\(firstID.uuidString)", + "lastAuthenticatedAt" : null, + "managedHomePath" : "/tmp/managed-home-1", + "updatedAt" : 20 + }, + { + "createdAt" : 30, + "email" : "first@example.com", + "id" : "\(secondID.uuidString)", + "lastAuthenticatedAt" : null, + "managedHomePath" : "/tmp/managed-home-2", + "updatedAt" : 40 + } + ], + "version" : 1 + } + """ + + try json.write(to: fileURL, atomically: true, encoding: .utf8) + + let store = FileManagedCodexAccountStore(fileURL: fileURL) + let loaded = try store.loadAccounts() + + #expect(loaded.accounts.count == 1) + #expect(loaded.accounts.first?.id == firstID) + #expect(loaded.accounts.first?.managedHomePath == "/tmp/managed-home-1") +} + +@Test +func `FileManagedCodexAccountStore drops duplicate IDs on load`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-duplicate-id-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let sharedID = UUID() + let json = """ + { + "accounts" : [ + { + "createdAt" : 10, + "email" : "first@example.com", + "id" : "\(sharedID.uuidString)", + "lastAuthenticatedAt" : null, + "managedHomePath" : "/tmp/managed-home-1", + "updatedAt" : 20 + }, + { + "createdAt" : 30, + "email" : "second@example.com", + "id" : "\(sharedID.uuidString)", + "lastAuthenticatedAt" : null, + "managedHomePath" : "/tmp/managed-home-2", + "updatedAt" : 40 + } + ], + "version" : 1 + } + """ + + try json.write(to: fileURL, atomically: true, encoding: .utf8) + + let store = FileManagedCodexAccountStore(fileURL: fileURL) + let loaded = try store.loadAccounts() + + #expect(loaded.accounts.count == 1) + #expect(loaded.accounts.first?.id == sharedID) + #expect(loaded.accounts.first?.email == "first@example.com") + #expect(loaded.accounts.first?.managedHomePath == "/tmp/managed-home-1") +} + +@Test +func `FileManagedCodexAccountStore ignores legacy active account key on load`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-legacy-active-key-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let accountID = UUID() + let danglingID = UUID() + let json = """ + { + "accounts" : [ + { + "createdAt" : 10, + "email" : "user@example.com", + "id" : "\(accountID.uuidString)", + "lastAuthenticatedAt" : null, + "managedHomePath" : "/tmp/managed-home", + "updatedAt" : 20 + } + ], + "activeAccountID" : "\(danglingID.uuidString)", + "version" : 1 + } + """ + + try json.write(to: fileURL, atomically: true, encoding: .utf8) + + let store = FileManagedCodexAccountStore(fileURL: fileURL) + let loaded = try store.loadAccounts() + + #expect(loaded.accounts.count == 1) + #expect(loaded.account(id: accountID)?.email == "user@example.com") +} + +@Test +func `FileManagedCodexAccountStore rejects unsupported on disk versions`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-unsupported-version-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let accountID = UUID() + let json = """ + { + "accounts" : [ + { + "createdAt" : 10, + "email" : "user@example.com", + "id" : "\(accountID.uuidString)", + "lastAuthenticatedAt" : null, + "managedHomePath" : "/tmp/managed-home", + "updatedAt" : 20 + } + ], + "version" : 999 + } + """ + + try json.write(to: fileURL, atomically: true, encoding: .utf8) + + let store = FileManagedCodexAccountStore(fileURL: fileURL) + + #expect(throws: FileManagedCodexAccountStoreError.unsupportedVersion(999)) { + try store.loadAccounts() + } +} + +@Test +func `FileManagedCodexAccountStore normalizes stored version to current schema`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-version-normalization-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let accountID = UUID() + let account = ManagedCodexAccount( + id: accountID, + email: "user@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 10, + updatedAt: 20, + lastAuthenticatedAt: nil) + let payload = ManagedCodexAccountSet( + version: 999, + accounts: [account]) + let store = FileManagedCodexAccountStore(fileURL: fileURL) + + try store.storeAccounts(payload) + let loaded = try store.loadAccounts() + let contents = try String(contentsOf: fileURL, encoding: .utf8) + + #expect(loaded.version == FileManagedCodexAccountStore.currentVersion) + #expect(contents.contains("\"version\" : 1")) + #expect(!contents.contains("\"version\" : 999")) +} diff --git a/Tests/CodexBarTests/MenuDescriptorCodexManagedFallbackTests.swift b/Tests/CodexBarTests/MenuDescriptorCodexManagedFallbackTests.swift new file mode 100644 index 000000000..88cdc718f --- /dev/null +++ b/Tests/CodexBarTests/MenuDescriptorCodexManagedFallbackTests.swift @@ -0,0 +1,114 @@ +import CodexBarCore +import Foundation +import Testing +@testable import CodexBar + +@MainActor +struct MenuDescriptorCodexManagedFallbackTests { + @Test + func `codex account section prefers managed fallback over ambient account`() throws { + let suite = "MenuDescriptorCodexManagedFallbackTests" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + + let settings = SettingsStore( + userDefaults: defaults, + configStore: testConfigStore(suiteName: suite), + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore(), + codexCookieStore: InMemoryCookieHeaderStore(), + claudeCookieStore: InMemoryCookieHeaderStore(), + cursorCookieStore: InMemoryCookieHeaderStore(), + opencodeCookieStore: InMemoryCookieHeaderStore(), + factoryCookieStore: InMemoryCookieHeaderStore(), + minimaxCookieStore: InMemoryMiniMaxCookieStore(), + minimaxAPITokenStore: InMemoryMiniMaxAPITokenStore(), + kimiTokenStore: InMemoryKimiTokenStore(), + kimiK2TokenStore: InMemoryKimiK2TokenStore(), + augmentCookieStore: InMemoryCookieHeaderStore(), + ampCookieStore: InMemoryCookieHeaderStore(), + copilotTokenStore: InMemoryCopilotTokenStore(), + tokenAccountStore: InMemoryTokenAccountStore()) + settings.statusChecksEnabled = false + + let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { + try? FileManager.default.removeItem(at: ambientHome) + try? FileManager.default.removeItem(at: managedHome) + } + + try Self.writeCodexAuthFile(homeURL: ambientHome, email: "ambient@example.com", plan: "plus") + try Self.writeCodexAuthFile(homeURL: managedHome, email: "managed@example.com", plan: "enterprise") + settings.codexActiveSource = .managedAccount(id: UUID()) + settings._test_activeManagedCodexRemoteHomePath = managedHome.path + + let fetcher = UsageFetcher(environment: ["CODEX_HOME": ambientHome.path]) + let store = UsageStore( + fetcher: fetcher, + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow(usedPercent: 10, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: RateWindow(usedPercent: 20, windowMinutes: 10080, resetsAt: nil, resetDescription: nil), + updatedAt: Date(), + identity: nil), + provider: .codex) + + let descriptor = MenuDescriptor.build( + provider: .codex, + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updateReady: false, + includeContextualActions: false) + + let lines = descriptor.sections + .flatMap(\.entries) + .compactMap { entry -> String? in + guard case let .text(text, _) = entry else { return nil } + return text + } + + #expect(lines.contains("Account: managed@example.com")) + #expect(lines.contains("Plan: Enterprise")) + #expect(!lines.contains("Account: ambient@example.com")) + #expect(!lines.contains("Plan: Plus")) + } + + 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))." + } +} diff --git a/Tests/CodexBarTests/OpenAIDashboardNavigationDelegateTests.swift b/Tests/CodexBarTests/OpenAIDashboardNavigationDelegateTests.swift index 94bed8275..64b1991bb 100644 --- a/Tests/CodexBarTests/OpenAIDashboardNavigationDelegateTests.swift +++ b/Tests/CodexBarTests/OpenAIDashboardNavigationDelegateTests.swift @@ -10,6 +10,12 @@ struct OpenAIDashboardNavigationDelegateTests { #expect(NavigationDelegate.shouldIgnoreNavigationError(error)) } + @Test + func `ignores WebKit frame load interrupted by policy change`() { + let error = NSError(domain: "WebKitErrorDomain", code: 102) + #expect(NavigationDelegate.shouldIgnoreNavigationError(error)) + } + @Test func `does not ignore non-cancelled URL errors`() { let error = NSError(domain: NSURLErrorDomain, code: NSURLErrorTimedOut) @@ -60,6 +66,29 @@ struct OpenAIDashboardNavigationDelegateTests { } } + @MainActor + @Test + func `frame load interrupted provisional failure is ignored until finish`() { + let webView = WKWebView() + var result: Result? + let delegate = NavigationDelegate { result = $0 } + + delegate.webView( + webView, + didFailProvisionalNavigation: nil, + withError: NSError(domain: "WebKitErrorDomain", code: 102)) + #expect(result == nil) + + delegate.webView(webView, didFinish: nil) + + switch result { + case .success?: + #expect(Bool(true)) + default: + #expect(Bool(false)) + } + } + @Test func `navigation timeout fails with timed out error`() async { final class DelegateBox: @unchecked Sendable { diff --git a/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift b/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift index 5f298a9e7..4f5a1d45c 100644 --- a/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift +++ b/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift @@ -168,6 +168,27 @@ struct OpenAIDashboardWebViewCacheTests { cache.clearAllForTesting() } + @Test + func `Evicted WebView should not be reused on next acquire`() async throws { + let cache = OpenAIDashboardWebViewCache() + let store = WKWebsiteDataStore.nonPersistent() + let url = try #require(URL(string: "about:blank")) + + let lease1 = try await cache.acquire(websiteDataStore: store, usageURL: url, logger: nil) + let webView1 = lease1.webView + lease1.release() + + cache.evict(websiteDataStore: store) + + let lease2 = try await cache.acquire(websiteDataStore: store, usageURL: url, logger: nil) + let webView2 = lease2.webView + + #expect(webView1 !== webView2, "Acquire after eviction should create a fresh WebView") + + lease2.release() + cache.clearAllForTesting() + } + // MARK: - Busy WebView Tests @Test diff --git a/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift b/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift index d5542e293..18cc20714 100644 --- a/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift +++ b/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift @@ -1,5 +1,6 @@ import CodexBarCore import Foundation +import SwiftUI import Testing @testable import CodexBar @@ -74,7 +75,7 @@ struct ProvidersPaneCoverageTests { @Test func `provider detail plan row formats open router as balance`() { - let row = ProviderDetailView.planRow(provider: .openrouter, planText: "Balance: $4.61") + let row = ProviderDetailView.planRow(provider: .openrouter, planText: "Balance: $4.61") #expect(row?.label == "Balance") #expect(row?.value == "$4.61") @@ -82,12 +83,58 @@ struct ProvidersPaneCoverageTests { @Test func `provider detail plan row keeps plan label for non open router`() { - let row = ProviderDetailView.planRow(provider: .codex, planText: "Pro") + let row = ProviderDetailView.planRow(provider: .codex, planText: "Pro") #expect(row?.label == "Plan") #expect(row?.value == "Pro") } + @Test + func `codex providers pane uses managed account fallback instead of ambient account`() throws { + let settings = Self.makeSettingsStore(suite: "ProvidersPaneCoverageTests-codex-managed-fallback") + let ambientHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + let managedHome = FileManager.default.temporaryDirectory.appendingPathComponent( + UUID().uuidString, + isDirectory: true) + defer { + try? FileManager.default.removeItem(at: ambientHome) + try? FileManager.default.removeItem(at: managedHome) + } + + try Self.writeCodexAuthFile(homeURL: ambientHome, email: "ambient@example.com", plan: "plus") + try Self.writeCodexAuthFile(homeURL: managedHome, email: "managed@example.com", plan: "enterprise") + let managedAccountID = UUID() + settings.codexActiveSource = .managedAccount(id: managedAccountID) + settings._test_activeManagedCodexAccount = ManagedCodexAccount( + id: managedAccountID, + email: "managed@example.com", + managedHomePath: managedHome.path, + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + + let store = UsageStore( + fetcher: UsageFetcher(environment: ["CODEX_HOME": ambientHome.path]), + browserDetection: BrowserDetection(cacheTTL: 0), + settings: settings, + startupBehavior: .testing) + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow(usedPercent: 12, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: RateWindow(usedPercent: 34, windowMinutes: 10080, resetsAt: nil, resetDescription: nil), + updatedAt: Date(), + identity: nil), + provider: .codex) + + let pane = ProvidersPane(settings: settings, store: store) + let model = pane._test_menuCardModel(for: .codex) + + #expect(model.email == "managed@example.com") + #expect(model.planText == "Enterprise") + } + private static func makeSettingsStore(suite: String) -> SettingsStore { let defaults = UserDefaults(suiteName: suite)! defaults.removePersistentDomain(forName: suite) @@ -119,4 +166,34 @@ struct ProvidersPaneCoverageTests { browserDetection: BrowserDetection(cacheTTL: 0), settings: 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))." + } } diff --git a/Tests/CodexBarTests/SettingsStoreTests.swift b/Tests/CodexBarTests/SettingsStoreTests.swift index a7cfb81b1..b568bfcb9 100644 --- a/Tests/CodexBarTests/SettingsStoreTests.swift +++ b/Tests/CodexBarTests/SettingsStoreTests.swift @@ -707,6 +707,35 @@ struct SettingsStoreTests { #expect(didChange == true) } + @Test + func `menu observation token updates on codex active source change`() async throws { + let suite = "SettingsStoreTests-observation-codex-active-source" + let defaults = try #require(UserDefaults(suiteName: suite)) + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + + let store = SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + + var didChange = false + + withObservationTracking { + _ = store.menuObservationToken + } onChange: { + Task { @MainActor in + didChange = true + } + } + + store.codexActiveSource = .liveSystem + try? await Task.sleep(nanoseconds: 50_000_000) + + #expect(didChange == true) + } + @Test func `provider order defaults to all cases`() throws { let suite = "SettingsStoreTests-providerOrder-default" diff --git a/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift new file mode 100644 index 000000000..eaca4ee99 --- /dev/null +++ b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift @@ -0,0 +1,655 @@ +import AppKit +import CodexBarCore +import Testing +@testable import CodexBar + +@Suite(.serialized) +@MainActor +struct StatusMenuCodexSwitcherTests { + private func disableMenuCardsForTesting() { + StatusItemController.menuCardRenderingEnabled = false + StatusItemController.menuRefreshEnabled = false + } + + private func makeStatusBarForTesting() -> NSStatusBar { + let env = ProcessInfo.processInfo.environment + if env["GITHUB_ACTIONS"] == "true" || env["CI"] == "true" { + return .system + } + return NSStatusBar() + } + + private func makeSettings() -> SettingsStore { + let suite = "StatusMenuCodexSwitcherTests-\(UUID().uuidString)" + let defaults = UserDefaults(suiteName: suite)! + defaults.removePersistentDomain(forName: suite) + let configStore = testConfigStore(suiteName: suite) + return SettingsStore( + userDefaults: defaults, + configStore: configStore, + zaiTokenStore: NoopZaiTokenStore(), + syntheticTokenStore: NoopSyntheticTokenStore()) + } + + private func enableOnlyCodex(_ settings: SettingsStore) { + let registry = ProviderRegistry.shared + for provider in UsageProvider.allCases { + guard let metadata = registry.metadata[provider] else { continue } + settings.setProviderEnabled(provider: provider, metadata: metadata, enabled: provider == .codex) + } + } + + private func makeManagedAccountStoreURL(accounts: [ManagedCodexAccount]) throws -> URL { + let storeURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + let store = FileManagedCodexAccountStore(fileURL: storeURL) + try store.storeAccounts(ManagedCodexAccountSet( + version: FileManagedCodexAccountStore.currentVersion, + accounts: accounts)) + return storeURL + } + + private func codexSwitcherButtons(in menu: NSMenu) -> [NSButton] { + guard let switcherView = menu.items.compactMap({ $0.view as? CodexAccountSwitcherView }).first + else { return [] } + return self.buttons(in: switcherView).sorted { $0.title < $1.title } + } + + private func buttons(in view: NSView) -> [NSButton] { + let directButtons = view.subviews.compactMap { $0 as? NSButton } + let nestedButtons = view.subviews.flatMap { self.buttons(in: $0) } + return directButtons + nestedButtons + } + + private func menuItem(titled title: String, in menu: NSMenu) -> NSMenuItem? { + menu.items.first { $0.title == title } + } + + private func representedIDs(in menu: NSMenu) -> [String] { + menu.items.compactMap { $0.representedObject as? String } + } + + private func installBlockingCodexProvider(on store: UsageStore, blocker: BlockingStatusMenuCodexFetchStrategy) { + let baseSpec = store.providerSpecs[.codex]! + store.providerSpecs[.codex] = Self.makeCodexProviderSpec(baseSpec: baseSpec) { + try await blocker.awaitResult() + } + } + + private static func makeCodexProviderSpec( + baseSpec: ProviderSpec, + loader: @escaping @Sendable () async throws -> UsageSnapshot) -> ProviderSpec + { + let baseDescriptor = baseSpec.descriptor + let strategy = StatusMenuTestCodexFetchStrategy(loader: loader) + let descriptor = ProviderDescriptor( + id: .codex, + metadata: baseDescriptor.metadata, + branding: baseDescriptor.branding, + tokenCost: baseDescriptor.tokenCost, + fetchPlan: ProviderFetchPlan( + sourceModes: [.auto, .cli, .oauth], + pipeline: ProviderFetchPipeline { _ in [strategy] }), + cli: baseDescriptor.cli) + return ProviderSpec( + style: baseSpec.style, + isEnabled: baseSpec.isEnabled, + descriptor: descriptor, + makeFetchContext: baseSpec.makeFetchContext) + } + + @Test + func `codex menu shows account switcher and add account action for multiple visible accounts`() throws { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodex(settings) + + let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let managedAccount = ManagedCodexAccount( + id: managedAccountID, + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) + defer { + settings._test_managedCodexAccountStoreURL = nil + settings._test_liveSystemCodexAccount = nil + try? FileManager.default.removeItem(at: storeURL) + } + + settings._test_managedCodexAccountStoreURL = storeURL + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings.codexActiveSource = .liveSystem + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + + let buttons = self.codexSwitcherButtons(in: menu) + #expect(buttons.map(\.title) == ["live@example.com", "managed@example.com"]) + #expect(self.menuItem(titled: "Add Account...", in: menu) != nil) + #expect(self.menuItem(titled: "Switch Account...", in: menu) == nil) + } + + @Test + func `codex menu hides account switcher when only one visible account exists`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodex(settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "solo@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + defer { settings._test_liveSystemCodexAccount = nil } + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + + #expect(self.codexSwitcherButtons(in: menu).isEmpty) + #expect(self.menuItem(titled: "Add Account...", in: menu) != nil) + } + + @Test + func `codex menu switcher selection activates the visible managed account`() throws { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodex(settings) + + let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let managedAccount = ManagedCodexAccount( + id: managedAccountID, + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) + defer { + settings._test_managedCodexAccountStoreURL = nil + settings._test_liveSystemCodexAccount = nil + try? FileManager.default.removeItem(at: storeURL) + } + + settings._test_managedCodexAccountStoreURL = storeURL + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings.codexActiveSource = .liveSystem + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + + let managedButton = try #require(self.codexSwitcherButtons(in: menu) + .first { $0.title == "managed@example.com" }) + managedButton.performClick(nil) + + #expect(settings.codexActiveSource == .managedAccount(id: managedAccountID)) + } + + @Test + func `codex menu switcher clears stale account state on the first click`() async throws { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + settings.costUsageEnabled = false + settings.codexCookieSource = .off + self.enableOnlyCodex(settings) + + let managedAccountID = try #require(UUID(uuidString: "AAAAAAAA-BBBB-CCCC-DDDD-111111111111")) + let managedAccount = ManagedCodexAccount( + id: managedAccountID, + email: "managed@example.com", + managedHomePath: "/tmp/managed-home", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 2) + let storeURL = try self.makeManagedAccountStoreURL(accounts: [managedAccount]) + defer { + settings._test_managedCodexAccountStoreURL = nil + settings._test_liveSystemCodexAccount = nil + try? FileManager.default.removeItem(at: storeURL) + } + + settings._test_managedCodexAccountStoreURL = storeURL + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings.codexActiveSource = .liveSystem + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + store._setSnapshotForTesting( + UsageSnapshot( + primary: RateWindow(usedPercent: 30, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: "live@example.com", + accountOrganization: nil, + loginMethod: "Pro")), + provider: .codex) + store.lastCodexAccountScopedRefreshGuard = store + .currentCodexAccountScopedRefreshGuard(preferCurrentSnapshot: false) + + let blocker = BlockingStatusMenuCodexFetchStrategy() + self.installBlockingCodexProvider(on: store, blocker: blocker) + + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + + let managedButton = try #require(self.codexSwitcherButtons(in: menu) + .first { $0.title == "managed@example.com" }) + managedButton.performClick(nil) + + await blocker.waitUntilStarted() + #expect(settings.codexActiveSource == .managedAccount(id: managedAccountID)) + #expect(store.snapshots[.codex] == nil) + + await blocker.resume(with: .success( + UsageSnapshot( + primary: RateWindow(usedPercent: 9, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: nil, + updatedAt: Date(), + identity: ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: "managed@example.com", + accountOrganization: nil, + loginMethod: "Pro")))) + for _ in 0..<10 where store.snapshots[.codex]?.accountEmail(for: .codex) != "managed@example.com" { + try? await Task.sleep(for: .milliseconds(20)) + } + #expect(store.snapshots[.codex]?.accountEmail(for: .codex) == "managed@example.com") + } + + @Test + func `codex open menu redraw picks up refreshed account data`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + settings.costUsageEnabled = true + self.enableOnlyCodex(settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + defer { settings._test_liveSystemCodexAccount = nil } + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let snapshot = UsageSnapshot( + primary: RateWindow(usedPercent: 10, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: nil, + tertiary: nil, + updatedAt: Date()) + store._setSnapshotForTesting(snapshot, provider: .codex) + store._setTokenSnapshotForTesting(nil, provider: .codex) + + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + #expect(self.representedIDs(in: menu).contains("menuCardCost") == false) + + store._setTokenSnapshotForTesting( + CostUsageTokenSnapshot( + sessionTokens: 123, + sessionCostUSD: 1.23, + last30DaysTokens: 456, + last30DaysCostUSD: 78.9, + daily: [ + CostUsageDailyReport.Entry( + date: "2025-12-23", + inputTokens: nil, + outputTokens: nil, + totalTokens: 456, + costUSD: 78.9, + modelsUsed: nil, + modelBreakdowns: nil), + ], + updatedAt: Date()), + provider: .codex) + + controller.refreshOpenMenuIfStillVisible(menu, provider: .codex) + + #expect(self.representedIDs(in: menu).contains("menuCardCost") == true) + } + + @Test + func `codex open menu redraw retries on next run loop while menu is tracking`() async { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + settings.costUsageEnabled = true + self.enableOnlyCodex(settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + defer { settings._test_liveSystemCodexAccount = nil } + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let snapshot = UsageSnapshot( + primary: RateWindow(usedPercent: 10, windowMinutes: 300, resetsAt: nil, resetDescription: nil), + secondary: nil, + tertiary: nil, + updatedAt: Date()) + store._setSnapshotForTesting(snapshot, provider: .codex) + store._setTokenSnapshotForTesting(nil, provider: .codex) + + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + #expect(self.representedIDs(in: menu).contains("menuCardCost") == false) + + controller._test_openMenuRefreshYieldOverride = { + store._setTokenSnapshotForTesting( + CostUsageTokenSnapshot( + sessionTokens: 123, + sessionCostUSD: 1.23, + last30DaysTokens: 456, + last30DaysCostUSD: 78.9, + daily: [ + CostUsageDailyReport.Entry( + date: "2025-12-23", + inputTokens: nil, + outputTokens: nil, + totalTokens: 456, + costUSD: 78.9, + modelsUsed: nil, + modelBreakdowns: nil), + ], + updatedAt: Date()), + provider: .codex) + } + + controller.refreshOpenMenuIfStillVisible(menu, provider: .codex) + await Task.yield() + await Task.yield() + + #expect(self.representedIDs(in: menu).contains("menuCardCost") == true) + } + + @Test + func `codex menu disables add account while managed authentication is in flight`() async throws { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodex(settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + defer { settings._test_liveSystemCodexAccount = nil } + + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + + let runner = BlockingManagedCodexLoginRunnerForStatusMenuTests() + let service = ManagedCodexAccountService( + store: InMemoryManagedCodexAccountStoreForStatusMenuTests(), + homeFactory: TestManagedCodexHomeFactoryForStatusMenuTests(root: root), + loginRunner: runner, + identityReader: StubManagedCodexIdentityReaderForStatusMenuTests(email: "managed@example.com")) + let coordinator = ManagedCodexAccountCoordinator(service: service) + let authTask = Task { try await coordinator.authenticateManagedAccount() } + await runner.waitUntilStarted() + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + managedCodexAccountCoordinator: coordinator, + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + + let addItem = try #require(self.menuItem(titled: "Add Account...", in: menu)) + #expect(addItem.isEnabled == false) + if #available(macOS 14.4, *) { + #expect(addItem.subtitle == "Managed Codex login in progress…") + } else { + #expect(addItem.toolTip?.contains("Managed Codex login in progress") == true) + } + + await runner.resume() + _ = try await authTask.value + } + + @Test + func `codex menu disables add account when managed store is unreadable`() { + self.disableMenuCardsForTesting() + let settings = self.makeSettings() + settings.statusChecksEnabled = false + settings.refreshFrequency = .manual + settings.mergeIcons = false + self.enableOnlyCodex(settings) + settings._test_liveSystemCodexAccount = ObservedSystemCodexAccount( + email: "live@example.com", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + settings._test_unreadableManagedCodexAccountStore = true + defer { + settings._test_liveSystemCodexAccount = nil + settings._test_unreadableManagedCodexAccountStore = false + } + + let fetcher = UsageFetcher() + let store = UsageStore(fetcher: fetcher, browserDetection: BrowserDetection(cacheTTL: 0), settings: settings) + let controller = StatusItemController( + store: store, + settings: settings, + account: fetcher.loadAccountInfo(), + updater: DisabledUpdaterController(), + preferencesSelection: PreferencesSelection(), + statusBar: self.makeStatusBarForTesting()) + + let menu = controller.makeMenu(for: .codex) + controller.menuWillOpen(menu) + + let addItem = try? #require(self.menuItem(titled: "Add Account...", in: menu)) + #expect(addItem?.isEnabled == false) + if #available(macOS 14.4, *) { + #expect(addItem?.subtitle == "Managed account storage unavailable") + } else { + #expect(addItem?.toolTip?.contains("Managed account storage unavailable") == true) + } + } +} + +private struct StatusMenuTestCodexFetchStrategy: ProviderFetchStrategy { + let loader: @Sendable () async throws -> UsageSnapshot + + var id: String { + "status-menu-test-codex" + } + + var kind: ProviderFetchKind { + .cli + } + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try await self.loader() + return self.makeResult(usage: snapshot, sourceLabel: "status-menu-test-codex") + } + + func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { + false + } +} + +private actor BlockingStatusMenuCodexFetchStrategy { + private var waiters: [CheckedContinuation, Never>] = [] + private var startedWaiters: [CheckedContinuation] = [] + private var didStart = false + + func awaitResult() async throws -> UsageSnapshot { + self.didStart = true + self.startedWaiters.forEach { $0.resume() } + self.startedWaiters.removeAll() + let result = await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + return try result.get() + } + + func waitUntilStarted() async { + if self.didStart { return } + await withCheckedContinuation { continuation in + self.startedWaiters.append(continuation) + } + } + + func resume(with result: Result) { + self.waiters.forEach { $0.resume(returning: result) } + self.waiters.removeAll() + } +} + +private actor BlockingManagedCodexLoginRunnerForStatusMenuTests: ManagedCodexLoginRunning { + private var waiters: [CheckedContinuation] = [] + private var startedWaiters: [CheckedContinuation] = [] + private var didStart = false + + func run(homePath _: String, timeout _: TimeInterval) async -> CodexLoginRunner.Result { + self.didStart = true + self.startedWaiters.forEach { $0.resume() } + self.startedWaiters.removeAll() + return await withCheckedContinuation { continuation in + self.waiters.append(continuation) + } + } + + func waitUntilStarted() async { + if self.didStart { return } + await withCheckedContinuation { continuation in + self.startedWaiters.append(continuation) + } + } + + func resume() { + let result = CodexLoginRunner.Result(outcome: .success, output: "ok") + self.waiters.forEach { $0.resume(returning: result) } + self.waiters.removeAll() + } +} + +private final class InMemoryManagedCodexAccountStoreForStatusMenuTests: ManagedCodexAccountStoring, +@unchecked Sendable { + private var snapshot = ManagedCodexAccountSet(version: 1, accounts: []) + + func loadAccounts() throws -> ManagedCodexAccountSet { + self.snapshot + } + + func storeAccounts(_ accounts: ManagedCodexAccountSet) throws { + self.snapshot = accounts + } + + func ensureFileExists() throws -> URL { + FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + } +} + +private struct TestManagedCodexHomeFactoryForStatusMenuTests: ManagedCodexHomeProducing, Sendable { + let root: URL + + func makeHomeURL() -> URL { + self.root.appendingPathComponent(UUID().uuidString, isDirectory: true) + } + + func validateManagedHomeForDeletion(_ url: URL) throws { + try ManagedCodexHomeFactory(root: self.root).validateManagedHomeForDeletion(url) + } +} + +private struct StubManagedCodexIdentityReaderForStatusMenuTests: ManagedCodexIdentityReading, Sendable { + let email: String + + func loadAccountInfo(homePath _: String) throws -> AccountInfo { + AccountInfo(email: self.email, plan: "Pro") + } +}