1use serde::{Deserialize, Serialize};
2
3use crate::pricing::UsdAmount;
4
5#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
6#[serde(rename_all = "snake_case")]
7pub enum BalanceSnapshotStatus {
8 #[default]
9 Unknown,
10 Ok,
11 Exhausted,
12 Stale,
13 Error,
14}
15
16impl BalanceSnapshotStatus {
17 pub fn as_str(self) -> &'static str {
18 match self {
19 BalanceSnapshotStatus::Unknown => "unknown",
20 BalanceSnapshotStatus::Ok => "ok",
21 BalanceSnapshotStatus::Exhausted => "exhausted",
22 BalanceSnapshotStatus::Stale => "stale",
23 BalanceSnapshotStatus::Error => "error",
24 }
25 }
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
29pub struct ProviderBalanceSnapshot {
30 pub provider_id: String,
31 #[serde(default, skip_serializing_if = "Option::is_none")]
32 pub station_name: Option<String>,
33 #[serde(default, skip_serializing_if = "Option::is_none")]
34 pub upstream_index: Option<usize>,
35 pub source: String,
36 pub fetched_at_ms: u64,
37 #[serde(default, skip_serializing_if = "Option::is_none")]
38 pub stale_after_ms: Option<u64>,
39 #[serde(default)]
40 pub stale: bool,
41 #[serde(default)]
42 pub status: BalanceSnapshotStatus,
43 #[serde(default, skip_serializing_if = "Option::is_none")]
44 pub exhausted: Option<bool>,
45 #[serde(
46 default = "default_exhaustion_affects_routing",
47 skip_serializing_if = "bool_is_true"
48 )]
49 pub exhaustion_affects_routing: bool,
50 #[serde(default, skip_serializing_if = "Option::is_none")]
51 pub plan_name: Option<String>,
52 #[serde(default, skip_serializing_if = "Option::is_none")]
53 pub total_balance_usd: Option<String>,
54 #[serde(default, skip_serializing_if = "Option::is_none")]
55 pub subscription_balance_usd: Option<String>,
56 #[serde(default, skip_serializing_if = "Option::is_none")]
57 pub paygo_balance_usd: Option<String>,
58 #[serde(default, skip_serializing_if = "Option::is_none")]
59 pub monthly_budget_usd: Option<String>,
60 #[serde(default, skip_serializing_if = "Option::is_none")]
61 pub monthly_spent_usd: Option<String>,
62 #[serde(default, skip_serializing_if = "Option::is_none")]
63 pub quota_period: Option<String>,
64 #[serde(default, skip_serializing_if = "Option::is_none")]
65 pub quota_remaining_usd: Option<String>,
66 #[serde(default, skip_serializing_if = "Option::is_none")]
67 pub quota_limit_usd: Option<String>,
68 #[serde(default, skip_serializing_if = "Option::is_none")]
69 pub quota_used_usd: Option<String>,
70 #[serde(default, skip_serializing_if = "Option::is_none")]
71 pub unlimited_quota: Option<bool>,
72 #[serde(default, skip_serializing_if = "Option::is_none")]
73 pub total_used_usd: Option<String>,
74 #[serde(default, skip_serializing_if = "Option::is_none")]
75 pub today_used_usd: Option<String>,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
77 pub total_requests: Option<u64>,
78 #[serde(default, skip_serializing_if = "Option::is_none")]
79 pub today_requests: Option<u64>,
80 #[serde(default, skip_serializing_if = "Option::is_none")]
81 pub total_tokens: Option<u64>,
82 #[serde(default, skip_serializing_if = "Option::is_none")]
83 pub today_tokens: Option<u64>,
84 #[serde(default, skip_serializing_if = "Option::is_none")]
85 pub error: Option<String>,
86}
87
88impl Default for ProviderBalanceSnapshot {
89 fn default() -> Self {
90 Self {
91 provider_id: String::new(),
92 station_name: None,
93 upstream_index: None,
94 source: String::new(),
95 fetched_at_ms: 0,
96 stale_after_ms: None,
97 stale: false,
98 status: BalanceSnapshotStatus::Unknown,
99 exhausted: None,
100 exhaustion_affects_routing: true,
101 plan_name: None,
102 total_balance_usd: None,
103 subscription_balance_usd: None,
104 paygo_balance_usd: None,
105 monthly_budget_usd: None,
106 monthly_spent_usd: None,
107 quota_period: None,
108 quota_remaining_usd: None,
109 quota_limit_usd: None,
110 quota_used_usd: None,
111 unlimited_quota: None,
112 total_used_usd: None,
113 today_used_usd: None,
114 total_requests: None,
115 today_requests: None,
116 total_tokens: None,
117 today_tokens: None,
118 error: None,
119 }
120 }
121}
122
123impl ProviderBalanceSnapshot {
124 pub fn new(
125 provider_id: impl Into<String>,
126 station_name: impl Into<String>,
127 upstream_index: usize,
128 source: impl Into<String>,
129 fetched_at_ms: u64,
130 stale_after_ms: Option<u64>,
131 ) -> Self {
132 let mut snapshot = Self {
133 provider_id: provider_id.into(),
134 station_name: Some(station_name.into()),
135 upstream_index: Some(upstream_index),
136 source: source.into(),
137 fetched_at_ms,
138 stale_after_ms,
139 ..Self::default()
140 };
141 snapshot.refresh_status(fetched_at_ms);
142 snapshot
143 }
144
145 pub fn with_error(mut self, error: impl Into<String>) -> Self {
146 self.error = Some(error.into());
147 self.exhausted = None;
148 self.refresh_status(self.fetched_at_ms);
149 self
150 }
151
152 pub fn refresh_status(&mut self, now_ms: u64) {
153 self.stale = self.stale_at(now_ms);
154 self.status = self.status_at(now_ms);
155 }
156
157 pub fn stale_at(&self, now_ms: u64) -> bool {
158 self.stale_after_ms
159 .is_some_and(|stale_after_ms| now_ms > stale_after_ms)
160 }
161
162 pub fn status_at(&self, now_ms: u64) -> BalanceSnapshotStatus {
163 let stale = self.stale_at(now_ms);
164 if self
165 .error
166 .as_deref()
167 .is_some_and(|value| !value.trim().is_empty())
168 {
169 BalanceSnapshotStatus::Error
170 } else if self.exhausted == Some(true) {
171 BalanceSnapshotStatus::Exhausted
172 } else if stale {
173 BalanceSnapshotStatus::Stale
174 } else if self.exhausted == Some(false) || self.has_amount_data() {
175 BalanceSnapshotStatus::Ok
176 } else {
177 BalanceSnapshotStatus::Unknown
178 }
179 }
180
181 pub fn routing_exhausted(&self) -> bool {
182 self.exhaustion_affects_routing && self.status == BalanceSnapshotStatus::Exhausted
183 }
184
185 pub fn routing_ignored_exhaustion(&self) -> bool {
186 self.status == BalanceSnapshotStatus::Exhausted && !self.exhaustion_affects_routing
187 }
188
189 fn has_amount_data(&self) -> bool {
190 self.total_balance_usd.is_some()
191 || self.subscription_balance_usd.is_some()
192 || self.paygo_balance_usd.is_some()
193 || self.monthly_budget_usd.is_some()
194 || self.monthly_spent_usd.is_some()
195 || self.quota_period.is_some()
196 || self.quota_remaining_usd.is_some()
197 || self.quota_limit_usd.is_some()
198 || self.quota_used_usd.is_some()
199 || self.unlimited_quota == Some(true)
200 || self.total_used_usd.is_some()
201 || self.today_used_usd.is_some()
202 }
203
204 pub fn amount_summary(&self) -> String {
205 let mut parts = Vec::new();
206 if let Some(plan) = self.plan_name.as_deref()
207 && !plan.trim().is_empty()
208 {
209 parts.push(format!("plan={plan}"));
210 }
211 if self.unlimited_quota == Some(true) {
212 parts.push("unlimited".to_string());
213 } else {
214 if let Some(total) = self.total_balance_usd.as_deref() {
215 parts.push(format!("total=${total}"));
216 }
217 if let Some(quota) = self.quota_summary() {
218 parts.push(quota);
219 }
220 match (
221 self.monthly_budget_usd.as_deref(),
222 self.monthly_spent_usd.as_deref(),
223 ) {
224 (Some(budget), Some(spent)) => {
225 if let Some(left) = left_from_budget_and_spent(budget, spent) {
226 parts.push(format!("left=${left} budget=${budget} spent=${spent}"));
227 } else {
228 parts.push(format!("budget=${budget} spent=${spent}"));
229 }
230 }
231 (Some(budget), None) => parts.push(format!("budget=${budget}")),
232 (None, Some(spent)) => parts.push(format!("used=${spent}")),
233 (None, None) => {}
234 }
235 if let Some(used) = self.total_used_usd.as_deref() {
236 parts.push(format!("used=${used}"));
237 }
238 if let Some(today) = self.today_used_usd.as_deref() {
239 parts.push(format!("today=${today}"));
240 }
241 if let Some(sub) = self.subscription_balance_usd.as_deref() {
242 parts.push(format!("sub=${sub}"));
243 }
244 if let Some(paygo) = self.paygo_balance_usd.as_deref() {
245 parts.push(format!("paygo=${paygo}"));
246 }
247 }
248 if let Some(requests) = self.total_requests {
249 parts.push(format!("req={requests}"));
250 }
251 if let Some(tokens) = self.total_tokens {
252 parts.push(format!("tok={tokens}"));
253 }
254 if parts.is_empty() {
255 "-".to_string()
256 } else {
257 parts.join(" ")
258 }
259 }
260
261 fn quota_summary(&self) -> Option<String> {
262 let period = self
263 .quota_period
264 .as_deref()
265 .map(str::trim)
266 .filter(|value| !value.is_empty());
267 let remaining = self
268 .quota_remaining_usd
269 .as_deref()
270 .map(str::trim)
271 .filter(|value| !value.is_empty());
272 let limit = self
273 .quota_limit_usd
274 .as_deref()
275 .map(str::trim)
276 .filter(|value| !value.is_empty());
277 let used = self
278 .quota_used_usd
279 .as_deref()
280 .map(str::trim)
281 .filter(|value| !value.is_empty());
282
283 if remaining.is_none() && limit.is_none() && used.is_none() {
284 return None;
285 }
286
287 let mut parts = Vec::new();
288 let quota_label = match period {
289 Some("quota") | None => "quota".to_string(),
290 Some(period) => format!("{period} quota"),
291 };
292 parts.push(quota_label);
293
294 match (remaining, limit, used) {
295 (Some(remaining), Some(limit), Some(used)) => {
296 parts.push(format!("left=${remaining} limit=${limit} used=${used}"))
297 }
298 (Some(remaining), Some(limit), None) => {
299 parts.push(format!("left=${remaining} limit=${limit}"))
300 }
301 (Some(remaining), None, Some(used)) => {
302 parts.push(format!("left=${remaining} used=${used}"))
303 }
304 (Some(remaining), None, None) => parts.push(format!("left=${remaining}")),
305 (None, Some(limit), Some(used)) => parts.push(format!("used=${used} limit=${limit}")),
306 (None, Some(limit), None) => parts.push(format!("limit=${limit}")),
307 (None, None, Some(used)) => parts.push(format!("used=${used}")),
308 (None, None, None) => {}
309 }
310
311 Some(parts.join(" "))
312 }
313}
314
315fn left_from_budget_and_spent(budget: &str, spent: &str) -> Option<String> {
316 let budget = UsdAmount::from_decimal_str(budget)?;
317 let spent = UsdAmount::from_decimal_str(spent)?;
318 Some(budget.saturating_sub(spent).format_usd())
319}
320
321#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
322pub struct StationRoutingBalanceSummary {
323 pub snapshots: usize,
324 #[serde(default)]
325 pub ok: usize,
326 #[serde(default)]
327 pub exhausted: usize,
328 #[serde(default)]
329 pub stale: usize,
330 #[serde(default)]
331 pub error: usize,
332 #[serde(default)]
333 pub unknown: usize,
334 #[serde(default)]
335 pub routing_snapshots: usize,
336 #[serde(default)]
337 pub routing_exhausted: usize,
338 #[serde(default)]
339 pub routing_ignored_exhausted: usize,
340}
341
342impl StationRoutingBalanceSummary {
343 pub fn from_snapshots(snapshots: Option<&[ProviderBalanceSnapshot]>) -> Self {
344 let mut out = Self::default();
345 let Some(snapshots) = snapshots else {
346 return out;
347 };
348
349 for snapshot in snapshots {
350 out.record(snapshot, snapshot.status);
351 }
352 out
353 }
354
355 pub fn from_snapshot_iter_at<'a>(
356 snapshots: impl IntoIterator<Item = &'a ProviderBalanceSnapshot>,
357 now_ms: u64,
358 ) -> Self {
359 let mut out = Self::default();
360 for snapshot in snapshots {
361 out.record(snapshot, snapshot.status_at(now_ms));
362 }
363 out
364 }
365
366 fn record(&mut self, snapshot: &ProviderBalanceSnapshot, status: BalanceSnapshotStatus) {
367 self.snapshots += 1;
368 match status {
369 BalanceSnapshotStatus::Ok => self.ok += 1,
370 BalanceSnapshotStatus::Exhausted => self.exhausted += 1,
371 BalanceSnapshotStatus::Stale => self.stale += 1,
372 BalanceSnapshotStatus::Error => self.error += 1,
373 BalanceSnapshotStatus::Unknown => self.unknown += 1,
374 }
375 if snapshot.exhaustion_affects_routing {
376 self.routing_snapshots += 1;
377 if status == BalanceSnapshotStatus::Exhausted {
378 self.routing_exhausted += 1;
379 }
380 } else if status == BalanceSnapshotStatus::Exhausted {
381 self.routing_ignored_exhausted += 1;
382 }
383 }
384
385 pub fn is_empty(&self) -> bool {
386 self.snapshots == 0
387 }
388}
389
390fn default_exhaustion_affects_routing() -> bool {
391 true
392}
393
394fn bool_is_true(value: &bool) -> bool {
395 *value
396}
397
398#[cfg(test)]
399mod tests {
400 use super::*;
401
402 #[test]
403 fn balance_snapshot_status_labels_are_stable() {
404 assert_eq!(BalanceSnapshotStatus::Unknown.as_str(), "unknown");
405 assert_eq!(BalanceSnapshotStatus::Ok.as_str(), "ok");
406 assert_eq!(BalanceSnapshotStatus::Exhausted.as_str(), "exhausted");
407 assert_eq!(BalanceSnapshotStatus::Stale.as_str(), "stale");
408 assert_eq!(BalanceSnapshotStatus::Error.as_str(), "error");
409 }
410
411 #[test]
412 fn provider_balance_amount_summary_formats_known_amounts() {
413 let snapshot = ProviderBalanceSnapshot {
414 plan_name: Some("monthly".to_string()),
415 total_balance_usd: Some("3.5".to_string()),
416 monthly_budget_usd: Some("5".to_string()),
417 monthly_spent_usd: Some("1.25".to_string()),
418 total_used_usd: Some("7".to_string()),
419 today_used_usd: Some("0.5".to_string()),
420 subscription_balance_usd: Some("2".to_string()),
421 paygo_balance_usd: Some("1.5".to_string()),
422 total_requests: Some(42),
423 total_tokens: Some(1234),
424 ..Default::default()
425 };
426
427 assert_eq!(
428 snapshot.amount_summary(),
429 "plan=monthly total=$3.5 left=$3.75 budget=$5 spent=$1.25 used=$7 today=$0.5 sub=$2 paygo=$1.5 req=42 tok=1234"
430 );
431 }
432
433 #[test]
434 fn provider_balance_amount_summary_keeps_wallet_with_quota() {
435 let snapshot = ProviderBalanceSnapshot {
436 plan_name: Some("rightcode".to_string()),
437 total_balance_usd: Some("3.25".to_string()),
438 quota_period: Some("daily".to_string()),
439 quota_remaining_usd: Some("7.5".to_string()),
440 quota_limit_usd: Some("20".to_string()),
441 quota_used_usd: Some("12.5".to_string()),
442 ..Default::default()
443 };
444
445 assert_eq!(
446 snapshot.amount_summary(),
447 "plan=rightcode total=$3.25 daily quota left=$7.5 limit=$20 used=$12.5"
448 );
449 }
450
451 #[test]
452 fn provider_balance_amount_summary_prioritizes_unlimited_quota() {
453 let snapshot = ProviderBalanceSnapshot {
454 plan_name: Some("cx".to_string()),
455 unlimited_quota: Some(true),
456 quota_used_usd: Some("106065.94".to_string()),
457 ..Default::default()
458 };
459
460 assert_eq!(snapshot.amount_summary(), "plan=cx unlimited");
461 }
462
463 #[test]
464 fn routing_exhausted_respects_snapshot_routing_flag() {
465 let mut snapshot = ProviderBalanceSnapshot {
466 exhausted: Some(true),
467 exhaustion_affects_routing: false,
468 ..Default::default()
469 };
470 snapshot.refresh_status(100);
471
472 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
473 assert!(!snapshot.routing_exhausted());
474 assert!(snapshot.routing_ignored_exhaustion());
475 }
476}