diff --git a/Sources/CodexBar/CodexLoginRunner.swift b/Sources/CodexBar/CodexLoginRunner.swift index 8f1f654f2..9d0b671bb 100644 --- a/Sources/CodexBar/CodexLoginRunner.swift +++ b/Sources/CodexBar/CodexLoginRunner.swift @@ -16,13 +16,14 @@ struct CodexLoginRunner { let output: String } - static func run(timeout: TimeInterval = 120) async -> 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/ManagedCodexAccountService.swift b/Sources/CodexBar/ManagedCodexAccountService.swift new file mode 100644 index 000000000..c851e98a2 --- /dev/null +++ b/Sources/CodexBar/ManagedCodexAccountService.swift @@ -0,0 +1,194 @@ +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 + } + + 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], + activeAccountID: account.id) + 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 } + let nextActive = if snapshot.activeAccountID == id { + remaining.last?.id + } else { + snapshot.activeAccountID + } + try self.store.storeAccounts(ManagedCodexAccountSet( + version: snapshot.version, + accounts: remaining, + activeAccountID: nextActive)) + + 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/MenuDescriptor.swift b/Sources/CodexBar/MenuDescriptor.swift index 05aa55fff..3d919d2bb 100644 --- a/Sources/CodexBar/MenuDescriptor.swift +++ b/Sources/CodexBar/MenuDescriptor.swift @@ -60,12 +60,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 +79,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 +309,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 +329,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 +340,7 @@ struct MenuDescriptor { provider: targetProvider, store: store, settings: store.settings, - account: account) + account: fallbackAccount) ProviderCatalog.implementation(for: targetProvider)? .appendActionMenuEntries(context: actionContext, entries: &entries) } diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index f3d5bc112..872cf84f7 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -364,7 +364,7 @@ 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), usageBarsShowUsed: self.settings.usageBarsShowUsed, diff --git a/Sources/CodexBar/ProviderRegistry.swift b/Sources/CodexBar/ProviderRegistry.swift index 1e26c4c86..882229bcd 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,20 @@ struct ProviderRegistry { env[key] = value } } + // Managed Codex selection only scopes remote account fetches such as identity, plan, + // quotas, and dashboard data. 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/CodexSettingsStore.swift b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift index 335dbf411..b1f87b114 100644 --- a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift +++ b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift @@ -2,6 +2,83 @@ import CodexBarCore import Foundation extension SettingsStore { + private enum ManagedCodexAccountStoreState { + case none + case active(ManagedCodexAccount) + case unreadable + } + + private static func failClosedManagedCodexHomePath(fileManager: FileManager = .default) -> String { + ManagedCodexHomeFactory.defaultRootURL(fileManager: fileManager) + .appendingPathComponent("managed-store-unreadable", isDirectory: true) + .path + } + + private func managedCodexAccountStoreState() -> ManagedCodexAccountStoreState { + #if DEBUG + if CodexManagedRemoteHomeTestingOverride.isUnreadable(for: self) { + return .unreadable + } + if let override = CodexManagedRemoteHomeTestingOverride.account(for: self) { + return .active(override) + } + #endif + + do { + let accounts = try FileManagedCodexAccountStore().loadAccounts() + guard let activeAccountID = accounts.activeAccountID, + let account = accounts.account(id: activeAccountID) + else { + return .none + } + return .active(account) + } catch { + return .unreadable + } + } + + var activeManagedCodexAccount: ManagedCodexAccount? { + guard case let .active(account) = self.managedCodexAccountStoreState() else { + return nil + } + return account + } + + var activeManagedCodexRemoteHomePath: String? { + #if DEBUG + if let override = CodexManagedRemoteHomeTestingOverride.homePath(for: self) { + return override + } + #endif + + switch self.managedCodexAccountStoreState() { + case let .active(account): + return account.managedHomePath + case .unreadable: + return Self.failClosedManagedCodexHomePath() + case .none: + return nil + } + } + + var activeManagedCodexCookieCacheScope: CookieHeaderCache.Scope? { + switch self.managedCodexAccountStoreState() { + case let .active(account): + .managedAccount(account.id) + case .unreadable: + .managedStoreUnreadable + case .none: + nil + } + } + + var hasUnreadableManagedCodexAccountStore: Bool { + if case .unreadable = self.managedCodexAccountStoreState() { + return true + } + return false + } + var codexUsageDataSource: CodexUsageDataSource { get { let source = self.configSnapshot.providerConfig(for: .codex)?.source @@ -23,6 +100,8 @@ extension SettingsStore { 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 +126,94 @@ extension SettingsStore { func ensureCodexCookieLoaded() {} } +#if DEBUG +private enum CodexManagedRemoteHomeTestingOverride { + private struct Override { + var account: ManagedCodexAccount? + var homePath: String? + var unreadableStore: Bool = false + } + + @MainActor + private static var values: [ObjectIdentifier: 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 + if override.account == nil, override.homePath == nil, !override.unreadableStore { + self.values.removeValue(forKey: key) + } else { + self.values[key] = override + } + } + + @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 + if override.account == nil, override.homePath == nil, !override.unreadableStore { + self.values.removeValue(forKey: key) + } else { + self.values[key] = override + } + } + + @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 + if override.account == nil, override.homePath == nil, !override.unreadableStore { + self.values.removeValue(forKey: key) + } else { + self.values[key] = override + } + } +} + +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) } + } +} +#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.hasUnreadableManagedCodexAccountStore) } private static func codexUsageDataSource(from source: ProviderSourceMode?) -> CodexUsageDataSource { diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift index 35b184897..2d73595ea 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexRefresh.swift @@ -6,10 +6,16 @@ extension UsageStore { nonisolated static let codexSnapshotWaitTimeoutSeconds: TimeInterval = 6 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 } do { - let credits = try await self.codexFetcher.loadLatestCredits( + let credits = try await self.codexCreditsFetcher().loadLatestCredits( keepCLISessionsAlive: self.settings.debugKeepCLISessionsAlive) await MainActor.run { self.credits = credits diff --git a/Sources/CodexBar/StatusItemController+Menu.swift b/Sources/CodexBar/StatusItemController+Menu.swift index 484be310a..27e9f4b77 100644 --- a/Sources/CodexBar/StatusItemController+Menu.swift +++ b/Sources/CodexBar/StatusItemController+Menu.swift @@ -1472,7 +1472,7 @@ 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), usageBarsShowUsed: self.settings.usageBarsShowUsed, diff --git a/Sources/CodexBar/UsageStore+Accessors.swift b/Sources/CodexBar/UsageStore+Accessors.swift index 98d77c6d8..634d1dbcd 100644 --- a/Sources/CodexBar/UsageStore+Accessors.swift +++ b/Sources/CodexBar/UsageStore+Accessors.swift @@ -31,7 +31,16 @@ 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() } } diff --git a/Sources/CodexBar/UsageStore+OpenAIWeb.swift b/Sources/CodexBar/UsageStore+OpenAIWeb.swift index af164cc62..aa60b3ff8 100644 --- a/Sources/CodexBar/UsageStore+OpenAIWeb.swift +++ b/Sources/CodexBar/UsageStore+OpenAIWeb.swift @@ -1,5 +1,522 @@ +import CodexBarCore import Foundation +// MARK: - OpenAI web lifecycle + +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) } + } + + 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) + } + + func applyOpenAIDashboardFailure(message: String) async { + await MainActor.run { + if self.settings.hasUnreadableManagedCodexAccountStore { + self.failClosedOpenAIDashboardSnapshot() + self.lastOpenAIDashboardError = message + return + } + 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?) async { + 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() async { + 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: " ") + if self.settings.hasUnreadableManagedCodexAccountStore { + self.failClosedOpenAIDashboardSnapshot() + return + } + self.openAIDashboard = self.lastOpenAIDashboardSnapshot + self.openAIDashboardRequiresLogin = true + } + } + + private func failClosedOpenAIDashboardSnapshot() { + self.openAIDashboard = nil + self.lastOpenAIDashboardSnapshot = nil + self.openAIDashboardRequiresLogin = true + } + + func refreshOpenAIDashboardIfNeeded(force: Bool = false) async { + guard self.isEnabled(.codex), self.settings.codexCookieSource.isEnabled else { + self.resetOpenAIWebState() + return + } + if self.settings.hasUnreadableManagedCodexAccountStore { + await self.failClosedRefreshForUnreadableManagedCodexStore() + 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 self.applyOpenAIDashboardMismatchFailure(signedInEmail: signedIn, expectedEmail: normalized) + 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 self.applyOpenAIDashboardLoginRequiredFailure() + } 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 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: " ") + } + } + + func importOpenAIDashboardCookiesIfNeeded(targetEmail: String?, force: Bool) async -> String? { + if self.settings.hasUnreadableManagedCodexAccountStore { + return await self.failClosedForUnreadableManagedCodexStore() + } + + 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: " ") + : [ + "Browser cookies do not match Codex account (\(normalizedTarget ?? "unknown")).", + "Found \(foundText).", + ].joined(separator: " ") + 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.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? { + if self.settings.hasUnreadableManagedCodexAccountStore { + return nil + } + + let managed = self.settings.activeManagedCodexAccount?.email + .trimmingCharacters(in: .whitespacesAndNewlines) + if let managed, !managed.isEmpty { return managed } + + 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 + } + + func codexCookieCacheScopeForOpenAIWeb() -> CookieHeaderCache.Scope? { + self.settings.activeManagedCodexCookieCacheScope + } +} + // MARK: - OpenAI web error messaging extension UsageStore { 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.swift b/Sources/CodexBar/UsageStore.swift index 47efc5b63..58328be2b 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -120,11 +120,17 @@ final class UsageStore { var historicalPaceRevision: Int = 0 @ObservationIgnored var lastCreditsSnapshot: CreditsSnapshot? @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 openAIWebAccountDidChange: Bool = false + @ObservationIgnored var _test_openAIDashboardCookieImportOverride: (@MainActor ( + String?, + Bool, + ProviderCookieSource, + CookieHeaderCache.Scope?, + @escaping (String) -> Void) async throws -> OpenAIDashboardBrowserCookieImporter.ImportResult)? @ObservationIgnored let codexFetcher: UsageFetcher @ObservationIgnored let claudeFetcher: any ClaudeUsageFetching @@ -134,11 +140,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 +311,7 @@ final class UsageStore { self.providerMetadata[provider]! } - private var codexBrowserCookieOrder: BrowserCookieImportOrder { + var codexBrowserCookieOrder: BrowserCookieImportOrder { self.metadata(for: .codex).browserCookieOrder ?? Browser.defaultImportOrder } @@ -630,449 +636,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 +1131,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..9c2dbcbd6 --- /dev/null +++ b/Sources/CodexBarCore/CodexManagedAccounts.swift @@ -0,0 +1,91 @@ +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 let activeAccountID: UUID? + + public init(version: Int, accounts: [ManagedCodexAccount], activeAccountID: UUID?) { + let sanitizedAccounts = Self.sanitizedAccounts(accounts) + self.version = version + self.accounts = sanitizedAccounts + self.activeAccountID = Self.validatedActiveAccountID(activeAccountID, accounts: sanitizedAccounts) + } + + 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), + activeAccountID: container.decodeIfPresent(UUID.self, forKey: .activeAccountID)) + } + + private static func validatedActiveAccountID(_ activeAccountID: UUID?, accounts: [ManagedCodexAccount]) -> UUID? { + guard let activeAccountID else { return nil } + return accounts.contains { $0.id == activeAccountID } ? activeAccountID : nil + } + + 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/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..78a53478d --- /dev/null +++ b/Sources/CodexBarCore/ManagedCodexAccountStore.swift @@ -0,0 +1,79 @@ +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, + activeAccountID: accounts.activeAccountID) + 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: [], activeAccountID: nil) + } + + 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..105aa83cc 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 @@ -669,10 +670,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 +799,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 +809,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/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..a5ae38511 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 } @@ -105,8 +105,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), @@ -144,3 +147,11 @@ public enum CodexOAuthCredentialsStore { return formatter.date(from: value) } } + +#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..5e6ca48f3 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift @@ -9,10 +9,16 @@ public struct CodexWebDashboardStrategy: ProviderFetchStrategy { public init() {} public func isAvailable(_ context: ProviderFetchContext) async -> Bool { - context.sourceMode.usesWeb + context.sourceMode.usesWeb && !Self.managedAccountStoreIsUnreadable(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 + } + // Ensure AppKit is initialized before using WebKit in a CLI. await MainActor.run { _ = NSApplication.shared @@ -41,6 +47,10 @@ public struct CodexWebDashboardStrategy: ProviderFetchStrategy { _ = error return true } + + private static func managedAccountStoreIsUnreadable(_ context: ProviderFetchContext) -> Bool { + context.settings?.codex?.managedAccountStoreUnreadable == true + } } private struct OpenAIWebCodexResult { diff --git a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift index c5d9af4f7..1bbe673fb 100644 --- a/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift +++ b/Sources/CodexBarCore/Providers/ProviderSettingsSnapshot.swift @@ -44,15 +44,18 @@ public struct ProviderSettingsSnapshot: Sendable { public let usageDataSource: CodexUsageDataSource public let cookieSource: ProviderCookieSource public let manualCookieHeader: String? + public let managedAccountStoreUnreadable: Bool public init( usageDataSource: CodexUsageDataSource, cookieSource: ProviderCookieSource, - manualCookieHeader: String?) + manualCookieHeader: String?, + managedAccountStoreUnreadable: Bool = false) { self.usageDataSource = usageDataSource self.cookieSource = cookieSource self.manualCookieHeader = manualCookieHeader + self.managedAccountStoreUnreadable = managedAccountStoreUnreadable } } diff --git a/Sources/CodexBarCore/UsageFetcher.swift b/Sources/CodexBarCore/UsageFetcher.swift index 98859b18e..fb2ccc9c8 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,7 +652,8 @@ 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") + let authURL = CodexHomeScope + .ambientHomeURL(env: self.environment) .appendingPathComponent("auth.json") guard let data = try? Data(contentsOf: authURL), let auth = try? JSONDecoder().decode(AuthFile.self, from: data), diff --git a/Tests/CodexBarTests/CLIWebFallbackTests.swift b/Tests/CodexBarTests/CLIWebFallbackTests.swift index 10fcf5396..a52560e7f 100644 --- a/Tests/CodexBarTests/CLIWebFallbackTests.swift +++ b/Tests/CodexBarTests/CLIWebFallbackTests.swift @@ -62,6 +62,20 @@ 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 `claude falls back when no session key`() { let context = self.makeContext() diff --git a/Tests/CodexBarTests/CodexManagedOpenAIWebTests.swift b/Tests/CodexBarTests/CodexManagedOpenAIWebTests.swift new file mode 100644 index 000000000..b67169cb1 --- /dev/null +++ b/Tests/CodexBarTests/CodexManagedOpenAIWebTests.swift @@ -0,0 +1,199 @@ +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 + 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 `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 + 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 `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?.contains("do not match Codex account") == 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) + } + + 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()) + } +} diff --git a/Tests/CodexBarTests/CodexManagedRoutingTests.swift b/Tests/CodexBarTests/CodexManagedRoutingTests.swift new file mode 100644 index 000000000..c5c22f18b --- /dev/null +++ b/Tests/CodexBarTests/CodexManagedRoutingTests.swift @@ -0,0 +1,257 @@ +import Foundation +import Testing +@testable import CodexBar +@testable import CodexBarCore + +@Suite(.serialized) +@MainActor +struct CodexManagedRoutingTests { + @Test + func `provider registry injects active managed home into codex env only`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-registry") + let managedHomePath = "/tmp/codex-managed-home" + settings._test_activeManagedCodexRemoteHomePath = managedHomePath + + 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"] == managedHomePath) + #expect(claudeEnv["CODEX_HOME"] == nil) + } + + @Test + func `provider registry fails closed when managed account store is unreadable`() { + let settings = self.makeSettingsStore(suite: "CodexManagedRoutingTests-unreadable-store") + settings._test_unreadableManagedCodexAccountStore = true + + 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 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 + try self.writeCodexAuthFile(homeURL: managedHome, email: "managed@example.com", plan: "pro") + + 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 + try self.writeCodexAuthFile(homeURL: managedHome, email: "token@example.com", plan: "team") + + 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 + try self.writeCodexAuthFile(homeURL: managedHome, email: "credits@example.com", plan: "enterprise") + + 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": [ + "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/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/ManagedCodexAccountServiceTests.swift b/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift new file mode 100644 index 000000000..db0f8fda8 --- /dev/null +++ b/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift @@ -0,0 +1,440 @@ +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(snapshot.activeAccountID == second.id) + #expect(second.managedHomePath.hasPrefix(root.standardizedFileURL.path + "/")) + } + + @Test + func `new authentication becomes active managed account`() 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], + activeAccountID: firstAccount.id)) + 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(store.snapshot.activeAccountID == authenticated.id) + #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], + activeAccountID: existingAccount.id)) + 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) + #expect(store.snapshot.activeAccountID == existingAccount.id) + } + + @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], + activeAccountID: alphaAccount.id)) + 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: [], activeAccountID: nil)) + 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) + #expect(store.snapshot.activeAccountID == nil) + } + + @Test + func `remove deletes managed home under managed root and clears active account`() 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], activeAccountID: account.id)) + 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.activeAccountID == nil) + #expect(store.snapshot.accounts.isEmpty) + #expect(FileManager.default.fileExists(atPath: home.path) == false) + } + + @Test + func `remove active account falls back to remaining managed account`() 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], + activeAccountID: second.id)) + 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(store.snapshot.activeAccountID == 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], activeAccountID: account.id)) + 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.activeAccountID == 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], activeAccountID: account.id)) + 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.activeAccountID == 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..d95e56332 --- /dev/null +++ b/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift @@ -0,0 +1,302 @@ +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], + activeAccountID: secondID) + 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 activeAccountRange = try #require(contents.range(of: "\"activeAccountID\"")) + let versionRange = try #require(contents.range(of: "\"version\"")) + + #expect(loaded.version == 1) + #expect(loaded.accounts.count == 2) + #expect(loaded.activeAccountID == secondID) + #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 < activeAccountRange.lowerBound) + #expect(activeAccountRange.lowerBound < versionRange.lowerBound) +} + +@Test +func `FileManagedCodexAccountStore preserves nil active account and 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) + #expect(initial.activeAccountID == nil) + + 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], + activeAccountID: nil) + + try store.storeAccounts(payload) + let loaded = try store.loadAccounts() + + #expect(loaded.version == 1) + #expect(loaded.accounts.count == 1) + #expect(loaded.activeAccountID == nil) + #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 + } + ], + "activeAccountID" : "\(accountID.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.first?.email == "mixed@example.com") + #expect(loaded.account(email: "mixed@example.com")?.id == accountID) +} + +@Test +func `FileManagedCodexAccountStore clears dangling active account IDs on load`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-dangling-active-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.activeAccountID == nil) +} + +@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 + } + ], + "activeAccountID" : "\(secondID.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.accounts.first?.id == firstID) + #expect(loaded.accounts.first?.managedHomePath == "/tmp/managed-home-1") + #expect(loaded.activeAccountID == nil) +} + +@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 + } + ], + "activeAccountID" : "\(sharedID.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.accounts.first?.id == sharedID) + #expect(loaded.accounts.first?.email == "first@example.com") + #expect(loaded.accounts.first?.managedHomePath == "/tmp/managed-home-1") + #expect(loaded.activeAccountID == sharedID) +} + +@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 + } + ], + "activeAccountID" : "\(accountID.uuidString)", + "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], + activeAccountID: accountID) + 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..dd8c0b32c --- /dev/null +++ b/Tests/CodexBarTests/MenuDescriptorCodexManagedFallbackTests.swift @@ -0,0 +1,111 @@ +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._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": [ + "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/ProvidersPaneCoverageTests.swift b/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift index d5542e293..548465c7c 100644 --- a/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift +++ b/Tests/CodexBarTests/ProvidersPaneCoverageTests.swift @@ -88,6 +88,44 @@ struct ProvidersPaneCoverageTests { #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") + settings._test_activeManagedCodexRemoteHomePath = managedHome.path + + 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 +157,32 @@ 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": [ + "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))." + } }