[0.64.0] Preserve Grok billing defaults

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<GrokCredentials, ProviderError> {
         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<ProviderFetchResult, ProviderError> {
+        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<ProviderFetchResult, ProviderError> {
         let allow_browser_cookie_fallback = !ctx.auto_prefer_web;
         for step in grok_auto_steps(
@@ -304,8 +336,8 @@ impl GrokProvider {
     ) -> Result<GrokBillingSnapshot, ProviderError> {
         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
             }
         }
     }

rust/src/providers/grok/tests.rs

@@ -1,6 +1,81 @@
 use super::*;
 use chrono::TimeZone;
 
+fn billing_response_with_percent(percent: f32) -> Vec<u8> {
+    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!(