From bdab6092fdb7d21eca52ad05c563d8affb28d560 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 19 Sep 2026 21:36:38 +0700 Subject: [PATCH 1/2] Port Grok usage reset coupons --- rust/src/providers/grok/billing.rs | 288 +++++++++++++++++++++++++++++ rust/src/providers/grok/mod.rs | 92 ++++++++- 2 files changed, 375 insertions(+), 5 deletions(-) diff --git a/rust/src/providers/grok/billing.rs b/rust/src/providers/grok/billing.rs index b69e9dec3c..035b780602 100644 --- a/rust/src/providers/grok/billing.rs +++ b/rust/src/providers/grok/billing.rs @@ -12,6 +12,15 @@ pub(super) struct GrokBillingSnapshot { pub(super) window_minutes: Option, } +/// One unused SuperGrok usage-limit reset coupon. The token ID is retained only +/// while parsing and is never attached to a public or persisted snapshot. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) struct GrokResetCoupon { + pub(super) token_id: String, + pub(super) granted_at: Option>, + pub(super) expires_at: DateTime, +} + pub(super) fn validate_grpc_headers(headers: &HeaderMap) -> Result<(), ProviderError> { if let Some(status) = headers .get("grpc-status") @@ -33,6 +42,231 @@ pub(super) fn parse_grpc_web_response(data: &[u8]) -> Result, +) -> Result, ProviderError> { + let mut payloads = grpc_web_data_frames(data); + if data.is_empty() || data == [0, 0, 0, 0, 0] { + return Ok(Vec::new()); + } + if payloads.is_empty() && looks_like_protobuf_payload(data) { + payloads.push(data.to_vec()); + } + if payloads.is_empty() { + return Err(ProviderError::Parse( + "Grok reset-credit response had no payload".to_string(), + )); + } + + let mut coupons = Vec::new(); + for payload in payloads { + parse_reset_coupon_container(&payload, now, &mut coupons)?; + } + coupons.sort_by_key(|coupon| coupon.expires_at); + Ok(coupons) +} + +fn parse_reset_coupon_container( + data: &[u8], + now: DateTime, + coupons: &mut Vec, +) -> Result<(), ProviderError> { + let mut index = 0; + while index < data.len() { + let (field, wire, next) = read_key(data, index).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit protobuf is malformed".to_string()) + })?; + index = next; + if field == 10 { + if wire != 2 { + return Err(ProviderError::Parse( + "Grok reset-credit record has an invalid wire type".to_string(), + )); + } + let (len, payload_start) = read_varint(data, index).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit record length is malformed".to_string()) + })?; + let len = usize::try_from(len).map_err(|_| { + ProviderError::Parse("Grok reset-credit record is too large".to_string()) + })?; + let payload_end = payload_start.checked_add(len).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit record length overflowed".to_string()) + })?; + if payload_end > data.len() { + return Err(ProviderError::Parse( + "Grok reset-credit record is truncated".to_string(), + )); + } + if let Some(coupon) = parse_reset_coupon(&data[payload_start..payload_end], now)? { + coupons.push(coupon); + } + index = payload_end; + } else { + index = skip_field(data, index, wire).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit protobuf is malformed".to_string()) + })?; + } + } + Ok(()) +} + +fn parse_reset_coupon( + data: &[u8], + now: DateTime, +) -> Result, ProviderError> { + let mut index = 0; + let mut token_id = None; + let mut granted_at = None; + let mut expires_at = None; + while index < data.len() { + let (field, wire, next) = read_key(data, index).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit record is malformed".to_string()) + })?; + index = next; + match field { + 10 => { + if wire != 2 { + return Err(ProviderError::Parse( + "Grok reset-credit token id has an invalid wire type".to_string(), + )); + } + let (len, start) = read_varint(data, index).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit token id is malformed".to_string()) + })?; + let len = usize::try_from(len).map_err(|_| { + ProviderError::Parse("Grok reset-credit token id is too large".to_string()) + })?; + let end = start.checked_add(len).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit token id length overflowed".to_string()) + })?; + if end > data.len() { + return Err(ProviderError::Parse( + "Grok reset-credit token id is truncated".to_string(), + )); + } + token_id = Some( + std::str::from_utf8(&data[start..end]) + .map_err(|_| { + ProviderError::Parse( + "Grok reset-credit token id is not UTF-8".to_string(), + ) + })? + .to_string(), + ); + index = end; + } + 20 | 30 => { + if wire != 2 { + return Err(ProviderError::Parse( + "Grok reset-credit timestamp has an invalid wire type".to_string(), + )); + } + let (len, start) = read_varint(data, index).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit timestamp is malformed".to_string()) + })?; + let len = usize::try_from(len).map_err(|_| { + ProviderError::Parse("Grok reset-credit timestamp is too large".to_string()) + })?; + let end = start.checked_add(len).ok_or_else(|| { + ProviderError::Parse( + "Grok reset-credit timestamp length overflowed".to_string(), + ) + })?; + if end > data.len() { + return Err(ProviderError::Parse( + "Grok reset-credit timestamp is truncated".to_string(), + )); + } + let timestamp = parse_timestamp_message(&data[start..end])?; + if field == 20 { + granted_at = timestamp; + } else { + expires_at = timestamp; + } + index = end; + } + _ => { + index = skip_field(data, index, wire).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit record is malformed".to_string()) + })?; + } + } + } + + let Some(token_id) = token_id.filter(|id| !id.trim().is_empty()) else { + return Ok(None); + }; + let Some(expires_at) = expires_at.filter(|expires_at| *expires_at > now) else { + return Ok(None); + }; + Ok(Some(GrokResetCoupon { + token_id, + granted_at, + expires_at, + })) +} + +fn parse_timestamp_message(data: &[u8]) -> Result>, ProviderError> { + let mut index = 0; + let mut seconds = None; + while index < data.len() { + let (field, wire, next) = read_key(data, index).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit timestamp is malformed".to_string()) + })?; + index = next; + if field == 1 { + if wire != 0 { + return Err(ProviderError::Parse( + "Grok reset-credit timestamp seconds has an invalid wire type".to_string(), + )); + } + let (value, next) = read_varint(data, index).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit timestamp seconds is malformed".to_string()) + })?; + seconds = Some(value); + index = next; + } else { + index = skip_field(data, index, wire).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit timestamp is malformed".to_string()) + })?; + } + } + let Some(seconds) = seconds else { + return Ok(None); + }; + #[allow( + clippy::cast_possible_wrap, + reason = "provider timestamps are bounded to the Unix-seconds range" + )] + let seconds = seconds as i64; + Ok((1_700_000_000..=2_100_000_000) + .contains(&seconds) + .then(|| Utc.timestamp_opt(seconds, 0).single()) + .flatten()) +} + +fn skip_field(data: &[u8], index: usize, wire: u64) -> Option { + match wire { + 0 => read_varint(data, index).map(|(_, next)| next), + 1 => index.checked_add(8).filter(|end| *end <= data.len()), + 2 => { + let (len, start) = read_varint(data, index)?; + let len = usize::try_from(len).ok()?; + start.checked_add(len).filter(|end| *end <= data.len()) + } + 5 => index.checked_add(4).filter(|end| *end <= data.len()), + _ => None, + } +} + fn parse_grpc_web_response_at( data: &[u8], now: DateTime, @@ -633,6 +867,60 @@ mod tests { } } + #[test] + fn reset_coupons_filter_expired_records_and_sort_by_expiry() { + let now = fixed_time(1_800_000_000); + let mut payload = Vec::new(); + payload.extend(reset_coupon_record("later", 1_900_000_000)); + payload.extend(reset_coupon_record("expired", 1_700_000_000)); + payload.extend(reset_coupon_record("earlier", 1_850_000_000)); + payload.extend(reset_coupon_record("", 1_950_000_000)); + + let coupons = parse_grpc_web_reset_coupons(&payload, now).unwrap(); + + assert_eq!( + coupons + .iter() + .map(|coupon| coupon.token_id.as_str()) + .collect::>(), + ["earlier", "later"] + ); + assert!(coupons.iter().all(|coupon| coupon.expires_at > now)); + } + + #[test] + fn reset_coupon_empty_payload_is_valid_and_malformed_payload_is_atomic() { + assert!( + parse_grpc_web_reset_coupons(&[0, 0, 0, 0, 0], fixed_time(1_800_000_000)) + .unwrap() + .is_empty() + ); + + let mut malformed = reset_coupon_record("valid", 1_900_000_000); + malformed.extend([0x52, 0x05, b'a']); + assert!(parse_grpc_web_reset_coupons(&malformed, fixed_time(1_800_000_000)).is_err()); + } + + #[test] + fn reset_coupon_truncated_timestamp_is_rejected() { + let mut record = length_field(10, b"valid"); + record.extend([0xf2, 0x01, 0x01, 0x08]); + let payload = length_field(10, &record); + + assert!(parse_grpc_web_reset_coupons(&payload, fixed_time(1_800_000_000)).is_err()); + } + + fn reset_coupon_record(token_id: &str, expires_at: u64) -> Vec { + let timestamp = { + let mut bytes = vec![0x08]; + bytes.extend(varint(expires_at)); + bytes + }; + let mut record = length_field(10, token_id.as_bytes()); + record.extend(length_field(30, ×tamp)); + length_field(10, &record) + } + fn fixed_time(seconds: i64) -> DateTime { Utc.timestamp_opt(seconds, 0).single().unwrap() } diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index a6fba75287..ac50997497 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -17,14 +17,16 @@ use std::path::PathBuf; use std::os::windows::process::CommandExt; use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderInventoryItem, + ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; use self::accounts::{GrokAuthKind, ParsedGrokAuthFile}; use self::billing::GrokBillingSnapshot; const BILLING_ENDPOINT: &str = "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig"; +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"; pub struct GrokProvider { @@ -103,7 +105,10 @@ impl GrokProvider { None } .or_else(|| credentials.login_method()); - Ok(result_from_billing( + let reset_credits = self + .fetch_remaining_resets(Some(&credentials.access_token), None) + .await; + let result = result_from_billing( billing, if kind == GrokAuthKind::Cli { "grok-cli" @@ -113,7 +118,11 @@ impl GrokProvider { credentials.email.clone(), credentials.team_id.clone(), plan, - )) + ); + Ok(match reset_credits { + Some(credits) => result.with_inventory_item(credits), + None => result, + }) } async fn fetch_cli_subscription_tier(&self, credentials: &GrokCredentials) -> Option { @@ -149,10 +158,15 @@ impl GrokProvider { let billing = self .fetch_billing(None, Some(cookie_header.to_string())) .await?; + let reset_credits = self.fetch_remaining_resets(None, Some(cookie_header)).await; // v0.56.0: a browser session is its own principal. Never enrich a // successful cookie billing result from ambient auth.json metadata, // which may belong to a different account or change during the fetch. - Ok(result_from_cookie_billing(billing)) + let result = result_from_cookie_billing(billing); + Ok(match reset_credits { + Some(credits) => result.with_inventory_item(credits), + None => result, + }) } async fn fetch_auto(&self, ctx: &FetchContext) -> Result { @@ -279,6 +293,74 @@ impl GrokProvider { billing::parse_grpc_web_response(&bytes) } + /// Fetch optional SuperGrok reset-credit inventory using the same principal + /// that produced the successful billing result. This is deliberately + /// best-effort: billing remains valid when this secondary endpoint is down, + /// malformed, unauthorized, or empty. + async fn fetch_remaining_resets( + &self, + access_token: Option<&str>, + cookie_header: Option<&str>, + ) -> Option { + let mut request = self + .client + .post(REMAINING_RESETS_ENDPOINT) + .body(vec![0, 0, 0, 0, 0]) + .timeout(std::time::Duration::from_secs(2)) + .header("Origin", "https://grok.com") + .header("Referer", "https://grok.com/?_s=usage") + .header("Accept", "*/*") + .header("Content-Type", "application/grpc-web+proto") + .header("x-grpc-web", "1") + .header("x-user-agent", "connect-es/2.1.1") + .header("User-Agent", "CodexBar"); + if let Some(access_token) = access_token { + request = request.header("Authorization", format!("Bearer {access_token}")); + } + if let Some(cookie_header) = cookie_header { + request = request.header("Cookie", cookie_header); + } + + let response = match request.send().await { + Ok(response) => response, + Err(error) => { + tracing::debug!("Grok reset-credit lookup failed: {error}"); + return None; + } + }; + if !response.status().is_success() { + tracing::debug!(status = %response.status(), "Grok reset-credit lookup returned a non-success status"); + return None; + } + let headers = response.headers().clone(); + let bytes = match response.bytes().await { + Ok(bytes) => bytes, + Err(error) => { + tracing::debug!("Grok reset-credit response read failed: {error}"); + return None; + } + }; + if let Err(error) = billing::validate_grpc_headers(&headers) { + tracing::debug!("Grok reset-credit RPC failed: {error}"); + return None; + } + let coupons = match billing::parse_grpc_web_reset_coupons(&bytes, Utc::now()) { + Ok(coupons) => coupons, + Err(error) => { + tracing::debug!("Grok reset-credit response was invalid: {error}"); + return None; + } + }; + let next_expiry = coupons.first().map(|coupon| coupon.expires_at); + let available_count = u32::try_from(coupons.len()).ok()?; + (!coupons.is_empty()).then_some(ProviderInventoryItem { + id: "reset-credits".to_string(), + title: "Limit Reset Credits".to_string(), + available_count, + next_expires_at: next_expiry, + }) + } + fn detect_cli_version() -> Option { let mut command = std::process::Command::new("grok"); command.arg("--version"); From 64a583ecd540c9c56adbe7be87eb2445c4b98003 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 19 Sep 2026 22:03:35 +0700 Subject: [PATCH 2/2] Harden Grok reset credit fetching --- rust/src/providers/grok/billing.rs | 138 ++++++++++++++++++++++++++--- rust/src/providers/grok/mod.rs | 138 +++++++++++++++++++++++------ rust/src/providers/grok/tests.rs | 11 +++ 3 files changed, 248 insertions(+), 39 deletions(-) diff --git a/rust/src/providers/grok/billing.rs b/rust/src/providers/grok/billing.rs index 035b780602..ea711c5c36 100644 --- a/rust/src/providers/grok/billing.rs +++ b/rust/src/providers/grok/billing.rs @@ -53,18 +53,7 @@ pub(super) fn parse_grpc_web_reset_coupons( data: &[u8], now: DateTime, ) -> Result, ProviderError> { - let mut payloads = grpc_web_data_frames(data); - if data.is_empty() || data == [0, 0, 0, 0, 0] { - return Ok(Vec::new()); - } - if payloads.is_empty() && looks_like_protobuf_payload(data) { - payloads.push(data.to_vec()); - } - if payloads.is_empty() { - return Err(ProviderError::Parse( - "Grok reset-credit response had no payload".to_string(), - )); - } + let payloads = grpc_web_reset_payloads(data)?; let mut coupons = Vec::new(); for payload in payloads { @@ -74,6 +63,97 @@ pub(super) fn parse_grpc_web_reset_coupons( Ok(coupons) } +fn grpc_web_reset_payloads(data: &[u8]) -> Result>, ProviderError> { + if data.is_empty() || data == [0, 0, 0, 0, 0] { + return Ok(Vec::new()); + } + + // A raw protobuf payload is retained as a compatibility fallback for the + // captured endpoint fixtures. Valid protobuf keys cannot begin with a + // gRPC-web data/trailer flag, so a leading 0/0x80 unambiguously selects + // framed parsing and makes truncated frames fail closed. + let is_framed = data.first().is_some_and(|flag| *flag == 0 || *flag == 0x80); + if !is_framed { + return looks_like_protobuf_payload(data) + .then(|| vec![data.to_vec()]) + .ok_or_else(|| { + ProviderError::Parse("Grok reset-credit response had no payload".to_string()) + }); + } + + let mut payloads = Vec::new(); + let mut index = 0; + while index < data.len() { + if index + 5 > data.len() { + return Err(ProviderError::Parse( + "Grok reset-credit gRPC-web frame is truncated".to_string(), + )); + } + let flags = data[index]; + let len = u32::from_be_bytes([ + data[index + 1], + data[index + 2], + data[index + 3], + data[index + 4], + ]) as usize; + let start = index + 5; + let end = start.checked_add(len).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit gRPC-web frame is too large".to_string()) + })?; + if end > data.len() { + return Err(ProviderError::Parse( + "Grok reset-credit gRPC-web frame is truncated".to_string(), + )); + } + let payload = &data[start..end]; + if flags & 0x80 != 0 { + validate_grpc_web_reset_trailer(payload)?; + } else { + payloads.push(payload.to_vec()); + } + index = end; + } + Ok(payloads) +} + +fn validate_grpc_web_reset_trailer(payload: &[u8]) -> Result<(), ProviderError> { + let text = std::str::from_utf8(payload).map_err(|_| { + ProviderError::Parse("Grok reset-credit gRPC-web trailer is not UTF-8".to_string()) + })?; + let mut grpc_status = None; + for line in text + .split(['\r', '\n']) + .filter(|line| !line.trim().is_empty()) + { + let Some((name, value)) = line.split_once(':') else { + continue; + }; + if name.trim().eq_ignore_ascii_case("grpc-status") { + let status = value.trim().parse::().map_err(|_| { + ProviderError::Parse("Grok reset-credit gRPC status is malformed".to_string()) + })?; + if grpc_status.is_some_and(|previous| previous != status) { + return Err(ProviderError::Parse( + "Grok reset-credit gRPC status is conflicting".to_string(), + )); + } + grpc_status = Some(status); + } + } + let status = grpc_status.ok_or_else(|| { + ProviderError::Parse("Grok reset-credit gRPC status is missing".to_string()) + })?; + if status != 0 { + if status == 16 { + return Err(ProviderError::AuthRequired); + } + return Err(ProviderError::Other(format!( + "Grok reset-credit RPC failed with status {status}" + ))); + } + Ok(()) +} + fn parse_reset_coupon_container( data: &[u8], now: DateTime, @@ -910,6 +990,29 @@ mod tests { assert!(parse_grpc_web_reset_coupons(&payload, fixed_time(1_800_000_000)).is_err()); } + #[test] + fn reset_coupon_nonzero_grpc_web_trailer_is_rejected() { + let payload = reset_coupon_record("valid", 1_900_000_000); + let mut framed = grpc_web_frame(0, &payload); + framed.extend(grpc_web_frame(0x80, b"grpc-status: 13\r\n")); + + assert!(matches!( + parse_grpc_web_reset_coupons(&framed, fixed_time(1_800_000_000)), + Err(ProviderError::Other(message)) if message.contains("status 13") + )); + } + + #[test] + fn reset_coupon_zero_grpc_web_trailer_is_accepted() { + let payload = reset_coupon_record("valid", 1_900_000_000); + let mut framed = grpc_web_frame(0, &payload); + framed.extend(grpc_web_frame(0x80, b"grpc-status: 0\r\n")); + + let coupons = parse_grpc_web_reset_coupons(&framed, fixed_time(1_800_000_000)).unwrap(); + assert_eq!(coupons.len(), 1); + assert_eq!(coupons[0].token_id, "valid"); + } + fn reset_coupon_record(token_id: &str, expires_at: u64) -> Vec { let timestamp = { let mut bytes = vec![0x08]; @@ -955,6 +1058,17 @@ mod tests { encoded } + fn grpc_web_frame(flags: u8, payload: &[u8]) -> Vec { + let mut frame = vec![flags]; + frame.extend( + u32::try_from(payload.len()) + .expect("test gRPC-web payload length fits u32") + .to_be_bytes(), + ); + frame.extend(payload); + frame + } + fn fixed32_field(value: f32) -> Vec { let mut encoded = vec![0x0d]; encoded.extend(value.to_le_bytes()); diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index ac50997497..1cc4dfb2c9 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -12,6 +12,7 @@ use chrono::{DateTime, Utc}; use reqwest::Client; use serde_json::Value; use std::path::PathBuf; +use std::time::{Duration, Instant}; #[cfg(windows)] use std::os::windows::process::CommandExt; @@ -28,6 +29,8 @@ const BILLING_ENDPOINT: &str = "https://grok.com/grok_api_v2.GrokBuildBilling/Ge 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"; +const RESET_CREDITS_TIMEOUT: Duration = Duration::from_secs(2); +const RESET_CREDITS_JOIN_GRACE: Duration = Duration::from_millis(250); pub struct GrokProvider { metadata: ProviderMetadata, @@ -87,7 +90,9 @@ impl GrokProvider { GrokAuthKind::Cli, ) }; - let result = self.fetch_with_auth(&credentials, kind).await?; + let result = self + .fetch_with_auth(&credentials, kind, &FetchContext::default()) + .await?; Ok(account_usage_from_result(&result)) } @@ -95,19 +100,28 @@ impl GrokProvider { &self, credentials: &GrokCredentials, kind: GrokAuthKind, + ctx: &FetchContext, ) -> Result { - let billing = self + let reset_task = + self.spawn_remaining_resets(ctx, Some(credentials.access_token.clone()), None); + let billing = match self .fetch_billing(Some(format!("Bearer {}", credentials.access_token)), None) - .await?; + .await + { + Ok(billing) => billing, + Err(error) => { + abort_remaining_resets(reset_task); + return Err(error); + } + }; let plan = if kind == GrokAuthKind::Cli { self.fetch_cli_subscription_tier(credentials).await } else { None } .or_else(|| credentials.login_method()); - let reset_credits = self - .fetch_remaining_resets(Some(&credentials.access_token), None) - .await; + let reset_credits = + join_remaining_resets(reset_task, ctx.requires_optional_usage_completeness).await; let result = result_from_billing( billing, if kind == GrokAuthKind::Cli { @@ -154,11 +168,21 @@ impl GrokProvider { async fn fetch_with_cookie( &self, cookie_header: &str, + ctx: &FetchContext, ) -> Result { - let billing = self + let reset_task = self.spawn_remaining_resets(ctx, None, Some(cookie_header.to_string())); + let billing = match self .fetch_billing(None, Some(cookie_header.to_string())) - .await?; - let reset_credits = self.fetch_remaining_resets(None, Some(cookie_header)).await; + .await + { + Ok(billing) => billing, + Err(error) => { + abort_remaining_resets(reset_task); + return Err(error); + } + }; + let reset_credits = + join_remaining_resets(reset_task, ctx.requires_optional_usage_completeness).await; // v0.56.0: a browser session is its own principal. Never enrich a // successful cookie billing result from ambient auth.json metadata, // which may belong to a different account or change during the fetch. @@ -183,12 +207,12 @@ impl GrokProvider { ) { match step { GrokAutoStep::AmbientOAuth => { - if let Some(result) = self.try_ambient(GrokAuthKind::OAuth).await { + if let Some(result) = self.try_ambient(GrokAuthKind::OAuth, ctx).await { return result; } } GrokAutoStep::AmbientCli => { - if let Some(result) = self.try_ambient(GrokAuthKind::Cli).await { + if let Some(result) = self.try_ambient(GrokAuthKind::Cli, ctx).await { return result; } } @@ -196,17 +220,17 @@ impl GrokProvider { if let Some(token) = ctx.api_key.as_deref() { let credentials = GrokCredentials::from_bearer(token); return self - .fetch_with_auth(&credentials, GrokAuthKind::OAuth) + .fetch_with_auth(&credentials, GrokAuthKind::OAuth, ctx) .await; } } GrokAutoStep::ManualCookie => { if let Some(cookie_header) = &ctx.manual_cookie_header { - return self.fetch_with_cookie(cookie_header).await; + return self.fetch_with_cookie(cookie_header, ctx).await; } } GrokAutoStep::CookieRefresh => { - return self.fetch_with_cookie_refresh().await; + return self.fetch_with_cookie_refresh(ctx).await; } } } @@ -216,9 +240,10 @@ impl GrokProvider { async fn try_ambient( &self, kind: GrokAuthKind, + ctx: &FetchContext, ) -> Option> { let credentials = Self::load_credentials(kind).ok()?; - match self.fetch_with_auth(&credentials, kind).await { + match self.fetch_with_auth(&credentials, kind, ctx).await { Ok(result) => Some(Ok(result)), Err(ProviderError::AuthRequired) => None, Err(error) => { @@ -231,11 +256,14 @@ impl GrokProvider { /// Cookie refresh path (upstream #2458): /// 1. Try last validated cached cookie header (background reuse) /// 2. On miss/auth failure: re-import browser cookies, validate, cache - async fn fetch_with_cookie_refresh(&self) -> Result { + async fn fetch_with_cookie_refresh( + &self, + ctx: &FetchContext, + ) -> Result { use crate::browser::cookie_cache::CookieHeaderCache; if let Some(cached) = CookieHeaderCache::load(ProviderId::Grok) { - match self.fetch_with_cookie(&cached.cookie_header).await { + match self.fetch_with_cookie(&cached.cookie_header, ctx).await { Ok(result) => return Ok(result), Err(err) if is_cookie_authentication_failure(&err) => { CookieHeaderCache::clear(ProviderId::Grok); @@ -245,7 +273,7 @@ impl GrokProvider { } let cookie_header = crate::providers::browser_cookie_header(&["grok.com"])?; - let result = self.fetch_with_cookie(&cookie_header).await?; + let result = self.fetch_with_cookie(&cookie_header, ctx).await?; // Best-effort cache write: failing to persist the cookie only costs a // re-read from the browser on the next fetch. let _cached = CookieHeaderCache::store(ProviderId::Grok, &cookie_header, "browser"); @@ -293,20 +321,41 @@ impl GrokProvider { billing::parse_grpc_web_response(&bytes) } + fn spawn_remaining_resets( + &self, + ctx: &FetchContext, + access_token: Option, + cookie_header: Option, + ) -> Option<( + tokio::task::JoinHandle>, + Instant, + )> { + if !ctx.include_credits { + return None; + } + + let client = self.client.clone(); + let started_at = Instant::now(); + let task = tokio::spawn(async move { + Self::fetch_remaining_resets(client, access_token.as_deref(), cookie_header.as_deref()) + .await + }); + Some((task, started_at)) + } + /// Fetch optional SuperGrok reset-credit inventory using the same principal /// that produced the successful billing result. This is deliberately /// best-effort: billing remains valid when this secondary endpoint is down, /// malformed, unauthorized, or empty. async fn fetch_remaining_resets( - &self, + client: Client, access_token: Option<&str>, cookie_header: Option<&str>, ) -> Option { - let mut request = self - .client + let mut request = client .post(REMAINING_RESETS_ENDPOINT) .body(vec![0, 0, 0, 0, 0]) - .timeout(std::time::Duration::from_secs(2)) + .timeout(RESET_CREDITS_TIMEOUT) .header("Origin", "https://grok.com") .header("Referer", "https://grok.com/?_s=usage") .header("Accept", "*/*") @@ -377,6 +426,39 @@ impl GrokProvider { } } +fn abort_remaining_resets( + task: Option<( + tokio::task::JoinHandle>, + Instant, + )>, +) { + if let Some((task, _started_at)) = task { + task.abort(); + } +} + +async fn join_remaining_resets( + task: Option<( + tokio::task::JoinHandle>, + Instant, + )>, + requires_optional_usage_completeness: bool, +) -> Option { + let (mut task, started_at) = task?; + let budget = if requires_optional_usage_completeness { + RESET_CREDITS_TIMEOUT.saturating_sub(started_at.elapsed()) + } else { + RESET_CREDITS_JOIN_GRACE + }; + match tokio::time::timeout(budget, &mut task).await { + Ok(Ok(inventory)) => inventory, + Ok(Err(_)) | Err(_) => { + task.abort(); + None + } + } +} + #[cfg(windows)] fn hide_windows_console(command: &mut std::process::Command) { const CREATE_NO_WINDOW: u32 = 0x08000000; @@ -407,13 +489,14 @@ impl Provider for GrokProvider { SourceMode::Auto => self.fetch_auto(ctx).await, SourceMode::Web => { if let Some(cookie_header) = &ctx.manual_cookie_header { - return self.fetch_with_cookie(cookie_header).await; + return self.fetch_with_cookie(cookie_header, ctx).await; } - self.fetch_with_cookie_refresh().await + self.fetch_with_cookie_refresh(ctx).await } SourceMode::Cli => { let credentials = Self::load_credentials(GrokAuthKind::Cli)?; - self.fetch_with_auth(&credentials, GrokAuthKind::Cli).await + self.fetch_with_auth(&credentials, GrokAuthKind::Cli, ctx) + .await } SourceMode::OAuth => { // Prefer the switched ~/.grok/auth.json over a leftover token @@ -428,14 +511,15 @@ impl Provider for GrokProvider { } }; match self - .fetch_with_auth(&credentials, GrokAuthKind::OAuth) + .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).await + self.fetch_with_auth(&fallback, GrokAuthKind::OAuth, ctx) + .await } else { Err(ProviderError::AuthRequired) } diff --git a/rust/src/providers/grok/tests.rs b/rust/src/providers/grok/tests.rs index 37df362e9d..b68e7bc104 100644 --- a/rust/src/providers/grok/tests.rs +++ b/rust/src/providers/grok/tests.rs @@ -121,6 +121,17 @@ fn is_cookie_auth_failure_only_auth_required() { assert!(!is_cookie_authentication_failure(&ProviderError::NoCookies)); } +#[test] +fn include_credits_false_does_not_start_reset_lookup() { + let provider = GrokProvider::new(); + let ctx = FetchContext { + include_credits: false, + ..FetchContext::default() + }; + + assert!(provider.spawn_remaining_resets(&ctx, None, None).is_none()); +} + #[test] fn cookie_billing_stays_siloed_from_auth_file_identity() { let result = result_from_cookie_billing(GrokBillingSnapshot {