diff --git a/Sources/CodexBarCore/UsageFetcher.swift b/Sources/CodexBarCore/UsageFetcher.swift index b560f005e..cee9482e7 100644 --- a/Sources/CodexBarCore/UsageFetcher.swift +++ b/Sources/CodexBarCore/UsageFetcher.swift @@ -324,6 +324,20 @@ private struct RPCCreditsSnapshot: Decodable, Encodable { let balance: String? } +private struct RPCRateLimitsErrorBody: Decodable { + let email: String? + let planType: String? + let rateLimit: CodexUsageResponse.RateLimitDetails? + let credits: CodexUsageResponse.CreditDetails? + + enum CodingKeys: String, CodingKey { + case email + case planType = "plan_type" + case rateLimit = "rate_limit" + case credits + } +} + private enum RPCWireError: Error, LocalizedError { case startFailed(String) case requestFailed(String) @@ -566,33 +580,40 @@ public struct UsageFetcher: Sendable { } private func loadRPCUsage() async throws -> UsageSnapshot { - let rpc = try CodexRPCClient(environment: self.environment) - defer { rpc.shutdown() } - - try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4") - // The app-server answers on a single stdout stream, so keep requests - // serialized to avoid starving one reader when multiple awaiters race - // for the same pipe. - let limits = try await rpc.fetchRateLimits().rateLimits - let account = try? await rpc.fetchAccount() + do { + let rpc = try CodexRPCClient(environment: self.environment) + defer { rpc.shutdown() } - let identity = ProviderIdentitySnapshot( - providerID: .codex, - accountEmail: account?.account.flatMap { details in - if case let .chatgpt(email, _) = details { email } else { nil } - }, - accountOrganization: nil, - loginMethod: account?.account.flatMap { details in - if case let .chatgpt(_, plan) = details { plan } else { nil } - }) - guard let state = CodexReconciledState.fromCLI( - primary: Self.makeWindow(from: limits.primary), - secondary: Self.makeWindow(from: limits.secondary), - identity: identity) - else { - throw UsageError.noRateLimitsFound + try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4") + // The app-server answers on a single stdout stream, so keep requests + // serialized to avoid starving one reader when multiple awaiters race + // for the same pipe. + let limits = try await rpc.fetchRateLimits().rateLimits + let account = try? await rpc.fetchAccount() + + let identity = ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: account?.account.flatMap { details in + if case let .chatgpt(email, _) = details { email } else { nil } + }, + accountOrganization: nil, + loginMethod: account?.account.flatMap { details in + if case let .chatgpt(_, plan) = details { plan } else { nil } + }) + guard let state = CodexReconciledState.fromCLI( + primary: Self.makeWindow(from: limits.primary), + secondary: Self.makeWindow(from: limits.secondary), + identity: identity) + else { + throw UsageError.noRateLimitsFound + } + return state.toUsageSnapshot() + } catch { + if let snapshot = Self.recoverUsageFromRPCError(error) { + return snapshot + } + throw error } - return state.toUsageSnapshot() } private func loadTTYUsage(keepCLISessionsAlive: Bool) async throws -> UsageSnapshot { @@ -625,13 +646,20 @@ public struct UsageFetcher: Sendable { } private func loadRPCCredits() async throws -> CreditsSnapshot { - let rpc = try CodexRPCClient(environment: self.environment) - defer { rpc.shutdown() } - try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4") - let limits = try await rpc.fetchRateLimits().rateLimits - guard let credits = limits.credits else { throw UsageError.noRateLimitsFound } - let remaining = Self.parseCredits(credits.balance) - return CreditsSnapshot(remaining: remaining, events: [], updatedAt: Date()) + do { + let rpc = try CodexRPCClient(environment: self.environment) + defer { rpc.shutdown() } + try await rpc.initialize(clientName: "codexbar", clientVersion: "0.5.4") + let limits = try await rpc.fetchRateLimits().rateLimits + guard let credits = limits.credits else { throw UsageError.noRateLimitsFound } + let remaining = Self.parseCredits(credits.balance) + return CreditsSnapshot(remaining: remaining, events: [], updatedAt: Date()) + } catch { + if let credits = Self.recoverCreditsFromRPCError(error) { + return credits + } + throw error + } } private func loadTTYCredits(keepCLISessionsAlive: Bool) async throws -> CreditsSnapshot { @@ -712,6 +740,16 @@ public struct UsageFetcher: Sendable { resetDescription: resetDescription) } + private static func makeWindow(from response: CodexUsageResponse.WindowSnapshot?) -> RateWindow? { + guard let response else { return nil } + let resetsAtDate = Date(timeIntervalSince1970: TimeInterval(response.resetAt)) + return RateWindow( + usedPercent: Double(response.usedPercent), + windowMinutes: response.limitWindowSeconds / 60, + resetsAt: resetsAtDate, + resetDescription: UsageFormatter.resetDescription(from: resetsAtDate)) + } + private static func makeTTYWindow( percentLeft: Int?, windowMinutes: Int, @@ -731,6 +769,80 @@ public struct UsageFetcher: Sendable { return val } + private static func recoverUsageFromRPCError(_ error: Error) -> UsageSnapshot? { + guard let body = self.decodeRateLimitsErrorBody(from: error) else { return nil } + let identity = ProviderIdentitySnapshot( + providerID: .codex, + accountEmail: self.normalizedCodexAccountField(body.email), + accountOrganization: nil, + loginMethod: self.normalizedCodexAccountField(body.planType)) + guard let state = CodexReconciledState.fromCLI( + primary: self.makeWindow(from: body.rateLimit?.primaryWindow), + secondary: self.makeWindow(from: body.rateLimit?.secondaryWindow), + identity: identity) + else { + return nil + } + return state.toUsageSnapshot() + } + + private static func recoverCreditsFromRPCError(_ error: Error) -> CreditsSnapshot? { + guard let credits = self.decodeRateLimitsErrorBody(from: error)?.credits else { return nil } + guard let remaining = credits.balance else { return nil } + return CreditsSnapshot(remaining: remaining, events: [], updatedAt: Date()) + } + + private static func decodeRateLimitsErrorBody(from error: Error) -> RPCRateLimitsErrorBody? { + guard case let RPCWireError.requestFailed(message) = error else { return nil } + guard let json = self.extractJSONObject(after: "body=", in: message) else { return nil } + guard let data = json.data(using: .utf8) else { return nil } + return try? JSONDecoder().decode(RPCRateLimitsErrorBody.self, from: data) + } + + private static func extractJSONObject(after marker: String, in text: String) -> String? { + guard let markerRange = text.range(of: marker) else { return nil } + let suffix = text[markerRange.upperBound...] + guard let start = suffix.firstIndex(of: "{") else { return nil } + + var depth = 0 + var inString = false + var isEscaped = false + var endIndex: String.Index? + + for index in suffix[start...].indices { + let character = suffix[index] + + if inString { + if isEscaped { + isEscaped = false + } else if character == "\\" { + isEscaped = true + } else if character == "\"" { + inString = false + } + continue + } + + switch character { + case "\"": + inString = true + case "{": + depth += 1 + case "}": + depth -= 1 + if depth == 0 { + endIndex = index + break + } + default: + break + } + } + + guard let endIndex else { return nil } + return String(suffix[start...endIndex]) + } + private static func normalizedCodexAccountField(_ value: String?) -> String? { guard let value = value?.trimmingCharacters(in: .whitespacesAndNewlines), !value.isEmpty else { return nil @@ -790,6 +902,14 @@ extension UsageFetcher { return state.toUsageSnapshot() } + public static func _recoverCodexRPCUsageFromErrorForTesting(_ message: String) -> UsageSnapshot? { + self.recoverUsageFromRPCError(RPCWireError.requestFailed(message)) + } + + public static func _recoverCodexRPCCreditsFromErrorForTesting(_ message: String) -> CreditsSnapshot? { + self.recoverCreditsFromRPCError(RPCWireError.requestFailed(message)) + } + private static func makeTestingWindow( _ value: (usedPercent: Double, windowMinutes: Int, resetsAt: Int?)) -> RateWindow diff --git a/Tests/CodexBarTests/CodexUsageFetcherFallbackTests.swift b/Tests/CodexBarTests/CodexUsageFetcherFallbackTests.swift index 5b64aa4e4..70956664a 100644 --- a/Tests/CodexBarTests/CodexUsageFetcherFallbackTests.swift +++ b/Tests/CodexBarTests/CodexUsageFetcherFallbackTests.swift @@ -3,6 +3,26 @@ import Foundation import Testing struct CodexUsageFetcherFallbackTests { + @Test + func `CLI usage recovers from RPC decode mismatch body payload`() { + let snapshot = UsageFetcher._recoverCodexRPCUsageFromErrorForTesting( + Self.decodeMismatchBodyMessage) + + #expect(snapshot?.primary?.usedPercent == 4) + #expect(snapshot?.primary?.windowMinutes == 300) + #expect(snapshot?.secondary?.usedPercent == 19) + #expect(snapshot?.secondary?.windowMinutes == 10080) + #expect(snapshot?.accountEmail(for: UsageProvider.codex) == "lukeforn@gmail.com") + #expect(snapshot?.loginMethod(for: UsageProvider.codex) == "prolite") + } + + @Test + func `CLI credits recover from RPC decode mismatch body payload`() { + let credits = UsageFetcher._recoverCodexRPCCreditsFromErrorForTesting(Self.decodeMismatchBodyMessage) + + #expect(credits?.remaining == 0) + } + @Test func `CLI usage falls back from RPC decode mismatch to TTY status`() async throws { let stubCLIPath = try self.makeDecodeMismatchStubCodexCLI() @@ -28,6 +48,39 @@ struct CodexUsageFetcherFallbackTests { #expect(credits.remaining == 42) } + private static let decodeMismatchBodyMessage = """ + failed to fetch codex rate limits: Decode error for https://chatgpt.com/backend-api/wham/usage: + unknown variant `prolite`, expected one of `guest`, `free`, `go`, `plus`, `pro`; + content-type=application/json; body={ + "user_id": "user-HjRmNJhdtyqaGzIB98OrOdJw", + "account_id": "user-HjRmNJhdtyqaGzIB98OrOdJw", + "email": "lukeforn@gmail.com", + "plan_type": "prolite", + "rate_limit": { + "allowed": true, + "limit_reached": false, + "primary_window": { + "used_percent": 4, + "limit_window_seconds": 18000, + "reset_after_seconds": 8657, + "reset_at": 1776216359 + }, + "secondary_window": { + "used_percent": 19, + "limit_window_seconds": 604800, + "reset_after_seconds": 187681, + "reset_at": 1776395384 + } + }, + "credits": { + "has_credits": false, + "unlimited": false, + "overage_limit_reached": false, + "balance": "0E-10" + } + } + """ + private func makeDecodeMismatchStubCodexCLI() throws -> String { let script = """ #!/usr/bin/python3