From 9609c9e14bd7d3d7f930960c059eda670da7ad2e Mon Sep 17 00:00:00 2001 From: Erfan Seyyedi Date: Mon, 27 Jul 2026 23:47:46 +0800 Subject: [PATCH 1/2] Keep limits available when usage history fails --- Sources/CodexLimits/CodexClient.swift | 91 +++++++++++-------- Sources/CodexLimits/CodexClientError.swift | 28 ++++++ Tests/CodexLimitsTests/CodexClientTests.swift | 65 +++++++++++++ 3 files changed, 148 insertions(+), 36 deletions(-) create mode 100644 Sources/CodexLimits/CodexClientError.swift 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}}}"# } From 0725253b1dd2f117f116be5a328c9980f8fdbd0f Mon Sep 17 00:00:00 2001 From: Erfan Seyyedi Date: Tue, 28 Jul 2026 09:59:41 +0800 Subject: [PATCH 2/2] Validate JSON-RPC error envelopes --- Sources/CodexLimits/CodexClient.swift | 52 +++++---- Tests/CodexLimitsTests/CodexClientTests.swift | 106 +++++++++++++----- 2 files changed, 107 insertions(+), 51 deletions(-) diff --git a/Sources/CodexLimits/CodexClient.swift b/Sources/CodexLimits/CodexClient.swift index d4c9196..b803127 100644 --- a/Sources/CodexLimits/CodexClient.swift +++ b/Sources/CodexLimits/CodexClient.swift @@ -26,7 +26,7 @@ enum CodexClient { do { let version = Bundle.main.infoDictionary?["CFBundleShortVersionString"] as? String ?? "0.1.0" try write( - #"{"id":1,"method":"initialize","params":{"clientInfo":{"name":"codex-limits","title":"Codex Limits","version":"\#(version)"},"capabilities":{"experimentalApi":true}}}"#, + #"{"id":\#(RequestID.initialize.rawValue),"method":"initialize","params":{"clientInfo":{"name":"codex-limits","title":"Codex Limits","version":"\#(version)"},"capabilities":{"experimentalApi":true}}}"#, to: input.fileHandleForWriting ) let fetchedAt = Date() @@ -170,32 +170,31 @@ enum CodexClient { let rawID = object["id"] as? Int, let id = RequestID(rawValue: rawID) else { continue } - if object["error"] != nil { + if object.keys.contains("error") { + guard (try? JSONDecoder().decode(RPCErrorEnvelope.self, from: data)) != nil else { + throw CodexClientError.invalidResponse + } switch id { - case .rateLimits: + case .initialize, .rateLimits: throw CodexClientError.invalidResponse case .usage: usageRequestFinished = true - default: - continue } - } - - switch id { - case .initialize: - try write(#"{"method":"initialized"}"#, to: input) - 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 .usage: - if object["error"] == nil { + } else { + switch id { + case .initialize: + try write(#"{"method":"initialized"}"#, to: input) + 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 .usage: usageResponse = data usageRequestFinished = true } @@ -219,6 +218,15 @@ private enum RequestID: Int { case usage = 3 } +private struct RPCErrorEnvelope: Decodable { + let error: RPCError +} + +private struct RPCError: Decodable { + let code: Double + let message: String +} + private struct RPCResponse: Decodable { let result: Result? } diff --git a/Tests/CodexLimitsTests/CodexClientTests.swift b/Tests/CodexLimitsTests/CodexClientTests.swift index e10f1c7..7e14107 100644 --- a/Tests/CodexLimitsTests/CodexClientTests.swift +++ b/Tests/CodexLimitsTests/CodexClientTests.swift @@ -39,19 +39,10 @@ final class CodexClientTests: XCTestCase { } 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 fetchedAt = Self.fetchedAt - let result = try await CodexClient.readSnapshot( - from: serverOutput.fileHandleForReading, - writingTo: clientInput.fileHandleForWriting, + let result = try await readSnapshot( + responses: [Self.usageErrorResponse, Self.rateLimitsResponse], fetchedAt: fetchedAt ) @@ -60,26 +51,41 @@ final class CodexClientTests: XCTestCase { XCTAssertEqual(result.fetchedAt, fetchedAt) } + func testUsageRPCErrorAfterRateLimitsDoesNotDiscardRateLimits() async throws { + let result = try await readSnapshot( + responses: [Self.rateLimitsResponse, Self.usageErrorResponse], + fetchedAt: Self.fetchedAt + ) + + XCTAssertEqual(result.mainLimit.window.remainingPercent, 80) + XCTAssertEqual(result.tokenHistory, []) + XCTAssertEqual(result.fetchedAt, Self.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) + try await assertInvalidResponse(responses: [rateLimitsError]) + } + + func testInitializationRPCErrorRemainsFailure() async throws { + let initializationError = #"{"id":1,"error":{"code":-32603,"message":"Initialization failed"}}"# + + try await assertInvalidResponse(responses: [initializationError]) + } + + func testMalformedUsageRPCErrorRemainsFailure() async throws { + let malformedErrors = [ + #"{"id":3,"error":null}"#, + #"{"id":3,"error":"bad envelope"}"#, + #"{"id":3,"error":{"message":"Missing code"}}"#, + #"{"id":3,"error":{"code":-32603}}"# + ] + + for malformedError in malformedErrors { + try await assertInvalidResponse( + responses: [Self.rateLimitsResponse, malformedError] ) - 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)") - } } } @@ -96,10 +102,52 @@ final class CodexClientTests: XCTestCase { try CodexClient.decode( rateLimitsResponse: Data(Self.rateLimitsResponse.utf8), usageResponse: usage, - fetchedAt: Date(timeIntervalSince1970: 1_900_000) + fetchedAt: Self.fetchedAt ) ) } + private func readSnapshot( + responses: [String], + fetchedAt: Date + ) async throws -> UsageSnapshot { + let serverOutput = Pipe() + let clientInput = Pipe() + try serverOutput.fileHandleForWriting.write( + contentsOf: Data((responses.joined(separator: "\n") + "\n").utf8) + ) + try serverOutput.fileHandleForWriting.close() + + return try await CodexClient.readSnapshot( + from: serverOutput.fileHandleForReading, + writingTo: clientInput.fileHandleForWriting, + fetchedAt: fetchedAt + ) + } + + private func assertInvalidResponse( + responses: [String], + file: StaticString = #filePath, + line: UInt = #line + ) async throws { + do { + _ = try await readSnapshot( + responses: responses, + fetchedAt: Self.fetchedAt + ) + XCTFail("Expected invalidResponse", file: file, line: line) + } catch let error as CodexClientError { + guard case .invalidResponse = error else { + return XCTFail( + "Expected invalidResponse, got \(error)", + file: file, + line: line + ) + } + } + } + + private static let fetchedAt = Date(timeIntervalSince1970: 1_900_000) + private static let usageErrorResponse = #"{"id":3,"error":{"code":-32603,"message":"Usage is temporarily unavailable"}}"# 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}}}"# }