diff --git a/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift b/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift index 0acd256..9c8d1ca 100644 --- a/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift +++ b/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift @@ -368,6 +368,8 @@ final class CodexProvider: ProviderProtocol { return 2 case .openCodeMultiAuth: return 1 + case .openCodeAnthropicAuthCodexCache: + return -1 case .codexLB: return 0 } @@ -379,6 +381,8 @@ final class CodexProvider: ProviderProtocol { return "OpenCode" case .openCodeMultiAuth: return "OpenCode Multi Auth" + case .openCodeAnthropicAuthCodexCache: + return "OpenCode Anthropic Auth" case .codexLB: return "Codex LB" case .codexAuth: @@ -427,6 +431,39 @@ final class CodexProvider: ProviderProtocol { } private func fetchUsageForAccount(_ account: OpenAIAuthAccount) async throws -> CodexAccountCandidate { + if let cachedUsage = account.cachedCodexUsage { + return buildCachedUsageCandidate(account: account, cachedUsage: cachedUsage) + } + + var account = account + var didRefresh = false + + if TokenManager.shared.openAIMultiAuthAccountNeedsRefresh(account) { + do { + account = try await TokenManager.shared.refreshOpenAIMultiAuthAccount(account) + didRefresh = true + logger.info("Codex retry will use refreshed OpenAI multi-auth token") + } catch { + logger.warning("OpenAI multi-auth token refresh before Codex request failed: \(error.localizedDescription)") + } + } + + do { + return try await fetchUsageForResolvedAccount(account) + } catch { + guard !didRefresh, + isUnauthorizedError(error), + TokenManager.shared.canRefreshOpenAIMultiAuthAccount(account) else { + throw error + } + + logger.info("Codex API returned 401 for OpenAI multi-auth account; refreshing token and retrying once") + let refreshedAccount = try await TokenManager.shared.refreshOpenAIMultiAuthAccount(account) + return try await fetchUsageForResolvedAccount(refreshedAccount) + } + } + + private func fetchUsageForResolvedAccount(_ account: OpenAIAuthAccount) async throws -> CodexAccountCandidate { let endpointConfiguration = TokenManager.shared.getCodexEndpointConfiguration() let url = try codexUsageURL(for: endpointConfiguration, account: account) @@ -469,6 +506,61 @@ final class CodexProvider: ProviderProtocol { ) } + private func isUnauthorizedError(_ error: Error) -> Bool { + guard case ProviderError.networkError(let message) = error else { + return false + } + return message.contains("HTTP 401") + } + + private func buildCachedUsageCandidate( + account: OpenAIAuthAccount, + cachedUsage: CodexCachedUsageSnapshot + ) -> CodexAccountCandidate { + let sourceLabels = account.sourceLabels.isEmpty ? [sourceLabel(account.source)] : account.sourceLabels + let authUsageSummary = sourceSummary(sourceLabels, fallback: "Unknown") + let primaryPercent = cachedUsage.primary?.utilization ?? 0 + let remaining = max(0, Int(round(100 - primaryPercent))) + let details = DetailedUsage( + dailyUsage: cachedUsage.primary?.utilization, + secondaryUsage: cachedUsage.secondary?.utilization, + primaryReset: cachedUsage.primary?.resetsAt, + codexPrimaryWindowLabel: cachedUsage.primary?.label, + codexPrimaryWindowHours: hours(fromWindowMs: cachedUsage.primary?.windowMs), + codexSecondaryWindowLabel: cachedUsage.secondary?.label, + codexSecondaryWindowHours: hours(fromWindowMs: cachedUsage.secondary?.windowMs), + sparkUsage: cachedUsage.sparkPrimary?.utilization, + sparkReset: cachedUsage.sparkPrimary?.resetsAt, + sparkSecondaryUsage: cachedUsage.sparkSecondary?.utilization, + sparkSecondaryReset: cachedUsage.sparkSecondary?.resetsAt, + sparkPrimaryWindowLabel: cachedUsage.sparkPrimary?.label, + sparkPrimaryWindowHours: hours(fromWindowMs: cachedUsage.sparkPrimary?.windowMs), + sparkSecondaryWindowLabel: cachedUsage.sparkSecondary?.label, + sparkSecondaryWindowHours: hours(fromWindowMs: cachedUsage.sparkSecondary?.windowMs), + creditsBalance: cachedUsage.creditsBalance, + planType: cachedUsage.planType, + email: account.email, + authSource: account.authSource, + authUsageSummary: authUsageSummary + ) + + logger.info( + """ + Codex cached usage loaded (\(authUsageSummary)): \ + account=\(account.email ?? account.accountId ?? "unknown", privacy: .private), \ + primary=\(primaryPercent, privacy: .public)% + """ + ) + + return CodexAccountCandidate( + accountId: account.accountId, + usage: ProviderUsage.quotaBased(remaining: remaining, entitlement: 100, overagePermitted: false), + details: details, + sourceLabels: sourceLabels, + source: account.source + ) + } + func decodeUsagePayload( data: Data, account: OpenAIAuthAccount, @@ -886,6 +978,11 @@ final class CodexProvider: ProviderProtocol { } } + private func hours(fromWindowMs windowMs: Int?) -> Int? { + guard let windowMs, windowMs > 0 else { return nil } + return max(1, Int(round(Double(windowMs) / 3_600_000.0))) + } + private func codexWindowMetadata(for window: RateLimitWindow, fallbackLabel: String) -> (label: String, hours: Int?) { if let seconds = window.limit_window_seconds, seconds > 0 { diff --git a/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift b/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift index a3465f1..7171fdc 100644 --- a/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift +++ b/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift @@ -372,10 +372,29 @@ struct CodexLBEncryptedAccount { let lastRefresh: String? } +struct CodexCachedUsageWindow { + let utilization: Double + let resetsAt: Date? + let label: String? + let windowMs: Int? +} + +struct CodexCachedUsageSnapshot { + let fetchedAt: Date? + let planType: String? + let primary: CodexCachedUsageWindow? + let secondary: CodexCachedUsageWindow? + let sparkPrimary: CodexCachedUsageWindow? + let sparkSecondary: CodexCachedUsageWindow? + let creditsBalance: Double? + let creditsUnlimited: Bool? +} + /// Auth source types for OpenAI (Codex) account discovery enum OpenAIAuthSource { case opencodeAuth case openCodeMultiAuth + case openCodeAnthropicAuthCodexCache case codexLB case codexAuth } @@ -395,6 +414,38 @@ struct OpenAIAuthAccount { let sourceLabels: [String] let source: OpenAIAuthSource let credentialType: OpenAICredentialType + let refreshToken: String? + let expiresAt: Date? + let idToken: String? + let cachedCodexUsage: CodexCachedUsageSnapshot? + + init( + accessToken: String, + accountId: String?, + externalUsageAccountId: String?, + email: String?, + authSource: String, + sourceLabels: [String], + source: OpenAIAuthSource, + credentialType: OpenAICredentialType, + refreshToken: String? = nil, + expiresAt: Date? = nil, + idToken: String? = nil, + cachedCodexUsage: CodexCachedUsageSnapshot? = nil + ) { + self.accessToken = accessToken + self.accountId = accountId + self.externalUsageAccountId = externalUsageAccountId + self.email = email + self.authSource = authSource + self.sourceLabels = sourceLabels + self.source = source + self.credentialType = credentialType + self.refreshToken = refreshToken + self.expiresAt = expiresAt + self.idToken = idToken + self.cachedCodexUsage = cachedCodexUsage + } } enum CodexEndpointMode: Equatable { @@ -744,6 +795,13 @@ final class TokenManager: @unchecked Sendable { /// Paths where oc-chatgpt-multi-auth account files were found private(set) var lastFoundOpenCodeMultiAuthPaths: [URL] = [] + /// Cached opencode-anthropic-auth Codex usage accounts + private var cachedOpenCodeAnthropicCodexAccounts: [OpenAIAuthAccount]? + private var openCodeAnthropicCodexAccountsCacheTimestamp: Date? + + /// Paths where opencode-anthropic-auth Codex account files were found + private(set) var lastFoundOpenCodeAnthropicCodexAccountPaths: [URL] = [] + /// Cached GitHub Copilot token accounts (OpenCode + VS Code) private var cachedCopilotAccounts: [CopilotAuthAccount]? private var copilotAccountsCacheTimestamp: Date? @@ -1238,6 +1296,9 @@ final class TokenManager: @unchecked Sendable { cachedAuth = nil cacheTimestamp = nil lastFoundAuthPath = nil + cachedOpenCodeAnthropicCodexAccounts = nil + openCodeAnthropicCodexAccountsCacheTimestamp = nil + lastFoundOpenCodeAnthropicCodexAccountPaths = [] } } @@ -1681,6 +1742,13 @@ final class TokenManager: @unchecked Sendable { ) } + func shouldIncludeCodexLBAccount(_ account: CodexLBEncryptedAccount) -> Bool { + guard let status = normalizedNonEmpty(account.status)?.lowercased() else { + return true + } + return status == "active" + } + func readCodexLBOpenAIAccounts() -> [OpenAIAuthAccount] { return queue.sync { if let cached = cachedCodexLBAccounts, @@ -1713,6 +1781,13 @@ final class TokenManager: @unchecked Sendable { var decodedAccounts: [OpenAIAuthAccount] = [] for encryptedAccount in encryptedAccounts { + guard shouldIncludeCodexLBAccount(encryptedAccount) else { + logger.info( + "Skipping inactive codex-lb OpenAI account with status \(encryptedAccount.status ?? "unknown", privacy: .public)" + ) + continue + } + do { let accessToken = try decryptCodexLBFernetToken( encryptedAccount.accessTokenEncrypted, @@ -1777,6 +1852,8 @@ final class TokenManager: @unchecked Sendable { return "OpenCode" case .openCodeMultiAuth: return "OpenCode Multi Auth" + case .openCodeAnthropicAuthCodexCache: + return "OpenCode Anthropic Auth" case .codexLB: return "Codex LB" case .codexAuth: @@ -1854,6 +1931,7 @@ final class TokenManager: @unchecked Sendable { case .opencodeAuth: return 3 case .codexAuth: return 2 case .openCodeMultiAuth: return 1 + case .openCodeAnthropicAuthCodexCache: return -1 case .codexLB: return 0 } } @@ -1947,7 +2025,11 @@ final class TokenManager: @unchecked Sendable { authSource: primary.authSource, sourceLabels: mergedSourceLabels, source: primary.source, - credentialType: primary.credentialType + credentialType: primary.credentialType, + refreshToken: normalizedNonEmpty(primary.refreshToken) ?? normalizedNonEmpty(fallback.refreshToken), + expiresAt: primary.expiresAt ?? fallback.expiresAt, + idToken: normalizedNonEmpty(primary.idToken) ?? normalizedNonEmpty(fallback.idToken), + cachedCodexUsage: primary.cachedCodexUsage ?? fallback.cachedCodexUsage ) } @@ -2133,6 +2215,9 @@ final class TokenManager: @unchecked Sendable { let accessToken: String let accountId: String? let email: String? + let refreshToken: String? + let expiresAt: Date? + let idToken: String? } private func valueForNormalizedKey(_ normalizedKeyName: String, in dict: [String: Any]) -> Any? { @@ -2272,6 +2357,9 @@ final class TokenManager: @unchecked Sendable { private func extractOpenAIMultiAuthPayload(from dict: [String: Any]) -> OpenAIMultiAuthPayload? { let accessKeys: Set = ["accesstoken", "access", "oauthtoken", "token"] + let refreshKeys: Set = ["refreshtoken", "oauthrefreshtoken", "refresh"] + let expiresKeys: Set = ["expiresat", "expires", "expiration", "expiry"] + let idTokenKeys: Set = ["idtoken"] let accountKeys: Set = ["accountid", "chatgptaccountid", "userid", "id"] let emailKeys: Set = ["email", "useremail", "login", "username"] @@ -2284,6 +2372,12 @@ final class TokenManager: @unchecked Sendable { let tokenAccountId = normalizedNonEmpty(accessTokenPayload?.auth?.chatGPTAccountId) let storedAccountId = findDirectStringValue(in: dict, matching: accountKeys) ?? findStringValue(in: dict, matching: accountKeys) + let expiresRaw = findDirectInt64Value(in: dict, matching: expiresKeys) + ?? findInt64Value(in: dict, matching: expiresKeys) + let expiresAt = expiresRaw.flatMap { dateFromEpoch($0) } ?? parseISO8601Date( + findDirectStringValue(in: dict, matching: expiresKeys) + ?? findStringValue(in: dict, matching: expiresKeys) + ) let email = normalizedNonEmpty(accessTokenPayload?.profile?.email) ?? normalizedNonEmpty(findDirectStringValue(in: dict, matching: emailKeys) ?? findStringValue(in: dict, matching: emailKeys)) @@ -2291,10 +2385,178 @@ final class TokenManager: @unchecked Sendable { return OpenAIMultiAuthPayload( accessToken: accessToken, accountId: tokenAccountId ?? normalizedNonEmpty(storedAccountId), - email: email + email: email, + refreshToken: normalizedNonEmpty(findDirectStringValue(in: dict, matching: refreshKeys) + ?? findStringValue(in: dict, matching: refreshKeys)), + expiresAt: expiresAt, + idToken: normalizedNonEmpty(findDirectStringValue(in: dict, matching: idTokenKeys) + ?? findStringValue(in: dict, matching: idTokenKeys)) ) } + private func openCodeAnthropicCodexAccountPaths() -> [URL] { + buildOpenCodeFilePaths( + envVarName: "XDG_CONFIG_HOME", + envRelativePathComponents: ["opencode", "opencode-anthropic-auth", "codex-accounts.json"], + fallbackRelativePathComponents: [ + [".config", "opencode", "opencode-anthropic-auth", "codex-accounts.json"] + ] + ) + } + + private func boolValue(_ value: Any?) -> Bool? { + if let bool = value as? Bool { + return bool + } + if let number = value as? NSNumber { + return number.boolValue + } + if let string = value as? String { + switch string.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() { + case "true", "yes", "1": + return true + case "false", "no", "0": + return false + default: + return nil + } + } + return nil + } + + private func directDoubleValue(in dict: [String: Any], matching keys: Set) -> Double? { + for (key, value) in dict where keys.contains(normalizedKey(key)) { + if let double = value as? Double { + return double + } + if let int = value as? Int { + return Double(int) + } + if let number = value as? NSNumber { + return number.doubleValue + } + if let string = value as? String { + return Double(string.trimmingCharacters(in: .whitespacesAndNewlines)) + } + } + return nil + } + + private func directBoolValue(in dict: [String: Any], matching keys: Set) -> Bool? { + for (key, value) in dict where keys.contains(normalizedKey(key)) { + if let bool = boolValue(value) { + return bool + } + } + return nil + } + + private func parseOpenCodeAnthropicCodexUsageWindow(_ value: Any?) -> CodexCachedUsageWindow? { + guard let dict = value as? [String: Any], + let utilization = directDoubleValue(in: dict, matching: ["utilization", "usedpercent", "usagepercent"]) else { + return nil + } + + let resetDate = parseISO8601Date(findDirectStringValue(in: dict, matching: ["resetsat", "resetat"])) + ?? dateFromEpoch(findDirectInt64Value(in: dict, matching: ["resetsat", "resetat"])) + let label = normalizedNonEmpty(findDirectStringValue(in: dict, matching: ["label", "windowlabel"])) + let windowMs = findDirectInt64Value(in: dict, matching: ["windowms"]).map { Int($0) } + + return CodexCachedUsageWindow( + utilization: utilization, + resetsAt: resetDate, + label: label, + windowMs: windowMs + ) + } + + private func parseOpenCodeAnthropicCodexUsage(_ value: Any?) -> CodexCachedUsageSnapshot? { + guard let dict = value as? [String: Any] else { + return nil + } + + let primary = parseOpenCodeAnthropicCodexUsageWindow(valueForNormalizedKey("primary", in: dict)) + let secondary = parseOpenCodeAnthropicCodexUsageWindow(valueForNormalizedKey("secondary", in: dict)) + let sparkPrimary = parseOpenCodeAnthropicCodexUsageWindow(valueForNormalizedKey("sparkprimary", in: dict)) + let sparkSecondary = parseOpenCodeAnthropicCodexUsageWindow(valueForNormalizedKey("sparksecondary", in: dict)) + + guard primary != nil || secondary != nil || sparkPrimary != nil || sparkSecondary != nil else { + return nil + } + + return CodexCachedUsageSnapshot( + fetchedAt: dateFromEpoch(findDirectInt64Value(in: dict, matching: ["fetchedat"])), + planType: normalizedNonEmpty(findDirectStringValue(in: dict, matching: ["plantype"])), + primary: primary, + secondary: secondary, + sparkPrimary: sparkPrimary, + sparkSecondary: sparkSecondary, + creditsBalance: directDoubleValue(in: dict, matching: ["creditsbalance", "balance"]), + creditsUnlimited: directBoolValue(in: dict, matching: ["creditsunlimited", "unlimited"]) + ) + } + + private func shouldIncludeOpenCodeAnthropicCodexAccount(_ accountDict: [String: Any]) -> Bool { + if let enabled = boolValue(valueForNormalizedKey("enabled", in: accountDict)), !enabled { + return false + } + guard let status = normalizedNonEmpty(findDirectStringValue(in: accountDict, matching: ["status"]))?.lowercased() else { + return true + } + return status == "active" + } + + func readOpenCodeAnthropicCodexAccountFiles(at paths: [URL]) -> [OpenAIAuthAccount] { + var accounts: [OpenAIAuthAccount] = [] + + for path in paths { + guard let dict = readJSONDictionary(at: path) else { continue } + let rawAccounts = valueForNormalizedKey("accounts", in: dict) as? [Any] ?? [dict] + var pathAccounts: [OpenAIAuthAccount] = [] + + for rawAccount in rawAccounts { + guard let accountDict = rawAccount as? [String: Any], + shouldIncludeOpenCodeAnthropicCodexAccount(accountDict), + let cachedUsage = parseOpenCodeAnthropicCodexUsage(valueForNormalizedKey("usage", in: accountDict)) else { + continue + } + + let accountId = normalizedNonEmpty(findDirectStringValue( + in: accountDict, + matching: ["accountid", "chatgptaccountid"] + )) ?? normalizedNonEmpty(findDirectStringValue(in: accountDict, matching: ["id"])) + let email = normalizedNonEmpty(findDirectStringValue( + in: accountDict, + matching: ["email", "useremail", "login", "username"] + )) + guard accountId != nil || email != nil else { + continue + } + + pathAccounts.append( + OpenAIAuthAccount( + accessToken: "", + accountId: accountId, + externalUsageAccountId: nil, + email: email, + authSource: path.path, + sourceLabels: [openAISourceLabel(for: .openCodeAnthropicAuthCodexCache)], + source: .openCodeAnthropicAuthCodexCache, + credentialType: .oauthBearer, + cachedCodexUsage: cachedUsage + ) + ) + } + + if !pathAccounts.isEmpty { + logger.info("Loaded \(pathAccounts.count) cached Codex account(s) from opencode-anthropic-auth at \(path.path)") + accounts.append(contentsOf: pathAccounts) + } + } + + return accounts + } + func readOpenAIMultiAuthFiles(at paths: [URL]) -> [OpenAIAuthAccount] { var accounts: [OpenAIAuthAccount] = [] @@ -2318,7 +2580,10 @@ final class TokenManager: @unchecked Sendable { authSource: path.path, sourceLabels: [openAISourceLabel(for: .openCodeMultiAuth)], source: .openCodeMultiAuth, - credentialType: .oauthBearer + credentialType: .oauthBearer, + refreshToken: payload.refreshToken, + expiresAt: payload.expiresAt, + idToken: payload.idToken ) ) } @@ -2352,6 +2617,219 @@ final class TokenManager: @unchecked Sendable { } } + private func readOpenCodeAnthropicCodexAccountFiles() -> [OpenAIAuthAccount] { + return queue.sync { + if let cached = cachedOpenCodeAnthropicCodexAccounts, + let timestamp = openCodeAnthropicCodexAccountsCacheTimestamp, + Date().timeIntervalSince(timestamp) < cacheValiditySeconds { + return cached + } + + let fileManager = FileManager.default + let paths = openCodeAnthropicCodexAccountPaths() + let accounts = readOpenCodeAnthropicCodexAccountFiles(at: paths) + let existingPaths = paths.filter { fileManager.fileExists(atPath: $0.path) } + + cachedOpenCodeAnthropicCodexAccounts = accounts + openCodeAnthropicCodexAccountsCacheTimestamp = Date() + lastFoundOpenCodeAnthropicCodexAccountPaths = existingPaths + return accounts + } + } + + private struct OpenAITokenRefreshResponse: Decodable { + let accessToken: String + let refreshToken: String? + let expiresIn: Double + let idToken: String? + + enum CodingKeys: String, CodingKey { + case accessToken = "access_token" + case refreshToken = "refresh_token" + case expiresIn = "expires_in" + case idToken = "id_token" + } + } + + private enum OpenAITokenRefreshError: LocalizedError { + case unsupportedSource + case missingRefreshToken + case invalidTokenURL + case invalidResponse + case httpStatus(Int) + case missingAccessToken + case storageUpdateFailed + + var errorDescription: String? { + switch self { + case .unsupportedSource: + return "OpenAI token refresh is only supported for OpenCode Multi Auth accounts" + case .missingRefreshToken: + return "OpenAI refresh token is missing" + case .invalidTokenURL: + return "OpenAI token refresh URL is invalid" + case .invalidResponse: + return "OpenAI token refresh returned an invalid response" + case .httpStatus(let status): + return "OpenAI token refresh failed with HTTP \(status)" + case .missingAccessToken: + return "OpenAI token refresh response did not include an access token" + case .storageUpdateFailed: + return "OpenAI multi-auth account storage could not be updated" + } + } + } + + func canRefreshOpenAIMultiAuthAccount(_ account: OpenAIAuthAccount) -> Bool { + account.source == .openCodeMultiAuth + && normalizedNonEmpty(account.refreshToken) != nil + && account.credentialType == .oauthBearer + } + + func openAIMultiAuthAccountNeedsRefresh(_ account: OpenAIAuthAccount, skew: TimeInterval = 60) -> Bool { + guard canRefreshOpenAIMultiAuthAccount(account), + let expiresAt = account.expiresAt else { + return false + } + + return expiresAt <= Date().addingTimeInterval(max(0, skew)) + } + + private func formURLEncodedBody(_ queryItems: [URLQueryItem]) -> Data { + var components = URLComponents() + components.queryItems = queryItems + return components.percentEncodedQuery?.data(using: .utf8) ?? Data() + } + + func refreshOpenAIMultiAuthAccount(_ account: OpenAIAuthAccount) async throws -> OpenAIAuthAccount { + guard account.source == .openCodeMultiAuth else { + throw OpenAITokenRefreshError.unsupportedSource + } + guard let refreshToken = normalizedNonEmpty(account.refreshToken) else { + throw OpenAITokenRefreshError.missingRefreshToken + } + guard let tokenURL = URL(string: "https://auth.openai.com/oauth/token") else { + throw OpenAITokenRefreshError.invalidTokenURL + } + + logger.info("Refreshing OpenAI multi-auth token for Codex usage") + + var request = URLRequest(url: tokenURL) + request.httpMethod = "POST" + request.setValue("application/x-www-form-urlencoded", forHTTPHeaderField: "Content-Type") + request.httpBody = formURLEncodedBody([ + URLQueryItem(name: "grant_type", value: "refresh_token"), + URLQueryItem(name: "refresh_token", value: refreshToken), + URLQueryItem(name: "client_id", value: "app_EMoamEEZ73f0CkXaXp7hrann") + ]) + + let (data, response) = try await URLSession.shared.data(for: request) + guard let httpResponse = response as? HTTPURLResponse else { + throw OpenAITokenRefreshError.invalidResponse + } + guard httpResponse.statusCode == 200 else { + logger.warning("OpenAI multi-auth token refresh failed with status \(httpResponse.statusCode)") + throw OpenAITokenRefreshError.httpStatus(httpResponse.statusCode) + } + + let refreshResponse = try JSONDecoder().decode(OpenAITokenRefreshResponse.self, from: data) + let accessToken = refreshResponse.accessToken.trimmingCharacters(in: .whitespacesAndNewlines) + guard !accessToken.isEmpty else { + throw OpenAITokenRefreshError.missingAccessToken + } + + let refreshedAccessPayload = decodeOpenAIAccessTokenPayload(accessToken) + let refreshedIDPayload = decodeOpenAIIDTokenPayload(refreshResponse.idToken) + let refreshedAccount = OpenAIAuthAccount( + accessToken: accessToken, + accountId: normalizedNonEmpty(refreshedAccessPayload?.auth?.chatGPTAccountId) ?? account.accountId, + externalUsageAccountId: account.externalUsageAccountId, + email: normalizedNonEmpty(refreshedIDPayload?.email) + ?? normalizedNonEmpty(refreshedAccessPayload?.profile?.email) + ?? account.email, + authSource: account.authSource, + sourceLabels: account.sourceLabels, + source: account.source, + credentialType: account.credentialType, + refreshToken: normalizedNonEmpty(refreshResponse.refreshToken) ?? refreshToken, + expiresAt: Date().addingTimeInterval(refreshResponse.expiresIn), + idToken: normalizedNonEmpty(refreshResponse.idToken) ?? account.idToken + ) + + try persistOpenAIMultiAuthRefresh(original: account, refreshed: refreshedAccount) + logger.info("OpenAI multi-auth token refresh succeeded for Codex usage") + return refreshedAccount + } + + private func persistOpenAIMultiAuthRefresh(original: OpenAIAuthAccount, refreshed: OpenAIAuthAccount) throws { + let storageURL = URL(fileURLWithPath: original.authSource) + try queue.sync { + let fileManager = FileManager.default + guard fileManager.fileExists(atPath: storageURL.path), + fileManager.isReadableFile(atPath: storageURL.path) else { + throw OpenAITokenRefreshError.storageUpdateFailed + } + + let data = try Data(contentsOf: storageURL) + guard var root = try JSONSerialization.jsonObject(with: data, options: []) as? [String: Any] else { + throw OpenAITokenRefreshError.storageUpdateFailed + } + + var updated = false + if var accounts = root["accounts"] as? [[String: Any]] { + for index in accounts.indices where openAIMultiAuthStorageAccount(accounts[index], matches: original) { + updateOpenAIMultiAuthStorageAccount(&accounts[index], with: refreshed) + updated = true + } + root["accounts"] = accounts + } else if openAIMultiAuthStorageAccount(root, matches: original) { + updateOpenAIMultiAuthStorageAccount(&root, with: refreshed) + updated = true + } + + guard updated else { + throw OpenAITokenRefreshError.storageUpdateFailed + } + + let output = try JSONSerialization.data(withJSONObject: root, options: [.prettyPrinted, .sortedKeys]) + try output.write(to: storageURL, options: .atomic) + + cachedOpenCodeMultiAuthAccounts = nil + openCodeMultiAuthAccountsCacheTimestamp = nil + } + } + + private func openAIMultiAuthStorageAccount(_ accountDict: [String: Any], matches account: OpenAIAuthAccount) -> Bool { + let accessKeys: Set = ["accesstoken", "access", "oauthtoken", "token"] + let refreshKeys: Set = ["refreshtoken", "oauthrefreshtoken", "refresh"] + let accountKeys: Set = ["accountid", "chatgptaccountid", "userid", "id"] + let emailKeys: Set = ["email", "useremail", "login", "username"] + + if let refreshToken = normalizedNonEmpty(account.refreshToken), + normalizedNonEmpty(findDirectStringValue(in: accountDict, matching: refreshKeys)) == refreshToken { + return true + } + + if normalizedNonEmpty(findDirectStringValue(in: accountDict, matching: accessKeys)) == account.accessToken { + return true + } + + let storedAccountId = normalizedNonEmpty(findDirectStringValue(in: accountDict, matching: accountKeys)) + let storedEmail = normalizedNonEmpty(findDirectStringValue(in: accountDict, matching: emailKeys))?.lowercased() + let accountIdMatches = storedAccountId != nil && storedAccountId == normalizedNonEmpty(account.accountId) + let emailMatches = storedEmail != nil && storedEmail == normalizedNonEmpty(account.email)?.lowercased() + return accountIdMatches && emailMatches + } + + private func updateOpenAIMultiAuthStorageAccount(_ accountDict: inout [String: Any], with account: OpenAIAuthAccount) { + accountDict["accessToken"] = account.accessToken + accountDict["refreshToken"] = account.refreshToken + accountDict["expiresAt"] = account.expiresAt.map { Int64($0.timeIntervalSince1970 * 1000) } + if let idToken = normalizedNonEmpty(account.idToken) { + accountDict["idToken"] = idToken + } + } + private func parseJSONDictionary(from data: Data) -> [String: Any]? { guard let json = try? JSONSerialization.jsonObject(with: data, options: []), let dict = json as? [String: Any] else { @@ -3165,6 +3643,11 @@ final class TokenManager: @unchecked Sendable { accounts.append(contentsOf: openCodeMultiAuthAccounts) } + let openCodeAnthropicCodexAccounts = readOpenCodeAnthropicCodexAccountFiles() + if !openCodeAnthropicCodexAccounts.isEmpty { + accounts.append(contentsOf: openCodeAnthropicCodexAccounts) + } + let codexLBAccounts = readCodexLBOpenAIAccounts() if !codexLBAccounts.isEmpty { accounts.append(contentsOf: codexLBAccounts) diff --git a/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift b/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift index 6a9a9a2..97db21c 100644 --- a/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift +++ b/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift @@ -369,6 +369,168 @@ final class TokenManagerTests: XCTestCase { XCTAssertEqual(accounts.map(\.sourceLabels), [["OpenCode Multi Auth"], ["OpenCode Multi Auth"]]) } + func testReadOpenAIMultiAuthFilesCapturesRefreshMetadata() throws { + let fileManager = FileManager.default + let tempDirectory = fileManager.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + let accountsPath = tempDirectory.appendingPathComponent("openai-codex-accounts.json") + + try fileManager.createDirectory(at: tempDirectory, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: tempDirectory) } + + let accessToken = makeTestJWT( + payload: #""" + { + "https://api.openai.com/auth": { + "chatgpt_account_id": "chatgpt-account-id" + }, + "https://api.openai.com/profile": { + "email": "user@example.com" + } + } + """# + ) + + let json = """ + { + "version": 3, + "accounts": [ + { + "accountId": "chatgpt-account-id", + "accountIdSource": "token", + "accessToken": "\(accessToken)", + "refreshToken": "refresh-token", + "expiresAt": 1770563557150 + } + ], + "activeIndex": 0 + } + """ + + try XCTUnwrap(json.data(using: .utf8)).write(to: accountsPath) + + let account = try XCTUnwrap(TokenManager.shared.readOpenAIMultiAuthFiles(at: [accountsPath]).first) + + XCTAssertEqual(account.refreshToken, "refresh-token") + let expiresAt = try XCTUnwrap(account.expiresAt) + XCTAssertEqual(expiresAt.timeIntervalSince1970, 1_770_563_557.15, accuracy: 0.01) + } + + func testReadOpenCodeAnthropicCodexAccountFilesCapturesCachedUsage() throws { + let fileManager = FileManager.default + let tempDirectory = fileManager.temporaryDirectory + .appendingPathComponent(UUID().uuidString, isDirectory: true) + let accountsPath = tempDirectory.appendingPathComponent("codex-accounts.json") + + try fileManager.createDirectory(at: tempDirectory, withIntermediateDirectories: true) + defer { try? fileManager.removeItem(at: tempDirectory) } + + let json = """ + { + "version": 1, + "accounts": [ + { + "id": "local-cache-id", + "accountID": "chatgpt-account-id", + "email": "user@example.com", + "enabled": true, + "usage": { + "fetchedAt": 1778410367180, + "planType": "pro", + "primary": { + "utilization": 25, + "resetsAt": "2026-05-10T15:15:02.180Z", + "label": "5h", + "windowMs": 18000000 + }, + "secondary": { + "utilization": 75, + "resetsAt": "2026-05-15T19:42:30.180Z", + "label": "Weekly", + "windowMs": 604800000 + }, + "sparkPrimary": { + "utilization": 10, + "resetsAt": "2026-05-10T15:52:47.715Z", + "label": "Spark 5h", + "windowMs": 18000000 + }, + "creditsBalance": "0", + "creditsUnlimited": false + } + }, + { + "id": "disabled-cache-id", + "email": "disabled@example.com", + "enabled": false, + "usage": { + "primary": { + "utilization": 50 + } + } + } + ] + } + """ + + try XCTUnwrap(json.data(using: .utf8)).write(to: accountsPath) + + let accounts = TokenManager.shared.readOpenCodeAnthropicCodexAccountFiles(at: [accountsPath]) + + XCTAssertEqual(accounts.count, 1) + let account = try XCTUnwrap(accounts.first) + XCTAssertEqual(account.accountId, "chatgpt-account-id") + XCTAssertEqual(account.email, "user@example.com") + XCTAssertEqual(account.authSource, accountsPath.path) + XCTAssertEqual(account.source, .openCodeAnthropicAuthCodexCache) + XCTAssertEqual(account.sourceLabels, ["OpenCode Anthropic Auth"]) + XCTAssertEqual(account.cachedCodexUsage?.planType, "pro") + XCTAssertEqual(account.cachedCodexUsage?.primary?.utilization, 25) + XCTAssertEqual(account.cachedCodexUsage?.primary?.label, "5h") + XCTAssertEqual(account.cachedCodexUsage?.primary?.windowMs, 18_000_000) + XCTAssertEqual(account.cachedCodexUsage?.secondary?.utilization, 75) + XCTAssertEqual(account.cachedCodexUsage?.sparkPrimary?.utilization, 10) + XCTAssertEqual(account.cachedCodexUsage?.creditsBalance, 0) + XCTAssertEqual(account.cachedCodexUsage?.creditsUnlimited, false) + } + + func testDedupeOpenAIAccountsPreservesOpenAIMultiAuthRefreshMetadata() throws { + let expiresAt = Date(timeIntervalSince1970: 1_770_563_557.15) + let firstAccount = OpenAIAuthAccount( + accessToken: "first-token", + accountId: "chatgpt-account-id", + externalUsageAccountId: nil, + email: "user@example.com", + authSource: "/tmp/project-a/openai-codex-accounts.json", + sourceLabels: ["OpenCode Multi Auth"], + source: .openCodeMultiAuth, + credentialType: .oauthBearer, + refreshToken: "refresh-token", + expiresAt: expiresAt, + idToken: "id-token" + ) + let duplicateAccount = OpenAIAuthAccount( + accessToken: "duplicate-token", + accountId: "chatgpt-account-id", + externalUsageAccountId: nil, + email: "user@example.com", + authSource: "/tmp/project-b/openai-codex-accounts.json", + sourceLabels: ["OpenCode Multi Auth"], + source: .openCodeMultiAuth, + credentialType: .oauthBearer, + refreshToken: "duplicate-refresh-token", + expiresAt: Date(timeIntervalSince1970: 1_770_563_600), + idToken: "duplicate-id-token" + ) + + let account = try XCTUnwrap(TokenManager.shared.dedupeOpenAIAccounts([firstAccount, duplicateAccount]).first) + + XCTAssertEqual(account.refreshToken, "refresh-token") + let mergedExpiresAt = try XCTUnwrap(account.expiresAt) + XCTAssertEqual(mergedExpiresAt.timeIntervalSince1970, expiresAt.timeIntervalSince1970, accuracy: 0.01) + XCTAssertEqual(account.idToken, "id-token") + } + func testCodexProviderUsesChatGPTAccountIDForCodexLBInExternalMode() { let provider = CodexProvider() let account = OpenAIAuthAccount( @@ -416,6 +578,46 @@ final class TokenManagerTests: XCTestCase { XCTAssertEqual(account.credentialType, .oauthBearer) } + func testShouldIncludeCodexLBAccountSkipsInactiveStatus() { + let activeAccount = CodexLBEncryptedAccount( + accountId: "active-id", + chatGPTAccountId: "active-chatgpt-id", + email: "active@example.com", + planType: "plus", + status: "active", + accessTokenEncrypted: Data([0x01]), + refreshTokenEncrypted: nil, + idTokenEncrypted: nil, + lastRefresh: nil + ) + let inactiveAccount = CodexLBEncryptedAccount( + accountId: "inactive-id", + chatGPTAccountId: "inactive-chatgpt-id", + email: "inactive@example.com", + planType: "plus", + status: "deactivated", + accessTokenEncrypted: Data([0x01]), + refreshTokenEncrypted: nil, + idTokenEncrypted: nil, + lastRefresh: nil + ) + let legacyAccount = CodexLBEncryptedAccount( + accountId: "legacy-id", + chatGPTAccountId: "legacy-chatgpt-id", + email: "legacy@example.com", + planType: nil, + status: nil, + accessTokenEncrypted: Data([0x01]), + refreshTokenEncrypted: nil, + idTokenEncrypted: nil, + lastRefresh: nil + ) + + XCTAssertTrue(TokenManager.shared.shouldIncludeCodexLBAccount(activeAccount)) + XCTAssertFalse(TokenManager.shared.shouldIncludeCodexLBAccount(inactiveAccount)) + XCTAssertTrue(TokenManager.shared.shouldIncludeCodexLBAccount(legacyAccount)) + } + func testCodexProviderKeepsDefaultAccountIDInDirectMode() { let provider = CodexProvider() let account = OpenAIAuthAccount(