Fix Grok CLI billing: use _x.ai/billing ACP extension method

CHANGELOG.md

@@ -9,6 +9,7 @@
 
 ### Fixed
 
+- Grok: restore CLI billing through the ACP extension method, including current usage, reset, plan, and prepaid balance fields (#4595). Thanks @noisy!
 - Menu bar: merged Current selection icons show the selected user plugin's quota instead of falling back to a built-in provider (#4604). Thanks @soimy!
 - CLI: expose OpenCode Go's local daily and per-model token/cost history through `codexbar cost` (#4606). Thanks @h3x89!
 

Sources/CodexBarCore/Providers/Grok/GrokRPCClient.swift

@@ -4,7 +4,7 @@ import Foundation
 ///
 /// The protocol mirrors Codex's app-server (newline-delimited JSON-RPC 2.0 over stdin/stdout),
 /// but uses `protocolVersion`/`clientCapabilities` for the `initialize` call instead of
-/// `clientInfo`. Billing is fetched via the `x.ai/billing` extension method.
+/// `clientInfo`. Billing uses the `_x.ai/billing` ACP extension method.
 final class GrokRPCClient: @unchecked Sendable {
     private static let log = CodexBarLog.logger(LogCategories.provider(.grok))
 
@@ -77,10 +77,24 @@ final class GrokRPCClient: @unchecked Sendable {
             timeout: self.initializeTimeoutSeconds)
     }
 
-    /// Calls `x.ai/billing` and returns the decoded response.
-    func fetchBilling() async throws -> GrokBillingResponse {
-        let message = try await self.request(method: "x.ai/billing", params: [:])
-        return try self.decodeResult(from: message)
+    /// Calls the ACP billing extension and decodes current or legacy response payloads.
+    func fetchBilling() async throws -> GrokRPCBilling {
+        let message = try await self.request(method: "_x.ai/billing", params: [:])
+        guard let result = message["result"] as? [String: Any] else {
+            throw GrokRPCError.malformed("missing result field")
+        }
+        let data = try JSONSerialization.data(withJSONObject: result)
+        if let credits = try? GrokCreditsProxyFetcher.parseSnapshot(data),
+           credits.usedPercent != nil || credits.prepaidBalanceUSD != nil
+        {
+            return .credits(credits.applying(subscriptionTier: result["subscription_tier"] as? String))
+        }
+        if let billing = try? JSONDecoder().decode(GrokBillingResponse.self, from: data),
+           billing.monthlyUsedPercent != nil
+        {
+            return .legacy(billing)
+        }
+        throw GrokRPCError.malformed(GrokStatusProbe.usageUnavailableMessage)
     }
 
     func shutdown() {
@@ -101,7 +115,12 @@ final class GrokRPCClient: @unchecked Sendable {
     {
         let id = self.nextID
         self.nextID += 1
-        try self.sendRequest(id: id, method: method, params: params)
+        try self.sendPayload([
+            "jsonrpc": "2.0",
+            "id": id,
+            "method": method,
+            "params": params ?? [:],
+        ])
 
         let resolvedTimeout = timeout ?? self.requestTimeoutSeconds
         let wrapped = try await RPCRequestTimeout.run(
@@ -137,27 +156,8 @@ final class GrokRPCClient: @unchecked Sendable {
         }
     }
 
-    private func sendRequest(id: Int, method: String, params: [String: Any]?) throws {
-        let paramsValue: Any = params ?? [:]
-        let payload: [String: Any] = [
-            "jsonrpc": "2.0",
-            "id": id,
-            "method": method,
-            "params": paramsValue,
-        ]
-        try self.sendPayload(payload)
-    }
-
     private func sendPayload(_ payload: [String: Any]) throws {
-        let raw = try JSONSerialization.data(withJSONObject: payload)
-        // Foundation's JSONSerialization escapes "/" as "\/" by default. Grok's
-        // ACP server treats the escaped form as a *different* method name (it does
-        // not unescape before lookup), so `x.ai/billing` becomes "Method not found"
-        // when sent as `x.ai\/billing`. Re-encode without slash escapes to match
-        // the on-the-wire shape the grok agent expects.
-        let unescaped = String(data: raw, encoding: .utf8)?
-            .replacingOccurrences(of: "\\/", with: "/")
-        var data = unescaped.flatMap { $0.data(using: .utf8) } ?? raw
+        var data = try JSONSerialization.data(withJSONObject: payload, options: [.withoutEscapingSlashes])
         if let preview = String(data: data.prefix(200), encoding: .utf8) {
             Self.log.debug("grok rpc -> \(preview)")
         }
@@ -182,24 +182,16 @@ final class GrokRPCClient: @unchecked Sendable {
         throw GrokRPCError.malformed("grok agent stdio closed stdout")
     }
 
-    private func decodeResult<T: Decodable>(from message: [String: Any]) throws -> T {
-        guard let result = message["result"] else {
-            throw GrokRPCError.malformed("missing result field")
-        }
-        let data = try JSONSerialization.data(withJSONObject: result)
-        let decoder = JSONDecoder()
-        return try decoder.decode(T.self, from: data)
-    }
-
     private func jsonID(_ value: Any?) -> Int? {
-        switch value {
-        case let int as Int: int
-        case let number as NSNumber: number.intValue
-        default: nil
-        }
+        (value as? NSNumber)?.intValue
     }
 }
 
+enum GrokRPCBilling: Sendable {
+    case credits(GrokWebBillingSnapshot)
+    case legacy(GrokBillingResponse)
+}
+
 public enum GrokRPCError: LocalizedError, Sendable {
     case binaryNotFound
     case startFailed(String)
@@ -232,7 +224,7 @@ public enum GrokRPCError: LocalizedError, Sendable {
     }
 }
 
-/// Schema for the `x.ai/billing` result.
+/// Legacy top-level billing schema; current `_x.ai/billing` results use the nested credits parser.
 ///
 /// All monetary amounts (`monthlyLimit`, `onDemandCap`, `includedUsed`, etc.) are
 /// represented as `{ val: <cents> }` in the wire format.

Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift

@@ -158,21 +158,26 @@ public struct GrokStatusProbe: Sendable {
         let credentials = try? GrokCredentialsStore.load(env: env)
 
         var billing: GrokBillingResponse?
+        var webBilling: GrokWebBillingSnapshot?
         var rpcError: Error?
         var billingAttempted = false
         do {
             let client = try GrokRPCClient(environment: env)
             defer { client.shutdown() }
             try await client.initialize()
             billingAttempted = true
-            billing = try await client.fetchBilling()
+            switch try await client.fetchBilling() {
+            case let .legacy(response): billing = response
+            case let .credits(response): webBilling = response
+            }
         } catch {
             rpcError = error
         }
 
-        let isIdentityOnly = billing == nil && self.identityOnlyFallback(credentials, billingAttempted, rpcError)
+        let hasBilling = billing != nil || webBilling != nil
+        let isIdentityOnly = !hasBilling && self.identityOnlyFallback(credentials, billingAttempted, rpcError)
         // Terminal CLI failures must reach the provider's web fallback without scanning discarded history.
-        guard billing != nil || isIdentityOnly else {
+        guard hasBilling || isIdentityOnly else {
             throw rpcError ?? GrokRPCError.notAuthenticated
         }
 
@@ -187,11 +192,11 @@ public struct GrokStatusProbe: Sendable {
             session: self.settingsTransport)
         return GrokUsageSnapshot(
             billing: billing,
-            webBilling: nil,
+            webBilling: webBilling,
             credentials: isIdentityOnly ? credentials : Self.credentialsForSnapshot(
                 credentials: credentials,
                 billing: billing,
-                webBilling: nil),
+                webBilling: webBilling),
             localSummary: localSummary,
             cliVersion: cliVersion,
             updatedAt: Date(),

Tests/CodexBarTests/BoundedChildProcessProofTests.swift

@@ -85,7 +85,10 @@ struct BoundedChildProcessProofTests {
         defer { client.shutdown() }
 
         try await client.initialize()
-        let billing = try await client.fetchBilling()
+        guard case let .legacy(billing) = try await client.fetchBilling() else {
+            Issue.record("Expected legacy billing response")
+            return
+        }
 
         #expect(billing.monthlyLimit?.val == 100)
         #expect(billing.usage?.totalUsed?.val == 25)

Tests/CodexBarTests/GrokFailedBillingWorkTests.swift

@@ -7,9 +7,22 @@ struct GrokFailedBillingWorkTests {
         case personalUnavailable, teamUnauthorized, initializationFailure, expiredTeam, teamFallback, billingSuccess
         case expiresDuringScan, expiresDuringVersion, acceptedIdentity
         case methodUnavailableAlternateMessage, methodMessageWrongCode
+        case nestedBillingSuccess, nestedBillingPeriodFallback, prepaidOnlyBillingSuccess, unknownBillingResult
+
+        var hasCredits: Bool {
+            self == .nestedBillingSuccess || self == .nestedBillingPeriodFallback || self == .prepaidOnlyBillingSuccess
+        }
+
+        var hasBilling: Bool {
+            self == .billingSuccess || self.hasCredits
+        }
+
+        var returnsBillingResult: Bool {
+            self == .billingSuccess || self.hasCredits || self == .unknownBillingResult
+        }
 
         var acceptsSnapshot: Bool {
-            self == .teamFallback || self == .billingSuccess || self == .acceptedIdentity
+            self == .teamFallback || self.hasBilling || self == .acceptedIdentity
                 || self == .methodUnavailableAlternateMessage
         }
 
@@ -21,11 +34,35 @@ struct GrokFailedBillingWorkTests {
             switch self {
             case .teamUnauthorized: "Unauthorized"
             case .initializationFailure: "Initialize failed"
-            case .methodUnavailableAlternateMessage: "Unsupported RPC method: x.ai/billing"
+            case .methodUnavailableAlternateMessage: "Unsupported RPC method: _x.ai/billing"
             default: "Method not found"
             }
         }
 
+        var billingResult: [String: Any] {
+            if self.hasCredits {
+                var config: [String: Any] = [
+                    "creditUsagePercent": 5.0,
+                    "billingPeriodEnd": "2026-10-22T17:24:18.000000+00:00",
+                ]
+                if self == .nestedBillingSuccess {
+                    config["currentPeriod"] = [
+                        "type": "USAGE_PERIOD_TYPE_WEEKLY",
+                        "start": "2026-10-15T17:24:18.000000+00:00",
+                        "end": "2026-10-22T17:24:18.000000+00:00",
+                    ]
+                    config["billingPeriodEnd"] = "2026-11-15T17:24:18Z"
+                } else if self == .prepaidOnlyBillingSuccess {
+                    config = ["prepaidBalance": ["val": 123]]
+                }
+                return ["config": config, "subscription_tier": "SuperGrok Heavy"]
+            } else if self == .unknownBillingResult {
+                return ["subscription_tier": "SuperGrok", "config": [:]]
+            } else {
+                return ["monthlyLimit": ["val": 100], "usage": ["totalUsed": ["val": 25]]]
+            }
+        }
+
         var errorCode: Int {
             switch self {
             case .teamUnauthorized, .methodMessageWrongCode: -32001
@@ -40,7 +77,7 @@ struct GrokFailedBillingWorkTests {
         let root = FileManager.default.temporaryDirectory.appendingPathComponent("grok-work-\(UUID().uuidString)")
         try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true)
         defer { try? FileManager.default.removeItem(at: root) }
-        if scenario != .personalUnavailable, scenario != .billingSuccess {
+        if scenario != .personalUnavailable, !scenario.hasBilling {
             let auth = try JSONSerialization.data(withJSONObject: [
                 "https://auth.x.ai::fake-client": [
                     "key": "synthetic-token", "refresh_token": "synthetic-refresh",
@@ -67,8 +104,8 @@ struct GrokFailedBillingWorkTests {
             id: 1, result: [:], error: scenario == .initializationFailure ? scenario.errorMessage : nil)
         let billing = try reply(
             id: 2,
-            result: ["monthlyLimit": ["val": 100], "usage": ["totalUsed": ["val": 25]]],
-            error: scenario == .billingSuccess ? nil : scenario.errorMessage)
+            result: scenario.billingResult,
+            error: scenario.returnsBillingResult ? nil : scenario.errorMessage)
         let script = """
         #!/bin/sh
         if [ "$1" = "--version" ]; then
@@ -79,6 +116,7 @@ struct GrokFailedBillingWorkTests {
         IFS= read -r initialize_request
         printf '%s\\n' '\(initialize)'
         IFS= read -r billing_request || exit 0
+        printf '%s\\n' "$billing_request" >> "$GROK_HOME/billing-requests.jsonl"
         printf '%s\\n' '\(billing)'
         """
         try script.write(to: binary, atomically: true, encoding: .utf8)
@@ -131,22 +169,51 @@ struct GrokFailedBillingWorkTests {
             #expect(snapshot.localSummary?.daily == summary.daily)
             #expect(snapshot.toUsageSnapshot().costUsage?.last30DaysTokens == 42)
             #expect(snapshot.cliVersion == "synthetic-version")
-            if scenario != .billingSuccess {
+            if !scenario.hasBilling {
                 #expect(snapshot.billing == nil)
                 #expect(snapshot.credentials?.email == "team@example.com")
                 #expect(snapshot.diagnostic == GrokStatusProbe.teamUsageUnavailableMessage)
+            } else if scenario.hasCredits {
+                let usage = snapshot.toUsageSnapshot()
+                if scenario == .prepaidOnlyBillingSuccess {
+                    #expect(usage.primary == nil)
+                    #expect(usage.providerCost?.balance == 1.23)
+                } else {
+                    let primary = try #require(usage.primary)
+                    #expect(primary.usedPercent == 5)
+                    #expect(primary.resetsAt == ISO8601DateParser.parse("2026-10-22T17:24:18Z"))
+                }
+                #expect(usage.identity?.loginMethod == "SuperGrok Heavy")
+                #expect(snapshot.diagnostic == nil)
             } else {
                 #expect(snapshot.billing?.monthlyUsedPercent == 25)
                 #expect(snapshot.diagnostic == nil)
             }
         } catch let error as GrokRPCError {
             #expect(!scenario.acceptsSnapshot)
-            guard case let .requestFailed(message, code) = error else {
-                Issue.record("Unexpected RPC error: \(error)")
-                return
+            if scenario == .unknownBillingResult {
+                guard case .malformed = error else {
+                    Issue.record("Expected malformed billing error: \(error)")
+                    return
+                }
+            } else {
+                guard case let .requestFailed(message, code) = error else {
+                    Issue.record("Unexpected RPC error: \(error)")
+                    return
+                }
+                #expect(message == scenario.errorMessage)
+                #expect(code == scenario.errorCode)
+            }
+        }
+        if scenario != .initializationFailure {
+            let requests = try String(
+                contentsOf: root.appendingPathComponent("billing-requests.jsonl"), encoding: .utf8)
+            let methods = try requests.split(separator: "\n").map { line in
+                let request = try #require(
+                    JSONSerialization.jsonObject(with: Data(line.utf8)) as? [String: Any])
+                return try #require(request["method"] as? String)
             }
-            #expect(message == scenario.errorMessage)
-            #expect(code == scenario.errorCode)
+            #expect(methods == ["_x.ai/billing"])
         }
         #expect(await calls.scans == (scenario.scansBeforeResult ? 1 : 0))
         #expect(await calls