diff --git a/rust/src/providers/grok/billing.rs b/rust/src/providers/grok/billing/mod.rs similarity index 78% rename from rust/src/providers/grok/billing.rs rename to rust/src/providers/grok/billing/mod.rs index b69e9dec3c..e77d5474c5 100644 --- a/rust/src/providers/grok/billing.rs +++ b/rust/src/providers/grok/billing/mod.rs @@ -3,6 +3,10 @@ use reqwest::header::HeaderMap; use crate::core::ProviderError; +mod reset_coupons; + +pub(super) use reset_coupons::parse_grpc_web_reset_coupons; + #[derive(Debug, Clone, Copy)] pub(super) struct GrokBillingSnapshot { pub(super) used_percent: Option, @@ -19,18 +23,109 @@ pub(super) fn validate_grpc_headers(headers: &HeaderMap) -> Result<(), ProviderE .and_then(|value| value.parse::().ok()) && status != 0 { + return map_grpc_status(status, "Grok RPC"); + } + Ok(()) +} + +pub(super) fn parse_grpc_web_response(data: &[u8]) -> Result { + parse_grpc_web_response_at(data, Utc::now()) +} + +/// Map a gRPC status code onto the provider error policy shared by the +/// billing and reset-credit endpoints. +pub(super) fn map_grpc_status(status: u16, context: &str) -> Result<(), ProviderError> { + if status != 0 { if status == 16 { return Err(ProviderError::AuthRequired); } return Err(ProviderError::Other(format!( - "Grok RPC failed with status {status}" + "{context} failed with status {status}" ))); } Ok(()) } -pub(super) fn parse_grpc_web_response(data: &[u8]) -> Result { - parse_grpc_web_response_at(data, Utc::now()) +/// Decode a length-prefixed field body: `read_varint -> try_from -> +/// checked_add -> bounds-check` in one place. +pub(super) fn read_length_field( + data: &[u8], + index: usize, + what: &str, +) -> Result<(usize, usize), ProviderError> { + let (len, start) = read_varint(data, index) + .ok_or_else(|| ProviderError::Parse(format!("Grok {what} is malformed")))?; + let len = usize::try_from(len) + .map_err(|_| ProviderError::Parse(format!("Grok {what} is too large")))?; + let end = start + .checked_add(len) + .ok_or_else(|| ProviderError::Parse(format!("Grok {what} length overflowed")))?; + if end > data.len() { + return Err(ProviderError::Parse(format!("Grok {what} is truncated"))); + } + Ok((start, end)) +} + +/// Decode a varint Unix-seconds timestamp with the shared epoch bounding. +pub(super) fn unix_seconds_timestamp(seconds: u64) -> Option> { + // Varint timestamps are Unix seconds inside the range checked below. + #[allow( + clippy::cast_possible_wrap, + reason = "varint timestamps are bounded to the Unix-seconds range checked below" + )] + let seconds = seconds as i64; + (1_700_000_000..=2_100_000_000) + .contains(&seconds) + .then(|| Utc.timestamp_opt(seconds, 0).single()) + .flatten() +} + +/// One parameterized gRPC-web frame walker. `on_malformed` decides the +/// malformed-frame policy: billing swallows malformed frames, the optional +/// reset lookup fails closed. +/// +/// Yields `(flags, payload)` for every frame, data and trailer alike; callers +/// split on the trailer flag. +pub(super) fn grpc_web_frames( + data: &[u8], + on_malformed: fn(&str) -> Option, +) -> Result, ProviderError> { + let mut frames = Vec::new(); + let mut index = 0; + while index < data.len() { + if index + 5 > data.len() { + return on_malformed("truncated").map_or(Ok(frames), Err); + } + let flags = data[index]; + let len = ((data[index + 1] as usize) << 24) + | ((data[index + 2] as usize) << 16) + | ((data[index + 3] as usize) << 8) + | (data[index + 4] as usize); + let start = index + 5; + let Some(end) = start.checked_add(len) else { + return on_malformed("frame is too large").map_or(Ok(frames), Err); + }; + if end > data.len() { + return on_malformed("truncated").map_or(Ok(frames), Err); + } + frames.push((flags, &data[start..end])); + index = end; + } + Ok(frames) +} + +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( @@ -127,16 +222,7 @@ fn looks_like_protobuf_payload(data: &[u8]) -> bool { } fn varint_timestamp(field: &VarintField) -> Option> { - // Varint timestamps are Unix seconds inside the range checked below. - #[allow( - clippy::cast_possible_wrap, - reason = "varint timestamps are bounded to the Unix-seconds range checked below" - )] - let seconds = field.value as i64; - (1_700_000_000..=2_100_000_000) - .contains(&field.value) - .then(|| Utc.timestamp_opt(seconds, 0).single()) - .flatten() + unix_seconds_timestamp(field.value) } fn current_period_window_minutes(scan: &ProtoScan, now: DateTime) -> Option { @@ -173,30 +259,12 @@ fn unique_varint_at_path(scan: &ProtoScan, path: &[u64]) -> Option { } fn grpc_web_data_frames(data: &[u8]) -> Vec> { - let mut frames = Vec::new(); - let mut index = 0; - while index < data.len() { - if index + 5 > data.len() { - return Vec::new(); - } - let flags = data[index]; - let len = ((data[index + 1] as usize) << 24) - | ((data[index + 2] as usize) << 16) - | ((data[index + 3] as usize) << 8) - | (data[index + 4] as usize); - let start = index + 5; - let Some(end) = start.checked_add(len) else { - return Vec::new(); - }; - if end > data.len() { - return Vec::new(); - } - if flags & 0x80 == 0 { - frames.push(data[start..end].to_vec()); - } - index = end; - } - frames + grpc_web_frames(data, |_| None) + .unwrap_or_default() + .into_iter() + .filter(|(flags, _)| flags & 0x80 == 0) + .map(|(_, payload)| payload.to_vec()) + .collect() } struct ProtoScan { @@ -633,6 +701,83 @@ 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()); + } + + #[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]; + 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() } @@ -667,6 +812,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/billing/reset_coupons.rs b/rust/src/providers/grok/billing/reset_coupons.rs new file mode 100644 index 0000000000..6cbb35194f --- /dev/null +++ b/rust/src/providers/grok/billing/reset_coupons.rs @@ -0,0 +1,258 @@ +//! SuperGrok reset-coupon parser for the GetRemainingResets gRPC-web endpoint. +//! +//! Independent of the billing response parser: own framing, own protobuf +//! shape, and a fail-closed failure policy so a partial inventory is never +//! published. The token ID stays private to this module and is never attached +//! to a public or persisted snapshot. + +use chrono::{DateTime, Utc}; + +use super::read_length_field; +use crate::core::ProviderError; + +/// 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(in crate::providers::grok) struct GrokResetCoupon { + pub(in crate::providers::grok) token_id: String, + pub(in crate::providers::grok) granted_at: Option>, + pub(in crate::providers::grok) expires_at: DateTime, +} + +pub(in crate::providers::grok) fn parse_grpc_web_reset_coupons( + data: &[u8], + now: DateTime, +) -> Result, ProviderError> { + let payloads = grpc_web_reset_payloads(data)?; + + 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) +} + +/// One empty unary gRPC-web request/response frame. +const EMPTY_GRPC_WEB_FRAME: [u8; 5] = [0, 0, 0, 0, 0]; + +fn grpc_web_reset_payloads(data: &[u8]) -> Result>, ProviderError> { + if data.is_empty() || data == EMPTY_GRPC_WEB_FRAME { + 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()) + }); + } + + super::grpc_web_frames(data, on_malformed_frame_fail).and_then(|frames| { + let mut payloads = Vec::new(); + for (flags, payload) in frames { + if flags & 0x80 != 0 { + validate_grpc_web_reset_trailer(payload)?; + } else { + payloads.push(payload.to_vec()); + } + } + Ok(payloads) + }) +} + +fn on_malformed_frame_fail(context: &str) -> Option { + Some(ProviderError::Parse(format!( + "Grok reset-credit gRPC-web frame is {context}" + ))) +} + +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()) + })?; + super::map_grpc_status(status, "Grok reset-credit RPC") +} + +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) = super::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 (start, end) = read_length_field(data, index, "record")?; + if let Some(coupon) = parse_reset_coupon(&data[start..end], now)? { + coupons.push(coupon); + } + index = 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) = super::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 (start, end) = read_length_field(data, index, "token id")?; + 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 (start, end) = read_length_field(data, index, "timestamp")?; + 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, + })) +} + +/// Decode one embedded timestamp message (field 1 varint, Unix seconds). +fn parse_timestamp_message(data: &[u8]) -> Result>, ProviderError> { + let mut index = 0; + let mut seconds = None; + while index < data.len() { + let (field, wire, next) = super::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) = super::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); + }; + Ok(super::unix_seconds_timestamp(seconds)) +} + +fn skip_field(data: &[u8], index: usize, wire: u64) -> Option { + match wire { + 0 => super::read_varint(data, index).map(|(_, next)| next), + 1 => index.checked_add(8).filter(|end| *end <= data.len()), + 2 => { + let (len, start) = super::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 looks_like_protobuf_payload(data: &[u8]) -> bool { + let Some(&first) = data.first() else { + return false; + }; + let field_number = first >> 3; + let wire_type = first & 0x07; + field_number > 0 && matches!(wire_type, 0 | 1 | 2 | 5) +} diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 976ba76b04..5565cea7eb 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -12,20 +12,25 @@ use chrono::{DateTime, Utc}; use reqwest::Client; use serde_json::Value; use std::path::PathBuf; +use std::time::Duration; #[cfg(windows)] 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"; +const RESET_CREDITS_TIMEOUT: Duration = Duration::from_secs(2); +const RESET_CREDITS_JOIN_GRACE: Duration = Duration::from_millis(250); pub struct GrokProvider { metadata: ProviderMetadata, @@ -64,6 +69,11 @@ impl GrokProvider { dirs::home_dir().map(|home| home.join(".grok").join("auth.json")) } + #[cfg(test)] + fn client_for_tests(&self) -> Client { + self.client.clone() + } + fn load_credentials(kind: GrokAuthKind) -> Result { let path = Self::auth_file_path() .ok_or_else(|| ProviderError::NotInstalled("Grok auth path not found".to_string()))?; @@ -86,7 +96,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)) } @@ -94,17 +106,34 @@ impl GrokProvider { &self, credentials: &GrokCredentials, kind: GrokAuthKind, + ctx: &FetchContext, ) -> Result { - let billing = self + let reset_lookup = GrokProvider::spawn_remaining_resets( + ctx, + Some(credentials.access_token.clone()), + None, + self.client.clone(), + ); + let billing = match self .fetch_billing(Some(format!("Bearer {}", credentials.access_token)), None) - .await?; + .await + { + Ok(billing) => billing, + Err(error) => { + reset_lookup.abort(); + return Err(error); + } + }; let plan = if kind == GrokAuthKind::Cli { self.fetch_cli_subscription_tier(credentials).await } else { None } .or_else(|| credentials.login_method()); - Ok(result_from_billing( + let reset_credits = reset_lookup + .join(ctx.requires_optional_usage_completeness) + .await; + let result = result_from_billing( billing, if kind == GrokAuthKind::Cli { "grok-cli" @@ -114,7 +143,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 { @@ -146,14 +179,35 @@ impl GrokProvider { async fn fetch_with_cookie( &self, cookie_header: &str, + ctx: &FetchContext, ) -> Result { - let billing = self + let reset_lookup = GrokProvider::spawn_remaining_resets( + ctx, + None, + Some(cookie_header.to_string()), + self.client.clone(), + ); + let billing = match self .fetch_billing(None, Some(cookie_header.to_string())) - .await?; + .await + { + Ok(billing) => billing, + Err(error) => { + reset_lookup.abort(); + return Err(error); + } + }; + let reset_credits = reset_lookup + .join(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. - 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 { @@ -170,12 +224,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; } } @@ -183,17 +237,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; } } } @@ -203,9 +257,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) => { @@ -218,11 +273,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); @@ -232,7 +290,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"); @@ -280,6 +338,94 @@ impl GrokProvider { billing::parse_grpc_web_response(&bytes) } + fn spawn_remaining_resets( + ctx: &FetchContext, + access_token: Option, + cookie_header: Option, + client: Client, + ) -> ResetLookup { + if !ctx.include_credits { + return ResetLookup::idle(); + } + ResetLookup { + task: Some(tokio::spawn(async move { + Self::fetch_remaining_resets( + client, + access_token.as_deref(), + cookie_header.as_deref(), + ) + .await + })), + } + } + + /// 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( + client: Client, + access_token: Option<&str>, + cookie_header: Option<&str>, + ) -> Option { + let mut request = client + .post(REMAINING_RESETS_ENDPOINT) + .body(vec![0, 0, 0, 0, 0]) + .timeout(RESET_CREDITS_TIMEOUT) + .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"); @@ -296,6 +442,47 @@ impl GrokProvider { } } +/// Best-effort reset-credit lookup running alongside the billing request. +/// `abort` and `join` own the handle state; `JoinHandle::abort` is already +/// idempotent on finished tasks. +struct ResetLookup { + task: Option>>, +} + +impl ResetLookup { + fn idle() -> Self { + Self { task: None } + } + + fn abort(self) { + if let Some(task) = self.task { + task.abort(); + } + } + + /// Wait for the lookup with the policy budget: under + /// `requires_optional_usage_completeness` the full timeout remains; + /// otherwise a short grace joins the already-running request. + async fn join( + self, + requires_optional_usage_completeness: bool, + ) -> Option { + let mut task = self.task?; + let budget = if requires_optional_usage_completeness { + RESET_CREDITS_TIMEOUT + } 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; @@ -326,13 +513,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 @@ -347,14 +535,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..71691756b9 100644 --- a/rust/src/providers/grok/tests.rs +++ b/rust/src/providers/grok/tests.rs @@ -121,6 +121,22 @@ 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 ctx = FetchContext { + include_credits: false, + ..FetchContext::default() + }; + + let lookup = GrokProvider::spawn_remaining_resets( + &ctx, + None, + None, + crate::providers::grok::GrokProvider::new().client_for_tests(), + ); + assert!(lookup.task.is_none()); +} + #[test] fn cookie_billing_stays_siloed_from_auth_file_identity() { let result = result_from_cookie_billing(GrokBillingSnapshot {