alex_core/
grok_billing.rs

1//! Grok web-billing via grok.com gRPC-web RPC (`GetGrokCreditsConfig`).
2//!
3//! Ported from CodexBar's `GrokWebBillingFetcher` — parses SuperGrok weekly credit
4//! utilization and reset time from the protobuf response. Does not perform HTTP.
5
6use std::borrow::Cow;
7
8/// Default endpoint used by grok.com's usage page.
9pub const GROK_CREDITS_ENDPOINT: &str =
10    "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig";
11
12/// Empty gRPC-web frame body (flags=0, length=0) sent for the no-arg RPC.
13pub const GROK_CREDITS_REQUEST_BODY: &[u8] = &[0x00, 0x00, 0x00, 0x00, 0x00];
14
15#[derive(Debug, Clone, PartialEq)]
16pub struct GrokWebBillingSnapshot {
17    /// Used percent of SuperGrok credits (0–100).
18    pub used_percent: f64,
19    /// Unix epoch seconds when the credit window resets, if known.
20    pub resets_at_s: Option<i64>,
21}
22
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum GrokWebBillingError {
25    EmptyResponse,
26    ParseFailed,
27    RpcFailed { status: i32, message: String },
28}
29
30impl std::fmt::Display for GrokWebBillingError {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        match self {
33            Self::EmptyResponse => write!(f, "Grok web billing returned no protobuf payload"),
34            Self::ParseFailed => write!(f, "Could not parse Grok web billing usage"),
35            Self::RpcFailed { status, message } => {
36                write!(f, "Grok web billing RPC failed with status {status}: {message}")
37            }
38        }
39    }
40}
41
42impl std::error::Error for GrokWebBillingError {}
43
44/// Window label derived from how far out the reset is (matches CodexBar labeling).
45pub fn window_label(resets_at_s: Option<i64>, now_s: i64) -> &'static str {
46    let Some(reset) = resets_at_s else {
47        return "7d";
48    };
49    let secs = (reset - now_s).max(0) as f64;
50    if secs <= 3600.0 {
51        return "7d";
52    }
53    let days = (secs / 86400.0).round() as i64;
54    if (4..=12).contains(&days) {
55        "7d"
56    } else if (20..=45).contains(&days) {
57        "30d"
58    } else {
59        "7d"
60    }
61}
62
63/// Parse a gRPC-web response body into a billing snapshot.
64///
65/// Accepts framed gRPC-web data frames and unframed protobuf payloads (as grok
66/// sometimes returns). `now_s` is used to prefer future reset timestamps.
67pub fn parse_grpc_web_response(
68    data: &[u8],
69    now_s: i64,
70) -> Result<GrokWebBillingSnapshot, GrokWebBillingError> {
71    validate_grpc_web_trailers(data)?;
72
73    let mut payloads = grpc_web_data_frames(data);
74    if payloads.is_empty() && looks_like_protobuf_payload(data) {
75        payloads = vec![data.to_vec()];
76    }
77    if payloads.is_empty() {
78        return Err(GrokWebBillingError::EmptyResponse);
79    }
80
81    let mut scan = ProtobufScan::default();
82    for payload in &payloads {
83        scan.merge(scan_protobuf(payload, 0, &[], 0).0);
84    }
85
86    let parsed_percent = scan
87        .fixed32_fields
88        .iter()
89        .filter(|f| {
90            f.path.last() == Some(&1)
91                && f.value.is_finite()
92                && (0.0..=100.0).contains(&f.value)
93        })
94        .min_by(|a, b| {
95            a.path
96                .len()
97                .cmp(&b.path.len())
98                .then(a.order.cmp(&b.order))
99        })
100        .map(|f| f.value as f64);
101
102    let reset_fields: Vec<(Vec<u64>, i64)> = scan
103        .varint_fields
104        .iter()
105        .filter_map(|f| {
106            let raw = f.value;
107            if (1_700_000_000..=2_100_000_000).contains(&raw) {
108                Some((f.path.clone(), raw as i64))
109            } else {
110                None
111            }
112        })
113        .collect();
114
115    let future_resets: Vec<(Vec<u64>, i64)> = reset_fields
116        .into_iter()
117        .filter(|(_, ts)| *ts > now_s)
118        .collect();
119
120    let preferred = future_resets
121        .iter()
122        .filter(|(path, _)| path.as_slice() == [1, 5, 1])
123        .map(|(_, ts)| *ts)
124        .min();
125    let reset = preferred.or_else(|| future_resets.iter().map(|(_, ts)| *ts).min());
126
127    let has_usage_period = scan.varint_fields.iter().any(|f| {
128        f.path.starts_with(&[1, 6])
129            || (f.path.as_slice() == [1, 8, 1] && (f.value == 1 || f.value == 2))
130    });
131    let no_usage_yet =
132        parsed_percent.is_none() && scan.fixed32_fields.is_empty() && reset.is_some() && has_usage_period;
133
134    let percent = match parsed_percent {
135        Some(p) => p,
136        None if no_usage_yet => 0.0,
137        None => return Err(GrokWebBillingError::ParseFailed),
138    };
139
140    Ok(GrokWebBillingSnapshot {
141        used_percent: percent,
142        resets_at_s: reset,
143    })
144}
145
146/// Validate gRPC status from response headers (keys lowercased).
147pub fn validate_grpc_status_headers(
148    headers: impl IntoIterator<Item = (impl AsRef<str>, impl AsRef<str>)>,
149) -> Result<(), GrokWebBillingError> {
150    let mut status: Option<i32> = None;
151    let mut message = String::new();
152    for (k, v) in headers {
153        let key = k.as_ref().trim().to_ascii_lowercase();
154        if key == "grpc-status" {
155            status = v.as_ref().trim().parse().ok();
156        } else if key == "grpc-message" {
157            message = percent_decode(v.as_ref().trim()).into_owned();
158        }
159    }
160    if let Some(s) = status {
161        if s != 0 {
162            return Err(GrokWebBillingError::RpcFailed { status: s, message });
163        }
164    }
165    Ok(())
166}
167
168pub fn looks_like_protobuf_payload(data: &[u8]) -> bool {
169    let Some(&first) = data.first() else {
170        return false;
171    };
172    let field_number = first >> 3;
173    let wire_type = first & 0x07;
174    field_number > 0 && matches!(wire_type, 0 | 1 | 2 | 5)
175}
176
177pub fn grpc_web_data_frames(data: &[u8]) -> Vec<Vec<u8>> {
178    let mut frames = Vec::new();
179    let mut index = 0;
180    while index < data.len() {
181        if index + 5 > data.len() {
182            return Vec::new();
183        }
184        let flags = data[index];
185        let length = u32::from_be_bytes([
186            data[index + 1],
187            data[index + 2],
188            data[index + 3],
189            data[index + 4],
190        ]) as usize;
191        let start = index + 5;
192        let end = start.saturating_add(length);
193        if end > data.len() {
194            return Vec::new();
195        }
196        if flags & 0x80 == 0 {
197            frames.push(data[start..end].to_vec());
198        }
199        index = end;
200    }
201    frames
202}
203
204fn validate_grpc_web_trailers(data: &[u8]) -> Result<(), GrokWebBillingError> {
205    let fields = grpc_web_trailer_fields(data);
206    if let Some(raw) = fields.get("grpc-status") {
207        if let Ok(status) = raw.parse::<i32>() {
208            if status != 0 {
209                let message = fields
210                    .get("grpc-message")
211                    .cloned()
212                    .unwrap_or_default();
213                return Err(GrokWebBillingError::RpcFailed { status, message });
214            }
215        }
216    }
217    Ok(())
218}
219
220pub fn grpc_web_trailer_fields(data: &[u8]) -> std::collections::HashMap<String, String> {
221    let mut fields = std::collections::HashMap::new();
222    let mut index = 0;
223    while index + 5 <= data.len() {
224        let flags = data[index];
225        let length = u32::from_be_bytes([
226            data[index + 1],
227            data[index + 2],
228            data[index + 3],
229            data[index + 4],
230        ]) as usize;
231        let start = index + 5;
232        let end = start.saturating_add(length);
233        if end > data.len() {
234            break;
235        }
236        if flags & 0x80 != 0 {
237            if let Ok(text) = std::str::from_utf8(&data[start..end]) {
238                for line in text.lines() {
239                    if line.is_empty() {
240                        continue;
241                    }
242                    let Some((key, value)) = line.split_once(':') else {
243                        continue;
244                    };
245                    let key = key.trim().to_ascii_lowercase();
246                    let value = percent_decode(value.trim()).into_owned();
247                    fields.insert(key, value);
248                }
249            }
250        }
251        index = end;
252    }
253    fields
254}
255
256fn percent_decode(s: &str) -> Cow<'_, str> {
257    if !s.contains('%') {
258        return Cow::Borrowed(s);
259    }
260    let bytes = s.as_bytes();
261    let mut out = Vec::with_capacity(bytes.len());
262    let mut i = 0;
263    while i < bytes.len() {
264        if bytes[i] == b'%' && i + 2 < bytes.len() {
265            if let (Some(h), Some(l)) = (from_hex(bytes[i + 1]), from_hex(bytes[i + 2])) {
266                out.push((h << 4) | l);
267                i += 3;
268                continue;
269            }
270        }
271        out.push(bytes[i]);
272        i += 1;
273    }
274    Cow::Owned(String::from_utf8_lossy(&out).into_owned())
275}
276
277fn from_hex(b: u8) -> Option<u8> {
278    match b {
279        b'0'..=b'9' => Some(b - b'0'),
280        b'a'..=b'f' => Some(b - b'a' + 10),
281        b'A'..=b'F' => Some(b - b'A' + 10),
282        _ => None,
283    }
284}
285
286#[derive(Default)]
287struct ProtobufScan {
288    fixed32_fields: Vec<Fixed32Field>,
289    varint_fields: Vec<VarintField>,
290}
291
292struct Fixed32Field {
293    path: Vec<u64>,
294    value: f32,
295    order: usize,
296}
297
298struct VarintField {
299    path: Vec<u64>,
300    value: u64,
301}
302
303impl ProtobufScan {
304    fn merge(&mut self, other: ProtobufScan) {
305        self.fixed32_fields.extend(other.fixed32_fields);
306        self.varint_fields.extend(other.varint_fields);
307    }
308}
309
310fn scan_protobuf(
311    data: &[u8],
312    depth: usize,
313    path: &[u64],
314    order: usize,
315) -> (ProtobufScan, usize) {
316    let mut scan = ProtobufScan::default();
317    let mut index = 0;
318    let mut next_order = order;
319
320    while index < data.len() {
321        let field_start = index;
322        let Some(key) = read_varint(data, &mut index) else {
323            index = field_start + 1;
324            continue;
325        };
326        if key == 0 {
327            index = field_start + 1;
328            continue;
329        }
330        let field_number = key >> 3;
331        let wire_type = key & 0x07;
332        let mut field_path = path.to_vec();
333        field_path.push(field_number);
334
335        match wire_type {
336            0 => {
337                if let Some(value) = read_varint(data, &mut index) {
338                    scan.varint_fields.push(VarintField {
339                        path: field_path,
340                        value,
341                    });
342                } else {
343                    index = field_start + 1;
344                }
345            }
346            1 => {
347                if index + 8 > data.len() {
348                    return (scan, next_order);
349                }
350                index += 8;
351            }
352            2 => {
353                let Some(length) = read_varint(data, &mut index) else {
354                    index = field_start + 1;
355                    continue;
356                };
357                if length as usize > data.len().saturating_sub(index) {
358                    index = field_start + 1;
359                    continue;
360                }
361                let start = index;
362                let end = index + length as usize;
363                if depth < 4 {
364                    let (nested, order) =
365                        scan_protobuf(&data[start..end], depth + 1, &field_path, next_order);
366                    scan.merge(nested);
367                    next_order = order;
368                }
369                index = end;
370            }
371            5 => {
372                if index + 4 > data.len() {
373                    return (scan, next_order);
374                }
375                let bits = u32::from_le_bytes([
376                    data[index],
377                    data[index + 1],
378                    data[index + 2],
379                    data[index + 3],
380                ]);
381                scan.fixed32_fields.push(Fixed32Field {
382                    path: field_path,
383                    value: f32::from_bits(bits),
384                    order: next_order,
385                });
386                next_order += 1;
387                index += 4;
388            }
389            _ => {
390                index = field_start + 1;
391            }
392        }
393    }
394
395    (scan, next_order)
396}
397
398fn read_varint(bytes: &[u8], index: &mut usize) -> Option<u64> {
399    let mut value: u64 = 0;
400    let mut shift: u64 = 0;
401    while *index < bytes.len() && shift < 64 {
402        let byte = bytes[*index];
403        *index += 1;
404        value |= u64::from(byte & 0x7F) << shift;
405        if byte & 0x80 == 0 {
406            return Some(value);
407        }
408        shift += 7;
409    }
410    None
411}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416
417    fn varint(mut value: u64) -> Vec<u8> {
418        let mut bytes = Vec::new();
419        loop {
420            let mut byte = (value & 0x7F) as u8;
421            value >>= 7;
422            if value != 0 {
423                byte |= 0x80;
424            }
425            bytes.push(byte);
426            if value == 0 {
427                break;
428            }
429        }
430        bytes
431    }
432
433    fn protobuf_payload(used_percent: f32, reset_epoch: u64) -> Vec<u8> {
434        let mut data = Vec::new();
435        data.push(0x0D); // field 1, fixed32
436        data.extend_from_slice(&used_percent.to_bits().to_le_bytes());
437        data.push(0x10); // field 2, varint
438        data.extend(varint(reset_epoch));
439        data
440    }
441
442    fn grpc_frame(payload: &[u8], flags: u8) -> Vec<u8> {
443        let mut data = vec![flags];
444        data.extend_from_slice(&(payload.len() as u32).to_be_bytes());
445        data.extend_from_slice(payload);
446        data
447    }
448
449    fn hex(s: &str) -> Vec<u8> {
450        (0..s.len())
451            .step_by(2)
452            .map(|i| u8::from_str_radix(&s[i..i + 2], 16).unwrap())
453            .collect()
454    }
455
456    #[test]
457    fn parses_grok_grpc_web_billing_frame() {
458        let reset = 1_800_000_000u64;
459        let payload = protobuf_payload(42.5, reset);
460        let data = grpc_frame(&payload, 0x00);
461
462        let snap = parse_grpc_web_response(&data, 1_799_000_000).unwrap();
463        assert!((snap.used_percent - 42.5).abs() < 1e-5);
464        assert_eq!(snap.resets_at_s, Some(reset as i64));
465    }
466
467    #[test]
468    fn parses_unframed_grok_billing_protobuf_payload() {
469        // Captured fixture from CodexBar GrokWebBillingFetcherTests.
470        let data = hex(
471            "0a3f0d7f6a9c3f12001a002206088097f3d0062a060880b191d2063a07080215a9389b3f3a07080115d6ea183c421208011206088097f3d0061a060880b191d206",
472        );
473        let snap = parse_grpc_web_response(&data, 1_780_000_000).unwrap();
474        assert!((snap.used_percent - 1.222000002861023).abs() < 1e-6);
475        assert_eq!(snap.resets_at_s, Some(1_782_864_000));
476    }
477
478    #[test]
479    fn parses_unframed_zero_percent_payload() {
480        let reset = 1_800_000_000u64;
481        let payload = protobuf_payload(0.0, reset);
482        assert!(grpc_web_data_frames(&payload).is_empty());
483
484        let snap = parse_grpc_web_response(&payload, 1_799_000_000).unwrap();
485        assert_eq!(snap.used_percent, 0.0);
486        assert_eq!(snap.resets_at_s, Some(reset as i64));
487    }
488
489    #[test]
490    fn does_not_treat_grpc_frame_prefix_as_raw_protobuf() {
491        assert!(!looks_like_protobuf_payload(&[0, 0, 0, 0, 10]));
492    }
493
494    #[test]
495    fn ignores_grpc_web_trailer_frames() {
496        let payload = protobuf_payload(12.25, 1_800_000_001);
497        let trailer = b"grpc-status: 0\r\n";
498        let data = {
499            let mut d = grpc_frame(&payload, 0x00);
500            d.extend(grpc_frame(trailer, 0x80));
501            d
502        };
503        let frames = grpc_web_data_frames(&data);
504        assert_eq!(frames, vec![payload]);
505    }
506
507    #[test]
508    fn rejects_reset_only_billing() {
509        let mut payload = vec![0x10]; // field 2, varint
510        payload.extend(varint(1_800_000_001));
511        let data = grpc_frame(&payload, 0x00);
512        assert_eq!(
513            parse_grpc_web_response(&data, 1_700_000_000),
514            Err(GrokWebBillingError::ParseFailed)
515        );
516    }
517
518    #[test]
519    fn parses_no_usage_yet_as_zero_percent() {
520        // Captured fixture from CodexBar (no fixed32 usage percent, has usage period).
521        let data: Vec<u8> = vec![
522            0x00, 0x00, 0x00, 0x00, 0x37, 0x0A, 0x35, 0x12, 0x00, 0x1A, 0x00, 0x22, 0x06, 0x08,
523            0x80, 0xDA, 0xCF, 0xCF, 0x06, 0x2A, 0x06, 0x08, 0x80, 0x97, 0xF3, 0xD0, 0x06, 0x32,
524            0x09, 0x0A, 0x05, 0x08, 0xEA, 0x0F, 0x10, 0x04, 0x12, 0x00, 0x32, 0x09, 0x0A, 0x05,
525            0x08, 0xEA, 0x0F, 0x10, 0x03, 0x12, 0x00, 0x32, 0x09, 0x0A, 0x05, 0x08, 0xEA, 0x0F,
526            0x10, 0x02, 0x12, 0x00, 0x80, 0x00, 0x00, 0x00, 0x0F, 0x67, 0x72, 0x70, 0x63, 0x2D,
527            0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x3A, 0x30, 0x0D, 0x0A,
528        ];
529        let snap = parse_grpc_web_response(&data, 1_768_000_000).unwrap();
530        assert_eq!(snap.used_percent, 0.0);
531        assert_eq!(snap.resets_at_s, Some(1_780_272_000));
532    }
533
534    #[test]
535    fn parses_omitted_zero_percent_with_current_billing_period() {
536        let data: Vec<u8> = vec![
537            0x00, 0x00, 0x00, 0x00, 0x2A, 0x0A, 0x28, 0x12, 0x00, 0x1A, 0x00, 0x22, 0x06, 0x08,
538            0x80, 0x97, 0xF3, 0xD0, 0x06, 0x2A, 0x06, 0x08, 0x80, 0xB1, 0x91, 0xD2, 0x06, 0x42,
539            0x12, 0x08, 0x01, 0x12, 0x06, 0x08, 0x80, 0x97, 0xF3, 0xD0, 0x06, 0x1A, 0x06, 0x08,
540            0x80, 0xB1, 0x91, 0xD2, 0x06, 0x80, 0x00, 0x00, 0x00, 0x0F, 0x67, 0x72, 0x70, 0x63,
541            0x2D, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x3A, 0x30, 0x0D, 0x0A,
542        ];
543        let snap = parse_grpc_web_response(&data, 1_781_000_000).unwrap();
544        assert_eq!(snap.used_percent, 0.0);
545        assert_eq!(snap.resets_at_s, Some(1_782_864_000));
546    }
547
548    #[test]
549    fn uses_billing_field_one_instead_of_earlier_unrelated_float() {
550        let mut payload = Vec::new();
551        payload.push(0x4D); // field 9, fixed32
552        payload.extend_from_slice(&7.0f32.to_bits().to_le_bytes());
553        payload.push(0x0D); // field 1, fixed32
554        payload.extend_from_slice(&42.0f32.to_bits().to_le_bytes());
555        payload.push(0x10); // field 2, varint
556        payload.extend(varint(1_800_000_001));
557
558        let snap = parse_grpc_web_response(&grpc_frame(&payload, 0x00), 1_700_000_000).unwrap();
559        assert_eq!(snap.used_percent, 42.0);
560    }
561
562    #[test]
563    fn chooses_future_billing_end_instead_of_recent_start() {
564        let recent_start = 1_800_000_000u64;
565        let billing_end = 1_802_592_000u64;
566        let mut payload = Vec::new();
567        payload.push(0x0D);
568        payload.extend_from_slice(&33.0f32.to_bits().to_le_bytes());
569        payload.push(0x10);
570        payload.extend(varint(recent_start));
571        payload.push(0x18);
572        payload.extend(varint(billing_end));
573
574        let snap =
575            parse_grpc_web_response(&grpc_frame(&payload, 0x00), (recent_start + 1800) as i64)
576                .unwrap();
577        assert_eq!(snap.resets_at_s, Some(billing_end as i64));
578    }
579
580    #[test]
581    fn trailer_rpc_failure_is_detected() {
582        let body = grpc_frame(b"grpc-status: 16\r\ngrpc-message: token%20expired\r\n", 0x80);
583        let fields = grpc_web_trailer_fields(&body);
584        assert_eq!(fields.get("grpc-status").map(String::as_str), Some("16"));
585        assert_eq!(
586            fields.get("grpc-message").map(String::as_str),
587            Some("token expired")
588        );
589        assert!(matches!(
590            parse_grpc_web_response(&body, 0),
591            Err(GrokWebBillingError::RpcFailed { status: 16, .. })
592        ));
593    }
594
595    #[test]
596    fn window_label_weekly_and_monthly() {
597        let now = 1_800_000_000i64;
598        assert_eq!(window_label(Some(now + 6 * 86400), now), "7d");
599        assert_eq!(window_label(Some(now + 30 * 86400), now), "30d");
600        assert_eq!(window_label(None, now), "7d");
601    }
602
603    #[test]
604    fn validate_grpc_headers() {
605        assert!(validate_grpc_status_headers([("grpc-status", "0")]).is_ok());
606        assert!(matches!(
607            validate_grpc_status_headers([
608                ("grpc-status", "16"),
609                ("grpc-message", "Invalid%20bearer%20token.")
610            ]),
611            Err(GrokWebBillingError::RpcFailed { status: 16, .. })
612        ));
613    }
614}
alex_core/grok_billing.rs

alex_core/
grok_billing.rs
