diff --git a/Sources/CodexLimits/CodexClient.swift b/Sources/CodexLimits/CodexClient.swift index 31a0683..d4c9196 100644 --- a/Sources/CodexLimits/CodexClient.swift +++ b/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.self, from: rateLimitsResponse).result, - let usageResult = try decoder.decode(RPCResponse.self, from: usageResponse).result else { + guard let rateResult = try decoder.decode( + RPCResponse.self, + from: rateLimitsResponse + ).result else { throw CodexClientError.invalidResponse } + let usageResult: UsageResult? + if let usageResponse { + guard let result = try decoder.decode( + RPCResponse.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: Decodable { let result: Result? } diff --git a/Sources/CodexLimits/CodexClientError.swift b/Sources/CodexLimits/CodexClientError.swift new file mode 100644 index 0000000..57872cb --- /dev/null +++ b/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." + } + } +} diff --git a/Tests/CodexLimitsTests/CodexClientTests.swift b/Tests/CodexLimitsTests/CodexClientTests.swift index e602e3f..e10f1c7 100644 --- a/Tests/CodexLimitsTests/CodexClientTests.swift +++ b/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}}}"# }