ai_usagebar/grok/
types.rs1use serde::Deserialize;
12
13use crate::error::{AppError, Result};
14use crate::usage::{GrokSnapshot, parse_amount};
15
16pub const SCOPE_TEAM: &str = "SCOPE_TEAM";
21pub const SCOPE_ORGANIZATION: &str = "SCOPE_ORGANIZATION";
22
23#[derive(Debug, Clone, Deserialize, Default)]
26#[serde(default)]
27pub struct Validation {
28 pub scope: Option<String>,
30 #[serde(rename = "scopeId")]
33 pub scope_id: Option<String>,
34 #[serde(rename = "teamId")]
35 pub team_id: Option<String>,
36}
37
38impl Validation {
39 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 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 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#[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 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 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 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 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 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 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 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}