diff --git a/Package.resolved b/Package.resolved index f84c0c217..268a030f0 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,6 +1,15 @@ { - "originHash" : "74bd6f3ab6e0b0cb0c2cddb00f2167c2ab0a1c00cd54ffc1a2899c7ef8c56367", + "originHash" : "923ccd04f0147a3e487f150ce686d3de211781782e8d9b295152f8484b87b0e6", "pins" : [ + { + "identity" : "appauth-ios", + "kind" : "remoteSourceControl", + "location" : "https://github.com/openid/AppAuth-iOS.git", + "state" : { + "revision" : "2781038865a80e2c425a1da12cc1327bcd56501f", + "version" : "1.7.6" + } + }, { "identity" : "commander", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index 83cbfca65..5dce42f8f 100644 --- a/Package.swift +++ b/Package.swift @@ -22,6 +22,7 @@ let package = Package( .package(url: "https://github.com/apple/swift-log", from: "1.9.1"), .package(url: "https://github.com/apple/swift-syntax", from: "600.0.1"), .package(url: "https://github.com/sindresorhus/KeyboardShortcuts", from: "2.4.0"), + .package(url: "https://github.com/openid/AppAuth-iOS.git", from: "1.7.6"), sweetCookieKitDependency, ], targets: { @@ -82,6 +83,7 @@ let package = Package( dependencies: [ .product(name: "Sparkle", package: "Sparkle"), .product(name: "KeyboardShortcuts", package: "KeyboardShortcuts"), + .product(name: "AppAuth", package: "AppAuth-iOS"), "CodexBarMacroSupport", "CodexBarCore", ], diff --git a/Scripts/package_app.sh b/Scripts/package_app.sh index 08ee481a0..e33e8dee0 100755 --- a/Scripts/package_app.sh +++ b/Scripts/package_app.sh @@ -317,7 +317,12 @@ elif [[ "$ALLOW_LLDB" == "1" ]]; then CODESIGN_ARGS=(--force --sign "$CODESIGN_ID") else CODESIGN_ID="${APP_IDENTITY:-Developer ID Application: Peter Steinberger (Y5PE65HELJ)}" - CODESIGN_ARGS=(--force --timestamp --options runtime --sign "$CODESIGN_ID") + if [[ "$CODESIGN_ID" == "CodexBar Development" ]]; then + echo "INFO: '${CODESIGN_ID}' has no Team ID; signing without hardened runtime for local dev stability." >&2 + CODESIGN_ARGS=(--force --sign "$CODESIGN_ID") + else + CODESIGN_ARGS=(--force --timestamp --options runtime --sign "$CODESIGN_ID") + fi fi function resign() { codesign "${CODESIGN_ARGS[@]}" "$1"; } # Sign innermost binaries first, then the framework root to seal resources diff --git a/Sources/CodexBar/CodexLoginRunner.swift b/Sources/CodexBar/CodexLoginRunner.swift index 8f1f654f2..193a3f05d 100644 --- a/Sources/CodexBar/CodexLoginRunner.swift +++ b/Sources/CodexBar/CodexLoginRunner.swift @@ -1,8 +1,14 @@ +import AppAuth +import AppKit import CodexBarCore -import Darwin import Foundation struct CodexLoginRunner { + enum Phase { + case requesting + case waitingBrowser + } + struct Result { enum Outcome { case success @@ -16,139 +22,187 @@ struct CodexLoginRunner { let output: String } - static func run(timeout: TimeInterval = 120) async -> Result { - await Task(priority: .userInitiated) { - var env = ProcessInfo.processInfo.environment - env["PATH"] = PathBuilder.effectivePATH( - purposes: [.rpc, .tty, .nodeTooling], - env: env, - loginPATH: LoginShellPathCache.shared.current) - - guard let executable = BinaryLocator.resolveCodexBinary( - env: env, - loginPATH: LoginShellPathCache.shared.current) - else { - return Result(outcome: .missingBinary, output: "") - } - - let process = Process() - process.executableURL = URL(fileURLWithPath: "/usr/bin/env") - process.arguments = [executable, "login"] - process.environment = env - - let stdout = Pipe() - let stderr = Pipe() - process.standardOutput = stdout - process.standardError = stderr - - var processGroup: pid_t? - do { - try process.run() - processGroup = self.attachProcessGroup(process) - } catch { - return Result(outcome: .launchFailed(error.localizedDescription), output: "") - } - - let timedOut = await self.wait(for: process, timeout: timeout) - if timedOut { - self.terminate(process, processGroup: processGroup) - } - - let output = await self.combinedOutput(stdout: stdout, stderr: stderr) - if timedOut { - return Result(outcome: .timedOut, output: output) - } - - let status = process.terminationStatus - if status == 0 { - return Result(outcome: .success, output: output) - } - return Result(outcome: .failed(status: status), output: output) - }.value - } - - private static func wait(for process: Process, timeout: TimeInterval) async -> Bool { - await withTaskGroup(of: Bool.self) { group -> Bool in + private static let authorizationEndpoint = URL(string: "https://auth.openai.com/oauth/authorize")! + private static let tokenEndpoint = URL(string: "https://auth.openai.com/oauth/token")! + private static let redirectPort: UInt16 = 1455 + private static let redirectURL = URL(string: "http://localhost:1455/auth/callback")! + private static let successURL = URL(string: "https://chatgpt.com/")! + private static let clientID = "app_EMoamEEZ73f0CkXaXp7hrann" + private static let scopes = ["openid", "profile", "email", "offline_access"] + private static let additionalParameters = [ + "id_token_add_organizations": "true", + "codex_cli_simplified_flow": "true", + "originator": "codex_cli_rs", + ] + + @MainActor + static func run( + timeout: TimeInterval = 120, + onPhaseChange: (@Sendable (Phase) -> Void)? = nil, + credentialSource: String? = nil) async -> Result + { + let coordinator = NativeOAuthCoordinator( + onPhaseChange: onPhaseChange, + credentialSource: credentialSource) + return await withTaskGroup(of: Result.self) { group in group.addTask { - process.waitUntilExit() - return false + await coordinator.run() } group.addTask { let nanos = UInt64(max(0, timeout) * 1_000_000_000) try? await Task.sleep(nanoseconds: nanos) - return true + return await coordinator.cancelForTimeout() } - let result = await group.next() ?? false + let result = await group.next() ?? Result(outcome: .timedOut, output: "Codex login timed out.") group.cancelAll() return result } } - private static func terminate(_ process: Process, processGroup: pid_t?) { - if let pgid = processGroup { - kill(-pgid, SIGTERM) - } - if process.isRunning { - process.terminate() - } + static func _authorizationRequestForTesting(listenerBaseURL: URL) -> OIDAuthorizationRequest { + Self.makeAuthorizationRequest(listenerBaseURL: listenerBaseURL) + } - let deadline = Date().addingTimeInterval(2.0) - while process.isRunning, Date() < deadline { - usleep(100_000) + @MainActor + private final class NativeOAuthCoordinator { + private let onPhaseChange: (@Sendable (Phase) -> Void)? + private let credentialSource: String? + private let lock = NSLock() + private var continuation: CheckedContinuation? + private var hasCompleted = false + private var redirectHandler: OIDRedirectHTTPHandler? + + init(onPhaseChange: (@Sendable (Phase) -> Void)?, credentialSource: String?) { + self.onPhaseChange = onPhaseChange + self.credentialSource = credentialSource } - if process.isRunning { - if let pgid = processGroup { - kill(-pgid, SIGKILL) + func run() async -> Result { + await withCheckedContinuation { continuation in + self.continuation = continuation + self.start() } - kill(process.processIdentifier, SIGKILL) } - } - private static func attachProcessGroup(_ process: Process) -> pid_t? { - let pid = process.processIdentifier - return setpgid(pid, pid) == 0 ? pid : nil - } + func cancelForTimeout() -> Result { + self.redirectHandler?.cancelHTTPListener() + let result = Result(outcome: .timedOut, output: "Codex login timed out.") + self.finish(result) + return result + } - private static func combinedOutput(stdout: Pipe, stderr: Pipe) async -> String { - async let out = self.readToEnd(stdout) - async let err = self.readToEnd(stderr) - let stdoutText = await out - let stderrText = await err + private func start() { + self.onPhaseChange?(.requesting) + + let configuration = OIDServiceConfiguration( + authorizationEndpoint: CodexLoginRunner.authorizationEndpoint, + tokenEndpoint: CodexLoginRunner.tokenEndpoint) + let redirectHandler = OIDRedirectHTTPHandler(successURL: CodexLoginRunner.successURL) + self.redirectHandler = redirectHandler + + var listenerError: NSError? + let listenerBaseURL = redirectHandler.startHTTPListener( + &listenerError, + withPort: CodexLoginRunner.redirectPort) + if let listenerError { + let message = listenerError.localizedDescription + self.finish(Result(outcome: .launchFailed(message), output: message)) + return + } - let merged: String = if !stdoutText.isEmpty, !stderrText.isEmpty { - [stdoutText, stderrText].joined(separator: "\n") - } else { - stdoutText + stderrText + let request = CodexLoginRunner.makeAuthorizationRequest( + configuration: configuration, + listenerBaseURL: listenerBaseURL) + + self.onPhaseChange?(.waitingBrowser) + redirectHandler.currentAuthorizationFlow = OIDAuthState + .authState(byPresenting: request) { authState, error in + NSRunningApplication.current.activate(options: [.activateAllWindows, .activateIgnoringOtherApps]) + redirectHandler.cancelHTTPListener() + + if let authState { + do { + let credentials = try Self.credentials(from: authState) + if let credentialSource = self.credentialSource { + try CodexOAuthCredentialsStore.save(credentials, rawSource: credentialSource) + } else { + try CodexOAuthCredentialsStore.save(credentials) + } + self.finish(Result(outcome: .success, output: "Codex OAuth login complete.")) + } catch { + self.finish(Result( + outcome: .launchFailed(error.localizedDescription), + output: error.localizedDescription)) + } + return + } + + let message = error?.localizedDescription ?? "Unknown OAuth login error." + self.finish(Result(outcome: .failed(status: 1), output: message)) + } } - let trimmed = merged.trimmingCharacters(in: .whitespacesAndNewlines) - let limited = trimmed.prefix(4000) - return limited.isEmpty ? "No output captured." : String(limited) - } - private static func readToEnd(_ pipe: Pipe, timeout: TimeInterval = 3.0) async -> String { - await withTaskGroup(of: String?.self) { group -> String in - group.addTask { - if #available(macOS 13.0, *) { - if let data = try? pipe.fileHandleForReading.readToEnd() { return self.decode(data) } - } - let data = pipe.fileHandleForReading.readDataToEndOfFile() - return Self.decode(data) + private func finish(_ result: Result) { + self.lock.lock() + defer { self.lock.unlock() } + guard !self.hasCompleted, let continuation = self.continuation else { return } + self.hasCompleted = true + self.continuation = nil + continuation.resume(returning: result) + } + + private static func credentials(from authState: OIDAuthState) throws -> CodexOAuthCredentials { + guard let tokenResponse = authState.lastTokenResponse, + let accessToken = tokenResponse.accessToken, + !accessToken.isEmpty + else { + throw CodexOAuthCredentialsError.missingTokens } - group.addTask { - let nanos = UInt64(max(0, timeout) * 1_000_000_000) - try? await Task.sleep(nanoseconds: nanos) - return nil + + let refreshToken = tokenResponse.refreshToken ?? authState.refreshToken ?? "" + if refreshToken.isEmpty { + throw CodexOAuthCredentialsError.missingTokens } - let result = await group.next() - group.cancelAll() - if let result, let text = result { return text } - return "" + + let idToken = tokenResponse.idToken + let accountId = CodexOAuthClaimResolver.accountID(accessToken: accessToken, idToken: idToken) + + return CodexOAuthCredentials( + accessToken: accessToken, + refreshToken: refreshToken, + idToken: idToken, + accountId: accountId, + lastRefresh: Date()) } } - private static func decode(_ data: Data) -> String { - guard let text = String(data: data, encoding: .utf8) else { return "" } - return text + private static func makeAuthorizationRequest(listenerBaseURL: URL) -> OIDAuthorizationRequest { + let configuration = OIDServiceConfiguration( + authorizationEndpoint: Self.authorizationEndpoint, + tokenEndpoint: Self.tokenEndpoint) + return self.makeAuthorizationRequest(configuration: configuration, listenerBaseURL: listenerBaseURL) + } + + private static func makeAuthorizationRequest( + configuration: OIDServiceConfiguration, + listenerBaseURL: URL) -> OIDAuthorizationRequest + { + _ = listenerBaseURL + let state = OIDAuthorizationRequest.generateState() + let codeVerifier = OIDAuthorizationRequest.generateCodeVerifier() + let codeChallenge = OIDAuthorizationRequest.codeChallengeS256(forVerifier: codeVerifier) + return OIDAuthorizationRequest( + configuration: configuration, + clientId: Self.clientID, + clientSecret: nil, + scope: Self.scopes.joined(separator: " "), + redirectURL: Self.redirectURL, + responseType: OIDResponseTypeCode, + state: state, + nonce: nil, + codeVerifier: codeVerifier, + codeChallenge: codeChallenge, + codeChallengeMethod: OIDOAuthorizationRequestCodeChallengeMethodS256, + additionalParameters: Self.additionalParameters) } } diff --git a/Sources/CodexBar/CodexOAuthAccountStore.swift b/Sources/CodexBar/CodexOAuthAccountStore.swift new file mode 100644 index 000000000..e29d9412e --- /dev/null +++ b/Sources/CodexBar/CodexOAuthAccountStore.swift @@ -0,0 +1,38 @@ +import CodexBarCore +import Foundation + +enum CodexOAuthAccountStore { + static func createProfileDirectory( + label: String, + rootDirectory: URL? = nil, + fileManager: FileManager = .default) throws -> URL + { + let root = rootDirectory ?? CodexBarConfigStore.defaultURL() + .deletingLastPathComponent() + .appendingPathComponent("codex-oauth-accounts", isDirectory: true) + try fileManager.createDirectory(at: root, withIntermediateDirectories: true) + + let slug = self.slug(for: label) + let directory = root.appendingPathComponent("\(slug)-\(UUID().uuidString)", isDirectory: true) + try fileManager.createDirectory(at: directory, withIntermediateDirectories: true) + return directory + } + + static func removeProfileDirectoryIfPresent(_ directory: URL, fileManager: FileManager = .default) { + try? fileManager.removeItem(at: directory) + } + + private static func slug(for label: String) -> String { + let trimmed = label.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + let slug = trimmed.unicodeScalars.map { scalar -> Character in + if CharacterSet.alphanumerics.contains(scalar) { + return Character(scalar) + } + return "-" + } + let collapsed = String(slug) + .replacingOccurrences(of: "--+", with: "-", options: .regularExpression) + .trimmingCharacters(in: CharacterSet(charactersIn: "-")) + return collapsed.isEmpty ? "codex-account" : collapsed + } +} diff --git a/Sources/CodexBar/MenuHighlightStyle.swift b/Sources/CodexBar/MenuHighlightStyle.swift index be76fe04a..ec06dbc23 100644 --- a/Sources/CodexBar/MenuHighlightStyle.swift +++ b/Sources/CodexBar/MenuHighlightStyle.swift @@ -1,7 +1,15 @@ import SwiftUI +private struct MenuItemHighlightedKey: EnvironmentKey { + static let defaultValue = false +} + +// swiftformat:disable:next environmentEntry extension EnvironmentValues { - @Entry var menuItemHighlighted: Bool = false + var menuItemHighlighted: Bool { + get { self[MenuItemHighlightedKey.self] } + set { self[MenuItemHighlightedKey.self] = newValue } + } } enum MenuHighlightStyle { diff --git a/Sources/CodexBar/PreferencesProviderSettingsRows.swift b/Sources/CodexBar/PreferencesProviderSettingsRows.swift index 414f41c55..e215257e6 100644 --- a/Sources/CodexBar/PreferencesProviderSettingsRows.swift +++ b/Sources/CodexBar/PreferencesProviderSettingsRows.swift @@ -205,6 +205,7 @@ struct ProviderSettingsTokenAccountsRowView: View { let descriptor: ProviderSettingsTokenAccountsDescriptor @State private var newLabel: String = "" @State private var newToken: String = "" + @State private var isBrowserLoginInFlight = false var body: some View { VStack(alignment: .leading, spacing: 8) { @@ -264,6 +265,23 @@ struct ProviderSettingsTokenAccountsRowView: View { .controlSize(.small) .disabled(self.newLabel.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty || self.newToken.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty) + + if let addAccountViaLogin = self.descriptor.addAccountViaLogin { + Button(self.isBrowserLoginInFlight ? "Opening..." : "Add with browser") { + let label = self.newLabel.trimmingCharacters(in: .whitespacesAndNewlines) + guard !label.isEmpty else { return } + self.isBrowserLoginInFlight = true + Task { @MainActor in + await addAccountViaLogin(label) + self.isBrowserLoginInFlight = false + self.newLabel = "" + } + } + .buttonStyle(.bordered) + .controlSize(.small) + .disabled(self.isBrowserLoginInFlight || + self.newLabel.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty) + } } HStack(spacing: 10) { diff --git a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift index e2dce0a7f..4b9a03f84 100644 --- a/Sources/CodexBar/PreferencesProvidersPane+Testing.swift +++ b/Sources/CodexBar/PreferencesProvidersPane+Testing.swift @@ -172,6 +172,7 @@ enum ProvidersPaneTestHarness { activeIndex: { 0 }, setActiveIndex: { _ in }, addAccount: { _, _ in }, + addAccountViaLogin: nil, removeAccount: { _ in }, openConfigFile: {}, reloadFromDisk: {}) diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index 877c78da9..79c362ea7 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -201,6 +201,9 @@ struct ProvidersPane: View { } } }, + addAccountViaLogin: provider == .codex ? { label in + await self.addCodexOAuthAccount(label: label) + } : nil, removeAccount: { accountID in self.settings.removeTokenAccount(provider: provider, accountID: accountID) Task { @MainActor in @@ -262,6 +265,57 @@ struct ProvidersPane: View { }) } + @MainActor + private func addCodexOAuthAccount(label: String) async { + let trimmedLabel = label.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmedLabel.isEmpty else { return } + + let profileDirectory: URL + do { + profileDirectory = try CodexOAuthAccountStore.createProfileDirectory(label: trimmedLabel) + } catch { + self.presentCodexOAuthAccountAlert(title: "Could not prepare account", message: error.localizedDescription) + return + } + + let result = await CodexLoginRunner.run( + timeout: 120, + onPhaseChange: nil, + credentialSource: profileDirectory.path) + + switch result.outcome { + case .success: + self.settings.addTokenAccount(provider: .codex, label: trimmedLabel, token: profileDirectory.path) + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await self.store.refreshProvider(.codex, allowDisabled: true) + } + case .timedOut: + CodexOAuthAccountStore.removeProfileDirectoryIfPresent(profileDirectory) + self.presentCodexOAuthAccountAlert(title: "Codex login timed out", message: result.output) + case .failed: + CodexOAuthAccountStore.removeProfileDirectoryIfPresent(profileDirectory) + self.presentCodexOAuthAccountAlert(title: "Codex login failed", message: result.output) + case .launchFailed: + CodexOAuthAccountStore.removeProfileDirectoryIfPresent(profileDirectory) + self.presentCodexOAuthAccountAlert(title: "Could not start Codex login", message: result.output) + case .missingBinary: + CodexOAuthAccountStore.removeProfileDirectoryIfPresent(profileDirectory) + self.presentCodexOAuthAccountAlert( + title: "Codex login unavailable", + message: "The native OAuth flow could not start.\n\n\(result.output)") + } + } + + @MainActor + private func presentCodexOAuthAccountAlert(title: String, message: String) { + let alert = NSAlert() + alert.alertStyle = .warning + alert.messageText = title + alert.informativeText = message + alert.addButton(withTitle: "OK") + alert.runModal() + } + func menuBarMetricPicker(for provider: UsageProvider) -> ProviderSettingsPickerDescriptor? { if provider == .zai { return nil } let options: [ProviderSettingsPickerOption] diff --git a/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift b/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift index c8847374c..39a9027d5 100644 --- a/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift +++ b/Sources/CodexBar/Providers/Codex/CodexLoginFlow.swift @@ -3,7 +3,15 @@ import CodexBarCore @MainActor extension StatusItemController { func runCodexLoginFlow() async { - let result = await CodexLoginRunner.run(timeout: 120) + let phaseHandler: @Sendable (CodexLoginRunner.Phase) -> Void = { [weak self] phase in + Task { @MainActor in + switch phase { + case .requesting: self?.loginPhase = .requesting + case .waitingBrowser: self?.loginPhase = .waitingBrowser + } + } + } + let result = await CodexLoginRunner.run(timeout: 120, onPhaseChange: phaseHandler) guard !Task.isCancelled else { return } self.loginPhase = .idle self.presentCodexLoginResult(result) diff --git a/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift b/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift index 35baa270d..243c6bf24 100644 --- a/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift +++ b/Sources/CodexBar/Providers/Codex/CodexProviderImplementation.swift @@ -27,6 +27,20 @@ struct CodexProviderImplementation: ProviderImplementation { .codex(context.settings.codexSettingsSnapshot(tokenOverride: context.tokenOverride)) } + @MainActor + func tokenAccountsVisibility(context: ProviderSettingsContext, support: TokenAccountSupport) -> Bool { + guard support.requiresManualCookieSource else { return true } + if !context.settings.tokenAccounts(for: context.provider).isEmpty { return true } + return context.settings.codexCookieSource == .manual + } + + @MainActor + func applyTokenAccountCookieSource(settings: SettingsStore) { + if settings.codexCookieSource != .manual { + settings.codexCookieSource = .manual + } + } + @MainActor func defaultSourceLabel(context: ProviderSourceLabelContext) -> String? { context.settings.codexUsageDataSource.rawValue @@ -47,9 +61,12 @@ struct CodexProviderImplementation: ProviderImplementation { @MainActor func sourceMode(context: ProviderSourceModeContext) -> ProviderSourceMode { switch context.settings.codexUsageDataSource { - case .auto: .auto - case .oauth: .oauth - case .cli: .cli + case .auto: + ProviderSourceMode.auto + case .oauth: + ProviderSourceMode.oauth + case .cli: + ProviderSourceMode.cli } } @@ -106,12 +123,17 @@ struct CodexProviderImplementation: ProviderImplementation { keychainDisabled: context.settings.debugDisableKeychainAccess) let cookieSubtitle: () -> String? = { - ProviderCookieSourceUI.subtitle( + let base = ProviderCookieSourceUI.subtitle( source: context.settings.codexCookieSource, keychainDisabled: context.settings.debugDisableKeychainAccess, auto: "Automatic imports browser cookies for dashboard extras.", manual: "Paste a Cookie header from a chatgpt.com request.", off: "Disable OpenAI dashboard cookie usage.") + let status = context.store.openAIDashboardCookieImportStatus? + .trimmingCharacters(in: .whitespacesAndNewlines) + guard let status, !status.isEmpty else { return base } + guard !base.isEmpty else { return status } + return "\(base)\n\(status)" } return [ @@ -155,7 +177,35 @@ struct CodexProviderImplementation: ProviderImplementation { kind: .secure, placeholder: "Cookie: …", binding: context.stringBinding(\.codexCookieHeader), - actions: [], + actions: [ + ProviderSettingsActionDescriptor( + id: "codex-test-cookie", + title: "Test cookie", + style: .bordered, + isVisible: nil, + perform: { + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await context.store.testOpenAIDashboardCookieNow() + } + }), + ProviderSettingsActionDescriptor( + id: "codex-save-cookie-account", + title: "Save as account", + style: .link, + isVisible: nil, + perform: { + let token = context.settings.codexCookieHeader + .trimmingCharacters(in: .whitespacesAndNewlines) + guard !token.isEmpty else { return } + context.settings.addTokenAccount( + provider: .codex, + label: "", + token: token) + await ProviderInteractionContext.$current.withValue(.userInitiated) { + await context.store.refreshProvider(.codex, allowDisabled: true) + } + }), + ], isVisible: { context.settings.codexCookieSource == .manual }, diff --git a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift index 335dbf411..acf06ae74 100644 --- a/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift +++ b/Sources/CodexBar/Providers/Codex/CodexSettingsStore.swift @@ -49,10 +49,13 @@ extension SettingsStore { extension SettingsStore { func codexSettingsSnapshot(tokenOverride: TokenAccountOverride?) -> ProviderSettingsSnapshot.CodexProviderSettings { - ProviderSettingsSnapshot.CodexProviderSettings( + let oauthCredentialSource = self.codexSnapshotOAuthCredentialSource(tokenOverride: tokenOverride) + return ProviderSettingsSnapshot.CodexProviderSettings( usageDataSource: self.codexUsageDataSource, cookieSource: self.codexSnapshotCookieSource(tokenOverride: tokenOverride), - manualCookieHeader: self.codexSnapshotCookieHeader(tokenOverride: tokenOverride)) + manualCookieHeader: oauthCredentialSource == nil ? self + .codexSnapshotCookieHeader(tokenOverride: tokenOverride) : nil, + oauthCredentialSource: oauthCredentialSource) } private static func codexUsageDataSource(from source: ProviderSourceMode?) -> CodexUsageDataSource { @@ -94,4 +97,18 @@ extension SettingsStore { if self.tokenAccounts(for: .codex).isEmpty { return fallback } return .manual } + + private func codexSnapshotOAuthCredentialSource(tokenOverride: TokenAccountOverride?) -> String? { + guard let account = ProviderTokenAccountSelection.selectedAccount( + provider: .codex, + settings: self, + override: tokenOverride) + else { + return nil + } + + let raw = account.token.trimmingCharacters(in: .whitespacesAndNewlines) + guard !raw.isEmpty else { return nil } + return (try? CodexOAuthCredentialsStore.load(rawSource: raw)) != nil ? raw : nil + } } diff --git a/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift b/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift index d5a85b8f7..419ff168c 100644 --- a/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift +++ b/Sources/CodexBar/Providers/Shared/ProviderSettingsDescriptors.swift @@ -97,6 +97,7 @@ struct ProviderSettingsTokenAccountsDescriptor: Identifiable { let activeIndex: () -> Int let setActiveIndex: (Int) -> Void let addAccount: (_ label: String, _ token: String) -> Void + let addAccountViaLogin: ((_ label: String) async -> Void)? let removeAccount: (_ accountID: UUID) -> Void let openConfigFile: () -> Void let reloadFromDisk: () -> Void diff --git a/Sources/CodexBar/SettingsStore+TokenAccounts.swift b/Sources/CodexBar/SettingsStore+TokenAccounts.swift index 1f8a0277b..3361aac6a 100644 --- a/Sources/CodexBar/SettingsStore+TokenAccounts.swift +++ b/Sources/CodexBar/SettingsStore+TokenAccounts.swift @@ -66,6 +66,33 @@ extension SettingsStore { ]) } + func updateTokenAccount(provider: UsageProvider, accountID: UUID, label: String? = nil, token: String? = nil) { + guard let data = self.tokenAccountsData(for: provider), !data.accounts.isEmpty else { return } + let updatedAccounts = data.accounts.map { account in + guard account.id == accountID else { return account } + let updatedLabel = label?.trimmingCharacters(in: .whitespacesAndNewlines) + let updatedToken = token?.trimmingCharacters(in: .whitespacesAndNewlines) + return ProviderTokenAccount( + id: account.id, + label: (updatedLabel?.isEmpty == false ? updatedLabel! : account.label), + token: (updatedToken?.isEmpty == false ? updatedToken! : account.token), + addedAt: account.addedAt, + lastUsed: account.lastUsed) + } + self.updateProviderConfig(provider: provider) { entry in + entry.tokenAccounts = ProviderTokenAccountData( + version: data.version, + accounts: updatedAccounts, + activeIndex: data.activeIndex) + } + CodexBarLog.logger(LogCategories.tokenAccounts).info( + "Token account updated", + metadata: [ + "provider": provider.rawValue, + "accountID": accountID.uuidString, + ]) + } + func removeTokenAccount(provider: UsageProvider, accountID: UUID) { guard let data = self.tokenAccountsData(for: provider), !data.accounts.isEmpty else { return } let filtered = data.accounts.filter { $0.id != accountID } diff --git a/Sources/CodexBar/SettingsStore.swift b/Sources/CodexBar/SettingsStore.swift index a45a78cf6..c102e78d6 100644 --- a/Sources/CodexBar/SettingsStore.swift +++ b/Sources/CodexBar/SettingsStore.swift @@ -148,9 +148,10 @@ final class SettingsStore { LaunchAtLoginManager.setEnabled(self.launchAtLogin) self.runInitialProviderDetectionIfNeeded() self.applyTokenCostDefaultIfNeeded() - if self.claudeUsageDataSource != .cli { self.claudeWebExtrasEnabled = false } - self.openAIWebAccessEnabled = self.codexCookieSource.isEnabled - Self.sharedDefaults?.set(self.debugDisableKeychainAccess, forKey: "debugDisableKeychainAccess") + if self.claudeUsageDataSource != .cli { + self.defaultsState.claudeWebExtrasEnabledRaw = false + } + self.defaultsState.openAIWebAccessEnabled = self.codexCookieSource.isEnabled KeychainAccessGate.isDisabled = self.debugDisableKeychainAccess } } @@ -166,24 +167,17 @@ extension SettingsStore { return stored } if let shared = Self.sharedDefaults?.object(forKey: "debugDisableKeychainAccess") as? Bool { - userDefaults.set(shared, forKey: "debugDisableKeychainAccess") return shared } return false }() let debugFileLoggingEnabled = userDefaults.object(forKey: "debugFileLoggingEnabled") as? Bool ?? false let debugLogLevelRaw = userDefaults.string(forKey: "debugLogLevel") ?? CodexBarLog.Level.verbose.rawValue - if userDefaults.string(forKey: "debugLogLevel") == nil { - userDefaults.set(debugLogLevelRaw, forKey: "debugLogLevel") - } let debugLoadingPatternRaw = userDefaults.string(forKey: "debugLoadingPattern") let debugKeepCLISessionsAlive = userDefaults.object(forKey: "debugKeepCLISessionsAlive") as? Bool ?? false let statusChecksEnabled = userDefaults.object(forKey: "statusChecksEnabled") as? Bool ?? true let sessionQuotaDefault = userDefaults.object(forKey: "sessionQuotaNotificationsEnabled") as? Bool let sessionQuotaNotificationsEnabled = sessionQuotaDefault ?? true - if sessionQuotaDefault == nil { - userDefaults.set(true, forKey: "sessionQuotaNotificationsEnabled") - } let usageBarsShowUsed = userDefaults.object(forKey: "usageBarsShowUsed") as? Bool ?? false let resetTimesShowAbsolute = userDefaults.object(forKey: "resetTimesShowAbsolute") as? Bool ?? false let menuBarShowsBrandIconWithPercent = userDefaults.object( @@ -209,10 +203,8 @@ extension SettingsStore { let claudeWebExtrasEnabledRaw = userDefaults.object(forKey: "claudeWebExtrasEnabled") as? Bool ?? false let creditsExtrasDefault = userDefaults.object(forKey: "showOptionalCreditsAndExtraUsage") as? Bool let showOptionalCreditsAndExtraUsage = creditsExtrasDefault ?? true - if creditsExtrasDefault == nil { userDefaults.set(true, forKey: "showOptionalCreditsAndExtraUsage") } let openAIWebAccessDefault = userDefaults.object(forKey: "openAIWebAccessEnabled") as? Bool let openAIWebAccessEnabled = openAIWebAccessDefault ?? true - if openAIWebAccessDefault == nil { userDefaults.set(true, forKey: "openAIWebAccessEnabled") } let jetbrainsIDEBasePath = userDefaults.string(forKey: "jetbrainsIDEBasePath") ?? "" let mergeIcons = userDefaults.object(forKey: "mergeIcons") as? Bool ?? true let switcherShowsIcons = userDefaults.object(forKey: "switcherShowsIcons") as? Bool ?? true diff --git a/Sources/CodexBar/StatusItemController+Actions.swift b/Sources/CodexBar/StatusItemController+Actions.swift index efe63ffd8..4f9b831e0 100644 --- a/Sources/CodexBar/StatusItemController+Actions.swift +++ b/Sources/CodexBar/StatusItemController+Actions.swift @@ -208,16 +208,16 @@ extension StatusItemController { return case .missingBinary: self.presentLoginAlert( - title: "Codex CLI not found", - message: "Install the Codex CLI (npm i -g @openai/codex) and try again.") + title: "Codex login helper unavailable", + message: "Install the Codex CLI (npm i -g @openai/codex) or update CodexBar and try again.") case let .launchFailed(message): - self.presentLoginAlert(title: "Could not start codex login", message: message) + self.presentLoginAlert(title: "Could not start Codex login", message: message) case .timedOut: self.presentLoginAlert( title: "Codex login timed out", message: self.trimmedLoginOutput(result.output)) case let .failed(status): - let statusLine = "codex login exited with status \(status)." + let statusLine = "Codex login failed with status \(status)." let message = self.trimmedLoginOutput(result.output.isEmpty ? statusLine : result.output) self.presentLoginAlert(title: "Codex login failed", message: message) } diff --git a/Sources/CodexBar/StatusItemController+Menu.swift b/Sources/CodexBar/StatusItemController+Menu.swift index 1ca78a835..3b47925ee 100644 --- a/Sources/CodexBar/StatusItemController+Menu.swift +++ b/Sources/CodexBar/StatusItemController+Menu.swift @@ -163,6 +163,12 @@ extension StatusItemController { let menuWidth = self.menuCardWidth(for: enabledProviders, menu: menu) let currentProvider = selectedProvider ?? enabledProviders.first ?? .codex let tokenAccountDisplay = isOverviewSelected ? nil : self.tokenAccountMenuDisplay(for: currentProvider) + if let tokenAccountDisplay, tokenAccountDisplay.showSwitcher { + self.primeTokenAccountSnapshotsIfNeeded( + provider: currentProvider, + accounts: tokenAccountDisplay.accounts, + menu: menu) + } let showAllTokenAccounts = tokenAccountDisplay?.showAll ?? false let openAIContext = self.openAIWebContext( currentProvider: currentProvider, @@ -601,11 +607,16 @@ extension StatusItemController { width: self.menuCardWidth(for: self.store.enabledProviders(), menu: menu), onSelect: { [weak self, weak menu] index in guard let self, let menu else { return } + let selectedAccount = display.accounts[index] self.settings.setActiveTokenAccountIndex(index, for: display.provider) + self.store.applyCachedTokenAccountState(provider: display.provider, accountID: selectedAccount.id) Task { @MainActor in await ProviderInteractionContext.$current.withValue(.userInitiated) { - await self.store.refresh() + await self.store.refreshTokenAccounts(provider: display.provider, accounts: display.accounts) } + self.populateMenu(menu, provider: display.provider) + self.markMenuFresh(menu) + self.applyIcon(phase: nil) } self.populateMenu(menu, provider: display.provider) self.markMenuFresh(menu) @@ -1363,7 +1374,7 @@ extension StatusItemController { let target = provider ?? self.store.enabledProviders().first ?? .codex let metadata = self.store.metadata(for: target) - let snapshot = snapshotOverride ?? self.store.snapshot(for: target) + let snapshot = snapshotOverride ?? self.preferredMenuSnapshot(for: target) let credits: CreditsSnapshot? let creditsError: String? let dashboard: OpenAIDashboardSnapshot? @@ -1415,6 +1426,55 @@ extension StatusItemController { return UsageMenuCardView.Model.make(input) } + func preferredMenuSnapshot(for provider: UsageProvider) -> UsageSnapshot? { + let liveSnapshot = self.store.snapshot(for: provider) + guard let selectedSnapshot = self.selectedTokenAccountCachedSnapshot(for: provider) else { + return liveSnapshot + } + guard let liveSnapshot else { return selectedSnapshot } + guard let selectedAccount = self.settings.selectedTokenAccount(for: provider) else { + return liveSnapshot + } + let liveEmail = liveSnapshot.accountEmail(for: provider)? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + let selectedEmail = selectedSnapshot.accountEmail(for: provider)? + .trimmingCharacters(in: .whitespacesAndNewlines) + .lowercased() + let selectedLabel = selectedAccount.label.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + if let liveEmail, let selectedEmail, liveEmail == selectedEmail { + return liveSnapshot + } + if let liveEmail, !selectedLabel.isEmpty, liveEmail == selectedLabel { + return liveSnapshot + } + return selectedSnapshot + } + + private func selectedTokenAccountCachedSnapshot(for provider: UsageProvider) -> UsageSnapshot? { + guard let selected = self.settings.selectedTokenAccount(for: provider) else { return nil } + return self.store.accountSnapshots[provider]? + .first(where: { $0.account.id == selected.id })? + .snapshot + } + + private func primeTokenAccountSnapshotsIfNeeded( + provider: UsageProvider, + accounts: [ProviderTokenAccount], + menu: NSMenu) + { + let cachedCount = self.store.accountSnapshots[provider]?.count ?? 0 + guard cachedCount < accounts.count else { return } + Task { @MainActor [weak self, weak menu] in + guard let self, let menu else { return } + await self.store.refreshTokenAccounts(provider: provider, accounts: accounts) + guard self.openMenus[ObjectIdentifier(menu)] != nil else { return } + self.populateMenu(menu, provider: provider) + self.markMenuFresh(menu) + self.applyIcon(phase: nil) + } + } + @objc private func menuCardNoOp(_ sender: NSMenuItem) { _ = sender } diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index bc25a5c86..0d0f91a09 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -34,7 +34,7 @@ extension UsageStore { defer { self.refreshingProviders.remove(provider) } let tokenAccounts = self.tokenAccounts(for: provider) - if self.shouldFetchAllTokenAccounts(provider: provider, accounts: tokenAccounts) { + if !tokenAccounts.isEmpty { await self.refreshTokenAccounts(provider: provider, accounts: tokenAccounts) return } else { @@ -78,12 +78,18 @@ extension UsageStore { switch outcome.result { case let .success(result): let scoped = result.usage.scoped(to: provider) + let credits = result.credits ?? result.dashboard?.toCreditsSnapshot() await MainActor.run { self.handleSessionQuotaTransition(provider: provider, snapshot: scoped) self.snapshots[provider] = scoped self.lastSourceLabels[provider] = result.sourceLabel self.errors[provider] = nil self.failureGates[provider]?.recordSuccess() + self.applyProviderSupplementaryData( + provider: provider, + credits: credits, + dashboard: result.dashboard, + preserveExisting: true) } if let runtime = self.providerRuntimes[provider] { let context = ProviderRuntimeContext( diff --git a/Sources/CodexBar/UsageStore+TokenAccounts.swift b/Sources/CodexBar/UsageStore+TokenAccounts.swift index 3e55ffa9f..7c84fea59 100644 --- a/Sources/CodexBar/UsageStore+TokenAccounts.swift +++ b/Sources/CodexBar/UsageStore+TokenAccounts.swift @@ -7,13 +7,24 @@ struct TokenAccountUsageSnapshot: Identifiable, Sendable { let snapshot: UsageSnapshot? let error: String? let sourceLabel: String? + let credits: CreditsSnapshot? + let dashboard: OpenAIDashboardSnapshot? - init(account: ProviderTokenAccount, snapshot: UsageSnapshot?, error: String?, sourceLabel: String?) { + init( + account: ProviderTokenAccount, + snapshot: UsageSnapshot?, + error: String?, + sourceLabel: String?, + credits: CreditsSnapshot? = nil, + dashboard: OpenAIDashboardSnapshot? = nil) + { self.id = account.id self.account = account self.snapshot = snapshot self.error = error self.sourceLabel = sourceLabel + self.credits = credits + self.dashboard = dashboard } } @@ -32,31 +43,43 @@ extension UsageStore { let selectedAccount = self.settings.selectedTokenAccount(for: provider) let limitedAccounts = self.limitedTokenAccounts(accounts, selected: selectedAccount) let effectiveSelected = selectedAccount ?? limitedAccounts.first - var snapshots: [TokenAccountUsageSnapshot] = [] - var selectedOutcome: ProviderFetchOutcome? - var selectedSnapshot: UsageSnapshot? + let prioritizedAccounts = self.prioritizedTokenAccounts(limitedAccounts, selected: effectiveSelected) + var snapshotsByAccountID: [UUID: TokenAccountUsageSnapshot] = [:] - for account in limitedAccounts { + for account in prioritizedAccounts { let override = TokenAccountOverride(provider: provider, account: account) - let outcome = await self.fetchOutcome(provider: provider, override: override) + let initialOutcome = await self.fetchOutcome(provider: provider, override: override) + let outcome = await self.repairTokenAccountOutcomeIfNeeded( + initialOutcome, + provider: provider, + account: account) let resolved = self.resolveAccountOutcome(outcome, provider: provider, account: account) - snapshots.append(resolved.snapshot) + snapshotsByAccountID[account.id] = resolved.snapshot if account.id == effectiveSelected?.id { - selectedOutcome = outcome - selectedSnapshot = resolved.usage + await self.applySelectedOutcome( + outcome, + provider: provider, + account: effectiveSelected, + fallbackSnapshot: resolved.usage) + } else { + await MainActor.run { + self.upsertTokenAccountSnapshot( + provider: provider, + account: account, + snapshot: resolved.snapshot.snapshot, + error: resolved.snapshot.error, + sourceLabel: resolved.snapshot.sourceLabel, + credits: resolved.snapshot.credits, + dashboard: resolved.snapshot.dashboard) + } } } await MainActor.run { - self.accountSnapshots[provider] = snapshots - } - - if let selectedOutcome { - await self.applySelectedOutcome( - selectedOutcome, - provider: provider, - account: effectiveSelected, - fallbackSnapshot: selectedSnapshot) + let orderedSnapshots = limitedAccounts.compactMap { snapshotsByAccountID[$0.id] } + if !orderedSnapshots.isEmpty { + self.accountSnapshots[provider] = orderedSnapshots + } } } @@ -74,12 +97,26 @@ extension UsageStore { return limited } + func prioritizedTokenAccounts( + _ accounts: [ProviderTokenAccount], + selected: ProviderTokenAccount?) -> [ProviderTokenAccount] + { + guard let selected else { return accounts } + guard let selectedIndex = accounts.firstIndex(where: { $0.id == selected.id }) else { return accounts } + if selectedIndex == 0 { return accounts } + + var prioritized = accounts + prioritized.remove(at: selectedIndex) + prioritized.insert(selected, at: 0) + return prioritized + } + func fetchOutcome( provider: UsageProvider, override: TokenAccountOverride?) async -> ProviderFetchOutcome { let descriptor = ProviderDescriptorRegistry.descriptor(for: provider) - let sourceMode = self.sourceMode(for: provider) + let sourceMode = self.sourceMode(for: provider, override: override) let snapshot = ProviderRegistry.makeSettingsSnapshot(settings: self.settings, tokenOverride: override) let env = ProviderRegistry.makeEnvironment( base: ProcessInfo.processInfo.environment, @@ -102,15 +139,106 @@ extension UsageStore { return await descriptor.fetchOutcome(context: context) } - func sourceMode(for provider: UsageProvider) -> ProviderSourceMode { - ProviderCatalog.implementation(for: provider)? + func sourceMode(for provider: UsageProvider, override: TokenAccountOverride? = nil) -> ProviderSourceMode { + if provider == .codex, + let account = ProviderTokenAccountSelection.selectedAccount( + provider: provider, + settings: self.settings, + override: override), + (try? CodexOAuthCredentialsStore.load(rawSource: account.token)) != nil + { + return .oauth + } + + if let support = TokenAccountSupportCatalog.support(for: provider), + support.requiresManualCookieSource, + override != nil || self.settings.selectedTokenAccount(for: provider) != nil + { + return .web + } + + return ProviderCatalog.implementation(for: provider)? .sourceMode(context: ProviderSourceModeContext(provider: provider, settings: self.settings)) ?? .auto } + private func repairTokenAccountOutcomeIfNeeded( + _ outcome: ProviderFetchOutcome, + provider: UsageProvider, + account: ProviderTokenAccount) async -> ProviderFetchOutcome + { + guard provider == .codex, + self.settings.codexCookieSource == .manual, + ProviderInteractionContext.current == .userInitiated, + case .failure = outcome.result + else { + return outcome + } + + if (try? CodexOAuthCredentialsStore.load(rawSource: account.token)) != nil { + return outcome + } + + guard let targetEmail = self.tokenAccountEmailHint(provider: provider, account: account), + !targetEmail.isEmpty + else { + return outcome + } + + let importer = OpenAIDashboardBrowserCookieImporter(browserDetection: self.browserDetection) + + do { + let result = try await importer.importBestCookies( + intoAccountEmail: targetEmail, + allowAnyAccount: false) + guard let cookieHeader = result.cookieHeader, + !cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + else { + return outcome + } + + await MainActor.run { + self.settings.updateTokenAccount(provider: provider, accountID: account.id, token: cookieHeader) + } + + let updatedAccount = await MainActor.run { + self.settings.tokenAccounts(for: provider).first(where: { $0.id == account.id }) + } ?? account + let override = TokenAccountOverride(provider: provider, account: updatedAccount) + return await self.fetchOutcome(provider: provider, override: override) + } catch { + return outcome + } + } + + private func tokenAccountEmailHint(provider: UsageProvider, account: ProviderTokenAccount) -> String? { + if let cached = self.tokenAccountUsageSnapshot(provider: provider, accountID: account.id) { + if let signedInEmail = cached.dashboard?.signedInEmail?.trimmingCharacters(in: .whitespacesAndNewlines), + !signedInEmail.isEmpty + { + return signedInEmail + } + if let email = cached.snapshot?.identity(for: provider)?.accountEmail? + .trimmingCharacters(in: .whitespacesAndNewlines), + !email.isEmpty + { + return email + } + } + + let label = account.label.trimmingCharacters(in: .whitespacesAndNewlines) + if label.contains("@") { + return label + } + + return nil + } + private struct ResolvedAccountOutcome { let snapshot: TokenAccountUsageSnapshot let usage: UsageSnapshot? + let credits: CreditsSnapshot? + let dashboard: OpenAIDashboardSnapshot? } private func resolveAccountOutcome( @@ -122,19 +250,26 @@ extension UsageStore { case let .success(result): let scoped = result.usage.scoped(to: provider) let labeled = self.applyAccountLabel(scoped, provider: provider, account: account) + let credits = result.credits ?? result.dashboard?.toCreditsSnapshot() let snapshot = TokenAccountUsageSnapshot( account: account, snapshot: labeled, error: nil, - sourceLabel: result.sourceLabel) - return ResolvedAccountOutcome(snapshot: snapshot, usage: labeled) + sourceLabel: result.sourceLabel, + credits: credits, + dashboard: result.dashboard) + return ResolvedAccountOutcome( + snapshot: snapshot, + usage: labeled, + credits: credits, + dashboard: result.dashboard) case let .failure(error): let snapshot = TokenAccountUsageSnapshot( account: account, snapshot: nil, error: error.localizedDescription, sourceLabel: nil) - return ResolvedAccountOutcome(snapshot: snapshot, usage: nil) + return ResolvedAccountOutcome(snapshot: snapshot, usage: nil, credits: nil, dashboard: nil) } } @@ -150,26 +285,59 @@ extension UsageStore { switch outcome.result { case let .success(result): let scoped = result.usage.scoped(to: provider) + let credits = result.credits ?? result.dashboard?.toCreditsSnapshot() let labeled: UsageSnapshot = if let account { self.applyAccountLabel(scoped, provider: provider, account: account) } else { scoped } await MainActor.run { + if let account { + self.upsertTokenAccountSnapshot( + provider: provider, + account: account, + snapshot: labeled, + error: nil, + sourceLabel: result.sourceLabel, + credits: credits, + dashboard: result.dashboard) + } self.handleSessionQuotaTransition(provider: provider, snapshot: labeled) self.snapshots[provider] = labeled self.lastSourceLabels[provider] = result.sourceLabel self.errors[provider] = nil self.failureGates[provider]?.recordSuccess() + self.applyProviderSupplementaryData( + provider: provider, + credits: credits, + dashboard: result.dashboard, + preserveExisting: true) } case let .failure(error): await MainActor.run { + let selectedAccountID = self.settings.selectedTokenAccount(for: provider)?.id + let isSelectedTokenAccountFailure = account?.id == selectedAccountID + if let account { + self.upsertTokenAccountSnapshot( + provider: provider, + account: account, + snapshot: nil, + error: error.localizedDescription, + sourceLabel: nil, + credits: nil, + dashboard: nil) + } let hadPriorData = self.snapshots[provider] != nil || fallbackSnapshot != nil let shouldSurface = self.failureGates[provider]? .shouldSurfaceError(onFailureWithPriorData: hadPriorData) ?? true - if shouldSurface { + if shouldSurface || isSelectedTokenAccountFailure { self.errors[provider] = error.localizedDescription self.snapshots.removeValue(forKey: provider) + self.applyProviderSupplementaryData( + provider: provider, + credits: nil, + dashboard: nil, + preserveExisting: false) } else { self.errors[provider] = nil } @@ -194,4 +362,65 @@ extension UsageStore { loginMethod: existing?.loginMethod) return snapshot.withIdentity(identity) } + + func upsertTokenAccountSnapshot( + provider: UsageProvider, + account: ProviderTokenAccount, + snapshot: UsageSnapshot?, + error: String?, + sourceLabel: String?, + credits: CreditsSnapshot?, + dashboard: OpenAIDashboardSnapshot?) + { + let entry = TokenAccountUsageSnapshot( + account: account, + snapshot: snapshot, + error: error, + sourceLabel: sourceLabel, + credits: credits, + dashboard: dashboard) + var snapshots = self.accountSnapshots[provider] ?? [] + if let index = snapshots.firstIndex(where: { $0.account.id == account.id }) { + snapshots[index] = entry + } else { + snapshots.append(entry) + } + self.accountSnapshots[provider] = snapshots + } + + func tokenAccountUsageSnapshot(provider: UsageProvider, accountID: UUID) -> TokenAccountUsageSnapshot? { + self.accountSnapshots[provider]? + .first(where: { $0.account.id == accountID }) + } + + func applyCachedTokenAccountState(provider: UsageProvider, accountID: UUID) { + guard let cached = self.tokenAccountUsageSnapshot(provider: provider, accountID: accountID) else { + self.snapshots.removeValue(forKey: provider) + self.errors.removeValue(forKey: provider) + self.applyProviderSupplementaryData( + provider: provider, + credits: nil, + dashboard: nil, + preserveExisting: false) + return + } + + if let snapshot = cached.snapshot { + self.snapshots[provider] = snapshot + self.errors[provider] = nil + } else { + self.snapshots.removeValue(forKey: provider) + if let error = cached.error, !error.isEmpty { + self.errors[provider] = error + } else { + self.errors.removeValue(forKey: provider) + } + } + + self.applyProviderSupplementaryData( + provider: provider, + credits: cached.credits, + dashboard: cached.dashboard, + preserveExisting: false) + } } diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 4bed59c54..8daefc1d8 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -165,6 +165,7 @@ final class UsageStore { @ObservationIgnored private var creditsFailureStreak: Int = 0 @ObservationIgnored private var lastOpenAIDashboardSnapshot: OpenAIDashboardSnapshot? @ObservationIgnored private var lastOpenAIDashboardTargetEmail: String? + @ObservationIgnored private var lastOpenAIDashboardTargetTokenAccountID: UUID? @ObservationIgnored private var lastOpenAIDashboardCookieImportAttemptAt: Date? @ObservationIgnored private var lastOpenAIDashboardCookieImportEmail: String? @ObservationIgnored private var openAIWebAccountDidChange: Bool = false @@ -716,12 +717,13 @@ extension UsageStore { self.failureGates[.codex]?.recordSuccess() self.lastSourceLabels[.codex] = "openai-web" } - if self.credits == nil, let credits = dash.toCreditsSnapshot() { + if let credits = dash.toCreditsSnapshot() { self.credits = credits self.lastCreditsSnapshot = credits self.lastCreditsError = nil self.creditsFailureStreak = 0 } + self.updateSelectedCodexTokenAccountSupplementaryData(dashboard: dash) } if let email = targetEmail, !email.isEmpty { @@ -750,7 +752,9 @@ extension UsageStore { } let targetEmail = self.codexAccountEmailForOpenAIDashboard() - self.handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: targetEmail) + let tokenAccountID = self.settings.selectedTokenAccount(for: .codex)?.id + self.handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: targetEmail, tokenAccountID: tokenAccountID) + let preferredStoreEmail = self.codexOpenAIWebStoreEmail() let now = Date() let minInterval = self.openAIWebRefreshIntervalSeconds() @@ -778,7 +782,7 @@ extension UsageStore { let normalized = targetEmail? .trimmingCharacters(in: .whitespacesAndNewlines) .lowercased() - var effectiveEmail = targetEmail + var effectiveEmail = targetEmail ?? preferredStoreEmail // Use a per-email persistent `WKWebsiteDataStore` so multiple dashboard sessions can coexist. // Strategy: @@ -891,23 +895,38 @@ extension UsageStore { /// Detect Codex account email changes and clear stale OpenAI web state so the UI can't show the wrong user. /// This does not delete other per-email WebKit cookie stores (we keep multiple accounts around). - func handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: String?) { + func handleOpenAIWebTargetEmailChangeIfNeeded(targetEmail: String?, tokenAccountID: UUID? = nil) { let normalized = targetEmail? .trimmingCharacters(in: .whitespacesAndNewlines) .lowercased() - guard let normalized, !normalized.isEmpty else { return } - + let normalizedOrNil: String? = if let normalized, !normalized.isEmpty { + normalized + } else { + nil + } let previous = self.lastOpenAIDashboardTargetEmail - self.lastOpenAIDashboardTargetEmail = normalized + let previousTokenAccountID = self.lastOpenAIDashboardTargetTokenAccountID + self.lastOpenAIDashboardTargetEmail = normalizedOrNil + self.lastOpenAIDashboardTargetTokenAccountID = tokenAccountID - if let previous, - !previous.isEmpty, - previous != normalized - { + let emailChanged = if let previous, !previous.isEmpty, let normalizedOrNil { + previous != normalizedOrNil + } else { + false + } + let tokenAccountChanged = if let previousTokenAccountID { + previousTokenAccountID != tokenAccountID + } else { + false + } + + if emailChanged || tokenAccountChanged { + let fromText = previous ?? "unknown" + let toText = normalizedOrNil ?? "unknown" let stamp = Date().formatted(date: .abbreviated, time: .shortened) self.logOpenAIWeb( - "[\(stamp)] Codex account changed: \(previous) → \(normalized); " + + "[\(stamp)] Codex account changed: \(fromText) → \(toText); " + "clearing OpenAI web snapshot") self.openAIWebAccountDidChange = true self.openAIDashboard = nil @@ -927,6 +946,16 @@ extension UsageStore { await self.refreshOpenAIDashboardIfNeeded(force: true) } + func testOpenAIDashboardCookieNow() async { + await MainActor.run { + self.openAIDashboardCookieImportStatus = "Testing OpenAI cookie…" + } + self.resetOpenAIWebDebugLog(context: "manual cookie test") + let targetEmail = self.codexAccountEmailForOpenAIDashboard() + _ = await self.importOpenAIDashboardCookiesIfNeeded(targetEmail: targetEmail, force: true) + await self.refreshOpenAIDashboardIfNeeded(force: true) + } + private func importOpenAIDashboardCookiesIfNeeded(targetEmail: String?, force: Bool) async -> String? { let normalizedTarget = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines) let allowAnyAccount = normalizedTarget == nil || normalizedTarget?.isEmpty == true @@ -968,7 +997,8 @@ extension UsageStore { switch cookieSource { case .manual: self.settings.ensureCodexCookieLoaded() - let manualHeader = self.settings.codexCookieHeader + let manualHeader = self.settings.selectedTokenAccount(for: .codex)?.token + ?? self.settings.codexCookieHeader guard CookieHeaderNormalizer.normalize(manualHeader) != nil else { throw OpenAIDashboardBrowserCookieImporter.ImportError.manualCookieHeaderInvalid } @@ -987,7 +1017,8 @@ extension UsageStore { sourceLabel: "Off", cookieCount: 0, signedInEmail: normalizedTarget, - matchesCodexEmail: true) + matchesCodexEmail: true, + cookieHeader: nil) } let effectiveEmail = result.signedInEmail? .trimmingCharacters(in: .whitespacesAndNewlines) @@ -1039,15 +1070,27 @@ extension UsageStore { } self.logOpenAIWeb("[\(stamp)] import mismatch: \(foundText)") await MainActor.run { - self.openAIDashboardCookieImportStatus = allowAnyAccount - ? [ - "No signed-in OpenAI web session found.", - "Found \(foundText).", - ].joined(separator: " ") - : [ - "Browser cookies do not match Codex account (\(normalizedTarget ?? "unknown")).", - "Found \(foundText).", - ].joined(separator: " ") + if cookieSource == .manual { + self.openAIDashboardCookieImportStatus = allowAnyAccount + ? [ + "Manual cookie does not map to a signed-in OpenAI dashboard session.", + "Found \(foundText).", + ].joined(separator: " ") + : [ + "Manual cookie does not match Codex account (\(normalizedTarget ?? "unknown")).", + "Found \(foundText).", + ].joined(separator: " ") + } else { + self.openAIDashboardCookieImportStatus = allowAnyAccount + ? [ + "No signed-in OpenAI web session found.", + "Found \(foundText).", + ].joined(separator: " ") + : [ + "Browser cookies do not match Codex account (\(normalizedTarget ?? "unknown")).", + "Found \(foundText).", + ].joined(separator: " ") + } // Treat mismatch like "not logged in" for the current Codex account. self.openAIDashboardRequiresLogin = true self.openAIDashboard = nil @@ -1095,6 +1138,7 @@ extension UsageStore { self.lastOpenAIDashboardError = nil self.lastOpenAIDashboardSnapshot = nil self.lastOpenAIDashboardTargetEmail = nil + self.lastOpenAIDashboardTargetTokenAccountID = nil self.openAIDashboardRequiresLogin = false self.openAIDashboardCookieImportStatus = nil self.openAIDashboardCookieImportDebugLog = nil @@ -1102,13 +1146,90 @@ extension UsageStore { self.lastOpenAIDashboardCookieImportEmail = nil } + func applyProviderSupplementaryData( + provider: UsageProvider, + credits: CreditsSnapshot?, + dashboard: OpenAIDashboardSnapshot?, + preserveExisting: Bool) + { + guard provider == .codex else { return } + + if let dashboard { + self.openAIDashboard = dashboard + self.lastOpenAIDashboardSnapshot = dashboard + self.lastOpenAIDashboardError = nil + self.openAIDashboardRequiresLogin = false + } else if !preserveExisting { + self.openAIDashboard = nil + self.lastOpenAIDashboardSnapshot = nil + self.lastOpenAIDashboardError = nil + self.openAIDashboardRequiresLogin = false + } + + let resolvedCredits = credits ?? dashboard?.toCreditsSnapshot() + if let resolvedCredits { + self.credits = resolvedCredits + self.lastCreditsSnapshot = resolvedCredits + self.lastCreditsError = nil + self.creditsFailureStreak = 0 + } else if !preserveExisting { + self.credits = nil + self.lastCreditsSnapshot = nil + self.lastCreditsError = nil + self.creditsFailureStreak = 0 + } + } + + private func updateSelectedCodexTokenAccountSupplementaryData(dashboard: OpenAIDashboardSnapshot) { + guard let account = self.settings.selectedTokenAccount(for: .codex) else { return } + let existingSnapshot = self.tokenAccountUsageSnapshot(provider: .codex, accountID: account.id) + self.upsertTokenAccountSnapshot( + provider: .codex, + account: account, + snapshot: existingSnapshot?.snapshot, + error: existingSnapshot?.error, + sourceLabel: existingSnapshot?.sourceLabel, + credits: dashboard.toCreditsSnapshot(), + dashboard: dashboard) + } + private func dashboardEmailMismatch(expected: String?, actual: String?) -> Bool { guard let expected, !expected.isEmpty else { return false } guard let raw = actual?.trimmingCharacters(in: .whitespacesAndNewlines), !raw.isEmpty else { return false } return raw.lowercased() != expected.lowercased() } + private func codexOpenAIWebStoreEmail() -> String? { + guard self.settings.codexCookieSource == .manual, + let selectedAccount = self.settings.selectedTokenAccount(for: .codex) + else { + return nil + } + + if let cachedEmail = self.tokenAccountUsageSnapshot(provider: .codex, accountID: selectedAccount.id)? + .dashboard? + .signedInEmail? + .trimmingCharacters(in: .whitespacesAndNewlines), + !cachedEmail.isEmpty + { + return cachedEmail + } + + let imported = self.lastOpenAIDashboardCookieImportEmail? + .trimmingCharacters(in: .whitespacesAndNewlines) + if let imported, !imported.isEmpty { + return imported + } + + return nil + } + func codexAccountEmailForOpenAIDashboard() -> String? { + if self.settings.codexCookieSource == .manual, + !self.settings.tokenAccounts(for: .codex).isEmpty + { + return nil + } let direct = self.snapshots[.codex]?.accountEmail(for: .codex)? .trimmingCharacters(in: .whitespacesAndNewlines) if let direct, !direct.isEmpty { return direct } diff --git a/Sources/CodexBarCLI/CLIConfigCommand.swift b/Sources/CodexBarCLI/CLIConfigCommand.swift index edc013874..ced76871c 100644 --- a/Sources/CodexBarCLI/CLIConfigCommand.swift +++ b/Sources/CodexBarCLI/CLIConfigCommand.swift @@ -2,7 +2,78 @@ import CodexBarCore import Commander import Foundation +private struct ConfigAddAccountPayload: Encodable { + let provider: String + let label: String + let activeIndex: Int + let count: Int +} + extension CodexBarCLI { + static func runConfigAddAccount(_ values: ParsedValues) { + let output = CLIOutputPreferences.from(values: values) + let store = CodexBarConfigStore() + let config: CodexBarConfig + do { + config = try store.loadOrCreateDefault() + } catch { + Self.exit(code: .failure, message: error.localizedDescription, output: output, kind: .config) + } + + guard let provider = Self.decodeSingleProvider(from: values) else { + Self.exit( + code: .failure, + message: "Error: --provider must be a single supported provider.", + output: output, + kind: .args) + } + guard TokenAccountSupportCatalog.support(for: provider) != nil else { + Self.exit( + code: .failure, + message: "Error: \(provider.rawValue) does not support token accounts.", + output: output, + kind: .args) + } + + let rawToken = values.options["token"]?.last?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "" + guard !rawToken.isEmpty else { + Self.exit(code: .failure, message: "Error: --token is required.", output: output, kind: .args) + } + + let rawLabel = values.options["label"]?.last?.trimmingCharacters(in: .whitespacesAndNewlines) + let updatedConfig = Self.configByAddingTokenAccount( + config, + provider: provider, + label: rawLabel, + token: rawToken, + activate: true) + + do { + try store.save(updatedConfig) + } catch { + Self.exit(code: .failure, message: error.localizedDescription, output: output, kind: .config) + } + + let providerConfig = updatedConfig.providerConfig(for: provider) + let accounts = providerConfig?.tokenAccounts?.accounts ?? [] + let addedAccount = accounts.last + + switch output.format { + case .text: + let label = addedAccount?.label ?? rawLabel ?? "Account" + print("Added \(provider.rawValue) account '\(label)' (\(accounts.count) total).") + case .json: + let payload = ConfigAddAccountPayload( + provider: provider.rawValue, + label: addedAccount?.label ?? rawLabel ?? "", + activeIndex: providerConfig?.tokenAccounts?.activeIndex ?? 0, + count: accounts.count) + Self.printJSON(payload, pretty: output.pretty) + } + + Self.exit(code: .success, output: output, kind: .config) + } + static func runConfigValidate(_ values: ParsedValues) { let output = CLIOutputPreferences.from(values: values) let config = Self.loadConfig(output: output) @@ -35,6 +106,42 @@ extension CodexBarCLI { Self.printJSON(config, pretty: output.pretty) Self.exit(code: .success, output: output, kind: .config) } + + private static func decodeSingleProvider(from values: ParsedValues) -> UsageProvider? { + guard let raw = values.options["provider"]?.last else { return nil } + guard case let .single(provider) = ProviderSelection(argument: raw) else { return nil } + return provider + } + + private static func configByAddingTokenAccount( + _ config: CodexBarConfig, + provider: UsageProvider, + label: String?, + token: String, + activate: Bool) -> CodexBarConfig + { + var updatedConfig = config + var providerConfig = updatedConfig.providerConfig(for: provider) ?? ProviderConfig(id: provider) + let existing = providerConfig.tokenAccounts + let accounts = existing?.accounts ?? [] + let fallbackLabel = label?.isEmpty == false ? label! : "Account \(accounts.count + 1)" + let account = ProviderTokenAccount( + id: UUID(), + label: fallbackLabel, + token: token, + addedAt: Date().timeIntervalSince1970, + lastUsed: nil) + let activeIndex = activate ? accounts.count : (existing?.clampedActiveIndex() ?? 0) + providerConfig.tokenAccounts = ProviderTokenAccountData( + version: existing?.version ?? 1, + accounts: accounts + [account], + activeIndex: activeIndex) + if TokenAccountSupportCatalog.support(for: provider)?.requiresManualCookieSource == true { + providerConfig.cookieSource = .manual + } + updatedConfig.setProviderConfig(providerConfig) + return updatedConfig + } } struct ConfigOptions: CommanderParsable { diff --git a/Sources/CodexBarCLI/CLIEntry.swift b/Sources/CodexBarCLI/CLIEntry.swift index 42957f717..4f206bb66 100644 --- a/Sources/CodexBarCLI/CLIEntry.swift +++ b/Sources/CodexBarCLI/CLIEntry.swift @@ -39,6 +39,8 @@ enum CodexBarCLI { await self.runUsage(invocation.parsedValues) case ["cost"]: await self.runCost(invocation.parsedValues) + case ["config", "add-account"]: + self.runConfigAddAccount(invocation.parsedValues) case ["config", "validate"]: self.runConfigValidate(invocation.parsedValues) case ["config", "dump"]: @@ -61,6 +63,7 @@ enum CodexBarCLI { let usageSignature = CommandSignature.describe(UsageOptions()) let costSignature = CommandSignature.describe(CostOptions()) let configSignature = CommandSignature.describe(ConfigOptions()) + let configAddAccountSignature = CommandSignature.describe(ConfigAddAccountOptions()) return [ CommandDescriptor( @@ -79,6 +82,11 @@ enum CodexBarCLI { discussion: nil, signature: CommandSignature(), subcommands: [ + CommandDescriptor( + name: "add-account", + abstract: "Add a token account to config", + discussion: nil, + signature: configAddAccountSignature), CommandDescriptor( name: "validate", abstract: "Validate config file", diff --git a/Sources/CodexBarCLI/CLIHelp.swift b/Sources/CodexBarCLI/CLIHelp.swift index 09e2a32bc..70fb636ae 100644 --- a/Sources/CodexBarCLI/CLIHelp.swift +++ b/Sources/CodexBarCLI/CLIHelp.swift @@ -74,10 +74,18 @@ extension CodexBarCLI { CodexBar \(version) Usage: + codexbar config add-account --provider --token + [--label