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