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 pub 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 carry_forward_amount_data_from(&mut self, previous: &Self) {
205 if self.has_amount_data() {
206 return;
207 }
208 self.plan_name = previous.plan_name.clone();
209 self.total_balance_usd = previous.total_balance_usd.clone();
210 self.subscription_balance_usd = previous.subscription_balance_usd.clone();
211 self.paygo_balance_usd = previous.paygo_balance_usd.clone();
212 self.monthly_budget_usd = previous.monthly_budget_usd.clone();
213 self.monthly_spent_usd = previous.monthly_spent_usd.clone();
214 self.quota_period = previous.quota_period.clone();
215 self.quota_remaining_usd = previous.quota_remaining_usd.clone();
216 self.quota_limit_usd = previous.quota_limit_usd.clone();
217 self.quota_used_usd = previous.quota_used_usd.clone();
218 self.unlimited_quota = previous.unlimited_quota;
219 self.total_used_usd = previous.total_used_usd.clone();
220 self.today_used_usd = previous.today_used_usd.clone();
221 self.total_requests = previous.total_requests;
222 self.today_requests = previous.today_requests;
223 self.total_tokens = previous.total_tokens;
224 self.today_tokens = previous.today_tokens;
225 }
226
227 pub fn amount_summary(&self) -> String {
228 let mut parts = Vec::new();
229 if let Some(plan) = self.plan_name.as_deref()
230 && !plan.trim().is_empty()
231 {
232 parts.push(format!("plan={plan}"));
233 }
234 if self.unlimited_quota == Some(true) {
235 parts.push("unlimited".to_string());
236 } else {
237 if let Some(total) = self.total_balance_usd.as_deref() {
238 parts.push(format!("total=${total}"));
239 }
240 if let Some(quota) = self.quota_summary() {
241 parts.push(quota);
242 }
243 match (
244 self.monthly_budget_usd.as_deref(),
245 self.monthly_spent_usd.as_deref(),
246 ) {
247 (Some(budget), Some(spent)) => {
248 if let Some(left) = left_from_budget_and_spent(budget, spent) {
249 parts.push(format!("left=${left} budget=${budget} spent=${spent}"));
250 } else {
251 parts.push(format!("budget=${budget} spent=${spent}"));
252 }
253 }
254 (Some(budget), None) => parts.push(format!("budget=${budget}")),
255 (None, Some(spent)) => parts.push(format!("used=${spent}")),
256 (None, None) => {}
257 }
258 if let Some(used) = self.total_used_usd.as_deref() {
259 parts.push(format!("used=${used}"));
260 }
261 if let Some(today) = self.today_used_usd.as_deref() {
262 parts.push(format!("today=${today}"));
263 }
264 if let Some(sub) = self.subscription_balance_usd.as_deref() {
265 parts.push(format!("sub=${sub}"));
266 }
267 if let Some(paygo) = self.paygo_balance_usd.as_deref() {
268 parts.push(format!("paygo=${paygo}"));
269 }
270 }
271 if let Some(requests) = self.total_requests {
272 parts.push(format!("req={requests}"));
273 }
274 if let Some(tokens) = self.total_tokens {
275 parts.push(format!("tok={tokens}"));
276 }
277 if parts.is_empty() {
278 "-".to_string()
279 } else {
280 parts.join(" ")
281 }
282 }
283
284 fn quota_summary(&self) -> Option<String> {
285 let period = self
286 .quota_period
287 .as_deref()
288 .map(str::trim)
289 .filter(|value| !value.is_empty());
290 let remaining = self
291 .quota_remaining_usd
292 .as_deref()
293 .map(str::trim)
294 .filter(|value| !value.is_empty());
295 let limit = self
296 .quota_limit_usd
297 .as_deref()
298 .map(str::trim)
299 .filter(|value| !value.is_empty());
300 let used = self
301 .quota_used_usd
302 .as_deref()
303 .map(str::trim)
304 .filter(|value| !value.is_empty());
305
306 if remaining.is_none() && limit.is_none() && used.is_none() {
307 return None;
308 }
309
310 let mut parts = Vec::new();
311 let quota_label = match period {
312 Some("quota") | None => "quota".to_string(),
313 Some(period) => format!("{period} quota"),
314 };
315 parts.push(quota_label);
316
317 match (remaining, limit, used) {
318 (Some(remaining), Some(limit), Some(used)) => {
319 parts.push(format!("left=${remaining} limit=${limit} used=${used}"))
320 }
321 (Some(remaining), Some(limit), None) => {
322 parts.push(format!("left=${remaining} limit=${limit}"))
323 }
324 (Some(remaining), None, Some(used)) => {
325 parts.push(format!("left=${remaining} used=${used}"))
326 }
327 (Some(remaining), None, None) => parts.push(format!("left=${remaining}")),
328 (None, Some(limit), Some(used)) => parts.push(format!("used=${used} limit=${limit}")),
329 (None, Some(limit), None) => parts.push(format!("limit=${limit}")),
330 (None, None, Some(used)) => parts.push(format!("used=${used}")),
331 (None, None, None) => {}
332 }
333
334 Some(parts.join(" "))
335 }
336}
337
338fn left_from_budget_and_spent(budget: &str, spent: &str) -> Option<String> {
339 let budget = UsdAmount::from_decimal_str(budget)?;
340 let spent = UsdAmount::from_decimal_str(spent)?;
341 Some(budget.saturating_sub(spent).format_usd())
342}
343
344#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
345pub struct StationRoutingBalanceSummary {
346 pub snapshots: usize,
347 #[serde(default)]
348 pub ok: usize,
349 #[serde(default)]
350 pub exhausted: usize,
351 #[serde(default)]
352 pub stale: usize,
353 #[serde(default)]
354 pub error: usize,
355 #[serde(default)]
356 pub unknown: usize,
357 #[serde(default)]
358 pub routing_snapshots: usize,
359 #[serde(default)]
360 pub routing_exhausted: usize,
361 #[serde(default)]
362 pub routing_ignored_exhausted: usize,
363}
364
365impl StationRoutingBalanceSummary {
366 pub fn from_snapshots(snapshots: Option<&[ProviderBalanceSnapshot]>) -> Self {
367 let mut out = Self::default();
368 let Some(snapshots) = snapshots else {
369 return out;
370 };
371
372 for snapshot in snapshots {
373 out.record(snapshot, snapshot.status);
374 }
375 out
376 }
377
378 pub fn from_snapshot_iter_at<'a>(
379 snapshots: impl IntoIterator<Item = &'a ProviderBalanceSnapshot>,
380 now_ms: u64,
381 ) -> Self {
382 let mut out = Self::default();
383 for snapshot in snapshots {
384 out.record(snapshot, snapshot.status_at(now_ms));
385 }
386 out
387 }
388
389 fn record(&mut self, snapshot: &ProviderBalanceSnapshot, status: BalanceSnapshotStatus) {
390 self.snapshots += 1;
391 match status {
392 BalanceSnapshotStatus::Ok => self.ok += 1,
393 BalanceSnapshotStatus::Exhausted => self.exhausted += 1,
394 BalanceSnapshotStatus::Stale => self.stale += 1,
395 BalanceSnapshotStatus::Error => self.error += 1,
396 BalanceSnapshotStatus::Unknown => self.unknown += 1,
397 }
398 if snapshot.exhaustion_affects_routing {
399 self.routing_snapshots += 1;
400 if status == BalanceSnapshotStatus::Exhausted {
401 self.routing_exhausted += 1;
402 }
403 } else if status == BalanceSnapshotStatus::Exhausted {
404 self.routing_ignored_exhausted += 1;
405 }
406 }
407
408 pub fn is_empty(&self) -> bool {
409 self.snapshots == 0
410 }
411}
412
413fn default_exhaustion_affects_routing() -> bool {
414 true
415}
416
417fn bool_is_true(value: &bool) -> bool {
418 *value
419}
420
421#[cfg(test)]
422mod tests {
423 use super::*;
424
425 #[test]
426 fn balance_snapshot_status_labels_are_stable() {
427 assert_eq!(BalanceSnapshotStatus::Unknown.as_str(), "unknown");
428 assert_eq!(BalanceSnapshotStatus::Ok.as_str(), "ok");
429 assert_eq!(BalanceSnapshotStatus::Exhausted.as_str(), "exhausted");
430 assert_eq!(BalanceSnapshotStatus::Stale.as_str(), "stale");
431 assert_eq!(BalanceSnapshotStatus::Error.as_str(), "error");
432 }
433
434 #[test]
435 fn provider_balance_amount_summary_formats_known_amounts() {
436 let snapshot = ProviderBalanceSnapshot {
437 plan_name: Some("monthly".to_string()),
438 total_balance_usd: Some("3.5".to_string()),
439 monthly_budget_usd: Some("5".to_string()),
440 monthly_spent_usd: Some("1.25".to_string()),
441 total_used_usd: Some("7".to_string()),
442 today_used_usd: Some("0.5".to_string()),
443 subscription_balance_usd: Some("2".to_string()),
444 paygo_balance_usd: Some("1.5".to_string()),
445 total_requests: Some(42),
446 total_tokens: Some(1234),
447 ..Default::default()
448 };
449
450 assert_eq!(
451 snapshot.amount_summary(),
452 "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"
453 );
454 }
455
456 #[test]
457 fn provider_balance_amount_summary_keeps_wallet_with_quota() {
458 let snapshot = ProviderBalanceSnapshot {
459 plan_name: Some("rightcode".to_string()),
460 total_balance_usd: Some("3.25".to_string()),
461 quota_period: Some("daily".to_string()),
462 quota_remaining_usd: Some("7.5".to_string()),
463 quota_limit_usd: Some("20".to_string()),
464 quota_used_usd: Some("12.5".to_string()),
465 ..Default::default()
466 };
467
468 assert_eq!(
469 snapshot.amount_summary(),
470 "plan=rightcode total=$3.25 daily quota left=$7.5 limit=$20 used=$12.5"
471 );
472 }
473
474 #[test]
475 fn provider_balance_amount_summary_prioritizes_unlimited_quota() {
476 let snapshot = ProviderBalanceSnapshot {
477 plan_name: Some("cx".to_string()),
478 unlimited_quota: Some(true),
479 quota_used_usd: Some("106065.94".to_string()),
480 ..Default::default()
481 };
482
483 assert_eq!(snapshot.amount_summary(), "plan=cx unlimited");
484 }
485
486 #[test]
487 fn routing_exhausted_respects_snapshot_routing_flag() {
488 let mut snapshot = ProviderBalanceSnapshot {
489 exhausted: Some(true),
490 exhaustion_affects_routing: false,
491 ..Default::default()
492 };
493 snapshot.refresh_status(100);
494
495 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
496 assert!(!snapshot.routing_exhausted());
497 assert!(snapshot.routing_ignored_exhaustion());
498 }
499}