Skip to main content

ai_usagebar/grok/
types.rs

1//! Wire types for xAI's Management API prepaid-balance endpoint.
2//!
3//! Confirmed against the official docs
4//! (<https://docs.x.ai/developers/rest-api-reference/management/billing>):
5//! `GET /v1/billing/teams/{team_id}/prepaid/balance` returns `{ changes[], total }`
6//! where `total.val` is a **string of USD cents**. It's an *inverted ledger*:
7//! a top-up shows as a negative value, so the remaining balance in dollars is
8//! `-cents / 100` (sign per community reverse-engineering — verify against a
9//! live account). `total` has no currency field (USD is implied).
10
11use serde::Deserialize;
12
13use crate::error::{AppError, Result};
14use crate::usage::{GrokSnapshot, parse_amount};
15
16/// The scope a management key is issued against
17/// (<https://docs.x.ai/developers/rest-api-reference/management/auth>).
18/// `scopeId` means different things per variant, which is exactly why it cannot
19/// be used as a team id unconditionally.
20pub const SCOPE_TEAM: &str = "SCOPE_TEAM";
21pub const SCOPE_ORGANIZATION: &str = "SCOPE_ORGANIZATION";
22
23/// `GET /auth/management-keys/validation` — used to discover the team id from
24/// the management key when the user hasn't configured one explicitly.
25#[derive(Debug, Clone, Deserialize, Default)]
26#[serde(default)]
27pub struct Validation {
28    /// `SCOPE_TEAM` or `SCOPE_ORGANIZATION`. Absent on older responses.
29    pub scope: Option<String>,
30    // `Option` so an explicit JSON `null` (which a non-team-scoped key may
31    // return) deserializes to `None` instead of failing the whole parse.
32    #[serde(rename = "scopeId")]
33    pub scope_id: Option<String>,
34    #[serde(rename = "teamId")]
35    pub team_id: Option<String>,
36}
37
38impl Validation {
39    /// Resolve the team to bill against.
40    ///
41    /// `scopeId` is only a team id when the key is **team-scoped**. For an
42    /// organization-scoped key it is the *organization* id, and using it as a
43    /// team would query a URL that does not identify the user's team — so we
44    /// ask for an explicit `team_id` instead of guessing.
45    pub fn resolved_team(&self) -> Result<String> {
46        let non_empty =
47            |s: &Option<String>| s.as_deref().filter(|v| !v.is_empty()).map(String::from);
48        let scope = self
49            .scope
50            .as_deref()
51            .unwrap_or("")
52            .trim()
53            .to_ascii_uppercase();
54
55        if scope == SCOPE_ORGANIZATION {
56            // The deprecated `teamId` is still authoritative when present.
57            return non_empty(&self.team_id).ok_or_else(|| {
58                AppError::Other(
59                    "grok: this management key is organization-scoped, so its `scopeId` is an \
60                     organization id, not a team. Set `team_id` under [grok] in config to the \
61                     team you want the prepaid balance for."
62                        .into(),
63                )
64            });
65        }
66
67        // Team-scoped (or a legacy response with no `scope` at all): `scopeId`
68        // is the team, with the deprecated `teamId` as fallback.
69        if scope.is_empty() || scope == SCOPE_TEAM {
70            return non_empty(&self.scope_id)
71                .or_else(|| non_empty(&self.team_id))
72                .ok_or_else(|| {
73                    AppError::Other(
74                        "grok: could not resolve team_id from the management key; \
75                         set `team_id` under [grok] in config"
76                            .into(),
77                    )
78                });
79        }
80
81        Err(AppError::Other(format!(
82            "grok: unrecognized management-key scope {scope:?}; \
83             set `team_id` under [grok] in config"
84        )))
85    }
86}
87
88/// `val` is **required**: a 200 error envelope must not read as a $0.00 balance.
89#[derive(Debug, Clone, Deserialize)]
90pub struct Amount {
91    pub val: String,
92}
93
94#[derive(Debug, Clone, Deserialize)]
95pub struct BalanceResp {
96    pub total: Amount,
97}
98
99pub fn to_snapshot(b: BalanceResp) -> Result<GrokSnapshot> {
100    let cents = parse_amount("grok", "total.val", &b.total.val)?;
101    // Inverted ledger: negative total => credit remaining.
102    Ok(GrokSnapshot {
103        balance: -cents / 100.0,
104    })
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn negative_total_is_positive_remaining_balance() {
113        // Docs example: a $10 top-up shows total.val = "-1000" (cents).
114        let b: BalanceResp = serde_json::from_str(r#"{"total":{"val":"-1000"}}"#).unwrap();
115        assert!((to_snapshot(b).unwrap().balance - 10.0).abs() < 1e-9);
116    }
117
118    #[test]
119    fn empty_or_malformed_total_is_a_schema_error_not_zero() {
120        // A 200 error envelope must not be read as "you have $0.00".
121        assert!(serde_json::from_str::<BalanceResp>("{}").is_err());
122        assert!(serde_json::from_str::<BalanceResp>(r#"{"total":{}}"#).is_err());
123        let b: BalanceResp = serde_json::from_str(r#"{"total":{"val":""}}"#).unwrap();
124        assert!(to_snapshot(b).is_err());
125        let b2: BalanceResp = serde_json::from_str(r#"{"total":{"val":"n/a"}}"#).unwrap();
126        assert!(to_snapshot(b2).is_err());
127    }
128
129    #[test]
130    fn team_scoped_key_uses_scope_id() {
131        let v = Validation {
132            scope: Some(SCOPE_TEAM.into()),
133            scope_id: Some("team-1".into()),
134            team_id: None,
135        };
136        assert_eq!(v.resolved_team().unwrap(), "team-1");
137    }
138
139    #[test]
140    fn organization_scoped_key_does_not_use_scope_id_as_a_team() {
141        // The regression this guards: scopeId here is an ORG id. Querying
142        // /v1/billing/teams/<org-id>/... is not this user's team.
143        let v = Validation {
144            scope: Some(SCOPE_ORGANIZATION.into()),
145            scope_id: Some("org-123".into()),
146            team_id: None,
147        };
148        let err = v.resolved_team().unwrap_err().to_string();
149        assert!(!err.contains("org-123"), "must not adopt the org id: {err}");
150        assert!(
151            err.contains("team_id"),
152            "must tell the user what to do: {err}"
153        );
154
155        // An explicit teamId alongside an org-scoped key is still usable.
156        let v2 = Validation {
157            scope: Some(SCOPE_ORGANIZATION.into()),
158            scope_id: Some("org-123".into()),
159            team_id: Some("team-9".into()),
160        };
161        assert_eq!(v2.resolved_team().unwrap(), "team-9");
162    }
163
164    #[test]
165    fn legacy_response_without_scope_still_resolves() {
166        // Explicit null must not fail the parse.
167        let v: Validation = serde_json::from_str(r#"{"scopeId":null,"teamId":"team-x"}"#).unwrap();
168        assert_eq!(v.resolved_team().unwrap(), "team-x");
169        let v2: Validation = serde_json::from_str(r#"{"scopeId":"team-y"}"#).unwrap();
170        assert_eq!(v2.resolved_team().unwrap(), "team-y");
171        // Nothing to go on → an actionable error, not a silently wrong URL.
172        let empty: Validation = serde_json::from_str("{}").unwrap();
173        assert!(empty.resolved_team().is_err());
174    }
175
176    #[test]
177    fn unknown_scope_is_rejected() {
178        let v = Validation {
179            scope: Some("SCOPE_GALAXY".into()),
180            scope_id: Some("x-1".into()),
181            team_id: None,
182        };
183        assert!(v.resolved_team().is_err());
184    }
185}