Keep limits available when usage history fails

Sources/CodexLimits/CodexClient.swift

@@ -1,25 +1,5 @@
 import Foundation
 
-enum CodexClientError: LocalizedError {
-    case cliNotFound
-    case invalidResponse
-    case mainLimitMissing
-    case timedOut
-
-    var errorDescription: String? {
-        switch self {
-        case .cliNotFound:
-            "Codex CLI was not found. Install it with Homebrew, sign in, and try again."
-        case .invalidResponse:
-            "Codex returned data this app could not read. Update Codex CLI and try again."
-        case .mainLimitMissing:
-            "Codex did not return a usable limit. Make sure Codex CLI is signed in."
-        case .timedOut:
-            "Codex took too long to respond. Try refreshing again."
-        }
-    }
-}
-
 enum CodexClient {
     private static let executablePaths = [
         "/opt/homebrew/bin/codex",
@@ -80,14 +60,28 @@ enum CodexClient {
 
     static func decode(
         rateLimitsResponse: Data,
-        usageResponse: Data,
+        usageResponse: Data?,
         fetchedAt: Date
     ) throws -> UsageSnapshot {
         let decoder = JSONDecoder()
-        guard let rateResult = try decoder.decode(RPCResponse<RateLimitsResult>.self, from: rateLimitsResponse).result,
-              let usageResult = try decoder.decode(RPCResponse<UsageResult>.self, from: usageResponse).result else {
+        guard let rateResult = try decoder.decode(
+            RPCResponse<RateLimitsResult>.self,
+            from: rateLimitsResponse
+        ).result else {
             throw CodexClientError.invalidResponse
         }
+        let usageResult: UsageResult?
+        if let usageResponse {
+            guard let result = try decoder.decode(
+                RPCResponse<UsageResult>.self,
+                from: usageResponse
+            ).result else {
+                throw CodexClientError.invalidResponse
+            }
+            usageResult = result
+        } else {
+            usageResult = nil
+        }
 
         let snapshots = rateResult.rateLimitsByLimitId ?? ["codex": rateResult.rateLimits]
         let mainSnapshot = snapshots["codex"] ?? rateResult.rateLimits
@@ -123,7 +117,7 @@ enum CodexClient {
         dateFormatter.calendar = Calendar(identifier: .gregorian)
         dateFormatter.timeZone = TimeZone(secondsFromGMT: 0)
         dateFormatter.dateFormat = "yyyy-MM-dd"
-        let tokenHistory = (usageResult.dailyUsageBuckets ?? []).compactMap { bucket -> TokenDay? in
+        let tokenHistory = (usageResult?.dailyUsageBuckets ?? []).compactMap { bucket -> TokenDay? in
             guard let date = dateFormatter.date(from: bucket.startDate) else { return nil }
             return TokenDay(date: date, tokens: bucket.tokens)
         }
@@ -160,35 +154,54 @@ enum CodexClient {
         try handle.write(contentsOf: Data((message + "\n").utf8))
     }
 
-    private static func readSnapshot(
+    static func readSnapshot(
         from output: FileHandle,
         writingTo input: FileHandle,
         fetchedAt: Date
     ) async throws -> UsageSnapshot {
         var rateLimitsResponse: Data?
         var usageResponse: Data?
+        var usageRequestFinished = false
 
         for try await line in output.bytes.lines {
             try Task.checkCancellation()
             let data = Data(line.utf8)
             guard let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
-                  let id = object["id"] as? Int else { continue }
+                  let rawID = object["id"] as? Int,
+                  let id = RequestID(rawValue: rawID) else { continue }
+
+            if object["error"] != nil {
+                switch id {
+                case .rateLimits:
+                    throw CodexClientError.invalidResponse
+                case .usage:
+                    usageRequestFinished = true
+                default:
+                    continue
+                }
+            }
 
-            if object["error"] != nil { throw CodexClientError.invalidResponse }
             switch id {
-            case 1:
+            case .initialize:
                 try write(#"{"method":"initialized"}"#, to: input)
-                try write(#"{"id":2,"method":"account/rateLimits/read"}"#, to: input)
-                try write(#"{"id":3,"method":"account/usage/read"}"#, to: input)
-            case 2:
+                try write(
+                    #"{"id":\#(RequestID.rateLimits.rawValue),"method":"account/rateLimits/read"}"#,
+                    to: input
+                )
+                try write(
+                    #"{"id":\#(RequestID.usage.rawValue),"method":"account/usage/read"}"#,
+                    to: input
+                )
+            case .rateLimits:
                 rateLimitsResponse = data
-            case 3:
-                usageResponse = data
-            default:
-                continue
+            case .usage:
+                if object["error"] == nil {
+                    usageResponse = data
+                    usageRequestFinished = true
+                }
             }
 
-            if let rateLimitsResponse, let usageResponse {
+            if let rateLimitsResponse, usageRequestFinished {
                 return try decode(
                     rateLimitsResponse: rateLimitsResponse,
                     usageResponse: usageResponse,
@@ -200,6 +213,12 @@ enum CodexClient {
     }
 }
 
+private enum RequestID: Int {
+    case initialize = 1
+    case rateLimits = 2
+    case usage = 3
+}
+
 private struct RPCResponse<Result: Decodable>: Decodable {
     let result: Result?
 }

Sources/CodexLimits/CodexClientError.swift

@@ -0,0 +1,28 @@
+//
+//  CodexClientError.swift
+//  CodexLimits
+//
+//  Created by Erfan on 27/7/26.
+//
+
+import Foundation
+
+enum CodexClientError: LocalizedError {
+    case cliNotFound
+    case invalidResponse
+    case mainLimitMissing
+    case timedOut
+
+    var errorDescription: String? {
+        switch self {
+        case .cliNotFound:
+            "Codex CLI was not found. Install it with Homebrew, sign in, and try again."
+        case .invalidResponse:
+            "Codex returned data this app could not read. Update Codex CLI and try again."
+        case .mainLimitMissing:
+            "Codex did not return a usable limit. Make sure Codex CLI is signed in."
+        case .timedOut:
+            "Codex took too long to respond. Try refreshing again."
+        }
+    }
+}

Tests/CodexLimitsTests/CodexClientTests.swift

@@ -37,4 +37,69 @@ final class CodexClientTests: XCTestCase {
         XCTAssertEqual(result.emergencyResetCount, 3)
         XCTAssertEqual(result.fetchedAt, fetchedAt)
     }
+
+    func testUsageRPCErrorDoesNotDiscardRateLimits() async throws {
+        let serverOutput = Pipe()
+        let clientInput = Pipe()
+        let rateLimits = Self.rateLimitsResponse
+        let usageError = #"{"id":3,"error":{"code":-32603,"message":"Usage is temporarily unavailable"}}"#
+        try serverOutput.fileHandleForWriting.write(
+            contentsOf: Data((usageError + "\n" + rateLimits + "\n").utf8)
+        )
+        try serverOutput.fileHandleForWriting.close()
+        let fetchedAt = Date(timeIntervalSince1970: 1_900_000)
+
+        let result = try await CodexClient.readSnapshot(
+            from: serverOutput.fileHandleForReading,
+            writingTo: clientInput.fileHandleForWriting,
+            fetchedAt: fetchedAt
+        )
+
+        XCTAssertEqual(result.mainLimit.window.remainingPercent, 80)
+        XCTAssertEqual(result.tokenHistory, [])
+        XCTAssertEqual(result.fetchedAt, fetchedAt)
+    }
+
+    func testRateLimitsRPCErrorRemainsFailure() async throws {
+        let serverOutput = Pipe()
+        let clientInput = Pipe()
+        let rateLimitsError = #"{"id":2,"error":{"code":-32603,"message":"Rate limits are temporarily unavailable"}}"#
+        try serverOutput.fileHandleForWriting.write(
+            contentsOf: Data((rateLimitsError + "\n").utf8)
+        )
+        try serverOutput.fileHandleForWriting.close()
+
+        do {
+            _ = try await CodexClient.readSnapshot(
+                from: serverOutput.fileHandleForReading,
+                writingTo: clientInput.fileHandleForWriting,
+                fetchedAt: Date(timeIntervalSince1970: 1_900_000)
+            )
+            XCTFail("Expected the rate-limits error to remain fatal")
+        } catch let error as CodexClientError {
+            guard case .invalidResponse = error else {
+                return XCTFail("Expected invalidResponse, got \(error)")
+            }
+        }
+    }
+
+    func testMalformedUsageResponseRemainsFailure() throws {
+        let usage = Data(#"""
+        {"id":3,"result":{
+          "dailyUsageBuckets":[
+            {"startDate":"2001-01-01","tokens":"not-a-number"}
+          ]
+        }}
+        """#.utf8)
+
+        XCTAssertThrowsError(
+            try CodexClient.decode(
+                rateLimitsResponse: Data(Self.rateLimitsResponse.utf8),
+                usageResponse: usage,
+                fetchedAt: Date(timeIntervalSince1970: 1_900_000)
+            )
+        )
+    }
+
+    private static let rateLimitsResponse = #"{"id":2,"result":{"rateLimits":{"limitId":"codex","primary":{"usedPercent":20,"windowDurationMins":10080,"resetsAt":2000000}},"rateLimitsByLimitId":{"codex":{"limitId":"codex","primary":{"usedPercent":20,"windowDurationMins":10080,"resetsAt":2000000}}},"rateLimitResetCredits":{"availableCount":3}}}"#
 }