diff --git a/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift index 0c944e0f13..28df122789 100644 --- a/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/Grok/GrokProviderDescriptor.swift @@ -126,12 +126,21 @@ struct GrokWebFetchStrategy: ProviderFetchStrategy { sourceLabel: String, authenticatedByAuthFile: Bool) + /// Browser-cookie import must stay limited to surfaces where a person explicitly asked for it: + /// the menu-bar app runtime, a `userInitiated` interaction (set only by explicit refresh + /// commands and app UI gestures), or the environment override. Scheduled and background work + /// must keep the default `.background` context so it can never reach Chromium Keychain prompts. static func canImportBrowserCookies(runtime: ProviderRuntime, env: [String: String]) -> Bool { - runtime == .app || env["CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT"] == "1" + runtime == .app || + ProviderInteractionContext.current == .userInitiated || + env["CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT"] == "1" } func isAvailable(_ context: ProviderFetchContext) async -> Bool { #if os(macOS) + if CookieHeaderCache.load(provider: .grok) != nil { + return true + } if Self.canImportBrowserCookies(runtime: context.runtime, env: context.env), GrokCookieImporter.hasSession(browserDetection: context.browserDetection) { @@ -201,13 +210,31 @@ struct GrokWebFetchStrategy: ProviderFetchStrategy { let browserCredentials = try? credentialsResult.get() #if os(macOS) + var cacheObservation = CookieHeaderCache.observeForConditionalMutation(provider: .grok) + var lastCookieError: Error? + if let cached = cacheObservation.entry { + do { + let snapshot = try await Self.fetchValidCookieHeader( + cached.cookieHeader, + credentials: browserCredentials, + preferTrailingAuthenticationFailure: true) + return (snapshot, cached.sourceLabel, false) + } catch { + guard Self.isCookieAuthenticationFailure(error) else { throw error } + if CookieHeaderCache.clearIfCurrent(provider: .grok, expected: cached) { + cacheObservation = cacheObservation.afterOwnedClear() + } + lastCookieError = error + } + } + if Self.canImportBrowserCookies(runtime: context.runtime, env: context.env) { - var lastCookieError: Error? do { let sessions = try GrokCookieImporter.importSessions(browserDetection: context.browserDetection) let (snapshot, sourceLabel) = try await Self.fetchFirstValidCookieSession( sessions, - credentials: browserCredentials) + credentials: browserCredentials, + cacheObservation: cacheObservation) return (snapshot, sourceLabel, false) } catch { lastCookieError = error @@ -242,6 +269,7 @@ struct GrokWebFetchStrategy: ProviderFetchStrategy { static func fetchFirstValidCookieSession( _ sessions: [GrokCookieImporter.SessionInfo], credentials: GrokCredentials? = nil, + cacheObservation: CookieHeaderCache.ConditionalMutationObservation? = nil, fetch: ((String, GrokCredentials?) async throws -> GrokWebBillingSnapshot)? = nil) async throws -> (GrokWebBillingSnapshot, String) { @@ -253,25 +281,83 @@ struct GrokWebFetchStrategy: ProviderFetchStrategy { 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 + do { + let snapshot = try await Self.fetchValidCookieHeader( + session.cookieHeader, + credentials: credentials, + fetch: fetchSnapshot) + if let cacheObservation { + CookieHeaderCache.storeIfObservationCurrent( + provider: .grok, + expected: cacheObservation, + cookieHeader: session.cookieHeader, + sourceLabel: session.sourceLabel) + } + return (snapshot, session.sourceLabel) + } catch { + if case GrokWebBillingError.teamUsageUnsupported = error { + teamUsageUnsupportedError = error } + lastError = error } } throw teamUsageUnsupportedError ?? lastError ?? GrokWebBillingError.missingCredentials } + /// `preferTrailingAuthenticationFailure` lets a cached-cookie caller surface a trailing + /// 401/403 over the team classification so stale sessions still trigger cache eviction. + /// Non-authentication trailing errors keep `teamUsageUnsupported` so team principals + /// degrade to identity-only data instead of failing outright. + static func fetchValidCookieHeader( + _ cookieHeader: String, + credentials: GrokCredentials? = nil, + preferTrailingAuthenticationFailure: Bool = false, + fetch: ((String, GrokCredentials?) async throws -> GrokWebBillingSnapshot)? = nil) async throws + -> GrokWebBillingSnapshot + { + let fetchSnapshot = fetch ?? { cookieHeader, credentials in + try await GrokWebBillingFetcher.fetch( + cookieHeader: cookieHeader, + credentials: credentials) + } + var lastError: Error? + var teamUsageUnsupportedError: Error? + for authCredentials in Self.cookieAuthAttempts(credentials: credentials) { + do { + return try await fetchSnapshot(cookieHeader, authCredentials) + } catch { + if case GrokWebBillingError.teamUsageUnsupported = error { + teamUsageUnsupportedError = error + } + lastError = error + } + } + if let teamUsageUnsupportedError { + let trailingAuthenticationFailure = preferTrailingAuthenticationFailure + && lastError.map(Self.isCookieAuthenticationFailure) == true + if !trailingAuthenticationFailure { + throw teamUsageUnsupportedError + } + } + throw lastError ?? GrokWebBillingError.missingCredentials + } + static func cookieAuthAttempts(credentials: GrokCredentials?) -> [GrokCredentials?] { guard let credentials, !credentials.isExpired else { return [nil] } return [credentials, nil] } + + static func isCookieAuthenticationFailure(_ error: Error) -> Bool { + guard let error = error as? GrokWebBillingError else { return false } + switch error { + case let .requestFailed(status, _): + return status == 401 || status == 403 + case let .rpcFailed(status, message): + return GrokWebBillingError.isAuthenticationFailure(status: status, message: message) + case .missingCredentials, .emptyResponse, .invalidResponse, .teamUsageUnsupported, .parseFailed: + return false + } + } #endif func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { diff --git a/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift b/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift index 73009016d2..afe19ac254 100644 --- a/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift +++ b/Tests/CodexBarTests/GrokWebBillingFetcherTests.swift @@ -54,15 +54,6 @@ struct GrokWebBillingFetcherTests { #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")) @@ -999,6 +990,142 @@ extension GrokWebBillingFetcherTests { } } +// MARK: - Browser cookie cache + +extension GrokWebBillingFetcherTests { + @Test + func `cli runtime imports browser cookies only when 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"])) + let userInitiated = ProviderInteractionContext.$current.withValue(.userInitiated) { + GrokWebFetchStrategy.canImportBrowserCookies(runtime: .cli, env: [:]) + } + #expect(userInitiated) + } + + @Test + func `web strategy is available from a cached browser session`() async { + let service = "com.steipete.codexbar.tests.grok-availability.\(UUID().uuidString)" + CookieHeaderCache.resetDisplayCacheForTesting() + defer { CookieHeaderCache.resetDisplayCacheForTesting() } + await KeychainCacheStore.withServiceOverrideForTesting(service) { + await KeychainCacheStore.withImplicitTestStoreForTesting { + CookieHeaderCache.store( + provider: .grok, + cookieHeader: "sso=cached-session", + sourceLabel: "Chrome") + let grokHome = FileManager.default.temporaryDirectory + .appendingPathComponent( + "CodexBar-GrokCachedAvailability-\(UUID().uuidString)", + isDirectory: true) + 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) + + #expect(await GrokWebFetchStrategy().isAvailable(context)) + } + } + } + + @Test + func `validated browser session stages a cache replacement for explicit refresh`() async throws { + let service = "com.steipete.codexbar.tests.grok-refresh.\(UUID().uuidString)" + CookieHeaderCache.resetDisplayCacheForTesting() + defer { CookieHeaderCache.resetDisplayCacheForTesting() } + try await KeychainCacheStore.withServiceOverrideForTesting(service) { + try await KeychainCacheStore.withImplicitTestStoreForTesting { + let gate = try #require(CookieHeaderCache.beginRefreshReadSuppression(provider: .grok)) + defer { CookieHeaderCache.endRefreshReadSuppression(gate) } + let observation = CookieHeaderCache.observeForConditionalMutation(provider: .grok) + let cookie = try #require(Self.cookie(name: "sso", value: "fresh-session")) + let sessions = [GrokCookieImporter.SessionInfo(cookies: [cookie], sourceLabel: "Chrome")] + + let result = try await GrokWebFetchStrategy.fetchFirstValidCookieSession( + sessions, + cacheObservation: observation) + { _, _ in + GrokWebBillingSnapshot(usedPercent: 7, resetsAt: nil) + } + + #expect(result.0.usedPercent == 7) + #expect(CookieHeaderCache.load(provider: .grok)?.cookieHeader == "sso=fresh-session") + let commit = CookieHeaderCache.commitRefreshReadSuppression(gate) + #expect(commit == CookieRefreshCommitSummary(stagedCount: 1, committedCount: 1, failedCount: 0)) + } + } + } + + @Test + func `cached cookie eviction is limited to authentication failures`() { + #expect(GrokWebFetchStrategy.isCookieAuthenticationFailure( + GrokWebBillingError.requestFailed(401, "expired"))) + #expect(GrokWebFetchStrategy.isCookieAuthenticationFailure( + GrokWebBillingError.rpcFailed(16, "unauthenticated"))) + #expect(!GrokWebFetchStrategy.isCookieAuthenticationFailure( + GrokWebBillingError.requestFailed(503, "unavailable"))) + #expect(!GrokWebFetchStrategy.isCookieAuthenticationFailure( + GrokWebBillingError.parseFailed)) + } + + @Test + func `cached team cookie surfaces trailing authentication failure`() async throws { + var attempts: [String] = [] + + await #expect { + _ = try await GrokWebFetchStrategy.fetchValidCookieHeader( + "sso=stale", + credentials: Self.credentials, + preferTrailingAuthenticationFailure: true) + { _, authCredentials in + if authCredentials != nil { + attempts.append("cookie+bearer") + throw GrokWebBillingError.teamUsageUnsupported + } + attempts.append("cookie-only") + throw GrokWebBillingError.requestFailed(401, "expired") + } + } throws: { error in + GrokWebFetchStrategy.isCookieAuthenticationFailure(error) + } + + #expect(attempts == ["cookie+bearer", "cookie-only"]) + } + + @Test + func `cached team cookie keeps team classification for non-auth trailing errors`() async { + await #expect { + _ = try await GrokWebFetchStrategy.fetchValidCookieHeader( + "sso=team-session", + credentials: Self.credentials, + preferTrailingAuthenticationFailure: true) + { _, authCredentials in + if authCredentials != nil { + throw GrokWebBillingError.teamUsageUnsupported + } + throw GrokWebBillingError.rpcFailed(9, "No personal team") + } + } throws: { error in + if case GrokWebBillingError.teamUsageUnsupported = error { + return true + } + return false + } + } +} + final class GrokWebBillingStubURLProtocol: URLProtocol { nonisolated(unsafe) static var requests: [URLRequest] = [] nonisolated(unsafe) static var requestBodies: [Data?] = [] diff --git a/docs/grok.md b/docs/grok.md index 3782904901..a51fafe4cf 100644 --- a/docs/grok.md +++ b/docs/grok.md @@ -43,8 +43,14 @@ browser session when the CLI surface does not expose billing. browser session, then retries that session with cookies only. - CodexBar imports Chrome only by default to avoid unrelated browser Keychain prompts. - - CLI/test runtime does not import browser cookies unless - `CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT=1` is set. + - Ordinary CLI/test runtime does not import browser cookies unless + `CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT=1` is set. An explicit + `codexbar cookie refresh --provider grok` also opts in for that refresh. + - Validated sessions are stored in the Keychain-backed cookie cache and are + reused first by later app and CLI fetches, so background work does not + re-open the Chromium Keychain gate. The cached cookie is evicted only on + authentication failures (HTTP 401/403 or gRPC auth statuses); a cached + team-limited session keeps degrading to identity-only data. - `~/.grok/auth.json` is still used for identity and as a last best-effort bearer-only probe after browser sessions fail. Expired tokens are not sent. - Parses the returned protobuf enough to recover used percent and diff --git a/docs/providers.md b/docs/providers.md index 2be0fdeb74..fa5780d530 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -445,7 +445,10 @@ scan fails, while provider/account configuration changes replace obsolete result - `grok agent stdio` (ACP) JSON-RPC `x.ai/billing` method; requires `grok login` (SuperGrok OAuth/OIDC). - Reads cached credentials from `~/.grok/auth.json` for identity (email, team). - Falls back to grok.com's billing gRPC-web endpoint via Chrome session cookies when the CLI does not expose billing. -- CLI/test runs do not import browser cookies unless `CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT=1` is set. +- Ordinary CLI/test runs do not import browser cookies unless `CODEXBAR_ALLOW_BROWSER_COOKIE_IMPORT=1` is set; + `codexbar cookie refresh --provider grok` opts in for its explicit refresh. +- Validated sessions are cached in the Keychain cookie cache and reused before any new browser import; + the cache is evicted only on authentication failures. - Local fallback aggregates `~/.grok/sessions/**/signals.json` token counts when the RPC is unavailable. - Status: link only to `https://status.x.ai` (no auto-polling yet). - Details: `docs/grok.md`.