import Foundation import Testing @testable import CodexBarCore @Suite(.serialized) struct GrokWebBillingFetcherTests { private final class AttemptCounter: @unchecked Sendable { private let lock = NSLock() private var value = 0 func increment() -> Int { self.lock.lock() defer { self.lock.unlock() } self.value += 1 return self.value } func current() -> Int { self.lock.lock() defer { self.lock.unlock() } return self.value } } @Test func `provider exposes cli and web source modes`() { #expect(GrokProviderDescriptor.descriptor.fetchPlan.sourceModes == [.auto, .cli, .web]) } @Test func `descriptor uses Credits label for primary usage window`() { let metadata = GrokProviderDescriptor.descriptor.metadata #expect(metadata.sessionLabel == "Credits") #expect(metadata.weeklyLabel == "On-demand") #expect(!metadata.supportsOpus) } @Test func `primaryLabel derives Weekly or Monthly from resetsAt`() { let now = Date() let in6Days = now.addingTimeInterval(6 * 86400) let in30Days = now.addingTimeInterval(30 * 86400) let in90Days = now.addingTimeInterval(90 * 86400) let lateWeeklyWindow = RateWindow( usedPercent: 25, windowMinutes: 7 * 24 * 60, resetsAt: now.addingTimeInterval(86400), resetDescription: nil) #expect(GrokProviderDescriptor.primaryLabel(resetsAt: in6Days, now: now) == "Weekly") #expect(GrokProviderDescriptor.primaryLabel(resetsAt: in30Days, now: now) == "Monthly") #expect(GrokProviderDescriptor.primaryLabel(resetsAt: in90Days, now: now) == nil) #expect(GrokProviderDescriptor.primaryLabel(window: lateWeeklyWindow, now: now) == "Weekly") #expect(GrokProviderDescriptor.primaryLabel(resetsAt: nil) == nil) } @Test func `cli runtime does not import browser cookies unless explicitly enabled`() { #expect(GrokWebFetchStrategy.canImportBrowserCookies(runtime: .app, env: [:])) #expect(!GrokWebFetchStrategy.canImportBrowserCookies(runtime: .cli, env: [:])) #expect(GrokWebFetchStrategy.canImportBrowserCookies( runtime: .cli, env: ["CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT": "1"])) } @Test func `web strategy tries later browser session when first cookie is stale`() async throws { let stale = try #require(Self.cookie(name: "sso", value: "stale")) let valid = try #require(Self.cookie(name: "sso", value: "valid")) let sessions = [ GrokCookieImporter.SessionInfo(cookies: [stale], sourceLabel: "Chrome Profile 1"), GrokCookieImporter.SessionInfo(cookies: [valid], sourceLabel: "Chrome Profile 2"), ] var attemptedHeaders: [String] = [] let result = try await GrokWebFetchStrategy.fetchFirstValidCookieSession(sessions) { cookieHeader, _ in attemptedHeaders.append(cookieHeader) guard cookieHeader.contains("valid") else { throw GrokWebBillingError.requestFailed(401, "stale") } return GrokWebBillingSnapshot( usedPercent: 12, resetsAt: Date(timeIntervalSince1970: 1_800_000_000)) } #expect(attemptedHeaders == ["sso=stale", "sso=valid"]) #expect(result.0.usedPercent == 12) #expect(result.1 == "Chrome Profile 2") } @Test func `cookie authenticated web billing does not reuse auth file identity`() { #expect(GrokWebFetchStrategy.credentialsForWebBillingSnapshot( credentials: Self.credentials, authenticatedByAuthFile: false) == nil) #expect(GrokWebFetchStrategy.credentialsForWebBillingSnapshot( credentials: Self.credentials, authenticatedByAuthFile: true)? .email == "grok@example.com") } @Test func `parses grok grpc web billing frame`() throws { let reset = UInt64(1_800_000_000) let payload = Self.protobufPayload(usedPercent: 42.5, resetEpoch: reset) let data = Self.grpcFrame(payload) let snapshot = try GrokWebBillingFetcher.parseGRPCWebResponse( data, now: Date(timeIntervalSince1970: 1_799_000_000)) #expect(snapshot.usedPercent == 42.5) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: TimeInterval(reset))) } @Test func `parses unframed grok billing protobuf payload`() throws { let hex = "0a3f0d7f6a9c3f12001a002206088097f3d0062a060880b191d2063a07080215a9389b3f3a07080115d6ea183c" + "421208011206088097f3d0061a060880b191d206" let data = try #require(Self.data(hexString: hex)) let snapshot = try GrokWebBillingFetcher.parseGRPCWebResponse( data, now: Date(timeIntervalSince1970: 1_780_000_000)) #expect(snapshot.usedPercent == 1.222000002861023) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: 1_782_864_000)) } @Test func `parses unframed zero percent payload that resembles an empty grpc frame`() throws { let reset = UInt64(1_800_000_000) let payload = Self.protobufPayload(usedPercent: 0, resetEpoch: reset) #expect(GrokWebBillingFetcher.grpcWebDataFrames(from: payload).isEmpty) let snapshot = try GrokWebBillingFetcher.parseGRPCWebResponse( payload, now: Date(timeIntervalSince1970: 1_799_000_000)) #expect(snapshot.usedPercent == 0) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: TimeInterval(reset))) } @Test func `does not treat grpc frame prefix as raw protobuf`() { #expect(!GrokWebBillingFetcher.looksLikeProtobufPayload(Data([0, 0, 0, 0, 10]))) } @Test func `web strategy tries cookie plus bearer before cookie only`() async throws { let cookie = try #require(Self.cookie(name: "sso", value: "session")) let sessions = [GrokCookieImporter.SessionInfo(cookies: [cookie], sourceLabel: "Chrome")] var attempts: [String] = [] let result = try await GrokWebFetchStrategy.fetchFirstValidCookieSession( sessions, credentials: Self.credentials) { _, authCredentials in let mode = authCredentials == nil ? "cookie-only" : "cookie+bearer" attempts.append(mode) guard mode == "cookie+bearer" else { throw GrokWebBillingError.requestFailed(401, "needs bearer") } return GrokWebBillingSnapshot(usedPercent: 9, resetsAt: nil) } #expect(attempts == ["cookie+bearer"]) #expect(result.0.usedPercent == 9) } @Test func `web strategy skips expired bearer for browser cookies`() async throws { let cookie = try #require(Self.cookie(name: "sso", value: "session")) let sessions = [GrokCookieImporter.SessionInfo(cookies: [cookie], sourceLabel: "Chrome")] let expired = GrokCredentials( accessToken: "expired-token", refreshToken: nil, scope: "https://auth.x.ai::client", authMode: "oidc", userId: nil, email: nil, firstName: nil, lastName: nil, teamId: nil, oidcIssuer: nil, oidcClientId: nil, expiresAt: .distantPast, createTime: nil) var attempts: [String] = [] _ = try await GrokWebFetchStrategy.fetchFirstValidCookieSession( sessions, credentials: expired) { _, authCredentials in attempts.append(authCredentials == nil ? "cookie-only" : "cookie+bearer") return GrokWebBillingSnapshot(usedPercent: 9, resetsAt: nil) } #expect(attempts == ["cookie-only"]) } @Test func `web strategy preserves malformed auth file error`() async throws { let grokHome = FileManager.default.temporaryDirectory .appendingPathComponent("CodexBar-GrokWebBilling-\(UUID().uuidString)", isDirectory: true) try FileManager.default.createDirectory(at: grokHome, withIntermediateDirectories: true) defer { try? FileManager.default.removeItem(at: grokHome) } try Data("not-json".utf8).write(to: grokHome.appendingPathComponent("auth.json")) let browserDetection = BrowserDetection(cacheTTL: 0) let context = ProviderFetchContext( runtime: .cli, sourceMode: .web, includeCredits: true, webTimeout: 1, webDebugDumpHTML: false, verbose: false, env: ["GROK_HOME": grokHome.path], settings: nil, fetcher: UsageFetcher(), claudeFetcher: ClaudeUsageFetcher(browserDetection: browserDetection), browserDetection: browserDetection) await #expect { _ = try await GrokWebFetchStrategy().fetch(context) } throws: { error in guard case GrokCredentialsError.decodeFailed = error else { return false } return true } } @Test func `status seven scope failure is not classified as bad credentials`() { #expect(!GrokWebBillingError.isAuthenticationFailure( status: 7, message: "OAuth2 access token lacks the required billing scope")) } @Test func `ignores grpc web trailer frames`() { let payload = Self.protobufPayload(usedPercent: 12.25, resetEpoch: 1_800_000_001) let trailer = Data("grpc-status: 0\r\n".utf8) let data = Self.grpcFrame(payload) + Self.grpcFrame(trailer, flags: 0x80) let frames = GrokWebBillingFetcher.grpcWebDataFrames(from: data) #expect(frames == [payload]) } @Test func `web fetch turns grpc unauthenticated trailer into reauth guidance`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) let body = Self.grpcFrame(Data("grpc-status: 16\r\ngrpc-message: token%20expired\r\n".utf8), flags: 0x80) #expect(GrokWebBillingFetcher.grpcWebTrailerFields(from: body)["grpc-status"] == "16") GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let url = try #require(request.url) let response = HTTPURLResponse( url: url, statusCode: 200, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! return (response, body) } await #expect { _ = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) } throws: { error in error.localizedDescription.contains("grok login") } } @Test func `web fetch turns grpc unauthenticated headers into reauth guidance`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let url = try #require(request.url) let response = HTTPURLResponse( url: url, statusCode: 200, httpVersion: nil, headerFields: [ "Content-Type": "application/grpc-web+proto", "grpc-status": "16", "grpc-message": "Invalid%20bearer%20token.", ])! return (response, Data()) } await #expect { _ = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) } throws: { error in error.localizedDescription.contains("grok login") } } @Test func `web fetch turns grpc permission denied bad credentials into reauth guidance`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) let message = "The OAuth2 access token could not be validated. [WKE=unauthenticated:bad-credentials]" let encodedMessage = message.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) ?? message let body = Self.grpcFrame( Data("grpc-status: 7\r\ngrpc-message: \(encodedMessage)\r\n".utf8), flags: 0x80) GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let url = try #require(request.url) let response = HTTPURLResponse( url: url, statusCode: 200, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! return (response, body) } await #expect { _ = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) } throws: { error in error.localizedDescription.contains("grok.com") && error.localizedDescription.contains("grok login") && !error.localizedDescription.contains("status 7") } } @Test func `rejects reset only billing because it cannot render usage`() { var payload = Data() payload.append(0x10) // field 2, varint reset timestamp payload.append(contentsOf: Self.varint(1_800_000_001)) #expect { _ = try GrokWebBillingFetcher.parseGRPCWebResponse(Self.grpcFrame(payload)) } throws: { error in guard case GrokWebBillingError.parseFailed = error else { return false } return true } } @Test func `parses grok no usage yet billing response as zero percent`() throws { let data = Data([ 0x00, 0x00, 0x00, 0x00, 0x37, 0x0A, 0x35, 0x12, 0x00, 0x1A, 0x00, 0x22, 0x06, 0x08, 0x80, 0xDA, 0xCF, 0xCF, 0x06, 0x2A, 0x06, 0x08, 0x80, 0x97, 0xF3, 0xD0, 0x06, 0x32, 0x09, 0x0A, 0x05, 0x08, 0xEA, 0x0F, 0x10, 0x04, 0x12, 0x00, 0x32, 0x09, 0x0A, 0x05, 0x08, 0xEA, 0x0F, 0x10, 0x03, 0x12, 0x00, 0x32, 0x09, 0x0A, 0x05, 0x08, 0xEA, 0x0F, 0x10, 0x02, 0x12, 0x00, 0x80, 0x00, 0x00, 0x00, 0x0F, 0x67, 0x72, 0x70, 0x63, 0x2D, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x3A, 0x30, 0x0D, 0x0A, ]) let snapshot = try GrokWebBillingFetcher.parseGRPCWebResponse( data, now: Date(timeIntervalSince1970: 1_768_000_000)) #expect(snapshot.usedPercent == 0) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: 1_780_272_000)) } @Test func `parses omitted zero percent with current billing period`() throws { let data = Data([ 0x00, 0x00, 0x00, 0x00, 0x2A, 0x0A, 0x28, 0x12, 0x00, 0x1A, 0x00, 0x22, 0x06, 0x08, 0x80, 0x97, 0xF3, 0xD0, 0x06, 0x2A, 0x06, 0x08, 0x80, 0xB1, 0x91, 0xD2, 0x06, 0x42, 0x12, 0x08, 0x01, 0x12, 0x06, 0x08, 0x80, 0x97, 0xF3, 0xD0, 0x06, 0x1A, 0x06, 0x08, 0x80, 0xB1, 0x91, 0xD2, 0x06, 0x80, 0x00, 0x00, 0x00, 0x0F, 0x67, 0x72, 0x70, 0x63, 0x2D, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x3A, 0x30, 0x0D, 0x0A, ]) let snapshot = try GrokWebBillingFetcher.parseGRPCWebResponse( data, now: Date(timeIntervalSince1970: 1_781_000_000)) #expect(snapshot.usedPercent == 0) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: 1_782_864_000)) } @Test func `uses billing field one instead of earlier unrelated float`() throws { var payload = Data() payload.append(0x4D) // field 9, fixed32 unrelated in-range float var unrelatedBits = Float(7).bitPattern.littleEndian withUnsafeBytes(of: &unrelatedBits) { payload.append(contentsOf: $0) } payload.append(0x0D) // field 1, fixed32 billing usage percent var usageBits = Float(42).bitPattern.littleEndian withUnsafeBytes(of: &usageBits) { payload.append(contentsOf: $0) } payload.append(0x10) // field 2, varint reset timestamp payload.append(contentsOf: Self.varint(1_800_000_001)) let snapshot = try GrokWebBillingFetcher.parseGRPCWebResponse(Self.grpcFrame(payload)) #expect(snapshot.usedPercent == 42) } @Test func `chooses future billing end instead of recent billing start`() throws { let recentStart = UInt64(1_800_000_000) let billingEnd = UInt64(1_802_592_000) var payload = Data() payload.append(0x0D) // field 1, fixed32 usage percent var percentBits = Float(33).bitPattern.littleEndian withUnsafeBytes(of: &percentBits) { payload.append(contentsOf: $0) } payload.append(0x10) // field 2, varint billing start payload.append(contentsOf: Self.varint(recentStart)) payload.append(0x18) // field 3, varint billing end payload.append(contentsOf: Self.varint(billingEnd)) let snapshot = try GrokWebBillingFetcher.parseGRPCWebResponse( Self.grpcFrame(payload), now: Date(timeIntervalSince1970: TimeInterval(recentStart + 1800))) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: TimeInterval(billingEnd))) } @Test func `web fetch posts grpc web request with bearer token`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) let reset = UInt64(1_800_000_002) GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let url = try #require(request.url) #expect(url == endpoint) #expect(request.httpMethod == "POST") #expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer token-123") #expect(request.value(forHTTPHeaderField: "Origin") == "https://grok.com") #expect(request.value(forHTTPHeaderField: "Referer") == "https://grok.com/?_s=usage") #expect(request.value(forHTTPHeaderField: "Accept") == "*/*") #expect(request.value(forHTTPHeaderField: "Content-Type") == "application/grpc-web+proto") #expect(request.value(forHTTPHeaderField: "x-grpc-web") == "1") #expect(request.timeoutInterval == 15) let response = HTTPURLResponse( url: url, statusCode: 200, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! let body = Self.grpcFrame(Self.protobufPayload(usedPercent: 55.5, resetEpoch: reset)) return (response, body) } let snapshot = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) #expect(GrokWebBillingStubURLProtocol.requests.count == 1) #expect(GrokWebBillingStubURLProtocol.requestBodies == [Data([0x00, 0x00, 0x00, 0x00, 0x00])]) #expect(snapshot.usedPercent == 55.5) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: TimeInterval(reset))) } @Test func `web fetch retries transient grpc timeout once`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) let reset = UInt64(1_800_000_005) let attempts = AttemptCounter() GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let attempt = attempts.increment() let url = try #require(request.url) let response = HTTPURLResponse( url: url, statusCode: 200, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! if attempt == 1 { let body = Self.grpcFrame( Data("grpc-status: 1\r\ngrpc-message: Timeout%20expired\r\n".utf8), flags: 0x80) return (response, body) } return (response, Self.grpcFrame(Self.protobufPayload(usedPercent: 25, resetEpoch: reset))) } let snapshot = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) #expect(attempts.current() == 2) #expect(GrokWebBillingStubURLProtocol.requests.count == 2) #expect(snapshot.usedPercent == 25) #expect(snapshot.resetsAt == Date(timeIntervalSince1970: TimeInterval(reset))) } @Test func `web fetch retries grpc deadline exceeded without message`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) let attempts = AttemptCounter() GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let attempt = attempts.increment() let url = try #require(request.url) let response = HTTPURLResponse( url: url, statusCode: 200, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! if attempt == 1 { return (response, Self.grpcFrame(Data("grpc-status: 4\r\n".utf8), flags: 0x80)) } return (response, Self.grpcFrame(Self.protobufPayload(usedPercent: 25, resetEpoch: 1_800_000_005))) } let snapshot = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) #expect(attempts.current() == 2) #expect(snapshot.usedPercent == 25) } @Test func `web fetch retries HTTP gateway timeout once`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) let attempts = AttemptCounter() GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let attempt = attempts.increment() let url = try #require(request.url) let statusCode = attempt == 1 ? 504 : 200 let response = HTTPURLResponse( url: url, statusCode: statusCode, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! if attempt == 1 { return (response, Data("gateway timeout".utf8)) } return (response, Self.grpcFrame(Self.protobufPayload(usedPercent: 25, resetEpoch: 1_800_000_005))) } let snapshot = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) #expect(attempts.current() == 2) #expect(snapshot.usedPercent == 25) } @Test func `web fetch can authenticate with browser cookies`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in #expect(request.value(forHTTPHeaderField: "Cookie") == "sso=session; sso-rw=session") #expect(request.value(forHTTPHeaderField: "Authorization") == nil) #expect(request.value(forHTTPHeaderField: "x-user-agent") == "connect-es/2.1.1") let response = HTTPURLResponse( url: endpoint, statusCode: 200, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! let body = Self.grpcFrame(Self.protobufPayload(usedPercent: 9, resetEpoch: 1_800_000_004)) return (response, body) } let snapshot = try await GrokWebBillingFetcher.fetch( cookieHeader: "sso=session; sso-rw=session", session: session, endpoint: endpoint) #expect(snapshot.usedPercent == 9) } @Test func `web fetch sends browser cookies with bearer credentials`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in #expect(request.value(forHTTPHeaderField: "Cookie") == "sso=session") #expect(request.value(forHTTPHeaderField: "Authorization") == "Bearer token-123") let response = HTTPURLResponse( url: endpoint, statusCode: 200, httpVersion: nil, headerFields: ["Content-Type": "application/grpc-web+proto"])! return (response, Self.protobufPayload(usedPercent: 9, resetEpoch: 1_800_000_004)) } let snapshot = try await GrokWebBillingFetcher.fetch( cookieHeader: "sso=session", credentials: Self.credentials, session: session, endpoint: endpoint) #expect(snapshot.usedPercent == 9) } @Test func `web fetch turns unauthorized response into reauth guidance`() async throws { defer { GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = nil } let config = URLSessionConfiguration.ephemeral config.protocolClasses = [GrokWebBillingStubURLProtocol.self] let session = URLSession(configuration: config) let endpoint = try #require(URL(string: "https://grok.test/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig")) GrokWebBillingStubURLProtocol.requests = [] GrokWebBillingStubURLProtocol.requestBodies = [] GrokWebBillingStubURLProtocol.handler = { request in let url = try #require(request.url) let response = HTTPURLResponse( url: url, statusCode: 401, httpVersion: nil, headerFields: ["Content-Type": "text/plain"])! return (response, Data("unauthorized".utf8)) } await #expect { _ = try await GrokWebBillingFetcher.fetch( credentials: Self.credentials, session: session, endpoint: endpoint) } throws: { error in error.localizedDescription.contains("grok login") } } @Test func `usage snapshot maps web billing when cli billing is absent`() { let snapshot = GrokUsageSnapshot( billing: nil, webBilling: GrokWebBillingSnapshot( usedPercent: 67.25, resetsAt: Date(timeIntervalSince1970: 1_800_000_003)), credentials: Self.credentials, localSummary: nil, cliVersion: nil, updatedAt: Date(timeIntervalSince1970: 1_799_000_000)) let usage = snapshot.toUsageSnapshot() #expect(usage.primary?.usedPercent == 67.25) #expect(usage.primary?.windowMinutes == nil) #expect(usage.primary?.resetsAt == Date(timeIntervalSince1970: 1_800_000_003)) #expect(usage.accountEmail(for: .grok) == "grok@example.com") #expect(usage.loginMethod(for: .grok) == "SuperGrok") } private static let credentials = GrokCredentials( accessToken: "token-123", refreshToken: "refresh-123", scope: "https://auth.x.ai::client", authMode: "oidc", userId: "user-123", email: "grok@example.com", firstName: "G", lastName: "Rok", teamId: "team-123", oidcIssuer: "https://auth.x.ai", oidcClientId: "client", expiresAt: Date(timeIntervalSince1970: 1_900_000_000), createTime: Date(timeIntervalSince1970: 1_799_000_000)) private static func protobufPayload(usedPercent: Float, resetEpoch: UInt64) -> Data { var data = Data() data.append(0x0D) // field 1, fixed32 var percentBits = usedPercent.bitPattern.littleEndian withUnsafeBytes(of: &percentBits) { data.append(contentsOf: $0) } data.append(0x10) // field 2, varint data.append(contentsOf: Self.varint(resetEpoch)) return data } private static func grpcFrame(_ payload: Data, flags: UInt8 = 0x00) -> Data { var data = Data([flags]) let length = UInt32(payload.count).bigEndian withUnsafeBytes(of: length) { data.append(contentsOf: $0) } data.append(payload) return data } private static func varint(_ value: UInt64) -> [UInt8] { var remaining = value var bytes: [UInt8] = [] repeat { var byte = UInt8(remaining & 0x7F) remaining >>= 7 if remaining != 0 { byte |= 0x80 } bytes.append(byte) } while remaining != 0 return bytes } private static func cookie(name: String, value: String) -> HTTPCookie? { HTTPCookie(properties: [ .domain: "grok.com", .path: "/", .name: name, .value: value, ]) } private static func data(hexString: String) -> Data? { var data = Data() var index = hexString.startIndex while index < hexString.endIndex { let next = hexString.index(index, offsetBy: 2, limitedBy: hexString.endIndex) ?? hexString.endIndex guard let byte = UInt8(hexString[index.. (HTTPURLResponse, Data))? override static func canInit(with _: URLRequest) -> Bool { true } override static func canonicalRequest(for request: URLRequest) -> URLRequest { request } override func startLoading() { Self.requests.append(self.request) Self.requestBodies.append(Self.readBody(from: self.request)) guard let handler = Self.handler else { self.client?.urlProtocol(self, didFailWithError: URLError(.badServerResponse)) return } do { let (response, data) = try handler(self.request) self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed) self.client?.urlProtocol(self, didLoad: data) self.client?.urlProtocolDidFinishLoading(self) } catch { self.client?.urlProtocol(self, didFailWithError: error) } } override func stopLoading() {} private static func readBody(from request: URLRequest) -> Data? { if let body = request.httpBody { return body } guard let stream = request.httpBodyStream else { return nil } stream.open() defer { stream.close() } var data = Data() var buffer = [UInt8](repeating: 0, count: 1024) while stream.hasBytesAvailable { let count = stream.read(&buffer, maxLength: buffer.count) if count > 0 { data.append(buffer, count: count) } else { break } } return data } }