Recover Codex CLI usage from prolite decode errors

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:

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": "[email protected]",
+          "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": "[email protected]",
+                            "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
+    }
+}