diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 5565cea7eb..7451f281ef 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -26,6 +26,7 @@ use self::accounts::{GrokAuthKind, ParsedGrokAuthFile}; use self::billing::GrokBillingSnapshot; const BILLING_ENDPOINT: &str = "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig"; +const BILLING_REQUEST_BODY: [u8; 7] = [0, 0, 0, 0, 2, 0x08, 0]; const REMAINING_RESETS_ENDPOINT: &str = "https://grok.com/prod_mc_billing.ConsumerUiSvc/GetRemainingResets"; const CLI_SETTINGS_ENDPOINT: &str = "https://cli-chat-proxy.grok.com/v1/settings"; @@ -35,6 +36,7 @@ const RESET_CREDITS_JOIN_GRACE: Duration = Duration::from_millis(250); pub struct GrokProvider { metadata: ProviderMetadata, client: Client, + billing_endpoint: String, } impl GrokProvider { @@ -57,6 +59,7 @@ impl GrokProvider { .timeout(std::time::Duration::from_secs(15)) .build() .unwrap_or_else(|_| Client::new()), + billing_endpoint: BILLING_ENDPOINT.to_string(), } } @@ -74,6 +77,12 @@ impl GrokProvider { self.client.clone() } + #[cfg(test)] + fn with_billing_endpoint_for_tests(mut self, endpoint: String) -> Self { + self.billing_endpoint = endpoint; + self + } + fn load_credentials(kind: GrokAuthKind) -> Result { let path = Self::auth_file_path() .ok_or_else(|| ProviderError::NotInstalled("Grok auth path not found".to_string()))?; @@ -210,6 +219,29 @@ impl GrokProvider { }) } + async fn fetch_with_oauth_fallback( + &self, + credentials: &GrokCredentials, + ctx: &FetchContext, + ) -> Result { + match self + .fetch_with_auth(credentials, GrokAuthKind::OAuth, ctx) + .await + { + Ok(result) => Ok(result), + Err(ProviderError::AuthRequired) => { + if let Some(token) = ctx.api_key.as_deref() { + let fallback = GrokCredentials::from_bearer(token); + self.fetch_with_auth(&fallback, GrokAuthKind::OAuth, ctx) + .await + } else { + Err(ProviderError::AuthRequired) + } + } + Err(error) => Err(error), + } + } + async fn fetch_auto(&self, ctx: &FetchContext) -> Result { let allow_browser_cookie_fallback = !ctx.auto_prefer_web; for step in grok_auto_steps( @@ -304,8 +336,8 @@ impl GrokProvider { ) -> Result { let mut request = self .client - .post(BILLING_ENDPOINT) - .body(vec![0, 0, 0, 0, 0]) + .post(&self.billing_endpoint) + .body(BILLING_REQUEST_BODY.to_vec()) .header("Origin", "https://grok.com") .header("Referer", "https://grok.com/?_s=usage") .header("Accept", "*/*") @@ -534,22 +566,7 @@ impl Provider for GrokProvider { GrokCredentials::from_bearer(token) } }; - match self - .fetch_with_auth(&credentials, GrokAuthKind::OAuth, ctx) - .await - { - Ok(result) => Ok(result), - Err(ProviderError::AuthRequired) => { - if let Some(token) = ctx.api_key.as_deref() { - let fallback = GrokCredentials::from_bearer(token); - self.fetch_with_auth(&fallback, GrokAuthKind::OAuth, ctx) - .await - } else { - Err(ProviderError::AuthRequired) - } - } - Err(error) => Err(error), - } + self.fetch_with_oauth_fallback(&credentials, ctx).await } } } diff --git a/rust/src/providers/grok/tests.rs b/rust/src/providers/grok/tests.rs index 71691756b9..586c20c637 100644 --- a/rust/src/providers/grok/tests.rs +++ b/rust/src/providers/grok/tests.rs @@ -1,6 +1,81 @@ use super::*; use chrono::TimeZone; +fn billing_response_with_percent(percent: f32) -> Vec { + let mut payload = vec![0x0a, 0x05, 0x0d]; + payload.extend(percent.to_le_bytes()); + + let mut response = vec![0x00]; + response.extend( + u32::try_from(payload.len()) + .expect("test gRPC-web payload length fits u32") + .to_be_bytes(), + ); + response.extend(payload); + response +} + +#[tokio::test] +async fn billing_request_encodes_explicit_false_in_a_nonempty_grpc_web_frame() { + let mut server = mockito::Server::new_async().await; + let endpoint = format!("{}/billing", server.url()); + let request = server + .mock("POST", "/billing") + .match_body(BILLING_REQUEST_BODY.to_vec()) + .with_status(200) + .with_body(billing_response_with_percent(37.0)) + .create_async() + .await; + let provider = GrokProvider::new().with_billing_endpoint_for_tests(endpoint); + + let snapshot = provider + .fetch_billing(Some("Bearer local-token".to_string()), None) + .await + .unwrap(); + + request.assert_async().await; + assert_eq!(BILLING_REQUEST_BODY, [0, 0, 0, 0, 2, 0x08, 0]); + assert_eq!(snapshot.used_percent, Some(37.0)); +} + +#[tokio::test] +async fn oauth_billing_auth_failure_falls_back_to_configured_local_token() { + let mut server = mockito::Server::new_async().await; + let endpoint = format!("{}/billing", server.url()); + let rejected = server + .mock("POST", "/billing") + .match_header("authorization", "Bearer expired-token") + .match_body(BILLING_REQUEST_BODY.to_vec()) + .with_status(401) + .create_async() + .await; + let fallback = server + .mock("POST", "/billing") + .match_header("authorization", "Bearer local-token") + .match_body(BILLING_REQUEST_BODY.to_vec()) + .with_status(200) + .with_body(billing_response_with_percent(42.0)) + .create_async() + .await; + let provider = GrokProvider::new().with_billing_endpoint_for_tests(endpoint); + let credentials = GrokCredentials::from_bearer("expired-token"); + let context = FetchContext { + include_credits: false, + api_key: Some("local-token".to_string()), + ..FetchContext::default() + }; + + let result = provider + .fetch_with_oauth_fallback(&credentials, &context) + .await + .unwrap(); + + rejected.assert_async().await; + fallback.assert_async().await; + assert_eq!(result.source_label, "grok-oauth"); + assert_eq!(result.usage.primary.used_percent, 42.0); +} + #[test] fn grok_plan_prefers_subscription_tier_display_names() { assert_eq!(