diff --git a/Sources/CodexBar/CodexAccountReconciliation.swift b/Sources/CodexBar/CodexAccountReconciliation.swift index 3b0683b97..aaf1751a2 100644 --- a/Sources/CodexBar/CodexAccountReconciliation.swift +++ b/Sources/CodexBar/CodexAccountReconciliation.swift @@ -4,12 +4,18 @@ import Foundation struct CodexVisibleAccount: Equatable, Sendable, Identifiable { let id: String let email: String + let workspaceLabel: String? let storedAccountID: UUID? let selectionSource: CodexActiveSource let isActive: Bool let isLive: Bool let canReauthenticate: Bool let canRemove: Bool + + var displayName: String { + guard let workspaceLabel, !workspaceLabel.isEmpty else { return self.email } + return "\(self.email) — \(workspaceLabel)" + } } struct CodexVisibleAccountProjection: Equatable, Sendable { @@ -40,7 +46,7 @@ enum CodexActiveSourceResolver { .liveSystem case let .managedAccount(id): if let activeStoredAccount = snapshot.activeStoredAccount { - self.matchesLiveSystemAccountEmail( + self.matchesLiveSystemAccountIdentity( storedAccount: activeStoredAccount, liveSystemAccount: snapshot.liveSystemAccount) ? .liveSystem : .managedAccount(id: id) } else { @@ -53,16 +59,18 @@ enum CodexActiveSourceResolver { resolvedSource: resolvedSource) } - private static func matchesLiveSystemAccountEmail( + private static func matchesLiveSystemAccountIdentity( storedAccount: ManagedCodexAccount, liveSystemAccount: ObservedSystemCodexAccount?) -> Bool { guard let liveSystemAccount else { return false } - return Self.normalizeEmail(storedAccount.email) == Self.normalizeEmail(liveSystemAccount.email) - } - - private static func normalizeEmail(_ email: String) -> String { - email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + return ManagedCodexAccount.identityKey( + email: storedAccount.email, + workspaceAccountID: storedAccount.workspaceAccountID, + workspaceLabel: storedAccount.workspaceLabel) == ManagedCodexAccount.identityKey( + email: liveSystemAccount.email, + workspaceAccountID: liveSystemAccount.workspaceAccountID, + workspaceLabel: liveSystemAccount.workspaceLabel) } } @@ -114,7 +122,10 @@ struct DefaultCodexAccountReconciler { nil } let matchingStoredAccountForLiveSystemAccount = liveSystemAccount.flatMap { - accounts.account(email: $0.email) + accounts.account( + email: $0.email, + workspaceAccountID: $0.workspaceAccountID, + workspaceLabel: $0.workspaceLabel) } return CodexAccountReconciliationSnapshot( @@ -150,6 +161,8 @@ struct DefaultCodexAccountReconciler { } return ObservedSystemCodexAccount( email: normalizedEmail, + workspaceLabel: ManagedCodexAccount.normalizeWorkspaceLabel(account.workspaceLabel), + workspaceAccountID: ManagedCodexAccount.normalizeWorkspaceAccountID(account.workspaceAccountID), codexHomePath: account.codexHomePath, observedAt: account.observedAt) } catch { @@ -165,13 +178,20 @@ struct DefaultCodexAccountReconciler { extension CodexVisibleAccountProjection { static func make(from snapshot: CodexAccountReconciliationSnapshot) -> CodexVisibleAccountProjection { let resolvedActiveSource = CodexActiveSourceResolver.resolve(from: snapshot).resolvedSource - var visibleByEmail: [String: CodexVisibleAccount] = [:] + var visibleByID: [String: CodexVisibleAccount] = [:] for storedAccount in snapshot.storedAccounts { let normalizedEmail = Self.normalizeVisibleEmail(storedAccount.email) - visibleByEmail[normalizedEmail] = CodexVisibleAccount( - id: normalizedEmail, + let workspaceLabel = ManagedCodexAccount.normalizeWorkspaceLabel(storedAccount.workspaceLabel) + let workspaceAccountID = ManagedCodexAccount.normalizeWorkspaceAccountID(storedAccount.workspaceAccountID) + let visibleID = Self.visibleAccountID( + email: normalizedEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) + visibleByID[visibleID] = CodexVisibleAccount( + id: visibleID, email: normalizedEmail, + workspaceLabel: workspaceLabel, storedAccountID: storedAccount.id, selectionSource: .managedAccount(id: storedAccount.id), isActive: false, @@ -182,10 +202,18 @@ extension CodexVisibleAccountProjection { if let liveSystemAccount = snapshot.liveSystemAccount { let normalizedEmail = Self.normalizeVisibleEmail(liveSystemAccount.email) - if let existing = visibleByEmail[normalizedEmail] { - visibleByEmail[normalizedEmail] = CodexVisibleAccount( + let workspaceLabel = ManagedCodexAccount.normalizeWorkspaceLabel(liveSystemAccount.workspaceLabel) + let workspaceAccountID = ManagedCodexAccount + .normalizeWorkspaceAccountID(liveSystemAccount.workspaceAccountID) + let visibleID = Self.visibleAccountID( + email: normalizedEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) + if let existing = visibleByID[visibleID] { + visibleByID[visibleID] = CodexVisibleAccount( id: existing.id, email: existing.email, + workspaceLabel: workspaceLabel ?? existing.workspaceLabel, storedAccountID: existing.storedAccountID, selectionSource: .liveSystem, isActive: existing.isActive, @@ -193,9 +221,10 @@ extension CodexVisibleAccountProjection { canReauthenticate: existing.canReauthenticate, canRemove: existing.canRemove) } else { - visibleByEmail[normalizedEmail] = CodexVisibleAccount( - id: normalizedEmail, + visibleByID[visibleID] = CodexVisibleAccount( + id: visibleID, email: normalizedEmail, + workspaceLabel: workspaceLabel, storedAccountID: nil, selectionSource: .liveSystem, isActive: false, @@ -205,17 +234,28 @@ extension CodexVisibleAccountProjection { } } - let activeEmail: String? = switch resolvedActiveSource { + let activeVisibleID: String? = switch resolvedActiveSource { case let .managedAccount(id): - snapshot.storedAccounts.first { $0.id == id }.map { Self.normalizeVisibleEmail($0.email) } + snapshot.storedAccounts.first { $0.id == id }.map { + Self.visibleAccountID( + email: Self.normalizeVisibleEmail($0.email), + workspaceAccountID: $0.workspaceAccountID, + workspaceLabel: $0.workspaceLabel) + } case .liveSystem: - snapshot.liveSystemAccount.map { Self.normalizeVisibleEmail($0.email) } + snapshot.liveSystemAccount.map { + Self.visibleAccountID( + email: Self.normalizeVisibleEmail($0.email), + workspaceAccountID: $0.workspaceAccountID, + workspaceLabel: $0.workspaceLabel) + } } - if let activeEmail, let current = visibleByEmail[activeEmail] { - visibleByEmail[activeEmail] = CodexVisibleAccount( + if let activeVisibleID, let current = visibleByID[activeVisibleID] { + visibleByID[activeVisibleID] = CodexVisibleAccount( id: current.id, email: current.email, + workspaceLabel: current.workspaceLabel, storedAccountID: current.storedAccountID, selectionSource: current.selectionSource, isActive: true, @@ -224,8 +264,11 @@ extension CodexVisibleAccountProjection { canRemove: current.canRemove) } - let visibleAccounts = visibleByEmail.values.sorted { lhs, rhs in - lhs.email < rhs.email + let visibleAccounts = visibleByID.values.sorted { lhs, rhs in + if lhs.email == rhs.email { + return (lhs.workspaceLabel ?? "") < (rhs.workspaceLabel ?? "") + } + return lhs.email < rhs.email } return CodexVisibleAccountProjection( @@ -238,11 +281,24 @@ extension CodexVisibleAccountProjection { private static func normalizeVisibleEmail(_ email: String) -> String { email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() } + + private static func visibleAccountID( + email: String, + workspaceAccountID: String?, + workspaceLabel: String?) -> String + { + ManagedCodexAccount.identityKey( + email: email, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) + } } private struct AccountIdentity: Equatable { let id: UUID let email: String + let workspaceLabel: String? + let workspaceAccountID: String? let managedHomePath: String let createdAt: TimeInterval let updatedAt: TimeInterval @@ -251,6 +307,8 @@ private struct AccountIdentity: Equatable { init(_ account: ManagedCodexAccount) { self.id = account.id self.email = account.email + self.workspaceLabel = account.workspaceLabel + self.workspaceAccountID = account.workspaceAccountID self.managedHomePath = account.managedHomePath self.createdAt = account.createdAt self.updatedAt = account.updatedAt diff --git a/Sources/CodexBar/CodexSystemAccountObserver.swift b/Sources/CodexBar/CodexSystemAccountObserver.swift index 584672c69..7953d930d 100644 --- a/Sources/CodexBar/CodexSystemAccountObserver.swift +++ b/Sources/CodexBar/CodexSystemAccountObserver.swift @@ -3,8 +3,24 @@ import Foundation struct ObservedSystemCodexAccount: Equatable, Sendable { let email: String + let workspaceLabel: String? + let workspaceAccountID: String? let codexHomePath: String let observedAt: Date + + init( + email: String, + workspaceLabel: String? = nil, + workspaceAccountID: String? = nil, + codexHomePath: String, + observedAt: Date) + { + self.email = email + self.workspaceLabel = workspaceLabel + self.workspaceAccountID = workspaceAccountID + self.codexHomePath = codexHomePath + self.observedAt = observedAt + } } protocol CodexSystemAccountObserving: Sendable { @@ -25,6 +41,8 @@ struct DefaultCodexSystemAccountObserver: CodexSystemAccountObserving { return ObservedSystemCodexAccount( email: rawEmail.lowercased(), + workspaceLabel: info.workspaceLabel, + workspaceAccountID: info.workspaceAccountID, codexHomePath: homeURL.path, observedAt: Date()) } diff --git a/Sources/CodexBar/ManagedCodexAccountService.swift b/Sources/CodexBar/ManagedCodexAccountService.swift index 91f9d32a7..98ae4b9ca 100644 --- a/Sources/CodexBar/ManagedCodexAccountService.swift +++ b/Sources/CodexBar/ManagedCodexAccountService.swift @@ -14,6 +14,10 @@ protocol ManagedCodexIdentityReading: Sendable { func loadAccountInfo(homePath: String) throws -> AccountInfo } +protocol ManagedCodexWorkspaceResolving: Sendable { + func resolveAccountInfo(homePath: String, fallback: AccountInfo) async -> AccountInfo +} + enum ManagedCodexAccountServiceError: Error, Equatable, Sendable { case loginFailed case missingEmail @@ -70,12 +74,55 @@ struct DefaultManagedCodexIdentityReader: ManagedCodexIdentityReading { } } +struct DefaultManagedCodexWorkspaceResolver: ManagedCodexWorkspaceResolving { + func resolveAccountInfo(homePath: String, fallback: AccountInfo) async -> AccountInfo { + let env = CodexHomeScope.scopedEnvironment( + base: ProcessInfo.processInfo.environment, + codexHome: homePath) + + guard let credentials = try? CodexOAuthCredentialsStore.load(env: env), + let workspaceAccountID = ManagedCodexAccount.normalizeWorkspaceAccountID(credentials.accountId) + else { + return fallback + } + + do { + let authoritativeIdentity = try await CodexOpenAIWorkspaceResolver.resolve(credentials: credentials) + if let authoritativeIdentity { + try? CodexOpenAIWorkspaceIdentityCache().store(authoritativeIdentity) + return CodexOpenAIWorkspaceResolver.mergeAuthoritativeIdentity( + into: fallback, + authoritativeIdentity: authoritativeIdentity) + } + } catch { + // Fail closed on the workspace label when we have a selected account id but cannot + // resolve the authoritative workspace name. This avoids persisting a misleading + // JWT-derived label such as "Personal" for a selected team workspace. + } + + if let cachedWorkspaceLabel = CodexOpenAIWorkspaceIdentityCache().workspaceLabel(for: workspaceAccountID) { + return AccountInfo( + email: fallback.email, + plan: fallback.plan, + workspaceLabel: cachedWorkspaceLabel, + workspaceAccountID: workspaceAccountID) + } + + return AccountInfo( + email: fallback.email, + plan: fallback.plan, + workspaceLabel: nil, + workspaceAccountID: workspaceAccountID) + } +} + @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 workspaceResolver: any ManagedCodexWorkspaceResolving private let fileManager: FileManager init( @@ -83,12 +130,14 @@ final class ManagedCodexAccountService { homeFactory: any ManagedCodexHomeProducing, loginRunner: any ManagedCodexLoginRunning, identityReader: any ManagedCodexIdentityReading, + workspaceResolver: any ManagedCodexWorkspaceResolving = DefaultManagedCodexWorkspaceResolver(), fileManager: FileManager = .default) { self.store = store self.homeFactory = homeFactory self.loginRunner = loginRunner self.identityReader = identityReader + self.workspaceResolver = workspaceResolver self.fileManager = fileManager } @@ -98,6 +147,7 @@ final class ManagedCodexAccountService { homeFactory: ManagedCodexHomeFactory(fileManager: fileManager), loginRunner: DefaultManagedCodexLoginRunner(), identityReader: DefaultManagedCodexIdentityReader(), + workspaceResolver: DefaultManagedCodexWorkspaceResolver(), fileManager: fileManager) } @@ -116,7 +166,8 @@ final class ManagedCodexAccountService { 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) + let baseInfo = try self.identityReader.loadAccountInfo(homePath: homeURL.path) + let info = await self.workspaceResolver.resolveAccountInfo(homePath: homeURL.path, fallback: baseInfo) guard let rawEmail = info.email?.trimmingCharacters(in: .whitespacesAndNewlines), !rawEmail.isEmpty else { throw ManagedCodexAccountServiceError.missingEmail } @@ -124,12 +175,16 @@ final class ManagedCodexAccountService { let now = Date().timeIntervalSince1970 let existing = self.reconciledExistingAccount( authenticatedEmail: rawEmail, + authenticatedWorkspaceAccountID: info.workspaceAccountID, + authenticatedWorkspaceLabel: info.workspaceLabel, existingAccountID: existingAccountID, snapshot: snapshot) account = ManagedCodexAccount( id: existing?.id ?? UUID(), email: rawEmail, + workspaceLabel: info.workspaceLabel, + workspaceAccountID: info.workspaceAccountID, managedHomePath: homeURL.path, createdAt: existing?.createdAt ?? now, updatedAt: now, @@ -138,7 +193,9 @@ final class ManagedCodexAccountService { let updatedSnapshot = ManagedCodexAccountSet( version: snapshot.version, - accounts: snapshot.accounts.filter { $0.id != account.id && $0.email != account.email } + [account]) + accounts: snapshot.accounts.filter { + $0.id != account.id && $0.identityKey != account.identityKey + } + [account]) try self.store.storeAccounts(updatedSnapshot) } catch { try? self.removeManagedHomeIfSafe(atPath: homeURL.path) @@ -178,19 +235,24 @@ final class ManagedCodexAccountService { private func reconciledExistingAccount( authenticatedEmail: String, + authenticatedWorkspaceAccountID: String?, + authenticatedWorkspaceLabel: String?, existingAccountID: UUID?, snapshot: ManagedCodexAccountSet) -> ManagedCodexAccount? { - if let existingByEmail = snapshot.account(email: authenticatedEmail) { - return existingByEmail + if let existingByIdentity = snapshot.account( + email: authenticatedEmail, + workspaceAccountID: authenticatedWorkspaceAccountID, + workspaceLabel: authenticatedWorkspaceLabel) + { + return existingByIdentity } 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() + return existingByID.identityKey == ManagedCodexAccount.identityKey( + email: authenticatedEmail, + workspaceAccountID: authenticatedWorkspaceAccountID, + workspaceLabel: authenticatedWorkspaceLabel) ? existingByID : nil } } diff --git a/Sources/CodexBar/PreferencesCodexAccountsSection.swift b/Sources/CodexBar/PreferencesCodexAccountsSection.swift index c1b620e07..e77067a09 100644 --- a/Sources/CodexBar/PreferencesCodexAccountsSection.swift +++ b/Sources/CodexBar/PreferencesCodexAccountsSection.swift @@ -105,7 +105,7 @@ struct CodexAccountsSectionView: View { Picker("", selection: selection) { ForEach(self.state.visibleAccounts) { account in - Text(account.email).tag(account.id) + Text(account.displayName).tag(account.id) } } .labelsHidden() @@ -127,7 +127,7 @@ struct CodexAccountsSectionView: View { .font(.subheadline.weight(.semibold)) .frame(width: ProviderSettingsMetrics.pickerLabelWidth, alignment: .leading) - Text(account.email) + Text(account.displayName) .font(.subheadline) Spacer(minLength: 0) @@ -192,7 +192,7 @@ private struct CodexAccountsSectionRowView: View { var body: some View { HStack(alignment: .center, spacing: 12) { HStack(alignment: .firstTextBaseline, spacing: 6) { - Text(self.account.email) + Text(self.account.displayName) .font(.subheadline.weight(.semibold)) if self.showsLiveBadge { Text("(Live)") diff --git a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift index 5e7b6298e..f00d76fd1 100644 --- a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift +++ b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift @@ -222,7 +222,10 @@ extension SettingsStore { } func selectAuthenticatedManagedCodexAccount(_ account: ManagedCodexAccount) { - let visibleAccountID = Self.codexVisibleAccountID(for: account.email) + let visibleAccountID = Self.codexVisibleAccountID( + for: account.email, + workspaceAccountID: account.workspaceAccountID, + workspaceLabel: account.workspaceLabel) if self.selectCodexVisibleAccount(id: visibleAccountID) { return } @@ -285,8 +288,15 @@ extension SettingsStore { return ProcessInfo.processInfo.environment } - private static func codexVisibleAccountID(for email: String) -> String { - email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + private static func codexVisibleAccountID( + for email: String, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil) -> String + { + ManagedCodexAccount.identityKey( + email: email, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) } } diff --git a/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift index c6f476717..999571393 100644 --- a/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift +++ b/Sources/CodexBar/Providers/Codex/UsageStore+CodexAccountState.swift @@ -85,9 +85,17 @@ extension UsageStore { func seedCodexAccountScopedRefreshGuard( source: CodexActiveSource? = nil, - accountEmail: String?) + accountEmail: String?, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil) { - guard let accountKey = Self.normalizeCodexAccountScopedKey(accountEmail) else { return } + guard let accountKey = Self.normalizeCodexAccountScopedKey( + email: accountEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) + else { + return + } self.lastCodexAccountScopedRefreshGuard = CodexAccountScopedRefreshGuard( source: source ?? self.settings.codexResolvedActiveSource, accountKey: accountKey) @@ -107,11 +115,16 @@ extension UsageStore { func currentCodexOpenAIWebRefreshGuard() -> CodexAccountScopedRefreshGuard { let accountKey: String? = switch self.settings.codexResolvedActiveSource { case .liveSystem: - Self - .normalizeCodexAccountScopedKey(self.settings.codexAccountReconciliationSnapshot.liveSystemAccount? - .email) + Self.normalizeCodexAccountScopedKey( + email: self.settings.codexAccountReconciliationSnapshot.liveSystemAccount?.email, + workspaceAccountID: self.settings.codexAccountReconciliationSnapshot.liveSystemAccount? + .workspaceAccountID, + workspaceLabel: self.settings.codexAccountReconciliationSnapshot.liveSystemAccount?.workspaceLabel) case .managedAccount: - Self.normalizeCodexAccountScopedKey(self.settings.activeManagedCodexAccount?.email) + Self.normalizeCodexAccountScopedKey( + email: self.settings.activeManagedCodexAccount?.email, + workspaceAccountID: self.settings.activeManagedCodexAccount?.workspaceAccountID, + workspaceLabel: self.settings.activeManagedCodexAccount?.workspaceLabel) } return CodexAccountScopedRefreshGuard( source: self.settings.codexResolvedActiveSource, @@ -129,7 +142,10 @@ extension UsageStore { return currentGuard.accountKey == expectedKey } - let resultKey = Self.normalizeCodexAccountScopedKey(usage.accountEmail(for: .codex)) + let resultKey = Self.normalizeCodexAccountScopedKey( + email: usage.accountEmail(for: .codex), + workspaceAccountID: usage.accountWorkspaceID(for: .codex), + workspaceLabel: usage.accountOrganization(for: .codex)) if let currentKey = currentGuard.accountKey { return resultKey == currentKey } @@ -174,10 +190,19 @@ extension UsageStore { guard currentGuard.source == expectedGuard.source else { return false } guard case .liveSystem = expectedGuard.source else { return false } guard currentGuard.accountKey == nil else { return false } - guard let dashboardKey = Self.normalizeCodexAccountScopedKey(dashboardAccountEmail) else { return false } - let currentTargetKey = Self.normalizeCodexAccountScopedKey(self.currentCodexOpenAIWebTargetEmail( - allowCurrentSnapshotFallback: true, - allowLastKnownLiveFallback: false)) + guard let dashboardKey = Self.normalizeCodexAccountScopedKey( + email: dashboardAccountEmail, + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard()) + else { + return false + } + let currentTargetKey = Self.normalizeCodexAccountScopedKey( + email: self.currentCodexOpenAIWebTargetEmail( + allowCurrentSnapshotFallback: true, + allowLastKnownLiveFallback: false), + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard()) if let currentTargetKey { return dashboardKey == currentTargetKey } @@ -195,9 +220,29 @@ extension UsageStore { allowLastKnownLiveFallback: Bool = true) -> String? { Self.normalizeCodexAccountScopedKey( - self.codexAccountScopedRefreshEmail( + email: self.codexAccountScopedRefreshEmail( preferCurrentSnapshot: preferCurrentSnapshot, - allowLastKnownLiveFallback: allowLastKnownLiveFallback)) + allowLastKnownLiveFallback: allowLastKnownLiveFallback), + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard()) + } + + func codexWorkspaceAccountIDForOpenAIDashboard() -> String? { + switch self.settings.codexResolvedActiveSource { + case .liveSystem: + self.settings.codexAccountReconciliationSnapshot.liveSystemAccount?.workspaceAccountID + case .managedAccount: + self.settings.activeManagedCodexAccount?.workspaceAccountID + } + } + + func codexWorkspaceLabelForOpenAIDashboard() -> String? { + switch self.settings.codexResolvedActiveSource { + case .liveSystem: + self.settings.codexAccountReconciliationSnapshot.liveSystemAccount?.workspaceLabel + case .managedAccount: + self.settings.activeManagedCodexAccount?.workspaceLabel + } } func codexAccountScopedRefreshEmail( @@ -239,7 +284,10 @@ extension UsageStore { private func clearCodexOpenAIWebStateForAccountTransition(targetEmail: String?) { self.invalidateOpenAIDashboardRefreshTask() if self.settings.codexCookieSource.isEnabled, - let normalizedTarget = Self.normalizeCodexAccountScopedEmail(targetEmail) + let normalizedTarget = Self.normalizeCodexAccountScopedKey( + email: targetEmail, + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard()) { let previous = self.lastOpenAIDashboardTargetEmail self.lastOpenAIDashboardTargetEmail = normalizedTarget @@ -272,7 +320,23 @@ extension UsageStore { return trimmed } - static func normalizeCodexAccountScopedKey(_ email: String?) -> String? { - self.normalizeCodexAccountScopedEmail(email)?.lowercased() + static func normalizeCodexAccountScopedKey( + email: String?, + workspaceAccountID: String? = nil, + workspaceLabel: String?) -> String? + { + guard let normalizedEmail = self.normalizeCodexAccountScopedEmail(email)?.lowercased() else { + return nil + } + if let normalizedWorkspaceAccountID = ManagedCodexAccount.normalizeWorkspaceAccountID(workspaceAccountID) { + return "\(normalizedEmail)\naccount:\(normalizedWorkspaceAccountID)" + } + let normalizedWorkspace = workspaceLabel? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + if let normalizedWorkspace, !normalizedWorkspace.isEmpty { + return "\(normalizedEmail)\n\(normalizedWorkspace)" + } + return normalizedEmail } } diff --git a/Sources/CodexBar/StatusItemController+SwitcherViews.swift b/Sources/CodexBar/StatusItemController+SwitcherViews.swift index 4c56de1ab..77a7f72ee 100644 --- a/Sources/CodexBar/StatusItemController+SwitcherViews.swift +++ b/Sources/CodexBar/StatusItemController+SwitcherViews.swift @@ -907,6 +907,9 @@ final class CodexAccountSwitcherView: NSView { private let unselectedBackground = NSColor.clear.cgColor private let selectedTextColor = NSColor.white private let unselectedTextColor = NSColor.secondaryLabelColor + private let buttonFont = NSFont.systemFont(ofSize: NSFont.smallSystemFontSize) + private let buttonHorizontalPadding: CGFloat = 14 + private let buttonSideInset: CGFloat = 6 init( accounts: [CodexVisibleAccount], @@ -954,17 +957,18 @@ final class CodexAccountSwitcherView: NSView { row.spacing = self.rowSpacing row.translatesAutoresizingMaskIntoConstraints = false + let buttonWidth = self.buttonWidth(for: rowAccounts.count) for account in rowAccounts { let button = PaddedToggleButton( - title: account.email, + title: self.compactButtonTitle(for: account, buttonWidth: buttonWidth), target: self, action: #selector(self.handleSelect)) button.identifier = NSUserInterfaceItemIdentifier(account.id) - button.toolTip = account.email + button.toolTip = account.displayName button.isBordered = false button.setButtonType(.toggle) button.controlSize = .small - button.font = NSFont.systemFont(ofSize: NSFont.smallSystemFontSize) + button.font = self.buttonFont button.wantsLayer = true button.layer?.cornerRadius = 6 row.addArrangedSubview(button) @@ -976,8 +980,8 @@ final class CodexAccountSwitcherView: NSView { self.addSubview(stack) NSLayoutConstraint.activate([ - stack.leadingAnchor.constraint(equalTo: self.leadingAnchor, constant: 6), - stack.trailingAnchor.constraint(equalTo: self.trailingAnchor, constant: -6), + stack.leadingAnchor.constraint(equalTo: self.leadingAnchor, constant: self.buttonSideInset), + stack.trailingAnchor.constraint(equalTo: self.trailingAnchor, constant: -self.buttonSideInset), stack.topAnchor.constraint(equalTo: self.topAnchor), stack.bottomAnchor.constraint(equalTo: self.bottomAnchor), stack.heightAnchor.constraint(equalToConstant: self.rowHeight * CGFloat(rows.count) + @@ -985,6 +989,91 @@ final class CodexAccountSwitcherView: NSView { ]) } + private func buttonWidth(for count: Int) -> CGFloat { + let contentWidth = self.bounds.width - (self.buttonSideInset * 2) + let spacing = self.rowSpacing * CGFloat(max(0, count - 1)) + guard count > 0 else { return contentWidth } + return max(44, floor((contentWidth - spacing) / CGFloat(count))) + } + + private func compactButtonTitle(for account: CodexVisibleAccount, buttonWidth: CGFloat) -> String { + let availableTextWidth = max(24, buttonWidth - self.buttonHorizontalPadding) + if self.textWidth(account.displayName) <= availableTextWidth { + return account.displayName + } + + guard let workspace = account.workspaceLabel, !workspace.isEmpty else { + return self.truncateTail(account.email, toFit: availableTextWidth) + } + + let separator = "|" + let separatorWidth = self.textWidth(separator) + let contentWidth = max(24, availableTextWidth - separatorWidth) + let minimumEmailWidth = min(contentWidth * 0.45, max(18, contentWidth * 0.3)) + let minimumWorkspaceWidth = min(contentWidth * 0.4, max(18, contentWidth * 0.25)) + var emailWidth = max(minimumEmailWidth, contentWidth * 0.58) + var workspaceWidth = max(minimumWorkspaceWidth, contentWidth - emailWidth) + + func makeTitle() -> String { + let email = self.truncateTail(account.email, toFit: emailWidth) + let workspace = self.truncateTail(workspace, toFit: workspaceWidth) + return "\(email)\(separator)\(workspace)" + } + + var title = makeTitle() + var attempts = 0 + while self.textWidth(title) > availableTextWidth, attempts < 16 { + let emailText = self.truncateTail(account.email, toFit: emailWidth) + let workspaceText = self.truncateTail(workspace, toFit: workspaceWidth) + let emailRenderedWidth = self.textWidth(emailText) + let workspaceRenderedWidth = self.textWidth(workspaceText) + + if emailRenderedWidth >= workspaceRenderedWidth, emailWidth > minimumEmailWidth { + emailWidth = max(minimumEmailWidth, emailWidth - 6) + } else if workspaceWidth > minimumWorkspaceWidth { + workspaceWidth = max(minimumWorkspaceWidth, workspaceWidth - 6) + } else { + break + } + + title = makeTitle() + attempts += 1 + } + + return title + } + + private func truncateTail(_ text: String, toFit width: CGFloat) -> String { + let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty else { return text } + if self.textWidth(trimmed) <= width { + return trimmed + } + + let ellipsis = "…" + let ellipsisWidth = self.textWidth(ellipsis) + guard ellipsisWidth < width else { return ellipsis } + + var candidate = "" + for character in trimmed { + let next = candidate + String(character) + if self.textWidth(next + ellipsis) > width { + break + } + candidate = next + } + + if candidate.isEmpty { + return ellipsis + } + return candidate + ellipsis + } + + private func textWidth(_ text: String) -> CGFloat { + let attributes: [NSAttributedString.Key: Any] = [.font: self.buttonFont] + return ceil((text as NSString).size(withAttributes: attributes).width) + } + private func updateButtonStyles() { for button in self.buttons { let selected = button.identifier?.rawValue == self.selectedAccountID @@ -1001,4 +1090,14 @@ final class CodexAccountSwitcherView: NSView { self.updateButtonStyles() self.onSelect(accountID) } + + #if DEBUG + func _test_buttonTitles() -> [String] { + self.buttons.map(\.title) + } + + func _test_buttonToolTips() -> [String?] { + self.buttons.map(\.toolTip) + } + #endif } diff --git a/Sources/CodexBar/UsageStore+OpenAIWeb.swift b/Sources/CodexBar/UsageStore+OpenAIWeb.swift index 238bcdfa9..82d597c84 100644 --- a/Sources/CodexBar/UsageStore+OpenAIWeb.swift +++ b/Sources/CodexBar/UsageStore+OpenAIWeb.swift @@ -42,7 +42,10 @@ extension UsageStore { { guard self.shouldApplyOpenAIDashboardRefreshTask(token: refreshTaskToken) else { return } let resolvedAccountEmail = targetEmail ?? dash.signedInEmail - let resolvedAccountKey = Self.normalizeCodexAccountScopedKey(resolvedAccountEmail) + let resolvedAccountKey = Self.normalizeCodexAccountScopedKey( + email: resolvedAccountEmail, + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard()) if let expectedGuard, !self.shouldApplyOpenAIDashboardResult( expectedGuard: expectedGuard, @@ -74,7 +77,10 @@ extension UsageStore { self.lastCreditsError = nil self.creditsFailureStreak = 0 } - self.seedCodexAccountScopedRefreshGuard(accountEmail: resolvedAccountEmail) + self.seedCodexAccountScopedRefreshGuard( + accountEmail: resolvedAccountEmail, + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard()) } if let email = targetEmail, !email.isEmpty { @@ -549,7 +555,12 @@ extension UsageStore { expectedGuard: CodexAccountScopedRefreshGuard?) -> String { let source = String(describing: expectedGuard?.source ?? self.settings.codexResolvedActiveSource) - let accountKey = Self.normalizeCodexAccountScopedKey(targetEmail ?? expectedGuard?.accountKey) ?? "unknown" + let accountKey = Self.normalizeCodexAccountScopedKey( + email: targetEmail, + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard()) + ?? expectedGuard?.accountKey + ?? "unknown" return "\(source)|\(accountKey)" } @@ -643,6 +654,8 @@ extension UsageStore { } return try await OpenAIDashboardFetcher().loadLatestDashboard( accountEmail: accountEmail, + workspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + workspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard(), logger: logger, debugDumpHTML: timeout != Self.openAIWebPrimaryFetchTimeout, timeout: timeout) @@ -761,12 +774,16 @@ extension UsageStore { result = try await importer.importManualCookies( cookieHeader: manualHeader, intoAccountEmail: normalizedTarget, + intoWorkspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + intoWorkspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard(), allowAnyAccount: allowAnyAccount, cacheScope: cacheScope, logger: log) case .auto: result = try await importer.importBestCookies( intoAccountEmail: normalizedTarget, + intoWorkspaceAccountID: self.codexWorkspaceAccountIDForOpenAIDashboard(), + intoWorkspaceLabel: self.codexWorkspaceLabelForOpenAIDashboard(), allowAnyAccount: allowAnyAccount, cacheScope: cacheScope, logger: log) diff --git a/Sources/CodexBarCore/CodexManagedAccounts.swift b/Sources/CodexBarCore/CodexManagedAccounts.swift index bfae7ce50..506794dcd 100644 --- a/Sources/CodexBarCore/CodexManagedAccounts.swift +++ b/Sources/CodexBarCore/CodexManagedAccounts.swift @@ -3,6 +3,8 @@ import Foundation public struct ManagedCodexAccount: Codable, Identifiable, Sendable { public let id: UUID public let email: String + public let workspaceLabel: String? + public let workspaceAccountID: String? public let managedHomePath: String public let createdAt: TimeInterval public let updatedAt: TimeInterval @@ -11,6 +13,8 @@ public struct ManagedCodexAccount: Codable, Identifiable, Sendable { public init( id: UUID, email: String, + workspaceLabel: String? = nil, + workspaceAccountID: String? = nil, managedHomePath: String, createdAt: TimeInterval, updatedAt: TimeInterval, @@ -18,6 +22,8 @@ public struct ManagedCodexAccount: Codable, Identifiable, Sendable { { self.id = id self.email = Self.normalizeEmail(email) + self.workspaceLabel = Self.normalizeWorkspaceLabel(workspaceLabel) + self.workspaceAccountID = Self.normalizeWorkspaceAccountID(workspaceAccountID) self.managedHomePath = managedHomePath self.createdAt = createdAt self.updatedAt = updatedAt @@ -28,11 +34,49 @@ public struct ManagedCodexAccount: Codable, Identifiable, Sendable { email.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() } + public static func normalizeWorkspaceLabel(_ workspaceLabel: String?) -> String? { + guard let trimmed = workspaceLabel?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } + + public static func normalizeWorkspaceAccountID(_ workspaceAccountID: String?) -> String? { + guard let trimmed = workspaceAccountID?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed.lowercased() + } + + public static func identityKey( + email: String, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil) -> String + { + let normalizedEmail = self.normalizeEmail(email) + if let normalizedWorkspaceAccountID = self.normalizeWorkspaceAccountID(workspaceAccountID) { + return "\(normalizedEmail)\naccount:\(normalizedWorkspaceAccountID)" + } + guard let normalizedWorkspace = self.normalizeWorkspaceLabel(workspaceLabel)?.lowercased() else { + return normalizedEmail + } + return "\(normalizedEmail)\n\(normalizedWorkspace)" + } + + public var identityKey: String { + Self.identityKey( + email: self.email, + workspaceAccountID: self.workspaceAccountID, + workspaceLabel: self.workspaceLabel) + } + 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), + workspaceLabel: container.decodeIfPresent(String.self, forKey: .workspaceLabel), + workspaceAccountID: container.decodeIfPresent(String.self, forKey: .workspaceAccountID), managedHomePath: container.decode(String.self, forKey: .managedHomePath), createdAt: container.decode(TimeInterval.self, forKey: .createdAt), updatedAt: container.decode(TimeInterval.self, forKey: .updatedAt), @@ -53,9 +97,16 @@ public struct ManagedCodexAccountSet: Codable, Sendable { 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 func account( + email: String, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil) -> ManagedCodexAccount? + { + let identityKey = ManagedCodexAccount.identityKey( + email: email, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) + return self.accounts.first { $0.identityKey == identityKey } } public init(from decoder: any Decoder) throws { @@ -67,13 +118,13 @@ public struct ManagedCodexAccountSet: Codable, Sendable { private static func sanitizedAccounts(_ accounts: [ManagedCodexAccount]) -> [ManagedCodexAccount] { var seenIDs: Set = [] - var seenEmails: Set = [] + var seenIdentityKeys: 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 } + guard seenIdentityKeys.insert(account.identityKey).inserted else { continue } sanitized.append(account) } diff --git a/Sources/CodexBarCore/OpenAIDashboardModels.swift b/Sources/CodexBarCore/OpenAIDashboardModels.swift index e4fe3a06f..6a9d6e401 100644 --- a/Sources/CodexBarCore/OpenAIDashboardModels.swift +++ b/Sources/CodexBarCore/OpenAIDashboardModels.swift @@ -118,6 +118,8 @@ extension OpenAIDashboardSnapshot { public func toUsageSnapshot( provider: UsageProvider = .codex, accountEmail: String? = nil, + accountOrganization: String? = nil, + accountWorkspaceID: String? = nil, accountPlan: String? = nil) -> UsageSnapshot? { guard let primaryLimit else { return nil } @@ -126,7 +128,8 @@ extension OpenAIDashboardSnapshot { let identity = ProviderIdentitySnapshot( providerID: provider, accountEmail: resolvedEmail, - accountOrganization: nil, + accountOrganization: accountOrganization, + accountWorkspaceID: accountWorkspaceID, loginMethod: resolvedPlan) return UsageSnapshot( primary: primaryLimit, diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift index 514c7f5f6..e270fa8a1 100644 --- a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardBrowserCookieImporter.swift @@ -3,6 +3,7 @@ import Foundation import SweetCookieKit import WebKit +// swiftlint:disable type_body_length @MainActor public struct OpenAIDashboardBrowserCookieImporter { public struct FoundAccount: Sendable, Hashable { @@ -80,10 +81,20 @@ public struct OpenAIDashboardBrowserCookieImporter { private struct ImportContext { let targetEmail: String? + let targetWorkspaceAccountID: String? + let targetWorkspaceLabel: String? let allowAnyAccount: Bool let cacheScope: CookieHeaderCache.Scope? } + private struct TargetMatchContext { + let targetEmail: String + let targetWorkspaceAccountID: String? + let targetWorkspaceLabel: String? + let candidateLabel: String + let log: (String) -> Void + } + private static let cookieDomains = ["chatgpt.com", "openai.com"] private static let cookieClient = BrowserCookieClient() private static let cookieImportOrder: BrowserCookieImportOrder = @@ -99,6 +110,8 @@ public struct OpenAIDashboardBrowserCookieImporter { public func importBestCookies( intoAccountEmail targetEmail: String?, + intoWorkspaceAccountID targetWorkspaceAccountID: String? = nil, + intoWorkspaceLabel targetWorkspaceLabel: String? = nil, allowAnyAccount: Bool = false, cacheScope: CookieHeaderCache.Scope? = nil, logger: ((String) -> Void)? = nil) async throws -> ImportResult @@ -111,6 +124,8 @@ public struct OpenAIDashboardBrowserCookieImporter { let normalizedTarget = targetEmail?.isEmpty == false ? targetEmail : nil let context = ImportContext( targetEmail: normalizedTarget, + targetWorkspaceAccountID: ManagedCodexAccount.normalizeWorkspaceAccountID(targetWorkspaceAccountID), + targetWorkspaceLabel: targetWorkspaceLabel, allowAnyAccount: allowAnyAccount, cacheScope: cacheScope) @@ -133,6 +148,8 @@ public struct OpenAIDashboardBrowserCookieImporter { return try await self.importManualCookies( cookieHeader: cached.cookieHeader, intoAccountEmail: context.targetEmail, + intoWorkspaceAccountID: context.targetWorkspaceAccountID, + intoWorkspaceLabel: context.targetWorkspaceLabel, allowAnyAccount: context.allowAnyAccount, cacheScope: cacheScope, logger: log) @@ -187,6 +204,8 @@ public struct OpenAIDashboardBrowserCookieImporter { public func importManualCookies( cookieHeader: String, intoAccountEmail targetEmail: String?, + intoWorkspaceAccountID targetWorkspaceAccountID: String? = nil, + intoWorkspaceLabel targetWorkspaceLabel: String? = nil, allowAnyAccount: Bool = false, cacheScope _: CookieHeaderCache.Scope? = nil, logger: ((String) -> Void)? = nil) async throws -> ImportResult @@ -196,6 +215,12 @@ public struct OpenAIDashboardBrowserCookieImporter { } let normalizedTarget = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) let allowAnyAccount = allowAnyAccount || normalizedTarget == nil || normalizedTarget?.isEmpty == true + let context = ImportContext( + targetEmail: normalizedTarget, + targetWorkspaceAccountID: ManagedCodexAccount.normalizeWorkspaceAccountID(targetWorkspaceAccountID), + targetWorkspaceLabel: targetWorkspaceLabel, + allowAnyAccount: allowAnyAccount, + cacheScope: nil) guard let normalized = CookieHeaderNormalizer.normalize(cookieHeader) else { throw ImportError.manualCookieHeaderInvalid @@ -206,16 +231,21 @@ public struct OpenAIDashboardBrowserCookieImporter { guard !cookies.isEmpty else { throw ImportError.manualCookieHeaderInvalid } let candidate = Candidate(label: "Manual", cookies: cookies) - switch await self.evaluateCandidate( - candidate, - targetEmail: normalizedTarget, - allowAnyAccount: allowAnyAccount, - log: log) - { + switch await self.evaluateCandidate(candidate, context: context, log: log) { case let .match(_, signedInEmail): - return try await self.persist(candidate: candidate, targetEmail: signedInEmail, logger: log) + return try await self.persist( + candidate: candidate, + targetEmail: signedInEmail, + workspaceAccountID: ManagedCodexAccount.normalizeWorkspaceAccountID(targetWorkspaceAccountID), + workspaceLabel: targetWorkspaceLabel, + logger: log) case let .loggedIn(_, signedInEmail): - return try await self.persist(candidate: candidate, targetEmail: signedInEmail, logger: log) + return try await self.persist( + candidate: candidate, + targetEmail: signedInEmail, + workspaceAccountID: ManagedCodexAccount.normalizeWorkspaceAccountID(targetWorkspaceAccountID), + workspaceLabel: targetWorkspaceLabel, + logger: log) case let .mismatch(_, signedInEmail): throw ImportError.noMatchingAccount(found: [FoundAccount(sourceLabel: "Manual", email: signedInEmail)]) case .unknown: @@ -396,16 +426,17 @@ public struct OpenAIDashboardBrowserCookieImporter { log: @escaping (String) -> Void, diagnostics: inout ImportDiagnostics) async -> ImportResult? { - switch await self.evaluateCandidate( - candidate, - targetEmail: context.targetEmail, - allowAnyAccount: context.allowAnyAccount, - log: log) - { + switch await self.evaluateCandidate(candidate, context: context, log: log) { case let .match(candidate, signedInEmail): log("Selected \(candidate.label) (matches Codex: \(signedInEmail))") guard let targetEmail = context.targetEmail else { return nil } - if let result = try? await self.persist(candidate: candidate, targetEmail: targetEmail, logger: log) { + if let result = try? await self.persist( + candidate: candidate, + targetEmail: targetEmail, + workspaceAccountID: context.targetWorkspaceAccountID, + workspaceLabel: context.targetWorkspaceLabel, + logger: log) + { self.cacheCookies(candidate: candidate, scope: context.cacheScope) return result } @@ -414,12 +445,19 @@ public struct OpenAIDashboardBrowserCookieImporter { await self.handleMismatch( candidate: candidate, signedInEmail: signedInEmail, + context: context, log: log, diagnostics: &diagnostics) return nil 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) { + if let result = try? await self.persist( + candidate: candidate, + targetEmail: signedInEmail, + workspaceAccountID: context.targetWorkspaceAccountID, + workspaceLabel: context.targetWorkspaceLabel, + logger: log) + { self.cacheCookies(candidate: candidate, scope: context.cacheScope) return result } @@ -442,12 +480,16 @@ public struct OpenAIDashboardBrowserCookieImporter { private func evaluateCandidate( _ candidate: Candidate, - targetEmail: String?, - allowAnyAccount: Bool, + context: ImportContext, log: @escaping (String) -> Void) async -> CandidateEvaluation { log("Trying candidate \(candidate.label) (\(candidate.cookies.count) cookies)") + let resolvedWorkspaceLabel = self.resolveWorkspaceLabel(from: candidate.cookies) + if let resolvedWorkspaceLabel, !resolvedWorkspaceLabel.isEmpty { + log("Candidate \(candidate.label) workspace: \(resolvedWorkspaceLabel)") + } + let apiEmail = await self.fetchSignedInEmailFromAPI(cookies: candidate.cookies, logger: log) if let apiEmail { log("Candidate \(candidate.label) API email: \(apiEmail)") @@ -455,13 +497,23 @@ public struct OpenAIDashboardBrowserCookieImporter { // Prefer the API email when available (fast; avoids WebKit hydration/timeout risks). if let apiEmail, !apiEmail.isEmpty { - if let targetEmail { - if apiEmail.lowercased() == targetEmail.lowercased() { + if let targetEmail = context.targetEmail { + if self.matchesTarget( + signedInEmail: apiEmail, + candidateWorkspaceAccountID: self.resolveWorkspaceAccountID(from: candidate.cookies), + candidateWorkspaceLabel: resolvedWorkspaceLabel, + context: TargetMatchContext( + targetEmail: targetEmail, + targetWorkspaceAccountID: context.targetWorkspaceAccountID, + targetWorkspaceLabel: context.targetWorkspaceLabel, + candidateLabel: candidate.label, + log: log)) + { return .match(candidate: candidate, signedInEmail: apiEmail) } return .mismatch(candidate: candidate, signedInEmail: apiEmail) } - if allowAnyAccount { return .loggedIn(candidate: candidate, signedInEmail: apiEmail) } + if context.allowAnyAccount { return .loggedIn(candidate: candidate, signedInEmail: apiEmail) } } if !self.hasSessionCookies(candidate.cookies) { @@ -482,13 +534,25 @@ public struct OpenAIDashboardBrowserCookieImporter { let resolvedEmail = signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) if let resolvedEmail, !resolvedEmail.isEmpty { - if let targetEmail { - if resolvedEmail.lowercased() == targetEmail.lowercased() { + if let targetEmail = context.targetEmail { + if self.matchesTarget( + signedInEmail: resolvedEmail, + candidateWorkspaceAccountID: self.resolveWorkspaceAccountID(from: candidate.cookies), + candidateWorkspaceLabel: resolvedWorkspaceLabel, + context: TargetMatchContext( + targetEmail: targetEmail, + targetWorkspaceAccountID: context.targetWorkspaceAccountID, + targetWorkspaceLabel: context.targetWorkspaceLabel, + candidateLabel: candidate.label, + log: log)) + { return .match(candidate: candidate, signedInEmail: resolvedEmail) } return .mismatch(candidate: candidate, signedInEmail: resolvedEmail) } - if allowAnyAccount { return .loggedIn(candidate: candidate, signedInEmail: resolvedEmail) } + if context.allowAnyAccount { + return .loggedIn(candidate: candidate, signedInEmail: resolvedEmail) + } } return .unknown(candidate: candidate) @@ -512,18 +576,122 @@ public struct OpenAIDashboardBrowserCookieImporter { return false } + private func matchesTarget( + signedInEmail: String, + candidateWorkspaceAccountID: String?, + candidateWorkspaceLabel: String?, + context: TargetMatchContext) -> Bool + { + guard signedInEmail.lowercased() == context.targetEmail.lowercased() else { return false } + + let normalizedTargetWorkspaceAccountID = ManagedCodexAccount.normalizeWorkspaceAccountID( + context.targetWorkspaceAccountID) + if let normalizedTargetWorkspaceAccountID { + let normalizedCandidateWorkspaceAccountID = ManagedCodexAccount.normalizeWorkspaceAccountID( + candidateWorkspaceAccountID) + guard let normalizedCandidateWorkspaceAccountID else { + context.log( + "Candidate \(context.candidateLabel) matched email but workspace account id is unknown; " + + "expected \(normalizedTargetWorkspaceAccountID)") + return false + } + if normalizedCandidateWorkspaceAccountID == normalizedTargetWorkspaceAccountID { + return true + } + context.log( + "Candidate \(context.candidateLabel) matched email but workspace account id mismatched " + + "(candidate=\(normalizedCandidateWorkspaceAccountID), " + + "target=\(normalizedTargetWorkspaceAccountID))") + return false + } + + let normalizedTargetWorkspace = self.normalizeWorkspaceLabel(context.targetWorkspaceLabel) + guard let normalizedTargetWorkspace else { return true } + + let normalizedCandidateWorkspace = self.normalizeWorkspaceLabel(candidateWorkspaceLabel) + guard let normalizedCandidateWorkspace else { + context.log( + "Candidate \(context.candidateLabel) matched email but workspace is unknown; " + + "expected \(normalizedTargetWorkspace)") + return false + } + + if normalizedCandidateWorkspace == normalizedTargetWorkspace { + return true + } + + context.log( + "Candidate \(context.candidateLabel) matched email but workspace mismatched " + + "(candidate=\(normalizedCandidateWorkspace), target=\(normalizedTargetWorkspace))") + return false + } + + private func normalizeWorkspaceLabel(_ label: String?) -> String? { + let trimmed = label?.trimmingCharacters(in: .whitespacesAndNewlines) + guard let trimmed, !trimmed.isEmpty else { return nil } + return trimmed.lowercased() + } + + private func resolveWorkspaceAccountID(from cookies: [HTTPCookie]) -> String? { + ManagedCodexAccount.normalizeWorkspaceAccountID( + cookies.first(where: { $0.name == "_account" })?.value) + } + + private func resolveWorkspaceLabel(from cookies: [HTTPCookie]) -> String? { + guard let accountID = self.resolveWorkspaceAccountID(from: cookies), + let sessionCookie = cookies.first(where: { $0.name == "oai-client-auth-session" })?.value, + let payload = self.decodeBase64URLJSONPayload(fromCookieValue: sessionCookie), + let json = try? JSONSerialization.jsonObject(with: payload) as? [String: Any], + let workspaces = json["workspaces"] as? [[String: Any]] + else { + return nil + } + + guard let workspace = workspaces.first(where: { + ManagedCodexAccount.normalizeWorkspaceAccountID($0["id"] as? String) == accountID + }) else { + return nil + } + + let name = (workspace["name"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + if let name, !name.isEmpty { return name } + + let kind = (workspace["kind"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + if kind == "personal" { return "Personal" } + return nil + } + + private func decodeBase64URLJSONPayload(fromCookieValue value: String) -> Data? { + let prefix = value.split(separator: ".", maxSplits: 1, omittingEmptySubsequences: false).first + .map(String.init) ?? value + guard !prefix.isEmpty else { return nil } + var base64 = prefix.replacingOccurrences(of: "-", with: "+").replacingOccurrences(of: "_", with: "/") + let remainder = base64.count % 4 + if remainder != 0 { + base64 += String(repeating: "=", count: 4 - remainder) + } + return Data(base64Encoded: base64) + } + private func handleMismatch( candidate: Candidate, signedInEmail: String, + context: ImportContext, log: @escaping (String) -> Void, diagnostics: inout ImportDiagnostics) async { log("Candidate \(candidate.label) mismatch (\(signedInEmail)); continuing browser search") diagnostics.mismatches.append(FoundAccount(sourceLabel: candidate.label, email: signedInEmail)) - // Mismatch still means we found a valid signed-in session. Persist it keyed by its email so if - // the user switches Codex accounts later, we can reuse this session immediately without another - // Keychain prompt. - await self.persistCookies(candidate: candidate, accountEmail: signedInEmail, logger: log) + // Mismatch still means we found a valid signed-in session. Persist it keyed by the + // candidate's resolved email/workspace so later account switches can reuse the right + // browser state without collapsing same-email workspaces together. + await self.persistCookies( + candidate: candidate, + accountEmail: signedInEmail, + workspaceAccountID: self.resolveWorkspaceAccountID(from: candidate.cookies) ?? context + .targetWorkspaceAccountID, + workspaceLabel: self.resolveWorkspaceLabel(from: candidate.cookies) ?? context.targetWorkspaceLabel, + logger: log) } private func fetchSignedInEmailFromAPI( @@ -591,9 +759,14 @@ public struct OpenAIDashboardBrowserCookieImporter { private func persist( candidate: Candidate, targetEmail: String, + workspaceAccountID: String?, + workspaceLabel: String?, logger: @escaping (String) -> Void) async throws -> ImportResult { - let persistent = OpenAIDashboardWebsiteDataStore.store(forAccountEmail: targetEmail) + let persistent = OpenAIDashboardWebsiteDataStore.store( + forAccountEmail: targetEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) await self.clearChatGPTCookies(in: persistent) await self.setCookies(candidate.cookies, into: persistent) @@ -697,11 +870,23 @@ public struct OpenAIDashboardBrowserCookieImporter { // MARK: - WebKit cookie store - private func persistCookies(candidate: Candidate, accountEmail: String, logger: (String) -> Void) async { - let store = OpenAIDashboardWebsiteDataStore.store(forAccountEmail: accountEmail) + private func persistCookies( + candidate: Candidate, + accountEmail: String, + workspaceAccountID: String?, + workspaceLabel: String?, + logger: (String) -> Void) async + { + let store = OpenAIDashboardWebsiteDataStore.store( + forAccountEmail: accountEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) await self.clearChatGPTCookies(in: store) await self.setCookies(candidate.cookies, into: store) - logger("Persisted cookies for \(accountEmail) (source=\(candidate.label))") + let workspaceSuffix = workspaceLabel.map { " [\($0)]" } ?? "" + logger( + "Persisted cookies for \(accountEmail)\(workspaceSuffix) " + + "(source=\(candidate.label))") } private func clearChatGPTCookies(in store: WKWebsiteDataStore) async { @@ -808,6 +993,8 @@ public struct OpenAIDashboardBrowserCookieImporter { public func importBestCookies( intoAccountEmail _: String?, + intoWorkspaceAccountID _: String? = nil, + intoWorkspaceLabel _: String? = nil, allowAnyAccount _: Bool = false, cacheScope _: CookieHeaderCache.Scope? = nil, logger _: ((String) -> Void)? = nil) async throws -> ImportResult @@ -818,6 +1005,8 @@ public struct OpenAIDashboardBrowserCookieImporter { public func importManualCookies( cookieHeader _: String, intoAccountEmail _: String?, + intoWorkspaceAccountID _: String? = nil, + intoWorkspaceLabel _: String? = nil, allowAnyAccount _: Bool = false, cacheScope _: CookieHeaderCache.Scope? = nil, logger _: ((String) -> Void)? = nil) async throws -> ImportResult @@ -825,4 +1014,5 @@ public struct OpenAIDashboardBrowserCookieImporter { throw ImportError.browserAccessDenied(details: "OpenAI web cookie import is only supported on macOS.") } } +// swiftlint:enable type_body_length #endif diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardFetcher.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardFetcher.swift index f9dfe030d..b18c8508b 100644 --- a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardFetcher.swift +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardFetcher.swift @@ -69,11 +69,16 @@ public struct OpenAIDashboardFetcher { public func loadLatestDashboard( accountEmail: String?, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil, logger: ((String) -> Void)? = nil, debugDumpHTML: Bool = false, timeout: TimeInterval = 60) async throws -> OpenAIDashboardSnapshot { - let store = OpenAIDashboardWebsiteDataStore.store(forAccountEmail: accountEmail) + let store = OpenAIDashboardWebsiteDataStore.store( + forAccountEmail: accountEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) return try await self.loadLatestDashboard( websiteDataStore: store, logger: logger, @@ -274,10 +279,20 @@ public struct OpenAIDashboardFetcher { return false } - public func clearSessionData(accountEmail: String?) async { - let store = OpenAIDashboardWebsiteDataStore.store(forAccountEmail: accountEmail) + public func clearSessionData( + accountEmail: String?, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil) async + { + let store = OpenAIDashboardWebsiteDataStore.store( + forAccountEmail: accountEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) OpenAIDashboardWebViewCache.shared.evict(websiteDataStore: store) - await OpenAIDashboardWebsiteDataStore.clearStore(forAccountEmail: accountEmail) + await OpenAIDashboardWebsiteDataStore.clearStore( + forAccountEmail: accountEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) } public func probeUsagePage( @@ -505,6 +520,8 @@ public struct OpenAIDashboardFetcher { public func loadLatestDashboard( accountEmail _: String?, + workspaceAccountID _: String? = nil, + workspaceLabel _: String? = nil, logger _: ((String) -> Void)? = nil, debugDumpHTML _: Bool = false, timeout _: TimeInterval = 60) async throws -> OpenAIDashboardSnapshot diff --git a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebsiteDataStore.swift b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebsiteDataStore.swift index e23ef66a1..2fd5eca97 100644 --- a/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebsiteDataStore.swift +++ b/Sources/CodexBarCore/OpenAIWeb/OpenAIDashboardWebsiteDataStore.swift @@ -6,42 +6,55 @@ import WebKit /// Per-account persistent `WKWebsiteDataStore` for the OpenAI dashboard scrape. /// /// Why: `WKWebsiteDataStore.default()` is a single shared cookie jar. If the user switches Codex accounts, -/// we want to keep multiple signed-in dashboard sessions around (one per email) without clearing cookies. +/// we want to keep multiple signed-in dashboard sessions around (one per email/workspace pair) without clearing +/// cookies. /// /// Implementation detail: macOS 14+ supports `WKWebsiteDataStore.dataStore(forIdentifier:)`, which creates -/// persistent isolated stores keyed by an identifier. We derive a stable UUID from the email so the same -/// account always maps to the same cookie store. +/// persistent isolated stores keyed by an identifier. We derive a stable UUID from the normalized email/workspace key +/// so the same account workspace always maps to the same cookie store. /// /// Important: We cache the `WKWebsiteDataStore` instances so the same object is returned for the same -/// account email. This ensures `OpenAIDashboardWebViewCache` can use object identity for cache lookups. +/// account key. This ensures `OpenAIDashboardWebViewCache` can use object identity for cache lookups. @MainActor public enum OpenAIDashboardWebsiteDataStore { - /// Cached data store instances keyed by normalized email. + /// Cached data store instances keyed by normalized account identity. /// Using the same instance ensures stable object identity for WebView cache lookups. private static var cachedStores: [String: WKWebsiteDataStore] = [:] - public static func store(forAccountEmail email: String?) -> WKWebsiteDataStore { - guard let normalized = normalizeEmail(email) else { return .default() } + public static func store( + forAccountEmail email: String?, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil) -> WKWebsiteDataStore + { + guard let normalized = normalizeAccountKey( + email: email, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) + else { return .default() } - // Return cached instance if available to maintain stable object identity if let cached = cachedStores[normalized] { return cached } - let id = Self.identifier(forNormalizedEmail: normalized) + let id = Self.identifier(forNormalizedKey: normalized) let store = WKWebsiteDataStore(forIdentifier: id) self.cachedStores[normalized] = store return store } - /// Clears the persistent cookie store for a single account email. + /// Clears the persistent cookie store for a single account identity. /// /// Note: this does *not* impact other accounts, and is safe to use when the stored session is "stuck" /// or signed in to a different account than expected. - public static func clearStore(forAccountEmail email: String?) async { - // Clear only ChatGPT/OpenAI domain data for the per-account store. - // Avoid deleting the entire persistent store (WebKit requires all WKWebViews using it to be released). - let store = self.store(forAccountEmail: email) + public static func clearStore( + forAccountEmail email: String?, + workspaceAccountID: String? = nil, + workspaceLabel: String? = nil) async + { + let store = self.store( + forAccountEmail: email, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) await withCheckedContinuation { cont in store.fetchDataRecords(ofTypes: WKWebsiteDataStore.allWebsiteDataTypes()) { records in let filtered = records.filter { record in @@ -54,8 +67,11 @@ public enum OpenAIDashboardWebsiteDataStore { } } - // Remove from cache so a fresh instance is created on next access - if let normalized = normalizeEmail(email) { + if let normalized = normalizeAccountKey( + email: email, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel) + { self.cachedStores.removeValue(forKey: normalized) } } @@ -69,13 +85,29 @@ public enum OpenAIDashboardWebsiteDataStore { // MARK: - Private - private static func normalizeEmail(_ email: String?) -> String? { - guard let raw = email?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return nil } - return raw.lowercased() + private static func normalizeAccountKey( + email: String?, + workspaceAccountID: String?, + workspaceLabel: String?) -> String? + { + guard let rawEmail = email?.trimmingCharacters(in: .whitespacesAndNewlines), !rawEmail.isEmpty else { + return nil + } + let normalizedEmail = rawEmail.lowercased() + if let normalizedWorkspaceAccountID = ManagedCodexAccount.normalizeWorkspaceAccountID(workspaceAccountID) { + return "\(normalizedEmail)\naccount:\(normalizedWorkspaceAccountID)" + } + let normalizedWorkspace = workspaceLabel? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + if let normalizedWorkspace, !normalizedWorkspace.isEmpty { + return "\(normalizedEmail)\n\(normalizedWorkspace)" + } + return normalizedEmail } - private static func identifier(forNormalizedEmail email: String) -> UUID { - let digest = SHA256.hash(data: Data(email.utf8)) + private static func identifier(forNormalizedKey key: String) -> UUID { + let digest = SHA256.hash(data: Data(key.utf8)) var bytes = Array(digest.prefix(16)) // Make it a well-formed UUID (v4 + RFC4122 variant) while staying deterministic. diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceIdentityCache.swift b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceIdentityCache.swift new file mode 100644 index 000000000..e64125c91 --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceIdentityCache.swift @@ -0,0 +1,117 @@ +import Foundation + +public struct CodexOpenAIWorkspaceIdentityCache: @unchecked Sendable { + public static let currentVersion = 1 + + private struct Payload: Codable, Sendable { + let version: Int + var labelsByWorkspaceAccountID: [String: String] + + init( + version: Int = CodexOpenAIWorkspaceIdentityCache.currentVersion, + labelsByWorkspaceAccountID: [String: String]) + { + self.version = version + self.labelsByWorkspaceAccountID = labelsByWorkspaceAccountID + } + } + + #if DEBUG + @TaskLocal static var taskFileURLOverride: URL? + #endif + + 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 workspaceLabel(for workspaceAccountID: String?) -> String? { + guard let normalizedWorkspaceAccountID = CodexOpenAIWorkspaceResolver + .normalizeWorkspaceAccountID(workspaceAccountID) + else { + return nil + } + + let payload = self.loadPayload() + return payload.labelsByWorkspaceAccountID[normalizedWorkspaceAccountID] + } + + public func store(_ identity: CodexOpenAIWorkspaceIdentity) throws { + guard let workspaceLabel = CodexOpenAIWorkspaceIdentity.normalizeWorkspaceLabel(identity.workspaceLabel) else { + return + } + + var payload = self.loadPayload() + payload.labelsByWorkspaceAccountID[identity.workspaceAccountID] = workspaceLabel + + let encoder = JSONEncoder() + encoder.outputFormatting = [.prettyPrinted, .sortedKeys] + let data = try encoder.encode(payload) + + 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() + } + + private func loadPayload() -> Payload { + guard self.fileManager.fileExists(atPath: self.fileURL.path), + let data = try? Data(contentsOf: self.fileURL), + let payload = try? JSONDecoder().decode(Payload.self, from: data), + payload.version == Self.currentVersion + else { + return Payload(labelsByWorkspaceAccountID: [:]) + } + + return payload + } + + private func applySecurePermissionsIfNeeded() throws { + #if os(macOS) + try self.fileManager.setAttributes([ + .posixPermissions: NSNumber(value: Int16(0o600)), + ], ofItemAtPath: self.fileURL.path) + #endif + } + + public static func defaultURL() -> URL { + #if DEBUG + if let override = self.taskFileURLOverride { + return override + } + #endif + + let base = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first + ?? FileManager.default.homeDirectoryForCurrentUser + return base + .appendingPathComponent("CodexBar", isDirectory: true) + .appendingPathComponent("codex-openai-workspaces.json") + } +} + +#if DEBUG +extension CodexOpenAIWorkspaceIdentityCache { + public static func withFileURLOverrideForTesting( + _ fileURL: URL, + operation: () throws -> T) rethrows -> T + { + try self.$taskFileURLOverride.withValue(fileURL) { + try operation() + } + } + + public static func withFileURLOverrideForTesting( + _ fileURL: URL, + operation: () async throws -> T) async rethrows -> T + { + try await self.$taskFileURLOverride.withValue(fileURL) { + try await operation() + } + } +} +#endif diff --git a/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceResolver.swift b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceResolver.swift new file mode 100644 index 000000000..bba68675d --- /dev/null +++ b/Sources/CodexBarCore/Providers/Codex/CodexOpenAIWorkspaceResolver.swift @@ -0,0 +1,110 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif + +public struct CodexOpenAIWorkspaceIdentity: Equatable, Sendable { + public let workspaceAccountID: String + public let workspaceLabel: String? + + public init(workspaceAccountID: String, workspaceLabel: String?) { + self.workspaceAccountID = Self.normalizeWorkspaceAccountID(workspaceAccountID) + self.workspaceLabel = Self.normalizeWorkspaceLabel(workspaceLabel) + } + + public static func normalizeWorkspaceAccountID(_ value: String) -> String { + value.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + } + + public static func normalizeWorkspaceLabel(_ value: String?) -> String? { + guard let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed + } +} + +public enum CodexOpenAIWorkspaceResolver { + private struct AccountsResponse: Decodable { + let items: [AccountItem] + } + + private struct AccountItem: Decodable { + let id: String + let name: String? + } + + private static let accountsURL = URL(string: "https://chatgpt.com/backend-api/accounts")! + + public static func resolve( + credentials: CodexOAuthCredentials, + session: URLSession = .shared) async throws -> CodexOpenAIWorkspaceIdentity? + { + guard let workspaceAccountID = normalizeWorkspaceAccountID(credentials.accountId) else { + return nil + } + + var request = URLRequest(url: self.accountsURL) + request.httpMethod = "GET" + request.timeoutInterval = 20 + request.setValue("Bearer \(credentials.accessToken)", forHTTPHeaderField: "Authorization") + request.setValue("codex-cli", forHTTPHeaderField: "User-Agent") + request.setValue("application/json", forHTTPHeaderField: "Accept") + request.setValue(workspaceAccountID, forHTTPHeaderField: "ChatGPT-Account-Id") + + let (data, response) = try await session.data(for: request) + guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else { + throw CodexOpenAIWorkspaceResolverError.invalidResponse + } + + let decoded = try JSONDecoder().decode(AccountsResponse.self, from: data) + if let account = decoded.items.first(where: { + Self.normalizeWorkspaceAccountID($0.id) == workspaceAccountID + }) { + let resolvedLabel = self.resolveWorkspaceLabel(from: account) + return CodexOpenAIWorkspaceIdentity( + workspaceAccountID: workspaceAccountID, + workspaceLabel: resolvedLabel) + } + + return CodexOpenAIWorkspaceIdentity( + workspaceAccountID: workspaceAccountID, + workspaceLabel: nil) + } + + public static func mergeAuthoritativeIdentity( + into accountInfo: AccountInfo, + authoritativeIdentity: CodexOpenAIWorkspaceIdentity?) -> AccountInfo + { + guard let authoritativeIdentity else { return accountInfo } + return AccountInfo( + email: accountInfo.email, + plan: accountInfo.plan, + workspaceLabel: authoritativeIdentity.workspaceLabel, + workspaceAccountID: authoritativeIdentity.workspaceAccountID) + } + + public static func normalizeWorkspaceAccountID(_ value: String?) -> String? { + guard let trimmed = value?.trimmingCharacters(in: .whitespacesAndNewlines), !trimmed.isEmpty else { + return nil + } + return trimmed.lowercased() + } + + private static func resolveWorkspaceLabel(from account: AccountItem) -> String? { + let name = account.name?.trimmingCharacters(in: .whitespacesAndNewlines) + if let name, !name.isEmpty { return name } + return "Personal" + } +} + +public enum CodexOpenAIWorkspaceResolverError: LocalizedError, Sendable { + case invalidResponse + + public var errorDescription: String? { + switch self { + case .invalidResponse: + "OpenAI account lookup returned an invalid response." + } + } +} diff --git a/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift index 884bb73f1..0c52ddc36 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift @@ -113,7 +113,18 @@ struct CodexCLIUsageStrategy: ProviderFetchStrategy { func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { let keepAlive = context.settings?.debugKeepCLISessionsAlive ?? false - let usage = try await context.fetcher.loadLatestUsage(keepCLISessionsAlive: keepAlive) + var usage = try await context.fetcher.loadLatestUsage(keepCLISessionsAlive: keepAlive) + if let credentials = try? CodexOAuthCredentialsStore.load(env: context.env) { + let resolvedWorkspaceIdentity = try? await CodexOpenAIWorkspaceResolver.resolve(credentials: credentials) + if let resolvedWorkspaceIdentity { + try? CodexOpenAIWorkspaceIdentityCache().store(resolvedWorkspaceIdentity) + } + if let workspaceIdentity = resolvedWorkspaceIdentity + ?? CodexOAuthFetchStrategy.cachedWorkspaceIdentity(for: credentials.accountId) + { + usage = CodexOAuthFetchStrategy.applyWorkspaceIdentity(workspaceIdentity, to: usage) + } + } let credits = await context.includeCredits ? (try? context.fetcher.loadLatestCredits(keepCLISessionsAlive: keepAlive)) : nil @@ -148,9 +159,14 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy { accessToken: credentials.accessToken, accountId: credentials.accountId, env: context.env) + let resolvedWorkspaceIdentity = try? await CodexOpenAIWorkspaceResolver.resolve(credentials: credentials) + if let resolvedWorkspaceIdentity { + try? CodexOpenAIWorkspaceIdentityCache().store(resolvedWorkspaceIdentity) + } + let workspaceIdentity = resolvedWorkspaceIdentity ?? Self.cachedWorkspaceIdentity(for: credentials.accountId) return self.makeResult( - usage: Self.mapUsage(usage, credentials: credentials), + usage: Self.mapUsage(usage, credentials: credentials, workspaceIdentity: workspaceIdentity), credits: Self.mapCredits(usage.credits), sourceLabel: "oauth") } @@ -160,7 +176,11 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy { return true } - private static func mapUsage(_ response: CodexUsageResponse, credentials: CodexOAuthCredentials) -> UsageSnapshot { + private static func mapUsage( + _ response: CodexUsageResponse, + credentials: CodexOAuthCredentials, + workspaceIdentity: CodexOpenAIWorkspaceIdentity?) -> UsageSnapshot + { let normalized = CodexRateWindowNormalizer.normalize( primary: Self.makeWindow(response.rateLimit?.primaryWindow), secondary: Self.makeWindow(response.rateLimit?.secondaryWindow)) @@ -170,7 +190,9 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy { let identity = ProviderIdentitySnapshot( providerID: .codex, accountEmail: Self.resolveAccountEmail(from: credentials), - accountOrganization: nil, + accountOrganization: workspaceIdentity?.workspaceLabel, + accountWorkspaceID: workspaceIdentity?.workspaceAccountID + ?? ManagedCodexAccount.normalizeWorkspaceAccountID(credentials.accountId), loginMethod: Self.resolvePlan(response: response, credentials: credentials)) return UsageSnapshot( @@ -222,13 +244,40 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy { let plan = (authDict?["chatgpt_plan_type"] as? String) ?? (payload["chatgpt_plan_type"] as? String) return plan?.trimmingCharacters(in: .whitespacesAndNewlines) } + + fileprivate static func applyWorkspaceIdentity( + _ workspaceIdentity: CodexOpenAIWorkspaceIdentity, + to usage: UsageSnapshot) -> UsageSnapshot + { + let existingIdentity = usage.identity(for: .codex) + let identity = ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: existingIdentity?.accountEmail, + accountOrganization: workspaceIdentity.workspaceLabel, + accountWorkspaceID: workspaceIdentity.workspaceAccountID, + loginMethod: existingIdentity?.loginMethod) + return usage.withIdentity(identity) + } + + fileprivate static func cachedWorkspaceIdentity(for workspaceAccountID: String?) -> CodexOpenAIWorkspaceIdentity? { + guard let normalizedWorkspaceAccountID = CodexOpenAIWorkspaceResolver + .normalizeWorkspaceAccountID(workspaceAccountID) + else { + return nil + } + + let cachedWorkspaceLabel = CodexOpenAIWorkspaceIdentityCache().workspaceLabel(for: normalizedWorkspaceAccountID) + return CodexOpenAIWorkspaceIdentity( + workspaceAccountID: normalizedWorkspaceAccountID, + workspaceLabel: cachedWorkspaceLabel) + } } #if DEBUG extension CodexOAuthFetchStrategy { static func _mapUsageForTesting(_ data: Data, credentials: CodexOAuthCredentials) throws -> UsageSnapshot { let usage = try JSONDecoder().decode(CodexUsageResponse.self, from: data) - return Self.mapUsage(usage, credentials: credentials) + return Self.mapUsage(usage, credentials: credentials, workspaceIdentity: nil) } } diff --git a/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift index 5a402ea5c..78d085a6f 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexWebDashboardStrategy.swift @@ -31,17 +31,29 @@ public struct CodexWebDashboardStrategy: ProviderFetchStrategy { _ = NSApplication.shared } - let accountEmail = context.fetcher.loadAccountInfo().email? + let baseAccountInfo = context.fetcher.loadAccountInfo() + let accountInfo = await Self.resolveAuthoritativeAccountInfo( + baseAccountInfo, + env: context.env) + let accountEmail = accountInfo.email? + .trimmingCharacters(in: .whitespacesAndNewlines) + let workspaceAccountID = accountInfo.workspaceAccountID? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + let workspaceLabel = accountInfo.workspaceLabel? .trimmingCharacters(in: .whitespacesAndNewlines) let options = OpenAIWebOptions( timeout: context.webTimeout, debugDumpHTML: context.webDebugDumpHTML, verbose: context.verbose) let result = try await Self.fetchOpenAIWebCodex( - accountEmail: accountEmail, - fetcher: context.fetcher, - options: options, - browserDetection: context.browserDetection) + OpenAIWebFetchRequest( + accountEmail: accountEmail, + workspaceAccountID: workspaceAccountID, + workspaceLabel: workspaceLabel, + fetcher: context.fetcher, + options: options, + browserDetection: context.browserDetection)) return self.makeResult( usage: result.usage, credits: result.credits, @@ -87,6 +99,15 @@ private struct OpenAIWebOptions { let verbose: Bool } +private struct OpenAIWebFetchRequest { + let accountEmail: String? + let workspaceAccountID: String? + let workspaceLabel: String? + let fetcher: UsageFetcher + let options: OpenAIWebOptions + let browserDetection: BrowserDetection +} + @MainActor private final class WebLogBuffer { private var lines: [String] = [] @@ -115,24 +136,54 @@ private final class WebLogBuffer { } extension CodexWebDashboardStrategy { + private static func resolveAuthoritativeAccountInfo( + _ accountInfo: AccountInfo, + env: [String: String]) async -> AccountInfo + { + guard let credentials = try? CodexOAuthCredentialsStore.load(env: env), + credentials.accountId?.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty == false + else { + return accountInfo + } + + do { + let authoritativeIdentity = try await CodexOpenAIWorkspaceResolver.resolve(credentials: credentials) + if let authoritativeIdentity { + try? CodexOpenAIWorkspaceIdentityCache().store(authoritativeIdentity) + return CodexOpenAIWorkspaceResolver.mergeAuthoritativeIdentity( + into: accountInfo, + authoritativeIdentity: authoritativeIdentity) + } + } catch { + // Keep the weaker fallback when the authoritative lookup is unavailable at runtime. + } + + if let cachedWorkspaceLabel = CodexOpenAIWorkspaceIdentityCache().workspaceLabel(for: credentials.accountId) { + return AccountInfo( + email: accountInfo.email, + plan: accountInfo.plan, + workspaceLabel: cachedWorkspaceLabel, + workspaceAccountID: credentials.accountId) + } + + return accountInfo + } + @MainActor fileprivate static func fetchOpenAIWebCodex( - accountEmail: String?, - fetcher: UsageFetcher, - options: OpenAIWebOptions, - browserDetection: BrowserDetection) async throws -> OpenAIWebCodexResult + _ request: OpenAIWebFetchRequest) async throws -> OpenAIWebCodexResult { - let logger = WebLogBuffer(verbose: options.verbose) + let logger = WebLogBuffer(verbose: request.options.verbose) let log: @MainActor (String) -> Void = { line in logger.append(line) } - let dashboard = try await Self.fetchOpenAIWebDashboard( - accountEmail: accountEmail, - fetcher: fetcher, - options: options, - browserDetection: browserDetection, - logger: log) - guard let usage = dashboard.toUsageSnapshot(provider: .codex, accountEmail: accountEmail) else { + let dashboard = try await Self.fetchOpenAIWebDashboard(request, logger: log) + guard let usage = dashboard.toUsageSnapshot( + provider: .codex, + accountEmail: request.accountEmail, + accountOrganization: request.workspaceLabel, + accountWorkspaceID: request.workspaceAccountID) + else { throw OpenAIWebCodexError.missingUsage } let credits = dashboard.toCreditsSnapshot() @@ -141,27 +192,31 @@ extension CodexWebDashboardStrategy { @MainActor fileprivate static func fetchOpenAIWebDashboard( - accountEmail: String?, - fetcher: UsageFetcher, - options: OpenAIWebOptions, - browserDetection: BrowserDetection, + _ request: OpenAIWebFetchRequest, logger: @MainActor @escaping (String) -> Void) async throws -> OpenAIDashboardSnapshot { - let trimmed = accountEmail?.trimmingCharacters(in: .whitespacesAndNewlines) - let fallback = fetcher.loadAccountInfo().email?.trimmingCharacters(in: .whitespacesAndNewlines) + let trimmed = request.accountEmail?.trimmingCharacters(in: .whitespacesAndNewlines) + let fallback = request.fetcher.loadAccountInfo().email?.trimmingCharacters(in: .whitespacesAndNewlines) let codexEmail = trimmed?.isEmpty == false ? trimmed : (fallback?.isEmpty == false ? fallback : nil) let allowAnyAccount = codexEmail == nil - let importResult = try await OpenAIDashboardBrowserCookieImporter(browserDetection: browserDetection) - .importBestCookies(intoAccountEmail: codexEmail, allowAnyAccount: allowAnyAccount, logger: logger) + let importResult = try await OpenAIDashboardBrowserCookieImporter(browserDetection: request.browserDetection) + .importBestCookies( + intoAccountEmail: codexEmail, + intoWorkspaceAccountID: request.workspaceAccountID, + intoWorkspaceLabel: request.workspaceLabel, + allowAnyAccount: allowAnyAccount, + logger: logger) let effectiveEmail = codexEmail ?? importResult.signedInEmail? .trimmingCharacters(in: .whitespacesAndNewlines) let dash = try await OpenAIDashboardFetcher().loadLatestDashboard( accountEmail: effectiveEmail, + workspaceAccountID: request.workspaceAccountID, + workspaceLabel: request.workspaceLabel, logger: logger, - debugDumpHTML: options.debugDumpHTML, - timeout: options.timeout) + debugDumpHTML: request.options.debugDumpHTML, + timeout: request.options.timeout) let cacheEmail = effectiveEmail ?? dash.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines) if let cacheEmail, !cacheEmail.isEmpty { OpenAIDashboardCacheStore.save(OpenAIDashboardCache(accountEmail: cacheEmail, snapshot: dash)) diff --git a/Sources/CodexBarCore/UsageFetcher.swift b/Sources/CodexBarCore/UsageFetcher.swift index 47051dffa..62886ca05 100644 --- a/Sources/CodexBarCore/UsageFetcher.swift +++ b/Sources/CodexBarCore/UsageFetcher.swift @@ -23,17 +23,20 @@ public struct ProviderIdentitySnapshot: Codable, Sendable { public let providerID: UsageProvider? public let accountEmail: String? public let accountOrganization: String? + public let accountWorkspaceID: String? public let loginMethod: String? public init( providerID: UsageProvider?, accountEmail: String?, accountOrganization: String?, + accountWorkspaceID: String? = nil, loginMethod: String?) { self.providerID = providerID self.accountEmail = accountEmail self.accountOrganization = accountOrganization + self.accountWorkspaceID = accountWorkspaceID?.trimmingCharacters(in: .whitespacesAndNewlines) self.loginMethod = loginMethod } @@ -43,6 +46,7 @@ public struct ProviderIdentitySnapshot: Codable, Sendable { providerID: provider, accountEmail: self.accountEmail, accountOrganization: self.accountOrganization, + accountWorkspaceID: self.accountWorkspaceID, loginMethod: self.loginMethod) } } @@ -69,6 +73,7 @@ public struct UsageSnapshot: Codable, Sendable { case identity case accountEmail case accountOrganization + case accountWorkspaceID case loginMethod } @@ -112,12 +117,14 @@ public struct UsageSnapshot: Codable, Sendable { } else { let email = try container.decodeIfPresent(String.self, forKey: .accountEmail) let organization = try container.decodeIfPresent(String.self, forKey: .accountOrganization) + let workspaceID = try container.decodeIfPresent(String.self, forKey: .accountWorkspaceID) let loginMethod = try container.decodeIfPresent(String.self, forKey: .loginMethod) - if email != nil || organization != nil || loginMethod != nil { + if email != nil || organization != nil || workspaceID != nil || loginMethod != nil { self.identity = ProviderIdentitySnapshot( providerID: nil, accountEmail: email, accountOrganization: organization, + accountWorkspaceID: workspaceID, loginMethod: loginMethod) } else { self.identity = nil @@ -137,6 +144,7 @@ public struct UsageSnapshot: Codable, Sendable { try container.encodeIfPresent(self.identity, forKey: .identity) try container.encodeIfPresent(self.identity?.accountEmail, forKey: .accountEmail) try container.encodeIfPresent(self.identity?.accountOrganization, forKey: .accountOrganization) + try container.encodeIfPresent(self.identity?.accountWorkspaceID, forKey: .accountWorkspaceID) try container.encodeIfPresent(self.identity?.loginMethod, forKey: .loginMethod) } @@ -205,6 +213,10 @@ public struct UsageSnapshot: Codable, Sendable { self.identity(for: provider)?.accountOrganization } + public func accountWorkspaceID(for provider: UsageProvider) -> String? { + self.identity(for: provider)?.accountWorkspaceID + } + public func loginMethod(for provider: UsageProvider) -> String? { self.identity(for: provider)?.loginMethod } @@ -242,10 +254,22 @@ public struct UsageSnapshot: Codable, Sendable { public struct AccountInfo: Equatable, Sendable { public let email: String? public let plan: String? + public let workspaceLabel: String? + public let workspaceAccountID: String? - public init(email: String?, plan: String?) { + public init( + email: String?, + plan: String?, + workspaceLabel: String? = nil, + workspaceAccountID: String? = nil) + { self.email = email self.plan = plan + self.workspaceLabel = workspaceLabel? + .trimmingCharacters(in: .whitespacesAndNewlines) + self.workspaceAccountID = workspaceAccountID? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() } } @@ -685,8 +709,26 @@ public struct UsageFetcher: Sendable { let email = (payload["email"] as? String) ?? (profileDict?["email"] as? String) + let workspaceAccountID = credentials.accountId?.trimmingCharacters(in: .whitespacesAndNewlines) + let workspaceLabel = CodexOpenAIWorkspaceIdentityCache().workspaceLabel(for: workspaceAccountID) + ?? Self.resolveOpenAIWorkspaceLabel(authPayload: authDict) + + return AccountInfo( + email: email, + plan: plan, + workspaceLabel: workspaceLabel, + workspaceAccountID: workspaceAccountID) + } + + private static func resolveOpenAIWorkspaceLabel(authPayload: [String: Any]?) -> String? { + guard let organizations = authPayload?["organizations"] as? [[String: Any]], !organizations.isEmpty else { + return nil + } - return AccountInfo(email: email, plan: plan) + let selected = organizations.first(where: { ($0["is_default"] as? Bool) == true }) ?? organizations.first + let title = (selected?["title"] as? String)?.trimmingCharacters(in: .whitespacesAndNewlines) + guard let title, !title.isEmpty else { return nil } + return title } // MARK: - Helpers diff --git a/Tests/CodexBarTests/CodexAccountReconciliationTests.swift b/Tests/CodexBarTests/CodexAccountReconciliationTests.swift index 89c7593a5..b779c6681 100644 --- a/Tests/CodexBarTests/CodexAccountReconciliationTests.swift +++ b/Tests/CodexBarTests/CodexAccountReconciliationTests.swift @@ -291,6 +291,7 @@ struct CodexAccountReconciliationTests { let accounts = ManagedCodexAccountSet(version: 1, accounts: [stored]) let live = ObservedSystemCodexAccount( email: "USER@example.com", + workspaceLabel: nil, codexHomePath: "/Users/test/.codex", observedAt: Date()) let reconciler = DefaultCodexAccountReconciler( @@ -305,6 +306,68 @@ struct CodexAccountReconciliationTests { #expect(projection.liveVisibleAccountID == "user@example.com") } + @Test + func `same email different workspace remains as separate visible accounts`() { + let stored = ManagedCodexAccount( + id: UUID(), + email: "user@example.com", + workspaceLabel: "Team Alpha", + managedHomePath: "/tmp/managed-a", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let accounts = ManagedCodexAccountSet(version: 1, accounts: [stored]) + let live = ObservedSystemCodexAccount( + email: "USER@example.com", + workspaceLabel: "Personal", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let reconciler = DefaultCodexAccountReconciler( + storeLoader: { accounts }, + systemObserver: StubSystemObserver(account: live), + activeSource: .managedAccount(id: stored.id)) + + let projection = reconciler.loadVisibleAccounts(environment: [:]) + + #expect(projection.visibleAccounts.count == 2) + #expect(Set(projection.visibleAccounts.map(\.displayName)) == [ + "user@example.com — Personal", + "user@example.com — Team Alpha", + ]) + } + + @Test + func `same email same label different workspace account ids remain separate visible accounts`() { + let stored = ManagedCodexAccount( + id: UUID(), + email: "user@example.com", + workspaceLabel: "Shared Team", + workspaceAccountID: "team-a", + managedHomePath: "/tmp/managed-a", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let accounts = ManagedCodexAccountSet(version: 1, accounts: [stored]) + let live = ObservedSystemCodexAccount( + email: "USER@example.com", + workspaceLabel: "Shared Team", + workspaceAccountID: "team-b", + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let reconciler = DefaultCodexAccountReconciler( + storeLoader: { accounts }, + systemObserver: StubSystemObserver(account: live), + activeSource: .managedAccount(id: stored.id)) + + let projection = reconciler.loadVisibleAccounts(environment: [:]) + + #expect(projection.visibleAccounts.count == 2) + #expect(Set(projection.visibleAccounts.map(\.id)) == [ + "user@example.com\naccount:team-a", + "user@example.com\naccount:team-b", + ]) + } + @Test func `matching live system account resolves merged row selection to live system`() { let stored = ManagedCodexAccount( @@ -336,6 +399,43 @@ struct CodexAccountReconciliationTests { #expect(projection.source(forVisibleAccountID: "user@example.com") == .liveSystem) } + @Test + func `matching live system account prefers live authoritative workspace label`() throws { + let workspaceAccountID = "team-123" + let stored = ManagedCodexAccount( + id: UUID(), + email: "user@example.com", + workspaceLabel: "Personal", + workspaceAccountID: workspaceAccountID, + managedHomePath: "/tmp/managed-a", + createdAt: 1, + updatedAt: 2, + lastAuthenticatedAt: 3) + let live = ObservedSystemCodexAccount( + email: "USER@example.com", + workspaceLabel: "Team Alpha", + workspaceAccountID: workspaceAccountID, + codexHomePath: "/Users/test/.codex", + observedAt: Date()) + let snapshot = CodexAccountReconciliationSnapshot( + storedAccounts: [stored], + activeStoredAccount: stored, + liveSystemAccount: live, + matchingStoredAccountForLiveSystemAccount: stored, + activeSource: .managedAccount(id: stored.id), + hasUnreadableAddedAccountStore: false) + + let projection = CodexVisibleAccountProjection.make(from: snapshot) + let visible = try #require(projection.visibleAccounts.first) + + #expect(projection.visibleAccounts.count == 1) + #expect(visible.id == "user@example.com\naccount:team-123") + #expect(visible.workspaceLabel == "Team Alpha") + #expect(visible.displayName == "user@example.com — Team Alpha") + #expect(visible.selectionSource == .liveSystem) + #expect(visible.isLive) + } + @Test func `missing managed source resolves to live system when live account exists`() { let live = ObservedSystemCodexAccount( diff --git a/Tests/CodexBarTests/CodexOpenAIWorkspaceResolverTests.swift b/Tests/CodexBarTests/CodexOpenAIWorkspaceResolverTests.swift new file mode 100644 index 000000000..a0ce9d3fa --- /dev/null +++ b/Tests/CodexBarTests/CodexOpenAIWorkspaceResolverTests.swift @@ -0,0 +1,97 @@ +import Foundation +#if canImport(FoundationNetworking) +import FoundationNetworking +#endif +import Testing +@testable import CodexBarCore + +@Suite(.serialized) +struct CodexOpenAIWorkspaceResolverTests { + @Test + func `resolver returns selected team workspace from accounts endpoint`() async throws { + let session = URLSession(configuration: ResolverURLProtocol.configuration(jsonObject: [ + "items": [ + ["id": "team-123", "name": "IDconcepts"], + ["id": "personal-456", "name": NSNull()], + ], + ])) + let credentials = CodexOAuthCredentials( + accessToken: "access-token", + refreshToken: "refresh-token", + idToken: nil, + accountId: "TEAM-123", + lastRefresh: nil) + + let resolved = try await CodexOpenAIWorkspaceResolver.resolve(credentials: credentials, session: session) + + #expect(resolved?.workspaceAccountID == "team-123") + #expect(resolved?.workspaceLabel == "IDconcepts") + #expect(ResolverURLProtocol.lastRequest?.value(forHTTPHeaderField: "ChatGPT-Account-Id") == "team-123") + } + + @Test + func `resolver maps unnamed selected account to personal`() async throws { + let session = URLSession(configuration: ResolverURLProtocol.configuration(jsonObject: [ + "items": [ + ["id": "personal-456", "name": NSNull()], + ], + ])) + let credentials = CodexOAuthCredentials( + accessToken: "access-token", + refreshToken: "refresh-token", + idToken: nil, + accountId: "personal-456", + lastRefresh: nil) + + let resolved = try await CodexOpenAIWorkspaceResolver.resolve(credentials: credentials, session: session) + + #expect(resolved?.workspaceAccountID == "personal-456") + #expect(resolved?.workspaceLabel == "Personal") + } +} + +private class ResolverURLProtocol: URLProtocol { + private static let lock = NSLock() + private nonisolated(unsafe) static var responseBody = Data() + private nonisolated(unsafe) static var responseStatusCode = 200 + nonisolated(unsafe) static var lastRequest: URLRequest? + + static func configuration(jsonObject: Any, statusCode: Int = 200) -> URLSessionConfiguration { + self.lock.lock() + self.responseBody = (try? JSONSerialization.data(withJSONObject: jsonObject)) ?? Data() + self.responseStatusCode = statusCode + self.lastRequest = nil + self.lock.unlock() + + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [ResolverURLProtocol.self] + return configuration + } + + override class func canInit(with request: URLRequest) -> Bool { + true + } + + override class func canonicalRequest(for request: URLRequest) -> URLRequest { + request + } + + override func startLoading() { + Self.lock.lock() + Self.lastRequest = self.request + let body = Self.responseBody + let statusCode = Self.responseStatusCode + Self.lock.unlock() + + let response = HTTPURLResponse( + url: self.request.url!, + statusCode: statusCode, + httpVersion: nil, + headerFields: ["Content-Type": "application/json"])! + self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) + self.client?.urlProtocol(self, didLoad: body) + self.client?.urlProtocolDidFinishLoading(self) + } + + override func stopLoading() {} +} diff --git a/Tests/CodexBarTests/CodexSystemAccountObserverTests.swift b/Tests/CodexBarTests/CodexSystemAccountObserverTests.swift index cf2928d86..d060e5b5f 100644 --- a/Tests/CodexBarTests/CodexSystemAccountObserverTests.swift +++ b/Tests/CodexBarTests/CodexSystemAccountObserverTests.swift @@ -44,25 +44,74 @@ struct CodexSystemAccountObserverTests { #expect(observed.observedAt >= before) } - private static func writeCodexAuthFile(homeURL: URL, email: String, plan: String) throws { + @Test + func `observer prefers cached authoritative workspace label over weak jwt label`() throws { + let home = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + let cacheURL = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { + try? FileManager.default.removeItem(at: home) + try? FileManager.default.removeItem(at: cacheURL) + } + + try Self.writeCodexAuthFile( + homeURL: home, + email: "user@example.com", + plan: "team", + accountId: "TEAM-123", + workspaceTitle: "Personal") + + try CodexOpenAIWorkspaceIdentityCache.withFileURLOverrideForTesting(cacheURL) { + try CodexOpenAIWorkspaceIdentityCache().store( + CodexOpenAIWorkspaceIdentity( + workspaceAccountID: "team-123", + workspaceLabel: "IDconcepts")) + + let observer = DefaultCodexSystemAccountObserver() + let account = try observer.loadSystemAccount(environment: ["CODEX_HOME": home.path]) + + #expect(account?.workspaceAccountID == "team-123") + #expect(account?.workspaceLabel == "IDconcepts") + } + } + + private static func writeCodexAuthFile( + homeURL: URL, + email: String, + plan: String, + accountId: String? = nil, + workspaceTitle: String? = nil) throws + { try FileManager.default.createDirectory(at: homeURL, withIntermediateDirectories: true) - let auth = [ - "tokens": [ - "accessToken": "access-token", - "refreshToken": "refresh-token", - "idToken": Self.fakeJWT(email: email, plan: plan), - ], + var tokens: [String: Any] = [ + "accessToken": "access-token", + "refreshToken": "refresh-token", + "idToken": Self.fakeJWT(email: email, plan: plan, workspaceTitle: workspaceTitle), ] + if let accountId { + tokens["account_id"] = accountId + } + let auth = ["tokens": tokens] let data = try JSONSerialization.data(withJSONObject: auth) try data.write(to: homeURL.appendingPathComponent("auth.json")) } - private static func fakeJWT(email: String, plan: String) -> String { + private static func fakeJWT(email: String, plan: String, workspaceTitle: String? = nil) -> String { let header = (try? JSONSerialization.data(withJSONObject: ["alg": "none"])) ?? Data() - let payload = (try? JSONSerialization.data(withJSONObject: [ + var payload: [String: Any] = [ "email": email, "chatgpt_plan_type": plan, - ])) ?? Data() + ] + if let workspaceTitle { + payload["https://api.openai.com/auth"] = [ + "organizations": [ + [ + "title": workspaceTitle, + "is_default": true, + ], + ], + ] + } + let payloadData = (try? JSONSerialization.data(withJSONObject: payload)) ?? Data() func base64URL(_ data: Data) -> String { data.base64EncodedString() @@ -71,6 +120,6 @@ struct CodexSystemAccountObserverTests { .replacingOccurrences(of: "/", with: "_") } - return "\(base64URL(header)).\(base64URL(payload))." + return "\(base64URL(header)).\(base64URL(payloadData))." } } diff --git a/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift b/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift index 67256f679..49c747c58 100644 --- a/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift +++ b/Tests/CodexBarTests/ManagedCodexAccountServiceTests.swift @@ -58,6 +58,97 @@ struct ManagedCodexAccountServiceTests { #expect(authenticated.email == "second@example.com") } + @Test + func `same email different workspace creates distinct managed accounts`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet(version: 1, accounts: [])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.accounts([ + AccountInfo(email: "user@example.com", plan: "Team", workspaceLabel: "Personal"), + AccountInfo(email: "user@example.com", plan: "Team", workspaceLabel: "Team Alpha"), + ])) + + let first = try await service.authenticateManagedAccount() + let second = try await service.authenticateManagedAccount() + + #expect(first.id != second.id) + #expect(store.snapshot.accounts.count == 2) + #expect(Set(store.snapshot.accounts.map { $0.workspaceLabel ?? "" }) == ["Personal", "Team Alpha"]) + } + + @Test + func `authoritative workspace resolver overrides weak personal label during authentication`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet(version: 1, accounts: [])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.accounts([ + AccountInfo( + email: "user@example.com", + plan: "Team", + workspaceLabel: "Personal", + workspaceAccountID: "team-123"), + ]), + workspaceResolver: StubManagedCodexWorkspaceResolver.accounts([ + AccountInfo( + email: "user@example.com", + plan: "Team", + workspaceLabel: "IDconcepts", + workspaceAccountID: "team-123"), + ])) + + let account = try await service.authenticateManagedAccount() + + #expect(account.workspaceLabel == "IDconcepts") + #expect(account.workspaceAccountID == "team-123") + #expect(store.snapshot.accounts.first?.workspaceLabel == "IDconcepts") + #expect(store.snapshot.accounts.first?.workspaceAccountID == "team-123") + } + + @Test + func `same email same label different workspace account ids create distinct managed accounts`() async throws { + let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) + defer { try? FileManager.default.removeItem(at: root) } + + let store = InMemoryManagedCodexAccountStore( + accounts: ManagedCodexAccountSet(version: 1, accounts: [])) + let service = ManagedCodexAccountService( + store: store, + homeFactory: TestManagedCodexHomeFactory(root: root), + loginRunner: StubManagedCodexLoginRunner.success, + identityReader: StubManagedCodexIdentityReader.accounts([ + AccountInfo( + email: "user@example.com", + plan: "Team", + workspaceLabel: "Shared Team", + workspaceAccountID: "team-a"), + AccountInfo( + email: "user@example.com", + plan: "Team", + workspaceLabel: "Shared Team", + workspaceAccountID: "team-b"), + ]), + workspaceResolver: StubManagedCodexWorkspaceResolver.passThrough) + + let first = try await service.authenticateManagedAccount() + let second = try await service.authenticateManagedAccount() + + #expect(first.id != second.id) + #expect(store.snapshot.accounts.count == 2) + #expect(Set(store.snapshot.accounts.compactMap(\.workspaceAccountID)) == ["team-a", "team-b"]) + } + @Test func `reauth keeps previous home when store write fails`() async throws { let root = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString, isDirectory: true) @@ -409,20 +500,49 @@ private enum TestManagedCodexAccountStoreError: Error, Equatable { private final class StubManagedCodexIdentityReader: ManagedCodexIdentityReading, @unchecked Sendable { private let lock = NSLock() - private var emails: [String] + private var accounts: [AccountInfo] - init(emails: [String]) { - self.emails = emails + init(accounts: [AccountInfo]) { + self.accounts = accounts } 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") + if self.accounts.isEmpty { + return AccountInfo(email: nil, plan: "Pro") + } + return self.accounts.removeFirst() } static func emails(_ emails: [String]) -> StubManagedCodexIdentityReader { - StubManagedCodexIdentityReader(emails: emails) + StubManagedCodexIdentityReader(accounts: emails.map { AccountInfo(email: $0, plan: "Pro") }) + } + + static func accounts(_ accounts: [AccountInfo]) -> StubManagedCodexIdentityReader { + StubManagedCodexIdentityReader(accounts: accounts) + } +} + +private actor StubManagedCodexWorkspaceResolver: ManagedCodexWorkspaceResolving { + private var accounts: [AccountInfo]? + + init(accounts: [AccountInfo]?) { + self.accounts = accounts + } + + func resolveAccountInfo(homePath _: String, fallback: AccountInfo) async -> AccountInfo { + guard var accounts = self.accounts, !accounts.isEmpty else { + return fallback + } + let next = accounts.removeFirst() + self.accounts = accounts + return next + } + + static let passThrough = StubManagedCodexWorkspaceResolver(accounts: nil) + + static func accounts(_ accounts: [AccountInfo]) -> StubManagedCodexWorkspaceResolver { + StubManagedCodexWorkspaceResolver(accounts: accounts) } } diff --git a/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift b/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift index 9925ad0c6..2302029ea 100644 --- a/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift +++ b/Tests/CodexBarTests/ManagedCodexAccountStoreTests.swift @@ -152,6 +152,44 @@ func `FileManagedCodexAccountStore drops duplicate canonical emails on load`() t #expect(loaded.accounts.first?.managedHomePath == "/tmp/managed-home-1") } +@Test +func `FileManagedCodexAccountStore keeps same email accounts when workspace differs`() throws { + let tempDir = FileManager.default.temporaryDirectory + let fileURL = tempDir.appendingPathComponent("codexbar-managed-codex-accounts-workspace-identity-test.json") + defer { try? FileManager.default.removeItem(at: fileURL) } + + let personalID = UUID() + let teamID = UUID() + let payload = ManagedCodexAccountSet( + version: 1, + accounts: [ + ManagedCodexAccount( + id: personalID, + email: "same@example.com", + workspaceLabel: "Personal", + managedHomePath: "/tmp/managed-home-1", + createdAt: 10, + updatedAt: 20, + lastAuthenticatedAt: nil), + ManagedCodexAccount( + id: teamID, + email: "same@example.com", + workspaceLabel: "Team Alpha", + managedHomePath: "/tmp/managed-home-2", + createdAt: 30, + updatedAt: 40, + lastAuthenticatedAt: nil), + ]) + let store = FileManagedCodexAccountStore(fileURL: fileURL) + + try store.storeAccounts(payload) + let loaded = try store.loadAccounts() + + #expect(loaded.accounts.count == 2) + #expect(loaded.account(email: "same@example.com", workspaceLabel: "personal")?.id == personalID) + #expect(loaded.account(email: "same@example.com", workspaceLabel: "team alpha")?.id == teamID) +} + @Test func `FileManagedCodexAccountStore drops duplicate IDs on load`() throws { let tempDir = FileManager.default.temporaryDirectory diff --git a/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift b/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift index 4f5a1d45c..20c96fcff 100644 --- a/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift +++ b/Tests/CodexBarTests/OpenAIDashboardWebViewCacheTests.swift @@ -32,6 +32,26 @@ struct OpenAIDashboardWebViewCacheTests { OpenAIDashboardWebsiteDataStore.clearCacheForTesting() } + @Test + func `WKWebsiteDataStore should return different instances for same email in different workspaces`() { + OpenAIDashboardWebsiteDataStore.clearCacheForTesting() + + let personal = OpenAIDashboardWebsiteDataStore.store( + forAccountEmail: "test@example.com", + workspaceLabel: "Personal") + let team = OpenAIDashboardWebsiteDataStore.store( + forAccountEmail: "test@example.com", + workspaceLabel: "Team Alpha") + let personalCaseVariant = OpenAIDashboardWebsiteDataStore.store( + forAccountEmail: "TEST@example.com", + workspaceLabel: "personal") + + #expect(personal !== team, "Same email with different workspaces should use separate stores") + #expect(personal === personalCaseVariant, "Workspace comparison should be case-insensitive") + + OpenAIDashboardWebsiteDataStore.clearCacheForTesting() + } + // MARK: - WebView Reuse Tests @Test diff --git a/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift index 97d197abf..63d7cc1a4 100644 --- a/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift +++ b/Tests/CodexBarTests/StatusMenuCodexSwitcherTests.swift @@ -163,6 +163,49 @@ struct StatusMenuCodexSwitcherTests { #expect(self.actionLabels(in: descriptor).contains("Add Account...")) } + @Test + func `codex switcher compacts same email pills while preserving email and workspace`() { + let accounts = [ + CodexVisibleAccount( + id: "pl.fr@yandex.com\naccount:personal", + email: "pl.fr@yandex.com", + workspaceLabel: "Personal", + storedAccountID: nil, + selectionSource: .liveSystem, + isActive: true, + isLive: true, + canReauthenticate: true, + canRemove: false), + CodexVisibleAccount( + id: "pl.fr@yandex.com\naccount:idconcepts", + email: "pl.fr@yandex.com", + workspaceLabel: "IDconcepts", + storedAccountID: nil, + selectionSource: .managedAccount(id: UUID()), + isActive: false, + isLive: false, + canReauthenticate: true, + canRemove: true), + ] + + let view = CodexAccountSwitcherView( + accounts: accounts, + selectedAccountID: accounts.first?.id, + width: 220, + onSelect: { _ in }) + + let titles = view._test_buttonTitles() + let toolTips = view._test_buttonToolTips() + + #expect(titles.count == 2) + #expect(titles[0] != titles[1]) + #expect(titles.allSatisfy { $0.contains("|") }) + #expect(titles.allSatisfy { $0.lowercased().contains("pl.") }) + #expect(titles[0].lowercased().contains("pers")) + #expect(titles[1].lowercased().contains("id")) + #expect(toolTips == accounts.map(\.displayName)) + } + @Test func `codex menu switcher selection activates the visible managed account`() throws { self.disableMenuCardsForTesting()