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