diff --git a/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift b/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift index 45c2d47..10ed124 100644 --- a/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift +++ b/CopilotMonitor/CopilotMonitor/Providers/CodexProvider.swift @@ -302,21 +302,25 @@ final class CodexProvider: ProviderProtocol { } private func fetchUsageForAccount(_ account: OpenAIAuthAccount) async throws -> CodexAccountCandidate { - let endpoint = "https://chatgpt.com/backend-api/wham/usage" - guard let url = URL(string: endpoint) else { - logger.error("Invalid Codex API endpoint URL") - throw ProviderError.networkError("Invalid endpoint URL") - } + let endpointConfiguration = TokenManager.shared.getCodexEndpointConfiguration() + let url = try codexUsageURL(for: endpointConfiguration) var request = URLRequest(url: url) request.httpMethod = "GET" request.setValue("Bearer \(account.accessToken)", forHTTPHeaderField: "Authorization") - if let accountId = account.accountId, !accountId.isEmpty { + let requestAccountId = codexRequestAccountID(for: account, endpointMode: endpointConfiguration.mode) + if let accountId = requestAccountId, !accountId.isEmpty { request.setValue(accountId, forHTTPHeaderField: "ChatGPT-Account-Id") } else { - logger.warning("Codex account ID missing for \(account.authSource), sending request without account header") + logger.warning( + "Codex account ID missing for \(account.authSource, privacy: .public) using endpoint source \(endpointConfiguration.source, privacy: .public); sending request without account header" + ) } + logger.debug( + "Codex endpoint resolved: url=\(url.absoluteString, privacy: .public), source=\(endpointConfiguration.source, privacy: .public), external_mode=\(self.isExternalEndpointMode(endpointConfiguration.mode) ? "YES" : "NO"), account_header=\(requestAccountId != nil ? "YES" : "NO")" + ) + let (data, response) = try await URLSession.shared.data(for: request) guard let httpResponse = response as? HTTPURLResponse else { @@ -437,6 +441,38 @@ final class CodexProvider: ProviderProtocol { ) } + func codexUsageURL(for configuration: CodexEndpointConfiguration) throws -> URL { + switch configuration.mode { + case .directChatGPT: + guard let url = URL(string: "https://chatgpt.com/backend-api/wham/usage") else { + logger.error("Default Codex usage URL is invalid; aborting request") + throw ProviderError.providerError("Default Codex usage URL is invalid") + } + return url + case .external(let usageURL): + return usageURL + } + } + + func codexRequestAccountID(for account: OpenAIAuthAccount, endpointMode: CodexEndpointMode) -> String? { + switch endpointMode { + case .directChatGPT: + return account.accountId + case .external: + if account.source == .codexLB { + return account.externalUsageAccountId ?? account.accountId + } + return account.accountId + } + } + + func isExternalEndpointMode(_ mode: CodexEndpointMode) -> Bool { + if case .external = mode { + return true + } + return false + } + private func isSameUsage(_ lhs: CodexAccountCandidate, _ rhs: CodexAccountCandidate) -> Bool { let primaryMatch = lhs.details.dailyUsage == rhs.details.dailyUsage let secondaryMatch = lhs.details.secondaryUsage == rhs.details.secondaryUsage diff --git a/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift b/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift index 6ee58f5..674607b 100644 --- a/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift +++ b/CopilotMonitor/CopilotMonitor/Services/TokenManager.swift @@ -305,8 +305,9 @@ struct CodexAuth: Codable { } /// codex-lb row payload from store.db/accounts table -private struct CodexLBEncryptedAccount { +struct CodexLBEncryptedAccount { let accountId: String? + let chatGPTAccountId: String? let email: String? let planType: String? let status: String? @@ -327,12 +328,23 @@ enum OpenAIAuthSource { struct OpenAIAuthAccount { let accessToken: String let accountId: String? + let externalUsageAccountId: String? let email: String? let authSource: String let sourceLabels: [String] let source: OpenAIAuthSource } +enum CodexEndpointMode: Equatable { + case directChatGPT + case external(usageURL: URL) +} + +struct CodexEndpointConfiguration: Equatable { + let mode: CodexEndpointMode + let source: String +} + /// Auth source types for Claude account discovery enum ClaudeAuthSource { case opencodeAuth @@ -902,6 +914,74 @@ final class TokenManager: @unchecked Sendable { return current as? String } + func codexEndpointConfiguration( + from configDictionary: [String: Any]?, + sourcePath: String? = nil + ) -> CodexEndpointConfiguration { + let defaultConfiguration = CodexEndpointConfiguration( + mode: .directChatGPT, + source: "Default ChatGPT usage endpoint" + ) + + guard let configDictionary else { + return defaultConfiguration + } + + if let explicitUsageURL = resolveConfigValue( + nestedString(in: configDictionary, path: ["opencode-bar", "codex", "usageURL"]) + ) { + if let resolvedURL = URL(string: explicitUsageURL), + let scheme = resolvedURL.scheme?.lowercased(), + scheme == "http" || scheme == "https", + resolvedURL.host != nil { + return CodexEndpointConfiguration( + mode: .external(usageURL: resolvedURL), + source: sourcePath ?? "opencode-bar.codex.usageURL" + ) + } + + logger.warning( + "Ignoring invalid Codex usage URL override '\(explicitUsageURL, privacy: .public)' from \(sourcePath ?? "opencode-bar.codex.usageURL", privacy: .public)" + ) + } + + if let baseURLString = resolveConfigValue( + nestedString(in: configDictionary, path: ["provider", "openai", "options", "baseURL"]) + ) { + if let baseURL = URL(string: baseURLString), + let scheme = baseURL.scheme?.lowercased(), + scheme == "http" || scheme == "https", + let host = baseURL.host { + var components = URLComponents() + components.scheme = scheme + components.host = host + components.port = baseURL.port + components.path = "/api/codex/usage" + + if let usageURL = components.url { + return CodexEndpointConfiguration( + mode: .external(usageURL: usageURL), + source: sourcePath ?? "provider.openai.options.baseURL" + ) + } + } + + logger.warning( + "Ignoring invalid OpenAI base URL '\(baseURLString, privacy: .public)' from \(sourcePath ?? "provider.openai.options.baseURL", privacy: .public)" + ) + } + + return defaultConfiguration + } + + func getCodexEndpointConfiguration() -> CodexEndpointConfiguration { + let config = readOpenCodeConfigJSON() + return codexEndpointConfiguration( + from: config, + sourcePath: lastFoundOpenCodeConfigPath?.path + ) + } + private struct SearchAPIKeyLookupSource { let dictionary: [String: Any]? let sourcePath: String? @@ -1172,6 +1252,7 @@ final class TokenManager: @unchecked Sendable { """ SELECT id, + chatgpt_account_id, email, plan_type, status, @@ -1184,6 +1265,33 @@ final class TokenManager: @unchecked Sendable { """ SELECT id, + chatgpt_account_id, + email, + NULL AS plan_type, + NULL AS status, + access_token_encrypted, + NULL AS refresh_token_encrypted, + NULL AS id_token_encrypted, + NULL AS last_refresh + FROM accounts + """, + """ + SELECT + id, + NULL AS chatgpt_account_id, + email, + plan_type, + status, + access_token_encrypted, + refresh_token_encrypted, + id_token_encrypted, + last_refresh + FROM accounts + """, + """ + SELECT + id, + NULL AS chatgpt_account_id, email, NULL AS plan_type, NULL AS status, @@ -1222,19 +1330,20 @@ final class TokenManager: @unchecked Sendable { throw CodexLBError.sqliteStepFailed(stepStatus) } - guard let accessTokenEncrypted = sqliteColumnData(statement, index: 4), !accessTokenEncrypted.isEmpty else { + guard let accessTokenEncrypted = sqliteColumnData(statement, index: 5), !accessTokenEncrypted.isEmpty else { continue } let account = CodexLBEncryptedAccount( accountId: sqliteColumnString(statement, index: 0), - email: sqliteColumnString(statement, index: 1), - planType: sqliteColumnString(statement, index: 2), - status: sqliteColumnString(statement, index: 3), + chatGPTAccountId: sqliteColumnString(statement, index: 1), + email: sqliteColumnString(statement, index: 2), + planType: sqliteColumnString(statement, index: 3), + status: sqliteColumnString(statement, index: 4), accessTokenEncrypted: accessTokenEncrypted, - refreshTokenEncrypted: sqliteColumnData(statement, index: 5), - idTokenEncrypted: sqliteColumnData(statement, index: 6), - lastRefresh: sqliteColumnString(statement, index: 7) + refreshTokenEncrypted: sqliteColumnData(statement, index: 6), + idTokenEncrypted: sqliteColumnData(statement, index: 7), + lastRefresh: sqliteColumnString(statement, index: 8) ) accounts.append(account) } @@ -1373,6 +1482,29 @@ final class TokenManager: @unchecked Sendable { return token } + func makeCodexLBOpenAIAccount( + from encryptedAccount: CodexLBEncryptedAccount, + accessToken: String, + authSourcePath: String + ) -> OpenAIAuthAccount { + let normalizedAccountId = encryptedAccount.accountId? + .trimmingCharacters(in: .whitespacesAndNewlines) + let normalizedChatGPTAccountId = encryptedAccount.chatGPTAccountId? + .trimmingCharacters(in: .whitespacesAndNewlines) + let normalizedEmail = encryptedAccount.email? + .trimmingCharacters(in: .whitespacesAndNewlines) + + return OpenAIAuthAccount( + accessToken: accessToken, + accountId: normalizedAccountId?.isEmpty == true ? nil : normalizedAccountId, + externalUsageAccountId: normalizedChatGPTAccountId?.isEmpty == true ? nil : normalizedChatGPTAccountId, + email: normalizedEmail?.isEmpty == true ? nil : normalizedEmail, + authSource: authSourcePath, + sourceLabels: [openAISourceLabel(for: .codexLB)], + source: .codexLB + ) + } + func readCodexLBOpenAIAccounts() -> [OpenAIAuthAccount] { return queue.sync { if let cached = cachedCodexLBAccounts, @@ -1410,25 +1542,20 @@ final class TokenManager: @unchecked Sendable { encryptedAccount.accessTokenEncrypted, key: fernetKey ) - let normalizedAccountId = encryptedAccount.accountId? - .trimmingCharacters(in: .whitespacesAndNewlines) - let normalizedEmail = encryptedAccount.email? - .trimmingCharacters(in: .whitespacesAndNewlines) decodedAccounts.append( - OpenAIAuthAccount( + makeCodexLBOpenAIAccount( + from: encryptedAccount, accessToken: accessToken, - accountId: normalizedAccountId?.isEmpty == true ? nil : normalizedAccountId, - email: normalizedEmail?.isEmpty == true ? nil : normalizedEmail, - authSource: candidate.databaseURL.path, - sourceLabels: [openAISourceLabel(for: .codexLB)], - source: .codexLB + authSourcePath: candidate.databaseURL.path ) ) // PII fields (email, account ID) kept at debug level to avoid // leaking personal info in production console logs. logger.debug( """ - codex-lb account loaded: id=\(normalizedAccountId ?? "nil"), \ + \(candidate.databaseURL.path, privacy: .public) codex-lb account loaded: \ + id=\(encryptedAccount.accountId?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "nil"), \ + chatgpt_account_id=\(encryptedAccount.chatGPTAccountId?.trimmingCharacters(in: .whitespacesAndNewlines) ?? "nil"), \ email=\(encryptedAccount.email ?? "unknown"), \ plan=\(encryptedAccount.planType ?? "unknown"), \ status=\(encryptedAccount.status ?? "unknown"), \ @@ -1607,6 +1734,7 @@ final class TokenManager: @unchecked Sendable { return OpenAIAuthAccount( accessToken: primary.accessToken, accountId: (primaryAccountId?.isEmpty == false) ? primaryAccountId : fallbackAccountId, + externalUsageAccountId: normalizedNonEmpty(primary.externalUsageAccountId) ?? normalizedNonEmpty(fallback.externalUsageAccountId), email: (primaryEmail?.isEmpty == false) ? primaryEmail : fallbackEmail, authSource: primary.authSource, sourceLabels: mergedSourceLabels, @@ -2646,6 +2774,7 @@ final class TokenManager: @unchecked Sendable { OpenAIAuthAccount( accessToken: access, accountId: auth.openai?.accountId, + externalUsageAccountId: nil, email: nil, authSource: authSource, sourceLabels: [openAISourceLabel(for: .opencodeAuth)], @@ -2673,6 +2802,7 @@ final class TokenManager: @unchecked Sendable { OpenAIAuthAccount( accessToken: access, accountId: codexAuth.tokens?.accountId, + externalUsageAccountId: nil, email: codexEmail, authSource: authSource, sourceLabels: [openAISourceLabel(for: .codexAuth)], diff --git a/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift b/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift index a72f414..be7fc69 100644 --- a/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift +++ b/CopilotMonitor/CopilotMonitorTests/TokenManagerTests.swift @@ -3,6 +3,105 @@ import XCTest final class TokenManagerTests: XCTestCase { + func testCodexEndpointConfigurationDefaultsToChatGPT() { + let configuration = TokenManager.shared.codexEndpointConfiguration(from: nil) + + XCTAssertEqual(configuration, CodexEndpointConfiguration( + mode: .directChatGPT, + source: "Default ChatGPT usage endpoint" + )) + } + + func testCodexEndpointConfigurationDerivesExternalUsageURLFromBaseURL() { + let configuration = TokenManager.shared.codexEndpointConfiguration( + from: [ + "provider": [ + "openai": [ + "options": [ + "baseURL": "https://codex.2631.eu/v1" + ] + ] + ] + ], + sourcePath: "/tmp/opencode.json" + ) + + XCTAssertEqual(configuration, CodexEndpointConfiguration( + mode: .external(usageURL: URL(string: "https://codex.2631.eu/api/codex/usage")!), + source: "/tmp/opencode.json" + )) + } + + func testCodexEndpointConfigurationPrefersExplicitUsageOverride() { + let configuration = TokenManager.shared.codexEndpointConfiguration( + from: [ + "provider": [ + "openai": [ + "options": [ + "baseURL": "https://codex.2631.eu/v1" + ] + ] + ], + "opencode-bar": [ + "codex": [ + "usageURL": "https://custom.example.com/api/codex/usage" + ] + ] + ], + sourcePath: "/tmp/opencode.json" + ) + + XCTAssertEqual(configuration, CodexEndpointConfiguration( + mode: .external(usageURL: URL(string: "https://custom.example.com/api/codex/usage")!), + source: "/tmp/opencode.json" + )) + } + + func testCodexEndpointConfigurationIgnoresMalformedUsageOverrideAndFallsBackToBaseURL() { + let configuration = TokenManager.shared.codexEndpointConfiguration( + from: [ + "provider": [ + "openai": [ + "options": [ + "baseURL": "https://codex.2631.eu/v1" + ] + ] + ], + "opencode-bar": [ + "codex": [ + "usageURL": "://bad-url" + ] + ] + ], + sourcePath: "/tmp/opencode.json" + ) + + XCTAssertEqual(configuration, CodexEndpointConfiguration( + mode: .external(usageURL: URL(string: "https://codex.2631.eu/api/codex/usage")!), + source: "/tmp/opencode.json" + )) + } + + func testCodexEndpointConfigurationFallsBackToDefaultWhenConfigIsMalformed() { + let configuration = TokenManager.shared.codexEndpointConfiguration( + from: [ + "provider": [ + "openai": [ + "options": [ + "baseURL": "://bad-url" + ] + ] + ] + ], + sourcePath: "/tmp/opencode.json" + ) + + XCTAssertEqual(configuration, CodexEndpointConfiguration( + mode: .directChatGPT, + source: "Default ChatGPT usage endpoint" + )) + } + func testReadClaudeAnthropicAuthFilesIncludesDisabledAccounts() throws { let fileManager = FileManager.default let tempDirectory = fileManager.temporaryDirectory @@ -65,4 +164,103 @@ final class TokenManagerTests: XCTestCase { let expiresAt = try XCTUnwrap(primaryAccount.expiresAt) XCTAssertEqual(expiresAt.timeIntervalSince1970, 1_770_563_557.15, accuracy: 0.01) } + + func testCodexProviderUsesChatGPTAccountIDForCodexLBInExternalMode() { + let provider = CodexProvider() + let account = OpenAIAuthAccount( + accessToken: "token", + accountId: "codex-lb-internal-id", + externalUsageAccountId: "chatgpt-account-id", + email: "user@example.com", + authSource: "codex-lb", + sourceLabels: ["Codex LB"], + source: .codexLB + ) + + let accountID = provider.codexRequestAccountID( + for: account, + endpointMode: .external(usageURL: URL(string: "https://codex.example.com/api/codex/usage")!) + ) + + XCTAssertEqual(accountID, "chatgpt-account-id") + } + + func testMakeCodexLBOpenAIAccountMapsChatGPTAccountIDToExternalUsageAccountID() { + let encryptedAccount = CodexLBEncryptedAccount( + accountId: "internal-id", + chatGPTAccountId: "chatgpt-id", + email: "user@example.com", + planType: "plus", + status: "active", + accessTokenEncrypted: Data([0x01]), + refreshTokenEncrypted: nil, + idTokenEncrypted: nil, + lastRefresh: "2026-03-22T10:00:00Z" + ) + + let account = TokenManager.shared.makeCodexLBOpenAIAccount( + from: encryptedAccount, + accessToken: "token", + authSourcePath: "/tmp/store.db" + ) + + XCTAssertEqual(account.accountId, "internal-id") + XCTAssertEqual(account.externalUsageAccountId, "chatgpt-id") + XCTAssertEqual(account.email, "user@example.com") + XCTAssertEqual(account.source, .codexLB) + } + + func testCodexProviderKeepsDefaultAccountIDInDirectMode() { + let provider = CodexProvider() + let account = OpenAIAuthAccount( + accessToken: "token", + accountId: "direct-account-id", + externalUsageAccountId: "chatgpt-account-id", + email: "user@example.com", + authSource: "codex-lb", + sourceLabels: ["Codex LB"], + source: .codexLB + ) + + let accountID = provider.codexRequestAccountID( + for: account, + endpointMode: .directChatGPT + ) + + XCTAssertEqual(accountID, "direct-account-id") + } + + func testCodexProviderKeepsRegularAccountIDForNonCodexLBExternalMode() { + let provider = CodexProvider() + let account = OpenAIAuthAccount( + accessToken: "token", + accountId: "openai-account-id", + externalUsageAccountId: nil, + email: "user@example.com", + authSource: "opencode-auth", + sourceLabels: ["OpenCode"], + source: .opencodeAuth + ) + + let accountID = provider.codexRequestAccountID( + for: account, + endpointMode: .external(usageURL: URL(string: "https://codex.example.com/api/codex/usage")!) + ) + + XCTAssertEqual(accountID, "openai-account-id") + } + + func testCodexProviderDoesNotInventExternalUsageIDForNonCodexSources() { + let account = OpenAIAuthAccount( + accessToken: "token", + accountId: "openai-account-id", + externalUsageAccountId: nil, + email: nil, + authSource: "opencode-auth", + sourceLabels: ["OpenCode"], + source: .opencodeAuth + ) + + XCTAssertNil(account.externalUsageAccountId) + } }