diff --git a/Sources/CodexBarCore/UsageFetcher.swift b/Sources/CodexBarCore/UsageFetcher.swift index 962dc956d..86bdc7990 100644 --- a/Sources/CodexBarCore/UsageFetcher.swift +++ b/Sources/CodexBarCore/UsageFetcher.swift @@ -316,12 +316,28 @@ private struct RPCRateLimitWindow: Decodable, Encodable { let usedPercent: Double let windowDurationMins: Int? let resetsAt: Int? + + static func from(_ window: CodexUsageResponse.WindowSnapshot?) -> Self? { + guard let window else { return nil } + return Self( + usedPercent: Double(window.usedPercent), + windowDurationMins: window.limitWindowSeconds / 60, + resetsAt: window.resetAt) + } } private struct RPCCreditsSnapshot: Decodable, Encodable { let hasCredits: Bool let unlimited: Bool let balance: String? + + static func from(_ credits: CodexUsageResponse.CreditDetails?) -> Self? { + guard let credits else { return nil } + return Self( + hasCredits: credits.hasCredits, + unlimited: credits.unlimited, + balance: credits.balance.map { String($0) }) + } } private enum RPCWireError: Error, LocalizedError { @@ -465,8 +481,15 @@ private final class CodexRPCClient: @unchecked Sendable { } func fetchRateLimits() async throws -> RPCRateLimitsResponse { - let message = try await self.request(method: "account/rateLimits/read") - return try self.decodeResult(from: message) + do { + let message = try await self.request(method: "account/rateLimits/read") + return try self.decodeResult(from: message) + } catch let RPCWireError.requestFailed(message) { + if let recovered = Self.recoverRateLimits(from: message) { + return recovered + } + throw RPCWireError.requestFailed(message) + } } func shutdown() { @@ -537,6 +560,76 @@ private final class CodexRPCClient: @unchecked Sendable { return try decoder.decode(T.self, from: data) } + private static func recoverRateLimits(from message: String) -> RPCRateLimitsResponse? { + guard let body = self.extractEmbeddedJSONBody(from: message), + let data = body.data(using: .utf8), + let usage = try? JSONDecoder().decode(CodexUsageResponse.self, from: data) + else { + return nil + } + + let primary = RPCRateLimitWindow.from(usage.rateLimit?.primaryWindow) + let secondary = RPCRateLimitWindow.from(usage.rateLimit?.secondaryWindow) + let credits = RPCCreditsSnapshot.from(usage.credits) + + guard primary != nil || secondary != nil || credits != nil else { + return nil + } + + return RPCRateLimitsResponse( + rateLimits: RPCRateLimitSnapshot( + primary: primary, + secondary: secondary, + credits: credits)) + } + + private static func extractEmbeddedJSONBody(from message: String) -> String? { + guard let markerRange = message.range(of: "body=") else { return nil } + let suffix = message[markerRange.upperBound...] + guard let jsonStart = suffix.firstIndex(where: { $0 == "{" || $0 == "[" }) else { return nil } + + var depth = 0 + var inString = false + var escaping = false + + for index in suffix[jsonStart...].indices { + let character = suffix[index] + + if escaping { + escaping = false + continue + } + + if character == "\\" { + escaping = true + continue + } + + if character == "\"" { + inString.toggle() + continue + } + + if inString { + continue + } + + switch character { + case "{", "[": + depth += 1 + case "}", "]": + depth -= 1 + if depth == 0 { + return String(suffix[jsonStart...index]) + } + default: + continue + } + } + + return nil + } + private func jsonID(_ value: Any?) -> Int? { switch value { case let int as Int: diff --git a/Tests/CodexBarTests/CodexUsageFetcherFallbackTests.swift b/Tests/CodexBarTests/CodexUsageFetcherFallbackTests.swift new file mode 100644 index 000000000..5ae5ae563 --- /dev/null +++ b/Tests/CodexBarTests/CodexUsageFetcherFallbackTests.swift @@ -0,0 +1,142 @@ +import Foundation +import Testing +@testable import CodexBarCore + +@Suite(.serialized) +struct CodexUsageFetcherFallbackTests { + @Test + func `recovers Codex CLI usage from rate limits decode error body`() async throws { + let stubCLIPath = try self.makeDecodeErrorCodexCLI() + let fetcher = UsageFetcher(environment: ["CODEX_CLI_PATH": stubCLIPath]) + + let usage = try await fetcher.loadLatestUsage() + + #expect(usage.primary?.usedPercent == 40) + #expect(usage.primary?.windowMinutes == 300) + #expect(usage.secondary?.usedPercent == 13) + #expect(usage.secondary?.windowMinutes == 10080) + } + + @Test + func `recovers Codex CLI credits from rate limits decode error body`() async throws { + let stubCLIPath = try self.makeDecodeErrorCodexCLI() + let fetcher = UsageFetcher(environment: ["CODEX_CLI_PATH": stubCLIPath]) + + let credits = try await fetcher.loadLatestCredits() + + #expect(credits.remaining == 19.5) + } + + @Test + func `ignores unrelated JSON error body when recovering rate limits`() async throws { + let stubCLIPath = try self.makeDecodeErrorCodexCLI( + embeddedBody: """ + { + "error": "unauthorized" + } + """) + let fetcher = UsageFetcher(environment: ["CODEX_CLI_PATH": stubCLIPath]) + + do { + _ = try await fetcher.loadLatestUsage() + Issue.record("Expected the original RPC failure for unrelated JSON body") + } catch let error as UsageError { + Issue.record("Expected original RPC failure, got \(error)") + } catch { + #expect(error.localizedDescription.contains("unauthorized")) + } + } + + private func makeDecodeErrorCodexCLI(embeddedBody: String? = nil) throws -> String { + let usageJSON = """ + { + "user_id": "user-123", + "account_id": "account-123", + "email": "prolite@example.com", + "plan_type": "prolite", + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": { + "used_percent": 40, + "limit_window_seconds": 18000, + "reset_after_seconds": 6263, + "reset_at": 1776214836 + }, + "secondary_window": { + "used_percent": 13, + "limit_window_seconds": 604800, + "reset_after_seconds": 575064, + "reset_at": 1776783636 + } + }, + "code_review_rate_limit": null, + "additional_rate_limits": {}, + "credits": { + "has_credits": true, + "unlimited": false, + "balance": "19.5" + } + } + """ + + let script = """ + #!/usr/bin/python3 + import json + import sys + + EMBEDDED_BODY = \(String(reflecting: embeddedBody ?? usageJSON)) + ERROR_MESSAGE = ( + "failed to fetch codex rate limits: " + "Decode error for https://chatgpt.com/backend-api/wham/usage/ " + "decoded body contained unknown variant 'prolite'; " + "content-type=application/json; body=" + EMBEDDED_BODY + ) + + if "app-server" not in sys.argv: + sys.exit(1) + + for line in sys.stdin: + if not line.strip(): + continue + + message = json.loads(line) + method = message.get("method") + if method == "initialized": + continue + + identifier = message.get("id") + if method == "initialize": + payload = {"id": identifier, "result": {}} + elif method == "account/rateLimits/read": + payload = { + "id": identifier, + "error": { + "message": ERROR_MESSAGE + } + } + elif method == "account/read": + payload = { + "id": identifier, + "result": { + "account": { + "type": "chatgpt", + "email": "prolite@example.com", + "planType": "prolite" + }, + "requiresOpenaiAuth": False + } + } + else: + payload = {"id": identifier, "result": {}} + + print(json.dumps(payload), flush=True) + """ + + let url = FileManager.default.temporaryDirectory + .appendingPathComponent("codex-decode-error-\(UUID().uuidString)", isDirectory: false) + try Data(script.utf8).write(to: url) + try FileManager.default.setAttributes([.posixPermissions: 0o755], ofItemAtPath: url.path) + return url.path + } +}