From f9833f59c40ce41d70d4f85996021d16cc585c22 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:59:17 +0700 Subject: [PATCH] Port upstream 0.69.0: show Grok product usage breakdown on the grok.com billing fallback Decode the per-product shares of the GetGrokCreditsConfig gRPC-web answer ([1, 7] entries) with a bounded shared protobuf reader, trust them only when they compose the same payload's config percent, and show them as plain detail rows. Billing frames now fail closed on truncated or reserved-flag frames. --- rust/src/providers/grok/billing/mod.rs | 237 +++++------- rust/src/providers/grok/billing/protobuf.rs | 182 ++++++++++ .../providers/grok/billing/reset_coupons.rs | 132 ++----- .../grok/billing/web_product_usage.rs | 91 +++++ .../grok/billing/web_product_usage_tests.rs | 336 ++++++++++++++++++ rust/src/providers/grok/credits_proxy.rs | 15 +- rust/src/providers/grok/mod.rs | 15 +- rust/src/providers/grok/product_usage.rs | 23 +- .../src/providers/grok/product_usage_tests.rs | 124 ++++++- rust/src/providers/grok/tests.rs | 7 + 10 files changed, 874 insertions(+), 288 deletions(-) create mode 100644 rust/src/providers/grok/billing/protobuf.rs create mode 100644 rust/src/providers/grok/billing/web_product_usage.rs create mode 100644 rust/src/providers/grok/billing/web_product_usage_tests.rs diff --git a/rust/src/providers/grok/billing/mod.rs b/rust/src/providers/grok/billing/mod.rs index e77d5474c5..98cb51f0e7 100644 --- a/rust/src/providers/grok/billing/mod.rs +++ b/rust/src/providers/grok/billing/mod.rs @@ -1,19 +1,29 @@ use chrono::{DateTime, TimeZone, Utc}; use reqwest::header::HeaderMap; +use super::product_usage::{GrokProductUsage, compose}; use crate::core::ProviderError; +mod protobuf; mod reset_coupons; +mod web_product_usage; +#[cfg(test)] +pub(super) mod web_product_usage_tests; +use protobuf::{FieldValue, ProtobufField, looks_like_protobuf_payload}; pub(super) use reset_coupons::parse_grpc_web_reset_coupons; +use web_product_usage::decode_product_usage; -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone)] pub(super) struct GrokBillingSnapshot { pub(super) used_percent: Option, pub(super) used_percent_is_wire_published: bool, pub(super) used_percent_is_implicit_zero: bool, pub(super) resets_at: Option>, pub(super) window_minutes: Option, + /// Product shares composing this same answer's wire percent; empty for + /// every other percent source. + pub(super) product_usage: Vec, } pub(super) fn validate_grpc_headers(headers: &HeaderMap) -> Result<(), ProviderError> { @@ -46,26 +56,6 @@ pub(super) fn map_grpc_status(status: u16, context: &str) -> Result<(), Provider Ok(()) } -/// 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. @@ -81,8 +71,7 @@ pub(super) fn unix_seconds_timestamp(seconds: u64) -> Option> { } /// One parameterized gRPC-web frame walker. `on_malformed` decides the -/// malformed-frame policy: billing swallows malformed frames, the optional -/// reset lookup fails closed. +/// malformed-frame policy; both callers fail closed, with their own message. /// /// Yields `(flags, payload)` for every frame, data and trailer alike; callers /// split on the trailer flag. @@ -114,20 +103,6 @@ pub(super) fn grpc_web_frames( 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( data: &[u8], now: DateTime, @@ -167,11 +142,11 @@ fn parse_grpc_web_response_at( .iter() .any(|field| field.value != first.value) }); - let parsed_percent = if scan.is_complete && !conflicting_percent { - valid_percent_fields.first().map(|field| field.value as f64) - } else { - None - }; + let percent_field = valid_percent_fields + .first() + .filter(|_| scan.is_complete && !conflicting_percent); + let parsed_percent = percent_field.map(|field| f64::from(field.value)); + let percent_path = percent_field.map(|field| field.path.as_slice()); let reset_fields: Vec<(&VarintField, DateTime)> = scan .varints @@ -203,24 +178,26 @@ fn parse_grpc_web_response_at( let used_percent_is_implicit_zero = no_usage_yet && payloads.len() == 1 && scan.is_complete && has_active_current_period; let used_percent = parsed_percent.or_else(|| used_percent_is_implicit_zero.then_some(0.0)); + // Shares are trusted only from one complete message whose aggregate is the + // config's own `[1, 1]` percent, never a nested or repeated lookalike. + let product_usage = match (parsed_percent, payloads.as_slice()) { + (Some(percent), [payload]) + if scan.is_complete && percent_path == Some([1, 1].as_slice()) => + { + compose(decode_product_usage(payload), percent) + } + _ => Vec::new(), + }; Ok(GrokBillingSnapshot { used_percent, used_percent_is_wire_published: parsed_percent.is_some(), used_percent_is_implicit_zero, resets_at, window_minutes, + product_usage, }) } -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) -} - fn varint_timestamp(field: &VarintField) -> Option> { unix_seconds_timestamp(field.value) } @@ -258,9 +235,21 @@ fn unique_varint_at_path(scan: &ProtoScan, path: &[u64]) -> Option { values.all(|value| value == first).then_some(first) } +/// Data-frame payloads of a billing response. Empty when any frame is +/// truncated or carries a compressed or reserved flag: such a response must +/// never supply a partial percent, product shares, or an implicit zero. fn grpc_web_data_frames(data: &[u8]) -> Vec> { - grpc_web_frames(data, |_| None) - .unwrap_or_default() + let Ok(frames) = grpc_web_frames(data, |context| { + Some(ProviderError::Parse(format!( + "Grok web billing frame is {context}" + ))) + }) else { + return Vec::new(); + }; + if frames.iter().any(|(flags, _)| !matches!(flags, 0 | 0x80)) { + return Vec::new(); + } + frames .into_iter() .filter(|(flags, _)| flags & 0x80 == 0) .map(|(_, payload)| payload.to_vec()) @@ -297,120 +286,62 @@ struct VarintField { } impl ProtoScan { + /// Walk one message. The first malformed field ends the scan and marks it + /// incomplete, so nothing after it can be mistaken for a real field. fn scan_message(&mut self, data: &[u8], path: &mut Vec, depth: usize) { if depth > 8 { self.is_complete = false; return; } - let mut i = 0; - while i < data.len() { - let field_start = i; - let Some((field, wire, next)) = read_key(data, i) else { + let mut index = 0; + while index < data.len() { + let Some(field) = ProtobufField::read(data, &mut index) else { self.is_complete = false; - i = field_start.saturating_add(1); - continue; + return; }; - i = next; - path.push(field); - let Some(next) = self.scan_field(data, i, path, depth, wire) else { - self.is_complete = false; - path.pop(); - i = field_start.saturating_add(1); - continue; - }; - i = next; + path.push(field.number); + let accepted = self.scan_field(field.value, path, depth); path.pop(); + if !accepted { + self.is_complete = false; + return; + } } } - fn scan_field( - &mut self, - data: &[u8], - i: usize, - path: &mut Vec, - depth: usize, - wire: u64, - ) -> Option { - if (path.as_slice() == [1, 1] && wire != 5) || (is_known_billing_message(path) && wire != 2) - { - return None; - } - match wire { - 0 => self.scan_varint(data, i, path), - 2 => self.scan_length_delimited(data, i, path, depth), - 5 => self.scan_fixed32(data, i, path), - 1 => i.checked_add(8).filter(|end| *end <= data.len()), - _ => None, - } - } - - fn scan_varint(&mut self, data: &[u8], i: usize, path: &[u64]) -> Option { - let (value, next) = read_varint(data, i)?; - self.varints.push(VarintField { - path: path.to_vec(), - value, - }); - Some(next) - } - - fn scan_length_delimited( - &mut self, - data: &[u8], - i: usize, - path: &mut Vec, - depth: usize, - ) -> Option { - let (len, next) = read_varint(data, i)?; - let start = next; - let len_usize = usize::try_from(len).ok()?; - let end = start.checked_add(len_usize)?; - if end > data.len() { - return None; - } - if depth < 4 && is_known_billing_message(path) { - self.scan_message(&data[start..end], path, depth + 1); - } - Some(end) - } - - fn scan_fixed32(&mut self, data: &[u8], i: usize, path: &[u64]) -> Option { - let end = i.checked_add(4)?; - if end > data.len() { - return None; - } - let bytes = [data[i], data[i + 1], data[i + 2], data[i + 3]]; - self.fixed32.push(Fixed32Field { - path: path.to_vec(), - value: f32::from_le_bytes(bytes), - order: self.order, - }); - self.order += 1; - Some(end) - } -} - -fn read_key(data: &[u8], i: usize) -> Option<(u64, u64, usize)> { - let (key, next) = read_varint(data, i)?; - let field = key >> 3; - (field > 0 && field <= 536_870_911).then_some((field, key & 0x07, next)) -} - -fn read_varint(data: &[u8], mut i: usize) -> Option<(u64, usize)> { - let mut value = 0u64; - let mut shift = 0; - while i < data.len() && shift < 64 { - let b = data[i]; - i += 1; - if shift == 63 && b > 1 { - return None; + /// Record one field; `false` when its wire type contradicts the descriptor. + fn scan_field(&mut self, value: FieldValue<'_>, path: &mut Vec, depth: usize) -> bool { + let known_message = is_known_billing_message(path); + let wire_matches_descriptor = match value { + FieldValue::Message(_) => path.as_slice() != [1, 1], + FieldValue::Fixed32(_) => !known_message, + FieldValue::Varint(_) | FieldValue::Fixed64 => { + !known_message && path.as_slice() != [1, 1] + } + }; + if !wire_matches_descriptor { + return false; } - value |= u64::from(b & 0x7f) << shift; - if b & 0x80 == 0 { - return Some((value, i)); + match value { + FieldValue::Varint(value) => self.varints.push(VarintField { + path: path.clone(), + value, + }), + FieldValue::Fixed32(value) => { + self.fixed32.push(Fixed32Field { + path: path.clone(), + value, + order: self.order, + }); + self.order += 1; + } + FieldValue::Message(bytes) if depth < 4 && known_message => { + self.scan_message(bytes, path, depth + 1); + } + FieldValue::Message(_) | FieldValue::Fixed64 => {} } - shift += 7; + true } - None } fn is_known_billing_message(path: &[u64]) -> bool { diff --git a/rust/src/providers/grok/billing/protobuf.rs b/rust/src/providers/grok/billing/protobuf.rs new file mode 100644 index 0000000000..5983027f12 --- /dev/null +++ b/rust/src/providers/grok/billing/protobuf.rs @@ -0,0 +1,182 @@ +//! Bounded protobuf wire reader shared by the Grok billing, product-usage and +//! reset-coupon parsers (upstream `GrokProtobufField`, v0.69.0). +//! +//! Every read is bounds-checked: varints stop after ten bytes and reject +//! overflow, and length-delimited values cannot extend past the enclosing +//! message. Unknown wire types and out-of-range field numbers are malformed. + +/// Largest field number protobuf allows (2^29 - 1). +const MAX_FIELD_NUMBER: u64 = 536_870_911; + +/// The payload of one decoded field. Fixed 64-bit values are consumed but +/// never interpreted; length-delimited bytes stay opaque until a caller knows +/// they hold a message. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(super) enum FieldValue<'a> { + Varint(u64), + Fixed64, + Message(&'a [u8]), + Fixed32(f32), +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub(super) struct ProtobufField<'a> { + pub(super) number: u64, + pub(super) value: FieldValue<'a>, +} + +impl<'a> ProtobufField<'a> { + /// Decode every field of `data`, or `None` when any of it is malformed. + pub(super) fn fields(data: &'a [u8]) -> Option> { + let mut fields = Vec::new(); + let mut index = 0; + while index < data.len() { + fields.push(Self::read(data, &mut index)?); + } + Some(fields) + } + + /// Decode one field at `index`, advancing it past the field. `index` is + /// left untouched when the field is malformed. + pub(super) fn read(data: &'a [u8], index: &mut usize) -> Option { + let (key, mut next) = read_varint(data, *index)?; + let number = key >> 3; + if number == 0 || number > MAX_FIELD_NUMBER { + return None; + } + let value = match key & 0x07 { + 0 => { + let (value, end) = read_varint(data, next)?; + next = end; + FieldValue::Varint(value) + } + 1 => { + next = advance(data, next, 8)?; + FieldValue::Fixed64 + } + 2 => { + let (len, start) = read_varint(data, next)?; + next = advance(data, start, usize::try_from(len).ok()?)?; + FieldValue::Message(&data[start..next]) + } + 5 => { + let end = advance(data, next, 4)?; + let bytes: [u8; 4] = data[next..end].try_into().ok()?; + next = end; + FieldValue::Fixed32(f32::from_le_bytes(bytes)) + } + _ => return None, + }; + *index = next; + Some(Self { number, value }) + } + + pub(super) fn varint(&self) -> Option { + match self.value { + FieldValue::Varint(value) => Some(value), + _ => None, + } + } + + pub(super) fn fixed32(&self) -> Option { + match self.value { + FieldValue::Fixed32(value) => Some(value), + _ => None, + } + } + + pub(super) fn message(&self) -> Option<&'a [u8]> { + match self.value { + FieldValue::Message(bytes) => Some(bytes), + _ => None, + } + } +} + +/// Whether `data` starts like a raw (unframed) protobuf message. Valid keys +/// cannot begin with field number 0, so a gRPC-web data flag never matches. +pub(super) fn looks_like_protobuf_payload(data: &[u8]) -> bool { + let Some(&first) = data.first() else { + return false; + }; + first >> 3 > 0 && matches!(first & 0x07, 0 | 1 | 2 | 5) +} + +fn advance(data: &[u8], index: usize, len: usize) -> Option { + index.checked_add(len).filter(|end| *end <= data.len()) +} + +fn read_varint(data: &[u8], mut index: usize) -> Option<(u64, usize)> { + let mut value = 0u64; + let mut shift = 0; + while index < data.len() && shift < 64 { + let byte = data[index]; + index += 1; + if shift == 63 && byte > 1 { + return None; + } + value |= u64::from(byte & 0x7f) << shift; + if byte & 0x80 == 0 { + return Some((value, index)); + } + shift += 7; + } + None +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn wire_reader_bounds_continuation_and_length_reads() { + let mut endless = vec![0x08]; + endless.extend([0x80; 4096]); + let mut index = 0; + assert_eq!(ProtobufField::read(&endless, &mut index), None); + assert_eq!(index, 0); + + let malformed: [&[u8]; 6] = [ + &[0x00], + &[0x02, 0], + &[0x09, 0], + &[0x15, 0, 0], + &[0x12, 0x7F, 0], + &[0x08], + ]; + for bytes in malformed { + assert_eq!(ProtobufField::fields(bytes), None, "{bytes:?}"); + } + } + + #[test] + fn wire_reader_accepts_maximum_integers_and_leaves_opaque_bytes_uninterpreted() { + let mut maximum = vec![0x08]; + maximum.extend([0xFF; 9]); + maximum.push(0x01); + maximum.extend([0x12, 0x01, 0xFF]); + + let fields = ProtobufField::fields(&maximum).unwrap(); + assert_eq!(fields.len(), 2); + assert_eq!(fields[0].varint(), Some(u64::MAX)); + assert_eq!(fields[1].message(), Some(&[0xFF][..])); + } + + #[test] + fn wire_reader_decodes_fixed_width_values_and_rejects_overflow() { + let mut bytes = vec![0x0d]; + bytes.extend(6.0f32.to_le_bytes()); + bytes.push(0x11); + bytes.extend([0; 8]); + let fields = ProtobufField::fields(&bytes).unwrap(); + assert_eq!(fields[0].fixed32(), Some(6.0)); + assert_eq!(fields[1].value, FieldValue::Fixed64); + assert_eq!(fields[1].varint(), None); + + let mut overflow = vec![0x08]; + overflow.extend([0xFF; 9]); + overflow.push(0x02); + assert_eq!(ProtobufField::fields(&overflow), None); + assert_eq!(ProtobufField::fields(&[0x08, 0x80]), None); + } +} diff --git a/rust/src/providers/grok/billing/reset_coupons.rs b/rust/src/providers/grok/billing/reset_coupons.rs index 6cbb35194f..f398ee2185 100644 --- a/rust/src/providers/grok/billing/reset_coupons.rs +++ b/rust/src/providers/grok/billing/reset_coupons.rs @@ -7,7 +7,7 @@ use chrono::{DateTime, Utc}; -use super::read_length_field; +use super::protobuf::{ProtobufField, looks_like_protobuf_payload}; use crate::core::ProviderError; /// One unused SuperGrok usage-limit reset coupon. The token ID is retained @@ -108,27 +108,15 @@ fn parse_reset_coupon_container( 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()) + let fields = ProtobufField::fields(data).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit protobuf is malformed".to_string()) + })?; + for field in fields.iter().filter(|field| field.number == 10) { + let record = field.message().ok_or_else(|| { + ProviderError::Parse("Grok reset-credit record has an invalid wire type".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()) - })?; + if let Some(coupon) = parse_reset_coupon(record, now)? { + coupons.push(coupon); } } Ok(()) @@ -138,25 +126,21 @@ fn parse_reset_coupon( data: &[u8], now: DateTime, ) -> Result, ProviderError> { - let mut index = 0; + let fields = ProtobufField::fields(data) + .ok_or_else(|| ProviderError::Parse("Grok reset-credit record is malformed".to_string()))?; 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 { + for field in &fields { + match field.number { 10 => { - if wire != 2 { - return Err(ProviderError::Parse( + let bytes = field.message().ok_or_else(|| { + 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]) + std::str::from_utf8(bytes) .map_err(|_| { ProviderError::Parse( "Grok reset-credit token id is not UTF-8".to_string(), @@ -164,28 +148,21 @@ fn parse_reset_coupon( })? .to_string(), ); - index = end; } 20 | 30 => { - if wire != 2 { - return Err(ProviderError::Parse( + let message = field.message().ok_or_else(|| { + 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 { + ) + })?; + let timestamp = parse_timestamp_message(message)?; + if field.number == 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()) - })?; } + _ => {} } } @@ -204,55 +181,16 @@ fn parse_reset_coupon( /// Decode one embedded timestamp message (field 1 varint, Unix seconds). fn parse_timestamp_message(data: &[u8]) -> Result>, ProviderError> { - let mut index = 0; + let fields = ProtobufField::fields(data).ok_or_else(|| { + ProviderError::Parse("Grok reset-credit timestamp is malformed".to_string()) + })?; 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()) - })?; - } + for field in fields.iter().filter(|field| field.number == 1) { + seconds = Some(field.varint().ok_or_else(|| { + ProviderError::Parse( + "Grok reset-credit timestamp seconds has an invalid wire type".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) + Ok(seconds.and_then(super::unix_seconds_timestamp)) } diff --git a/rust/src/providers/grok/billing/web_product_usage.rs b/rust/src/providers/grok/billing/web_product_usage.rs new file mode 100644 index 0000000000..87a0f7f0f8 --- /dev/null +++ b/rust/src/providers/grok/billing/web_product_usage.rs @@ -0,0 +1,91 @@ +//! Per-product credit shares carried by the grok.com `GetGrokCreditsConfig` +//! gRPC-web answer, as repeated `[1, 7]` entries `{1: product id, 2: float +//! percent}` (an omitted percent means 0). Follows upstream +//! `GrokWebBillingFetcher.decodeProductUsage` (v0.69.0). +//! +//! Decoding is all-or-nothing: any malformed or duplicate entry drops the +//! whole list so a partial list can never pass the composition check as if it +//! were complete. + +use std::collections::HashSet; + +use super::protobuf::ProtobufField; +use crate::providers::grok::product_usage::GrokProductUsage; + +/// Field of the response message that holds the billing config. +const CONFIG_FIELD: u64 = 1; +/// Field of the config message that repeats one entry per product. +const PRODUCT_FIELD: u64 = 7; +/// Field of the config message that carries the aggregate credit percent. +const AGGREGATE_FIELD: u64 = 1; + +/// Decode the product shares of one complete response message. Empty unless +/// the message holds exactly one config with exactly one aggregate field. +pub(super) fn decode_product_usage(payload: &[u8]) -> Vec { + decode(payload).unwrap_or_default() +} + +fn decode(payload: &[u8]) -> Option> { + let root = ProtobufField::fields(payload)?; + let mut configs = root.iter().filter(|field| field.number == CONFIG_FIELD); + let config = configs.next()?.message()?; + if configs.next().is_some() { + return None; + } + let config = ProtobufField::fields(config)?; + if config + .iter() + .filter(|field| field.number == AGGREGATE_FIELD) + .count() + != 1 + { + return None; + } + + let mut products = Vec::new(); + let mut seen_ids = HashSet::new(); + for field in config.iter().filter(|field| field.number == PRODUCT_FIELD) { + let (id, percent) = decode_entry(field.message()?)?; + if !seen_ids.insert(id) { + return None; + } + match product_name(id) { + Some(name) => products.push(GrokProductUsage { + product: name.to_string(), + used_percent: percent, + }), + // An unnamed product may only be idle; a nonzero share would make + // the named shares an incomplete list. + None if percent > 0.0 => return None, + None => {} + } + } + Some(products) +} + +/// Only product ids verified against live CLI-proxy samples are named. +fn product_name(id: u64) -> Option<&'static str> { + match id { + 2 => Some("GrokBuild"), + 4 => Some("GrokChat"), + _ => None, + } +} + +fn decode_entry(entry: &[u8]) -> Option<(u64, f64)> { + let fields = ProtobufField::fields(entry)?; + let mut ids = fields.iter().filter(|field| field.number == 1); + let id = ids.next()?.varint()?; + if ids.next().is_some() { + return None; + } + let mut percentages = fields.iter().filter(|field| field.number == 2); + let percent = match percentages.next() { + Some(field) => f64::from(field.fixed32()?), + None => 0.0, + }; + if percentages.next().is_some() || !percent.is_finite() || percent < 0.0 { + return None; + } + Some((id, percent)) +} diff --git a/rust/src/providers/grok/billing/web_product_usage_tests.rs b/rust/src/providers/grok/billing/web_product_usage_tests.rs new file mode 100644 index 0000000000..5124c012dc --- /dev/null +++ b/rust/src/providers/grok/billing/web_product_usage_tests.rs @@ -0,0 +1,336 @@ +//! Product breakdown decoding on the grok.com gRPC-web billing answer. The +//! live frame and the case list follow upstream `GrokWebBillingProductUsageTests` +//! (v0.69.0). + +use chrono::{DateTime, TimeZone, Utc}; + +use super::*; + +/// Live `GetGrokCreditsConfig` frame captured 2026-09-26: 6% used, split into +/// GrokChat 4 and GrokBuild 2 (checked in verbatim upstream). +const LIVE_FRAME_HEX: &str = concat!( + "000000005f0a5d0d0000c04012001a00220c08a5d2c0d5061088ccb580022a0c08a5c7e5d5061088ccb58002", + "3a07080415000080403a0708021500000040421e0802120c08a5d2c0d5061088ccb580021a0c08a5c7e5d506", + "1088ccb58002580162006801800000000f677270632d7374617475733a300d0a", +); + +const PERIOD_START: u64 = 1_789_929_765; +const PERIOD_END: u64 = 1_790_534_565; + +fn now() -> DateTime { + Utc.timestamp_opt(1_790_456_400, 0).single().unwrap() +} + +fn product(name: &str, used_percent: f64) -> GrokProductUsage { + GrokProductUsage { + product: name.to_string(), + used_percent, + } +} + +fn chat_and_build() -> Vec { + vec![product("GrokChat", 4.0), product("GrokBuild", 2.0)] +} + +fn parse(data: &[u8]) -> GrokBillingSnapshot { + parse_grpc_web_response_at(data, now()).unwrap() +} + +pub(in crate::providers::grok) fn live_frame() -> Vec { + hex_bytes(LIVE_FRAME_HEX) +} + +fn hex_bytes(hex: &str) -> Vec { + let digits = hex.as_bytes(); + assert!( + digits.len().is_multiple_of(2), + "hex fixture length must be even" + ); + digits + .chunks(2) + .map(|pair| u8::from_str_radix(std::str::from_utf8(pair).unwrap(), 16).unwrap()) + .collect() +} + +fn varint(mut value: u64) -> Vec { + let mut bytes = Vec::new(); + loop { + let low = u8::try_from(value & 0x7f).unwrap(); + value >>= 7; + if value == 0 { + bytes.push(low); + return bytes; + } + bytes.push(low | 0x80); + } +} + +fn varint_field(number: u64, value: u64) -> Vec { + let mut bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes +} + +fn fixed32(number: u64, value: f32) -> Vec { + let mut bytes = varint((number << 3) | 5); + bytes.extend(value.to_le_bytes()); + bytes +} + +fn message(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(u64::try_from(value.len()).unwrap())); + bytes.extend(value); + bytes +} + +fn frame(payload: &[u8]) -> Vec { + let mut bytes = vec![0]; + bytes.extend(u32::try_from(payload.len()).unwrap().to_be_bytes()); + bytes.extend(payload); + bytes +} + +fn entry(id: Option, percent: Option) -> Vec { + let mut bytes = Vec::new(); + if let Some(id) = id { + bytes.extend(varint_field(1, id)); + } + if let Some(percent) = percent { + bytes.extend(fixed32(2, percent)); + } + bytes +} + +fn share(id: u64, percent: f32) -> Vec { + entry(Some(id), Some(percent)) +} + +fn payload(aggregate: Option, entries: &[Vec], extra: &[u8]) -> Vec { + let mut config = Vec::new(); + if let Some(aggregate) = aggregate { + config.extend(fixed32(1, aggregate)); + } + config.extend(message(5, &varint_field(1, PERIOD_END))); + let mut current_period = varint_field(1, 2); + current_period.extend(message(2, &varint_field(1, PERIOD_START))); + current_period.extend(message(3, &varint_field(1, PERIOD_END))); + config.extend(message(8, ¤t_period)); + for entry in entries { + config.extend(message(7, entry)); + } + config.extend(extra); + message(1, &config) +} + +/// `data` must parse to the same billing as `baseline` (which carries no +/// products) except for the product list. +fn expect_same_billing(data: &[u8], baseline: &[u8], products: Vec) { + let actual = parse(data); + let without_products = parse(baseline); + assert_eq!(actual.used_percent, without_products.used_percent); + assert_eq!(actual.resets_at, without_products.resets_at); + assert_eq!( + actual.used_percent_is_wire_published, + without_products.used_percent_is_wire_published + ); + assert_eq!( + actual.used_percent_is_implicit_zero, + without_products.used_percent_is_implicit_zero + ); + assert_eq!(actual.product_usage, products); +} + +/// `data` must publish no breakdown. Win-CodexBar's percent parser is stricter +/// than upstream's: a payload whose scan is incomplete also withholds the +/// percent, so it may report 6 or nothing but never an implicit zero. +fn expect_no_breakdown(data: &[u8]) { + let parsed = parse(data); + assert!(parsed.product_usage.is_empty()); + assert!(matches!(parsed.used_percent, None | Some(6.0))); + assert!(!parsed.used_percent_is_implicit_zero); +} + +#[test] +fn live_billing_frame_decodes_the_product_breakdown() { + let parsed = parse(&live_frame()); + + assert_eq!(parsed.used_percent, Some(6.0)); + assert!(parsed.used_percent_is_wire_published); + assert_eq!(parsed.product_usage, chat_and_build()); +} + +#[test] +fn unnamed_products_keep_only_complete_named_shares() { + let named = [share(4, 4.0), share(2, 2.0)]; + let baseline = frame(&payload(Some(6.0), &[], &[])); + + let with_busy_unnamed = [named.to_vec(), vec![share(7, 1.0)]].concat(); + expect_same_billing( + &frame(&payload(Some(6.0), &with_busy_unnamed, &[])), + &baseline, + Vec::new(), + ); + let with_idle_unnamed = [vec![entry(Some(7), None)], named.to_vec()].concat(); + expect_same_billing( + &frame(&payload(Some(6.0), &with_idle_unnamed, &[])), + &baseline, + chat_and_build(), + ); +} + +#[test] +fn missing_duplicate_and_malformed_product_ids_drop_the_breakdown() { + let baseline = frame(&payload(Some(6.0), &[], &[])); + let invalid_entries: Vec>> = vec![ + vec![entry(None, Some(4.0))], + vec![share(4, 4.0), share(4, 2.0)], + // The percent sent as a varint instead of a fixed32. + vec![ + [entry(Some(4), None), varint_field(2, 1)].concat(), + share(2, 6.0), + ], + ]; + for entries in invalid_entries { + expect_same_billing( + &frame(&payload(Some(6.0), &entries, &[])), + &baseline, + Vec::new(), + ); + } + let malformed_entries: Vec>> = vec![ + // Truncated field 1 varint inside an entry. + vec![vec![0x08]], + // The id sent as a fixed32 instead of a varint. + vec![fixed32(1, 4.0), share(2, 6.0)], + ]; + for entries in malformed_entries { + expect_no_breakdown(&frame(&payload(Some(6.0), &entries, &[]))); + } + // A product entry sent as a varint instead of a message. + expect_no_breakdown(&frame(&payload(Some(6.0), &[], &varint_field(7, 4)))); +} + +#[test] +fn omitted_percentages_default_to_zero_and_unknown_entry_fields_are_skipped() { + let baseline = frame(&payload(Some(6.0), &[], &[])); + let mut chat = share(4, 6.0); + chat.extend(varint_field(3, 42)); + chat.extend(message(4, &[0xFF])); + chat.extend(fixed32(5, 3.0)); + // Field 6, fixed64. + chat.extend([0x31]); + chat.extend([0; 8]); + + expect_same_billing( + &frame(&payload(Some(6.0), &[chat, entry(Some(2), None)], &[])), + &baseline, + vec![product("GrokChat", 6.0), product("GrokBuild", 0.0)], + ); +} + +#[test] +fn invalid_or_noncomposing_percentages_drop_the_breakdown() { + let baseline = frame(&payload(Some(6.0), &[], &[])); + let cases = [ + vec![share(4, -1.0), share(2, 7.0)], + vec![share(4, f32::NAN), share(2, 2.0)], + vec![share(2, 2.0)], + ]; + for entries in cases { + expect_same_billing( + &frame(&payload(Some(6.0), &entries, &[])), + &baseline, + Vec::new(), + ); + } +} + +#[test] +fn products_require_one_complete_payload_with_a_published_config_aggregate() { + let named = [share(4, 4.0), share(2, 2.0)]; + let baseline = frame(&payload(Some(6.0), &[], &[])); + expect_same_billing(&baseline, &baseline, Vec::new()); + + // An implicit zero publishes no percent, so shares have nothing to compose. + let implicit_baseline = frame(&payload(None, &[], &[])); + let implicit_zero = frame(&payload(None, &[entry(Some(4), None)], &[])); + let implicit = parse(&implicit_baseline); + assert_eq!(implicit.used_percent, Some(0.0)); + assert!(implicit.used_percent_is_implicit_zero); + expect_same_billing(&implicit_zero, &implicit_baseline, Vec::new()); + + let two_frames = [ + frame(&payload(Some(6.0), &named, &[])), + frame(&payload(Some(6.0), &[], &[])), + ] + .concat(); + let two_frames_baseline = [baseline.clone(), baseline.clone()].concat(); + expect_same_billing(&two_frames, &two_frames_baseline, Vec::new()); + + // A percent nested under another message is not the config aggregate. + let nested_aggregate = message(2, &fixed32(1, 6.0)); + expect_same_billing( + &frame(&payload(None, &named, &nested_aggregate)), + &frame(&payload(None, &[], &nested_aggregate)), + Vec::new(), + ); + + // Malformed bytes elsewhere in the payload make the whole scan incomplete. + let malformed_other_field = [0x62, 0x02, 0x08]; + expect_same_billing( + &frame(&payload(Some(6.0), &named, &malformed_other_field)), + &frame(&payload(Some(6.0), &[], &malformed_other_field)), + Vec::new(), + ); +} + +#[test] +fn duplicate_scalar_fields_cannot_relabel_or_reweight_product_shares() { + let baseline = frame(&payload(Some(6.0), &[], &[])); + let duplicated_id = [share(2, 6.0), varint_field(1, 4)].concat(); + let duplicated_percent = [share(4, 1.0), fixed32(2, 6.0)].concat(); + for entry in [duplicated_id, duplicated_percent] { + expect_same_billing( + &frame(&payload(Some(6.0), &[entry], &[])), + &baseline, + Vec::new(), + ); + } + let repeated_aggregate = frame(&payload(Some(6.0), &[share(4, 6.0)], &fixed32(1, 6.0))); + assert!(parse(&repeated_aggregate).product_usage.is_empty()); +} + +#[test] +fn compressed_or_reserved_frame_flags_fail_closed() { + for flag in [1u8, 2, 3, 0x81] { + let mut data = frame(&payload(Some(6.0), &[share(4, 6.0)], &[])); + data[0] = flag; + // 0x81 also reads as a (malformed) raw protobuf key, so it may parse + // to an unavailable reading instead of an error; either way nothing + // from the frame is published. + if let Ok(parsed) = parse_grpc_web_response_at(&data, now()) { + assert_eq!(parsed.used_percent, None, "flag {flag:#x}"); + assert!(parsed.product_usage.is_empty(), "flag {flag:#x}"); + assert!(!parsed.used_percent_is_implicit_zero, "flag {flag:#x}"); + } + } +} + +#[test] +fn truncated_framing_and_overflowing_product_values_cannot_supply_shares() { + let valid = frame(&payload(Some(6.0), &[share(4, 6.0)], &[])); + for suffix in [vec![0], vec![0, 0xFF, 0xFF, 0xFF, 0xFF]] { + let data = [valid.clone(), suffix].concat(); + assert!(parse_grpc_web_response_at(&data, now()).is_err()); + } + + let mut overflow = vec![0x08]; + overflow.extend([0xFF; 9]); + overflow.push(0x02); + let mut oversized_length = vec![0x12]; + oversized_length.extend(varint(u64::MAX)); + for entry in [overflow, oversized_length] { + expect_no_breakdown(&frame(&payload(Some(6.0), &[entry], &[]))); + } +} diff --git a/rust/src/providers/grok/credits_proxy.rs b/rust/src/providers/grok/credits_proxy.rs index 8172b8d501..65b35501b2 100644 --- a/rust/src/providers/grok/credits_proxy.rs +++ b/rust/src/providers/grok/credits_proxy.rs @@ -15,7 +15,7 @@ use crate::core::ProviderError; use crate::providers::{BoundedBodyError, read_bounded_response}; use super::billing::GrokBillingSnapshot; -use super::product_usage::{GrokProductUsage, LossyProductUsage}; +use super::product_usage::LossyProductUsage; use super::{GrokCredentials, GrokProvider, grok_plan_display_name}; pub(super) const CREDITS_PROXY_ENDPOINT: &str = @@ -32,9 +32,6 @@ const MAX_RESPONSE_BYTES: usize = 256 * 1024; pub(super) struct BearerBilling { pub(super) billing: GrokBillingSnapshot, pub(super) subscription_tier: Option, - /// Product shares of the same payload's wire credit percent; empty for - /// every other percent source (on-demand ratio, period-only, gRPC-web). - pub(super) product_usage: Vec, } #[derive(Deserialize)] @@ -158,9 +155,11 @@ pub(super) fn parse_credits_response( used_percent_is_implicit_zero: false, resets_at, window_minutes, + // Shares of this payload's wire percent; empty for every other + // percent source (on-demand ratio, period-only). + product_usage, }, subscription_tier, - product_usage, }) } @@ -177,11 +176,12 @@ fn adopt_grpc_percent(proxy: BearerBilling, grpc: GrokBillingSnapshot) -> Bearer used_percent: grpc.used_percent, used_percent_is_wire_published: grpc.used_percent_is_wire_published, used_percent_is_implicit_zero: grpc.used_percent_is_implicit_zero, + // The grok.com percent is a different total than any proxy share + // list, so only the grok.com breakdown may accompany it. + product_usage: grpc.product_usage, ..proxy.billing }, subscription_tier: proxy.subscription_tier, - // The grok.com percent is a different total than any proxy share list. - product_usage: Vec::new(), } } @@ -204,7 +204,6 @@ impl GrokProvider { return Ok(BearerBilling { billing, subscription_tier: None, - product_usage: Vec::new(), }); } }; diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 7c72b7421d..e78e3ceb9e 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -138,7 +138,6 @@ impl GrokProvider { let BearerBilling { billing, subscription_tier, - product_usage, } = match self.fetch_bearer_billing(credentials).await { Ok(bearer) => bearer, Err(error) => { @@ -167,12 +166,6 @@ impl GrokProvider { credentials.team_id.clone(), plan, ); - // Reset-credit enrichment adds inventory only, so the breakdown survives it. - let result = product_usage::display_details(&product_usage) - .into_iter() - .fold(result, |result, detail| { - result.with_display_detail(Some(detail)) - }); Ok(match reset_credits { Some(credits) => result.with_inventory_item(credits), None => result, @@ -808,7 +801,13 @@ fn result_from_billing( usage.account_email = email; usage.account_organization = team_id; usage.login_method = login_method; - ProviderFetchResult::new(usage, source_label) + // Reset-credit enrichment adds inventory only, so the breakdown survives it. + product_usage::display_details(&billing.product_usage) + .into_iter() + .fold( + ProviderFetchResult::new(usage, source_label), + |result, detail| result.with_display_detail(Some(detail)), + ) } /// Whether a cookie-path error should invalidate the cached browser session. diff --git a/rust/src/providers/grok/product_usage.rs b/rust/src/providers/grok/product_usage.rs index 6483e398be..b6f56be2f0 100644 --- a/rust/src/providers/grok/product_usage.rs +++ b/rust/src/providers/grok/product_usage.rs @@ -62,16 +62,25 @@ impl LossyProductUsage { /// Shares whose sum matches the raw (unclamped) credit percent from the /// same payload; anything else is dropped. pub(super) fn composing(self, credit_usage_percent: f64) -> Vec { - let products = self.0.unwrap_or_default(); - let sum: f64 = products.iter().map(|entry| entry.used_percent).sum(); - if products.is_empty() || (sum - credit_usage_percent).abs() > COMPOSITION_TOLERANCE_PERCENT - { - return Vec::new(); - } - products + compose(self.0.unwrap_or_default(), credit_usage_percent) } } +/// Shares whose sum matches the raw (unclamped) credit percent of the same +/// payload; anything else, including an empty list, is dropped. Shared by the +/// credits proxy and the grok.com gRPC-web answer (upstream +/// `GrokProductUsage.composing`, v0.69.0). +pub(super) fn compose( + products: Vec, + credit_usage_percent: f64, +) -> Vec { + let sum: f64 = products.iter().map(|entry| entry.used_percent).sum(); + if products.is_empty() || (sum - credit_usage_percent).abs() > COMPOSITION_TOLERANCE_PERCENT { + return Vec::new(); + } + products +} + fn product_label(product: &str) -> &str { match product { "GrokBuild" => "Grok Build", diff --git a/rust/src/providers/grok/product_usage_tests.rs b/rust/src/providers/grok/product_usage_tests.rs index 0cf53e2365..e465021ffb 100644 --- a/rust/src/providers/grok/product_usage_tests.rs +++ b/rust/src/providers/grok/product_usage_tests.rs @@ -41,11 +41,11 @@ fn live_multi_product_capture_keeps_the_composed_breakdown() { let parsed = parse(LIVE_MULTI_PRODUCT_BODY); assert_eq!(parsed.billing.used_percent, Some(6.0)); assert_eq!( - parsed.product_usage, + parsed.billing.product_usage, vec![product("GrokChat", 4.0), product("GrokBuild", 2.0)] ); assert_eq!( - rows(&parsed.product_usage), + rows(&parsed.billing.product_usage), vec![ ( "grok.product.GrokChat".to_string(), @@ -66,7 +66,10 @@ fn weekly_capture_with_one_product_keeps_it() { let parsed = parse( r#"{"config":{"creditUsagePercent":1.0,"productUsage":[{"product":"GrokBuild","usagePercent":1.0}]}}"#, ); - assert_eq!(parsed.product_usage, vec![product("GrokBuild", 1.0)]); + assert_eq!( + parsed.billing.product_usage, + vec![product("GrokBuild", 1.0)] + ); } #[test] @@ -74,7 +77,7 @@ fn malformed_products_drop_the_breakdown_but_not_total_or_period() { let baseline = parse(&format!( r#"{{"config":{{"creditUsagePercent":42,{BASE_PERIOD}}}}}"# )); - assert!(baseline.product_usage.is_empty()); + assert!(baseline.billing.product_usage.is_empty()); let fragments = [ r#""productUsage":null"#, r#""productUsage":{}"#, @@ -92,7 +95,7 @@ fn malformed_products_drop_the_breakdown_but_not_total_or_period() { let parsed = parse(&format!( r#"{{"config":{{"creditUsagePercent":42,{BASE_PERIOD},{fragment}}}}}"# )); - assert!(parsed.product_usage.is_empty(), "{fragment}"); + assert!(parsed.billing.product_usage.is_empty(), "{fragment}"); assert_eq!(parsed.billing.used_percent, Some(42.0), "{fragment}"); assert_eq!(parsed.billing.resets_at, baseline.billing.resets_at); assert_eq!( @@ -109,12 +112,12 @@ fn one_malformed_entry_cannot_hide_behind_the_tolerance() { let parsed = parse( r#"{"config":{"creditUsagePercent":6,"productUsage":[{"product":"GrokChat","usagePercent":5},{"product":"GrokBuild","usagePercent":"1"}]}}"#, ); - assert!(parsed.product_usage.is_empty()); + assert!(parsed.billing.product_usage.is_empty()); assert_eq!(parsed.billing.used_percent, Some(6.0)); let exact_remainder = parse( r#"{"config":{"creditUsagePercent":5,"productUsage":[{"product":"GrokChat","usagePercent":5},{"usagePercent":1}]}}"#, ); - assert!(exact_remainder.product_usage.is_empty()); + assert!(exact_remainder.billing.product_usage.is_empty()); assert_eq!(exact_remainder.billing.used_percent, Some(5.0)); } @@ -132,7 +135,7 @@ fn products_attach_only_to_the_wire_credit_percent() { r#"{{"config":{{{fields},"productUsage":[{{"product":"GrokBuild","usagePercent":3}}]}}}}"# )); assert_eq!(parsed.billing.used_percent, percent, "{fields}"); - assert!(parsed.product_usage.is_empty(), "{fields}"); + assert!(parsed.billing.product_usage.is_empty(), "{fields}"); } } @@ -145,7 +148,7 @@ fn shares_that_do_not_compose_the_percent_are_dropped() { let parsed = parse(&format!( r#"{{"config":{{"creditUsagePercent":30,{BASE_PERIOD},"productUsage":{products}}}}}"# )); - assert!(parsed.product_usage.is_empty(), "{products}"); + assert!(parsed.billing.product_usage.is_empty(), "{products}"); assert_eq!(parsed.billing.used_percent, Some(30.0)); assert_eq!(parsed.billing.window_minutes, Some(10080)); } @@ -157,20 +160,20 @@ fn shares_within_rounding_of_the_percent_are_kept() { r#"{"config":{"creditUsagePercent":10,"productUsage":[{"product":"GrokBuild","usagePercent":6.0},{"product":"GrokChat","usagePercent":3.6}]}}"#, ); assert_eq!( - rounded.product_usage, + rounded.billing.product_usage, vec![product("GrokBuild", 6.0), product("GrokChat", 3.6)] ); let exactly_one_point_off = parse( r#"{"config":{"creditUsagePercent":10,"productUsage":[{"product":"GrokBuild","usagePercent":9.0}]}}"#, ); assert_eq!( - exactly_one_point_off.product_usage, + exactly_one_point_off.billing.product_usage, vec![product("GrokBuild", 9.0)] ); let just_outside = parse( r#"{"config":{"creditUsagePercent":10,"productUsage":[{"product":"GrokBuild","usagePercent":8.9}]}}"#, ); - assert!(just_outside.product_usage.is_empty()); + assert!(just_outside.billing.product_usage.is_empty()); } #[test] @@ -179,7 +182,10 @@ fn shares_compose_the_raw_unclamped_percent() { r#"{"config":{"creditUsagePercent":120,"productUsage":[{"product":"GrokBuild","usagePercent":120}]}}"#, ); assert_eq!(parsed.billing.used_percent, Some(100.0)); - assert_eq!(parsed.product_usage, vec![product("GrokBuild", 120.0)]); + assert_eq!( + parsed.billing.product_usage, + vec![product("GrokBuild", 120.0)] + ); } #[test] @@ -218,9 +224,9 @@ fn product_names_are_trimmed_and_imagine_is_labelled() { let parsed = parse( r#"{"config":{"creditUsagePercent":3,"productUsage":[{"product":" GrokImagine ","usagePercent":2.5},{"product":"Other","usagePercent":0.5}]}}"#, ); - assert_eq!(parsed.product_usage[0].product, "GrokImagine"); + assert_eq!(parsed.billing.product_usage[0].product, "GrokImagine"); assert_eq!( - rows(&parsed.product_usage)[0], + rows(&parsed.billing.product_usage)[0], ( "grok.product.GrokImagine".to_string(), "Grok Imagine".to_string(), @@ -336,3 +342,91 @@ async fn grpc_percent_never_borrows_the_proxy_products() { assert_eq!(result.usage.primary.used_percent, 12.0); assert!(result.display_details().is_empty()); } + +fn shown(result: &ProviderFetchResult) -> Vec<(String, String)> { + result + .display_details() + .iter() + .map(|row| (row.title().to_string(), row.value().to_string())) + .collect() +} + +fn live_breakdown() -> Vec<(String, String)> { + vec![ + ("Grok Chat".to_string(), "4%".to_string()), + ("Grok Build".to_string(), "2%".to_string()), + ] +} + +#[tokio::test] +async fn grpc_breakdown_is_adopted_with_the_grpc_percent_on_a_period_only_proxy_answer() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(200) + .with_body( + r#"{"config":{"currentPeriod":{"start":"2026-08-06T00:00:00Z","end":"2026-08-13T00:00:00Z"},"productUsage":[{"product":"GrokBuild","usagePercent":12}]}}"#, + ) + .create_async() + .await; + server + .mock("POST", "/billing") + .with_status(200) + .with_body(super::billing::web_product_usage_tests::live_frame()) + .create_async() + .await; + + let result = provider_for(&server) + .fetch_with_auth( + &GrokCredentials::from_bearer("token-123"), + GrokAuthKind::OAuth, + &context_without_credits(), + ) + .await + .unwrap(); + + assert_eq!(result.usage.primary.used_percent, 6.0); + assert_eq!(shown(&result), live_breakdown()); +} + +#[tokio::test] +async fn grpc_only_bearer_fallback_shows_the_breakdown() { + let mut server = mockito::Server::new_async().await; + server + .mock("GET", "/credits") + .with_status(500) + .create_async() + .await; + server + .mock("POST", "/billing") + .with_status(200) + .with_body(super::billing::web_product_usage_tests::live_frame()) + .create_async() + .await; + + let result = provider_for(&server) + .fetch_with_auth( + &GrokCredentials::from_bearer("token-123"), + GrokAuthKind::OAuth, + &context_without_credits(), + ) + .await + .unwrap(); + + assert_eq!(result.usage.primary.used_percent, 6.0); + assert_eq!(shown(&result), live_breakdown()); +} + +#[test] +fn cookie_billing_result_shows_the_decoded_product_rows() { + let billing = super::billing::parse_grpc_web_response( + &super::billing::web_product_usage_tests::live_frame(), + ) + .unwrap(); + let result = result_from_cookie_billing(billing); + + assert_eq!(result.source_label, "grok-browser"); + assert_eq!(result.usage.primary.used_percent, 6.0); + assert!(result.usage.secondary.is_none()); + assert_eq!(shown(&result), live_breakdown()); +} diff --git a/rust/src/providers/grok/tests.rs b/rust/src/providers/grok/tests.rs index 585c0255a8..ee12ae6bc5 100644 --- a/rust/src/providers/grok/tests.rs +++ b/rust/src/providers/grok/tests.rs @@ -230,6 +230,7 @@ fn cookie_billing_stays_siloed_from_auth_file_identity() { used_percent_is_implicit_zero: false, resets_at: None, window_minutes: None, + product_usage: Vec::new(), }); assert_eq!(result.source_label, "grok-browser"); assert!(result.usage.account_email.is_none()); @@ -247,6 +248,7 @@ fn billing_snapshot_uses_full_weekly_cycle_for_pace() { used_percent_is_implicit_zero: false, resets_at: Some(resets), window_minutes: Some(crate::core::WEEKLY_WINDOW_MINUTES), + product_usage: Vec::new(), }, "web", None, @@ -278,6 +280,7 @@ fn monthly_cycle_stays_monthly_with_six_days_remaining() { used_percent_is_implicit_zero: false, resets_at: Some(resets), window_minutes: Some(monthly_minutes), + product_usage: Vec::new(), }, "cli", None, @@ -298,6 +301,7 @@ fn reset_distance_alone_does_not_invent_a_cadence() { used_percent_is_implicit_zero: false, resets_at: Some(resets), window_minutes: None, + product_usage: Vec::new(), }, "web", None, @@ -318,6 +322,7 @@ fn period_only_billing_is_informational_not_zero_usage() { used_percent_is_implicit_zero: false, resets_at: Some(resets), window_minutes: None, + product_usage: Vec::new(), }, "cli", Some("user@example.com".into()), @@ -346,6 +351,7 @@ fn unpublished_zero_does_not_reach_the_usage_surface() { used_percent_is_implicit_zero: false, resets_at: Some(Utc.timestamp_opt(1_789_000_000, 0).single().unwrap()), window_minutes: None, + product_usage: Vec::new(), }, "grok-web", None, @@ -365,6 +371,7 @@ fn account_usage_marks_informational_windows_unavailable() { used_percent_is_implicit_zero: false, resets_at: None, window_minutes: Some(crate::core::WEEKLY_WINDOW_MINUTES), + product_usage: Vec::new(), }, "grok-cli", None,