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 #[serde(default, skip_serializing_if = "Option::is_none")]
36 pub provider_endpoint_key: Option<String>,
37 pub source: String,
38 pub fetched_at_ms: u64,
39 #[serde(default, skip_serializing_if = "Option::is_none")]
40 pub stale_after_ms: Option<u64>,
41 #[serde(default)]
42 pub stale: bool,
43 #[serde(default)]
44 pub status: BalanceSnapshotStatus,
45 #[serde(default, skip_serializing_if = "Option::is_none")]
46 pub exhausted: Option<bool>,
47 #[serde(
48 default = "default_exhaustion_affects_routing",
49 skip_serializing_if = "bool_is_true"
50 )]
51 pub exhaustion_affects_routing: bool,
52 #[serde(default, skip_serializing_if = "Option::is_none")]
53 pub plan_name: Option<String>,
54 #[serde(default, skip_serializing_if = "Option::is_none")]
55 pub total_balance_usd: Option<String>,
56 #[serde(default, skip_serializing_if = "Option::is_none")]
57 pub subscription_balance_usd: Option<String>,
58 #[serde(default, skip_serializing_if = "Option::is_none")]
59 pub paygo_balance_usd: Option<String>,
60 #[serde(default, skip_serializing_if = "Option::is_none")]
61 pub monthly_budget_usd: Option<String>,
62 #[serde(default, skip_serializing_if = "Option::is_none")]
63 pub monthly_spent_usd: Option<String>,
64 #[serde(default, skip_serializing_if = "Option::is_none")]
65 pub quota_period: Option<String>,
66 #[serde(default, skip_serializing_if = "Option::is_none")]
67 pub quota_remaining_usd: Option<String>,
68 #[serde(default, skip_serializing_if = "Option::is_none")]
69 pub quota_limit_usd: Option<String>,
70 #[serde(default, skip_serializing_if = "Option::is_none")]
71 pub quota_used_usd: Option<String>,
72 #[serde(default, skip_serializing_if = "Option::is_none")]
73 pub unlimited_quota: Option<bool>,
74 #[serde(default, skip_serializing_if = "Option::is_none")]
75 pub total_used_usd: Option<String>,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
77 pub today_used_usd: Option<String>,
78 #[serde(default, skip_serializing_if = "Option::is_none")]
79 pub total_requests: Option<u64>,
80 #[serde(default, skip_serializing_if = "Option::is_none")]
81 pub today_requests: Option<u64>,
82 #[serde(default, skip_serializing_if = "Option::is_none")]
83 pub total_tokens: Option<u64>,
84 #[serde(default, skip_serializing_if = "Option::is_none")]
85 pub today_tokens: Option<u64>,
86 #[serde(default, skip_serializing_if = "Option::is_none")]
87 pub subscription_expires_at: Option<String>,
88 #[serde(default, skip_serializing_if = "Vec::is_empty")]
89 pub usage_windows: Vec<ProviderUsageWindow>,
90 #[serde(default, skip_serializing_if = "Option::is_none")]
91 pub usage_rate: Option<ProviderUsageRateSnapshot>,
92 #[serde(default, skip_serializing_if = "Vec::is_empty")]
93 pub usage_model_stats: Vec<ProviderUsageModelStat>,
94 #[serde(default, skip_serializing_if = "Vec::is_empty")]
95 pub usage_alerts: Vec<ProviderUsageAlert>,
96 #[serde(default, skip_serializing_if = "Option::is_none")]
97 pub error: Option<String>,
98}
99
100#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
101pub struct ProviderUsageRateSnapshot {
102 #[serde(default, skip_serializing_if = "Option::is_none")]
103 pub average_duration_ms: Option<String>,
104 #[serde(default, skip_serializing_if = "Option::is_none")]
105 pub rpm: Option<String>,
106 #[serde(default, skip_serializing_if = "Option::is_none")]
107 pub tpm: Option<String>,
108}
109
110impl ProviderUsageRateSnapshot {
111 pub fn is_empty(&self) -> bool {
112 self.average_duration_ms.is_none() && self.rpm.is_none() && self.tpm.is_none()
113 }
114}
115
116#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
117pub struct ProviderUsageWindow {
118 pub period: String,
119 #[serde(default, skip_serializing_if = "Option::is_none")]
120 pub used_usd: Option<String>,
121 #[serde(default, skip_serializing_if = "Option::is_none")]
122 pub limit_usd: Option<String>,
123 #[serde(default, skip_serializing_if = "Option::is_none")]
124 pub remaining_usd: Option<String>,
125 #[serde(default, skip_serializing_if = "Option::is_none")]
126 pub unlimited: Option<bool>,
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
130pub struct ProviderUsageModelStat {
131 pub model: String,
132 #[serde(default, skip_serializing_if = "Option::is_none")]
133 pub request_count: Option<u64>,
134 #[serde(default, skip_serializing_if = "Option::is_none")]
135 pub input_tokens: Option<u64>,
136 #[serde(default, skip_serializing_if = "Option::is_none")]
137 pub output_tokens: Option<u64>,
138 #[serde(default, skip_serializing_if = "Option::is_none")]
139 pub total_tokens: Option<u64>,
140 #[serde(default, skip_serializing_if = "Option::is_none")]
141 pub input_cost_usd: Option<String>,
142 #[serde(default, skip_serializing_if = "Option::is_none")]
143 pub output_cost_usd: Option<String>,
144 #[serde(default, skip_serializing_if = "Option::is_none")]
145 pub total_cost_usd: Option<String>,
146}
147
148#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
149#[serde(rename_all = "snake_case")]
150pub enum ProviderUsageAlertKind {
151 DailyUsage80,
152 DailyUsage95,
153 LowBalance,
154 SubscriptionExpiringSoon,
155 SubscriptionExpired,
156}
157
158impl ProviderUsageAlertKind {
159 pub fn as_str(self) -> &'static str {
160 match self {
161 ProviderUsageAlertKind::DailyUsage80 => "daily_usage_80",
162 ProviderUsageAlertKind::DailyUsage95 => "daily_usage_95",
163 ProviderUsageAlertKind::LowBalance => "low_balance",
164 ProviderUsageAlertKind::SubscriptionExpiringSoon => "subscription_expiring_soon",
165 ProviderUsageAlertKind::SubscriptionExpired => "subscription_expired",
166 }
167 }
168}
169
170#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
171pub struct ProviderUsageAlert {
172 pub kind: ProviderUsageAlertKind,
173 pub message: String,
174}
175
176impl Default for ProviderBalanceSnapshot {
177 fn default() -> Self {
178 Self {
179 provider_id: String::new(),
180 station_name: None,
181 upstream_index: None,
182 provider_endpoint_key: None,
183 source: String::new(),
184 fetched_at_ms: 0,
185 stale_after_ms: None,
186 stale: false,
187 status: BalanceSnapshotStatus::Unknown,
188 exhausted: None,
189 exhaustion_affects_routing: true,
190 plan_name: None,
191 total_balance_usd: None,
192 subscription_balance_usd: None,
193 paygo_balance_usd: None,
194 monthly_budget_usd: None,
195 monthly_spent_usd: None,
196 quota_period: None,
197 quota_remaining_usd: None,
198 quota_limit_usd: None,
199 quota_used_usd: None,
200 unlimited_quota: None,
201 total_used_usd: None,
202 today_used_usd: None,
203 total_requests: None,
204 today_requests: None,
205 total_tokens: None,
206 today_tokens: None,
207 subscription_expires_at: None,
208 usage_windows: Vec::new(),
209 usage_rate: None,
210 usage_model_stats: Vec::new(),
211 usage_alerts: Vec::new(),
212 error: None,
213 }
214 }
215}
216
217impl ProviderBalanceSnapshot {
218 pub fn new(
219 provider_id: impl Into<String>,
220 station_name: impl Into<String>,
221 upstream_index: usize,
222 source: impl Into<String>,
223 fetched_at_ms: u64,
224 stale_after_ms: Option<u64>,
225 ) -> Self {
226 let mut snapshot = Self {
227 provider_id: provider_id.into(),
228 station_name: Some(station_name.into()),
229 upstream_index: Some(upstream_index),
230 source: source.into(),
231 fetched_at_ms,
232 stale_after_ms,
233 ..Self::default()
234 };
235 snapshot.refresh_status(fetched_at_ms);
236 snapshot
237 }
238
239 pub fn with_error(mut self, error: impl Into<String>) -> Self {
240 self.error = Some(error.into());
241 self.exhausted = None;
242 self.refresh_status(self.fetched_at_ms);
243 self
244 }
245
246 pub fn with_provider_endpoint_key(mut self, key: impl Into<String>) -> Self {
247 let key = key.into();
248 if !key.trim().is_empty() {
249 self.provider_endpoint_key = Some(key);
250 }
251 self
252 }
253
254 pub fn refresh_status(&mut self, now_ms: u64) {
255 self.stale = self.stale_at(now_ms);
256 self.status = self.status_at(now_ms);
257 }
258
259 pub fn stale_at(&self, now_ms: u64) -> bool {
260 self.stale_after_ms
261 .is_some_and(|stale_after_ms| now_ms > stale_after_ms)
262 }
263
264 pub fn status_at(&self, now_ms: u64) -> BalanceSnapshotStatus {
265 let stale = self.stale_at(now_ms);
266 if self
267 .error
268 .as_deref()
269 .is_some_and(|value| !value.trim().is_empty())
270 {
271 BalanceSnapshotStatus::Error
272 } else if self.exhausted == Some(true) {
273 BalanceSnapshotStatus::Exhausted
274 } else if stale {
275 BalanceSnapshotStatus::Stale
276 } else if self.exhausted == Some(false) || self.has_amount_data() {
277 BalanceSnapshotStatus::Ok
278 } else {
279 BalanceSnapshotStatus::Unknown
280 }
281 }
282
283 pub fn routing_exhausted(&self) -> bool {
284 self.exhaustion_affects_routing && self.status == BalanceSnapshotStatus::Exhausted
285 }
286
287 pub fn routing_ignored_exhaustion(&self) -> bool {
288 self.status == BalanceSnapshotStatus::Exhausted && !self.exhaustion_affects_routing
289 }
290
291 pub fn has_amount_data(&self) -> bool {
292 self.total_balance_usd.is_some()
293 || self.subscription_balance_usd.is_some()
294 || self.paygo_balance_usd.is_some()
295 || self.monthly_budget_usd.is_some()
296 || self.monthly_spent_usd.is_some()
297 || self.quota_period.is_some()
298 || self.quota_remaining_usd.is_some()
299 || self.quota_limit_usd.is_some()
300 || self.quota_used_usd.is_some()
301 || self.unlimited_quota == Some(true)
302 || self.total_used_usd.is_some()
303 || self.today_used_usd.is_some()
304 }
305
306 pub fn carry_forward_amount_data_from(&mut self, previous: &Self) {
307 if self.has_amount_data() {
308 return;
309 }
310 self.plan_name = previous.plan_name.clone();
311 self.total_balance_usd = previous.total_balance_usd.clone();
312 self.subscription_balance_usd = previous.subscription_balance_usd.clone();
313 self.paygo_balance_usd = previous.paygo_balance_usd.clone();
314 self.monthly_budget_usd = previous.monthly_budget_usd.clone();
315 self.monthly_spent_usd = previous.monthly_spent_usd.clone();
316 self.quota_period = previous.quota_period.clone();
317 self.quota_remaining_usd = previous.quota_remaining_usd.clone();
318 self.quota_limit_usd = previous.quota_limit_usd.clone();
319 self.quota_used_usd = previous.quota_used_usd.clone();
320 self.unlimited_quota = previous.unlimited_quota;
321 self.total_used_usd = previous.total_used_usd.clone();
322 self.today_used_usd = previous.today_used_usd.clone();
323 self.total_requests = previous.total_requests;
324 self.today_requests = previous.today_requests;
325 self.total_tokens = previous.total_tokens;
326 self.today_tokens = previous.today_tokens;
327 self.subscription_expires_at = previous.subscription_expires_at.clone();
328 self.usage_windows = previous.usage_windows.clone();
329 self.usage_rate = previous.usage_rate.clone();
330 self.usage_model_stats = previous.usage_model_stats.clone();
331 self.usage_alerts = previous.usage_alerts.clone();
332 }
333
334 pub fn amount_summary(&self) -> String {
335 let mut parts = Vec::new();
336 if let Some(plan) = self.plan_name.as_deref()
337 && !plan.trim().is_empty()
338 {
339 parts.push(format!("plan={plan}"));
340 }
341 if self.unlimited_quota == Some(true) {
342 parts.push("unlimited".to_string());
343 } else {
344 if let Some(total) = self.total_balance_usd.as_deref() {
345 parts.push(format!("total=${total}"));
346 }
347 if let Some(quota) = self.quota_summary() {
348 parts.push(quota);
349 }
350 match (
351 self.monthly_budget_usd.as_deref(),
352 self.monthly_spent_usd.as_deref(),
353 ) {
354 (Some(budget), Some(spent)) => {
355 if let Some(left) = left_from_budget_and_spent(budget, spent) {
356 parts.push(format!("left=${left} budget=${budget} spent=${spent}"));
357 } else {
358 parts.push(format!("budget=${budget} spent=${spent}"));
359 }
360 }
361 (Some(budget), None) => parts.push(format!("budget=${budget}")),
362 (None, Some(spent)) => parts.push(format!("used=${spent}")),
363 (None, None) => {}
364 }
365 if let Some(used) = self.total_used_usd.as_deref() {
366 parts.push(format!("used=${used}"));
367 }
368 if let Some(today) = self.today_used_usd.as_deref() {
369 parts.push(format!("today=${today}"));
370 }
371 if let Some(sub) = self.subscription_balance_usd.as_deref() {
372 parts.push(format!("sub=${sub}"));
373 }
374 if let Some(paygo) = self.paygo_balance_usd.as_deref() {
375 parts.push(format!("paygo=${paygo}"));
376 }
377 }
378 if let Some(requests) = self.total_requests {
379 parts.push(format!("req={requests}"));
380 }
381 if let Some(tokens) = self.total_tokens {
382 parts.push(format!("tok={tokens}"));
383 }
384 if !self.usage_alerts.is_empty() {
385 let alerts = self
386 .usage_alerts
387 .iter()
388 .map(|alert| alert.kind.as_str())
389 .collect::<Vec<_>>()
390 .join(",");
391 parts.push(format!("alerts={alerts}"));
392 }
393 if parts.is_empty() {
394 "-".to_string()
395 } else {
396 parts.join(" ")
397 }
398 }
399
400 fn quota_summary(&self) -> Option<String> {
401 let period = self
402 .quota_period
403 .as_deref()
404 .map(str::trim)
405 .filter(|value| !value.is_empty());
406 let remaining = self
407 .quota_remaining_usd
408 .as_deref()
409 .map(str::trim)
410 .filter(|value| !value.is_empty());
411 let limit = self
412 .quota_limit_usd
413 .as_deref()
414 .map(str::trim)
415 .filter(|value| !value.is_empty());
416 let used = self
417 .quota_used_usd
418 .as_deref()
419 .map(str::trim)
420 .filter(|value| !value.is_empty());
421
422 if remaining.is_none() && limit.is_none() && used.is_none() {
423 return None;
424 }
425
426 let mut parts = Vec::new();
427 let quota_label = match period {
428 Some("quota") | None => "quota".to_string(),
429 Some(period) => format!("{period} quota"),
430 };
431 parts.push(quota_label);
432
433 match (remaining, limit, used) {
434 (Some(remaining), Some(limit), Some(used)) => {
435 parts.push(format!("left=${remaining} limit=${limit} used=${used}"))
436 }
437 (Some(remaining), Some(limit), None) => {
438 parts.push(format!("left=${remaining} limit=${limit}"))
439 }
440 (Some(remaining), None, Some(used)) => {
441 parts.push(format!("left=${remaining} used=${used}"))
442 }
443 (Some(remaining), None, None) => parts.push(format!("left=${remaining}")),
444 (None, Some(limit), Some(used)) => parts.push(format!("used=${used} limit=${limit}")),
445 (None, Some(limit), None) => parts.push(format!("limit=${limit}")),
446 (None, None, Some(used)) => parts.push(format!("used=${used}")),
447 (None, None, None) => {}
448 }
449
450 Some(parts.join(" "))
451 }
452}
453
454fn left_from_budget_and_spent(budget: &str, spent: &str) -> Option<String> {
455 let budget = UsdAmount::from_decimal_str(budget)?;
456 let spent = UsdAmount::from_decimal_str(spent)?;
457 Some(budget.saturating_sub(spent).format_usd())
458}
459
460#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
461pub struct StationRoutingBalanceSummary {
462 pub snapshots: usize,
463 #[serde(default)]
464 pub ok: usize,
465 #[serde(default)]
466 pub exhausted: usize,
467 #[serde(default)]
468 pub stale: usize,
469 #[serde(default)]
470 pub error: usize,
471 #[serde(default)]
472 pub unknown: usize,
473 #[serde(default)]
474 pub routing_snapshots: usize,
475 #[serde(default)]
476 pub routing_exhausted: usize,
477 #[serde(default)]
478 pub routing_ignored_exhausted: usize,
479}
480
481impl StationRoutingBalanceSummary {
482 pub fn from_snapshots(snapshots: Option<&[ProviderBalanceSnapshot]>) -> Self {
483 let mut out = Self::default();
484 let Some(snapshots) = snapshots else {
485 return out;
486 };
487
488 for snapshot in snapshots {
489 out.record(snapshot, snapshot.status);
490 }
491 out
492 }
493
494 pub fn from_snapshot_iter_at<'a>(
495 snapshots: impl IntoIterator<Item = &'a ProviderBalanceSnapshot>,
496 now_ms: u64,
497 ) -> Self {
498 let mut out = Self::default();
499 for snapshot in snapshots {
500 out.record(snapshot, snapshot.status_at(now_ms));
501 }
502 out
503 }
504
505 fn record(&mut self, snapshot: &ProviderBalanceSnapshot, status: BalanceSnapshotStatus) {
506 self.snapshots += 1;
507 match status {
508 BalanceSnapshotStatus::Ok => self.ok += 1,
509 BalanceSnapshotStatus::Exhausted => self.exhausted += 1,
510 BalanceSnapshotStatus::Stale => self.stale += 1,
511 BalanceSnapshotStatus::Error => self.error += 1,
512 BalanceSnapshotStatus::Unknown => self.unknown += 1,
513 }
514 if snapshot.exhaustion_affects_routing {
515 self.routing_snapshots += 1;
516 if status == BalanceSnapshotStatus::Exhausted {
517 self.routing_exhausted += 1;
518 }
519 } else if status == BalanceSnapshotStatus::Exhausted {
520 self.routing_ignored_exhausted += 1;
521 }
522 }
523
524 pub fn is_empty(&self) -> bool {
525 self.snapshots == 0
526 }
527}
528
529fn default_exhaustion_affects_routing() -> bool {
530 true
531}
532
533fn bool_is_true(value: &bool) -> bool {
534 *value
535}
536
537#[cfg(test)]
538mod tests {
539 use super::*;
540
541 #[test]
542 fn balance_snapshot_status_labels_are_stable() {
543 assert_eq!(BalanceSnapshotStatus::Unknown.as_str(), "unknown");
544 assert_eq!(BalanceSnapshotStatus::Ok.as_str(), "ok");
545 assert_eq!(BalanceSnapshotStatus::Exhausted.as_str(), "exhausted");
546 assert_eq!(BalanceSnapshotStatus::Stale.as_str(), "stale");
547 assert_eq!(BalanceSnapshotStatus::Error.as_str(), "error");
548 }
549
550 #[test]
551 fn provider_balance_amount_summary_formats_known_amounts() {
552 let snapshot = ProviderBalanceSnapshot {
553 plan_name: Some("monthly".to_string()),
554 total_balance_usd: Some("3.5".to_string()),
555 monthly_budget_usd: Some("5".to_string()),
556 monthly_spent_usd: Some("1.25".to_string()),
557 total_used_usd: Some("7".to_string()),
558 today_used_usd: Some("0.5".to_string()),
559 subscription_balance_usd: Some("2".to_string()),
560 paygo_balance_usd: Some("1.5".to_string()),
561 total_requests: Some(42),
562 total_tokens: Some(1234),
563 ..Default::default()
564 };
565
566 assert_eq!(
567 snapshot.amount_summary(),
568 "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"
569 );
570 }
571
572 #[test]
573 fn provider_balance_amount_summary_keeps_wallet_with_quota() {
574 let snapshot = ProviderBalanceSnapshot {
575 plan_name: Some("rightcode".to_string()),
576 total_balance_usd: Some("3.25".to_string()),
577 quota_period: Some("daily".to_string()),
578 quota_remaining_usd: Some("7.5".to_string()),
579 quota_limit_usd: Some("20".to_string()),
580 quota_used_usd: Some("12.5".to_string()),
581 ..Default::default()
582 };
583
584 assert_eq!(
585 snapshot.amount_summary(),
586 "plan=rightcode total=$3.25 daily quota left=$7.5 limit=$20 used=$12.5"
587 );
588 }
589
590 #[test]
591 fn provider_balance_amount_summary_prioritizes_unlimited_quota() {
592 let snapshot = ProviderBalanceSnapshot {
593 plan_name: Some("cx".to_string()),
594 unlimited_quota: Some(true),
595 quota_used_usd: Some("106065.94".to_string()),
596 ..Default::default()
597 };
598
599 assert_eq!(snapshot.amount_summary(), "plan=cx unlimited");
600 }
601
602 #[test]
603 fn routing_exhausted_respects_snapshot_routing_flag() {
604 let mut snapshot = ProviderBalanceSnapshot {
605 exhausted: Some(true),
606 exhaustion_affects_routing: false,
607 ..Default::default()
608 };
609 snapshot.refresh_status(100);
610
611 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
612 assert!(!snapshot.routing_exhausted());
613 assert!(snapshot.routing_ignored_exhaustion());
614 }
615
616 #[test]
617 fn provider_balance_amount_summary_includes_usage_alerts() {
618 let snapshot = ProviderBalanceSnapshot {
619 usage_alerts: vec![ProviderUsageAlert {
620 kind: ProviderUsageAlertKind::LowBalance,
621 message: "remaining balance is low".to_string(),
622 }],
623 ..Default::default()
624 };
625
626 assert_eq!(snapshot.amount_summary(), "alerts=low_balance");
627 }
628}