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,