From f1a5d3691e3d663c6497b9bfff16dff3b946063f Mon Sep 17 00:00:00 2001 From: Vincent Peng Date: Wed, 15 Jul 2026 16:27:13 +1000 Subject: [PATCH 1/2] Handle Grok team principal billing fallback --- .../CodexBar/PreferencesProvidersPane.swift | 4 +- Sources/CodexBar/UsageStore+Accessors.swift | 7 + .../UsageStore+BackgroundRefresh.swift | 1 + .../UsageStore+CodexResetCredits.swift | 1 + Sources/CodexBar/UsageStore+Refresh.swift | 3 + Sources/CodexBar/UsageStore.swift | 3 + Sources/CodexBarCLI/CLIPayloads.swift | 6 + Sources/CodexBarCLI/CLIUsageCommand.swift | 7 +- .../Codex/CodexProviderDescriptor.swift | 3 +- .../Providers/Grok/GrokAuth.swift | 15 +- .../Grok/GrokProviderDescriptor.swift | 44 ++++- .../Providers/Grok/GrokStatusProbe.swift | 66 ++++++- .../Grok/GrokWebBillingFetcher.swift | 42 ++++- .../Providers/ProviderFetchPlan.swift | 10 +- Tests/CodexBarTests/CLISnapshotTests.swift | 4 +- .../CodexUserFacingErrorTests.swift | 15 ++ Tests/CodexBarTests/GrokAuthTests.swift | 63 +++++++ .../GrokWebBillingFetcherTests.swift | 169 +++++++++++++++++- docs/grok.md | 18 +- 19 files changed, 455 insertions(+), 26 deletions(-) diff --git a/Sources/CodexBar/PreferencesProvidersPane.swift b/Sources/CodexBar/PreferencesProvidersPane.swift index 9be5a35ab5..4f8705dcb9 100644 --- a/Sources/CodexBar/PreferencesProvidersPane.swift +++ b/Sources/CodexBar/PreferencesProvidersPane.swift @@ -325,7 +325,9 @@ struct ProvidersPane: View { } func providerErrorDisplay(_ provider: UsageProvider) -> ProviderErrorDisplay? { - guard let full = self.store.error(for: provider), !full.isEmpty else { return nil } + guard let full = self.store.error(for: provider) ?? self.store.diagnostic(for: provider), + !full.isEmpty + else { return nil } let preview = self.store.userFacingError(for: provider) ?? full return ProviderErrorDisplay( preview: self.truncated(preview, prefix: ""), diff --git a/Sources/CodexBar/UsageStore+Accessors.swift b/Sources/CodexBar/UsageStore+Accessors.swift index 5759af4abe..a54aeb1fd3 100644 --- a/Sources/CodexBar/UsageStore+Accessors.swift +++ b/Sources/CodexBar/UsageStore+Accessors.swift @@ -38,11 +38,18 @@ extension UsageStore { self.errors[provider] } + func diagnostic(for provider: UsageProvider) -> String? { + self.diagnostics[provider] + } + func userFacingError(for provider: UsageProvider) -> String? { if let raw = self.errors[provider] { guard provider == .codex else { return raw } return CodexUIErrorMapper.userFacingMessage(raw) } + if let diagnostic = self.diagnostics[provider] { + return diagnostic + } return self.unavailableMessage(for: provider) } diff --git a/Sources/CodexBar/UsageStore+BackgroundRefresh.swift b/Sources/CodexBar/UsageStore+BackgroundRefresh.swift index 5e7a929504..68cbe900e7 100644 --- a/Sources/CodexBar/UsageStore+BackgroundRefresh.swift +++ b/Sources/CodexBar/UsageStore+BackgroundRefresh.swift @@ -25,6 +25,7 @@ extension UsageStore { self.snapshots.removeValue(forKey: provider) self.lastKnownResetSnapshots.removeValue(forKey: provider) self.errors[provider] = nil + self.diagnostics[provider] = nil if provider == .gemini { self.clearGeminiConsumerTierDeprecationObservation() } diff --git a/Sources/CodexBar/UsageStore+CodexResetCredits.swift b/Sources/CodexBar/UsageStore+CodexResetCredits.swift index 6dd2b9542e..a5d84af753 100644 --- a/Sources/CodexBar/UsageStore+CodexResetCredits.swift +++ b/Sources/CodexBar/UsageStore+CodexResetCredits.swift @@ -117,6 +117,7 @@ extension ProviderFetchOutcome { sourceLabel: result.sourceLabel, strategyID: result.strategyID, strategyKind: result.strategyKind, + diagnostic: result.diagnostic, claudeOAuthKeychainPersistentRefHash: result.claudeOAuthKeychainPersistentRefHash, claudeOAuthHistoryOwnerIdentifier: result.claudeOAuthHistoryOwnerIdentifier, claudeOAuthKeychainCredentialMismatch: result.claudeOAuthKeychainCredentialMismatch, diff --git a/Sources/CodexBar/UsageStore+Refresh.swift b/Sources/CodexBar/UsageStore+Refresh.swift index 0d3eb3e42d..de9c85030f 100644 --- a/Sources/CodexBar/UsageStore+Refresh.swift +++ b/Sources/CodexBar/UsageStore+Refresh.swift @@ -328,6 +328,7 @@ extension UsageStore { } } + self.diagnostics[provider] = nil let claudeAuthStateBeforeFetch = provider == .claude ? await Self.captureClaudeRefreshAuthState(invalidateCredentialsFile: true) : nil @@ -551,6 +552,7 @@ extension UsageStore { } self.lastSourceLabels[provider] = result.sourceLabel self.errors[provider] = nil + self.diagnostics[provider] = result.diagnostic if let tokenAccount = currentTokenAccount { self.cacheTokenAccountSnapshot( provider: provider, @@ -954,6 +956,7 @@ extension UsageStore { let shouldNotifyPermissionPrompt = Self.isPermissionPromptWaiting(error) await MainActor.run { guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return } + self.diagnostics[provider] = nil if provider == .gemini, Self.isGeminiConsumerTierDeprecationError(error) { // This is a durable provider migration signal, not a transient fetch failure. // Surface it immediately so a cached snapshot cannot hide the required handoff. diff --git a/Sources/CodexBar/UsageStore.swift b/Sources/CodexBar/UsageStore.swift index 21203bcb00..bf41d667e8 100644 --- a/Sources/CodexBar/UsageStore.swift +++ b/Sources/CodexBar/UsageStore.swift @@ -11,6 +11,7 @@ extension UsageStore { var menuObservationToken: Int { _ = self.snapshots _ = self.errors + _ = self.diagnostics _ = self.knownLimitsAvailabilityByProvider _ = self.lastSourceLabels _ = self.lastFetchAttempts @@ -44,6 +45,7 @@ extension UsageStore { var iconObservationToken: Int { _ = self.snapshots _ = self.errors + _ = self.diagnostics _ = self.knownLimitsAvailabilityByProvider _ = self.credits _ = self.lastCreditsError @@ -153,6 +155,7 @@ final class UsageStore { var snapshots: [UsageProvider: UsageSnapshot] = [:] var errors: [UsageProvider: String] = [:] + var diagnostics: [UsageProvider: String] = [:] var geminiObservedConsumerTierDeprecation = false var knownLimitsAvailabilityByProvider: [UsageProvider: UsageLimitsAvailability] = [:] var lastSourceLabels: [UsageProvider: String] = [:] diff --git a/Sources/CodexBarCLI/CLIPayloads.swift b/Sources/CodexBarCLI/CLIPayloads.swift index 5d530edbc6..07177e8e3b 100644 --- a/Sources/CodexBarCLI/CLIPayloads.swift +++ b/Sources/CodexBarCLI/CLIPayloads.swift @@ -15,6 +15,7 @@ struct ProviderPayload: Encodable { let credits: CreditsSnapshot? let antigravityPlanInfo: AntigravityPlanInfoSummary? let openaiDashboard: OpenAIDashboardSnapshot? + let diagnostic: String? let error: ProviderErrorPayload? let pace: ProviderPacePayload? @@ -28,6 +29,7 @@ struct ProviderPayload: Encodable { case credits case antigravityPlanInfo case openaiDashboard + case diagnostic case error case pace } @@ -44,6 +46,7 @@ struct ProviderPayload: Encodable { antigravityPlanInfo: AntigravityPlanInfoSummary?, openaiDashboard: OpenAIDashboardSnapshot?, error: ProviderErrorPayload?, + diagnostic: String? = nil, pace: ProviderPacePayload? = nil) { self.provider = provider.rawValue @@ -56,6 +59,7 @@ struct ProviderPayload: Encodable { self.credits = credits self.antigravityPlanInfo = antigravityPlanInfo self.openaiDashboard = openaiDashboard + self.diagnostic = diagnostic self.error = error self.pace = pace } @@ -72,6 +76,7 @@ struct ProviderPayload: Encodable { antigravityPlanInfo: AntigravityPlanInfoSummary?, openaiDashboard: OpenAIDashboardSnapshot?, error: ProviderErrorPayload?, + diagnostic: String? = nil, pace: ProviderPacePayload? = nil) { self.provider = providerID @@ -84,6 +89,7 @@ struct ProviderPayload: Encodable { self.credits = credits self.antigravityPlanInfo = antigravityPlanInfo self.openaiDashboard = openaiDashboard + self.diagnostic = diagnostic self.error = error self.pace = pace } diff --git a/Sources/CodexBarCLI/CLIUsageCommand.swift b/Sources/CodexBarCLI/CLIUsageCommand.swift index 3ebee626f3..b1c9230aec 100644 --- a/Sources/CodexBarCLI/CLIUsageCommand.swift +++ b/Sources/CodexBarCLI/CLIUsageCommand.swift @@ -48,6 +48,7 @@ private struct UsageSuccessRenderInput { let dashboard: OpenAIDashboardSnapshot? let effectiveSourceMode: ProviderSourceMode let command: UsageCommandContext + let diagnostic: String? let notes: [String] } @@ -287,6 +288,7 @@ extension CodexBarCLI { credits: CreditsSnapshot?, antigravityPlanInfo: AntigravityPlanInfoSummary?, dashboard: OpenAIDashboardSnapshot?, + diagnostic: String?, weeklyWorkDays: Int?) -> ProviderPayload { ProviderPayload( @@ -301,6 +303,7 @@ extension CodexBarCLI { antigravityPlanInfo: antigravityPlanInfo, openaiDashboard: dashboard, error: nil, + diagnostic: diagnostic, pace: CLIRenderer.providerPacePayload(provider: provider, snapshot: usage, weeklyWorkDays: weeklyWorkDays)) } @@ -354,6 +357,7 @@ extension CodexBarCLI { credits: input.credits, antigravityPlanInfo: input.antigravityPlanInfo, dashboard: input.dashboard, + diagnostic: input.diagnostic, weeklyWorkDays: input.command.weeklyWorkDays)) } } @@ -458,7 +462,7 @@ extension CodexBarCLI { let notes = Self.usageTextNotes( provider: provider, sourceMode: effectiveSourceMode, - resolvedSourceLabel: source) + resolvedSourceLabel: source) + (result.diagnostic.map { [$0] } ?? []) Self.appendSuccessRenderOutput( UsageSuccessRenderInput( @@ -474,6 +478,7 @@ extension CodexBarCLI { dashboard: dashboard, effectiveSourceMode: effectiveSourceMode, command: command, + diagnostic: result.diagnostic, notes: notes), output: &output) case let .failure(error): diff --git a/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift index 5a15ee161a..3dabff36de 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexProviderDescriptor.swift @@ -321,7 +321,8 @@ struct CodexOAuthFetchStrategy: ProviderFetchStrategy { dashboard: oauthResult.dashboard, sourceLabel: oauthResult.sourceLabel, strategyID: oauthResult.strategyID, - strategyKind: oauthResult.strategyKind) + strategyKind: oauthResult.strategyKind, + diagnostic: oauthResult.diagnostic) } private static func fetchResetCreditsIfRequested( diff --git a/Sources/CodexBarCore/Providers/Grok/GrokAuth.swift b/Sources/CodexBarCore/Providers/Grok/GrokAuth.swift index 9544808d1f..86eefa7f37 100644 --- a/Sources/CodexBarCore/Providers/Grok/GrokAuth.swift +++ b/Sources/CodexBarCore/Providers/Grok/GrokAuth.swift @@ -10,6 +10,9 @@ public struct GrokCredentials: Sendable { public let firstName: String? public let lastName: String? public let teamId: String? + /// The principal scope reported by Grok's cached OIDC credential, when available. + /// Keep this optional because older auth.json entries do not include it. + public let principalType: String? public let oidcIssuer: String? public let oidcClientId: String? public let expiresAt: Date? @@ -28,7 +31,8 @@ public struct GrokCredentials: Sendable { oidcIssuer: String?, oidcClientId: String?, expiresAt: Date?, - createTime: Date?) + createTime: Date?, + principalType: String? = nil) { self.accessToken = accessToken self.refreshToken = refreshToken @@ -39,6 +43,7 @@ public struct GrokCredentials: Sendable { self.firstName = firstName self.lastName = lastName self.teamId = teamId + self.principalType = principalType self.oidcIssuer = oidcIssuer self.oidcClientId = oidcClientId self.expiresAt = expiresAt @@ -56,6 +61,11 @@ public struct GrokCredentials: Sendable { return Date() >= expiresAt } + public var isTeamPrincipal: Bool { + self.principalType?.trimmingCharacters(in: .whitespacesAndNewlines) + .caseInsensitiveCompare("team") == .orderedSame + } + public var loginMethod: String? { switch self.authMode?.lowercased() { case "oidc": "SuperGrok" @@ -150,7 +160,8 @@ public enum GrokCredentialsStore { oidcIssuer: (entry["oidc_issuer"] as? String)?.nilIfEmpty, oidcClientId: (entry["oidc_client_id"] as? String)?.nilIfEmpty, expiresAt: Self.parseDate(entry["expires_at"]), - createTime: Self.parseDate(entry["create_time"])) + createTime: Self.parseDate(entry["create_time"]), + principalType: (entry["principal_type"] as? String)?.nilIfEmpty) } private static func selectPreferredEntry(in root: [String: Any]) -> (scope: String, entry: [String: Any])? { diff --git a/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift index 6aaf76f6a1..41f7234c2e 100644 --- a/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift @@ -90,7 +90,8 @@ struct GrokCLIFetchStrategy: ProviderFetchStrategy { let snap = try await probe.fetch(env: context.env) return self.makeResult( usage: snap.toUsageSnapshot(), - sourceLabel: "grok-cli") + sourceLabel: "grok-cli", + diagnostic: snap.diagnostic) } func shouldFallback(on _: Error, context: ProviderFetchContext) -> Bool { @@ -101,6 +102,10 @@ struct GrokCLIFetchStrategy: ProviderFetchStrategy { struct GrokWebFetchStrategy: ProviderFetchStrategy { let id: String = "grok.web" let kind: ProviderFetchKind = .web + typealias WebBillingFetch = @Sendable () async throws -> ( + snapshot: GrokWebBillingSnapshot, + sourceLabel: String, + authenticatedByAuthFile: Bool) static func canImportBrowserCookies(runtime: ProviderRuntime, env: [String: String]) -> Bool { runtime == .app || env["CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT"] == "1" @@ -118,7 +123,36 @@ struct GrokWebFetchStrategy: ProviderFetchStrategy { } func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { - let (webBilling, sourceLabel, authenticatedByAuthFile) = try await self.fetchWebBilling(context: context) + try await self.fetch(context, webBilling: { [self] in + try await self.fetchWebBilling(context: context) + }) + } + + func fetch( + _ context: ProviderFetchContext, + webBilling fetchWebBilling: @escaping WebBillingFetch) async throws -> ProviderFetchResult + { + let webBilling: GrokWebBillingSnapshot + let sourceLabel: String + let authenticatedByAuthFile: Bool + do { + (webBilling, sourceLabel, authenticatedByAuthFile) = try await fetchWebBilling() + } catch GrokWebBillingError.teamUsageUnsupported { + guard let credentials = try? GrokCredentialsStore.load(env: context.env), + !credentials.isExpired, + credentials.isTeamPrincipal + else { + throw GrokWebBillingError.teamUsageUnsupported + } + let identitySnapshot = GrokStatusProbe.identityOnlySnapshot( + credentials: credentials, + localSummary: GrokLocalSessionScanner.summarize(env: context.env), + cliVersion: GrokStatusProbe.detectVersion(env: context.env)) + return self.makeResult( + usage: identitySnapshot.toUsageSnapshot(), + sourceLabel: "grok-web", + diagnostic: identitySnapshot.diagnostic) + } let credentials = Self.credentialsForWebBillingSnapshot( credentials: try? GrokCredentialsStore.load(env: context.env), authenticatedByAuthFile: authenticatedByAuthFile) @@ -198,17 +232,21 @@ struct GrokWebFetchStrategy: ProviderFetchStrategy { credentials: credentials) } var lastError: Error? + var teamUsageUnsupportedError: Error? for session in sessions { for authCredentials in Self.cookieAuthAttempts(credentials: credentials) { do { let snapshot = try await fetchSnapshot(session.cookieHeader, authCredentials) return (snapshot, session.sourceLabel) } catch { + if case GrokWebBillingError.teamUsageUnsupported = error { + teamUsageUnsupportedError = error + } lastError = error } } } - throw lastError ?? GrokWebBillingError.missingCredentials + throw teamUsageUnsupportedError ?? lastError ?? GrokWebBillingError.missingCredentials } static func cookieAuthAttempts(credentials: GrokCredentials?) -> [GrokCredentials?] { diff --git a/Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift b/Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift index 5473d24928..a56f801a2a 100644 --- a/Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift +++ b/Sources/CodexBarCore/Providers/Grok/GrokStatusProbe.swift @@ -6,6 +6,7 @@ public struct GrokUsageSnapshot: Sendable { public let credentials: GrokCredentials? public let localSummary: GrokLocalSessionSummary? public let cliVersion: String? + public let diagnostic: String? public let updatedAt: Date public init( @@ -14,13 +15,15 @@ public struct GrokUsageSnapshot: Sendable { credentials: GrokCredentials?, localSummary: GrokLocalSessionSummary?, cliVersion: String?, - updatedAt: Date) + updatedAt: Date, + diagnostic: String? = nil) { self.billing = billing self.webBilling = webBilling self.credentials = credentials self.localSummary = localSummary self.cliVersion = cliVersion + self.diagnostic = diagnostic self.updatedAt = updatedAt } @@ -62,6 +65,9 @@ public struct GrokUsageSnapshot: Sendable { } public struct GrokStatusProbe: Sendable { + public static let teamUsageUnavailableMessage = + "Grok team usage is unavailable from the current billing surface; identity is still available." + public init() {} public static func detectVersion(env: [String: String] = ProcessInfo.processInfo.environment) -> String? { @@ -89,10 +95,12 @@ public struct GrokStatusProbe: Sendable { var billing: GrokBillingResponse? 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() } catch { rpcError = error @@ -106,6 +114,19 @@ public struct GrokStatusProbe: Sendable { // identity field, so a stale `~/.grok/sessions/` directory must not // suppress the auth-required hint. CLI-only fetches need a billing // response; the provider pipeline owns the separate web fallback. + if billing == nil, + let credentials, + Self.shouldUseIdentityOnlyFallback( + credentials: credentials, + billingAttempted: billingAttempted, + error: rpcError) + { + return Self.identityOnlySnapshot( + credentials: credentials, + localSummary: localSummary, + cliVersion: cliVersion) + } + if billing == nil { throw rpcError ?? GrokRPCError.notAuthenticated } @@ -122,6 +143,47 @@ public struct GrokStatusProbe: Sendable { updatedAt: Date()) } + static func identityOnlySnapshot( + credentials: GrokCredentials, + localSummary: GrokLocalSessionSummary?, + cliVersion: String?, + updatedAt: Date = .init()) -> GrokUsageSnapshot + { + GrokUsageSnapshot( + billing: nil, + webBilling: nil, + credentials: credentials, + localSummary: localSummary, + cliVersion: cliVersion, + updatedAt: updatedAt, + diagnostic: GrokStatusProbe.teamUsageUnavailableMessage) + } + + static func isBillingMethodUnavailable(_ error: Error?) -> Bool { + guard let error, + case let GrokRPCError.requestFailed(message) = error + else { + return false + } + let normalized = message.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + return normalized == "method not found" || normalized.hasPrefix("method not found:") + } + + static func shouldUseIdentityOnlyFallback( + credentials: GrokCredentials?, + billingAttempted: Bool, + error: Error?) -> Bool + { + guard billingAttempted, + let credentials, + !credentials.isExpired, + credentials.isTeamPrincipal + else { + return false + } + return Self.isBillingMethodUnavailable(error) + } + static func credentialsForSnapshot( credentials: GrokCredentials?, billing: GrokBillingResponse?, @@ -140,7 +202,7 @@ public struct GrokStatusProbe: Sendable { return status == 401 || status == 403 case let .rpcFailed(status, _): return status == 16 - case .missingCredentials, .emptyResponse, .invalidResponse, .parseFailed: + case .missingCredentials, .emptyResponse, .invalidResponse, .teamUsageUnsupported, .parseFailed: return false } } diff --git a/Sources/CodexBarCore/Providers/Grok/GrokWebBillingFetcher.swift b/Sources/CodexBarCore/Providers/Grok/GrokWebBillingFetcher.swift index 4ab67ac8e2..fef62a741a 100644 --- a/Sources/CodexBarCore/Providers/Grok/GrokWebBillingFetcher.swift +++ b/Sources/CodexBarCore/Providers/Grok/GrokWebBillingFetcher.swift @@ -19,6 +19,7 @@ public enum GrokWebBillingError: LocalizedError, Sendable { case invalidResponse case requestFailed(Int, String) case rpcFailed(Int, String) + case teamUsageUnsupported case parseFailed public var errorDescription: String? { @@ -41,6 +42,8 @@ public enum GrokWebBillingError: LocalizedError, Sendable { } else { "Grok web billing RPC failed with status \(status): \(message)" } + case .teamUsageUnsupported: + "Grok team usage is unavailable from the current billing surface." case .parseFailed: "Could not parse Grok web billing usage." } @@ -76,6 +79,7 @@ public enum GrokWebBillingFetcher { try await self.fetch( authorizationHeader: "Bearer \(credentials.accessToken)", cookieHeader: nil, + principalType: credentials.isExpired ? nil : credentials.principalType, transport: transport, endpoint: endpoint) } @@ -104,6 +108,7 @@ public enum GrokWebBillingFetcher { return try await self.fetch( authorizationHeader: authorizationHeader, cookieHeader: cookieHeader, + principalType: credentials.flatMap { $0.isExpired ? nil : $0.principalType }, transport: transport, endpoint: endpoint) } @@ -111,6 +116,7 @@ public enum GrokWebBillingFetcher { private static func fetch( authorizationHeader: String?, cookieHeader: String?, + principalType: String?, transport: any ProviderHTTPTransport, endpoint: URL) async throws -> GrokWebBillingSnapshot { @@ -120,15 +126,33 @@ public enum GrokWebBillingFetcher { cookieHeader: cookieHeader, transport: transport, endpoint: endpoint) - } catch where self.shouldRetry(error) { - return try await self.fetchOnce( - authorizationHeader: authorizationHeader, - cookieHeader: cookieHeader, - transport: transport, - endpoint: endpoint) + } catch { + if self.shouldRetry(error) { + do { + return try await self.fetchOnce( + authorizationHeader: authorizationHeader, + cookieHeader: cookieHeader, + transport: transport, + endpoint: endpoint) + } catch { + throw self.classified(error, principalType: principalType) + } + } + throw self.classified(error, principalType: principalType) } } + private static func classified(_ error: Error, principalType: String?) -> Error { + guard principalType?.trimmingCharacters(in: .whitespacesAndNewlines) + .caseInsensitiveCompare("team") == .orderedSame, + case let GrokWebBillingError.rpcFailed(status, message) = error, + self.isTeamBillingUnavailable(status: status, message: message) + else { + return error + } + return GrokWebBillingError.teamUsageUnsupported + } + private static func fetchOnce( authorizationHeader: String?, cookieHeader: String?, @@ -279,6 +303,12 @@ public enum GrokWebBillingFetcher { throw GrokWebBillingError.rpcFailed(status, fields["grpc-message"] ?? "") } + static func isTeamBillingUnavailable(status: Int, message: String) -> Bool { + guard status == 9 else { return false } + let normalized = message.trimmingCharacters(in: .whitespacesAndNewlines).lowercased() + return normalized == "no personal team" || normalized == "no personal team." + } + static func grpcHeaderFields(from headers: [AnyHashable: Any]) -> [String: String] { var fields: [String: String] = [:] for (key, value) in headers { diff --git a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift index c654cfcd4d..a78dbe474e 100644 --- a/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift +++ b/Sources/CodexBarCore/Providers/ProviderFetchPlan.swift @@ -101,6 +101,8 @@ public struct ProviderFetchResult: Sendable { public let sourceLabel: String public let strategyID: String public let strategyKind: ProviderFetchKind + /// Optional live diagnostic retained alongside an otherwise usable snapshot. + public let diagnostic: String? /// Transient account ownership evidence for plan-utilization history. /// The raw Keychain reference never enters the persisted usage snapshot. public let claudeOAuthKeychainPersistentRefHash: String? @@ -121,6 +123,7 @@ public struct ProviderFetchResult: Sendable { sourceLabel: String, strategyID: String, strategyKind: ProviderFetchKind, + diagnostic: String? = nil, claudeOAuthKeychainPersistentRefHash: String? = nil, claudeOAuthHistoryOwnerIdentifier: String? = nil, claudeOAuthKeychainCredentialMismatch: Bool = false, @@ -133,6 +136,7 @@ public struct ProviderFetchResult: Sendable { self.sourceLabel = sourceLabel self.strategyID = strategyID self.strategyKind = strategyKind + self.diagnostic = diagnostic self.claudeOAuthKeychainPersistentRefHash = claudeOAuthKeychainPersistentRefHash self.claudeOAuthHistoryOwnerIdentifier = claudeOAuthHistoryOwnerIdentifier self.claudeOAuthKeychainCredentialMismatch = claudeOAuthKeychainCredentialMismatch @@ -198,7 +202,8 @@ extension ProviderFetchStrategy { usage: UsageSnapshot, credits: CreditsSnapshot? = nil, dashboard: OpenAIDashboardSnapshot? = nil, - sourceLabel: String) -> ProviderFetchResult + sourceLabel: String, + diagnostic: String? = nil) -> ProviderFetchResult { ProviderFetchResult( usage: usage, @@ -206,7 +211,8 @@ extension ProviderFetchStrategy { dashboard: dashboard, sourceLabel: sourceLabel, strategyID: self.id, - strategyKind: self.kind) + strategyKind: self.kind, + diagnostic: diagnostic) } } diff --git a/Tests/CodexBarTests/CLISnapshotTests.swift b/Tests/CodexBarTests/CLISnapshotTests.swift index a445474fb7..d352d74442 100644 --- a/Tests/CodexBarTests/CLISnapshotTests.swift +++ b/Tests/CodexBarTests/CLISnapshotTests.swift @@ -820,7 +820,8 @@ struct CLISnapshotTests { credits: nil, antigravityPlanInfo: nil, openaiDashboard: nil, - error: nil) + error: nil, + diagnostic: "Grok team usage is unavailable from the current billing surface.") let encoder = JSONEncoder() encoder.dateEncodingStrategy = .secondsSince1970 let data = try encoder.encode(payload) @@ -833,6 +834,7 @@ struct CLISnapshotTests { #expect(json.contains("\"version\":\"1.2.3\"")) #expect(json.contains("\"status\"")) #expect(json.contains("status.example.com")) + #expect(json.contains("Grok team usage is unavailable from the current billing surface.")) #expect(json.contains("\"primary\"")) #expect(json.contains("\"windowMinutes\":300")) #expect(json.contains("1700000000")) diff --git a/Tests/CodexBarTests/CodexUserFacingErrorTests.swift b/Tests/CodexBarTests/CodexUserFacingErrorTests.swift index ec8f569e14..695b26aec0 100644 --- a/Tests/CodexBarTests/CodexUserFacingErrorTests.swift +++ b/Tests/CodexBarTests/CodexUserFacingErrorTests.swift @@ -180,6 +180,21 @@ struct CodexUserFacingErrorTests { #expect(store.userFacingError(for: .claude) == "Claude probe failed with debug detail") } + @Test + func `successful provider diagnostic does not make usage stale`() { + let settings = self.makeSettingsStore(suite: "CodexUserFacingErrorTests-success-diagnostic") + let store = self.makeUsageStore(settings: settings) + store.diagnostics[.grok] = GrokStatusProbe.teamUsageUnavailableMessage + + #expect(store.userFacingError(for: .grok) == GrokStatusProbe.teamUsageUnavailableMessage) + #expect(!store.isStale(provider: .grok)) + + let pane = ProvidersPane(settings: settings, store: store) + let display = pane._test_providerErrorDisplay(for: .grok) + #expect(display?.preview == GrokStatusProbe.teamUsageUnavailableMessage) + #expect(display?.full == GrokStatusProbe.teamUsageUnavailableMessage) + } + @Test func `providers pane codex model uses sanitized values`() { let settings = self.makeSettingsStore(suite: "CodexUserFacingErrorTests-pane-model") diff --git a/Tests/CodexBarTests/GrokAuthTests.swift b/Tests/CodexBarTests/GrokAuthTests.swift index a2ed52e554..d37b74dc64 100644 --- a/Tests/CodexBarTests/GrokAuthTests.swift +++ b/Tests/CodexBarTests/GrokAuthTests.swift @@ -16,6 +16,7 @@ struct GrokAuthTests { "first_name": "Ada", "last_name": "Lovelace", "team_id": "team-uuid", + "principal_type": "Team", "refresh_token": "refresh-secret", "expires_at": "2026-05-22T19:31:33.384327Z", "oidc_issuer": "https://auth.x.ai", @@ -30,6 +31,8 @@ struct GrokAuthTests { #expect(creds.refreshToken == "refresh-secret") #expect(creds.email == "user@example.com") #expect(creds.teamId == "team-uuid") + #expect(creds.principalType == "Team") + #expect(creds.isTeamPrincipal) #expect(creds.authMode == "oidc") #expect(creds.displayName == "Ada Lovelace") #expect(creds.loginMethod == "SuperGrok") @@ -143,6 +146,66 @@ struct GrokAuthTests { #expect(!GrokStatusProbe.shouldSurfaceRemoteAuthError(GrokWebBillingError.parseFailed)) } + @Test + func `team method unavailable is classified without broadening other rpc failures`() { + #expect(GrokStatusProbe.isBillingMethodUnavailable( + GrokRPCError.requestFailed("Method not found"))) + #expect(GrokStatusProbe.isBillingMethodUnavailable( + GrokRPCError.requestFailed("Method not found: x.ai/billing"))) + #expect(!GrokStatusProbe.isBillingMethodUnavailable( + GrokRPCError.requestFailed("Authentication required"))) + #expect(!GrokStatusProbe.isBillingMethodUnavailable(nil)) + } + + @Test + func `team identity fallback requires an attempted billing call`() throws { + let json = #"{"https://auth.x.ai::client":{"key":"token","principal_type":"Team"}}"# + let credentials = try GrokCredentialsStore.parse(data: Data(json.utf8)) + let methodNotFound = GrokRPCError.requestFailed("Method not found") + + #expect(GrokStatusProbe.shouldUseIdentityOnlyFallback( + credentials: credentials, + billingAttempted: true, + error: methodNotFound)) + #expect(!GrokStatusProbe.shouldUseIdentityOnlyFallback( + credentials: credentials, + billingAttempted: false, + error: methodNotFound)) + } + + @Test + func `principal type matching is case and whitespace insensitive`() throws { + let json = #"{"https://auth.x.ai::client":{"key":"token","principal_type":" team "}}"# + let credentials = try GrokCredentialsStore.parse(data: Data(json.utf8)) + #expect(credentials.isTeamPrincipal) + } + + @Test + func `identity-only team snapshot retains identity and diagnostic`() throws { + let json = #""" + { + "https://auth.x.ai::client": { + "key": "token", + "email": "team@example.com", + "team_id": "team-123", + "principal_type": "Team" + } + } + """# + let credentials = try GrokCredentialsStore.parse(data: Data(json.utf8)) + let snapshot = GrokStatusProbe.identityOnlySnapshot( + credentials: credentials, + localSummary: nil, + cliVersion: "0.1.210", + updatedAt: Date(timeIntervalSince1970: 1_800_000_000)) + + let usage = snapshot.toUsageSnapshot() + #expect(usage.primary == nil) + #expect(usage.accountEmail(for: .grok) == "team@example.com") + #expect(usage.accountOrganization(for: .grok) == "team-123") + #expect(snapshot.diagnostic == GrokStatusProbe.teamUsageUnavailableMessage) + } + @Test func `falls back to legacy when OIDC entry has no key`() throws { // A stale/partial OIDC record must not shadow a healthy legacy session. diff --git a/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift b/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift index 408b6a7f87..3d51d38ea1 100644 --- a/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift +++ b/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift @@ -170,6 +170,27 @@ struct GrokWebBillingFetcherTests { #expect(result.0.usedPercent == 9) } + @Test + func `cookie session loop preserves team unsupported billing`() async throws { + let cookie = try #require(Self.cookie(name: "sso", value: "team-session")) + let sessions = [GrokCookieImporter.SessionInfo(cookies: [cookie], sourceLabel: "Chrome")] + + await #expect { + _ = try await GrokWebFetchStrategy.fetchFirstValidCookieSession( + sessions, + credentials: Self.credentials) + { _, authCredentials in + if authCredentials != nil { + throw GrokWebBillingError.teamUsageUnsupported + } + throw GrokWebBillingError.rpcFailed(9, "No personal team") + } + } throws: { error in + guard case GrokWebBillingError.teamUsageUnsupported = error else { return false } + return true + } + } + @Test func `web strategy skips expired bearer for browser cookies`() async throws { let cookie = try #require(Self.cookie(name: "sso", value: "session")) @@ -238,6 +259,149 @@ struct GrokWebBillingFetcherTests { message: "OAuth2 access token lacks the required billing scope")) } + @Test + func `only a team principal with no personal team gets unsupported billing guidance`() { + #expect(GrokWebBillingFetcher.isTeamBillingUnavailable( + status: 9, + message: "No personal team")) + #expect(GrokWebBillingFetcher.isTeamBillingUnavailable( + status: 9, + message: " no PERSONAL team ")) + #expect(GrokWebBillingFetcher.isTeamBillingUnavailable( + status: 9, + message: "No personal team.")) + #expect(!GrokWebBillingFetcher.isTeamBillingUnavailable( + status: 9, + message: "Permission denied")) + #expect(!GrokWebBillingFetcher.isTeamBillingUnavailable( + status: 7, + message: "No personal team")) + #expect(GrokWebBillingError.teamUsageUnsupported.errorDescription?.contains("identity") == false) + } + + @Test + func `team principal status nine response is classified as unsupported billing`() 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 = "No personal team.".addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) + ?? "No personal team." + let body = Self.grpcFrame( + Data("grpc-status: 9\r\ngrpc-message: \(message)\r\n".utf8), + flags: 0x80) + + 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 + guard case GrokWebBillingError.teamUsageUnsupported = error else { return false } + return true + } + + let expiredCredentials = GrokCredentials( + accessToken: "expired-token", + refreshToken: nil, + scope: Self.credentials.scope, + authMode: Self.credentials.authMode, + userId: Self.credentials.userId, + email: Self.credentials.email, + firstName: Self.credentials.firstName, + lastName: Self.credentials.lastName, + teamId: Self.credentials.teamId, + oidcIssuer: Self.credentials.oidcIssuer, + oidcClientId: Self.credentials.oidcClientId, + expiresAt: .distantPast, + createTime: Self.credentials.createTime, + principalType: Self.credentials.principalType) + + await #expect { + _ = try await GrokWebBillingFetcher.fetch( + cookieHeader: "sso=team-session", + credentials: expiredCredentials, + session: session, + endpoint: endpoint) + } throws: { error in + guard case let GrokWebBillingError.rpcFailed(status, message) = error else { return false } + return status == 9 && message == "No personal team." + } + + await #expect { + _ = try await GrokWebBillingFetcher.fetch( + credentials: expiredCredentials, + session: session, + endpoint: endpoint) + } throws: { error in + guard case let GrokWebBillingError.rpcFailed(status, message) = error else { return false } + return status == 9 && message == "No personal team." + } + } + + @Test + func `web strategy publishes identity-only result for team billing`() async throws { + let grokHome = FileManager.default.temporaryDirectory + .appendingPathComponent("CodexBar-GrokTeamFallback-\(UUID().uuidString)", isDirectory: true) + try FileManager.default.createDirectory(at: grokHome, withIntermediateDirectories: true) + defer { try? FileManager.default.removeItem(at: grokHome) } + let auth = #""" + { + "https://auth.x.ai::client": { + "key": "team-token", + "email": "team@example.com", + "team_id": "team-123", + "principal_type": "Team" + } + } + """# + try Data(auth.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, + "GROK_CLI_PATH": grokHome.appendingPathComponent("missing-grok").path, + "PATH": grokHome.path, + ], + settings: nil, + fetcher: UsageFetcher(), + claudeFetcher: ClaudeUsageFetcher(browserDetection: browserDetection), + browserDetection: browserDetection) + + let result = try await GrokWebFetchStrategy().fetch(context) { + throw GrokWebBillingError.teamUsageUnsupported + } + + #expect(result.sourceLabel == "grok-web") + #expect(result.diagnostic == GrokStatusProbe.teamUsageUnavailableMessage) + #expect(result.usage.primary == nil) + #expect(result.usage.accountEmail(for: .grok) == "team@example.com") + #expect(result.usage.accountOrganization(for: .grok) == "team-123") + } + @Test func `ignores grpc web trailer frames`() { let payload = Self.protobufPayload(usedPercent: 12.25, resetEpoch: 1_800_000_001) @@ -635,7 +799,9 @@ struct GrokWebBillingFetcherTests { #expect(attempts.current() == 2) #expect(snapshot.usedPercent == 25) } +} +extension GrokWebBillingFetcherTests { @Test func `web fetch can authenticate with browser cookies`() async throws { defer { @@ -776,7 +942,8 @@ struct GrokWebBillingFetcherTests { oidcIssuer: "https://auth.x.ai", oidcClientId: "client", expiresAt: Date(timeIntervalSince1970: 1_900_000_000), - createTime: Date(timeIntervalSince1970: 1_799_000_000)) + createTime: Date(timeIntervalSince1970: 1_799_000_000), + principalType: "Team") private static func protobufPayload(usedPercent: Float, resetEpoch: UInt64) -> Data { var data = Data() diff --git a/docs/grok.md b/docs/grok.md index 452102be59..3782904901 100644 --- a/docs/grok.md +++ b/docs/grok.md @@ -15,16 +15,21 @@ browser session when the CLI surface does not expose billing. ## Data sources + fallback order -1) **`~/.grok/auth.json` (primary; always works for SuperGrok subscribers)** - - Reads `email`, `team_id`, `first_name`/`last_name`, plan-hint (`auth_mode`) - for the identity row in the menu. +1) **`~/.grok/auth.json` (primary identity source)** + - Reads `email`, `team_id`, `first_name`/`last_name`, plan-hint (`auth_mode`), + and the optional `principal_type` for the identity row in the menu. + - Team principals are recognized on the CLI and web billing paths. Until Grok + exposes a supported team usage surface, CodexBar keeps the identity row and + reports that team usage is unavailable instead of exposing the personal-team + rejection verbatim. 2) **`grok agent stdio` ACP JSON-RPC** (best-effort, currently disabled in grok 0.1.210) - We spawn `grok agent stdio` and call `initialize` + `x.ai/billing` (no params). - **Known limitation:** in grok 0.1.210 the `x.ai/billing` extension method is only wired in the interactive TUI; the agent-stdio surface returns - `-32601 Method not found`. The provider degrades silently to identity-only - when this happens. When xAI exposes billing on the agent protocol, no - code change is required. + `-32601 Method not found`. Personal/unknown principals continue to the web + fallback, while a team principal degrades to identity-only with an explicit + unsupported-team-usage diagnostic. When xAI exposes billing on the agent + protocol, no code change is required. - One non-obvious quirk: grok's ACP parser does not unescape `\/` in method names. `Foundation.JSONSerialization.data` defaults to escaping forward slashes, so payloads must be re-encoded with `\/` → `/` before being @@ -60,6 +65,7 @@ browser session when the CLI surface does not expose billing. `https://accounts.x.ai/sign-in` (legacy session). - Required fields per entry: `key` (bearer token), `refresh_token`, `expires_at`, `auth_mode`, `email`, `team_id`, `user_id`, `first_name`/`last_name`. + `principal_type` is optional because older auth files do not include it. - Tokens are issued by `grok login` and expire after ~7 days; refresh is handled by the CLI itself (CodexBar does not refresh; it just reads the cached credential). From 09b3792e92bf43b24f3fe6ace4f1ddcb35d56fd0 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 16 Jul 2026 18:09:53 +0100 Subject: [PATCH 2/2] refactor: clarify Grok auth state --- .../Providers/Grok/GrokProviderDescriptor.swift | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift index 7090942868..02379d308c 100644 --- a/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift @@ -142,14 +142,14 @@ struct GrokWebFetchStrategy: ProviderFetchStrategy { do { (webBilling, sourceLabel, authenticatedByAuthFile) = try await fetchWebBilling() } catch GrokWebBillingError.teamUsageUnsupported { - guard let credentials = try? GrokCredentialsStore.load(env: context.env), - !credentials.isExpired, - credentials.isTeamPrincipal + guard let authState = try? GrokCredentialsStore.load(env: context.env), + !authState.isExpired, + authState.isTeamPrincipal else { throw GrokWebBillingError.teamUsageUnsupported } let identitySnapshot = GrokStatusProbe.identityOnlySnapshot( - credentials: credentials, + credentials: authState, localSummary: GrokLocalSessionScanner.summarize(env: context.env), cliVersion: GrokStatusProbe.detectVersion(env: context.env)) return self.makeResult(